"""Kimi Delta Attention (KDA) forward — chunk form, custom Triton kernels (SM90).

Implements the same math as reference.py (chunk-parallel KDA) without calling
any forbidden library entry points. Pipeline:

  1. local cumsum of per-channel log-decay g within each chunk
  2. intra-chunk: Aqk (lower-tri incl diag) + Akk via WY / tril solve
  3. w = A @ (beta ⊙ exp(g) ⊙ k),  u = A @ (beta ⊙ v)
  4. inter-chunk gated-delta state recurrence → v_new, h
  5. o = (q ⊙ exp(g)) @ h + Aqk @ v_new

Specialized for BT=64, K=V=128, H==HV (no GVA), no initial/final state.
"""
from __future__ import annotations

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

# log2(e) so exp(x) = exp2(x * LOG2E)  (matches FLA's RCP_LN2 cumsum path)
LOG2E = 1.4426950408889634


# ---------------------------------------------------------------------------
# Kernel 1: in-chunk cumsum of g, convert to log2 space (scale by LOG2E)
# g_in:  [B, T, H, K] fp32   (raw log-decay, natural)
# g_out: [B, T, H, K] fp32   (cumsum * LOG2E so exp2(g_out) == exp(cumsum_nat))
# ---------------------------------------------------------------------------
@triton.jit
def _local_cumsum_kernel(
    g_in,
    g_out,
    T,
    H: tl.constexpr,
    K: tl.constexpr,
    BT: tl.constexpr,
    BK: tl.constexpr,
    SCALE: tl.constexpr,
):
    i_bh = tl.program_id(0)
    i_t = tl.program_id(1)
    i_k = tl.program_id(2)
    i_b = i_bh // H
    i_h = i_bh % H

    bos = i_b * T
    offs_t = i_t * BT + tl.arange(0, BT)
    offs_k = i_k * BK + tl.arange(0, BK)
    mask_t = offs_t < T
    mask_k = offs_k < K
    mask = mask_t[:, None] & mask_k[None, :]

    ptrs = g_in + ((bos + offs_t)[:, None] * H + i_h) * K + offs_k[None, :]
    vals = tl.load(ptrs, mask=mask, other=0.0).to(tl.float32) * SCALE
    # inclusive cumsum along the chunk time axis (axis 0)
    out = tl.cumsum(vals, axis=0)

    out_ptrs = g_out + ((bos + offs_t)[:, None] * H + i_h) * K + offs_k[None, :]
    tl.store(out_ptrs, out, mask=mask)


