KDA CUTLASS Kimi K3 (1M)

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