"""Kimi Delta Attention (KDA) chunk-form forward for SM120. Hybrid: torch/cuBLAS GEMMs + Triton solve_tril and inter-chunk recurrence. Matches reference.py math. Does not import FLA operators. """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl OP_TYPE = "linear_attention" SUPPORTED_PRECISIONS = ["bf16"] HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"] _WS: dict = {} _STREAMS: list | None = None RCP_LN2 = 1.4426950408889634 # 1/ln(2) for exp2(x*RCP_LN2) == exp(x) def _streams(): global _STREAMS if _STREAMS is None: _STREAMS = [torch.cuda.Stream(), torch.cuda.Stream()] return _STREAMS # --------------------------------------------------------------------------- # Prep: load (B,T,H,*) -> write packed (N,BT,*) kg_pos/neg, qg, kg_st, beta, g_last, v # Grid: (NT, B*H). Uses tl.cumsum + exp2 (g scaled by RCP_LN2). # --------------------------------------------------------------------------- @triton.jit def _prep_pack_kernel( q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr, kgpos_ptr, kgneg_ptr, qg_ptr, kgst_ptr, beta_out_ptr, glast_ptr, vout_ptr, scale, rcp_ln2, T, H, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, s_qb, s_qt, s_qh, s_qk, s_kb, s_kt, s_kh, s_kk, s_vb, s_vt, s_vh, s_vv, s_gb, s_gt, s_gh, s_gk, s_bb, s_bt, s_bh, s_n, s_c, s_k, # packed (N,BT,K) s_vn, s_vc, s_vv2, # packed v (N,BT,V) s_bn, s_bc, s_gln, s_glk, ): i_n = tl.program_id(0) pid = tl.program_id(1) i_b = pid // H i_h = pid % H NT = T // BT nflat = (i_b * H + i_h) * NT + i_n t0 = i_n * BT offs_c = tl.arange(0, BT) offs_k = tl.arange(0, K) offs_v = tl.arange(0, V) g = tl.load( g_ptr + i_b * s_gb + i_h * s_gh + (t0 + offs_c[:, None]) * s_gt + offs_k[None, :] * s_gk ) gcs = tl.cumsum(g, axis=0) * rcp_ln2 k = tl.load( k_ptr + i_b * s_kb + i_h * s_kh + (t0 + offs_c[:, None]) * s_kt + offs_k[None, :] * s_kk ).to(tl.float32) q = tl.load( q_ptr + i_b * s_qb + i_h * s_qh + (t0 + offs_c[:, None]) * s_qt + offs_k[None, :] * s_qk ).to(tl.float32) * scale beta = tl.load( beta_ptr + i_b * s_bb + i_h * s_bh + (t0 + offs_c) * s_bt ).to(tl.float32) v = tl.load( v_ptr + i_b * s_vb + i_h * s_vh + (t0 + offs_c[:, None]) * s_vt + offs_v[None, :] * s_vv ) eg = tl.exp2(gcs) kg_pos = (k * eg).to(tl.bfloat16) kg_neg = (k * tl.exp2(-gcs)).to(tl.bfloat16) qg = (q * eg).to(tl.bfloat16) g_last = tl.sum(tl.where(offs_c[:, None] == (BT - 1), gcs, 0.0), axis=0) kg_st = (k * tl.exp2(g_last[None, :] - gcs)).to(tl.bfloat16) base = nflat * s_n tl.store(kgpos_ptr + base + offs_c[:, None] * s_c + offs_k[None, :] * s_k, kg_pos) tl.store(kgneg_ptr + base + offs_c[:, None] * s_c + offs_k[None, :] * s_k, kg_neg) tl.store(qg_ptr + base + offs_c[:, None] * s_c + offs_k[None, :] * s_k, qg) tl.store(kgst_ptr + base + offs_c[:, None] * s_c + offs_k[None, :] * s_k, kg_st) tl.store(beta_out_ptr + nflat * s_bn + offs_c * s_bc, beta) tl.store(glast_ptr + nflat * s_gln + offs_k * s_glk, g_last) tl.store(vout_ptr + nflat * s_vn + offs_c[:, None] * s_vc + offs_v[None, :] * s_vv2, v) # --------------------------------------------------------------------------- # Solve A + form w/u. A enters as raw (kg_pos @ kg_neg.T) fp32. # Applies beta_row, -strict_lower, solve, (A+I)*beta_col; then # w = A @ kg_pos, u = A @ v. Does NOT store A (only w, u). # --------------------------------------------------------------------------- @triton.jit def _solve_wu_kernel( A_ptr, beta_ptr, kgpos_ptr, v_ptr, w_ptr, u_ptr, N, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, s_an, s_ar, s_ac, s_bn, s_bc, s_kn, s_kc, s_kk, s_vn, s_vc, s_vv, s_wn, s_wc, s_wk, s_un, s_uc, s_uv, ): n = tl.program_id(0) if n >= N: return offs = tl.arange(0, BT) offs_k = tl.arange(0, K) offs_v = tl.arange(0, V) A = tl.load(A_ptr + n * s_an + offs[:, None] * s_ar + offs[None, :] * s_ac) beta = tl.load(beta_ptr + n * s_bn + offs * s_bc) A = A * beta[:, None] A = -tl.where(offs[:, None] > offs[None, :], A, 0.0) for i in range(1, BT): row = tl.sum(tl.where(offs[:, None] == i, A, 0.0), axis=0) delta = tl.sum(row[:, None] * A, axis=0) A = tl.where((offs[:, None] == i) & (offs[None, :] < i), A + delta[None, :], A) A = A + tl.where(offs[:, None] == offs[None, :], 1.0, 0.0) A = A * beta[None, :] Ab = A.to(tl.bfloat16) kg = tl.load(kgpos_ptr + n * s_kn + offs[:, None] * s_kc + offs_k[None, :] * s_kk) v = tl.load(v_ptr + n * s_vn + offs[:, None] * s_vc + offs_v[None, :] * s_vv) w = tl.dot(Ab, kg) u = tl.dot(Ab, v) tl.store(w_ptr + n * s_wn + offs[:, None] * s_wc + offs_k[None, :] * s_wk, w) tl.store(u_ptr + n * s_un + offs[:, None] * s_uc + offs_v[None, :] * s_uv, u) @triton.jit def _tril_bf16_kernel(A_ptr, N, BT: tl.constexpr, s_an, s_ar, s_ac): n = tl.program_id(0) if n >= N: return offs = tl.arange(0, BT) A = tl.load(A_ptr + n * s_an + offs[:, None] * s_ar + offs[None, :] * s_ac) A = tl.where(offs[:, None] >= offs[None, :], A, 0.0) tl.store(A_ptr + n * s_an + offs[:, None] * s_ar + offs[None, :] * s_ac, A) # --------------------------------------------------------------------------- # Inter-chunk fused recurrence + output. # Grid: (cdiv(V, BV), B*H) # Stores: w/u/Aqk/qg/kg as bf16 for bandwidth; glast fp32. # --------------------------------------------------------------------------- @triton.jit def _inter_kernel( w_ptr, u_ptr, Aqk_ptr, qg_ptr, kg_ptr, glast_ptr, o_ptr, T, H, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr, BK: tl.constexpr, s_wb, s_wh, s_wn, s_wc, s_wk, s_ub, s_uh, s_un, s_uc, s_uv, s_ab, s_ah, s_an, s_ar, s_ac, s_qgb, s_qgh, s_qgn, s_qgc, s_qgk, s_kgb, s_kgh, s_kgn, s_kgc, s_kgk, s_glb, s_glh, s_gln, s_glk, s_ob, s_ot, s_oh, s_ov, ): i_v = tl.program_id(0) pid = tl.program_id(1) i_b = pid // H i_h = pid % H NT = T // BT offs_c = tl.arange(0, BT) offs_v = i_v * BV + tl.arange(0, BV) mask_v = offs_v < V offs_k0 = tl.arange(0, BK) offs_k1 = BK + tl.arange(0, BK) b_h0 = tl.zeros([BK, BV], dtype=tl.float32) b_h1 = tl.zeros([BK, BV], dtype=tl.float32) for i_n in range(NT): wb = i_b * s_wb + i_h * s_wh + i_n * s_wn ub = i_b * s_ub + i_h * s_uh + i_n * s_un ab = i_b * s_ab + i_h * s_ah + i_n * s_an qgb = i_b * s_qgb + i_h * s_qgh + i_n * s_qgn kgb = i_b * s_kgb + i_h * s_kgh + i_n * s_kgn glb = i_b * s_glb + i_h * s_glh + i_n * s_gln w0 = tl.load(w_ptr + wb + offs_c[:, None] * s_wc + offs_k0[None, :] * s_wk) w1 = tl.load(w_ptr + wb + offs_c[:, None] * s_wc + offs_k1[None, :] * s_wk) u = tl.load( u_ptr + ub + offs_c[:, None] * s_uc + offs_v[None, :] * s_uv, mask=mask_v[None, :], other=0.0, ).to(tl.float32) Aqk = tl.load(Aqk_ptr + ab + offs_c[:, None] * s_ar + offs_c[None, :] * s_ac) qg0 = tl.load(qg_ptr + qgb + offs_c[:, None] * s_qgc + offs_k0[None, :] * s_qgk) qg1 = tl.load(qg_ptr + qgb + offs_c[:, None] * s_qgc + offs_k1[None, :] * s_qgk) h0b = b_h0.to(tl.bfloat16) h1b = b_h1.to(tl.bfloat16) vnew = u - (tl.dot(w0, h0b) + tl.dot(w1, h1b)) o = tl.dot(qg0, h0b) + tl.dot(qg1, h1b) + tl.dot(Aqk, vnew.to(tl.bfloat16)) t0 = i_n * BT tl.store( o_ptr + i_b * s_ob + i_h * s_oh + (t0 + offs_c[:, None]) * s_ot + offs_v[None, :] * s_ov, o.to(tl.bfloat16), mask=mask_v[None, :], ) # state: S *= exp(g_last); S += kg.T @ vnew # g_last stored already scaled by RCP_LN2 so we use exp2 gl0 = tl.load(glast_ptr + glb + offs_k0 * s_glk) gl1 = tl.load(glast_ptr + glb + offs_k1 * s_glk) b_h0 = b_h0 * tl.exp2(gl0)[:, None] b_h1 = b_h1 * tl.exp2(gl1)[:, None] kg0 = tl.load(kg_ptr + kgb + offs_c[:, None] * s_kgc + offs_k0[None, :] * s_kgk) kg1 = tl.load(kg_ptr + kgb + offs_c[:, None] * s_kgc + offs_k1[None, :] * s_kgk) vn = vnew.to(tl.bfloat16) b_h0 = b_h0 + tl.dot(tl.trans(kg0), vn) b_h1 = b_h1 + tl.dot(tl.trans(kg1), vn) def _workspace(B, T, H, K, V, BT, device): key = (str(device), B, T, H, K, V, BT) if key in _WS: return _WS[key] NT = T // BT N = B * H * NT d = dict( kg_pos=torch.empty(N, BT, K, device=device, dtype=torch.bfloat16), kg_neg=torch.empty(N, BT, K, device=device, dtype=torch.bfloat16), qg=torch.empty(N, BT, K, device=device, dtype=torch.bfloat16), kg_st=torch.empty(N, BT, K, device=device, dtype=torch.bfloat16), beta=torch.empty(N, BT, device=device, dtype=torch.float32), g_last=torch.empty(N, K, device=device, dtype=torch.float32), A=torch.empty(N, BT, BT, device=device, dtype=torch.float32), Aqk=torch.empty(N, BT, BT, device=device, dtype=torch.bfloat16), w=torch.empty(N, BT, K, device=device, dtype=torch.bfloat16), u=torch.empty(N, BT, V, device=device, dtype=torch.bfloat16), v_buf=torch.empty(N, BT, V, device=device, dtype=torch.bfloat16), ) _WS[key] = d return d def _kda_fwd(q, k, v, g, beta, scale, chunk_size=64): dtype = v.dtype B, T, H, K = q.shape V = v.shape[-1] BT = chunk_size assert T % BT == 0 NT = T // BT N = B * H * NT device = q.device q = q.contiguous() k = k.contiguous() v = v.contiguous() g = g.contiguous() beta = beta.contiguous() ws = _workspace(B, T, H, K, V, BT, device) kg_pos = ws["kg_pos"] kg_neg = ws["kg_neg"] qg = ws["qg"] kg_st = ws["kg_st"] beta_f = ws["beta"] g_last_t = ws["g_last"] A = ws["A"] Aqk = ws["Aqk"] w = ws["w"] u = ws["u"] v_buf = ws["v_buf"] # ---- fused prep: pack + cumsum + exp2 + products ---- _prep_pack_kernel[(NT, B * H)]( q, k, v, g, beta, kg_pos, kg_neg, qg, kg_st, beta_f, g_last_t, v_buf, scale, RCP_LN2, T, H, K, V, BT, q.stride(0), q.stride(1), q.stride(2), q.stride(3), k.stride(0), k.stride(1), k.stride(2), k.stride(3), v.stride(0), v.stride(1), v.stride(2), v.stride(3), g.stride(0), g.stride(1), g.stride(2), g.stride(3), beta.stride(0), beta.stride(1), beta.stride(2), kg_pos.stride(0), kg_pos.stride(1), kg_pos.stride(2), v_buf.stride(0), v_buf.stride(1), v_buf.stride(2), beta_f.stride(0), beta_f.stride(1), g_last_t.stride(0), g_last_t.stride(1), num_warps=4, num_stages=1, ) # ---- A_kk path and Aqk path are independent; overlap on two streams ---- s_akk, s_aqk = _streams() cur = torch.cuda.current_stream() s_akk.wait_stream(cur) s_aqk.wait_stream(cur) with torch.cuda.stream(s_akk): A.copy_(torch.bmm(kg_pos, kg_neg.transpose(1, 2)).float()) # Fused solve+w/u is better at moderate N; still fine at large N. _solve_wu_kernel[(N,)]( A, beta_f, kg_pos, v_buf, w, u, N, K, V, BT, A.stride(0), A.stride(1), A.stride(2), beta_f.stride(0), beta_f.stride(1), kg_pos.stride(0), kg_pos.stride(1), kg_pos.stride(2), v_buf.stride(0), v_buf.stride(1), v_buf.stride(2), w.stride(0), w.stride(1), w.stride(2), u.stride(0), u.stride(1), u.stride(2), num_warps=4, ) with torch.cuda.stream(s_aqk): Aqk.copy_(torch.bmm(qg, kg_neg.transpose(1, 2))) _tril_bf16_kernel[(N,)]( Aqk, N, BT, Aqk.stride(0), Aqk.stride(1), Aqk.stride(2), num_warps=4, ) cur.wait_stream(s_akk) cur.wait_stream(s_aqk) # reshape views (B,H,NT,BT,*) def view5(x, last): return x.view(B, H, NT, BT, last) w_v = view5(w, K) u_v = view5(u, V) Aqk_v = view5(Aqk, BT) qg_v = view5(qg, K) kg_v = view5(kg_st, K) gl_v = g_last_t.view(B, H, NT, K) o = torch.empty(B, T, H, V, device=device, dtype=dtype) BV, BK = 64, 64 _inter_kernel[(triton.cdiv(V, BV), B * H)]( w_v, u_v, Aqk_v, qg_v, kg_v, gl_v, o, T, H, K, V, BT, BV, BK, w_v.stride(0), w_v.stride(1), w_v.stride(2), w_v.stride(3), w_v.stride(4), u_v.stride(0), u_v.stride(1), u_v.stride(2), u_v.stride(3), u_v.stride(4), Aqk_v.stride(0), Aqk_v.stride(1), Aqk_v.stride(2), Aqk_v.stride(3), Aqk_v.stride(4), qg_v.stride(0), qg_v.stride(1), qg_v.stride(2), qg_v.stride(3), qg_v.stride(4), kg_v.stride(0), kg_v.stride(1), kg_v.stride(2), kg_v.stride(3), kg_v.stride(4), gl_v.stride(0), gl_v.stride(1), gl_v.stride(2), gl_v.stride(3), o.stride(0), o.stride(1), o.stride(2), o.stride(3), num_warps=4, num_stages=1, ) return o class Model(nn.Module): """KDA forward (chunk form). No learned parameters.""" 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 _kda_fwd(q, k, v, g, beta, scale=self.scale, chunk_size=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]