"""Kimi Delta Attention forward (chunk form) -- hand-written CUDA kernels for B200.

Three kernels (mma.sync.m16n8k16 bf16 tensor cores, cp.async smem pipelines,
ldmatrix fragment loads), written from scratch for BT=64, K=V=128:

  K1 prep (one block per chunk*head): per-column log2 cumsum of the decay,
     exp2 factor tables, decay-weighted q/k variants, the intra-chunk
     attention matrix Aqk = tril(qq kk^T), the strictly-lower KK^T matrix,
     an exact blocked inverse of (I + beta tril(K e^dg K^T)) (per-warp serial
     16x16 diagonal inverses + block substitution on tensor cores), and the
     WY-representation vectors w (negated) and u.
  K2 scan (two blocks per head): the sequential inter-chunk state recurrence
     h <- lam*h + kg^T (u - w h), state h^T kept in fp32 accumulator
     fragments, per-chunk states and delta-corrected values stored transposed.
  K3 out (one block per chunk*head): o = qg @ h + Aqk @ v_new.

The Model wraps the three launches in a cached CUDA graph so steady-state
calls replay with zero launch overhead.
"""
from __future__ import annotations

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

OP_TYPE = "linear_attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

_CUDA_SRC = r"""
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cstdint>

using bf16 = __nv_bfloat16;

// Common helpers: mma.sync m16n8k16 bf16, ldmatrix, cp.async, stmatrix.

#define DEVI __device__ __forceinline__

DEVI uint32_t smem_u32(const void* p) {
    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

// ---- cp.async (16B) ----
DEVI void cp_async16(void* dst_smem, const void* src_gmem) {
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::
                 "r"(smem_u32(dst_smem)), "l"(src_gmem));
}
DEVI void cp_async_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N>
DEVI void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); }

// ---- ldmatrix ----
DEVI void ldmatrix_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, uint32_t addr) {
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                 : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(addr));
}
DEVI void ldmatrix_x4_trans(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, uint32_t addr) {
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                 : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(addr));
}
DEVI void ldmatrix_x2(uint32_t& r0, uint32_t& r1, uint32_t addr) {
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
                 : "=r"(r0), "=r"(r1) : "r"(addr));
}
DEVI void ldmatrix_x2_trans(uint32_t& r0, uint32_t& r1, uint32_t addr) {
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];\n"
                 : "=r"(r0), "=r"(r1) : "r"(addr));
}

// ---- stmatrix (SM90+) ----
DEVI void stmatrix_x4(uint32_t addr, uint32_t r0, uint32_t r1, uint32_t r2, uint32_t r3) {
    asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1,%2,%3,%4};\n" ::
                 "r"(addr), "r"(r0), "r"(r1), "r"(r2), "r"(r3));
}
DEVI void stmatrix_x4_trans(uint32_t addr, uint32_t r0, uint32_t r1, uint32_t r2, uint32_t r3) {
    asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1,%2,%3,%4};\n" ::
                 "r"(addr), "r"(r0), "r"(r1), "r"(r2), "r"(r3));
}

// ---- mma.sync m16n8k16 bf16 ----
// D (f32 4regs) = A (bf16 4regs) @ B (bf16 2regs) + C
DEVI void mma16816(float& d0, float& d1, float& d2, float& d3,
                   uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                   uint32_t b0, uint32_t b1,
                   float c0, float c1, float c2, float c3) {
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
        : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1),
          "f"(c0), "f"(c1), "f"(c2), "f"(c3));
}

// pack two f32 into bf16x2 register
DEVI uint32_t pack_bf16x2(float lo, float hi) {
    uint32_t r;
    asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(r) : "f"(hi), "f"(lo));
    return r;
}

DEVI void ldg128(uint4& v, const void* p) {
    v = *reinterpret_cast<const uint4*>(p);
}
DEVI void stg128(void* p, const uint4& v) {
    *reinterpret_cast<uint4*>(p) = v;
}

// KDA per-chunk preparation kernel (K1), mma.sync bf16 + smem blocked solve.
//
// Per block = one (chunk, batch*head), 128 threads:
//   1. cp.async loads: q,k,g,v(->sW),beta
//   2. column pass (thread=column): gc cumsum (log2), ep=exp2(gc), em=exp2(-gc)
//      (bf16, into sEP/sEM), lam=exp2(gn) (vector + global)
//   3. row pass (vectorized, 4 cols/lane): kq=k*ep, kk=k*em, qq=q*scale*ep (smem)
//      qg=qq, kg=kk*lam -> global (coalesced); RHS-u = beta*v (sW in place)
//   4. mma: Aqk = tril(qq @ kk^T) -> staged in smem -> global
//      negA0 = -beta_row*stril(kq @ kk^T) -> sA bf16 + fp32 diag -> sD
//   5. Dinv per warp (register-shfl sweep), Ltilde = Dinv_i @ negA0_im
//   6. two solve passes (u, then wneg with RHS=-beta*kq), 4 rounds each; copyout
//
// Outputs (chunk-tiled): wneg,u,kg,qg (B,H,NT,64,128); Aqk (B,H,NT,64,64);
//                        lam (B,H,NT,128) fp32


constexpr float RCP_LN2f_P = 1.4426950408889634f;

constexpr int PQ_P = 136;    // bf16 tile pitch
constexpr int PG_P = 136;    // f32 g pitch (matches alias size)
constexpr int PA_P = 72;     // A block pitch
constexpr int PDB_P = 24;    // Dinv bf16 pitch

constexpr int OFF_Q_P = 0;                        // 64*PQ_P bf16 (qq)
constexpr int OFF_K_P = OFF_Q_P + 64 * PQ_P * 2;      // kq
constexpr int OFF_KK_P = OFF_K_P + 64 * PQ_P * 2;     // kk
constexpr int OFF_W_P = OFF_KK_P + 64 * PQ_P * 2;     // RHS / W
constexpr int OFF_G_P = OFF_W_P + 64 * PQ_P * 2;      // 64*PG_P f32 g / packed (ep,em) pairs
constexpr int SZ_G_P = 64 * PG_P * 4;
// column pass packs (ep,em) bf16x2 into the same 4B word of g it just read.
// After the row pass the region is dead and hosts Aqk staging / sD / sDb / sA.
constexpr int OFF_AST_P = OFF_G_P;                  // Aqk staging 64*PA_P bf16
constexpr int OFF_D_P = OFF_G_P;                    // diag f32 (after Aqk copyout)
constexpr int OFF_DB_P = OFF_D_P + 4 * 16 * 17 * 4;
constexpr int OFF_A_P = OFF_G_P + 64 * PQ_P * 2;      // negA0 / Ltilde (second half)
static_assert(4 * 16 * 17 * 4 + 4 * 16 * PDB_P * 2 <= 64 * PQ_P * 2, "d/db alias");
static_assert(64 * PQ_P * 2 + 64 * PA_P * 2 <= SZ_G_P, "A alias");
constexpr int OFF_VEC_P = OFF_G_P + SZ_G_P;
constexpr int OFF_LAMS_P = OFF_VEC_P;               // 128 f32
constexpr int OFF_BETA_P = OFF_LAMS_P + 128 * 4;    // 64 f32
constexpr int SMEM_TOTAL_P = OFF_BETA_P + 64 * 4;

__global__ void __launch_bounds__(128, 2)
kda_prep_kernel(
    const bf16* __restrict__ q,
    const bf16* __restrict__ k,
    const bf16* __restrict__ v,
    const float* __restrict__ g,
    const bf16* __restrict__ beta,
    bf16* __restrict__ wneg,
    bf16* __restrict__ u,
    bf16* __restrict__ kg,
    bf16* __restrict__ qg,
    bf16* __restrict__ Aqk,
    float* __restrict__ lam,
    int* __restrict__ flag_prep,
    float scale,
    int T, int H, int NT) {
    constexpr int BT = 64, K = 128;
    extern __shared__ char smem_raw[];
    bf16* sQ = reinterpret_cast<bf16*>(smem_raw + OFF_Q_P);
    bf16* sK = reinterpret_cast<bf16*>(smem_raw + OFF_K_P);
    bf16* sKK = reinterpret_cast<bf16*>(smem_raw + OFF_KK_P);
    bf16* sW = reinterpret_cast<bf16*>(smem_raw + OFF_W_P);
    float* sG = reinterpret_cast<float*>(smem_raw + OFF_G_P);
    uint32_t* sEPM = reinterpret_cast<uint32_t*>(smem_raw + OFF_G_P);   // packed (ep,em)
    bf16* sAst = reinterpret_cast<bf16*>(smem_raw + OFF_AST_P);
    float* sD = reinterpret_cast<float*>(smem_raw + OFF_D_P);
    bf16* sDb = reinterpret_cast<bf16*>(smem_raw + OFF_DB_P);
    bf16* sA = reinterpret_cast<bf16*>(smem_raw + OFF_A_P);
    float* sLam = reinterpret_cast<float*>(smem_raw + OFF_LAMS_P);
    float* sBeta = reinterpret_cast<float*>(smem_raw + OFF_BETA_P);

    const int i_t = blockIdx.y, i_bh = blockIdx.x;   // bh fastest: chunk 0 of all heads first
    const int i_b = i_bh / H, i_h = i_bh % H;
    const int tid = threadIdx.x;
    const int warp = tid >> 5, lane = tid & 31;

    const int strideT_K = H * K;
    const int64_t tok0 = (int64_t)i_b * T + i_t * BT;
    const bf16* gq = q + tok0 * strideT_K + i_h * K;
    const bf16* gk = k + tok0 * strideT_K + i_h * K;
    const bf16* gv = v + tok0 * strideT_K + i_h * K;
    const float* gg = g + tok0 * strideT_K + i_h * K;
    const bf16* gbeta = beta + tok0 * H + i_h;

    const int64_t cbase = ((int64_t)i_bh * NT + i_t);
    bf16* gwneg = wneg + cbase * BT * K;
    bf16* gu = u + cbase * BT * K;
    bf16* gkg = kg + cbase * BT * K;
    bf16* gqg = qg + cbase * BT * K;
    bf16* gAqk = Aqk + cbase * BT * BT;
    float* glam = lam + cbase * K;

    // ------------------------------------------------------------------ loads
    {
        const int rg = tid >> 5, og = (tid & 31) * 4;
        #pragma unroll
        for (int i = 0; i < 16; i++)
            cp_async16(sG + (rg + i * 4) * PG_P + og, gg + (int64_t)(rg + i * 4) * strideT_K + og);
        cp_async_commit();
        const int r = tid >> 4, o = (tid & 15) * 8;
        #pragma unroll
        for (int i = 0; i < 8; i++) {
            cp_async16(sQ + (r + i * 8) * PQ_P + o, gq + (int64_t)(r + i * 8) * strideT_K + o);
            cp_async16(sK + (r + i * 8) * PQ_P + o, gk + (int64_t)(r + i * 8) * strideT_K + o);
            cp_async16(sW + (r + i * 8) * PQ_P + o, gv + (int64_t)(r + i * 8) * strideT_K + o);
        }
        cp_async_commit();
    }
    float betav = 0.f;
    if (tid < 64) betav = __bfloat162float(gbeta[(int64_t)tid * H]);
    cp_async_wait<1>();     // g ready; q/k/v still in flight
    __syncthreads();
    if (tid < 64) sBeta[tid] = betav;

    // ------------------------------------------------- column pass (gc, ep, em, lam)
    {
        const int c = tid;
        float acc = 0.f;
        float epv = 1.f;
        #pragma unroll
        for (int r = 0; r < BT; r++) {
            acc += sG[r * PG_P + c] * RCP_LN2f_P;
            epv = exp2f(acc);
            const float emv = exp2f(-acc);
            // pack (ep,em) into the exact word just read: thread-local, race-free
            sEPM[r * PG_P + c] = pack_bf16x2(epv, emv);
        }
        sLam[c] = epv;
        glam[c] = epv;
    }
    cp_async_wait<0>();     // q/k/v ready
    __syncthreads();

    // ------------------------------------------------- row pass (vectorized)
    {
        // warp w -> rows 16w..16w+15 ; lane -> 4 cols [4*lane, 4*lane+4)
        const int c4 = 4 * lane;
        float lam4[4];
        #pragma unroll
        for (int e = 0; e < 4; e++) lam4[e] = sLam[c4 + e];
        #pragma unroll
        for (int rr = 0; rr < 16; rr++) {
            const int r = warp * 16 + rr;
            const float bet = sBeta[r];
            uint2 kp = *reinterpret_cast<const uint2*>(sK + r * PQ_P + c4);
            uint2 qp = *reinterpret_cast<const uint2*>(sQ + r * PQ_P + c4);
            uint2 vp = *reinterpret_cast<const uint2*>(sW + r * PQ_P + c4);
            uint4 epm = *reinterpret_cast<const uint4*>(sEPM + r * PG_P + c4);
            float kf[4], qf[4], vf[4], ep[4], em[4];
            {
                float2 t;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&kp.x)); kf[0] = t.x; kf[1] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&kp.y)); kf[2] = t.x; kf[3] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&qp.x)); qf[0] = t.x; qf[1] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&qp.y)); qf[2] = t.x; qf[3] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&vp.x)); vf[0] = t.x; vf[1] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&vp.y)); vf[2] = t.x; vf[3] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&epm.x)); ep[0] = t.x; em[0] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&epm.y)); ep[1] = t.x; em[1] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&epm.z)); ep[2] = t.x; em[2] = t.y;
                t = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&epm.w)); ep[3] = t.x; em[3] = t.y;
            }
            float kq[4], kk[4], qq[4], kgv[4];
            #pragma unroll
            for (int e = 0; e < 4; e++) {
                kq[e] = kf[e] * ep[e];
                kk[e] = kf[e] * em[e];
                qq[e] = qf[e] * scale * ep[e];
                kgv[e] = kk[e] * lam4[e];
            }
            uint2 kqp = {pack_bf16x2(kq[0], kq[1]), pack_bf16x2(kq[2], kq[3])};
            uint2 kkp = {pack_bf16x2(kk[0], kk[1]), pack_bf16x2(kk[2], kk[3])};
            uint2 qqp = {pack_bf16x2(qq[0], qq[1]), pack_bf16x2(qq[2], qq[3])};
            uint2 kgp = {pack_bf16x2(kgv[0], kgv[1]), pack_bf16x2(kgv[2], kgv[3])};
            uint2 vbp = {pack_bf16x2(vf[0] * bet, vf[1] * bet), pack_bf16x2(vf[2] * bet, vf[3] * bet)};
            *reinterpret_cast<uint2*>(sK + r * PQ_P + c4) = kqp;
            *reinterpret_cast<uint2*>(sKK + r * PQ_P + c4) = kkp;
            *reinterpret_cast<uint2*>(sQ + r * PQ_P + c4) = qqp;
            *reinterpret_cast<uint2*>(sW + r * PQ_P + c4) = vbp;
            *reinterpret_cast<uint2*>(gqg + r * K + c4) = qqp;
            *reinterpret_cast<uint2*>(gkg + r * K + c4) = kgp;
        }
    }
    __syncthreads();

    const int vw = warp * 16;
    // --------------------------------------------------------- mma pass 1: Aqk
    {
        uint32_t aQ[8][4];
        int rowa = vw + (lane & 15), cola = ((lane >> 4) & 1) * 8;
        #pragma unroll
        for (int kk8 = 0; kk8 < 8; kk8++)
            ldmatrix_x4(aQ[kk8][0], aQ[kk8][1], aQ[kk8][2], aQ[kk8][3],
                        smem_u32(sQ + rowa * PQ_P + kk8 * 16 + cola));
        float accQ[8][4] = {};
        #pragma unroll
        for (int kk8 = 0; kk8 < 8; kk8 += 2) {
            uint32_t bw[8][4];
            #pragma unroll
            for (int n = 0; n < 8; n++) {
                int row = n * 8 + (lane & 7), col = kk8 * 16 + ((lane >> 3) & 3) * 8;
                ldmatrix_x4(bw[n][0], bw[n][1], bw[n][2], bw[n][3], smem_u32(sKK + row * PQ_P + col));
            }
            #pragma unroll
            for (int n = 0; n < 8; n++)
                mma16816(accQ[n][0], accQ[n][1], accQ[n][2], accQ[n][3],
                         aQ[kk8][0], aQ[kk8][1], aQ[kk8][2], aQ[kk8][3], bw[n][0], bw[n][1],
                         accQ[n][0], accQ[n][1], accQ[n][2], accQ[n][3]);
            #pragma unroll
            for (int n = 0; n < 8; n++)
                mma16816(accQ[n][0], accQ[n][1], accQ[n][2], accQ[n][3],
                         aQ[kk8 + 1][0], aQ[kk8 + 1][1], aQ[kk8 + 1][2], aQ[kk8 + 1][3], bw[n][2], bw[n][3],
                         accQ[n][0], accQ[n][1], accQ[n][2], accQ[n][3]);
        }
        __syncthreads();   // EP/EM fully consumed (row pass done long ago; barrier for staging reuse)
        const int r0 = vw + (lane >> 2);
        #pragma unroll
        for (int n = 0; n < 8; n++) {
            const int c0 = n * 8 + 2 * (lane & 3);
            accQ[n][0] = (r0 >= c0) ? accQ[n][0] : 0.f;
            accQ[n][1] = (r0 >= c0 + 1) ? accQ[n][1] : 0.f;
            accQ[n][2] = (r0 + 8 >= c0) ? accQ[n][2] : 0.f;
            accQ[n][3] = (r0 + 8 >= c0 + 1) ? accQ[n][3] : 0.f;
            *reinterpret_cast<uint32_t*>(sAst + r0 * PA_P + c0) = pack_bf16x2(accQ[n][0], accQ[n][1]);
            *reinterpret_cast<uint32_t*>(sAst + (r0 + 8) * PA_P + c0) = pack_bf16x2(accQ[n][2], accQ[n][3]);
        }
    }
    __syncthreads();
    // coalesced Aqk copyout: 64 rows x 128B
    {
        const int r = tid >> 3, o = (tid & 7) * 8;
        #pragma unroll
        for (int i = 0; i < 4; i++) {
            uint4 val = *reinterpret_cast<uint4*>(sAst + (r + i * 16) * PA_P + o);
            stg128(gAqk + (r + i * 16) * BT + o, val);
        }
    }
    __syncthreads();   // sAst -> becomes sD/sDb space

    // --------------------------------------------------------- mma pass 2: negA0
    {
        uint32_t aK[8][4];
        int rowa = vw + (lane & 15), cola = ((lane >> 4) & 1) * 8;
        #pragma unroll
        for (int kk8 = 0; kk8 < 8; kk8++)
            ldmatrix_x4(aK[kk8][0], aK[kk8][1], aK[kk8][2], aK[kk8][3],
                        smem_u32(sK + rowa * PQ_P + kk8 * 16 + cola));
        float accK[8][4] = {};
        #pragma unroll
        for (int kk8 = 0; kk8 < 8; kk8 += 2) {
            uint32_t bw[8][4];
            #pragma unroll
            for (int n = 0; n < 8; n++) {
                int row = n * 8 + (lane & 7), col = kk8 * 16 + ((lane >> 3) & 3) * 8;
                ldmatrix_x4(bw[n][0], bw[n][1], bw[n][2], bw[n][3], smem_u32(sKK + row * PQ_P + col));
            }
            #pragma unroll
            for (int n = 0; n < 8; n++)
                mma16816(accK[n][0], accK[n][1], accK[n][2], accK[n][3],
                         aK[kk8][0], aK[kk8][1], aK[kk8][2], aK[kk8][3], bw[n][0], bw[n][1],
                         accK[n][0], accK[n][1], accK[n][2], accK[n][3]);
            #pragma unroll
            for (int n = 0; n < 8; n++)
                mma16816(accK[n][0], accK[n][1], accK[n][2], accK[n][3],
                         aK[kk8 + 1][0], aK[kk8 + 1][1], aK[kk8 + 1][2], aK[kk8 + 1][3], bw[n][2], bw[n][3],
                         accK[n][0], accK[n][1], accK[n][2], accK[n][3]);
        }
        const int r0 = vw + (lane >> 2);
        const float b0 = sBeta[r0], b1 = sBeta[r0 + 8];
        #pragma unroll
        for (int n = 0; n < 8; n++) {
            const int c0 = n * 8 + 2 * (lane & 3);
            accK[n][0] = (r0 > c0) ? -b0 * accK[n][0] : 0.f;
            accK[n][1] = (r0 > c0 + 1) ? -b0 * accK[n][1] : 0.f;
            accK[n][2] = (r0 + 8 > c0) ? -b1 * accK[n][2] : 0.f;
            accK[n][3] = (r0 + 8 > c0 + 1) ? -b1 * accK[n][3] : 0.f;
            *reinterpret_cast<uint32_t*>(sA + r0 * PA_P + c0) = pack_bf16x2(accK[n][0], accK[n][1]);
            *reinterpret_cast<uint32_t*>(sA + (r0 + 8) * PA_P + c0) = pack_bf16x2(accK[n][2], accK[n][3]);
        }
        float* sDw = sD + warp * 16 * 17;
        #pragma unroll
        for (int nn = 0; nn < 2; nn++) {
            const int n = 2 * warp + nn;
            const int c0 = n * 8 + 2 * (lane & 3) - vw;
            const int rl = (lane >> 2);
            sDw[rl * 17 + c0] = accK[n][0];
            sDw[rl * 17 + c0 + 1] = accK[n][1];
            sDw[(rl + 8) * 17 + c0] = accK[n][2];
            sDw[(rl + 8) * 17 + c0 + 1] = accK[n][3];
        }
    }
    __syncwarp();

    // ---------------------------------- Dinv (register sweep with shfl broadcast)
    {
        float* sDw = sD + warp * 16 * 17;
        const int row = lane >> 1, cg = (lane & 1) * 8;
        float a[8], x[8];
        #pragma unroll
        for (int e = 0; e < 8; e++) {
            a[e] = sDw[row * 17 + cg + e];
            x[e] = (cg + e == row) ? 1.f : 0.f;
        }
        #pragma unroll
        for (int j = 0; j < 15; j++) {
            // a_j = negA0[row][j]: reg j%8 of lane 2*row + (j>=8)
            float aj = __shfl_sync(0xffffffffu, a[j & 7], 2 * row + (j >> 3));
            const int src = 2 * j + (lane & 1);
            float xj[8];
            #pragma unroll
            for (int e = 0; e < 8; e++) xj[e] = __shfl_sync(0xffffffffu, x[e], src);
            #pragma unroll
            for (int e = 0; e < 8; e++) x[e] += aj * xj[e];
        }
        bf16* sDbw = sDb + warp * 16 * PDB_P;
        #pragma unroll
        for (int e = 0; e < 8; e += 2)
            *reinterpret_cast<uint32_t*>(sDbw + row * PDB_P + cg + e) = pack_bf16x2(x[e], x[e + 1]);
    }
    __syncthreads();

    // --------------------------------------------- Ltilde_im = Dinv_i @ negA0_im
    {
        const int nblk = (warp == 3) ? 0 : 2;
        const int bi[2] = {warp == 0 ? 1 : 2, 3};
        const int bm[2] = {warp == 0 ? 0 : (warp == 1 ? 0 : 1), warp};
        #pragma unroll
        for (int z = 0; z < 2; z++) {
            if (z >= nblk) break;
            const int i = bi[z], m = bm[z];
            uint32_t da[4];
            {
                int rowa = (lane & 15), cola = ((lane >> 4) & 1) * 8;
                ldmatrix_x4(da[0], da[1], da[2], da[3], smem_u32(sDb + i * 16 * PDB_P + rowa * PDB_P + cola));
            }
            float acc[2][4] = {};
            #pragma unroll
            for (int s = 0; s < 2; s++) {
                uint32_t b0, b1;
                int rowb = (lane & 15);
                ldmatrix_x2_trans(b0, b1, smem_u32(sA + (i * 16 + rowb) * PA_P + m * 16 + s * 8));
                mma16816(acc[s][0], acc[s][1], acc[s][2], acc[s][3], da[0], da[1], da[2], da[3], b0, b1,
                         acc[s][0], acc[s][1], acc[s][2], acc[s][3]);
            }
            __syncwarp();
            const int rl = (lane >> 2), cl = 2 * (lane & 3);
            #pragma unroll
            for (int s = 0; s < 2; s++) {
                *reinterpret_cast<uint32_t*>(sA + (i * 16 + rl) * PA_P + m * 16 + s * 8 + cl) =
                    pack_bf16x2(acc[s][0], acc[s][1]);
                *reinterpret_cast<uint32_t*>(sA + (i * 16 + rl + 8) * PA_P + m * 16 + s * 8 + cl) =
                    pack_bf16x2(acc[s][2], acc[s][3]);
            }
        }
    }
    __syncthreads();

    // ---------------------------- merged solve: W = [sW: u | sQ: wneg], 4 rounds
    // build RHS-w = -beta*kq into sQ (qq dead after Aqk mma)
    {
        const int c4 = 4 * lane;
        #pragma unroll
        for (int rr = 0; rr < 16; rr++) {
            const int r = warp * 16 + rr;
            const float nb = -sBeta[r];
            uint2 kp = *reinterpret_cast<const uint2*>(sK + r * PQ_P + c4);
            float2 t0 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&kp.x));
            float2 t1 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&kp.y));
            uint2 outp = {pack_bf16x2(t0.x * nb, t0.y * nb), pack_bf16x2(t1.x * nb, t1.y * nb)};
            *reinterpret_cast<uint2*>(sQ + r * PQ_P + c4) = outp;
        }
    }
    __syncthreads();
    {
        const int cw = warp * 32;
        #pragma unroll
        for (int i = 0; i < 4; i++) {
            float accU[4][4] = {};
            float accW[4][4] = {};
            uint32_t da[4];
            {
                int rowa = (lane & 15), cola = ((lane >> 4) & 1) * 8;
                ldmatrix_x4(da[0], da[1], da[2], da[3], smem_u32(sDb + i * 16 * PDB_P + rowa * PDB_P + cola));
                #pragma unroll
                for (int n = 0; n < 4; n += 2) {
                    uint32_t b0, b1, b2, b3;
                    int m4 = lane >> 3;
                    int rowb = i * 16 + (m4 & 1) * 8 + (lane & 7);
                    int colb = cw + n * 8 + (m4 >> 1) * 8;
                    ldmatrix_x4_trans(b0, b1, b2, b3, smem_u32(sW + rowb * PQ_P + colb));
                    mma16816(accU[n][0], accU[n][1], accU[n][2], accU[n][3], da[0], da[1], da[2], da[3], b0, b1,
                             accU[n][0], accU[n][1], accU[n][2], accU[n][3]);
                    mma16816(accU[n + 1][0], accU[n + 1][1], accU[n + 1][2], accU[n + 1][3], da[0], da[1], da[2], da[3], b2, b3,
                             accU[n + 1][0], accU[n + 1][1], accU[n + 1][2], accU[n + 1][3]);
                    ldmatrix_x4_trans(b0, b1, b2, b3, smem_u32(sQ + rowb * PQ_P + colb));
                    mma16816(accW[n][0], accW[n][1], accW[n][2], accW[n][3], da[0], da[1], da[2], da[3], b0, b1,
                             accW[n][0], accW[n][1], accW[n][2], accW[n][3]);
                    mma16816(accW[n + 1][0], accW[n + 1][1], accW[n + 1][2], accW[n + 1][3], da[0], da[1], da[2], da[3], b2, b3,
                             accW[n + 1][0], accW[n + 1][1], accW[n + 1][2], accW[n + 1][3]);
                }
            }
            #pragma unroll
            for (int m = 0; m < 3; m++) {
                if (m >= i) break;
                uint32_t la[4];
                int rowa = (lane & 15), cola = ((lane >> 4) & 1) * 8;
                ldmatrix_x4(la[0], la[1], la[2], la[3], smem_u32(sA + (i * 16 + rowa) * PA_P + m * 16 + cola));
                #pragma unroll
                for (int n = 0; n < 4; n += 2) {
                    uint32_t b0, b1, b2, b3;
                    int m4 = lane >> 3;
                    int rowb = m * 16 + (m4 & 1) * 8 + (lane & 7);
                    int colb = cw + n * 8 + (m4 >> 1) * 8;
                    ldmatrix_x4_trans(b0, b1, b2, b3, smem_u32(sW + rowb * PQ_P + colb));
                    mma16816(accU[n][0], accU[n][1], accU[n][2], accU[n][3], la[0], la[1], la[2], la[3], b0, b1,
                             accU[n][0], accU[n][1], accU[n][2], accU[n][3]);
                    mma16816(accU[n + 1][0], accU[n + 1][1], accU[n + 1][2], accU[n + 1][3], la[0], la[1], la[2], la[3], b2, b3,
                             accU[n + 1][0], accU[n + 1][1], accU[n + 1][2], accU[n + 1][3]);
                    ldmatrix_x4_trans(b0, b1, b2, b3, smem_u32(sQ + rowb * PQ_P + colb));
                    mma16816(accW[n][0], accW[n][1], accW[n][2], accW[n][3], la[0], la[1], la[2], la[3], b0, b1,
                             accW[n][0], accW[n][1], accW[n][2], accW[n][3]);
                    mma16816(accW[n + 1][0], accW[n + 1][1], accW[n + 1][2], accW[n + 1][3], la[0], la[1], la[2], la[3], b2, b3,
                             accW[n + 1][0], accW[n + 1][1], accW[n + 1][2], accW[n + 1][3]);
                }
            }
            __syncthreads();
            {
                const int rl = lane >> 2, cl = 2 * (lane & 3);
                #pragma unroll
                for (int n = 0; n < 4; n++) {
                    *reinterpret_cast<uint32_t*>(sW + (i * 16 + rl) * PQ_P + cw + n * 8 + cl) =
                        pack_bf16x2(accU[n][0], accU[n][1]);
                    *reinterpret_cast<uint32_t*>(sW + (i * 16 + rl + 8) * PQ_P + cw + n * 8 + cl) =
                        pack_bf16x2(accU[n][2], accU[n][3]);
                    *reinterpret_cast<uint32_t*>(sQ + (i * 16 + rl) * PQ_P + cw + n * 8 + cl) =
                        pack_bf16x2(accW[n][0], accW[n][1]);
                    *reinterpret_cast<uint32_t*>(sQ + (i * 16 + rl + 8) * PQ_P + cw + n * 8 + cl) =
                        pack_bf16x2(accW[n][2], accW[n][3]);
                }
            }
            __syncthreads();
        }
    }
    // copyout u and wneg
    {
        const int r = tid >> 4, o = (tid & 15) * 8;
        #pragma unroll
        for (int i2 = 0; i2 < 8; i2++) {
            uint4 val = *reinterpret_cast<uint4*>(sW + (r + i2 * 8) * PQ_P + o);
            stg128(gu + (r + i2 * 8) * K + o, val);
            uint4 val2 = *reinterpret_cast<uint4*>(sQ + (r + i2 * 8) * PQ_P + o);
            stg128(gwneg + (r + i2 * 8) * K + o, val2);
        }
    }
    // publish chunk readiness
    __syncthreads();
    if (tid == 0) {
        __threadfence();
        atomicExch(flag_prep + cbase * 32, 1);
    }
}

void kda_prep(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
              torch::Tensor beta, torch::Tensor wneg, torch::Tensor u, torch::Tensor kg,
              torch::Tensor qg, torch::Tensor Aqk, torch::Tensor lam, torch::Tensor flag_prep,
              double scale) {
    int B = q.size(0), T = q.size(1), H = q.size(2);
    int NT = T / 64;
    dim3 grid(B * H, NT);
    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(kda_prep_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_TOTAL_P);
        configured = true;
    }
    auto stream = at::cuda::getCurrentCUDAStream();
    kda_prep_kernel<<<grid, 128, SMEM_TOTAL_P, stream>>>(
        reinterpret_cast<bf16*>(q.data_ptr()),
        reinterpret_cast<bf16*>(k.data_ptr()),
        reinterpret_cast<bf16*>(v.data_ptr()),
        g.data_ptr<float>(),
        reinterpret_cast<bf16*>(beta.data_ptr()),
        reinterpret_cast<bf16*>(wneg.data_ptr()),
        reinterpret_cast<bf16*>(u.data_ptr()),
        reinterpret_cast<bf16*>(kg.data_ptr()),
        reinterpret_cast<bf16*>(qg.data_ptr()),
        reinterpret_cast<bf16*>(Aqk.data_ptr()),
        lam.data_ptr<float>(),
        flag_prep.data_ptr<int>(),
        (float)scale, T, H, NT);
}


// KDA inter-chunk state scan (w-form), mma.sync m16n8k16 bf16.
//
//   per (bh, i_v):  hT (64 v-rows, 128 k-cols) fp32 accumulator, per-warp 16 rows.
//   loop chunks:  store h (as (K,V) via stmatrix.trans staging, warp-local cols)
//                 vnT = uT + hT @ wnegT     (w pre-negated by producer)
//                 store vn (as (BT,V) via stmatrix.trans staging, warp-local cols)
//                 hT  = lam * hT + vnT @ kg
//
// Inputs (chunk-tiled): wneg,kg (B,H,NT,64,K) bf16; u (B,H,NT,64,V) bf16; lam (B,H,NT,K) f32
// Outputs: ht (B,H,NT,V,K) bf16 (transposed);  vnt (B,H,NT,V,BT) bf16 (transposed)


constexpr int PW_S = 136;   // 128 + 8 pad  (272 B rows)
constexpr int PU_S = 72;    // 64 + 8 pad   (144 B rows)

constexpr int SZ_W_S = 64 * PW_S;
constexpr int SZ_U_S = 64 * PU_S;
constexpr int SZ_LAM_S = 256;
constexpr int STAGE_S = SZ_W_S * 2 + SZ_U_S + SZ_LAM_S;
constexpr int OFF_W_S = 0;
constexpr int OFF_KG_S = SZ_W_S;
constexpr int OFF_U_S = SZ_W_S * 2;
constexpr int OFF_LAM_S = SZ_W_S * 2 + SZ_U_S;
constexpr int NSTAGE_S = 4;
constexpr int OFF_HST_S = STAGE_S * NSTAGE_S;   // hT staging: 64 v-rows x PW_S
constexpr int SZ_HST_S = 64 * PW_S;
constexpr int OFF_VNST_S = OFF_HST_S + SZ_HST_S; // vnT staging: 64 v-rows x PU_S
constexpr int SZ_VNST_S = 64 * PU_S;
constexpr int SMEM_TOTAL_S = (OFF_VNST_S + SZ_VNST_S) * 2;

__global__ void __launch_bounds__(128, 1)
kda_scan_kernel(
    const bf16* __restrict__ wneg,
    const bf16* __restrict__ u,
    const bf16* __restrict__ kg,
    const float* __restrict__ lam,
    bf16* __restrict__ ht,
    bf16* __restrict__ vn,
    const int* __restrict__ flag_prep,
    int* __restrict__ flag_scan,
    int NT) {
    constexpr int K = 128, V = 128, BT = 64, BV = 64;
    extern __shared__ bf16 smem[];

    const int i_v = blockIdx.x;
    const int i_bh = blockIdx.y;
    const int tid = threadIdx.x;
    const int warp = tid >> 5, lane = tid & 31;
    const int v0 = i_v * BV;
    const int vw = warp * 16;

    // chunk-tiled inputs: rows are contiguous (stride K / V within a chunk tile)
    const int wr = tid >> 4, wo = (tid & 15) * 8;   // w,kg: 16 chunks/row
    const int ur = tid >> 3, uo = (tid & 7) * 8;    // u: 8 chunks/row (64-wide slice)
    const int64_t cb = (int64_t)i_bh * NT;
    const bf16* gw = wneg + cb * BT * K + wr * K + wo;
    const bf16* gkg = kg + cb * BT * K + wr * K + wo;
    const bf16* gu = u + cb * BT * V + v0 + ur * V + uo;
    const float* glam = lam + cb * K + tid * 4;
    // transposed outputs: row v, cols k / c ; this block owns v rows [v0, v0+64)
    bf16* ght = ht + (cb * V + v0) * K;
    bf16* gvnt = vn + (cb * V + v0) * 64;

    bf16* sw_base = smem + OFF_W_S + wr * PW_S + wo;
    bf16* skg_base = smem + OFF_KG_S + wr * PW_S + wo;
    bf16* su_base = smem + OFF_U_S + ur * PU_S + uo;

    const int* fprep = flag_prep + cb * 32;
    int* fscan = flag_scan + cb * 32;
    // single-thread poll; consumers use cp.async.cg (L2-direct) so one acquire
    // observation plus the block barrier is sufficient ordering.
    auto wait_prep = [&](int it) {
        if (tid == 0) {
            int v;
            do {
                asm volatile("ld.global.acquire.gpu.b32 %0, [%1];" : "=r"(v) : "l"(fprep + it * 32));
                if (!v) __nanosleep(128);
            } while (!v);
        }
    };
    auto fill_stage2 = [&](int it, int s) {
        const int64_t offK = (int64_t)it * BT * K;
        const int64_t offV = (int64_t)it * BT * V;
        bf16* sb = smem + s * STAGE_S;
        #pragma unroll
        for (int i = 0; i < 8; i++)
            cp_async16(sw_base + s * STAGE_S + i * 8 * PW_S, gw + offK + i * 8 * K);
        #pragma unroll
        for (int i = 0; i < 8; i++)
            cp_async16(skg_base + s * STAGE_S + i * 8 * PW_S, gkg + offK + i * 8 * K);
        #pragma unroll
        for (int i = 0; i < 4; i++)
            cp_async16(su_base + s * STAGE_S + i * 16 * PU_S, gu + offV + i * 16 * V);
        if (tid < 32)
            cp_async16(reinterpret_cast<float*>(sb + OFF_LAM_S) + tid * 4, glam + (int64_t)it * K);
        cp_async_commit();
    };

    float hacc[16][4];
    #pragma unroll
    for (int j = 0; j < 16; j++)
        #pragma unroll
        for (int e = 0; e < 4; e++) hacc[j][e] = 0.f;

    wait_prep(0);
    if (NT > 2) wait_prep(1);   // cover both prologue fills with one barrier
    __syncthreads();
    fill_stage2(0, 0);
    if (NT > 1) fill_stage2(1, 1);
    // preload one more stage for depth-3 pipeline
    if (NT > 2) { wait_prep(2); __syncthreads(); fill_stage2(2, 2); }

    for (int it = 0; it < NT; ++it) {
        const int s = it % NSTAGE_S;
        bf16* st = smem + s * STAGE_S;
        bf16* hst = smem + OFF_HST_S;
        bf16* vst = smem + OFF_VNST_S;

        if (it + 3 < NT) wait_prep(it + 3);   // tid0 poll (pre-barrier)
        cp_async_wait<2>();
        __syncthreads();
        // stage s ready; stage (it+3)%4 (= (it-1)%4) has no readers -> refill (depth-3).
        if (it + 3 < NT) fill_stage2(it + 3, (it + 3) % NSTAGE_S);
        else cp_async_commit();

        // ---- lam preload (registers)
        const float* sl = reinterpret_cast<const float*>(st + OFF_LAM_S);
        float2 lam2[16];
        #pragma unroll
        for (int j = 0; j < 16; j++)
            lam2[j] = *reinterpret_cast<const float2*>(sl + 8 * j + 2 * (lane & 3));

        // ---- pack h A-frags (h-store payload) then lam-scale hacc in one pass
        uint32_t hA[8][4];
        #pragma unroll
        for (int kk = 0; kk < 8; kk++) {
            hA[kk][0] = pack_bf16x2(hacc[2 * kk][0], hacc[2 * kk][1]);
            hA[kk][1] = pack_bf16x2(hacc[2 * kk][2], hacc[2 * kk][3]);
            hA[kk][2] = pack_bf16x2(hacc[2 * kk + 1][0], hacc[2 * kk + 1][1]);
            hA[kk][3] = pack_bf16x2(hacc[2 * kk + 1][2], hacc[2 * kk + 1][3]);
            hacc[2 * kk][0] *= lam2[2 * kk].x; hacc[2 * kk][1] *= lam2[2 * kk].y;
            hacc[2 * kk][2] *= lam2[2 * kk].x; hacc[2 * kk][3] *= lam2[2 * kk].y;
            hacc[2 * kk + 1][0] *= lam2[2 * kk + 1].x; hacc[2 * kk + 1][1] *= lam2[2 * kk + 1].y;
            hacc[2 * kk + 1][2] *= lam2[2 * kk + 1].x; hacc[2 * kk + 1][3] *= lam2[2 * kk + 1].y;
        }

        // ---- h store: non-trans stmatrix into (v,k) staging rows (warp-local rows)
        {
            int m = lane >> 3, r = lane & 7;
            #pragma unroll
            for (int kk = 0; kk < 8; kk++) {
                bf16* base = hst + (vw + (m & 1) * 8 + r) * PW_S + kk * 16 + (m >> 1) * 8;
                stmatrix_x4(smem_u32(base), hA[kk][0], hA[kk][1], hA[kk][2], hA[kk][3]);
            }
        }

        // ---- vn acc init = u^T
        float vacc[8][4];
        {
            bf16* su = st + OFF_U_S;
            #pragma unroll
            for (int cc = 0; cc < 4; cc++) {
                uint32_t q0, q1, q2, q3;
                int row = 16 * cc + (lane & 15), colh = (lane >> 4) * 8;
                ldmatrix_x4_trans(q0, q1, q2, q3, smem_u32(su + row * PU_S + vw + colh));
                float2 f0 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&q0));
                float2 f1 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&q1));
                float2 f2 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&q2));
                float2 f3 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&q3));
                vacc[2 * cc][0] = f0.x; vacc[2 * cc][1] = f0.y;
                vacc[2 * cc][2] = f2.x; vacc[2 * cc][3] = f2.y;
                vacc[2 * cc + 1][0] = f1.x; vacc[2 * cc + 1][1] = f1.y;
                vacc[2 * cc + 1][2] = f3.x; vacc[2 * cc + 1][3] = f3.y;
            }
        }

        // ---- mma#1: vnT += hT @ wnegT (kk outer, n inner)
        {
            bf16* sw = st + OFF_W_S;
            #pragma unroll
            for (int kk = 0; kk < 8; kk += 2) {
                uint32_t bw[8][4];
                #pragma unroll
                for (int n = 0; n < 8; n++) {
                    int row = n * 8 + (lane & 7), col = kk * 16 + ((lane >> 3) & 3) * 8;
                    ldmatrix_x4(bw[n][0], bw[n][1], bw[n][2], bw[n][3], smem_u32(sw + row * PW_S + col));
                }
                #pragma unroll
                for (int n = 0; n < 8; n++)
                    mma16816(vacc[n][0], vacc[n][1], vacc[n][2], vacc[n][3],
                             hA[kk][0], hA[kk][1], hA[kk][2], hA[kk][3], bw[n][0], bw[n][1],
                             vacc[n][0], vacc[n][1], vacc[n][2], vacc[n][3]);
                #pragma unroll
                for (int n = 0; n < 8; n++)
                    mma16816(vacc[n][0], vacc[n][1], vacc[n][2], vacc[n][3],
                             hA[kk + 1][0], hA[kk + 1][1], hA[kk + 1][2], hA[kk + 1][3], bw[n][2], bw[n][3],
                             vacc[n][0], vacc[n][1], vacc[n][2], vacc[n][3]);
            }
        }

        // ---- pack vn; stmatrix.trans to (c,v) staging; warp-local copyout
        uint32_t vA[4][4];
        #pragma unroll
        for (int kk = 0; kk < 4; kk++) {
            vA[kk][0] = pack_bf16x2(vacc[2 * kk][0], vacc[2 * kk][1]);
            vA[kk][1] = pack_bf16x2(vacc[2 * kk][2], vacc[2 * kk][3]);
            vA[kk][2] = pack_bf16x2(vacc[2 * kk + 1][0], vacc[2 * kk + 1][1]);
            vA[kk][3] = pack_bf16x2(vacc[2 * kk + 1][2], vacc[2 * kk + 1][3]);
        }
        {
            int m = lane >> 3, r = lane & 7;
            #pragma unroll
            for (int kk = 0; kk < 4; kk++) {
                bf16* base = vst + (vw + (m & 1) * 8 + r) * PU_S + kk * 16 + (m >> 1) * 8;
                stmatrix_x4(smem_u32(base), vA[kk][0], vA[kk][1], vA[kk][2], vA[kk][3]);
            }
        }

        // ---- mma#2: hT += vnT @ kg
        {
            bf16* sk = st + OFF_KG_S;
            #pragma unroll
            for (int kk = 0; kk < 4; kk++) {
                #pragma unroll
                for (int n = 0; n < 16; n += 2) {
                    uint32_t b0, b1, b2, b3;
                    int m = lane >> 3;
                    int row = kk * 16 + (m & 1) * 8 + (lane & 7);
                    int col = n * 8 + (m >> 1) * 8;
                    ldmatrix_x4_trans(b0, b1, b2, b3, smem_u32(sk + row * PW_S + col));
                    mma16816(hacc[n][0], hacc[n][1], hacc[n][2], hacc[n][3],
                             vA[kk][0], vA[kk][1], vA[kk][2], vA[kk][3], b0, b1,
                             hacc[n][0], hacc[n][1], hacc[n][2], hacc[n][3]);
                    mma16816(hacc[n + 1][0], hacc[n + 1][1], hacc[n + 1][2], hacc[n + 1][3],
                             vA[kk][0], vA[kk][1], vA[kk][2], vA[kk][3], b2, b3,
                             hacc[n + 1][0], hacc[n + 1][1], hacc[n + 1][2], hacc[n + 1][3]);
                }
            }
        }
        // ---- coalesced copyout (hT: 64 rows x 256B; vnT: 64 rows x 128B)
        __syncthreads();
        {
            bf16* dsth = ght + (int64_t)it * V * K;
            #pragma unroll
            for (int i = 0; i < 8; i++) {
                int c = tid + i * 128;             // 1024 chunks: row c>>4, off (c&15)*8
                int r = c >> 4, o = (c & 15) * 8;
                uint4 val = *reinterpret_cast<uint4*>(hst + r * PW_S + o);
                stg128(dsth + r * K + o, val);
            }
            bf16* dstv = gvnt + (int64_t)it * V * 64;
            #pragma unroll
            for (int i = 0; i < 4; i++) {
                int c = tid + i * 128;             // 512 chunks: row c>>3, off (c&7)*8
                int r = c >> 3, o = (c & 7) * 8;
                uint4 val = *reinterpret_cast<uint4*>(vst + r * PU_S + o);
                stg128(dstv + r * 64 + o, val);
            }
        }
        // publish this chunk's h/vn (2 producer blocks per bh -> counter to 2)
        __syncthreads();
        if (tid == 0) {
            __threadfence();
            atomicAdd(fscan + it * 32, 1);
        }
    }
}

void kda_scan(torch::Tensor wneg, torch::Tensor u, torch::Tensor kg, torch::Tensor lam,
              torch::Tensor ht, torch::Tensor vn, torch::Tensor flag_prep, torch::Tensor flag_scan) {
    int B = wneg.size(0), H = wneg.size(1), NT = wneg.size(2);
    dim3 grid(2, B * H);
    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(kda_scan_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_TOTAL_S);
        configured = true;
    }
    auto stream = at::cuda::getCurrentCUDAStream();
    kda_scan_kernel<<<grid, 128, SMEM_TOTAL_S, stream>>>(
        reinterpret_cast<bf16*>(wneg.data_ptr()),
        reinterpret_cast<bf16*>(u.data_ptr()),
        reinterpret_cast<bf16*>(kg.data_ptr()),
        lam.data_ptr<float>(),
        reinterpret_cast<bf16*>(ht.data_ptr()),
        reinterpret_cast<bf16*>(vn.data_ptr()),
        flag_prep.data_ptr<int>(),
        flag_scan.data_ptr<int>(),
        NT);
}


// KDA output kernel (K3):  o = qg @ h + Aqk @ vn   per chunk.
//
// Inputs (chunk-tiled): qg (B,H,NT,64,128), Aqk (B,H,NT,64,64),
//                       ht (B,H,NT,128v,128k)  [transposed state],
//                       vnt (B,H,NT,128v,64c)  [transposed v_new]
// Output: o (B,T,H,V) bf16


constexpr int PQ_O = 136;   // qg pitch
constexpr int PH_O = 136;   // hT pitch (128 rows x 128 cols)
constexpr int PA_O = 72;    // Aqk pitch
constexpr int PV_O = 72;    // vnT pitch (128 rows x 64 cols)
constexpr int PO_O = 136;   // o staging pitch

constexpr int OFF_QG_O = 0;                       // 64*PQ_O bf16
constexpr int OFF_H_O = OFF_QG_O + 64 * PQ_O * 2;     // 128*PH_O bf16
constexpr int OFF_AQ_O = OFF_H_O + 128 * PH_O * 2;    // 64*PA_O bf16
constexpr int OFF_VN_O = OFF_AQ_O + 64 * PA_O * 2;    // 128*PV_O bf16
constexpr int OFF_O_O = OFF_VN_O + 128 * PV_O * 2;    // 64*PO_O bf16 staging
constexpr int SMEM_TOTAL_O = OFF_O_O + 64 * PO_O * 2;

__global__ void __launch_bounds__(128, 2)
kda_out_kernel(
    const bf16* __restrict__ qg,
    const bf16* __restrict__ Aqk,
    const bf16* __restrict__ ht,
    const bf16* __restrict__ vnt,
    bf16* __restrict__ o,
    const int* __restrict__ flag_scan,
    int T, int H, int NT) {
    constexpr int BT = 64, K = 128, V = 128;
    extern __shared__ char smem_raw[];
    bf16* sQG = reinterpret_cast<bf16*>(smem_raw + OFF_QG_O);
    bf16* sH = reinterpret_cast<bf16*>(smem_raw + OFF_H_O);
    bf16* sAQ = reinterpret_cast<bf16*>(smem_raw + OFF_AQ_O);
    bf16* sVN = reinterpret_cast<bf16*>(smem_raw + OFF_VN_O);
    bf16* sO = reinterpret_cast<bf16*>(smem_raw + OFF_O_O);

    const int i_t = blockIdx.y, i_bh = blockIdx.x;
    const int i_b = i_bh / H, i_h = i_bh % H;
    const int tid = threadIdx.x;
    const int warp = tid >> 5, lane = tid & 31;

    const int64_t cbase = (int64_t)i_bh * NT + i_t;
    const bf16* gqg = qg + cbase * BT * K;
    const bf16* gAqk = Aqk + cbase * BT * BT;
    const bf16* ghT = ht + cbase * V * K;
    const bf16* gvnT = vnt + cbase * V * BT;
    bf16* go = o + ((int64_t)i_b * T + i_t * BT) * H * V + i_h * V;

    // wait for the scan to publish this chunk's h / vn (both V-halves).
    // tid0-only poll: loads below go through cp.async.cg (L2-direct), so the
    // single acquire observation + block barrier is sufficient.
    if (tid == 0) {
        const int* f = flag_scan + cbase * 32;
        int v;
        do {
            asm volatile("ld.global.acquire.gpu.b32 %0, [%1];" : "=r"(v) : "l"(f));
            if (v < 2) __nanosleep(1024);
        } while (v < 2);
    }
    __syncthreads();
    // ------------------------------------------------------------------ loads
    {
        // qg: 64 rows x 256B = 1024 chunks
        const int r = tid >> 4, oo = (tid & 15) * 8;
        #pragma unroll
        for (int i = 0; i < 8; i++)
            cp_async16(sQG + (r + i * 8) * PQ_O + oo, gqg + (r + i * 8) * K + oo);
        // hT: 128 rows x 256B = 2048 chunks
        #pragma unroll
        for (int i = 0; i < 16; i++)
            cp_async16(sH + (r + i * 8) * PH_O + oo, ghT + (r + i * 8) * K + oo);
        // Aqk: 64 x 128B
        const int r2 = tid >> 3, o2 = (tid & 7) * 8;
        #pragma unroll
        for (int i = 0; i < 4; i++)
            cp_async16(sAQ + (r2 + i * 16) * PA_O + o2, gAqk + (r2 + i * 16) * BT + o2);
        // vnT: 128 x 128B
        #pragma unroll
        for (int i = 0; i < 8; i++)
            cp_async16(sVN + (r2 + i * 16) * PV_O + o2, gvnT + (r2 + i * 16) * BT + o2);
        cp_async_commit();
    }
    cp_async_wait<0>();
    __syncthreads();

    const int vw = warp * 16;   // warp's 16 token rows
    float acc[16][4];
    #pragma unroll
    for (int n = 0; n < 16; n++)
        #pragma unroll
        for (int e = 0; e < 4; e++) acc[n][e] = 0.f;

    // ---- o += qg @ h : A = qg rows, B^T = hT (v,k) row-major -> non-trans ldm
    {
        uint32_t a[8][4];
        int rowa = vw + (lane & 15), cola = ((lane >> 4) & 1) * 8;
        #pragma unroll
        for (int kk8 = 0; kk8 < 8; kk8++)
            ldmatrix_x4(a[kk8][0], a[kk8][1], a[kk8][2], a[kk8][3],
                        smem_u32(sQG + rowa * PQ_O + kk8 * 16 + cola));
        #pragma unroll
        for (int kk8 = 0; kk8 < 8; kk8 += 2) {
            #pragma unroll
            for (int n = 0; n < 16; n += 2) {
                uint32_t b0, b1, b2, b3;
                // x4: m0/m1 = strip n (kfrag kk8 halves), m2/m3 = strip n+1? No:
                // use pattern: rows = n*8 + (lane&7), col = kk8*16 + ((lane>>3)&3)*8
                // -> (b0,b1) kfrag kk8 strip n_a ... replicate scan mma#1 pattern per strip.
                int row = (n + ((lane >> 4) & 1)) * 8 + (lane & 7);
                int col = kk8 * 16 + ((lane >> 3) & 1) * 8;
                ldmatrix_x4(b0, b1, b2, b3, smem_u32(sH + row * PH_O + col));
                // matrices: m0=(strip n, k lo) m1=(strip n, k hi) m2=(strip n+1, k lo) m3=(strip n+1, k hi)
                mma16816(acc[n][0], acc[n][1], acc[n][2], acc[n][3],
                         a[kk8][0], a[kk8][1], a[kk8][2], a[kk8][3], b0, b1,
                         acc[n][0], acc[n][1], acc[n][2], acc[n][3]);
                mma16816(acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3],
                         a[kk8][0], a[kk8][1], a[kk8][2], a[kk8][3], b2, b3,
                         acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3]);
                ldmatrix_x4(b0, b1, b2, b3, smem_u32(sH + row * PH_O + col + 16));
                mma16816(acc[n][0], acc[n][1], acc[n][2], acc[n][3],
                         a[kk8 + 1][0], a[kk8 + 1][1], a[kk8 + 1][2], a[kk8 + 1][3], b0, b1,
                         acc[n][0], acc[n][1], acc[n][2], acc[n][3]);
                mma16816(acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3],
                         a[kk8 + 1][0], a[kk8 + 1][1], a[kk8 + 1][2], a[kk8 + 1][3], b2, b3,
                         acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3]);
            }
        }
    }

    // ---- o += Aqk @ vn : A = Aqk rows, B^T = vnT (v,c) row-major -> non-trans
    {
        uint32_t a[4][4];
        int rowa = vw + (lane & 15), cola = ((lane >> 4) & 1) * 8;
        #pragma unroll
        for (int kk8 = 0; kk8 < 4; kk8++)
            ldmatrix_x4(a[kk8][0], a[kk8][1], a[kk8][2], a[kk8][3],
                        smem_u32(sAQ + rowa * PA_O + kk8 * 16 + cola));
        #pragma unroll
        for (int kk8 = 0; kk8 < 4; kk8 += 2) {
            #pragma unroll
            for (int n = 0; n < 16; n += 2) {
                uint32_t b0, b1, b2, b3;
                int row = (n + ((lane >> 4) & 1)) * 8 + (lane & 7);
                int col = kk8 * 16 + ((lane >> 3) & 1) * 8;
                ldmatrix_x4(b0, b1, b2, b3, smem_u32(sVN + row * PV_O + col));
                mma16816(acc[n][0], acc[n][1], acc[n][2], acc[n][3],
                         a[kk8][0], a[kk8][1], a[kk8][2], a[kk8][3], b0, b1,
                         acc[n][0], acc[n][1], acc[n][2], acc[n][3]);
                mma16816(acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3],
                         a[kk8][0], a[kk8][1], a[kk8][2], a[kk8][3], b2, b3,
                         acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3]);
                ldmatrix_x4(b0, b1, b2, b3, smem_u32(sVN + row * PV_O + col + 16));
                mma16816(acc[n][0], acc[n][1], acc[n][2], acc[n][3],
                         a[kk8 + 1][0], a[kk8 + 1][1], a[kk8 + 1][2], a[kk8 + 1][3], b0, b1,
                         acc[n][0], acc[n][1], acc[n][2], acc[n][3]);
                mma16816(acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3],
                         a[kk8 + 1][0], a[kk8 + 1][1], a[kk8 + 1][2], a[kk8 + 1][3], b2, b3,
                         acc[n + 1][0], acc[n + 1][1], acc[n + 1][2], acc[n + 1][3]);
            }
        }
    }

    // ---- stage o (bf16) and coalesced copyout to (B,T,H,V)
    {
        const int r0 = vw + (lane >> 2);
        #pragma unroll
        for (int n = 0; n < 16; n++) {
            const int c0 = n * 8 + 2 * (lane & 3);
            *reinterpret_cast<uint32_t*>(sO + r0 * PO_O + c0) = pack_bf16x2(acc[n][0], acc[n][1]);
            *reinterpret_cast<uint32_t*>(sO + (r0 + 8) * PO_O + c0) = pack_bf16x2(acc[n][2], acc[n][3]);
        }
    }
    __syncthreads();
    {
        const int r = tid >> 4, oo = (tid & 15) * 8;
        const int64_t strideT_V = (int64_t)H * V;
        #pragma unroll
        for (int i = 0; i < 8; i++) {
            uint4 val = *reinterpret_cast<uint4*>(sO + (r + i * 8) * PO_O + oo);
            stg128(go + (r + i * 8) * strideT_V + oo, val);
        }
    }
}

void kda_out(torch::Tensor qg, torch::Tensor Aqk, torch::Tensor ht, torch::Tensor vnt,
             torch::Tensor o, torch::Tensor flag_scan) {
    int B = o.size(0), T = o.size(1), H = o.size(2);
    int NT = T / 64;
    dim3 grid(B * H, NT);
    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(kda_out_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_TOTAL_O);
        configured = true;
    }
    auto stream = at::cuda::getCurrentCUDAStream();
    kda_out_kernel<<<grid, 128, SMEM_TOTAL_O, stream>>>(
        reinterpret_cast<bf16*>(qg.data_ptr()),
        reinterpret_cast<bf16*>(Aqk.data_ptr()),
        reinterpret_cast<bf16*>(ht.data_ptr()),
        reinterpret_cast<bf16*>(vnt.data_ptr()),
        reinterpret_cast<bf16*>(o.data_ptr()),
        flag_scan.data_ptr<int>(),
        T, H, NT);
}

"""

