"""Kimi Delta Attention (KDA) forward, chunk form -- custom Triton kernels. Reimplements the chunk-parallel gated delta rule (WY representation) as a self-contained set of kernels. 1. chunk_local_kernel (parallel over B*H*NT chunks): in-chunk cumsum of g; intra-chunk K-K matrix M (strictly lower) and its inverse R=(I-M)^-1 via the doubling product (I+M)(I+M^2)...(I+M^32); w = Aw @ (exp(gc)*k), u = Aw @ v where Aw = R * beta[col]. Also folds the inter-chunk recurrence into an affine map S_{n+1} = P_n S_n + Q_n with P_n = diag(exp(gc_last)) - kd^T w, Q_n = kd^T u, kd = exp(gc_last - gc) * k. 2. scan_kernel (sequential over chunks, parallel over B*H and V-blocks): one matmul per step: S_{n+1} = P_n @ S_n + Q_n; stores entry states S_n. 3. output_kernel (fully parallel over B*H*NT and V-blocks): v_i = u - w @ S_n; o_n = (q*exp(gc)) @ S_n + tril(Aqk) @ v_i. """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl # Launch configs (tuned for SM90 / H100, K=V=128, chunk=64). _SCAN_BV = 16 # V-block for the sequential inter-chunk scan _SCAN_WARPS = 8 _SCAN_STAGES = 3 _OUT_BV = 128 # V-block for the parallel output kernel _OUT_WARPS = 4 _OUT_STAGES = 1 _INV_ITERS = tl.constexpr(5) # doubling steps: (I+M)..(I+M^32) exact for 64x64 _CL_WARPS = 4 _CL_STAGES = 2 @triton.jit def chunk_local_kernel( k_ptr, v_ptr, g_ptr, beta_ptr, w_ptr, u_ptr, gc_ptr, P_ptr, Q_ptr, B, T, H, NT, s_kb, s_kt, s_kh, s_kd, # k / g / w strides (B,T,H,K) s_vb, s_vt, s_vh, s_vd, # v / u strides (B,T,H,V) s_bb, s_bt, s_bh, # beta strides (B,T,H) s_pn, s_qn, # P / Q chunk strides within one bh BT: tl.constexpr, K: tl.constexpr, V: tl.constexpr, ): pid = tl.program_id(0) n = pid % NT bh = pid // NT h = bh % H b = bh // H t0 = n * BT rc = tl.arange(0, BT) rk = tl.arange(0, K) rv = tl.arange(0, V) base_k = b * s_kb + t0 * s_kt + h * s_kh base_v = b * s_vb + t0 * s_vt + h * s_vh kk = rc[:, None] * s_kt + rk[None, :] * s_kd vv = rc[:, None] * s_vt + rv[None, :] * s_vd k = tl.load(k_ptr + base_k + kk).to(tl.float32) g = tl.load(g_ptr + base_k + kk).to(tl.float32) v = tl.load(v_ptr + base_v + vv).to(tl.float32) beta = tl.load(beta_ptr + b * s_bb + (t0 + rc) * s_bt + h * s_bh).to(tl.float32) gc = tl.cumsum(g, axis=0) # (BT, K) eg = tl.exp(gc) kc = k * eg # exp(gc)*k kr = k * tl.exp(-gc) # exp(-gc)*k Araw = tl.dot(kc.to(tl.bfloat16), tl.trans(kr).to(tl.bfloat16)) # (BT,BT) col = rc[None, :] row = rc[:, None] M = tl.where(col < row, -beta[:, None] * Araw, 0.0) eye = (row == col).to(tl.float32) R = eye + M cur = M for _ in range(_INV_ITERS): cur = tl.dot(cur.to(tl.bfloat16), cur.to(tl.bfloat16)) R = R + tl.dot(R.to(tl.bfloat16), cur.to(tl.bfloat16)) Aw = R * beta[None, :] # multiply column j by beta[j] Aw_b = Aw.to(tl.bfloat16) w = tl.dot(Aw_b, kc.to(tl.bfloat16)) # (BT,K) u = tl.dot(Aw_b, v.to(tl.bfloat16)) # (BT,V) tl.store(w_ptr + base_k + kk, w.to(w_ptr.dtype.element_ty)) tl.store(u_ptr + base_v + vv, u.to(u_ptr.dtype.element_ty)) tl.store(gc_ptr + base_k + kk, gc.to(gc_ptr.dtype.element_ty)) # ---- affine inter-chunk transition (P_n, Q_n) ---- gc_last = tl.sum(tl.where(row == (BT - 1), gc, 0.0), axis=0) # (K,) a = tl.exp(gc_last) kd = tl.exp(gc_last[None, :] - gc) * k # (BT,K) kdT = tl.trans(kd).to(tl.bfloat16) # (K,BT) P = -tl.dot(kdT, w.to(tl.bfloat16)) # (K,K) P = P + tl.where(rk[:, None] == rk[None, :], a[:, None], 0.0) Q = tl.dot(kdT, u.to(tl.bfloat16)) # (K,V) p_base = bh * (NT * s_pn) + n * s_pn + rk[:, None] * K + rk[None, :] q_base = bh * (NT * s_qn) + n * s_qn + rk[:, None] * V + rv[None, :] tl.store(P_ptr + p_base, P.to(P_ptr.dtype.element_ty)) tl.store(Q_ptr + q_base, Q.to(Q_ptr.dtype.element_ty)) @triton.jit def scan_kernel( P_ptr, Q_ptr, S_ptr, B, H, NT, s_pn, s_qn, s_sn, s_sk, s_sv, K: tl.constexpr, V: tl.constexpr, BV: tl.constexpr, ): pid_bh = tl.program_id(0) pid_v = tl.program_id(1) v0 = pid_v * BV rk = tl.arange(0, K) rk2 = tl.arange(0, K) rv = tl.arange(0, BV) S = tl.zeros((K, BV), dtype=tl.float32) p_off = pid_bh * (NT * s_pn) + rk[:, None] * K + rk2[None, :] q_off = pid_bh * (NT * s_qn) + rk[:, None] * V + (v0 + rv)[None, :] s_off = pid_bh * (NT * s_sn) + rk[:, None] * s_sk + (v0 + rv)[None, :] * s_sv for n in range(NT): tl.store(S_ptr + s_off + n * s_sn, S.to(S_ptr.dtype.element_ty)) P = tl.load(P_ptr + p_off + n * s_pn) # (K,K) Q = tl.load(Q_ptr + q_off + n * s_qn).to(tl.float32) # (K,BV) S = tl.dot(P, S.to(P.dtype)) + Q @triton.jit def output_kernel( q_ptr, k_ptr, gc_ptr, w_ptr, u_ptr, S_ptr, o_ptr, B, T, H, NT, scale, s_kb, s_kt, s_kh, s_kd, s_vb, s_vt, s_vh, s_vd, s_sn, s_sk, s_sv, BT: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BV: tl.constexpr, ): pid = tl.program_id(0) pid_v = tl.program_id(1) n = pid % NT bh = pid // NT h = bh % H b = bh // H v0 = pid_v * BV t0 = n * BT rc = tl.arange(0, BT) rk = tl.arange(0, K) rv = tl.arange(0, BV) base_k = b * s_kb + t0 * s_kt + h * s_kh base_v = b * s_vb + t0 * s_vt + h * s_vh kk = rc[:, None] * s_kt + rk[None, :] * s_kd vv = rc[:, None] * s_vt + (v0 + rv)[None, :] * s_vd s_base = bh * (NT * s_sn) + n * s_sn + rk[:, None] * s_sk + (v0 + rv)[None, :] * s_sv q = tl.load(q_ptr + base_k + kk).to(tl.float32) * scale k = tl.load(k_ptr + base_k + kk).to(tl.float32) gc = tl.load(gc_ptr + base_k + kk).to(tl.float32) w = tl.load(w_ptr + base_k + kk).to(tl.bfloat16) u = tl.load(u_ptr + base_v + vv).to(tl.float32) S = tl.load(S_ptr + s_base) eg = tl.exp(gc) qc = (q * eg).to(tl.bfloat16) kr = (k * tl.exp(-gc)).to(tl.bfloat16) Sb = S.to(tl.bfloat16) Aqk = tl.dot(qc, tl.trans(kr)) Aqk = tl.where(rc[None, :] <= rc[:, None], Aqk, 0.0) v_i = u - tl.dot(w, Sb) o = tl.dot(qc, Sb) + tl.dot(Aqk.to(tl.bfloat16), v_i.to(tl.bfloat16)) tl.store(o_ptr + base_v + vv, o.to(o_ptr.dtype.element_ty)) def _kda_forward(q, k, v, g, beta, scale, chunk_size=64): B, T, H, K = q.shape V = v.shape[-1] BT = chunk_size NT = T // BT assert T % BT == 0 BH = B * H q = q.contiguous(); k = k.contiguous(); v = v.contiguous() g = g.contiguous(); beta = beta.contiguous() w = torch.empty(B, T, H, K, dtype=torch.bfloat16, device=q.device) u = torch.empty(B, T, H, V, dtype=torch.bfloat16, device=q.device) gc = torch.empty(B, T, H, K, dtype=torch.bfloat16, device=q.device) P = torch.empty(BH, NT, K, K, dtype=torch.bfloat16, device=q.device) Q = torch.empty(BH, NT, K, V, dtype=torch.bfloat16, device=q.device) Sstate = torch.empty(BH, NT, K, V, dtype=torch.bfloat16, device=q.device) o = torch.empty(B, T, H, V, dtype=torch.bfloat16, device=q.device) s_kb, s_kt, s_kh, s_kd = k.stride() s_vb, s_vt, s_vh, s_vd = v.stride() s_bb, s_bt, s_bh = beta.stride() s_sn, s_sk, s_sv = Sstate.stride()[1:] s_pn = P.stride()[1] s_qn = Q.stride()[1] chunk_local_kernel[(BH * NT,)]( k, v, g, beta, w, u, gc, P, Q, B, T, H, NT, s_kb, s_kt, s_kh, s_kd, s_vb, s_vt, s_vh, s_vd, s_bb, s_bt, s_bh, s_pn, s_qn, BT=BT, K=K, V=V, num_warps=_CL_WARPS, num_stages=_CL_STAGES, ) scan_kernel[(BH, V // _SCAN_BV)]( P, Q, Sstate, B, H, NT, s_pn, s_qn, s_sn, s_sk, s_sv, K=K, V=V, BV=_SCAN_BV, num_warps=_SCAN_WARPS, num_stages=_SCAN_STAGES, ) BVO = _OUT_BV output_kernel[(BH * NT, V // BVO)]( q, k, gc, w, u, Sstate, o, B, T, H, NT, scale, s_kb, s_kt, s_kh, s_kd, s_vb, s_vt, s_vh, s_vd, s_sn, s_sk, s_sv, BT=BT, K=K, V=V, BV=BVO, num_warps=_OUT_WARPS, num_stages=_OUT_STAGES, ) return o class Model(nn.Module): def __init__(self, B, T, H, K, V, chunk_size=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) def forward(self, q, k, v, g, beta): return _kda_forward(q, k, v, g, beta, self.scale, self.chunk_size) 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]