# ---------------------------------------------------------------------------
# Kernel 2: diagonal sub-chunk Aqk / Akk (BC x BC), token-serial within subchunk
# Computes lower-tri Aqk (incl diag) and strict-lower Akk (row-scaled by beta).
# Akk written as raw L (not yet inverted) into Akkd fp32 buffer of width BC.
# ---------------------------------------------------------------------------
@triton.jit
def _intra_diag_kernel(
    q,
    k,
    g,
    beta,
    Aqk,
    Akkd,
    scale,
    T,
    H: tl.constexpr,
    K: tl.constexpr,
    BT: tl.constexpr,
    BC: tl.constexpr,
    BK: tl.constexpr,
):
    """One program = one (batch, head, chunk, sub-chunk)."""
    i_bh = tl.program_id(0)
    i_tc = tl.program_id(1)  # chunk index * NC + subchunk index
    NC: tl.constexpr = BT // BC
    i_t = i_tc // NC
    i_s = i_tc % NC
    i_b = i_bh // H
    i_h = i_bh % H

    bos = i_b * T
    i_ti = i_t * BT + i_s * BC  # absolute start of this sub-chunk

    q += (bos * H + i_h) * K
    k += (bos * H + i_h) * K
    g += (bos * H + i_h) * K
    beta += bos * H + i_h
    Aqk += (bos * H + i_h) * BT
    Akkd += (bos * H + i_h) * BC

    o_i = tl.arange(0, BC)
    o_k = tl.arange(0, BK)
    m_k = o_k < K

    # Midpoint offset for numerical stability (same trick as FLA safe path)
    mid = BC // 2
    b_gn = tl.load(g + (i_ti + mid) * H * K + o_k, mask=m_k, other=0.0)

    # Load full sub-chunk blocks
    p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
    p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
    p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
    p_b = tl.make_block_ptr(beta, (T,), (H,), (i_ti,), (BC,), (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_beta = tl.load(p_b, boundary_check=(0,)).to(tl.float32)

    b_gm = b_g - b_gn[None, :]
    b_gq = tl.math.exp2(b_gm)
    b_gk = tl.math.exp2(-b_gm)

    b_kgt = tl.trans(b_k * b_gk)
    b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale
    b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None]

    m_Aqk = o_i[:, None] >= o_i[None, :]
    m_Akk = o_i[:, None] > o_i[None, :]
    m_I = o_i[:, None] == o_i[None, :]

    b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0)
    b_Akk = tl.where(m_Akk, b_Akk, 0.0)

    p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_ti, i_s * BC), (BC, BC), (1, 0))
    tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))

    # Forward substitution on diagonal: invert (I - L) entirely in registers.
    # Row i of -L is extracted via masked sum so we never reload from HBM.
    b_Ai = -b_Akk
    for i in range(2, BC):
        b_a = -tl.sum(tl.where((o_i == i)[:, None], b_Akk, 0.0), 0)
        b_a = tl.where(o_i < i, b_a, 0.0)
        b_a += tl.sum(b_a[:, None] * b_Ai, 0)
        b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
    b_Ai += m_I
    p_Akkd = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_ti, 0), (BC, BC), (1, 0))
    tl.store(p_Akkd, b_Ai.to(tl.float32), boundary_check=(0, 1))