_CPP_SRC = r"""
void kda_prep(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
              torch::Tensor beta, torch::Tensor wneg, torch::Tensor u, torch::Tensor kg,
              torch::Tensor qg, torch::Tensor Aqk, torch::Tensor lam, torch::Tensor flag_prep,
              double scale);
void kda_scan(torch::Tensor wneg, torch::Tensor u, torch::Tensor kg, torch::Tensor lam,
              torch::Tensor ht, torch::Tensor vn, torch::Tensor flag_prep, torch::Tensor flag_scan);
void kda_out(torch::Tensor qg, torch::Tensor Aqk, torch::Tensor ht, torch::Tensor vnt,
             torch::Tensor o, torch::Tensor flag_scan);
"""

_ext = load_inline(
    name="kda_b200_kernels",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["kda_prep", "kda_scan", "kda_out"],
    extra_cuda_cflags=["-arch=sm_100a", "-O3", "--use_fast_math"],
    verbose=False,
)


def _kda_run_seq(q, k, v, g, beta, scale, ws):
    """Sequential launch (warmup / fallback): flags still flow, just no overlap."""
    ws["fp"].zero_()
    ws["fs"].zero_()
    _ext.kda_prep(q, k, v, g, beta, ws["wneg"], ws["u"], ws["kg"], ws["qg"],
                  ws["Aqk"], ws["lam"], ws["fp"], scale)
    _ext.kda_scan(ws["wneg"], ws["u"], ws["kg"], ws["lam"], ws["ht"], ws["vnt"],
                  ws["fp"], ws["fs"])
    _ext.kda_out(ws["qg"], ws["Aqk"], ws["ht"], ws["vnt"], ws["o"], ws["fs"])
    return ws["o"]


