"""Kimi Delta Attention forward (chunk form) — custom kernel implementation. Pipeline (four Triton kernels): 1. kda_prep_kernel: per (batch*head, chunk) — Phase A (K-block loop): chunk-local cumsum of g, A_s = -tri0(beta_row * (k e^g)(k e^-g)^T) accumulated, F = tri1(scale * (q e^g)(k e^-g)^T) accumulated, stores of qg = scale*q*e^g, k~ = k*e^{g_last - g}, gl = log2(g_last..) per-chunk decay in the log domain, and a (beta*e^g*k) scratch for phase W. Solve: X = (I - A_s)^{-1} via 16x16 blocked forward substitution using exact nilpotent doubling (A_s^16 = 0 on each diagonal block). Phase W: w = X @ (beta*e^g*k), u = X @ (beta*v), P = qg - F @ w (fixup matrix for the grouped scan). 2. scan_grouped: per (batch*head, group, V-tile) — sequential local scan over the group's chunks with zero entry state; writes o_local and the group-exit state F_g (bf16). 3. combine_states: per (batch*head, V-tile) — tiny sequential pass over groups: E_{g+1} = Dg_g * E_g + F_g (entry state of each group). 4. fixup_output: per (batch*head, chunk, V-tile) for chunks not in group 0: o += (P_c * D_c[None,:]) @ E_g, D_c = prod of chunk decays from group entry. Derivation of the group correction (entry state E of group g): S_c = S_loc_c + D_c x E => o - o_loc = (qg - F w) @ (D_c x E). The recurrence X = (I - A_s)^{-1} with A_s = L(beta * (k e^g)(k e^-g)^T) (L = strictly lower part, beta scaling ROWS like the reference's `A * beta[..., None]`) matches the reference forward-substitution exactly: M = A_s + A_s @ M <=> M + I = (I - A_s)^{-1}. """ import torch import triton import triton.language as tl _LOG2E = tl.constexpr(1.4426950408889634) @triton.jit def _solve16(a_ptr, base, stride, o_i, m_I, invert_diag: tl.constexpr): """X_nn = (I - A_nn)^{-1} for a strictly-lower BC x BC block stored at a_ptr + base (row stride = stride). A_nn is strictly lower triangular with BC = 16, so A_nn^16 = 0 and (I - A_nn)^{-1} = sum_{p=0}^{15} A_nn^p = prod_{j=0}^{3} (I + A_nn^{2^j}). Exact nilpotent doubling: 6 small GEMMs, no serial dependency chain.""" b_A = tl.load(a_ptr + base + o_i[:, None] * stride + o_i[None, :]) b_A = tl.where(o_i[:, None] > o_i[None, :], b_A, 0.0) b_A2 = tl.dot(b_A, b_A, input_precision="tf32") b_A4 = tl.dot(b_A2, b_A2, input_precision="tf32") b_A8 = tl.dot(b_A4, b_A4, input_precision="tf32") b_P = tl.where(m_I, 1.0 + 0.0, 0.0) + b_A b_P = tl.dot(b_P, tl.where(m_I, 1.0, 0.0) + b_A2, input_precision="tf32") b_P = tl.dot(b_P, tl.where(m_I, 1.0, 0.0) + b_A4, input_precision="tf32") b_P = tl.dot(b_P, tl.where(m_I, 1.0, 0.0) + b_A8, input_precision="tf32") return b_P @triton.jit def kda_prep_kernel( q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr, w_ptr, u_ptr, qg_ptr, kk_ptr, f_ptr, gl_ptr, p_ptr, asave_ptr, x_ptr, kh_ptr, scale, T, NT, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BK: tl.constexpr, ): i_t = tl.program_id(0) i_bh = tl.program_id(1) i_b = i_bh // H i_h = i_bh % H o_r = tl.arange(0, BT) o_i = tl.arange(0, BC) o_bk = tl.arange(0, BK) m_I = o_i[:, None] == o_i[None, :] t0 = i_b * T + i_t * BT row_in = t0 + o_r nch = i_bh * NT + i_t b_beta = tl.load(beta_ptr + row_in * H + i_h).to(tl.float32) # ---------------- Phase A: A_s / F accumulation + qg/kk/gl/kh stores --- b_A = tl.zeros((BT, BT), dtype=tl.float32) b_F = tl.zeros((BT, BT), dtype=tl.float32) for kb in range(0, K, BK): p_gk = g_ptr + (row_in[:, None] * H + i_h) * K + kb + o_bk[None, :] p_kk = k_ptr + (row_in[:, None] * H + i_h) * K + kb + o_bk[None, :] p_qq = q_ptr + (row_in[:, None] * H + i_h) * K + kb + o_bk[None, :] b_g = tl.load(p_gk) * _LOG2E b_gc = tl.cumsum(b_g, axis=0) b_eg = tl.math.exp2(b_gc) b_egm = tl.math.exp2(-b_gc) b_k = tl.load(p_kk).to(tl.float32) b_q = tl.load(p_qq).to(tl.float32) b_kh = b_k * b_eg # k e^gamma b_kt = b_k * b_egm # k e^-gamma b_A += tl.dot(b_kh, tl.trans(b_kt), input_precision="tf32") b_F += tl.dot(b_q * b_eg, tl.trans(b_kt), input_precision="tf32") b_gl = tl.sum(b_g, axis=0) # log2 total decay of this K-block tl.store(gl_ptr + nch * K + kb + o_bk, b_gl) tl.store(qg_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :], (b_q * b_eg * scale).to(tl.bfloat16)) tl.store(kk_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :], (b_k * tl.math.exp2(b_gl[None, :] - b_gc)).to(tl.bfloat16)) tl.store(kh_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :], (b_kh * b_beta[:, None]).to(tl.bfloat16)) b_A = -b_A * b_beta[:, None] b_A = tl.where(o_r[:, None] > o_r[None, :], b_A, 0.0) a_base = asave_ptr + nch * (BT * BT) tl.store(a_base + o_r[:, None] * BT + o_r[None, :], b_A) b_F = b_F * scale b_F = tl.where(o_r[:, None] >= o_r[None, :], b_F, 0.0) tl.store(f_ptr + nch * (BT * BT) + o_r[:, None] * BT + o_r[None, :], b_F.to(tl.bfloat16)) # ---------------- Solve: X = (I - A_s)^{-1} --------------------------- b_X00 = _solve16(asave_ptr, nch * (BT * BT), BT, o_i, m_I, True) b_X11 = _solve16(asave_ptr, nch * (BT * BT) + BC * BT + BC, BT, o_i, m_I, True) b_X22 = _solve16(asave_ptr, nch * (BT * BT) + 2 * BC * BT + 2 * BC, BT, o_i, m_I, True) b_X33 = _solve16(asave_ptr, nch * (BT * BT) + 3 * BC * BT + 3 * BC, BT, o_i, m_I, True) b_A10 = tl.load(a_base + (BC + o_i)[:, None] * BT + o_i[None, :]) b_X10 = tl.dot(tl.dot(b_X11, b_A10, input_precision="tf32"), b_X00, input_precision="tf32") b_A20 = tl.load(a_base + (2 * BC + o_i)[:, None] * BT + o_i[None, :]) b_A21 = tl.load(a_base + (2 * BC + o_i)[:, None] * BT + BC + o_i[None, :]) b_X20 = tl.dot(b_X22, tl.dot(b_A20, b_X00, input_precision="tf32") + tl.dot(b_A21, b_X10, input_precision="tf32"), input_precision="tf32") b_A30 = tl.load(a_base + (3 * BC + o_i)[:, None] * BT + o_i[None, :]) b_A31 = tl.load(a_base + (3 * BC + o_i)[:, None] * BT + BC + o_i[None, :]) b_A32 = tl.load(a_base + (3 * BC + o_i)[:, None] * BT + 2 * BC + o_i[None, :]) b_X30 = tl.dot(b_X33, tl.dot(b_A30, b_X00, input_precision="tf32") + tl.dot(b_A31, b_X10, input_precision="tf32") + tl.dot(b_A32, b_X20, input_precision="tf32"), input_precision="tf32") b_X21 = tl.dot(tl.dot(b_X22, b_A21, input_precision="tf32"), b_X11, input_precision="tf32") b_X31 = tl.dot(b_X33, tl.dot(b_A31, b_X11, input_precision="tf32") + tl.dot(b_A32, b_X21, input_precision="tf32"), input_precision="tf32") b_X32 = tl.dot(tl.dot(b_X33, b_A32, input_precision="tf32"), b_X22, input_precision="tf32") pX = x_ptr + nch * (BT * BT) tl.store(pX + (0 * BC + o_i)[:, None] * BT + (0 * BC + o_i)[None, :], b_X00.to(tl.bfloat16)) tl.store(pX + (1 * BC + o_i)[:, None] * BT + (0 * BC + o_i)[None, :], b_X10.to(tl.bfloat16)) tl.store(pX + (1 * BC + o_i)[:, None] * BT + (1 * BC + o_i)[None, :], b_X11.to(tl.bfloat16)) tl.store(pX + (2 * BC + o_i)[:, None] * BT + (0 * BC + o_i)[None, :], b_X20.to(tl.bfloat16)) tl.store(pX + (2 * BC + o_i)[:, None] * BT + (1 * BC + o_i)[None, :], b_X21.to(tl.bfloat16)) tl.store(pX + (2 * BC + o_i)[:, None] * BT + (2 * BC + o_i)[None, :], b_X22.to(tl.bfloat16)) tl.store(pX + (3 * BC + o_i)[:, None] * BT + (0 * BC + o_i)[None, :], b_X30.to(tl.bfloat16)) tl.store(pX + (3 * BC + o_i)[:, None] * BT + (1 * BC + o_i)[None, :], b_X31.to(tl.bfloat16)) tl.store(pX + (3 * BC + o_i)[:, None] * BT + (2 * BC + o_i)[None, :], b_X32.to(tl.bfloat16)) tl.store(pX + (3 * BC + o_i)[:, None] * BT + (3 * BC + o_i)[None, :], b_X33.to(tl.bfloat16)) tl.debug_barrier() # --------------- Phase W: w = X@kh, u = X@(beta*v), P = qg - F w ------ b_X = tl.load(pX + o_r[:, None] * BT + o_r[None, :]) b_X = tl.where(o_r[:, None] >= o_r[None, :], b_X, 0.0) b_Fr = tl.load(f_ptr + nch * (BT * BT) + o_r[:, None] * BT + o_r[None, :]) for kb in range(0, K, BK): b_khb = tl.load(kh_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :]) b_w = tl.dot(b_X, b_khb) tl.store(w_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :], b_w.to(tl.bfloat16)) b_qgb = tl.load(qg_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :]).to(tl.float32) b_P = b_qgb - tl.dot(b_Fr, b_w.to(tl.bfloat16)) tl.store(p_ptr + (nch * BT + o_r)[:, None] * K + kb + o_bk[None, :], b_P.to(tl.bfloat16)) o_bv = tl.arange(0, 64) for vb in range(0, V, 64): b_vb = tl.load(v_ptr + (row_in[:, None] * H + i_h) * V + vb + o_bv[None, :]) b_vb = (b_vb.to(tl.float32) * b_beta[:, None]).to(tl.bfloat16) b_u = tl.dot(b_X, b_vb) tl.store(u_ptr + (nch * BT + o_r)[:, None] * V + vb + o_bv[None, :], b_u.to(tl.bfloat16)) @triton.jit def scan_grouped( w_ptr, u_ptr, qg_ptr, kk_ptr, f_ptr, gl_ptr, o_ptr, fs_ptr, T, NT, GS, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, num_stages: tl.constexpr = 3, ): """Local scan of one group of GS chunks with zero entry state. Writes o_local to o_ptr and the group-exit state to fs (bf16). """ i_v = tl.program_id(0) i_g = tl.program_id(1) i_bh = tl.program_id(2) i_b = i_bh // H i_h = i_bh % H o_r = tl.arange(0, BT) o_k = tl.arange(0, K) o_v = i_v * BV + tl.arange(0, BV) it0 = i_g * GS it1 = min(it0 + GS, NT) b_S = tl.zeros((K, BV), dtype=tl.float32) for it in tl.range(it0, it1, num_stages=num_stages): base = (i_bh * NT + it) * BT b_w = tl.load(w_ptr + (base + o_r)[:, None] * K + o_k[None, :]) b_Sb = b_S.to(tl.bfloat16) b_vnew = tl.load(u_ptr + (base + o_r)[:, None] * V + o_v[None, :]).to(tl.float32) b_vnew -= tl.dot(b_w, b_Sb) b_vnb = b_vnew.to(tl.bfloat16) b_qg = tl.load(qg_ptr + (base + o_r)[:, None] * K + o_k[None, :]) b_o = tl.dot(b_qg, b_Sb) b_F = tl.load(f_ptr + (i_bh * NT + it) * (BT * BT) + o_r[:, None] * BT + o_r[None, :]) b_o += tl.dot(b_F, b_vnb) t0 = i_b * T + it * BT tl.store(o_ptr + ((t0 + o_r)[:, None] * H + i_h) * V + o_v[None, :], b_o.to(tl.bfloat16)) b_gl = tl.load(gl_ptr + (i_bh * NT + it) * K + o_k) b_kk = tl.load(kk_ptr + (base + o_r)[:, None] * K + o_k[None, :]) b_S = b_S * tl.math.exp2(b_gl)[:, None] + tl.dot(tl.trans(b_kk), b_vnb) # fs: (NGP, B*H, K, V) bf16 fs_off = ((i_g * tl.num_programs(2) + i_bh) * K + o_k[:, None]) * V + o_v[None, :] tl.store(fs_ptr + fs_off, b_S.to(tl.bfloat16)) @triton.jit def combine_states( gl_ptr, fs_ptr, es_ptr, T, NT, GS, NGP, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, ): """E_g = entry state of group g (E_0 = 0; E_{g+1} = Dg_g * E_g + F_g).""" i_v = tl.program_id(0) i_bh = tl.program_id(1) o_k = tl.arange(0, K) o_v = i_v * BV + tl.arange(0, BV) b_E = tl.zeros((K, BV), dtype=tl.float32) nb = tl.num_programs(1) for g in range(0, NGP - 1): it0 = g * GS it1 = min(it0 + GS, NT) b_glsum = tl.zeros((K,), dtype=tl.float32) for it in range(it0, it1): b_glsum += tl.load(gl_ptr + (i_bh * NT + it) * K + o_k) fs_off = ((g * nb + i_bh) * K + o_k[:, None]) * V + o_v[None, :] b_F = tl.load(fs_ptr + fs_off).to(tl.float32) b_E = b_E * tl.math.exp2(b_glsum)[:, None] + b_F es_off = (((g + 1) * nb + i_bh) * K + o_k[:, None]) * V + o_v[None, :] tl.store(es_ptr + es_off, b_E.to(tl.bfloat16)) @triton.jit def fixup_output( p_ptr, gl_ptr, es_ptr, o_ptr, T, NT, GS, NGP, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, num_stages: tl.constexpr = 3, ): """o += (P_c * D_c[None, :]) @ E_g for chunks outside group 0.""" i_v = tl.program_id(0) i_t = tl.program_id(1) + GS # skip group 0 i_bh = tl.program_id(2) if i_t >= NT: return i_g = i_t // GS i_b = i_bh // H i_h = i_bh % H o_r = tl.arange(0, BT) o_k = tl.arange(0, K) o_v = i_v * BV + tl.arange(0, BV) b_glsum = tl.zeros((K,), dtype=tl.float32) it0 = i_g * GS for it in tl.range(it0, i_t, num_stages=num_stages): b_glsum += tl.load(gl_ptr + (i_bh * NT + it) * K + o_k) b_D = tl.math.exp2(b_glsum) base = (i_bh * NT + i_t) * BT b_P = tl.load(p_ptr + (base + o_r)[:, None] * K + o_k[None, :]) nb = tl.num_programs(2) es_off = ((i_g * nb + i_bh) * K + o_k[:, None]) * V + o_v[None, :] b_E = tl.load(es_ptr + es_off) b_corr = tl.dot(b_P * b_D[None, :].to(tl.bfloat16), b_E) t0 = i_b * T + i_t * BT o_off = ((t0 + o_r)[:, None] * H + i_h) * V + o_v[None, :] b_o = tl.load(o_ptr + o_off).to(tl.float32) tl.store(o_ptr + o_off, (b_o + b_corr).to(tl.bfloat16)) def _pick_groups(NT, BH, VT, target_occupancy=228): """Choose chunks-per-group GS so that NG groups x BH x VT tiles cover the GPU; GS=1 when NT is small (enough blocks already).""" NG = max(1, round(target_occupancy / (BH * VT))) NG = min(NG, NT) GS = (NT + NG - 1) // NG NG = (NT + GS - 1) // GS return GS, NG def _fwd(q, k, v, g, beta, scale, chunk_size): B, T, H, K = q.shape V = v.shape[-1] BT = chunk_size assert T % BT == 0 NT = T // BT NCH = B * H * NT dev = q.device w = torch.empty(NCH, BT, K, dtype=torch.bfloat16, device=dev) u = torch.empty(NCH, BT, V, dtype=torch.bfloat16, device=dev) qg = torch.empty(NCH, BT, K, dtype=torch.bfloat16, device=dev) kk = torch.empty(NCH, BT, K, dtype=torch.bfloat16, device=dev) F = torch.empty(NCH, BT, BT, dtype=torch.bfloat16, device=dev) gl = torch.empty(NCH, K, dtype=torch.float32, device=dev) P = torch.empty(NCH, BT, K, dtype=torch.bfloat16, device=dev) asave = torch.empty(NCH, BT, BT, dtype=torch.float32, device=dev) X = torch.empty(NCH, BT, BT, dtype=torch.bfloat16, device=dev) kh = torch.empty(NCH, BT, K, dtype=torch.bfloat16, device=dev) o = torch.empty_like(v) kda_prep_kernel[(NT, B * H)]( q, k, v, g, beta, w, u, qg, kk, F, gl, P, asave, X, kh, scale, T, NT, H=H, K=K, V=V, BT=BT, BC=16, BK=64, num_warps=8, num_stages=1, ) BV = 32 VT = V // BV GS, NG = _pick_groups(NT, B * H, VT) fs = torch.empty(NG, B * H, K, V, dtype=torch.bfloat16, device=dev) es = torch.empty(NG, B * H, K, V, dtype=torch.bfloat16, device=dev) scan_grouped[(VT, NG, B * H)]( w, u, qg, kk, F, gl, o, fs, T, NT, GS, H=H, K=K, V=V, BT=BT, BV=BV, num_warps=4, ) combine_states[(VT, B * H)]( gl, fs, es, T, NT, GS, NG, H=H, K=K, V=V, BT=BT, BV=BV, num_warps=4, ) fixup_output[(VT, max(NT - GS, 0), B * H)]( P, gl, es, o, T, NT, GS, NG, H=H, K=K, V=V, BT=BT, BV=BV, num_warps=4, ) return o class Model(torch.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) def forward(self, q, k, v, g, beta): return _fwd(q, k, v, g, beta, self.scale, self.chunk_size) # Module-level shape shims (mirrors reference interface). 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]