KDA CUTLASS Kimi K3 (1M)

"""Kimi Delta Attention (KDA) forward, chunk form, RTX PRO 6000 (SM120).

Pipeline (per (b,h)):
  Kernel A (Triton, parallel over chunks): in-chunk cumsum of g, decay factor
    application, K-K interaction, exact unit-triangular inverse in WY form
    (Neumann factorization), w/u sides, decayed q, chunk-end-decayed k,
    causal A_qk.
  Kernel B (Triton, sequential over chunks, V-tiled): delta-rule state
    recurrence with fp32 state in registers. Emits h^T and vn^T slices.
  Kernel C (Triton, parallel over chunks): o = qd @ h + Aqk @ vn.

bf16 tensor-core dots with fp32 accumulation everywhere in the recurrence.
"""
from __future__ import annotations

import os

import torch
import torch.nn as nn
import triton
import triton.language as tl

os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0a")

# ===========================================================================
# Kernel A: intra-chunk
# ===========================================================================
@triton.jit
def _inv32_strict(L):
    """Invert (I + L) for strictly-lower-triangular L (32,32) via exact Neumann
    factorization: (I+L)^{-1} = (I-L)(I+L^2)(I+L^4)(I+L^8)(I+L^16) (L^32 = 0).
    fp32 (ieee) dots -- same values as forward substitution."""
    rows = tl.arange(0, 32)
    eye = tl.where(rows[:, None] == rows[None, :], 1.0, 0.0).to(tl.float32)
    L2 = tl.dot(L, L, input_precision="ieee")
    L4 = tl.dot(L2, L2, input_precision="ieee")
    L8 = tl.dot(L4, L4, input_precision="ieee")
    L16 = tl.dot(L8, L8, input_precision="ieee")
    P = (eye - L) + L2 - tl.dot(L, L2, input_precision="ieee")
    P = P + tl.dot(P, L4, input_precision="ieee")
    P = P + tl.dot(P, L8, input_precision="ieee")
    P = P + tl.dot(P, L16, input_precision="ieee")
    return P