def _kda_run_overlap(q, k, v, g, beta, scale, ws, streams, events):
    """prep overlaps the scan (flag handshake); out runs after the scan.

    The scan is launched FIRST so its <=32 blocks (1 per SM, big smem) are
    resident before prep floods the machine; prep always has >=100 free SMs,
    so the flag wait cannot deadlock.
    """
    s_scan, s_prep, _ = streams
    e0, e1, e2, _ = events
    cur = torch.cuda.current_stream()
    ws["fp"].zero_()
    ws["fs"].zero_()
    e0.record(cur)
    s_scan.wait_event(e0)
    s_prep.wait_event(e0)
    with torch.cuda.stream(s_scan):
        _ext.kda_scan(ws["wneg"], ws["u"], ws["kg"], ws["lam"], ws["ht"], ws["vnt"],
                      ws["fp"], ws["fs"])
    with torch.cuda.stream(s_prep):
        _ext.kda_prep(q, k, v, g, beta, ws["wneg"], ws["u"], ws["kg"], ws["qg"],
                      ws["Aqk"], ws["lam"], ws["fp"], scale)
        # out chases the scan through flags once prep's SMs are free;
        # scan is already resident, prep is done -> no starvation possible.
        _ext.kda_out(ws["qg"], ws["Aqk"], ws["ht"], ws["vnt"], ws["o"], ws["fs"])
    e1.record(s_scan)
    e2.record(s_prep)
    cur.wait_event(e1)
    cur.wait_event(e2)
    return ws["o"]


