KDA CUTLASS Claude Opus 4.8

"""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]