"""Kimi Delta Attention forward (chunk form) — custom CUDA (SM90) kernels.

Three kernels launched from one C++ op:
  K1 prep : per-chunk cumsum(g), Akk build (tensor cores), exact fp32 triangular
            inverse, w/u/Aqk GEMMs, decay-folded kg/qg.
  K2 relay: sequential inter-chunk state recurrence per (b,h,v-half).
  K3 out  : o = qg @ h + Aqk @ v_new, parallel over chunks.
Written from scratch (mma.sync m16n8k16 + cp.async); no FLA imports.
"""
import os

import torch
import torch.nn as nn

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

_CUDA_SRC = r'''
// Kimi Delta Attention forward (chunk form) — custom CUDA kernels for SM90.
//
// Pipeline (chunk size BT=64, K=V=128):
//   K1 prep  (grid NT*BH):  gc = cumsum(g)/ln2; Akk = (k*2^(gc-gn)) @ (k*2^(gn-gc))^T;
//                           X = (I - (-beta*tril(Akk,-1)))^{-1} exact fp32;
//                           w = X @ (beta*2^gc*k); u = X @ (beta*v);
//                           Aqk = tril((q*scale*2^(gc-gn)) @ (k*2^(gn-gc))^T);
//                           kg = k*2^(gl-gc); qg = q*scale*2^gc; egl = 2^gl.
//   K2 relay (grid BH*NV):  per (b,h) sequential over chunks:
//                           h_i = S (stored bf16); vn = u - w@S;
//                           S = diag(2^gl) S + kg^T @ vn.
//   K3 out   (grid NT*BH):  o = qg @ h_i + Aqk @ vn.
//
// Layouts (intermediates, chunk-head-major):
//   w,u,kg,qg,vn: [BH, NT, 64, 128] bf16 ; Aqk: [BH, NT, 64, 64] bf16
//   h: [BH, NT, 128, 128] bf16 ; egl: [BH, NT, 128] f32
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cstdint>

using bf16 = __nv_bfloat16;

#define DEVI __device__ __forceinline__

constexpr float RCP_LN2 = 1.4426950408889634f;
constexpr int BT = 64;    // chunk size
constexpr int DK = 128;   // head dim K
constexpr int DV = 128;   // head dim V

// smem row strides (elements) chosen so consecutive rows land in different
// 16B bank-groups for ldmatrix (stride*elt_size ≡ 16 mod 128 works: 272B, 144B)
constexpr int S128 = 136;  // for 128-wide bf16 tiles (272B rows)
constexpr int S64 = 72;    // for 64-wide bf16 tiles (144B rows)
constexpr int SG = 132;    // for 128-wide f32 g tile (528B rows)
constexpr int SA = 68;     // for 64-wide f32 Akk tile (272B rows)

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

DEVI void cp_async16(void* dst, const void* src) {
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem_u32(dst)), "l"(src));
}
DEVI void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N>
DEVI void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }

// ldmatrix x4: quad order (r0..r3) = (R0-7|C0-15lo, R8-15|C0-7? ...) — we use the
// canonical pairing: r0=(rows0-7,cols0-7), r1=(rows8-15,cols0-7), r2=(rows0-7,cols8-15),
// r3=(rows8-15,cols8-15) where "rows" are the 8 addressed 16B lines per quad.
// addr(lane): lanes 0-7 -> quad0 rows, 8-15 -> quad1, 16-23 -> quad2, 24-31 -> quad3.
DEVI void ldm_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, const void* p) {
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
               : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
               : "r"(smem_u32(p)));
}
DEVI void ldm_x4_trans(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, const void* p) {
  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"(smem_u32(p)));
}

// D = A@B + C, m16n8k16 bf16 -> f32
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) {
  asm volatile(
      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
      : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
      : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// ---------------------------------------------------------------------------
// Fragment loaders. All operate on bf16 smem tiles with row stride `ld` elems.
// A-frag for mma m16 (rows m0..m0+15) x k16 (cols k0..k0+15), m-major source.
DEVI void load_A(uint32_t (&a)[4], const bf16* s, int m0, int k0, int ld, int lane) {
  const bf16* p = s + (m0 + (lane & 15)) * ld + k0 + ((lane >> 4) << 3);
  ldm_x4(a[0], a[1], a[2], a[3], p);
}
// A-frag from k-major source (i.e. we want A = srcT): src rows = k, cols = m.
// quads: r0=(m0-7,k0-7): addr rows k0+l cols m0 ; r1=(m8-15,k0-7): rows k0+l-8, col m0+8?? no:
// r1 = A rows 8-15, k0-7 -> src rows k0-7, cols m0+8.. so lanes8-15 addr src row k0+(l-8), col m0+8.
// r2 = A rows0-7, k8-15 -> src rows k8-15, col m0: lanes16-23 addr row k8+(l-16), col m0.
// r3 -> lanes24-31 row k8+(l-24), col m0+8.
DEVI void load_A_trans(uint32_t (&a)[4], const bf16* s, int m0, int k0, int ld, int lane) {
  int q = lane >> 3, r = lane & 7;
  int row = k0 + ((q & 2) << 2) + r;      // +8 for quads 2,3
  int col = m0 + ((q & 1) << 3);          // +8 for quads 1,3
  ldm_x4_trans(a[0], a[1], a[2], a[3], s + row * ld + col);
}
// B-frag (k16 x n16 => two mma B pairs) from n-major source (rows n, cols k):
// r0=b0(n0-7): (k0-7 x n0-7): lanes0-7 addr src row n0+l col k0
// r1=b1(n0-7): (k8-15): lanes8-15 addr row n0+(l-8) col k0+8
// r2=b0(n8-15): lanes16-23 row n8+(l-16) col k0 ; r3=b1(n8-15): lanes24-31 row n8.. col k0+8
DEVI void load_B_nmajor(uint32_t (&bb)[4], const bf16* s, int k0, int n0, int ld, int lane) {
  int q = lane >> 3, r = lane & 7;
  int row = n0 + ((q & 2) << 2) + r;      // +8 for quads 2,3
  int col = k0 + ((q & 1) << 3);          // +8 for quads 1,3
  ldm_x4(bb[0], bb[1], bb[2], bb[3], s + row * ld + col);
}
// B-frag (k16 x n16) from k-major source (rows k, cols n): use trans.
// r0=b0(n0-7)=(k0-7,n0-7): lanes0-7 addr src row k0+l col n0
// r1=b1(n0-7)=(k8-15,n0-7): lanes8-15 addr row k8+(l-8) col n0
// r2=b0(n8-15): lanes16-23 row k0+(l-16) col n8 ; r3: lanes24-31 row k8+(l-24) col n8
DEVI void load_B_kmajor(uint32_t (&bb)[4], const bf16* s, int k0, int n0, int ld, int lane) {
  int q = lane >> 3, r = lane & 7;
  int row = k0 + ((q & 1) << 3) + r;      // +8 for quads 1,3
  int col = n0 + ((q & 2) << 2);          // +8 for quads 2,3
  ldm_x4_trans(bb[0], bb[1], bb[2], bb[3], s + row * ld + col);
}

DEVI uint32_t pack_bf16(float lo, float hi) {
  __nv_bfloat162 t = __floats2bfloat162_rn(lo, hi);
  return *reinterpret_cast<uint32_t*>(&t);
}

// C-frag (m16n8) coordinates for lane: values c0,c1 at (row g, col 2t/2t+1), c2,c3 at row g+8.
// g = lane>>2, t = lane&3.

// ===========================================================================
// K1: prep kernel.
// smem layout (dynamic):
//   sG   f32 [64][SG]                 33.8 KB   (g -> gc in-place)
//   sK   bf16 [64][S128]              17.4 KB
//   sQ   bf16 [64][S128]              17.4 KB   (later: v lands here for u GEMM)
//   sOpA bf16 [64][S128]              17.4 KB   (kp -> qp -> beta*2^gc*k -> beta*v)
//   sOpB bf16 [64][S128]              17.4 KB   (kn)
//   sAkk f32 [64][SA]                 17.4 KB   (raw kk acc -> X in-place)
//   sXb  bf16 [64][S64]                9.2 KB
//   sTmp f32 [32][34]                  4.4 KB   (merge scratch)
//   sBeta f32 [64], sGn f32[128], sEgl f32[128]
// ===========================================================================
struct PrepSmem {
  float sG[64 * SG];
  bf16 sK[64 * S128];
  bf16 sQ[64 * S128];
  bf16 sOpA[64 * S128];
  bf16 sOpB[64 * S128];
  float sAkk[64 * SA];
  bf16 sXb[64 * S64];
  float sTmp[32 * 34];
  float sBeta[64];
  float sGn[128];
};

__global__ void __launch_bounds__(256, 1) 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__ w, bf16* __restrict__ u, bf16* __restrict__ kg, bf16* __restrict__ qg,
    bf16* __restrict__ Aqk, float* __restrict__ egl,
    int T, int H, int NT, float scale) {
  extern __shared__ char smem_raw[];
  PrepSmem& sm = *reinterpret_cast<PrepSmem*>(smem_raw);

  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int wid = tid >> 5;
  const int bh = blockIdx.x / NT;
  const int it = blockIdx.x % NT;
  const int b = bh / H;
  const int h = bh % H;
  const int t0 = it * BT;

  // ---- stage A: async loads ----
  // g rows: 512B each (128 f32); 32 x 16B ops per row; 64 rows = 2048 ops / 256 thr = 8
  {
    const float* gsrc = g + (((int64_t)b * T + t0) * H + h) * DK;
    const int64_t rowstride = (int64_t)H * DK;
#pragma unroll
    for (int i = 0; i < 8; i++) {
      int idx = tid + i * 256;         // 16B unit index
      int r = idx >> 5, c = (idx & 31) << 2;  // 32 x 16B per row, c in f32 elems
      cp_async16(&sm.sG[r * SG + c], gsrc + r * rowstride + c);
    }
  }
  // k rows: 256B each (128 bf16); 16 x 16B per row; 64 rows = 1024 ops / 256 = 4
  {
    const bf16* ksrc = k + (((int64_t)b * T + t0) * H + h) * DK;
    const bf16* qsrc = q + (((int64_t)b * T + t0) * H + h) * DK;
    const int64_t rowstride = (int64_t)H * DK;
#pragma unroll
    for (int i = 0; i < 4; i++) {
      int idx = tid + i * 256;
      int r = idx >> 4, c = (idx & 15) << 3;  // c in bf16 elems
      cp_async16(&sm.sK[r * S128 + c], ksrc + r * rowstride + c);
      cp_async16(&sm.sQ[r * S128 + c], qsrc + r * rowstride + c);
    }
  }
  cp_commit();
  // beta (sync loads, small)
  if (tid < 64) {
    sm.sBeta[tid] = __bfloat162float(beta[(((int64_t)b * T + t0 + tid) * H) + h]);
  }
  cp_wait<0>();
  __syncthreads();

  // ---- stage B: cumsum (log2 space) ----
  // 128 threads: col each; two-phase over rows with 256 threads:
  // phase 1: thread t (t<256): col=t&127, half=t>>7 -> rows [32*half, 32*half+32)
  {
    int col = tid & 127, half = tid >> 7;
    int r0 = half * 32;
    float acc = 0.f;
#pragma unroll
    for (int r = 0; r < 32; r++) {
      acc += sm.sG[(r0 + r) * SG + col];
      sm.sG[(r0 + r) * SG + col] = acc;
    }
    __syncthreads();
    if (half == 1) {
      float base = sm.sG[31 * SG + col];
#pragma unroll
      for (int r = 32; r < 64; r++) sm.sG[r * SG + col] += base;
    }
    __syncthreads();
    // scale to log2 and stash gn (row 31)
    if (tid < 128) sm.sGn[tid] = sm.sG[31 * SG + tid] * RCP_LN2;
    __syncthreads();
#pragma unroll
    for (int r = 0; r < 32; r++) {
      sm.sG[(r0 + r) * SG + col] *= RCP_LN2;
    }
  }
  __syncthreads();

  // ---- stage C: kp, kn ----
  for (int idx = tid; idx < 64 * 128; idx += 256) {
    int r = idx >> 7, c = idx & 127;
    float gc = sm.sG[r * SG + c];
    float kf = __bfloat162float(sm.sK[r * S128 + c]);
    float ep = exp2f(gc - sm.sGn[c]);
    float en = exp2f(sm.sGn[c] - gc);
    sm.sOpA[r * S128 + c] = __float2bfloat16_rn(kf * ep);
    sm.sOpB[r * S128 + c] = __float2bfloat16_rn(kf * en);
  }
  __syncthreads();

  // ---- stage D: Akk = kp @ kn^T (raw, fp32) ----
  // 8 warps: 4x2 grid of 16x32 tiles. out[m][n] with n from kn rows (n-major B source)
  {
    int m0 = (wid & 3) * 16, n0 = (wid >> 2) * 32;
    float acc[4][4] = {};
#pragma unroll
    for (int kb = 0; kb < 8; kb++) {
      uint32_t a[4], b0[4], b1[4];
      load_A(a, sm.sOpA, m0, kb * 16, S128, lane);
      load_B_nmajor(b0, sm.sOpB, kb * 16, n0, S128, lane);
      load_B_nmajor(b1, sm.sOpB, kb * 16, n0 + 16, S128, lane);
#pragma unroll
      for (int nn = 0; nn < 2; nn++) {
        mma16816(acc[nn][0], acc[nn][1], acc[nn][2], acc[nn][3], a[0], a[1], a[2], a[3],
                 b0[nn * 2], b0[nn * 2 + 1]);
        mma16816(acc[nn + 2][0], acc[nn + 2][1], acc[nn + 2][2], acc[nn + 2][3], a[0], a[1], a[2],
                 a[3], b1[nn * 2], b1[nn * 2 + 1]);
      }
    }
    // store raw acc to sAkk. tile order: acc[0]=(n0..n0+7), acc[1]=(n0+8..15), acc[2]=(n0+16..23), acc[3]=(n0+24..31)
    int g_ = lane >> 2, t_ = lane & 3;
#pragma unroll
    for (int nn = 0; nn < 4; nn++) {
      int c = n0 + nn * 8 + 2 * t_;
      sm.sAkk[(m0 + g_) * SA + c] = acc[nn][0];
      sm.sAkk[(m0 + g_) * SA + c + 1] = acc[nn][1];
      sm.sAkk[(m0 + g_ + 8) * SA + c] = acc[nn][2];
      sm.sAkk[(m0 + g_ + 8) * SA + c + 1] = acc[nn][3];
    }
  }
  __syncthreads();

  // ---- stage C2: qp = q * 2^(gc-gn) * scale into sOpA ----
  for (int idx = tid; idx < 64 * 128; idx += 256) {
    int r = idx >> 7, c = idx & 127;
    float gc = sm.sG[r * SG + c];
    float qf = __bfloat162float(sm.sQ[r * S128 + c]);
    sm.sOpA[r * S128 + c] = __float2bfloat16_rn(qf * exp2f(gc - sm.sGn[c]) * scale);
  }
  __syncthreads();

  // ---- stage E: warps 0-3: invert; warps 4-7: Aqk mma ----
  float aqk_acc[8][4];  // warp tile 16x64 for warps 4-7
  if (wid >= 4) {
    int m0 = (wid - 4) * 16;
#pragma unroll
    for (int i = 0; i < 8; i++)
#pragma unroll
      for (int j = 0; j < 4; j++) aqk_acc[i][j] = 0.f;
#pragma unroll
    for (int kb = 0; kb < 8; kb++) {
      uint32_t a[4], bb[4];
      load_A(a, sm.sOpA, m0, kb * 16, S128, lane);
#pragma unroll
      for (int nn = 0; nn < 4; nn++) {
        load_B_nmajor(bb, sm.sOpB, kb * 16, nn * 16, S128, lane);
        mma16816(aqk_acc[nn * 2][0], aqk_acc[nn * 2][1], aqk_acc[nn * 2][2], aqk_acc[nn * 2][3],
                 a[0], a[1], a[2], a[3], bb[0], bb[1]);
        mma16816(aqk_acc[nn * 2 + 1][0], aqk_acc[nn * 2 + 1][1], aqk_acc[nn * 2 + 1][2],
                 aqk_acc[nn * 2 + 1][3], a[0], a[1], a[2], a[3], bb[2], bb[3]);
      }
    }
    // mask r>=c and store to global Aqk[bh][it]
    bf16* dst = Aqk + ((int64_t)bh * NT + it) * BT * BT;
    int g_ = lane >> 2, t_ = lane & 3;
#pragma unroll
    for (int nn = 0; nn < 8; nn++) {
      int c = nn * 8 + 2 * t_;
      int r1 = m0 + g_, r2 = m0 + g_ + 8;
      uint32_t p1 = pack_bf16(r1 >= c ? aqk_acc[nn][0] : 0.f, r1 >= c + 1 ? aqk_acc[nn][1] : 0.f);
      uint32_t p2 = pack_bf16(r2 >= c ? aqk_acc[nn][2] : 0.f, r2 >= c + 1 ? aqk_acc[nn][3] : 0.f);
      *reinterpret_cast<uint32_t*>(dst + r1 * BT + c) = p1;
      *reinterpret_cast<uint32_t*>(dst + r2 * BT + c) = p2;
    }
  } else {
    // ---- inversion: X = (I - M)^{-1}, M = -beta_r * strict_lower(Akk) ----
    // (a) diagonal 16x16 blocks, one per warp (wid 0-3), lane = column j (lanes 16-31 idle)
    int r0 = wid * 16;
    if (lane < 16) {
      int j = lane;
      float x[16];
#pragma unroll
      for (int i = 0; i < 16; i++) x[i] = 0.f;
      x[j] = 1.f;  // xx[m]: 0 for m<j, 1 at j, X[m][j] for m>j
#pragma unroll
      for (int i = 1; i < 16; i++) {
        float bi = sm.sBeta[r0 + i];
        float t = 0.f;
#pragma unroll
        for (int m = 0; m < 16; m++) {
          if (m < i) t += -bi * sm.sAkk[(r0 + i) * SA + r0 + m] * x[m];
        }
        if (i > j) x[i] = t;
      }
      // write X block back into sAkk diag (including unit diagonal, zero upper)
#pragma unroll
      for (int i = 0; i < 16; i++) sm.sAkk[(r0 + i) * SA + r0 + j] = (i == j) ? 1.f : (i > j ? x[i] : 0.f);
    }
    // zero the 16x16 upper blocks (0:16,16:32) and (32:48,48:64) that the merge
    // steps read as parts of "complete" 32x32 triangular factors
#pragma unroll
    for (int e = 0; e < 4; e++) {
      int idx = tid + e * 128;         // 0..511
      int blk = idx >> 8;              // 0 or 1
      int rr = (idx >> 4) & 15, cc = idx & 15;
      sm.sAkk[(blk * 32 + rr) * SA + blk * 32 + 16 + cc] = 0.f;
    }
    // barrier for warps 0-3 only
    asm volatile("bar.sync 1, 128;\n");
    // (b) level-32 merge: pairs (block1,block0) and (block3,block2)
    // X10 = X11 @ M10 @ X00 ; T = M10 @ X00 first (16x16), M10 = -beta*Akk_raw
    // 128 threads: 2 pairs x 256 entries = 512 -> 4 entries per thread
    {
      int tid4 = tid;  // 0..127
#pragma unroll
      for (int e = 0; e < 4; e++) {
        int idx = tid4 + e * 128;        // 0..511
        int pair = idx >> 8;             // 0 or 1
        int rr = (idx >> 4) & 15, cc = idx & 15;
        int rbase = pair * 32 + 16, cbase = pair * 32;
        float acc = 0.f;
        float bi = sm.sBeta[rbase + rr];
#pragma unroll
        for (int m = 0; m < 16; m++)
          acc += -bi * sm.sAkk[(rbase + rr) * SA + cbase + m] * sm.sAkk[(cbase + m) * SA + cbase + cc];
        sm.sTmp[pair * 16 * 17 + rr * 17 + cc] = acc;
      }
      asm volatile("bar.sync 1, 128;\n");
#pragma unroll
      for (int e = 0; e < 4; e++) {
        int idx = tid4 + e * 128;
        int pair = idx >> 8;
        int rr = (idx >> 4) & 15, cc = idx & 15;
        int rbase = pair * 32 + 16, cbase = pair * 32;
        float acc = 0.f;
#pragma unroll
        for (int m = 0; m < 16; m++)
          acc += sm.sAkk[(rbase + rr) * SA + rbase + m] * sm.sTmp[pair * 16 * 17 + m * 17 + cc];
        sm.sTmp[544 + pair * 16 * 17 + rr * 17 + cc] = acc;  // stage in tmp (avoid RAW smem hazards)
      }
      asm volatile("bar.sync 1, 128;\n");
#pragma unroll
      for (int e = 0; e < 4; e++) {
        int idx = tid4 + e * 128;
        int pair = idx >> 8;
        int rr = (idx >> 4) & 15, cc = idx & 15;
        int rbase = pair * 32 + 16, cbase = pair * 32;
        sm.sAkk[(rbase + rr) * SA + cbase + cc] = sm.sTmp[544 + pair * 16 * 17 + rr * 17 + cc];
      }
      asm volatile("bar.sync 1, 128;\n");
    }
    // (c) level-64 merge: X[32:64][0:32] = Xbot @ (-beta*A[32:][:32]) @ Xtop
    // T2 = M @ Xtop (32x32): 1024 entries / 128 thr = 8 each
    {
#pragma unroll
      for (int e = 0; e < 8; e++) {
        int idx = tid + e * 128;
        int rr = idx >> 5, cc = idx & 31;
        float bi = sm.sBeta[32 + rr];
        float acc = 0.f;
#pragma unroll
        for (int m = 0; m < 32; m++)
          acc += -bi * sm.sAkk[(32 + rr) * SA + m] * sm.sAkk[m * SA + cc];
        sm.sTmp[rr * 34 + cc] = acc;
      }
      asm volatile("bar.sync 1, 128;\n");
#pragma unroll
      for (int e = 0; e < 8; e++) {
        int idx = tid + e * 128;
        int rr = idx >> 5, cc = idx & 31;
        float acc = 0.f;
#pragma unroll
        for (int m = 0; m < 32; m++)
          acc += sm.sAkk[(32 + rr) * SA + 32 + m] * sm.sTmp[m * 34 + cc];
        ((float*)sm.sXb)[rr * 36 + cc] = acc;  // sXb (9.2KB) as f32 scratch: 32*36*4=4.6KB
      }
      asm volatile("bar.sync 1, 128;\n");
#pragma unroll
      for (int e = 0; e < 8; e++) {
        int idx = tid + e * 128;
        int rr = idx >> 5, cc = idx & 31;
        sm.sAkk[(32 + rr) * SA + cc] = ((float*)sm.sXb)[rr * 36 + cc];
      }
      // zero strict-upper 32x32 block of X (rows 0-31, cols 32-63)
#pragma unroll
      for (int e = 0; e < 8; e++) {
        int idx = tid + e * 128;
        int rr = idx >> 5, cc = idx & 31;
        sm.sAkk[rr * SA + 32 + cc] = 0.f;
      }
    }
  }
  __syncthreads();

  // ---- stage F: X -> bf16, build (beta * 2^gc * k) into sOpA ----
  for (int idx = tid; idx < 64 * 64; idx += 256) {
    int r = idx >> 6, c = idx & 63;
    sm.sXb[r * S64 + c] = __float2bfloat16_rn(sm.sAkk[r * SA + c]);
  }
  for (int idx = tid; idx < 64 * 128; idx += 256) {
    int r = idx >> 7, c = idx & 127;
    float gc = sm.sG[r * SG + c];
    float kf = __bfloat162float(sm.sK[r * S128 + c]);
    sm.sOpA[r * S128 + c] = __float2bfloat16_rn(kf * exp2f(gc) * sm.sBeta[r]);
    // also produce kg and egl here (independent of X)
    float kgf = kf * exp2f(sm.sG[63 * SG + c] - gc);
    kg[(((int64_t)bh * NT + it) * BT + r) * DK + c] = __float2bfloat16_rn(kgf);
    float qf = __bfloat162float(sm.sQ[r * S128 + c]);
    qg[(((int64_t)bh * NT + it) * BT + r) * DK + c] = __float2bfloat16_rn(qf * exp2f(gc) * scale);
  }
  if (tid < 128) egl[((int64_t)bh * NT + it) * DK + tid] = exp2f(sm.sG[63 * SG + tid]);
  __syncthreads();

  // ---- stage G: w = Xb @ sOpA  (m=64 rows, n=128 K-cols, k=64) ----
  {
    int m0 = (wid & 3) * 16, n0 = (wid >> 2) * 64;
    float acc[8][4] = {};
#pragma unroll
    for (int kb = 0; kb < 4; kb++) {
      uint32_t a[4], bb[4];
      load_A(a, sm.sXb, m0, kb * 16, S64, lane);
#pragma unroll
      for (int nn = 0; nn < 4; nn++) {
        load_B_kmajor(bb, sm.sOpA, kb * 16, n0 + nn * 16, S128, lane);
        mma16816(acc[nn * 2][0], acc[nn * 2][1], acc[nn * 2][2], acc[nn * 2][3], a[0], a[1], a[2],
                 a[3], bb[0], bb[1]);
        mma16816(acc[nn * 2 + 1][0], acc[nn * 2 + 1][1], acc[nn * 2 + 1][2], acc[nn * 2 + 1][3],
                 a[0], a[1], a[2], a[3], bb[2], bb[3]);
      }
    }
    bf16* dst = w + ((int64_t)bh * NT + it) * BT * DK;
    int g_ = lane >> 2, t_ = lane & 3;
#pragma unroll
    for (int nn = 0; nn < 8; nn++) {
      int c = n0 + nn * 8 + 2 * t_;
      *reinterpret_cast<uint32_t*>(dst + (m0 + g_) * DK + c) = pack_bf16(acc[nn][0], acc[nn][1]);
      *reinterpret_cast<uint32_t*>(dst + (m0 + g_ + 8) * DK + c) = pack_bf16(acc[nn][2], acc[nn][3]);
    }
  }

  // ---- stage H: v loads into sQ, then u = Xb @ (beta * v) ----
  // (load v synchronously via cp.async now; overlap with w GEMM above was possible but
  //  keep simple: issue before GEMM G would race sQ use in stage F. sQ free now.)
  {
    const bf16* vsrc = v + (((int64_t)b * T + t0) * H + h) * DV;
    const int64_t rowstride = (int64_t)H * DV;
#pragma unroll
    for (int i = 0; i < 4; i++) {
      int idx = tid + i * 256;
      int r = idx >> 4, c = (idx & 15) << 3;
      cp_async16(&sm.sQ[r * S128 + c], vsrc + r * rowstride + c);
    }
    cp_commit();
    cp_wait<0>();
    __syncthreads();
    for (int idx = tid; idx < 64 * 128; idx += 256) {
      int r = idx >> 7, c = idx & 127;
      float vf = __bfloat162float(sm.sQ[r * S128 + c]);
      sm.sOpA[r * S128 + c] = __float2bfloat16_rn(vf * sm.sBeta[r]);
    }
    __syncthreads();
    int m0 = (wid & 3) * 16, n0 = (wid >> 2) * 64;
    float acc[8][4] = {};
#pragma unroll
    for (int kb = 0; kb < 4; kb++) {
      uint32_t a[4], bb[4];
      load_A(a, sm.sXb, m0, kb * 16, S64, lane);
#pragma unroll
      for (int nn = 0; nn < 4; nn++) {
        load_B_kmajor(bb, sm.sOpA, kb * 16, n0 + nn * 16, S128, lane);
        mma16816(acc[nn * 2][0], acc[nn * 2][1], acc[nn * 2][2], acc[nn * 2][3], a[0], a[1], a[2],
                 a[3], bb[0], bb[1]);
        mma16816(acc[nn * 2 + 1][0], acc[nn * 2 + 1][1], acc[nn * 2 + 1][2], acc[nn * 2 + 1][3],
                 a[0], a[1], a[2], a[3], bb[2], bb[3]);
      }
    }
    bf16* dst = u + ((int64_t)bh * NT + it) * BT * DV;
    int g_ = lane >> 2, t_ = lane & 3;
#pragma unroll
    for (int nn = 0; nn < 8; nn++) {
      int c = n0 + nn * 8 + 2 * t_;
      *reinterpret_cast<uint32_t*>(dst + (m0 + g_) * DV + c) = pack_bf16(acc[nn][0], acc[nn][1]);
      *reinterpret_cast<uint32_t*>(dst + (m0 + g_ + 8) * DV + c) = pack_bf16(acc[nn][2], acc[nn][3]);
    }
  }
}

// ===========================================================================
// K2: relay kernel. grid (BH * NV), NV = DV/BV, BV = 64.
// ===========================================================================
constexpr int BV = 64;
struct RelaySmem {
  bf16 sW[2][64 * S128];
  bf16 sU[2][64 * S64];
  bf16 sKg[2][64 * S128];
  float sEgl[2][128];
  bf16 sS[128 * S64];   // bf16 mirror of state (K x BV)
  bf16 sVn[64 * S64];
};

__global__ void __launch_bounds__(256, 1) kda_relay_kernel(
    const bf16* __restrict__ w, const bf16* __restrict__ u, const bf16* __restrict__ kg,
    const float* __restrict__ egl, bf16* __restrict__ hst, bf16* __restrict__ vn,
    int NT) {
  extern __shared__ char smem_raw[];
  RelaySmem& sm = *reinterpret_cast<RelaySmem*>(smem_raw);

  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int wid = tid >> 5;
  const int bh = blockIdx.x >> 1;   // NV = 2
  const int iv = blockIdx.x & 1;
  const int v0 = iv * BV;

  const bf16* wp = w + (int64_t)bh * NT * BT * DK;
  const bf16* up = u + (int64_t)bh * NT * BT * DV + v0;
  const bf16* kgp = kg + (int64_t)bh * NT * BT * DK;
  const float* eglp = egl + (int64_t)bh * NT * DK;
  bf16* hp = hst + (int64_t)bh * NT * DK * DV + v0;
  bf16* vnp = vn + (int64_t)bh * NT * BT * DV + v0;

  // S fragments: 128(K) x 64(V) over 8 warps as 4x2 grid of 32x32 tiles.
  // Each warp: 2(m) x 4(n) mma-tiles of 16x8 -> acc[2][4][4].
  const int smw = (wid & 3) * 32;   // K-row base of this warp's S tile
  const int snw = (wid >> 2) * 32;  // V-col base
  float S[2][4][4] = {};

  // zero mirror
  for (int idx = tid; idx < 128 * S64; idx += 256) sm.sS[idx] = __float2bfloat16_rn(0.f);

  // prefetch chunk 0 into buf 0
  auto prefetch = [&](int i, int buf) {
    const bf16* wsrc = wp + (int64_t)i * BT * DK;
    const bf16* usrc = up + (int64_t)i * BT * DV;
    const bf16* kgsrc = kgp + (int64_t)i * BT * DK;
#pragma unroll
    for (int j = 0; j < 4; j++) {
      int idx = tid + j * 256;
      int r = idx >> 4, c = (idx & 15) << 3;
      cp_async16(&sm.sW[buf][r * S128 + c], wsrc + r * DK + c);
      cp_async16(&sm.sKg[buf][r * S128 + c], kgsrc + r * DK + c);
    }
    // u tile 64 x 64 (BV): 8 x 16B per row -> 512 ops / 256 thr = 2
#pragma unroll
    for (int j = 0; j < 2; j++) {
      int idx = tid + j * 256;
      int r = idx >> 3, c = (idx & 7) << 3;
      cp_async16(&sm.sU[buf][r * S64 + c], usrc + r * DV + c);
    }
    if (tid < 32) cp_async16(&sm.sEgl[buf][tid * 4], eglp + (int64_t)i * DK + tid * 4);
    cp_commit();
  };
  prefetch(0, 0);
  __syncthreads();  // mirror zero visible

  for (int i = 0; i < NT; i++) {
    int buf = i & 1;
    if (i + 1 < NT) {
      prefetch(i + 1, buf ^ 1);  // one extra group in flight
      cp_wait<1>();              // wait for chunk i's group
    } else {
      cp_wait<0>();
    }
    __syncthreads();

    // ---- store h_i = mirror (bf16 of S before update) ----
    {
      bf16* hdst = hp + (int64_t)i * DK * DV;
      // 128 rows x 64 cols, 16B vectors: 8 per row -> 1024 ops / 256 thr = 4
#pragma unroll
      for (int j = 0; j < 4; j++) {
        int idx = tid + j * 256;
        int r = idx >> 3, c = (idx & 7) << 3;
        *reinterpret_cast<uint4*>(hdst + r * DV + c) =
            *reinterpret_cast<const uint4*>(&sm.sS[r * S64 + c]);
      }
    }

    // ---- vn = u - w @ S ---- out 64x64: warps 4x2? use 4x2 grid of 16x32
    {
      int m0 = (wid & 3) * 16, n0 = (wid >> 2) * 32;
      float acc[4][4] = {};
#pragma unroll
      for (int kb = 0; kb < 8; kb++) {
        uint32_t a[4], b0[4], b1[4];
        load_A(a, sm.sW[buf], m0, kb * 16, S128, lane);
        load_B_kmajor(b0, sm.sS, kb * 16, n0, S64, lane);
        load_B_kmajor(b1, sm.sS, kb * 16, n0 + 16, S64, lane);
#pragma unroll
        for (int nn = 0; nn < 2; nn++) {
          mma16816(acc[nn][0], acc[nn][1], acc[nn][2], acc[nn][3], a[0], a[1], a[2], a[3],
                   b0[nn * 2], b0[nn * 2 + 1]);
          mma16816(acc[nn + 2][0], acc[nn + 2][1], acc[nn + 2][2], acc[nn + 2][3], a[0], a[1],
                   a[2], a[3], b1[nn * 2], b1[nn * 2 + 1]);
        }
      }
      // vn = u - acc; write to sVn + global
      bf16* vdst = vnp + (int64_t)i * BT * DV;
      int g_ = lane >> 2, t_ = lane & 3;
#pragma unroll
      for (int nn = 0; nn < 4; nn++) {
        int c = n0 + nn * 8 + 2 * t_;
        int r1 = m0 + g_, r2 = m0 + g_ + 8;
        float v00 = __bfloat162float(sm.sU[buf][r1 * S64 + c]) - acc[nn][0];
        float v01 = __bfloat162float(sm.sU[buf][r1 * S64 + c + 1]) - acc[nn][1];
        float v10 = __bfloat162float(sm.sU[buf][r2 * S64 + c]) - acc[nn][2];
        float v11 = __bfloat162float(sm.sU[buf][r2 * S64 + c + 1]) - acc[nn][3];
        uint32_t p0 = pack_bf16(v00, v01), p1 = pack_bf16(v10, v11);
        *reinterpret_cast<uint32_t*>(&sm.sVn[r1 * S64 + c]) = p0;
        *reinterpret_cast<uint32_t*>(&sm.sVn[r2 * S64 + c]) = p1;
        *reinterpret_cast<uint32_t*>(vdst + r1 * DV + c) = p0;
        *reinterpret_cast<uint32_t*>(vdst + r2 * DV + c) = p1;
      }
    }
    __syncthreads();  // sVn ready; h store done; safe to update S mirror later

    // ---- decay S and accumulate kg^T @ vn ----
    {
      // decay: rows of S tile: lane rows smw + (mm*16) + g, +8
      int g_ = lane >> 2;
#pragma unroll
      for (int mm = 0; mm < 2; mm++) {
        float d0 = sm.sEgl[buf][smw + mm * 16 + g_];
        float d1 = sm.sEgl[buf][smw + mm * 16 + g_ + 8];
#pragma unroll
        for (int nn = 0; nn < 4; nn++) {
          S[mm][nn][0] *= d0;
          S[mm][nn][1] *= d0;
          S[mm][nn][2] *= d1;
          S[mm][nn][3] *= d1;
        }
      }
      // A = kg^T (m = K rows), from k-major kg source (rows = chunk j) -> load_A_trans
#pragma unroll
      for (int kb = 0; kb < 4; kb++) {
        uint32_t a0[4], a1[4], bb[4];
        load_A_trans(a0, sm.sKg[buf], smw, kb * 16, S128, lane);
        load_A_trans(a1, sm.sKg[buf], smw + 16, kb * 16, S128, lane);
        load_B_kmajor(bb, sm.sVn, kb * 16, snw, S64, lane);
        mma16816(S[0][0][0], S[0][0][1], S[0][0][2], S[0][0][3], a0[0], a0[1], a0[2], a0[3], bb[0], bb[1]);
        mma16816(S[0][1][0], S[0][1][1], S[0][1][2], S[0][1][3], a0[0], a0[1], a0[2], a0[3], bb[2], bb[3]);
        mma16816(S[1][0][0], S[1][0][1], S[1][0][2], S[1][0][3], a1[0], a1[1], a1[2], a1[3], bb[0], bb[1]);
        mma16816(S[1][1][0], S[1][1][1], S[1][1][2], S[1][1][3], a1[0], a1[1], a1[2], a1[3], bb[2], bb[3]);
        load_B_kmajor(bb, sm.sVn, kb * 16, snw + 16, S64, lane);
        mma16816(S[0][2][0], S[0][2][1], S[0][2][2], S[0][2][3], a0[0], a0[1], a0[2], a0[3], bb[0], bb[1]);
        mma16816(S[0][3][0], S[0][3][1], S[0][3][2], S[0][3][3], a0[0], a0[1], a0[2], a0[3], bb[2], bb[3]);
        mma16816(S[1][2][0], S[1][2][1], S[1][2][2], S[1][2][3], a1[0], a1[1], a1[2], a1[3], bb[0], bb[1]);
        mma16816(S[1][3][0], S[1][3][1], S[1][3][2], S[1][3][3], a1[0], a1[1], a1[2], a1[3], bb[2], bb[3]);
      }
      // rewrite mirror (this warp's own 32x32 region)
      int t_ = lane & 3;
#pragma unroll
      for (int mm = 0; mm < 2; mm++) {
        int r1 = smw + mm * 16 + g_, r2 = r1 + 8;
#pragma unroll
        for (int nn = 0; nn < 4; nn++) {
          int c = snw + nn * 8 + 2 * t_;
          *reinterpret_cast<uint32_t*>(&sm.sS[r1 * S64 + c]) = pack_bf16(S[mm][nn][0], S[mm][nn][1]);
          *reinterpret_cast<uint32_t*>(&sm.sS[r2 * S64 + c]) = pack_bf16(S[mm][nn][2], S[mm][nn][3]);
        }
      }
    }
    __syncthreads();
  }
}

// ===========================================================================
// K3: out kernel. o = qg @ h + Aqk @ vn.  grid (NT*BH).
// ===========================================================================
struct OutSmem {
  bf16 sQg[64 * S128];
  bf16 sH[128 * S128];
  bf16 sAqk[64 * S64];
  bf16 sVn[64 * S128];
};

__global__ void __launch_bounds__(256, 2) kda_out_kernel(
    const bf16* __restrict__ qg, const bf16* __restrict__ hst, const bf16* __restrict__ Aqk,
    const bf16* __restrict__ vn, bf16* __restrict__ o, int T, int H, int NT) {
  extern __shared__ char smem_raw[];
  OutSmem& sm = *reinterpret_cast<OutSmem*>(smem_raw);

  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int wid = tid >> 5;
  const int bh = blockIdx.x / NT;
  const int it = blockIdx.x % NT;
  const int b = bh / H;
  const int h = bh % H;

  {
    const bf16* qgsrc = qg + ((int64_t)bh * NT + it) * BT * DK;
    const bf16* hsrc = hst + ((int64_t)bh * NT + it) * DK * DV;
    const bf16* asrc = Aqk + ((int64_t)bh * NT + it) * BT * BT;
    const bf16* vsrc = vn + ((int64_t)bh * NT + it) * BT * DV;
#pragma unroll
    for (int j = 0; j < 4; j++) {
      int idx = tid + j * 256;
      int r = idx >> 4, c = (idx & 15) << 3;
      cp_async16(&sm.sQg[r * S128 + c], qgsrc + r * DK + c);
      cp_async16(&sm.sVn[r * S128 + c], vsrc + r * DV + c);
    }
#pragma unroll
    for (int j = 0; j < 8; j++) {
      int idx = tid + j * 256;
      int r = idx >> 4, c = (idx & 15) << 3;
      cp_async16(&sm.sH[r * S128 + c], hsrc + r * DV + c);
    }
#pragma unroll
    for (int j = 0; j < 2; j++) {
      int idx = tid + j * 256;
      int r = idx >> 3, c = (idx & 7) << 3;
      cp_async16(&sm.sAqk[r * S64 + c], asrc + r * BT + c);
    }
    cp_commit();
    cp_wait<0>();
  }
  __syncthreads();

  // out 64x128 over 8 warps: 4x2 grid of 16x64
  int m0 = (wid & 3) * 16, n0 = (wid >> 2) * 64;
  float acc[8][4] = {};
  // qg @ h : k = 128 (h k-major)
#pragma unroll
  for (int kb = 0; kb < 8; kb++) {
    uint32_t a[4], bb[4];
    load_A(a, sm.sQg, m0, kb * 16, S128, lane);
#pragma unroll
    for (int nn = 0; nn < 4; nn++) {
      load_B_kmajor(bb, sm.sH, kb * 16, n0 + nn * 16, S128, lane);
      mma16816(acc[nn * 2][0], acc[nn * 2][1], acc[nn * 2][2], acc[nn * 2][3], a[0], a[1], a[2],
               a[3], bb[0], bb[1]);
      mma16816(acc[nn * 2 + 1][0], acc[nn * 2 + 1][1], acc[nn * 2 + 1][2], acc[nn * 2 + 1][3],
               a[0], a[1], a[2], a[3], bb[2], bb[3]);
    }
  }
  // + Aqk @ vn : k = 64 (vn k-major = chunk rows)
#pragma unroll
  for (int kb = 0; kb < 4; kb++) {
    uint32_t a[4], bb[4];
    load_A(a, sm.sAqk, m0, kb * 16, S64, lane);
#pragma unroll
    for (int nn = 0; nn < 4; nn++) {
      load_B_kmajor(bb, sm.sVn, kb * 16, n0 + nn * 16, S128, lane);
      mma16816(acc[nn * 2][0], acc[nn * 2][1], acc[nn * 2][2], acc[nn * 2][3], a[0], a[1], a[2],
               a[3], bb[0], bb[1]);
      mma16816(acc[nn * 2 + 1][0], acc[nn * 2 + 1][1], acc[nn * 2 + 1][2], acc[nn * 2 + 1][3],
               a[0], a[1], a[2], a[3], bb[2], bb[3]);
    }
  }
  // store o (B,T,H,V)
  bf16* odst = o + (((int64_t)b * T + it * BT) * H + h) * DV;
  const int64_t rowstride = (int64_t)H * DV;
  int g_ = lane >> 2, t_ = lane & 3;
#pragma unroll
  for (int nn = 0; nn < 8; nn++) {
    int c = n0 + nn * 8 + 2 * t_;
    *reinterpret_cast<uint32_t*>(odst + (m0 + g_) * rowstride + c) = pack_bf16(acc[nn][0], acc[nn][1]);
    *reinterpret_cast<uint32_t*>(odst + (m0 + g_ + 8) * rowstride + c) = pack_bf16(acc[nn][2], acc[nn][3]);
  }
}

// ===========================================================================
// ===========================================================================
// Host side
// ===========================================================================
static inline int64_t align256(int64_t x) { return (x + 255) & ~255LL; }

std::vector<torch::Tensor> kda_fwd(torch::Tensor q, torch::Tensor k, torch::Tensor v,
                                   torch::Tensor g, torch::Tensor beta, double scale,
                                   bool want_debug) {
  TORCH_CHECK(q.is_cuda() && q.dtype() == torch::kBFloat16);
  TORCH_CHECK(g.dtype() == torch::kFloat32);
  auto B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3);
  auto V = v.size(3);
  TORCH_CHECK(K == 128 && V == 128, "K=V=128 only");
  TORCH_CHECK(T % BT == 0);
  TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() && g.is_contiguous() &&
              beta.is_contiguous());
  int NT = T / BT;
  int BH = B * H;
  int64_t nchunk = (int64_t)BH * NT;

  auto opts = q.options();
  // workspace slab
  int64_t off = 0;
  int64_t o_w = off; off += align256(nchunk * BT * DK * 2);
  int64_t o_u = off; off += align256(nchunk * BT * DV * 2);
  int64_t o_kg = off; off += align256(nchunk * BT * DK * 2);
  int64_t o_qg = off; off += align256(nchunk * BT * DK * 2);
  int64_t o_aqk = off; off += align256(nchunk * BT * BT * 2);
  int64_t o_egl = off; off += align256(nchunk * DK * 4);
  int64_t o_h = off; off += align256(nchunk * DK * DV * 2);
  int64_t o_vn = off; off += align256(nchunk * BT * DV * 2);
  auto slab = torch::empty({off}, opts.dtype(torch::kUInt8));
  char* base = reinterpret_cast<char*>(slab.data_ptr());

  bf16* w_p = reinterpret_cast<bf16*>(base + o_w);
  bf16* u_p = reinterpret_cast<bf16*>(base + o_u);
  bf16* kg_p = reinterpret_cast<bf16*>(base + o_kg);
  bf16* qg_p = reinterpret_cast<bf16*>(base + o_qg);
  bf16* aqk_p = reinterpret_cast<bf16*>(base + o_aqk);
  float* egl_p = reinterpret_cast<float*>(base + o_egl);
  bf16* h_p = reinterpret_cast<bf16*>(base + o_h);
  bf16* vn_p = reinterpret_cast<bf16*>(base + o_vn);

  auto o = torch::empty({B, T, H, V}, opts);
  auto stream = at::cuda::getCurrentCUDAStream();

  {
    static int smem_set = 0;
    if (!smem_set) {
      cudaFuncSetAttribute(kda_prep_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                           sizeof(PrepSmem));
      cudaFuncSetAttribute(kda_relay_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                           sizeof(RelaySmem));
      cudaFuncSetAttribute(kda_out_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                           sizeof(OutSmem));
      smem_set = 1;
    }
  }

  kda_prep_kernel<<<BH * NT, 256, sizeof(PrepSmem), stream>>>(
      reinterpret_cast<const bf16*>(q.data_ptr()), reinterpret_cast<const bf16*>(k.data_ptr()),
      reinterpret_cast<const bf16*>(v.data_ptr()), g.data_ptr<float>(),
      reinterpret_cast<const bf16*>(beta.data_ptr()), w_p, u_p, kg_p, qg_p, aqk_p, egl_p, (int)T,
      (int)H, NT, (float)scale);
  kda_relay_kernel<<<BH * 2, 256, sizeof(RelaySmem), stream>>>(w_p, u_p, kg_p, egl_p, h_p, vn_p,
                                                               NT);
  kda_out_kernel<<<BH * NT, 256, sizeof(OutSmem), stream>>>(qg_p, h_p, aqk_p, vn_p,
                                                            reinterpret_cast<bf16*>(o.data_ptr()),
                                                            (int)T, (int)H, NT);

  std::vector<torch::Tensor> ret;
  ret.push_back(o);
  if (want_debug) {
    auto mk = [&](int64_t offset, std::vector<int64_t> shape, torch::Dtype dt) {
      auto t = torch::empty(shape, opts.dtype(dt));
      cudaMemcpyAsync(t.data_ptr(), base + offset, t.numel() * t.element_size(),
                      cudaMemcpyDeviceToDevice, stream);
      return t;
    };
    ret.push_back(mk(o_w, {BH, NT, BT, DK}, torch::kBFloat16));
    ret.push_back(mk(o_u, {BH, NT, BT, DV}, torch::kBFloat16));
    ret.push_back(mk(o_kg, {BH, NT, BT, DK}, torch::kBFloat16));
    ret.push_back(mk(o_qg, {BH, NT, BT, DK}, torch::kBFloat16));
    ret.push_back(mk(o_aqk, {BH, NT, BT, BT}, torch::kBFloat16));
    ret.push_back(mk(o_egl, {BH, NT, DK}, torch::kFloat32));
    ret.push_back(mk(o_h, {BH, NT, DK, DV}, torch::kBFloat16));
    ret.push_back(mk(o_vn, {BH, NT, BT, DV}, torch::kBFloat16));
  }
  return ret;
}



'''

