"""Fused Kimi Delta Attention forward for fixed K=V=128.

The chunk formulation in the reference is an algebraic blocking of the token
recurrence.  This implementation fuses that recurrence and keeps each slice of
the 128x128 state in registers for the complete sequence.
"""
from __future__ import annotations

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


@triton.jit
def _kda_recurrent_fwd(
    q_ptr,
    k_ptr,
    v_ptr,
    g_ptr,
    beta_ptr,
    o_ptr,
    T: tl.constexpr,
    H: tl.constexpr,
    scale: tl.constexpr,
    BLOCK_V: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    bh = tl.program_id(0)
    vb = tl.program_id(1)
    b = bh // H
    h = bh - b * H

    rk = tl.arange(0, BLOCK_K)
    rv = vb * BLOCK_V + tl.arange(0, BLOCK_V)
    state = tl.zeros((BLOCK_K, BLOCK_V), dtype=tl.float32)

    q_bh = q_ptr + b * T * H * BLOCK_K + h * BLOCK_K
    k_bh = k_ptr + b * T * H * BLOCK_K + h * BLOCK_K
    g_bh = g_ptr + b * T * H * BLOCK_K + h * BLOCK_K
    v_bh = v_ptr + b * T * H * BLOCK_K + h * BLOCK_K
    o_bh = o_ptr + b * T * H * BLOCK_K + h * BLOCK_K
    beta_bh = beta_ptr + b * T * H + h

    for t in range(T):
        q = tl.load(q_bh + t * H * BLOCK_K + rk).to(tl.float32)
        k = tl.load(k_bh + t * H * BLOCK_K + rk).to(tl.float32)
        decay = tl.exp(tl.load(g_bh + t * H * BLOCK_K + rk))
        value = tl.load(v_bh + t * H * BLOCK_K + rv).to(tl.float32)
        beta = tl.load(beta_bh + t * H).to(tl.float32)

        state *= decay[:, None]
        residual = beta * (value - tl.sum(k[:, None] * state, axis=0))
        state += k[:, None] * residual[None, :]
        out = scale * tl.sum(q[:, None] * state, axis=0)
        tl.store(o_bh + t * H * BLOCK_K + rv, out)


@triton.jit
def _local_gate_scan(
    g_ptr,
    gc_ptr,
    T,
    H: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    """64-wide local prefix sum of the per-token log gates."""
    chunk = tl.program_id(0)
    bh = tl.program_id(1)
    kb = tl.program_id(2)
    b = bh // H
    h = bh - b * H
    rt = tl.arange(0, 64)[:, None]
    rk = kb * BLOCK_K + tl.arange(0, BLOCK_K)[None, :]
    t = chunk * 64 + rt
    offsets = ((b * T + t) * H + h) * 128 + rk
    x = tl.load(g_ptr + offsets)
    x = tl.cumsum(x, axis=0)
    tl.store(gc_ptr + offsets, x)


@triton.jit
def _build_chunk_matrices(
    q_ptr,
    k_ptr,
    gc_ptr,
    beta_ptr,
    a_ptr,
    aq_ptr,
    T,
    H: tl.constexpr,
    scale: tl.constexpr,
):
    """Build the two causal 64x64 products and the WY inverse."""
    chunk = tl.program_id(0)
    bh = tl.program_id(1)
    b = bh // H
    h = bh - b * H
    rc = tl.arange(0, 64)
    rk = tl.arange(0, 128)
    t = chunk * 64 + rc
    offs = ((b * T + t[:, None]) * H + h) * 128 + rk[None, :]
    q = tl.load(q_ptr + offs)
    k = tl.load(k_ptr + offs)
    gc = tl.load(gc_ptr + offs)
    ep = tl.exp(gc)
    kp = (k * ep).to(tl.bfloat16)
    kn = (k / ep).to(tl.bfloat16)
    qp = (q * ep).to(tl.bfloat16)

    kk = tl.dot(kp, tl.trans(kn))
    qk = tl.dot(qp, tl.trans(kn)) * scale
    row = rc[:, None]
    col = rc[None, :]
    beta = tl.load(beta_ptr + (b * T + t) * H + h).to(tl.float32)
    inv = -tl.where(row > col, kk * beta[:, None], 0.0)

    # Forward substitution for inv(I + tril(beta * K K^T, -1)).
    for i in range(1, 64):
        r = tl.sum(tl.where(row == i, inv, 0.0), axis=0)
        r = tl.where(rc < i, r, 0.0)
        r += tl.sum(r[:, None] * inv, axis=0)
        inv = tl.where(row == i, r[None, :], inv)
    inv = (inv + (row == col)) * beta[None, :]
    qk = tl.where(row >= col, qk, 0.0)

    base = (bh * (T // 64) + chunk) * 64 * 64
    m_offs = rc[:, None] * 64 + rc[None, :]
    tl.store(a_ptr + base + m_offs, inv.to(tl.bfloat16))
    tl.store(aq_ptr + base + m_offs, qk.to(tl.bfloat16))


@triton.jit
def _solve_diagonal_tiles(
    q_ptr,
    k_ptr,
    gc_ptr,
    beta_ptr,
    diag_ptr,
    aq_ptr,
    T,
    H: tl.constexpr,
    scale: tl.constexpr,
):
    chunk = tl.program_id(0)
    tile = tl.program_id(1)
    bh = tl.program_id(2)
    b = bh // H
    h = bh - b * H
    r = tl.arange(0, 16)
    d = tl.arange(0, 128)
    t = chunk * 64 + tile * 16 + r
    offs = ((b * T + t[:, None]) * H + h) * 128 + d[None, :]
    q = tl.load(q_ptr + offs)
    k = tl.load(k_ptr + offs)
    gc = tl.load(gc_ptr + offs)
    ep = tl.exp(gc)
    kp = (k * ep).to(tl.bfloat16)
    kn = (k / ep).to(tl.bfloat16)
    qp = (q * ep).to(tl.bfloat16)
    kk = tl.dot(kp, tl.trans(kn))
    qk = tl.dot(qp, tl.trans(kn)) * scale
    beta = tl.load(beta_ptr + (b * T + t) * H + h).to(tl.float32)
    row, col = r[:, None], r[None, :]
    inv = -tl.where(row > col, kk * beta[:, None], 0.0)
    for i in range(1, 16):
        current = tl.sum(tl.where(row == i, inv, 0.0), axis=0)
        current = tl.where(r < i, current, 0.0)
        current += tl.sum(current[:, None] * inv, axis=0)
        inv = tl.where(row == i, current[None, :], inv)
    inv += row == col

    db = ((bh * (T // 64) + chunk) * 4 + tile) * 16 * 16
    tl.store(diag_ptr + db + row * 16 + col, inv)
    ab = (bh * (T // 64) + chunk) * 64 * 64
    aoffs = (tile * 16 + row) * 64 + tile * 16 + col
    tl.store(aq_ptr + ab + aoffs, tl.where(row >= col, qk, 0.0).to(tl.bfloat16))


@triton.jit
def _compose_chunk_inverse(
    q_ptr,
    k_ptr,
    gc_ptr,
    beta_ptr,
    diag_ptr,
    a_ptr,
    aq_ptr,
    T,
    H: tl.constexpr,
    scale: tl.constexpr,
):
    chunk = tl.program_id(0)
    bh = tl.program_id(1)
    b = bh // H
    h = bh - b * H
    r = tl.arange(0, 16)
    d = tl.arange(0, 64)

    m10 = tl.zeros((16, 16), tl.float32)
    m20 = tl.zeros((16, 16), tl.float32)
    m21 = tl.zeros((16, 16), tl.float32)
    m30 = tl.zeros((16, 16), tl.float32)
    m31 = tl.zeros((16, 16), tl.float32)
    m32 = tl.zeros((16, 16), tl.float32)
    q10 = tl.zeros((16, 16), tl.float32)
    q20 = tl.zeros((16, 16), tl.float32)
    q21 = tl.zeros((16, 16), tl.float32)
    q30 = tl.zeros((16, 16), tl.float32)
    q31 = tl.zeros((16, 16), tl.float32)
    q32 = tl.zeros((16, 16), tl.float32)

    for kb in range(2):
        ds = kb * 64 + d
        t0 = chunk * 64 + r
        t1 = t0 + 16
        t2 = t0 + 32
        t3 = t0 + 48
        o0 = ((b * T + t0[:, None]) * H + h) * 128 + ds[None, :]
        o1 = ((b * T + t1[:, None]) * H + h) * 128 + ds[None, :]
        o2 = ((b * T + t2[:, None]) * H + h) * 128 + ds[None, :]
        o3 = ((b * T + t3[:, None]) * H + h) * 128 + ds[None, :]
        k0, k1 = tl.load(k_ptr + o0), tl.load(k_ptr + o1)
        k2, k3 = tl.load(k_ptr + o2), tl.load(k_ptr + o3)
        q1, q2 = tl.load(q_ptr + o1), tl.load(q_ptr + o2)
        q3 = tl.load(q_ptr + o3)
        g0, g1 = tl.load(gc_ptr + o0), tl.load(gc_ptr + o1)
        g2, g3 = tl.load(gc_ptr + o2), tl.load(gc_ptr + o3)
        p0 = (k0 * tl.exp(g0)).to(tl.bfloat16)
        p1 = (k1 * tl.exp(g1)).to(tl.bfloat16)
        p2 = (k2 * tl.exp(g2)).to(tl.bfloat16)
        p3 = (k3 * tl.exp(g3)).to(tl.bfloat16)
        n0 = (k0 * tl.exp(-g0)).to(tl.bfloat16)
        n1 = (k1 * tl.exp(-g1)).to(tl.bfloat16)
        n2 = (k2 * tl.exp(-g2)).to(tl.bfloat16)
        qp1 = (q1 * tl.exp(g1)).to(tl.bfloat16)
        qp2 = (q2 * tl.exp(g2)).to(tl.bfloat16)
        qp3 = (q3 * tl.exp(g3)).to(tl.bfloat16)
        n0t, n1t, n2t = tl.trans(n0), tl.trans(n1), tl.trans(n2)
        m10 += tl.dot(p1, n0t)
        m20 += tl.dot(p2, n0t)
        m21 += tl.dot(p2, n1t)
        m30 += tl.dot(p3, n0t)
        m31 += tl.dot(p3, n1t)
        m32 += tl.dot(p3, n2t)
        q10 += tl.dot(qp1, n0t)
        q20 += tl.dot(qp2, n0t)
        q21 += tl.dot(qp2, n1t)
        q30 += tl.dot(qp3, n0t)
        q31 += tl.dot(qp3, n1t)
        q32 += tl.dot(qp3, n2t)

    db = (bh * (T // 64) + chunk) * 4 * 16 * 16
    matoffs = r[:, None] * 16 + r[None, :]
    d00 = tl.load(diag_ptr + db + 0 * 256 + matoffs)
    d11 = tl.load(diag_ptr + db + 1 * 256 + matoffs)
    d22 = tl.load(diag_ptr + db + 2 * 256 + matoffs)
    d33 = tl.load(diag_ptr + db + 3 * 256 + matoffs)
    base_t = b * T + chunk * 64
    b0 = tl.load(beta_ptr + (base_t + r) * H + h).to(tl.float32)
    b1 = tl.load(beta_ptr + (base_t + 16 + r) * H + h).to(tl.float32)
    b2 = tl.load(beta_ptr + (base_t + 32 + r) * H + h).to(tl.float32)
    b3 = tl.load(beta_ptr + (base_t + 48 + r) * H + h).to(tl.float32)
    c10, c20, c21 = m10 * b1[:, None], m20 * b2[:, None], m21 * b2[:, None]
    c30, c31, c32 = m30 * b3[:, None], m31 * b3[:, None], m32 * b3[:, None]
    d10 = -tl.dot(tl.dot(d11, c10), d00)
    d21 = -tl.dot(tl.dot(d22, c21), d11)
    d20 = -tl.dot(d22, tl.dot(c20, d00) + tl.dot(c21, d10))
    d32 = -tl.dot(tl.dot(d33, c32), d22)
    d31 = -tl.dot(d33, tl.dot(c31, d11) + tl.dot(c32, d21))
    d30 = -tl.dot(d33, tl.dot(c30, d00) + tl.dot(c31, d10) + tl.dot(c32, d20))

    ab = (bh * (T // 64) + chunk) * 64 * 64
    # The full WY matrix is D @ diag(beta).  Only its lower blocks are consumed.
    tl.store(a_ptr + ab + (0 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (d00 * b0[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (1 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (d10 * b0[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (1 * 16 + r[:, None]) * 64 + 1 * 16 + r[None, :], (d11 * b1[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (2 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (d20 * b0[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (2 * 16 + r[:, None]) * 64 + 1 * 16 + r[None, :], (d21 * b1[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (2 * 16 + r[:, None]) * 64 + 2 * 16 + r[None, :], (d22 * b2[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (3 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (d30 * b0[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (3 * 16 + r[:, None]) * 64 + 1 * 16 + r[None, :], (d31 * b1[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (3 * 16 + r[:, None]) * 64 + 2 * 16 + r[None, :], (d32 * b2[None, :]).to(tl.bfloat16))
    tl.store(a_ptr + ab + (3 * 16 + r[:, None]) * 64 + 3 * 16 + r[None, :], (d33 * b3[None, :]).to(tl.bfloat16))

    tl.store(aq_ptr + ab + (1 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (q10 * scale).to(tl.bfloat16))
    tl.store(aq_ptr + ab + (2 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (q20 * scale).to(tl.bfloat16))
    tl.store(aq_ptr + ab + (2 * 16 + r[:, None]) * 64 + 1 * 16 + r[None, :], (q21 * scale).to(tl.bfloat16))
    tl.store(aq_ptr + ab + (3 * 16 + r[:, None]) * 64 + 0 * 16 + r[None, :], (q30 * scale).to(tl.bfloat16))
    tl.store(aq_ptr + ab + (3 * 16 + r[:, None]) * 64 + 1 * 16 + r[None, :], (q31 * scale).to(tl.bfloat16))
    tl.store(aq_ptr + ab + (3 * 16 + r[:, None]) * 64 + 2 * 16 + r[None, :], (q32 * scale).to(tl.bfloat16))


@triton.jit
def _form_w_u(
    k_ptr,
    v_ptr,
    gc_ptr,
    a_ptr,
    w_ptr,
    u_ptr,
    kg_ptr,
    T,
    H: tl.constexpr,
):
    chunk = tl.program_id(0)
    bh = tl.program_id(1)
    b = bh // H
    h = bh - b * H
    rc = tl.arange(0, 64)
    rd = tl.arange(0, 128)
    t = chunk * 64 + rc
    offs = ((b * T + t[:, None]) * H + h) * 128 + rd[None, :]
    k = tl.load(k_ptr + offs)
    v = tl.load(v_ptr + offs)
    gc = tl.load(gc_ptr + offs)
    ke = (k * tl.exp(gc)).to(tl.bfloat16)
    glast = tl.sum(tl.where(rc[:, None] == 63, gc, 0.0), axis=0)
    kg = (k * tl.exp(glast[None, :] - gc)).to(tl.bfloat16)
    base = (bh * (T // 64) + chunk) * 64 * 64
    a = tl.load(a_ptr + base + rc[:, None] * 64 + rc[None, :])
    a = tl.where(rc[:, None] >= rc[None, :], a, 0.0)
    w = tl.dot(a, ke)
    u = tl.dot(a, v)
    tl.store(w_ptr + offs, w.to(tl.bfloat16))
    tl.store(u_ptr + offs, u.to(tl.bfloat16))
    tl.store(kg_ptr + offs, kg)


@triton.jit
def _chunk_state_output(
    q_ptr,
    k_ptr,
    gc_ptr,
    w_ptr,
    u_ptr,
    aq_ptr,
    o_ptr,
    T,
    NT,
    H: tl.constexpr,
    BLOCK_V: tl.constexpr,
    scale: tl.constexpr,
):
    bh = tl.program_id(0)
    vb = tl.program_id(1)
    b = bh // H
    h = bh - b * H
    rc = tl.arange(0, 64)
    rk = tl.arange(0, 128)
    rv = vb * BLOCK_V + tl.arange(0, BLOCK_V)
    state = tl.zeros((128, BLOCK_V), tl.float32)

    for chunk in range(NT):
        t = chunk * 64 + rc
        qk_offs = ((b * T + t[:, None]) * H + h) * 128 + rk[None, :]
        v_offs = ((b * T + t[:, None]) * H + h) * 128 + rv[None, :]
        q = tl.load(q_ptr + qk_offs)
        k = tl.load(k_ptr + qk_offs)
        gc = tl.load(gc_ptr + qk_offs)
        qg = (q * tl.exp(gc) * scale).to(tl.bfloat16)
        kg = (k * tl.exp(-gc)).to(tl.bfloat16)
        w = tl.load(w_ptr + qk_offs)
        u = tl.load(u_ptr + v_offs)
        state_bf = state.to(tl.bfloat16)
        vn = u.to(tl.float32) - tl.dot(w, state_bf)
        out = tl.dot(qg, state_bf)

        base = (bh * NT + chunk) * 64 * 64
        aq = tl.load(aq_ptr + base + rc[:, None] * 64 + rc[None, :])
        vn_bf = vn.to(tl.bfloat16)
        out += tl.dot(aq, vn_bf)
        tl.store(o_ptr + v_offs, out.to(tl.bfloat16))

        state += tl.dot(tl.trans(kg), vn_bf)
        g_last = tl.sum(tl.where(rc[:, None] == 63, gc, 0.0), axis=0)
        state *= tl.exp(g_last)[:, None]


@triton.jit
def _chunk_state_scan(
    kg_ptr,
    gc_ptr,
    w_ptr,
    u_ptr,
    h_ptr,
    vn_ptr,
    T,
    NT,
    H: tl.constexpr,
    BLOCK_V: tl.constexpr,
):
    """The only sequential portion: one short scan over sequence chunks."""
    bh = tl.program_id(0)
    vb = tl.program_id(1)
    b = bh // H
    hidx = bh - b * H
    rc = tl.arange(0, 64)
    rk = tl.arange(0, 64)
    rv = vb * BLOCK_V + tl.arange(0, BLOCK_V)
    state0 = tl.zeros((64, BLOCK_V), tl.float32)
    state1 = tl.zeros((64, BLOCK_V), tl.float32)

    for chunk in range(NT):
        hbase = (bh * NT + chunk) * 128 * 128
        tl.store(h_ptr + hbase + rk[:, None] * 128 + rv[None, :], state0)
        tl.store(h_ptr + hbase + (rk[:, None] + 64) * 128 + rv[None, :], state1)

        t = chunk * 64 + rc
        off0 = ((b * T + t[:, None]) * H + hidx) * 128 + rk[None, :]
        off1 = off0 + 64
        voff = ((b * T + t[:, None]) * H + hidx) * 128 + rv[None, :]
        w0 = tl.load(w_ptr + off0)
        w1 = tl.load(w_ptr + off1)
        pred = tl.dot(w0, state0.to(tl.bfloat16))
        pred += tl.dot(w1, state1.to(tl.bfloat16))
        vn = tl.load(u_ptr + voff).to(tl.float32) - pred
        vn_bf = vn.to(tl.bfloat16)
        tl.store(vn_ptr + voff, vn_bf)

        kg0 = tl.load(kg_ptr + off0)
        kg1 = tl.load(kg_ptr + off1)
        last0 = tl.load(
            gc_ptr + ((b * T + chunk * 64 + 63) * H + hidx) * 128 + rk
        )
        last1 = tl.load(
            gc_ptr + ((b * T + chunk * 64 + 63) * H + hidx) * 128 + 64 + rk
        )
        state0 *= tl.exp(last0)[:, None]
        state1 *= tl.exp(last1)[:, None]
        state0 += tl.dot(tl.trans(kg0), vn_bf)
        state1 += tl.dot(tl.trans(kg1), vn_bf)


@triton.jit
def _parallel_chunk_output(
    q_ptr,
    gc_ptr,
    vn_ptr,
    h_ptr,
    aq_ptr,
    o_ptr,
    T,
    NT,
    H: tl.constexpr,
    BLOCK_V: tl.constexpr,
    scale: tl.constexpr,
):
    """All chunk reads run in parallel after the state scan."""
    chunk = tl.program_id(0)
    bh = tl.program_id(1)
    vb = tl.program_id(2)
    b = bh // H
    hidx = bh - b * H
    rc = tl.arange(0, 64)
    rk = tl.arange(0, 64)
    rv = vb * BLOCK_V + tl.arange(0, BLOCK_V)
    t = chunk * 64 + rc
    off0 = ((b * T + t[:, None]) * H + hidx) * 128 + rk[None, :]
    off1 = off0 + 64
    voff = ((b * T + t[:, None]) * H + hidx) * 128 + rv[None, :]

    q0 = tl.load(q_ptr + off0)
    q1 = tl.load(q_ptr + off1)
    g0 = tl.load(gc_ptr + off0)
    g1 = tl.load(gc_ptr + off1)
    qg0 = (q0 * tl.exp(g0)).to(tl.bfloat16)
    qg1 = (q1 * tl.exp(g1)).to(tl.bfloat16)
    hbase = (bh * NT + chunk) * 128 * 128
    hs0 = tl.load(h_ptr + hbase + rk[:, None] * 128 + rv[None, :])
    hs1 = tl.load(h_ptr + hbase + (rk[:, None] + 64) * 128 + rv[None, :])
    out = (tl.dot(qg0, hs0) + tl.dot(qg1, hs1)) * scale

    abase = (bh * NT + chunk) * 64 * 64
    aq = tl.load(aq_ptr + abase + rc[:, None] * 64 + rc[None, :])
    aq = tl.where(rc[:, None] >= rc[None, :], aq, 0.0)
    vn = tl.load(vn_ptr + voff)
    out += tl.dot(aq, vn)
    tl.store(o_ptr + voff, out.to(tl.bfloat16))


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)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        g: torch.Tensor,
        beta: torch.Tensor,
    ) -> torch.Tensor:
        assert self.K == 128 and self.V == 128
        nt = self.T // 64
        gc = torch.empty_like(g)
        a = torch.empty(
            self.B * self.H, nt, 64, 64, dtype=torch.bfloat16, device=q.device
        )
        aq = torch.empty_like(a)
        diag = torch.empty(
            self.B * self.H, nt, 4, 16, 16, dtype=torch.float32, device=q.device
        )
        w = torch.empty_like(k)
        u = torch.empty_like(v)
        kg = torch.empty_like(k)
        h = torch.empty(
            self.B * self.H, nt, 128, 128, dtype=torch.bfloat16, device=q.device
        )
        vn = torch.empty_like(v)
        o = torch.empty_like(v)

        _local_gate_scan[(nt, self.B * self.H, 4)](
            g, gc, self.T, H=self.H, BLOCK_K=32, num_warps=4
        )
        _solve_diagonal_tiles[(nt, 4, self.B * self.H)](
            q,
            k,
            gc,
            beta,
            diag,
            aq,
            self.T,
            H=self.H,
            scale=self.scale,
            num_warps=2,
            num_stages=2,
        )
        _compose_chunk_inverse[(nt, self.B * self.H)](
            q,
            k,
            gc,
            beta,
            diag,
            a,
            aq,
            self.T,
            H=self.H,
            scale=self.scale,
            num_warps=2 if self.B * self.H == 4 else 1,
            num_stages=2,
        )
        _form_w_u[(nt, self.B * self.H)](
            k, v, gc, a, w, u, kg, self.T, H=self.H, num_warps=8, num_stages=2
        )
        state_bv = 16
        _chunk_state_scan[(self.B * self.H, self.V // state_bv)](
            kg,
            gc,
            w,
            u,
            h,
            vn,
            self.T,
            nt,
            H=self.H,
            BLOCK_V=state_bv,
            num_warps=2,
            num_stages=2,
        )
        out_bv = 128
        _parallel_chunk_output[(nt, self.B * self.H, self.V // out_bv)](
            q,
            gc,
            vn,
            h,
            aq,
            o,
            self.T,
            nt,
            H=self.H,
            BLOCK_V=out_bv,
            scale=self.scale,
            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]