# ---------------------------------------------------------------------------
# Kernel 3: off-diagonal Aqk / Akk + merge 4x4 block inverse into full Akk
# ---------------------------------------------------------------------------
@triton.jit
def _inter_solve_kernel(
    q,
    k,
    g,
    beta,
    Aqk,
    Akkd,
    Akk,
    scale,
    T,
    H: tl.constexpr,
    K: 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
    bos = i_b * T

    i_tc0 = i_t * BT
    i_tc1 = i_t * BT + BC
    i_tc2 = i_t * BT + 2 * BC
    i_tc3 = i_t * BT + 3 * BC

    q += (bos * H + i_h) * K
    k += (bos * H + i_h) * K
    g += (bos * H + i_h) * K
    beta += bos * H + i_h
    Aqk += (bos * H + i_h) * BT
    Akk += (bos * H + i_h) * BT
    Akkd += (bos * H + i_h) * BC

    o_i = tl.arange(0, BC)

    b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32)
    b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32)

    for i_k in range(tl.cdiv(K, BK)):
        o_k = i_k * BK + tl.arange(0, BK)
        m_k = o_k < K

        p_k0 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
        p_g0 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0))
        b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
        b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)

        # block 1 vs 0
        p_q1 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
        p_k1 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
        p_g1 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0))
        b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
        b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
        b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
        b_gn1 = tl.load(g + i_tc1 * H * K + o_k, mask=m_k, other=0.0).to(tl.float32)
        b_gqn = tl.math.exp2(b_g1 - b_gn1[None, :])
        b_kgt = tl.trans(b_k0 * tl.math.exp2(b_gn1[None, :] - b_g0))
        b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt)
        b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)

        # block 2 vs 0,1
        p_q2 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
        p_k2 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
        p_g2 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0))
        b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
        b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
        b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
        b_gn2 = tl.load(g + i_tc2 * H * K + o_k, mask=m_k, other=0.0).to(tl.float32)
        b_gqn2 = tl.math.exp2(b_g2 - b_gn2[None, :])
        b_qg2 = b_q2 * b_gqn2
        b_kg2 = b_k2 * b_gqn2
        b_kgt = tl.trans(b_k0 * tl.math.exp2(b_gn2[None, :] - b_g0))
        b_Aqk20 += tl.dot(b_qg2, b_kgt)
        b_Akk20 += tl.dot(b_kg2, b_kgt)
        b_kgt = tl.trans(b_k1 * tl.math.exp2(b_gn2[None, :] - b_g1))
        b_Aqk21 += tl.dot(b_qg2, b_kgt)
        b_Akk21 += tl.dot(b_kg2, b_kgt)

        # block 3 vs 0,1,2
        p_q3 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
        p_k3 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
        p_g3 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0))
        b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
        b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
        b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
        b_gn3 = tl.load(g + i_tc3 * H * K + o_k, mask=m_k, other=0.0).to(tl.float32)
        b_gqn3 = tl.math.exp2(b_g3 - b_gn3[None, :])
        b_qg3 = b_q3 * b_gqn3
        b_kg3 = b_k3 * b_gqn3
        b_kgt = tl.trans(b_k0 * tl.math.exp2(b_gn3[None, :] - b_g0))
        b_Aqk30 += tl.dot(b_qg3, b_kgt)
        b_Akk30 += tl.dot(b_kg3, b_kgt)
        b_kgt = tl.trans(b_k1 * tl.math.exp2(b_gn3[None, :] - b_g1))
        b_Aqk31 += tl.dot(b_qg3, b_kgt)
        b_Akk31 += tl.dot(b_kg3, b_kgt)
        b_kgt = tl.trans(b_k2 * tl.math.exp2(b_gn3[None, :] - b_g2))
        b_Aqk32 += tl.dot(b_qg3, b_kgt)
        b_Akk32 += tl.dot(b_kg3, b_kgt)

    # store off-diagonal Aqk (with scale) and apply beta to Akk off-diagonals
    p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
    tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
    b_b1 = tl.load(tl.make_block_ptr(beta, (T,), (H,), (i_tc1,), (BC,), (0,)), boundary_check=(0,)).to(tl.float32)
    b_Akk10 = b_Akk10 * b_b1[:, None]

    p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
    p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
    tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
    b_b2 = tl.load(tl.make_block_ptr(beta, (T,), (H,), (i_tc2,), (BC,), (0,)), boundary_check=(0,)).to(tl.float32)
    b_Akk20 = b_Akk20 * b_b2[:, None]
    b_Akk21 = b_Akk21 * b_b2[:, None]

    p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
    p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
    p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0))
    tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1))
    b_b3 = tl.load(tl.make_block_ptr(beta, (T,), (H,), (i_tc3,), (BC,), (0,)), boundary_check=(0,)).to(tl.float32)
    b_Akk30 = b_Akk30 * b_b3[:, None]
    b_Akk31 = b_Akk31 * b_b3[:, None]
    b_Akk32 = b_Akk32 * b_b3[:, None]

    # load diagonal inverses (already solved in _intra_diag_kernel)
    b_Ai00 = tl.load(tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)), boundary_check=(0, 1)).to(tl.float32)
    b_Ai11 = tl.load(tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)), boundary_check=(0, 1)).to(tl.float32)
    b_Ai22 = tl.load(tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)), boundary_check=(0, 1)).to(tl.float32)
    b_Ai33 = tl.load(tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)), boundary_check=(0, 1)).to(tl.float32)

    # Block-triangular inverse merge:
    #   Ai10 = -Ai11 @ Akk10 @ Ai00
    #   Ai21 = -Ai22 @ Akk21 @ Ai11
    #   Ai20 = -Ai22 @ (Akk20 @ Ai00 + Akk21 @ Ai10)
    #   ...
    b_Ai10 = -tl.dot(tl.dot(b_Ai11, b_Akk10), b_Ai00)
    b_Ai21 = -tl.dot(tl.dot(b_Ai22, b_Akk21), b_Ai11)
    b_Ai20 = -tl.dot(b_Ai22, tl.dot(b_Akk20, b_Ai00) + tl.dot(b_Akk21, b_Ai10))
    b_Ai32 = -tl.dot(tl.dot(b_Ai33, b_Akk32), b_Ai22)
    b_Ai31 = -tl.dot(b_Ai33, tl.dot(b_Akk31, b_Ai11) + tl.dot(b_Akk32, b_Ai21))
    b_Ai30 = -tl.dot(
        b_Ai33,
        tl.dot(b_Akk30, b_Ai00) + tl.dot(b_Akk31, b_Ai10) + tl.dot(b_Akk32, b_Ai20),
    )

    # store full Akk inverse (without column beta — applied in w/u kernel)
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)), b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)), b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)), b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)), b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)), b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)), b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)), b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)), b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)), b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1))
    tl.store(tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)), b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))


