"""Kimi Delta Attention (KDA) forward, chunk form, RTX PRO 6000 (SM120). Pipeline (per (b,h)): Kernel A (Triton, parallel over chunks): in-chunk cumsum of g, decay factor application, K-K interaction, exact unit-triangular inverse in WY form (Neumann factorization), w/u sides, decayed q, chunk-end-decayed k, causal A_qk. Kernel B (Triton, sequential over chunks, V-tiled): delta-rule state recurrence with fp32 state in registers. Emits h^T and vn^T slices. Kernel C (Triton, parallel over chunks): o = qd @ h + Aqk @ vn. bf16 tensor-core dots with fp32 accumulation everywhere in the recurrence. """ from __future__ import annotations import os import torch import torch.nn as nn import triton import triton.language as tl os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0a") # =========================================================================== # Kernel A: intra-chunk # =========================================================================== @triton.jit def _inv32_strict(L): """Invert (I + L) for strictly-lower-triangular L (32,32) via exact Neumann factorization: (I+L)^{-1} = (I-L)(I+L^2)(I+L^4)(I+L^8)(I+L^16) (L^32 = 0). fp32 (ieee) dots -- same values as forward substitution.""" rows = tl.arange(0, 32) eye = tl.where(rows[:, None] == rows[None, :], 1.0, 0.0).to(tl.float32) L2 = tl.dot(L, L, input_precision="ieee") L4 = tl.dot(L2, L2, input_precision="ieee") L8 = tl.dot(L4, L4, input_precision="ieee") L16 = tl.dot(L8, L8, input_precision="ieee") P = (eye - L) + L2 - tl.dot(L, L2, input_precision="ieee") P = P + tl.dot(P, L4, input_precision="ieee") P = P + tl.dot(P, L8, input_precision="ieee") P = P + tl.dot(P, L16, input_precision="ieee") return P @triton.jit def kda_intra_kernel( q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr, w_ptr, u_ptr, qd_ptr, aqk_ptr, kd_ptr, egl_ptr, scale, T, H, NT, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, HB: tl.constexpr, ): i_t = tl.program_id(0) i_bh = tl.program_id(1) i_b = i_bh // H i_h = i_bh % H t0 = i_t * BT r = tl.arange(0, HB) dK = tl.arange(0, K) row0 = t0 + r row1 = t0 + HB + r off0 = ((i_b * T + row0) * H + i_h)[:, None] * K + dK[None, :] off1 = ((i_b * T + row1) * H + i_h)[:, None] * K + dK[None, :] offb0 = (i_b * T + row0) * H + i_h offb1 = (i_b * T + row1) * H + i_h cbase = (i_bh * NT + i_t) * BT off_o0 = (cbase + r)[:, None] * K + dK[None, :] off_o1 = (cbase + HB + r)[:, None] * K + dK[None, :] # ---- decay factors (fp32 cumsum, eager bf16 conversion) ---- g0 = tl.load(g_ptr + off0) gc0 = tl.cumsum(g0, axis=0) s0 = tl.sum(g0, axis=0) g1 = tl.load(g_ptr + off1) gc1 = tl.cumsum(g1, axis=0) + s0[None, :] glast = s0 + tl.sum(g1, axis=0) egl = tl.exp(glast) tl.store(egl_ptr + (i_bh * NT + i_t) * K + dK, egl) e0 = tl.exp(gc0) e1 = tl.exp(gc1) k0 = tl.load(k_ptr + off0) k1 = tl.load(k_ptr + off1) kb0 = (k0.to(tl.float32) * e0).to(tl.bfloat16) kb1 = (k1.to(tl.float32) * e1).to(tl.bfloat16) kf0 = (k0.to(tl.float32) * tl.exp(-gc0)).to(tl.bfloat16) kf1 = (k1.to(tl.float32) * tl.exp(-gc1)).to(tl.bfloat16) b0 = tl.load(beta_ptr + offb0).to(tl.float32) b1 = tl.load(beta_ptr + offb1).to(tl.float32) # kd = k*exp(glast - gc) = kf * egl (early store, pre-transposed (K,BT) # layout so the state kernel needs no per-iteration transpose) kd0 = (kf0.to(tl.float32) * egl[None, :]).to(tl.bfloat16) kd1 = (kf1.to(tl.float32) * egl[None, :]).to(tl.bfloat16) kdT_base = (i_bh * NT + i_t) * BT * K tl.store(kd_ptr + kdT_base + r[:, None] + dK[None, :] * BT, kd0) tl.store(kd_ptr + kdT_base + (HB + r)[:, None] + dK[None, :] * BT, kd1) # ---- decayed q (early store) ---- q0 = tl.load(q_ptr + off0) q1 = tl.load(q_ptr + off1) qd0 = (q0.to(tl.float32) * e0 * scale).to(tl.bfloat16) qd1 = (q1.to(tl.float32) * e1 * scale).to(tl.bfloat16) tl.store(qd_ptr + off_o0, qd0) tl.store(qd_ptr + off_o1, qd1) # ---- Aqk (causal, quadrant dots; strictly-upper quadrant left unwritten) ---- causal = r[:, None] >= r[None, :] aq00 = tl.where(causal, tl.dot(qd0, tl.trans(kf0)), 0.0).to(tl.bfloat16) aq10 = tl.dot(qd1, tl.trans(kf0)).to(tl.bfloat16) aq11 = tl.where(causal, tl.dot(qd1, tl.trans(kf1)), 0.0).to(tl.bfloat16) abase = (i_bh * NT + i_t) * BT * BT tl.store(aqk_ptr + abase + r[:, None] * BT + r[None, :], aq00) tl.store(aqk_ptr + abase + BT // 2 * BT + r[:, None] * BT + r[None, :], aq10) tl.store(aqk_ptr + abase + BT // 2 * BT + BT // 2 + r[:, None] * BT + r[None, :], aq11) # ---- A = strictly-lower K-K interaction, inverse in WY form ---- strict = r[:, None] > r[None, :] L00 = tl.where(strict, tl.dot(kb0, tl.trans(kf0)) * b0[:, None], 0.0) L11 = tl.where(strict, tl.dot(kb1, tl.trans(kf1)) * b1[:, None], 0.0) L10 = tl.dot(kb1, tl.trans(kf0)) * b1[:, None] M00 = _inv32_strict(L00) M11 = _inv32_strict(L11) C10 = -tl.dot(tl.dot(M11, L10), M00) M00b = M00.to(tl.bfloat16) M11b = M11.to(tl.bfloat16) C10b = C10.to(tl.bfloat16) # ---- w = M @ (beta * exp(g) * k); u = M @ (beta * v) ---- rw0 = (kb0.to(tl.float32) * b0[:, None]).to(tl.bfloat16) rw1 = (kb1.to(tl.float32) * b1[:, None]).to(tl.bfloat16) tl.store(w_ptr + off_o0, tl.dot(M00b, rw0).to(tl.bfloat16)) tl.store(w_ptr + off_o1, (tl.dot(C10b, rw0) + tl.dot(M11b, rw1)).to(tl.bfloat16)) v0 = tl.load(v_ptr + off0) v1 = tl.load(v_ptr + off1) ru0 = (v0.to(tl.float32) * b0[:, None]).to(tl.bfloat16) ru1 = (v1.to(tl.float32) * b1[:, None]).to(tl.bfloat16) tl.store(u_ptr + off_o0, tl.dot(M00b, ru0).to(tl.bfloat16)) tl.store(u_ptr + off_o1, (tl.dot(C10b, ru0) + tl.dot(M11b, ru1)).to(tl.bfloat16)) # =========================================================================== # Kernel B: Triton inter-chunk recurrence (fallback path) # =========================================================================== @triton.jit def kda_state_kernel( w_ptr, u_ptr, kd_ptr, egl_ptr, ht_ptr, vnt_ptr, H, NT, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, ): i_v = tl.program_id(0) i_bh = tl.program_id(1) rT = tl.arange(0, BT) dK = tl.arange(0, K) dV = i_v * BV + tl.arange(0, BV) S = tl.zeros((K, BV), dtype=tl.float32) for i in range(NT): hbase = (i_bh * NT + i) * V * K tl.store(ht_ptr + hbase + dV[None, :] * K + dK[:, None], S.to(tl.bfloat16)) cbase = (i_bh * NT + i) * BT w_i = tl.load(w_ptr + (cbase + rT)[:, None] * K + dK[None, :]) u_i = tl.load(u_ptr + (cbase + rT)[:, None] * V + dV[None, :]) vn = u_i.to(tl.float32) - tl.dot(w_i, S.to(tl.bfloat16)) tl.store(vnt_ptr + (i_bh * NT + i) * V * BT + dV[None, :] * BT + rT[:, None], vn.to(tl.bfloat16)) kd_i = tl.load(kd_ptr + (i_bh * NT + i) * BT * K + dK[:, None] * BT + rT[None, :]) egl = tl.load(egl_ptr + (i_bh * NT + i) * K + dK) S = S * egl[:, None] + tl.dot(kd_i, vn.to(tl.bfloat16)) # =========================================================================== # Kernel C: output # =========================================================================== @triton.jit def kda_output_kernel( qd_ptr, aqk_ptr, ht_ptr, vnt_ptr, o_ptr, T, H, NT, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, ): i_t = tl.program_id(0) i_bh = tl.program_id(1) i_b = i_bh // H i_h = i_bh % H t0 = i_t * BT rT = tl.arange(0, BT) dK = tl.arange(0, K) dV = tl.arange(0, V) cbase = (i_bh * NT + i_t) * BT qd = tl.load(qd_ptr + (cbase + rT)[:, None] * K + dK[None, :]) # causal mask applied here: the intra kernel leaves the strictly-upper # quadrant of aqk unwritten. Aqk = tl.load(aqk_ptr + (i_bh * NT + i_t) * BT * BT + rT[:, None] * BT + rT[None, :]).to(tl.float32) causal = rT[:, None] >= rT[None, :] Aqk = tl.where(causal, Aqk, 0.0).to(tl.bfloat16) # h^T buffer (V rows, K cols); need (K, V) hT = tl.load(ht_ptr + (i_bh * NT + i_t) * V * K + dV[:, None] * K + dK[None, :]) vnT = tl.load(vnt_ptr + (i_bh * NT + i_t) * V * BT + dV[:, None] * BT + rT[None, :]) o = tl.dot(qd, tl.trans(hT)) + tl.dot(Aqk, tl.trans(vnT)) o_off = ((i_b * T + t0 + rT) * H + i_h)[:, None] * V + dV[None, :] tl.store(o_ptr + o_off, o.to(tl.bfloat16)) # =========================================================================== # Host wrapper # =========================================================================== class Model(nn.Module): def __init__(self, B: int, T: int, H: int, K: int, V: int, chunk_size: int = 64): super().__init__() self.B, self.T, self.H, self.K, self.V = B, T, H, K, V self.chunk_size = chunk_size self.scale = float(K) ** -0.5 self.register_buffer("_dummy", torch.zeros(1), persistent=False) assert K == 128 and V == 128 and chunk_size == 64, "kernel specialized for K=V=128, BT=64" assert T % chunk_size == 0 self.NT = T // chunk_size self._scratch = None self._cuda = None def _alloc(self, device): B, T, H, K, V, NT = self.B, self.T, self.H, self.K, self.V, self.NT BT = self.chunk_size f = dict(device=device, dtype=torch.bfloat16) return dict( w=torch.empty((B * H * NT, BT, K), **f), u=torch.empty((B * H * NT, BT, V), **f), qd=torch.empty((B * H * NT, BT, K), **f), aqk=torch.empty((B * H * NT, BT, BT), **f), kd=torch.empty((B * H * NT, BT, K), **f), egl=torch.empty((B * H * NT, K), device=device, dtype=torch.float32), ht=torch.empty((B * H * NT, V, K), **f), vnt=torch.empty((B * H * NT, V, BT), **f), o=torch.empty((B, T, H, V), **f), ) def forward(self, q, k, v, g, beta): B, T, H, K, V, NT = self.B, self.T, self.H, self.K, self.V, self.NT BT = self.chunk_size BH = B * H dev = q.device if self._scratch is None or self._scratch["w"].device != dev: self._scratch = self._alloc(dev) s = self._scratch o = s["o"] kda_intra_kernel[(NT, BH)]( q, k, v, g, beta, s["w"], s["u"], s["qd"], s["aqk"], s["kd"], s["egl"], self.scale, T, H, NT, K=K, V=V, BT=BT, HB=BT // 2, num_warps=8, num_stages=3, ) BV = 16 kda_state_kernel[(V // BV, BH)]( s["w"], s["u"], s["kd"], s["egl"], s["ht"], s["vnt"], H, NT, K=K, V=V, BT=BT, BV=BV, num_warps=4, num_stages=3, ) kda_output_kernel[(NT, BH)]( s["qd"], s["aqk"], s["ht"], s["vnt"], o, T, H, NT, K=K, V=V, BT=BT, num_warps=8, num_stages=2, ) return o B = 2 T = 1024 H = 8 K = 128 V = 128 CHUNK_SIZE = 64 def get_inputs(): torch.manual_seed(0) q = torch.randn(B, T, H, K, dtype=torch.bfloat16) * 0.1 k = torch.randn(B, T, H, K, dtype=torch.bfloat16) * 0.1 v = torch.randn(B, T, H, V, dtype=torch.bfloat16) * 0.1 g = (torch.randn(B, T, H, K, dtype=torch.float32) * 0.1 - 0.05) beta = torch.sigmoid(torch.randn(B, T, H, dtype=torch.bfloat16)) return [q, k, v, g, beta] def get_init_inputs(): return [B, T, H, K, V, CHUNK_SIZE]