def _make_ws(B, T, H, K, V, dev):
    BT, NT = 64, T // 64
    mk = lambda *s, dt=torch.bfloat16: torch.empty(*s, device=dev, dtype=dt)
    return {
        "wneg": mk(B, H, NT, BT, K), "u": mk(B, H, NT, BT, K),
        "kg": mk(B, H, NT, BT, K), "qg": mk(B, H, NT, BT, K),
        "Aqk": mk(B, H, NT, BT, BT), "lam": mk(B, H, NT, K, dt=torch.float32),
        "ht": mk(B, H, NT, V, K), "vnt": mk(B, H, NT, V, BT),
        "o": mk(B, T, H, V),
        "fp": torch.zeros(B * H * NT * 32, device=dev, dtype=torch.int32),
        "fs": torch.zeros(B * H * NT * 32, device=dev, dtype=torch.int32),
    }


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)
        self._graphs = {}

    def forward(self, q, k, v, g, beta):
        assert self.K == 128 and self.V == 128 and self.chunk_size == 64
        q = q.contiguous(); k = k.contiguous(); v = v.contiguous()
        g = g.contiguous(); beta = beta.contiguous()
        B, T, H, K = q.shape
        V = v.shape[-1]
        key = (q.data_ptr(), k.data_ptr(), v.data_ptr(), g.data_ptr(),
               beta.data_ptr(), B, T, H, K, V)
        ent = self._graphs.get(key)
        if ent is None:
            try:
                ws = _make_ws(B, T, H, K, V, q.device)
                streams = [torch.cuda.Stream(priority=-1), torch.cuda.Stream(), torch.cuda.Stream()]
                events = [torch.cuda.Event() for _ in range(4)]
                # warmup (sequential) on a side stream, then capture the overlapped form
                s = torch.cuda.Stream()
                s.wait_stream(torch.cuda.current_stream())
                with torch.cuda.stream(s):
                    _kda_run_seq(q, k, v, g, beta, self.scale, ws)
                torch.cuda.current_stream().wait_stream(s)
                gr = torch.cuda.CUDAGraph()
                with torch.cuda.graph(gr):
                    _kda_run_overlap(q, k, v, g, beta, self.scale, ws, streams, events)
                ent = (gr, ws)
                self._graphs[key] = ent
                if len(self._graphs) > 64:  # bound the cache
                    self._graphs.pop(next(iter(self._graphs)))
            except Exception:
                ws = _make_ws(B, T, H, K, V, q.device)
                return _kda_run_seq(q, k, v, g, beta, self.scale, ws)
        gr, ws = ent
        gr.replay()
        return ws["o"]


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