_CPP_SRC = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> kda_fwd(torch::Tensor q, torch::Tensor k, torch::Tensor v,
                                   torch::Tensor g, torch::Tensor beta, double scale,
                                   bool want_debug);
'''

_ext = None


def _get_ext():
    global _ext
    if _ext is None:
        from torch.utils.cpp_extension import load_inline
        _ext = load_inline(
            name="kda_kbh_v1",
            cpp_sources=[_CPP_SRC],
            cuda_sources=[_CUDA_SRC],
            functions=["kda_fwd"],
            extra_cuda_cflags=[
                "-O3", "-std=c++17", "--use_fast_math", "-lineinfo",
                "-gencode=arch=compute_90a,code=sm_90a",
            ],
            extra_cflags=["-O3", "-std=c++17"],
            verbose=False,
        )
    return _ext


RCP_LN2 = 1.4426950408889634


def _torch_fallback(q, k, v, g, beta, scale, BT):
    """Pure-torch chunked path for unsupported shapes (not used for graded shapes)."""
    B, T, H, K = q.shape
    V = v.shape[-1]
    NT = T // BT
    f = lambda x: x.float()
    qc = f(q).view(B, NT, BT, H, K).permute(0, 3, 1, 2, 4)
    kc = f(k).view(B, NT, BT, H, K).permute(0, 3, 1, 2, 4)
    vc = f(v).view(B, NT, BT, H, V).permute(0, 3, 1, 2, 4)
    gc = f(g).view(B, NT, BT, H, K).permute(0, 3, 1, 2, 4).cumsum(-2) * RCP_LN2
    bc = f(beta).view(B, NT, BT, H).permute(0, 3, 1, 2)
    gn = gc[..., BT // 2 - 1:BT // 2, :]
    kp = kc * torch.exp2(gc - gn)
    kn = kc * torch.exp2(gn - gc)
    Akk = kp @ kn.transpose(-1, -2)
    tril = torch.tril(torch.ones(BT, BT, device=q.device, dtype=torch.bool), -1)
    M = torch.where(tril, -Akk * bc[..., None], torch.zeros((), device=q.device))
    eye = torch.eye(BT, device=q.device)
    X = torch.linalg.solve_triangular(eye - M, eye.expand_as(M), upper=False, unitriangular=True)
    w = X @ (kc * torch.exp2(gc) * bc[..., None])
    u = X @ (vc * bc[..., None])
    qp = qc * torch.exp2(gc - gn) * scale
    trile = torch.tril(torch.ones(BT, BT, device=q.device, dtype=torch.bool), 0)
    Aqk = torch.where(trile, qp @ kn.transpose(-1, -2), torch.zeros((), device=q.device))
    gl = gc[..., -1:, :]
    kg = kc * torch.exp2(gl - gc)
    qg = qc * torch.exp2(gc) * scale
    egl = torch.exp2(gl.squeeze(-2))
    S = q.new_zeros(B, H, K, V, dtype=torch.float32)
    o = torch.empty_like(vc)
    for i in range(NT):
        vn = u[:, :, i] - w[:, :, i] @ S
        o[:, :, i] = qg[:, :, i] @ S + Aqk[:, :, i] @ vn
        S = S * egl[:, :, i, :, None] + kg[:, :, i].transpose(-1, -2) @ vn
    return o.permute(0, 2, 3, 1, 4).reshape(B, T, H, V).to(v.dtype)


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, k, v, g, beta):
        if (
            q.is_cuda and self.chunk_size == 64 and q.shape[-1] == 128
            and v.shape[-1] == 128 and q.shape[1] % 64 == 0
        ):
            qc = q.contiguous()
            kc = k.contiguous()
            vc = v.contiguous()
            gc = g.contiguous().float()
            bc = beta.contiguous()
            return _get_ext().kda_fwd(qc, kc, vc, gc, bc, self.scale, False)[0]
        return _torch_fallback(q, k, v, g, beta, self.scale, 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]