@triton.jit
def kda_intra_kernel(
    q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr,
    w_ptr, u_ptr, qd_ptr, aqk_ptr, kd_ptr, egl_ptr,
    scale,
    T, H, NT,
    K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, HB: tl.constexpr,
):
    i_t = tl.program_id(0)
    i_bh = tl.program_id(1)
    i_b = i_bh // H
    i_h = i_bh % H
    t0 = i_t * BT

    r = tl.arange(0, HB)
    dK = tl.arange(0, K)

    row0 = t0 + r
    row1 = t0 + HB + r
    off0 = ((i_b * T + row0) * H + i_h)[:, None] * K + dK[None, :]
    off1 = ((i_b * T + row1) * H + i_h)[:, None] * K + dK[None, :]
    offb0 = (i_b * T + row0) * H + i_h
    offb1 = (i_b * T + row1) * H + i_h

    cbase = (i_bh * NT + i_t) * BT
    off_o0 = (cbase + r)[:, None] * K + dK[None, :]
    off_o1 = (cbase + HB + r)[:, None] * K + dK[None, :]

    # ---- decay factors (fp32 cumsum, eager bf16 conversion) ----
    g0 = tl.load(g_ptr + off0)
    gc0 = tl.cumsum(g0, axis=0)
    s0 = tl.sum(g0, axis=0)
    g1 = tl.load(g_ptr + off1)
    gc1 = tl.cumsum(g1, axis=0) + s0[None, :]
    glast = s0 + tl.sum(g1, axis=0)
    egl = tl.exp(glast)
    tl.store(egl_ptr + (i_bh * NT + i_t) * K + dK, egl)

    e0 = tl.exp(gc0)
    e1 = tl.exp(gc1)
    k0 = tl.load(k_ptr + off0)
    k1 = tl.load(k_ptr + off1)
    kb0 = (k0.to(tl.float32) * e0).to(tl.bfloat16)
    kb1 = (k1.to(tl.float32) * e1).to(tl.bfloat16)
    kf0 = (k0.to(tl.float32) * tl.exp(-gc0)).to(tl.bfloat16)
    kf1 = (k1.to(tl.float32) * tl.exp(-gc1)).to(tl.bfloat16)

    b0 = tl.load(beta_ptr + offb0).to(tl.float32)
    b1 = tl.load(beta_ptr + offb1).to(tl.float32)

    # kd = k*exp(glast - gc) = kf * egl (early store, pre-transposed (K,BT)
    # layout so the state kernel needs no per-iteration transpose)
    kd0 = (kf0.to(tl.float32) * egl[None, :]).to(tl.bfloat16)
    kd1 = (kf1.to(tl.float32) * egl[None, :]).to(tl.bfloat16)
    kdT_base = (i_bh * NT + i_t) * BT * K
    tl.store(kd_ptr + kdT_base + r[:, None] + dK[None, :] * BT, kd0)
    tl.store(kd_ptr + kdT_base + (HB + r)[:, None] + dK[None, :] * BT, kd1)

    # ---- decayed q (early store) ----
    q0 = tl.load(q_ptr + off0)
    q1 = tl.load(q_ptr + off1)
    qd0 = (q0.to(tl.float32) * e0 * scale).to(tl.bfloat16)
    qd1 = (q1.to(tl.float32) * e1 * scale).to(tl.bfloat16)
    tl.store(qd_ptr + off_o0, qd0)
    tl.store(qd_ptr + off_o1, qd1)

    # ---- Aqk (causal, quadrant dots; strictly-upper quadrant left unwritten) ----
    causal = r[:, None] >= r[None, :]
    aq00 = tl.where(causal, tl.dot(qd0, tl.trans(kf0)), 0.0).to(tl.bfloat16)
    aq10 = tl.dot(qd1, tl.trans(kf0)).to(tl.bfloat16)
    aq11 = tl.where(causal, tl.dot(qd1, tl.trans(kf1)), 0.0).to(tl.bfloat16)
    abase = (i_bh * NT + i_t) * BT * BT
    tl.store(aqk_ptr + abase + r[:, None] * BT + r[None, :], aq00)
    tl.store(aqk_ptr + abase + BT // 2 * BT + r[:, None] * BT + r[None, :], aq10)
    tl.store(aqk_ptr + abase + BT // 2 * BT + BT // 2 + r[:, None] * BT + r[None, :], aq11)

    # ---- A = strictly-lower K-K interaction, inverse in WY form ----
    strict = r[:, None] > r[None, :]
    L00 = tl.where(strict, tl.dot(kb0, tl.trans(kf0)) * b0[:, None], 0.0)
    L11 = tl.where(strict, tl.dot(kb1, tl.trans(kf1)) * b1[:, None], 0.0)
    L10 = tl.dot(kb1, tl.trans(kf0)) * b1[:, None]

    M00 = _inv32_strict(L00)
    M11 = _inv32_strict(L11)
    C10 = -tl.dot(tl.dot(M11, L10), M00)
    M00b = M00.to(tl.bfloat16)
    M11b = M11.to(tl.bfloat16)
    C10b = C10.to(tl.bfloat16)

    # ---- w = M @ (beta * exp(g) * k); u = M @ (beta * v) ----
    rw0 = (kb0.to(tl.float32) * b0[:, None]).to(tl.bfloat16)
    rw1 = (kb1.to(tl.float32) * b1[:, None]).to(tl.bfloat16)
    tl.store(w_ptr + off_o0, tl.dot(M00b, rw0).to(tl.bfloat16))
    tl.store(w_ptr + off_o1, (tl.dot(C10b, rw0) + tl.dot(M11b, rw1)).to(tl.bfloat16))

    v0 = tl.load(v_ptr + off0)
    v1 = tl.load(v_ptr + off1)
    ru0 = (v0.to(tl.float32) * b0[:, None]).to(tl.bfloat16)
    ru1 = (v1.to(tl.float32) * b1[:, None]).to(tl.bfloat16)
    tl.store(u_ptr + off_o0, tl.dot(M00b, ru0).to(tl.bfloat16))
    tl.store(u_ptr + off_o1, (tl.dot(C10b, ru0) + tl.dot(M11b, ru1)).to(tl.bfloat16))


# ===========================================================================
# Kernel B: Triton inter-chunk recurrence (fallback path)
# ===========================================================================
@triton.jit
def kda_state_kernel(
    w_ptr, u_ptr, kd_ptr, egl_ptr,
    ht_ptr, vnt_ptr,
    H, NT,
    K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BV: tl.constexpr,
):
    i_v = tl.program_id(0)
    i_bh = tl.program_id(1)

    rT = tl.arange(0, BT)
    dK = tl.arange(0, K)
    dV = i_v * BV + tl.arange(0, BV)

    S = tl.zeros((K, BV), dtype=tl.float32)
    for i in range(NT):
        hbase = (i_bh * NT + i) * V * K
        tl.store(ht_ptr + hbase + dV[None, :] * K + dK[:, None], S.to(tl.bfloat16))

        cbase = (i_bh * NT + i) * BT
        w_i = tl.load(w_ptr + (cbase + rT)[:, None] * K + dK[None, :])
        u_i = tl.load(u_ptr + (cbase + rT)[:, None] * V + dV[None, :])
        vn = u_i.to(tl.float32) - tl.dot(w_i, S.to(tl.bfloat16))
        tl.store(vnt_ptr + (i_bh * NT + i) * V * BT + dV[None, :] * BT + rT[:, None], vn.to(tl.bfloat16))

        kd_i = tl.load(kd_ptr + (i_bh * NT + i) * BT * K + dK[:, None] * BT + rT[None, :])
        egl = tl.load(egl_ptr + (i_bh * NT + i) * K + dK)
        S = S * egl[:, None] + tl.dot(kd_i, vn.to(tl.bfloat16))


# ===========================================================================
# Kernel C: output
# ===========================================================================
@triton.jit
def kda_output_kernel(
    qd_ptr, aqk_ptr, ht_ptr, vnt_ptr, o_ptr,
    T, H, NT,
    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
    t0 = i_t * BT

    rT = tl.arange(0, BT)
    dK = tl.arange(0, K)
    dV = tl.arange(0, V)

    cbase = (i_bh * NT + i_t) * BT
    qd = tl.load(qd_ptr + (cbase + rT)[:, None] * K + dK[None, :])
    # causal mask applied here: the intra kernel leaves the strictly-upper
    # quadrant of aqk unwritten.
    Aqk = tl.load(aqk_ptr + (i_bh * NT + i_t) * BT * BT + rT[:, None] * BT + rT[None, :]).to(tl.float32)
    causal = rT[:, None] >= rT[None, :]
    Aqk = tl.where(causal, Aqk, 0.0).to(tl.bfloat16)

    # h^T buffer (V rows, K cols); need (K, V)
    hT = tl.load(ht_ptr + (i_bh * NT + i_t) * V * K + dV[:, None] * K + dK[None, :])
    vnT = tl.load(vnt_ptr + (i_bh * NT + i_t) * V * BT + dV[:, None] * BT + rT[None, :])

    o = tl.dot(qd, tl.trans(hT)) + tl.dot(Aqk, tl.trans(vnT))

    o_off = ((i_b * T + t0 + rT) * H + i_h)[:, None] * V + dV[None, :]
    tl.store(o_ptr + o_off, o.to(tl.bfloat16))


# ===========================================================================
# Host wrapper
# ===========================================================================
class Model(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)
        assert K == 128 and V == 128 and chunk_size == 64, "kernel specialized for K=V=128, BT=64"
        assert T % chunk_size == 0
        self.NT = T // chunk_size
        self._scratch = None
        self._cuda = None

    def _alloc(self, device):
        B, T, H, K, V, NT = self.B, self.T, self.H, self.K, self.V, self.NT
        BT = self.chunk_size
        f = dict(device=device, dtype=torch.bfloat16)
        return dict(
            w=torch.empty((B * H * NT, BT, K), **f),
            u=torch.empty((B * H * NT, BT, V), **f),
            qd=torch.empty((B * H * NT, BT, K), **f),
            aqk=torch.empty((B * H * NT, BT, BT), **f),
            kd=torch.empty((B * H * NT, BT, K), **f),
            egl=torch.empty((B * H * NT, K), device=device, dtype=torch.float32),
            ht=torch.empty((B * H * NT, V, K), **f),
            vnt=torch.empty((B * H * NT, V, BT), **f),
            o=torch.empty((B, T, H, V), **f),
        )

    def forward(self, q, k, v, g, beta):
        B, T, H, K, V, NT = self.B, self.T, self.H, self.K, self.V, self.NT
        BT = self.chunk_size
        BH = B * H
        dev = q.device
        if self._scratch is None or self._scratch["w"].device != dev:
            self._scratch = self._alloc(dev)
        s = self._scratch
        o = s["o"]

        kda_intra_kernel[(NT, BH)](
            q, k, v, g, beta,
            s["w"], s["u"], s["qd"], s["aqk"], s["kd"], s["egl"],
            self.scale, T, H, NT,
            K=K, V=V, BT=BT, HB=BT // 2,
            num_warps=8, num_stages=3,
        )
        BV = 16
        kda_state_kernel[(V // BV, BH)](
            s["w"], s["u"], s["kd"], s["egl"], s["ht"], s["vnt"],
            H, NT,
            K=K, V=V, BT=BT, BV=BV,
            num_warps=4, num_stages=3,
        )
        kda_output_kernel[(NT, BH)](
            s["qd"], s["aqk"], s["ht"], s["vnt"], o,
            T, H, NT,
            K=K, V=V, BT=BT,
            num_warps=8, num_stages=2,
        )
        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]