# ---------------------------------------------------------------------------
# Kernel 4: w = A @ (beta * exp2(g) * k), u = A @ (beta * v), kg = k * exp2(g_last-g)
# ---------------------------------------------------------------------------
@triton.jit
def _wu_kernel(
    q,
    k,
    v,
    beta,
    A,
    g,
    w,
    u,
    kg,
    qg,
    scale,
    T,
    H: tl.constexpr,
    K: tl.constexpr,
    V: tl.constexpr,
    BT: tl.constexpr,
    BK: tl.constexpr,
    BV: tl.constexpr,
):
    i_t = tl.program_id(0)
    i_bh = tl.program_id(1)
    i_b = i_bh // H
    i_h = i_bh % H
    bos = i_b * T

    q += (bos * H + i_h) * K
    k += (bos * H + i_h) * K
    v += (bos * H + i_h) * V
    w += (bos * H + i_h) * K
    u += (bos * H + i_h) * V
    kg += (bos * H + i_h) * K
    qg += (bos * H + i_h) * K
    g += (bos * H + i_h) * K
    beta += bos * H + i_h
    A += (bos * H + i_h) * BT

    p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,))
    b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)

    p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
    b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.bfloat16)

    last_idx = i_t * BT + BT - 1

    for i_v in range(tl.cdiv(V, BV)):
        p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        p_u = tl.make_block_ptr(u, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32)
        b_vb = (b_v * b_b[:, None]).to(tl.bfloat16)
        b_u = tl.dot(b_A, b_vb)
        tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))

    for i_k in range(tl.cdiv(K, BK)):
        p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        p_w = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        p_kg = tl.make_block_ptr(kg, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        p_qg = tl.make_block_ptr(qg, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32)
        b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32)
        b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
        o_k = i_k * BK + tl.arange(0, BK)
        m_k = o_k < K
        b_gn = tl.load(g + last_idx * H * K + o_k, mask=m_k, other=0.0).to(tl.float32)

        b_eg = tl.math.exp2(b_g)
        b_kb = (b_k * b_b[:, None] * b_eg).to(tl.bfloat16)
        b_w = tl.dot(b_A, b_kb)
        tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))

        b_kg = b_k * tl.math.exp2(b_gn[None, :] - b_g)
        tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))

        b_qg = (b_q * b_eg * scale).to(tl.bfloat16)
        tl.store(p_qg, b_qg, boundary_check=(0, 1))


