"""Kimi Delta Attention (KDA) forward, chunk form — custom Triton kernels for B200 (SM100). Decomposition: 1. prepare: per (b,h,chunk) compute cumsum(g), the WY inverse A=(I-B)^-1 @ diag(beta), w = A @ (e^{gc} k), u = A @ v, and the folded chunk transition M_i = diag(e^{gc_last}) - kg_i^T w_i , c_i = kg_i^T u_i , where kg_i = e^{gc_last - gc} k. This turns the inter-chunk recurrence into a single affine map S_{i+1} = M_i S_i + c_i (one matmul per step in the scan). 2. scan: per (b,h, v-tile) sequential over chunks, S_{i+1}=M_i S_i + c_i, emit state h_i=S_i. 3. fwd_o: per (b,h,chunk) parallel, v_new=u-w@h_i, o = (q e^{gc} scale) @ h_i + Aqk @ v_new. A CUDA graph is captured for the (fixed-pointer) benchmark path; correctness checks with fresh input tensors fall back to eager launches. """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl # ---------------------------------------------------------------------------- # Kernel 1: prepare WY representation + folded transition (per (b,h,chunk)) # ---------------------------------------------------------------------------- @triton.jit def _prepare_kernel( k_ptr, v_ptr, g_ptr, beta_ptr, w_ptr, u_ptr, M_ptr, c_ptr, T, NT, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, NITER: tl.constexpr, ): i_t = tl.program_id(0) i_bh = tl.program_id(1) i_b = i_bh // H i_h = i_bh % H base_k = (i_b * T * H + i_h) * K base_v = (i_b * T * H + i_h) * V row0 = i_t * BT p_k = tl.make_block_ptr(k_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) p_g = tl.make_block_ptr(g_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) p_v = tl.make_block_ptr(v_ptr + base_v, (T, V), (H * V, 1), (row0, 0), (BT, V), (1, 0)) b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) b_v = tl.load(p_v, boundary_check=(0, 1)) p_beta = tl.make_block_ptr(beta_ptr + (i_b * T * H + i_h), (T,), (H,), (row0,), (BT,), (0,)) b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) gc = tl.cumsum(b_g, axis=0) # [BT, K] eg = tl.exp(gc) egn = 1.0 / eg # e^{-gc} kge = (b_k * eg).to(tl.bfloat16) # e^{gc} k kgen = b_k * egn # e^{-gc} k (fp32) # B (strict lower): B[c,i] = -beta[c] * Smat = tl.dot(kge, kgen.to(tl.bfloat16).T) # [BT, BT] Bm = -b_beta[:, None] * Smat row = tl.arange(0, BT) col = tl.arange(0, BT) Bm = tl.where(row[:, None] > col[None, :], Bm, 0.0).to(tl.bfloat16) # T_full = (I - B)^-1 = sum_{j>=0} B^j = prod (I + B^{2^k}), B nilpotent (B^BT=0) eye = (row[:, None] == col[None, :]).to(tl.float32) Tf = (eye + Bm).to(tl.bfloat16) Bp = Bm for _ in range(NITER): Bp = tl.dot(Bp, Bp).to(tl.bfloat16) Tf = (Tf + tl.dot(Tf, Bp)).to(tl.bfloat16) A = (Tf * b_beta[None, :]).to(tl.bfloat16) # column-scale by beta b_w = tl.dot(A, kge).to(tl.bfloat16) # [BT, K] b_u = tl.dot(A, b_v).to(tl.bfloat16) # [BT, V] gc_last = tl.sum(b_g, axis=0) # [K] (= last row of cumsum) D = tl.exp(gc_last) # [K] diagonal decay kg = (D[None, :] * kgen).to(tl.bfloat16) # e^{gc_last - gc} k # folded transition: M = diag(D) - kg^T w , c = kg^T u P = tl.dot(kg.T, b_w) # [K, K] ok = tl.arange(0, K) M = tl.where(ok[:, None] == ok[None, :], D[:, None] - P, -P) cmat = tl.dot(kg.T, b_u) # [K, V] p_w = tl.make_block_ptr(w_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) p_u = tl.make_block_ptr(u_ptr + base_v, (T, V), (H * V, 1), (row0, 0), (BT, V), (1, 0)) tl.store(p_w, b_w, boundary_check=(0, 1)) tl.store(p_u, b_u, boundary_check=(0, 1)) base_M = ((i_b * NT + i_t) * H + i_h) * K * K base_c = ((i_b * NT + i_t) * H + i_h) * K * V p_M = tl.make_block_ptr(M_ptr + base_M, (K, K), (K, 1), (0, 0), (K, K), (1, 0)) p_c = tl.make_block_ptr(c_ptr + base_c, (K, V), (V, 1), (0, 0), (K, V), (1, 0)) tl.store(p_M, M.to(M_ptr.dtype.element_ty), boundary_check=(0, 1)) tl.store(p_c, cmat.to(c_ptr.dtype.element_ty), boundary_check=(0, 1)) # ---------------------------------------------------------------------------- # Kernel 2: sequential affine scan S_{i+1} = M_i S_i + c_i (per (b,h, v-tile)) # ---------------------------------------------------------------------------- @triton.jit def _scan_kernel( M_ptr, c_ptr, h_ptr, NT, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BV: tl.constexpr, ): i_bh = tl.program_id(0) i_v = tl.program_id(1) i_b = i_bh // H i_h = i_bh % H v0 = i_v * BV S = tl.zeros((K, BV), dtype=tl.float32) for i_t in range(NT): base_M = ((i_b * NT + i_t) * H + i_h) * K * K base_c = ((i_b * NT + i_t) * H + i_h) * K * V base_h = ((i_b * NT + i_t) * H + i_h) * K * V p_h = tl.make_block_ptr(h_ptr + base_h, (K, V), (V, 1), (0, v0), (K, BV), (1, 0)) tl.store(p_h, S.to(h_ptr.dtype.element_ty), boundary_check=(0, 1)) p_M = tl.make_block_ptr(M_ptr + base_M, (K, K), (K, 1), (0, 0), (K, K), (1, 0)) p_c = tl.make_block_ptr(c_ptr + base_c, (K, V), (V, 1), (0, v0), (K, BV), (1, 0)) b_M = tl.load(p_M, boundary_check=(0, 1)) b_c = tl.load(p_c, boundary_check=(0, 1)).to(tl.float32) S = tl.dot(b_M, S.to(b_M.dtype)) + b_c # ---------------------------------------------------------------------------- # Kernel 3: output (one program per (b,h,chunk)) # ---------------------------------------------------------------------------- @triton.jit def _output_kernel( q_ptr, k_ptr, g_ptr, w_ptr, u_ptr, h_ptr, o_ptr, scale, T, NT, H: tl.constexpr, 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 row0 = i_t * BT base_k = (i_b * T * H + i_h) * K base_v = (i_b * T * H + i_h) * V base_h = ((i_b * NT + i_t) * H + i_h) * K * V p_q = tl.make_block_ptr(q_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) p_k = tl.make_block_ptr(k_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) p_g = tl.make_block_ptr(g_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) b_gc = tl.cumsum(b_g, axis=0) eg = tl.exp(b_gc) qg = (b_q * eg * scale).to(tl.bfloat16) # [BT, K] kn = (b_k / eg).to(tl.bfloat16) # [BT, K] (= k e^{-gc}) p_h = tl.make_block_ptr(h_ptr + base_h, (K, V), (V, 1), (0, 0), (K, V), (1, 0)) b_h = tl.load(p_h, boundary_check=(0, 1)) # v_new = u - w @ h p_w = tl.make_block_ptr(w_ptr + base_k, (T, K), (H * K, 1), (row0, 0), (BT, K), (1, 0)) p_u = tl.make_block_ptr(u_ptr + base_v, (T, V), (H * V, 1), (row0, 0), (BT, V), (1, 0)) b_w = tl.load(p_w, boundary_check=(0, 1)) b_u = tl.load(p_u, boundary_check=(0, 1)).to(tl.float32) v_new = (b_u - tl.dot(b_w, b_h)).to(tl.bfloat16) # o1 = qg @ h o = tl.dot(qg, b_h) # Aqk = qg @ kn^T, masked to lower-incl-diag Aqk = tl.dot(qg, kn.T) row = tl.arange(0, BT) col = tl.arange(0, BT) Aqk = tl.where(col[None, :] <= row[:, None], Aqk, 0.0).to(tl.bfloat16) o += tl.dot(Aqk, v_new) p_o = tl.make_block_ptr(o_ptr + base_v, (T, V), (H * V, 1), (row0, 0), (BT, V), (1, 0)) tl.store(p_o, o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1)) def _launch(q, k, v, g, beta, scale, BT, buf): B, T, H, K = q.shape V = v.shape[-1] NT = T // BT BH = B * H w, u, M, c, h, o = buf # Thin problems (few b*h programs per wave) prefer more warps per program. thin = (NT * BH) < 256 pnw, onw, ons = (8, 8, 1) if thin else (4, 4, 3) _prepare_kernel[(NT, BH)]( k, v, g, beta, w, u, M, c, T, NT, H, K, V, BT, 5, num_warps=pnw, num_stages=2, ) BV = 16 _scan_kernel[(BH, V // BV)]( M, c, h, NT, H, K, V, BV, num_warps=4, num_stages=4, ) _output_kernel[(NT, BH)]( q, k, g, w, u, h, o, scale, T, NT, H, K, V, BT, num_warps=onw, num_stages=ons, ) 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) self._buf = None self._graph = None self._gptrs = None def _alloc(self, dev): B, T, H, K, V = self.B, self.T, self.H, self.K, self.V NT = T // self.chunk_size bf = torch.bfloat16 self._buf = ( torch.empty((B, T, H, K), dtype=bf, device=dev), # w torch.empty((B, T, H, V), dtype=bf, device=dev), # u torch.empty((B, NT, H, K, K), dtype=bf, device=dev), # M torch.empty((B, NT, H, K, V), dtype=bf, device=dev), # c torch.empty((B, NT, H, K, V), dtype=bf, device=dev), # h torch.empty((B, T, H, V), dtype=bf, device=dev), # o ) def _capture(self, ins): s = torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): for _ in range(3): _launch(*ins, self.scale, self.chunk_size, self._buf) torch.cuda.current_stream().wait_stream(s) gph = torch.cuda.CUDAGraph() with torch.cuda.graph(gph): _launch(*ins, self.scale, self.chunk_size, self._buf) self._graph = gph def forward(self, q, k, v, g, beta): if self._buf is None: self._alloc(q.device) ins = [q.contiguous(), k.contiguous(), v.contiguous(), g.contiguous(), beta.contiguous()] ptrs = tuple(x.data_ptr() for x in ins) if self._graph not in (None, False) and ptrs == self._gptrs: self._graph.replay() return self._buf[-1] o = _launch(*ins, self.scale, self.chunk_size, self._buf) if self._graph is None: try: self._capture(ins) self._gptrs = ptrs except Exception: self._graph = False 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]