# ---------------------------------------------------------------------------
# Kernel 5: inter-chunk state recurrence
#   v_new = u - w @ S
#   store S at each chunk start into h
#   S = S * exp2(g_last) + kg.T @ v_new
# ---------------------------------------------------------------------------
@triton.jit
def _fwd_h_kernel(
    kg,
    w,
    u,
    g,
    v_new,
    h,
    T,
    H: tl.constexpr,
    K: tl.constexpr,
    V: tl.constexpr,
    BT: tl.constexpr,
    BV: tl.constexpr,
):
    i_v = tl.program_id(0)
    i_bh = tl.program_id(1)
    i_b = i_bh // H
    i_h = i_bh % H
    bos = i_b * T
    NT = tl.cdiv(T, BT)

    kg += (bos * H + i_h) * K
    w += (bos * H + i_h) * K
    u += (bos * H + i_h) * V
    v_new += (bos * H + i_h) * V
    g += (bos * H + i_h) * K
    h += (i_b * NT * H + i_h) * K * V

    b_h1 = tl.zeros([64, BV], dtype=tl.float32)
    b_h2 = tl.zeros([64, BV], dtype=tl.float32)

    for i_t in range(NT):
        p_h1 = tl.make_block_ptr(h + i_t * H * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
        p_h2 = tl.make_block_ptr(h + i_t * H * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
        tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))
        tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1))

        p_w1 = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0))
        p_w2 = tl.make_block_ptr(w, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0))
        b_w1 = tl.load(p_w1, boundary_check=(0, 1))
        b_w2 = tl.load(p_w2, boundary_check=(0, 1))
        b_v = tl.dot(b_w1, b_h1.to(b_w1.dtype)) + tl.dot(b_w2, b_h2.to(b_w2.dtype))

        p_u = tl.make_block_ptr(u, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        b_v = tl.load(p_u, boundary_check=(0, 1)) - b_v
        p_vn = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        tl.store(p_vn, b_v.to(p_vn.dtype.element_ty), boundary_check=(0, 1))

        last_idx = i_t * BT + BT - 1
        o_k1 = tl.arange(0, 64)
        o_k2 = 64 + o_k1
        b_gk1 = tl.load(g + last_idx * H * K + o_k1).to(tl.float32)
        b_gk2 = tl.load(g + last_idx * H * K + o_k2).to(tl.float32)
        b_h1 *= tl.math.exp2(b_gk1)[:, None]
        b_h2 *= tl.math.exp2(b_gk2)[:, None]

        b_v = b_v.to(kg.dtype.element_ty)
        p_kg1 = tl.make_block_ptr(kg, (K, T), (1, H * K), (0, i_t * BT), (64, BT), (0, 1))
        p_kg2 = tl.make_block_ptr(kg, (K, T), (1, H * K), (64, i_t * BT), (64, BT), (0, 1))
        b_kg1 = tl.load(p_kg1, boundary_check=(0, 1))
        b_kg2 = tl.load(p_kg2, boundary_check=(0, 1))
        b_h1 += tl.dot(b_kg1, b_v)
        b_h2 += tl.dot(b_kg2, b_v)


# ---------------------------------------------------------------------------
# Kernel 6: o = (q * exp2(g) * scale) @ h + Aqk @ v_new
# ---------------------------------------------------------------------------
@triton.jit
def _fwd_o_kernel(
    qg,
    v_new,
    A,
    h,
    o,
    T,
    H: tl.constexpr,
    K: tl.constexpr,
    V: tl.constexpr,
    BT: tl.constexpr,
    BV: tl.constexpr,
):
    i_v = tl.program_id(0)
    i_t = tl.program_id(1)
    i_bh = tl.program_id(2)
    i_b = i_bh // H
    i_h = i_bh % H
    bos = i_b * T
    NT = tl.cdiv(T, BT)

    qg += (bos * H + i_h) * K
    v_new += (bos * H + i_h) * V
    A += (bos * H + i_h) * BT
    o += (bos * H + i_h) * V
    h += (i_b * NT * H + i_h + i_t * H) * K * V

    p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
    b_A = tl.load(p_A, boundary_check=(0, 1))

    p_v = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
    b_v = tl.load(p_v, boundary_check=(0, 1))
    b_o = tl.dot(b_A, b_v)

    p_qg1 = tl.make_block_ptr(qg, (T, K), (H * K, 1), (i_t * BT, 0), (BT, 64), (1, 0))
    p_qg2 = tl.make_block_ptr(qg, (T, K), (H * K, 1), (i_t * BT, 64), (BT, 64), (1, 0))
    b_qg1 = tl.load(p_qg1, boundary_check=(0, 1))
    b_qg2 = tl.load(p_qg2, boundary_check=(0, 1))

    p_h1 = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
    p_h2 = tl.make_block_ptr(h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
    b_h1 = tl.load(p_h1, boundary_check=(0, 1))
    b_h2 = tl.load(p_h2, boundary_check=(0, 1))
    b_o += tl.dot(b_qg1, b_h1.to(tl.bfloat16))
    b_o += tl.dot(b_qg2, b_h2.to(tl.bfloat16))

    p_o = tl.make_block_ptr(o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))


# ---------------------------------------------------------------------------
# Host glue — buffer cache to avoid re-allocating intermediates every call
# ---------------------------------------------------------------------------
_BUF_CACHE: dict = {}


def _bufs(B, T, H, K, V, BT, BC, NT, device, dtype):
    key = (B, T, H, K, V, BT, BC, NT, device, dtype, 3)
    hit = _BUF_CACHE.get(key)
    if hit is not None:
        return hit
    # Aqk/Akk upper triangle must stay zero; kernels only rewrite lower/diag.
    bufs = {
        "g_cs": torch.empty(B, T, H, K, device=device, dtype=torch.float32),
        "Aqk": torch.zeros(B, T, H, BT, device=device, dtype=dtype),
        "Akkd": torch.empty(B, T, H, BC, device=device, dtype=torch.float32),
        "Akk": torch.zeros(B, T, H, BT, device=device, dtype=dtype),
        "w": torch.empty(B, T, H, K, device=device, dtype=dtype),
        "u": torch.empty(B, T, H, V, device=device, dtype=dtype),
        "kg": torch.empty(B, T, H, K, device=device, dtype=dtype),
        "qg": torch.empty(B, T, H, K, device=device, dtype=dtype),
        "v_new": torch.empty(B, T, H, V, device=device, dtype=dtype),
        "h": torch.empty(B, NT, H, K, V, device=device, dtype=dtype),
    }
    _BUF_CACHE[key] = bufs
    return bufs


def _kimi_delta_attn_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float,
    chunk_size: int = 64,
) -> torch.Tensor:
    B, T, H, K = q.shape
    V = v.shape[-1]
    BT = chunk_size
    assert T % BT == 0
    assert K == 128 and V == 128 and BT == 64, "specialized for K=V=128, BT=64"
    NT = T // BT
    BC = 16
    NC = BT // BC
    BK = 64
    BV = 64
    device = q.device
    dtype = v.dtype
    b = _bufs(B, T, H, K, V, BT, BC, NT, device, dtype)

    # 1. local cumsum of g → log2 space
    _local_cumsum_kernel[(B * H, NT, triton.cdiv(K, BK))](
        g, b["g_cs"], T, H=H, K=K, BT=BT, BK=BK, SCALE=LOG2E, num_warps=4,
    )

    # 2. diagonal sub-chunk Aqk/Akk + local solve
    _intra_diag_kernel[(B * H, NT * NC)](
        q, k, b["g_cs"], beta, b["Aqk"], b["Akkd"], scale,
        T, H=H, K=K, BT=BT, BC=BC, BK=128,
        num_warps=4,
    )

    # 3. off-diagonal + merge inverse
    _inter_solve_kernel[(NT, B * H)](
        q, k, b["g_cs"], beta, b["Aqk"], b["Akkd"], b["Akk"], scale,
        T, H=H, K=K, BT=BT, BC=BC, BK=BK,
        num_warps=4,
    )

    # 4. w, u, kg, qg
    _wu_kernel[(NT, B * H)](
        q, k, v, beta, b["Akk"], b["g_cs"], b["w"], b["u"], b["kg"], b["qg"], scale,
        T, H=H, K=K, V=V, BT=BT, BK=BK, BV=BV,
        num_warps=4, num_stages=2,
    )

    # 5. inter-chunk state
    _fwd_h_kernel[(triton.cdiv(V, BV), B * H)](
        b["kg"], b["w"], b["u"], b["g_cs"], b["v_new"], b["h"],
        T, H=H, K=K, V=V, BT=BT, BV=BV,
        num_warps=4, num_stages=2,
    )

    # 6. output
    o = torch.empty(B, T, H, V, device=device, dtype=dtype)
    _fwd_o_kernel[(triton.cdiv(V, BV), NT, B * H)](
        b["qg"], b["v_new"], b["Aqk"], b["h"], o,
        T, H=H, K=K, V=V, BT=BT, BV=BV,
        num_warps=4, num_stages=2,
    )
    return o


class Model(nn.Module):
    """KDA forward (chunk form). No learned parameters; all inputs are activations."""

    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: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        g: torch.Tensor,
        beta: torch.Tensor,
    ) -> torch.Tensor:
        return _kimi_delta_attn_fwd(q, k, v, g, beta, scale=self.scale, chunk_size=self.chunk_size)


# Module-level shape shims (overridden by check.py / benchmark.py per shape).
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]