"""DeepSeek-NSA-style block-sparse attention (bench semantics) in raw CUDA.

Semantics (see reference.py): per query t, block importance = mean over causal keys
of the block of (q.k)/sqrt(D); top-8 blocks union the last-64-token sliding window;
softmax attention over the selected keys.

Design (key-major sparse attention, 3 kernels, replayed as a CUDA graph with programmatic
dependent launches between the kernels):
  1. select_kernel      : per query tile (64 queries). Importance is linear
                          (mean(q.k_j) == q.mean(k_j)), so each tile publishes the mean key of
                          its own block exactly split into 3 bf16 parts (hi/mid/lo); tiles are
                          dispatched in block order and wait, chunk by chunk, for the flags of
                          the lower tiles of their head (with a deadlock-free recompute
                          fallback), then score all full blocks with bf16 mma.sync (3 splits,
                          fp32 accumulate) from smem-staged mean keys; the partial current block
                          is scored with the same MMA and a causal prefix mask. Exact top-8 with
                          the reference tie-break (64-bit sortable keys), then a bucketed scatter
                          of (query, slot) pairs per key block.
  2. gather_attn_kernel : one CTA per (head, key block): K/V block in smem, queries that
                          selected it gathered 16 at a time, mma.sync QK^T / PV, writes
                          normalized per-(query, block) partials (m, l, o) flash-decoding style.
  3. window_combine     : per query tile (two output-half CTAs for D=128): dense causal
                          attention for t < 512 (all blocks are selected there), otherwise the
                          sliding-window blocks (bt-1, bt) with the per-query mask, then merges
                          the 8 partials and writes bf16 O; also resets the flags/counters.
"""
from __future__ import annotations

import math
import os

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

B, H, S, D = 1, 16, 1024, 64
BLOCK_SIZE = 64
TOP_N_BLOCKS = 8
SLIDING_WINDOW = 64

CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <stdint.h>
#include <math.h>

namespace nsa {

constexpr int BLK = 64;          // block_size == sliding_window == query tile
constexpr int TOPN = 8;
constexpr int DENSE_TILES = 8;   // query tiles bt < 8 (t < 512) attend to every causal key
constexpr float LOG2E = 1.4426950408889634f;
#ifndef PARTIAL_CHUNK_BYTES
#define PARTIAL_CHUNK_BYTES (80ll << 20)
#endif
#ifndef NSA_SPIN_LIMIT
#define NSA_SPIN_LIMIT 200000
#endif
#ifndef WIN_NSTAGE64
#define WIN_NSTAGE64 2
#endif

#define DEVI __device__ __forceinline__

DEVI float neg_inf() { return __int_as_float(0xff800000u); }

DEVI uint32_t smem_u32(const void* p) { return static_cast<uint32_t>(__cvta_generic_to_shared(p)); }
// Programmatic dependent launch: the next kernel in the stream is launched while this one
// drains (implicit trigger at CTA exit); it waits for full completion before consuming outputs.
#ifdef NSA_NO_PDL
DEVI void pdl_wait() {}
#else
DEVI void pdl_wait() { asm volatile("griddepcontrol.wait;\n" ::: "memory"); }
#endif

DEVI void cp_async16(uint32_t dst, const void* src, bool valid) {
  const int sz = valid ? 16 : 0;
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" [REDACTED: IP]"r"(dst), "l"(src), "r"(sz) : "memory");
}
DEVI void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::: "memory"); }
template <int N>
DEVI void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" [REDACTED: IP]"n"(N) : "memory"); }

DEVI void ldsm_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 ldsm_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));
}
// D = A(16x16 bf16, row) * B(16x8 bf16, col) + D, fp32 accumulate
DEVI void mma16816(float* c, const uint32_t* a, uint32_t b0, uint32_t b1) {
  asm(  // not volatile: pure register op, the compiler may interleave/reorder it
      "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"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
}
DEVI uint32_t pack_bf16x2(float lo, float hi) {
  __nv_bfloat162 t = __floats2bfloat162_rn(lo, hi);
  return *reinterpret_cast<uint32_t*>(&t);
}
DEVI uint32_t pack_half2(float lo, float hi) {
  __half2 t = __floats2half2_rn(lo, hi);
  return *reinterpret_cast<uint32_t*>(&t);
}
DEVI float2 unpack_half2(uint32_t u) { return __half22float2(*reinterpret_cast<__half2*>(&u)); }
DEVI float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
DEVI uint32_t ldg32(const __nv_bfloat16* p) { return __ldg(reinterpret_cast<const uint32_t*>(p)); }
// XOR-swizzled 16B-chunk address inside a [rows][D] bf16 tile (conflict-free ldmatrix)
DEVI uint32_t swz(int row, int chunk, int rowbytes) { return row * rowbytes + ((chunk ^ (row & 7)) << 4); }

// Exact top-8 under the total order (imp, bi) — identical to python's
// sorted(..., reverse=True)[:8] on (imp, bi) tuples. A candidate is a 64-bit key
// (order-preserving float bits << 32 | (bi + 1)); the list is kept sorted descending
// and insertion is a branch-free 8-wide select network (no serial dependency chain).
DEVI uint64_t topk_key(float imp, int bi) {
  const uint32_t f = __float_as_uint(imp);
  const uint32_t s = (f & 0x80000000u) ? ~f : (f | 0x80000000u);
  return ((uint64_t)s << 32) | (uint32_t)(bi + 1);
}
DEVI int topk_block(uint64_t key) { return (int)(uint32_t)key - 1; }
DEVI void topk_insert(uint64_t (&kk)[TOPN], uint64_t x) {
  bool gt[TOPN];
#pragma unroll
  for (int i = 0; i < TOPN; ++i) gt[i] = kk[i] > x;
#pragma unroll
  for (int i = TOPN - 1; i >= 1; --i) kk[i] = gt[i] ? kk[i] : (gt[i - 1] ? x : kk[i - 1]);
  kk[0] = gt[0] ? kk[0] : x;
}
// ------------------------------------------------------------------------------------
// Kernel 1 (fused): block-mean keys + importance scoring + exact top-8 + bucketed scatter.
// grid (n_blocks * B*H) CTAs of 128 threads, tiles in increasing block order (block-major
// blockIdx). Tile bt first publishes the split mean key of its own block, then (for bt >= 8)
// waits chunk by chunk for the flags of the lower tiles of its head before scoring them.
// Lower tiles are dispatched first; a bounded spin with a local recompute fallback keeps
// the wait deadlock-free under any scheduling order.
// ------------------------------------------------------------------------------------
template <int D>
__global__ void __launch_bounds__(128)
select_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k,
              __nv_bfloat16* __restrict__ kbar3, int* __restrict__ flags,
              int16_t* __restrict__ sel, int* __restrict__ count, uint32_t* __restrict__ entries,
              int S, int n_blocks, int BH, int cap, float scale) {
  constexpr int CPR = D / 8, ROWB = D * 2, KTILE = BLK * ROWB;   // K tile: 64 rows
#ifndef SEL_PASS64
#define SEL_PASS64 16
#endif
  constexpr int PASS = (D == 64) ? SEL_PASS64 : 16;                 // full blocks scored per pass
  constexpr int NT = PASS / 8;                                      // n-tiles per pass
  constexpr int IMPC = PASS + 4;                                    // importance row stride (floats)
  constexpr int CHUNK = 3 * PASS * ROWB;                            // staged split mean keys (3 x PASS rows)
#ifndef SEL_NBUF128
#define SEL_NBUF128 1
#endif
#ifndef SEL_NBUF64
#define SEL_NBUF64 2
#endif
  constexpr int NBUF = (D == 64) ? SEL_NBUF64 : SEL_NBUF128;        // chunk staging buffers
  constexpr int REGA = (KTILE > BLK * IMPC * 4) ? KTILE : BLK * IMPC * 4;
  extern __shared__ __align__(128) uint8_t smem[];
  uint8_t* regA = smem;                                             // K tile, later imp matrix, finally merge keys
  uint8_t* regB = smem + REGA;                                      // kbar3 chunk staging (NBUF buffers)
  float* imp = reinterpret_cast<float*>(regA);
  float* s_cur = reinterpret_cast<float*>(regB + NBUF * CHUNK);     // [64]
  int* s_cnt = reinterpret_cast<int*>(s_cur + 64);                  // [n_blocks]
  int* s_base = s_cnt + n_blocks;                                   // [n_blocks]

  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, c = lane & 3;
  const int bt = blockIdx.x / BH, bh = blockIdx.x % BH;  // block-major: lower tiles are dispatched first
  const __nv_bfloat16* qh = q + (size_t)bh * S * D;
  const __nv_bfloat16* kh = k + (size_t)bh * S * D;
  const uint32_t regA_u = smem_u32(regA), regB_u = smem_u32(regB);

  // ---- phase A: K tile of this block -> smem; publish its split mean key ----
  const int nvalid = min(BLK, S - bt * BLK);
  for (int id = tid; id < BLK * CPR; id += 128) {
    const int row = id / CPR, ch = id % CPR;
    const bool ok = row < nvalid;
    cp_async16(regA_u + swz(row, ch, ROWB), kh + (size_t)(bt * BLK + (ok ? row : 0)) * D + ch * 8, ok);
  }
  cp_async_commit();
  cp_async_wait<0>();
  __syncthreads();
  {
    constexpr int NKH = 128 / D, KPT = BLK / NKH;   // threads per column, rows per thread
    const int col = tid % D, kh_ = tid / D;
    float acc = 0.f;
#pragma unroll 8
    for (int j = 0; j < KPT; ++j) {
      const int row = kh_ * KPT + j;
      const __nv_bfloat16* e = reinterpret_cast<const __nv_bfloat16*>(regA + swz(row, col >> 3, ROWB)) + (col & 7);
      acc += bf2f(*e);
    }
    if (NKH == 2) {
      // combine the two row-halves through s_cur (64 floats): half 1 writes, half 0 adds
      if (kh_ == 1) s_cur[col] = acc;
      __syncthreads();
      if (kh_ == 0) acc += s_cur[col];
    }
    if (kh_ == 0) {
      const float m = acc / (float)nvalid;
      const __nv_bfloat16 hi = __float2bfloat16(m);
      const float r1 = m - bf2f(hi);
      const __nv_bfloat16 mid = __float2bfloat16(r1);
      const float r2 = r1 - bf2f(mid);
      const __nv_bfloat16 lo = __float2bfloat16(r2);
      __nv_bfloat16* out = kbar3 + ((size_t)bh * 3 * n_blocks + bt) * D + col;
      out[0] = hi;
      out[(size_t)n_blocks * D] = mid;
      out[2 * (size_t)n_blocks * D] = lo;
    }
  }
  __syncthreads();
  if (tid == 0) {  // cumulative release: publishes every thread's writes ordered by the barrier
    asm volatile("st.release.gpu.s32 [%0], %1;\n" [REDACTED: IP]"l"(flags + bh * n_blocks + bt), "r"(1) : "memory");
  }
  if (bt < DENSE_TILES) return;   // dense tiles: nothing to select

  // ---- Q fragments (standard m16n8k16 A layout; matches ldmatrix'd B operands); issued
  //      after the K tile so they do not delay the mean keys everybody waits for, and
  //      before the flag wait so their latency overlaps it ----
  const int row0 = warp * 16 + g, row1 = row0 + 8;
  uint32_t af[D / 16][4];
  {
    const __nv_bfloat16* q0 = qh + (size_t)min(bt * BLK + row0, S - 1) * D + 2 * c;
    const __nv_bfloat16* q1 = qh + (size_t)min(bt * BLK + row1, S - 1) * D + 2 * c;
#pragma unroll
    for (int s = 0; s < D / 16; ++s) {
      af[s][0] = ldg32(q0 + 16 * s);
      af[s][1] = ldg32(q1 + 16 * s);
      af[s][2] = ldg32(q0 + 16 * s + 8);
      af[s][3] = ldg32(q1 + 16 * s + 8);
    }
  }

  for (int i = tid; i < n_blocks; i += 128) s_cnt[i] = 0;
  // scan role: lane l of warp w owns query row qi = 16w + (l & 15) (a row this warp scored)
  // and half hf = l >> 4 of each pass; the two halves merge with shuffles.
  const int qi = warp * 16 + (lane & 15), hf = lane >> 4;
  const int t = bt * BLK + qi;
  const bool active = t < S;

  const __nv_bfloat16* kb3 = kbar3 + (size_t)bh * 3 * n_blocks * D;
  // Wait until the tiles of chunk [p0, p0+PASS) of this head have published their mean keys.
  // Lower tiles have lower blockIdx and are dispatched first; should a flag still be missing
  // after a long spin (scheduler anomaly), recompute that mean key here (identical values),
  // so the wait can never deadlock. Waiting per chunk (instead of for all lower tiles up
  // front) lets the scoring of low blocks overlap the mean-key phase of higher tiles.
  auto wait_flags = [&](int p0) {
    for (int i = p0 + tid; i < min(p0 + PASS, bt); i += 128) {
      const int* f = flags + bh * n_blocks + i;
      int v, spins = 0;
      do {
        asm volatile("ld.acquire.gpu.s32 %0, [%1];\n" : "=r"(v) : "l"(f) : "memory");
        if (v == 0) { __nanosleep(100); if (++spins > NSA_SPIN_LIMIT) break; }
      } while (v == 0);
      if (v == 0) {  // fallback: compute mean key of block i ourselves (thread-serial, rare path)
        const int nv = min(BLK, S - i * BLK);
        for (int col = 0; col < D; ++col) {
          float acc = 0.f;
          for (int r = 0; r < nv; ++r) acc += bf2f(kh[(size_t)(i * BLK + r) * D + col]);
          const float m = acc / (float)nv;
          const __nv_bfloat16 hi = __float2bfloat16(m);
          const float r1 = m - bf2f(hi);
          const __nv_bfloat16 mid = __float2bfloat16(r1);
          const float r2 = r1 - bf2f(mid);
          __nv_bfloat16* out = kbar3 + ((size_t)bh * 3 * n_blocks + i) * D + col;
          out[0] = hi;
          out[(size_t)n_blocks * D] = mid;
          out[2 * (size_t)n_blocks * D] = __float2bfloat16(r2);
        }
        __threadfence();
      }
    }
  };
  // one cp.async group is committed per call (possibly empty) so that wait_group at the top
  // of a pass always covers that pass's chunk. The caller guarantees (barrier) that the flags
  // of the chunk have been observed by the whole CTA.
  auto issue_chunk = [&](int p0) {
    if (p0 < bt) {
      const int nbp = min(PASS, bt - p0);
      const uint32_t buf = regB_u + ((p0 / PASS) % NBUF) * CHUNK;
      for (int id = tid; id < 3 * PASS * CPR; id += 128) {
        const int row = id / CPR, ch = id % CPR;          // row = sp*PASS + local block
        const int sp = row / PASS, bl = row % PASS;
        const bool ok = bl < nbp;
        cp_async16(buf + swz(row, ch, ROWB), kb3 + ((size_t)sp * n_blocks + p0 + (ok ? bl : 0)) * D + ch * 8, ok);
      }
    }
    cp_async_commit();
  };
  wait_flags(0);
  __syncthreads();
  issue_chunk(0);
  // ---- current (partial) block: mean of causal scores q_t . k_j, j in [64*bt, t] ----
  {
    float acc[8][4];
#pragma unroll
    for (int jn = 0; jn < 8; ++jn) { acc[jn][0] = acc[jn][1] = acc[jn][2] = acc[jn][3] = 0.f; }
    const int krow = (lane & 7) + ((lane >> 4) & 1) * 8, kch = (lane >> 3) & 1;
#pragma unroll
    for (int s = 0; s < D / 16; ++s) {
      uint32_t bk[4][4];
#pragma unroll
      for (int jj = 0; jj < 4; ++jj)
        ldsm_x4(bk[jj][0], bk[jj][1], bk[jj][2], bk[jj][3], regA_u + swz(16 * jj + krow, 2 * s + kch, ROWB));
#pragma unroll
      for (int jj = 0; jj < 4; ++jj) {
        mma16816(acc[2 * jj], af[s], bk[jj][0], bk[jj][1]);
        mma16816(acc[2 * jj + 1], af[s], bk[jj][2], bk[jj][3]);
      }
    }
    float sum0 = 0.f, sum1 = 0.f;
#pragma unroll
    for (int jn = 0; jn < 8; ++jn) {
      const int k0 = 8 * jn + 2 * c;
      sum0 += (k0 <= row0 ? acc[jn][0] : 0.f) + (k0 + 1 <= row0 ? acc[jn][1] : 0.f);
      sum1 += (k0 <= row1 ? acc[jn][2] : 0.f) + (k0 + 1 <= row1 ? acc[jn][3] : 0.f);
    }
    sum0 += __shfl_xor_sync(0xffffffffu, sum0, 1);
    sum0 += __shfl_xor_sync(0xffffffffu, sum0, 2);
    sum1 += __shfl_xor_sync(0xffffffffu, sum1, 1);
    sum1 += __shfl_xor_sync(0xffffffffu, sum1, 2);
    __syncthreads();  // everyone done with the K tile and the s_cur scratch
    if (c == 0) {
      s_cur[row0] = sum0 * scale / (float)(row0 + 1);
      s_cur[row1] = sum1 * scale / (float)(row1 + 1);
    }
  }

  uint64_t kk[TOPN];
#pragma unroll
  for (int i = 0; i < TOPN; ++i) kk[i] = 0ull;  // smaller than any real key

  for (int p0 = 0; p0 < bt; p0 += PASS) {
    const int nbp = min(PASS, bt - p0);
    const uint32_t bufc = regB_u + ((p0 / PASS) % NBUF) * CHUNK;
    wait_flags(p0 + PASS);  // flags of the next chunk (one pass ahead)
    cp_async_wait<0>();
    __syncthreads();  // chunk p0 visible; next chunk's flags observed; previous pass finished
    if (NBUF == 2) issue_chunk(p0 + PASS);  // refill the other buffer; covered by this pass
    // ---- importance MMA: rows = this warp's 16 queries; NT n-tiles (PASS blocks) x 3 splits ----
    {
      float acc[NT][4];
#pragma unroll
      for (int u = 0; u < NT; ++u) { acc[u][0] = acc[u][1] = acc[u][2] = acc[u][3] = 0.f; }
      const int krow = (lane & 7) + ((lane >> 4) & 1) * 8, kch = (lane >> 3) & 1;
#pragma unroll
      for (int s = 0; s < D / 16; ++s) {
        uint32_t bf[3][NT / 2][4];  // [split][jj] : r0=(rows 0-7, chunk 2s) r1=(rows 0-7, 2s+1) r2=(rows 8-15, 2s) r3=(rows 8-15, 2s+1)
#pragma unroll
        for (int sp = 0; sp < 3; ++sp)
#pragma unroll
          for (int jj = 0; jj < NT / 2; ++jj)
            ldsm_x4(bf[sp][jj][0], bf[sp][jj][1], bf[sp][jj][2], bf[sp][jj][3],
                    bufc + swz(sp * PASS + 16 * jj + krow, 2 * s + kch, ROWB));
#pragma unroll
        for (int sp = 0; sp < 3; ++sp)
#pragma unroll
          for (int jj = 0; jj < NT / 2; ++jj) {
            mma16816(acc[2 * jj], af[s], bf[sp][jj][0], bf[sp][jj][1]);
            mma16816(acc[2 * jj + 1], af[s], bf[sp][jj][2], bf[sp][jj][3]);
          }
      }
#pragma unroll
      for (int u = 0; u < NT; ++u) {
        float* wr = imp + row0 * IMPC + 8 * u + 2 * c;
        *reinterpret_cast<float2*>(wr) = make_float2(acc[u][0] * scale, acc[u][1] * scale);
        *reinterpret_cast<float2*>(wr + 8 * IMPC) = make_float2(acc[u][2] * scale, acc[u][3] * scale);
      }
    }
    __syncwarp();  // imp rows of this warp are complete (warp-private)
    if (NBUF == 1) {
      __syncthreads();          // single buffer: everyone must be done reading it
      issue_chunk(p0 + PASS);
    }
    // ---- scan: lane (qi, hf) keeps the top-8 of its half of the pass ----
    {
      const int lo = hf * (PASS / 2);
      const int hi = min(lo + PASS / 2, nbp);
      const float4* rowp = reinterpret_cast<const float4*>(imp + qi * IMPC);
      for (int b4 = lo; b4 < hi; b4 += 4) {
        const float4 v = rowp[b4 >> 2];
        const float vv[4] = {v.x, v.y, v.z, v.w};
#pragma unroll
        for (int e = 0; e < 4; ++e)
          if (b4 + e < hi) topk_insert(kk, topk_key(vv[e], p0 + b4 + e));
      }
    }
    __syncwarp();  // this warp's imp rows are rewritten by its next pass
  }
  // ---- merge of the two halves (exact 64-bit order) via shuffles, then the current block ----
  int pos[TOPN], tb[TOPN];
  const bool emitter = (hf == 0) && active;
  {
    uint64_t other[TOPN];
#pragma unroll
    for (int i = 0; i < TOPN; ++i) {
      const uint32_t lo = __shfl_xor_sync(0xffffffffu, (uint32_t)kk[i], 16);
      const uint32_t hi = __shfl_xor_sync(0xffffffffu, (uint32_t)(kk[i] >> 32), 16);
      other[i] = ((uint64_t)hi << 32) | lo;
    }
    if (hf == 0) {
#pragma unroll
      for (int i = 0; i < TOPN; ++i) topk_insert(kk, other[i]);
    }
  }
  __syncwarp();
  if (hf == 0) {
    topk_insert(kk, topk_key(s_cur[qi], bt));
#pragma unroll
    for (int i = 0; i < TOPN; ++i) tb[i] = topk_block(kk[i]);
  }
  if (emitter) {
    uint4 u;
    u.x = (uint32_t)(uint16_t)tb[0] | ((uint32_t)(uint16_t)tb[1] << 16);
    u.y = (uint32_t)(uint16_t)tb[2] | ((uint32_t)(uint16_t)tb[3] << 16);
    u.z = (uint32_t)(uint16_t)tb[4] | ((uint32_t)(uint16_t)tb[5] << 16);
    u.w = (uint32_t)(uint16_t)tb[6] | ((uint32_t)(uint16_t)tb[7] << 16);
    *reinterpret_cast<uint4*>(sel + ((size_t)bh * S + t) * TOPN) = u;
#pragma unroll
    for (int r = 0; r < TOPN; ++r) pos[r] = (tb[r] >= 0 && tb[r] <= bt - 2) ? atomicAdd(&s_cnt[tb[r]], 1) : -1;
  }
  __syncthreads();
  for (int i = tid; i < n_blocks; i += 128)
    if (s_cnt[i] > 0) s_base[i] = atomicAdd(&count[bh * n_blocks + i], s_cnt[i]);
  __syncthreads();
  if (emitter) {
#pragma unroll
    for (int r = 0; r < TOPN; ++r)
      if (pos[r] >= 0)
        entries[((size_t)bh * n_blocks + tb[r]) * cap + s_base[tb[r]] + pos[r]] = ((uint32_t)t << 3) | (uint32_t)r;
  }
}

// ------------------------------------------------------------------------------------
// Kernel 2: key-major gathered attention. grid (heads, n_blocks - 2), NW warps.
// CTA = one key block; each warp processes groups of 16 gathered queries.
// Writes normalized partials: pml[(bh,t,r)] = (m2, l), po[(bh,t,r)][D] fp16 in the
// quad-interleaved C-fragment order ([i][c][8 halves]) so each 16B store/load
// instruction of a quad covers 64 contiguous bytes.
// ------------------------------------------------------------------------------------
#ifndef G_MINB
#define G_MINB 2
#endif
#ifndef G_MINB128
#define G_MINB128 1
#endif
template <int D, int NW>
__global__ void __launch_bounds__(NW * 32, (D == 64) ? G_MINB : G_MINB128)
gather_attn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k,
                   const __nv_bfloat16* __restrict__ v, const int* __restrict__ count,
                   const uint32_t* __restrict__ entries, float2* __restrict__ pml,
                   __half* __restrict__ po, int S, int n_blocks, int cap, float scale_log2) {
  constexpr int CPR = D / 8, ROWB = D * 2, TILE = BLK * ROWB, QST = 16 * ROWB, NT = NW * 32;
#ifndef NQB64
#define NQB64 1
#endif
#ifndef NQB128
#define NQB128 1
#endif
  constexpr int NQB = (D == 64) ? NQB64 : NQB128;  // Q staging buffers per warp (prefetch depth)
  extern __shared__ __align__(128) uint8_t smem[];
  const int bi = blockIdx.y, bh = blockIdx.x;  // heads fastest: heavy (low) blocks of all heads first
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, c = lane & 3;
  const uint32_t sK_u = smem_u32(smem), sV_u = sK_u + TILE, sQ0_u = sK_u + 2 * TILE + warp * (NQB * QST);
  const uint32_t* ent = entries + ((size_t)bh * n_blocks + bi) * cap;
  const __nv_bfloat16* qbase = q + (size_t)bh * S * D;
  // Wait for the select kernel first (griddepcontrol.wait also drains outstanding loads, so
  // issuing anything before it would serialize the prologue), then issue the pair count, the
  // first entries (speculatively, bounded by the bucket capacity) and the K/V tile loads back
  // to back so their latencies overlap.
  pdl_wait();  // (the dependent grid is launched when this CTA exits: implicit trigger)
  const int n = __ldg(count + bh * n_blocks + bi);
  uint32_t e_cur = 0, e_nxt = 0, e_nn = 0;
  {
    const int i0 = warp * 16 + lane, i1 = i0 + NW * 16, i2 = i1 + NW * 16;
    if (lane < 16) {
      if (i0 < cap) e_cur = __ldg(ent + i0);
      if (i1 < cap) e_nxt = __ldg(ent + i1);
      if (NQB == 2 && i2 < cap) e_nn = __ldg(ent + i2);
    }
  }
  {
    const __nv_bfloat16* kp = k + ((size_t)bh * S + (size_t)bi * BLK) * D;
    const __nv_bfloat16* vp = v + ((size_t)bh * S + (size_t)bi * BLK) * D;
    for (int id = tid; id < BLK * CPR; id += NT) {
      const int row = id / CPR, ch = id % CPR;
      cp_async16(sK_u + swz(row, ch, ROWB), kp + row * D + ch * 8, true);
      cp_async16(sV_u + swz(row, ch, ROWB), vp + row * D + ch * 8, true);
    }
  }
  if (n == 0) {  // uniform across the CTA
    cp_async_commit();
    cp_async_wait<0>();
    return;
  }
  const int ngroups = (n + 15) >> 4;

  // entry for lane's row of a group (lanes >= 16 hold nothing); validity kept separately
  auto load_entries = [&](int gidx, uint32_t& e_out, bool& valid_out) {
    const int idx = gidx * 16 + lane;
    valid_out = (lane < 16) && (idx < n);
    e_out = valid_out ? __ldg(ent + idx) : 0u;
  };
  auto issue_q = [&](uint32_t e, uint32_t sQ_u) {
#pragma unroll
    for (int kk = 0; kk < CPR / 2; ++kk) {
      const int id = lane + 32 * kk;
      const int row = id / CPR, ch = id % CPR;
      const uint32_t er = __shfl_sync(0xffffffffu, e, row);
      const int tr = (int)(er >> 3);
      cp_async16(sQ_u + swz(row, ch, ROWB), qbase + (size_t)tr * D + ch * 8, true);
    }
    cp_async_commit();
  };

  // Software pipeline over this warp's groups (stride NW): bucket entries are loaded
  // NQB+1 groups ahead, the Q rows are gathered NQB groups ahead (one cp.async group
  // committed per iteration so wait_group<NQB-1> always covers the buffer we read next).
  bool v_cur = (lane < 16) && (warp * 16 + lane < n);
  bool v_nxt = (lane < 16) && ((warp + NW) * 16 + lane < n);
  bool v_nn = (NQB == 2) && (lane < 16) && ((warp + 2 * NW) * 16 + lane < n);
  e_cur = v_cur ? e_cur : 0u;
  e_nxt = v_nxt ? e_nxt : 0u;
  e_nn = v_nn ? e_nn : 0u;
  if (warp < ngroups) issue_q(e_cur, sQ0_u);        // group A (also carries this thread's K/V chunks)
  else cp_async_commit();
  if (NQB == 2) {
    if (warp + NW < ngroups) issue_q(e_nxt, sQ0_u + QST);  // group B
    else cp_async_commit();
    cp_async_wait<1>();
  } else {
    cp_async_wait<0>();
  }
  __syncthreads();

  int it = 0;
  for (int gi = warp; gi < ngroups; gi += NW, ++it) {
    const uint32_t sQ_u = sQ0_u + (NQB == 2 ? (it & 1) * QST : 0);
    uint32_t af[D / 16][4];
    {
      const int row = (lane & 7) + ((lane >> 3) & 1) * 8, chs = lane >> 4;
#pragma unroll
      for (int s = 0; s < D / 16; ++s)
        ldsm_x4(af[s][0], af[s][1], af[s][2], af[s][3], sQ_u + swz(row, 2 * s + chs, ROWB));
    }
    const uint32_t e0 = __shfl_sync(0xffffffffu, e_cur, g), e1 = __shfl_sync(0xffffffffu, e_cur, g + 8);
    const bool val0 = __shfl_sync(0xffffffffu, (int)v_cur, g) != 0;
    const bool val1 = __shfl_sync(0xffffffffu, (int)v_cur, g + 8) != 0;
    __syncwarp();
    uint32_t e_new = 0;
    bool v_new = false;
    if (NQB == 1) {
      // entries for gi+2NW; Q of gi+NW (entries loaded one iteration ago) into the single buffer
      if (gi + 2 * NW < ngroups) load_entries(gi + 2 * NW, e_new, v_new);
      if (gi + NW < ngroups) issue_q(e_nxt, sQ_u); else cp_async_commit();
    } else {
      // entries for gi+3NW; Q of gi+2NW (entries loaded one iteration ago) into this buffer,
      // which is free now that its fragments sit in registers
      if (gi + 3 * NW < ngroups) load_entries(gi + 3 * NW, e_new, v_new);
      if (gi + 2 * NW < ngroups) issue_q(e_nn, sQ_u); else cp_async_commit();
    }

    // ---- S = Q K^T (16 x 64) ----
    float s[8][4];
#pragma unroll
    for (int j = 0; j < 8; ++j) { s[j][0] = s[j][1] = s[j][2] = s[j][3] = 0.f; }
    {
      const int krow = (lane & 7) + ((lane >> 4) & 1) * 8, kch = (lane >> 3) & 1;
#pragma unroll
      for (int ks = 0; ks < D / 16; ++ks) {
        uint32_t bk[4][4];
#pragma unroll
        for (int jj = 0; jj < 4; ++jj)
          ldsm_x4(bk[jj][0], bk[jj][1], bk[jj][2], bk[jj][3], sK_u + swz(16 * jj + krow, 2 * ks + kch, ROWB));
#pragma unroll
        for (int jj = 0; jj < 4; ++jj) {
          mma16816(s[2 * jj], af[ks], bk[jj][0], bk[jj][1]);
          mma16816(s[2 * jj + 1], af[ks], bk[jj][2], bk[jj][3]);
        }
      }
    }
    // ---- softmax over the full block (all 64 keys valid & causal) ----
    float mx0 = neg_inf(), mx1 = neg_inf();
#pragma unroll
    for (int j = 0; j < 8; ++j) {
      mx0 = fmaxf(mx0, fmaxf(s[j][0], s[j][1]));
      mx1 = fmaxf(mx1, fmaxf(s[j][2], s[j][3]));
    }
    mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1));
    mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2));
    mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1));
    mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2));
    const float m0 = mx0 * scale_log2, m1 = mx1 * scale_log2;
    float l0 = 0.f, l1 = 0.f;
    uint32_t pa[4][4];
#pragma unroll
    for (int j = 0; j < 8; ++j) {
      const float p0 = exp2f(fmaf(s[j][0], scale_log2, -m0));
      const float p1 = exp2f(fmaf(s[j][1], scale_log2, -m0));
      const float p2 = exp2f(fmaf(s[j][2], scale_log2, -m1));
      const float p3 = exp2f(fmaf(s[j][3], scale_log2, -m1));
      l0 += p0 + p1;
      l1 += p2 + p3;
      pa[j >> 1][(j & 1) * 2 + 0] = pack_bf16x2(p0, p1);
      pa[j >> 1][(j & 1) * 2 + 1] = pack_bf16x2(p2, p3);
    }
    l0 += __shfl_xor_sync(0xffffffffu, l0, 1);
    l0 += __shfl_xor_sync(0xffffffffu, l0, 2);
    l1 += __shfl_xor_sync(0xffffffffu, l1, 1);
    l1 += __shfl_xor_sync(0xffffffffu, l1, 2);
    // ---- O = P V (16 x D) ----
    float o[D / 8][4];
#pragma unroll
    for (int j = 0; j < D / 8; ++j) { o[j][0] = o[j][1] = o[j][2] = o[j][3] = 0.f; }
    {
      const int vrow = (lane & 7) + ((lane >> 3) & 1) * 8, vch = lane >> 4;
#pragma unroll
      for (int kk = 0; kk < 4; ++kk) {
        uint32_t bv[D / 16][4];
#pragma unroll
        for (int jp = 0; jp < D / 16; ++jp)
          ldsm_x4_trans(bv[jp][0], bv[jp][1], bv[jp][2], bv[jp][3], sV_u + swz(16 * kk + vrow, 2 * jp + vch, ROWB));
#pragma unroll
        for (int jp = 0; jp < D / 16; ++jp) {
          mma16816(o[2 * jp], pa[kk], bv[jp][0], bv[jp][1]);
          mma16816(o[2 * jp + 1], pa[kk], bv[jp][2], bv[jp][3]);
        }
      }
    }
    // ---- write normalized partials ----
    if (val0) {
      const int t0 = (int)(e0 >> 3), r0 = (int)(e0 & 7);
      const size_t pidx = ((size_t)bh * S + t0) * TOPN + r0;
      const float inv = 1.f / l0;
      uint4* dst = reinterpret_cast<uint4*>(po + pidx * D) + c;
#pragma unroll
      for (int i = 0; i < D / 32; ++i) {
        uint4 w;
        w.x = pack_half2(o[4 * i][0] * inv, o[4 * i][1] * inv);
        w.y = pack_half2(o[4 * i + 1][0] * inv, o[4 * i + 1][1] * inv);
        w.z = pack_half2(o[4 * i + 2][0] * inv, o[4 * i + 2][1] * inv);
        w.w = pack_half2(o[4 * i + 3][0] * inv, o[4 * i + 3][1] * inv);
        dst[4 * i] = w;
      }
      if (c == 0) pml[pidx] = make_float2(m0, l0);
    }
    if (val1) {
      const int t1 = (int)(e1 >> 3), r1 = (int)(e1 & 7);
      const size_t pidx = ((size_t)bh * S + t1) * TOPN + r1;
      const float inv = 1.f / l1;
      uint4* dst = reinterpret_cast<uint4*>(po + pidx * D) + c;
#pragma unroll
      for (int i = 0; i < D / 32; ++i) {
        uint4 w;
        w.x = pack_half2(o[4 * i][2] * inv, o[4 * i][3] * inv);
        w.y = pack_half2(o[4 * i + 1][2] * inv, o[4 * i + 1][3] * inv);
        w.z = pack_half2(o[4 * i + 2][2] * inv, o[4 * i + 2][3] * inv);
        w.w = pack_half2(o[4 * i + 3][2] * inv, o[4 * i + 3][3] * inv);
        dst[4 * i] = w;
      }
      if (c == 0) pml[pidx] = make_float2(m1, l1);
    }
    e_cur = e_nxt; v_cur = v_nxt;
    if (NQB == 1) { e_nxt = e_new; v_nxt = v_new; }
    else { e_nxt = e_nn; v_nxt = v_nn; e_nn = e_new; v_nn = v_new; }
    if (NQB == 2) cp_async_wait<1>(); else cp_async_wait<0>();
    __syncwarp();
  }
}

// ------------------------------------------------------------------------------------
// Kernel 3: sliding-window attention (+ dense causal for t < 512) and partial merge.
// grid (heads * output-halves, n_blocks), 128 threads = 4 warps x 16 query rows.
// ------------------------------------------------------------------------------------
DEVI bool has_block(uint4 u, int b) {
  const uint32_t w[4] = {u.x, u.y, u.z, u.w};
  bool f = false;
#pragma unroll
  for (int i = 0; i < 4; ++i) {
    f |= ((int)(int16_t)(w[i] & 0xffffu) == b);
    f |= ((int)(int16_t)(w[i] >> 16) == b);
  }
  return f;
}
DEVI int sel_at(uint4 u, int r) {
  const uint32_t w = (r < 2) ? u.x : (r < 4) ? u.y : (r < 6) ? u.z : u.w;
  return (r & 1) ? (int)(int16_t)(w >> 16) : (int)(int16_t)(w & 0xffffu);
}

template <int D>
__global__ void __launch_bounds__(128, 3)
window_combine_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k,
                      const __nv_bfloat16* __restrict__ v, const int16_t* __restrict__ sel,
                      const float2* __restrict__ pml, const __half* __restrict__ po,
                      __nv_bfloat16* __restrict__ out, int* __restrict__ flags, int* __restrict__ count,
                      int S, int n_blocks, float scale_log2) {
  constexpr int DH = (D == 128) ? 2 : 1;        // output-half CTAs per query tile
  constexpr int DO = D / DH;                    // output columns handled by this CTA
  constexpr int CPR = D / 8, ROWB = D * 2, TILE_K = BLK * ROWB;
  constexpr int CPRV = DO / 8, ROWBV = DO * 2, TILE_V = BLK * ROWBV;
  constexpr int NSTAGE = (D == 64) ? WIN_NSTAGE64 : 1;
  constexpr int MB = (D == 64) ? 8 : 4;         // partial-merge batch (slots loaded together)
  extern __shared__ __align__(128) uint8_t smem[];
  const int bt = blockIdx.y, dh = blockIdx.x % DH, bh = blockIdx.x / DH;  // heads fastest
  const bool dense = bt < DENSE_TILES;
  const int nblk = dense ? bt + 1 : 2;
  const int blk0 = dense ? 0 : bt - 1;
  const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, g = lane >> 2, c = lane & 3;
  const __nv_bfloat16* qb = q + (size_t)bh * S * D;
  const __nv_bfloat16* kb = k + (size_t)bh * S * D;
  const __nv_bfloat16* vb = v + (size_t)bh * S * D + dh * DO;
  const uint32_t sKV_u = smem_u32(smem);  // stage st: K at st*(TILE_K+TILE_V), V at +TILE_K

  auto issue_kv = [&](int kbi, int stage) {
    const uint32_t sk = sKV_u + stage * (TILE_K + TILE_V), sv = sk + TILE_K;
    for (int id = tid; id < BLK * CPR; id += 128) {
      const int row = id / CPR, ch = id % CPR;
      const int key = kbi * BLK + row;
      const bool valid = key < S;
      cp_async16(sk + swz(row, ch, ROWB), kb + (size_t)(valid ? key : 0) * D + ch * 8, valid);
    }
    for (int id = tid; id < BLK * CPRV; id += 128) {
      const int row = id / CPRV, ch = id % CPRV;
      const int key = kbi * BLK + row;
      const bool valid = key < S;
      cp_async16(sv + swz(row, ch, ROWBV), vb + (size_t)(valid ? key : 0) * D + ch * 8, valid);
    }
  };
  pdl_wait();  // wait for the gather kernel (its outputs are consumed below)
  issue_kv(blk0, 0);
  cp_async_commit();

  const int t0 = bt * BLK + warp * 16 + g, t1 = t0 + 8;
  const bool ok0 = t0 < S, ok1 = t1 < S;
  // Q A-fragments straight from global (standard m16n8k16 layout, matches ldmatrix'd K)
  uint32_t af[D / 16][4];
  {
    const __nv_bfloat16* q0 = qb + (size_t)min(t0, S - 1) * D + 2 * c;
    const __nv_bfloat16* q1 = qb + (size_t)min(t1, S - 1) * D + 2 * c;
#pragma unroll
    for (int s = 0; s < D / 16; ++s) {
      af[s][0] = ldg32(q0 + 16 * s);
      af[s][1] = ldg32(q1 + 16 * s);
      af[s][2] = ldg32(q0 + 16 * s + 8);
      af[s][3] = ldg32(q1 + 16 * s + 8);
    }
  }
  uint4 sel0 = make_uint4(0, 0, 0, 0), sel1 = make_uint4(0, 0, 0, 0);
  bool prev0 = false, prev1 = false;
  if (!dense) {
    if (ok0) sel0 = __ldg(reinterpret_cast<const uint4*>(sel + ((size_t)bh * S + t0) * TOPN));
    if (ok1) sel1 = __ldg(reinterpret_cast<const uint4*>(sel + ((size_t)bh * S + t1) * TOPN));
    prev0 = has_block(sel0, bt - 1);
    prev1 = has_block(sel1, bt - 1);
  }

  float m2[2] = {neg_inf(), neg_inf()}, l[2] = {0.f, 0.f};
  float o[DO / 8][4];
#pragma unroll
  for (int j = 0; j < DO / 8; ++j) { o[j][0] = o[j][1] = o[j][2] = o[j][3] = 0.f; }

  for (int ib = 0; ib < nblk; ++ib) {
    if (NSTAGE == 1 && ib > 0) {
      __syncthreads();  // everyone done with the previous block
      issue_kv(blk0 + ib, 0);
      cp_async_commit();
    }
    cp_async_wait<0>();
    __syncthreads();
    if (NSTAGE == 2 && ib + 1 < nblk) {
      issue_kv(blk0 + ib + 1, (ib + 1) & 1);
      cp_async_commit();
    }
    const int stage = (NSTAGE == 2) ? (ib & 1) : 0;
    const uint32_t sK_u = sKV_u + stage * (TILE_K + TILE_V), sV_u = sK_u + TILE_K;
    const int kbi = blk0 + ib, kbase = kbi * BLK;

    float s[8][4];
#pragma unroll
    for (int j = 0; j < 8; ++j) { s[j][0] = s[j][1] = s[j][2] = s[j][3] = 0.f; }
    {
      const int krow = (lane & 7) + ((lane >> 4) & 1) * 8, kch = (lane >> 3) & 1;
#pragma unroll
      for (int ks = 0; ks < D / 16; ++ks) {
        uint32_t bk[4][4];
#pragma unroll
        for (int jj = 0; jj < 4; ++jj)
          ldsm_x4(bk[jj][0], bk[jj][1], bk[jj][2], bk[jj][3], sK_u + swz(16 * jj + krow, 2 * ks + kch, ROWB));
#pragma unroll
        for (int jj = 0; jj < 4; ++jj) {
          mma16816(s[2 * jj], af[ks], bk[jj][0], bk[jj][1]);
          mma16816(s[2 * jj + 1], af[ks], bk[jj][2], bk[jj][3]);
        }
      }
    }
    // per-row key validity range [lo, hi]
    int lo[2], hi[2];
    {
      const bool cur = (kbi == bt);
      hi[0] = cur ? t0 : 0x7fffffff;
      hi[1] = cur ? t1 : 0x7fffffff;
      lo[0] = (cur || dense || prev0) ? -1 : (t0 - (BLK - 1));
      lo[1] = (cur || dense || prev1) ? -1 : (t1 - (BLK - 1));
    }
    float mx[2] = {neg_inf(), neg_inf()};
#pragma unroll
    for (int j = 0; j < 8; ++j) {
#pragma unroll
      for (int e = 0; e < 4; ++e) {
        const int key = kbase + 8 * j + 2 * c + (e & 1);
        const int r = e >> 1;
        const bool valid = (key >= lo[r]) && (key <= hi[r]);
        s[j][e] = valid ? s[j][e] : neg_inf();
        mx[r] = fmaxf(mx[r], s[j][e]);
      }
    }
    float alpha[2], msafe[2];
#pragma unroll
    for (int r = 0; r < 2; ++r) {
      mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 1));
      mx[r] = fmaxf(mx[r], __shfl_xor_sync(0xffffffffu, mx[r], 2));
      const float mnew = fmaxf(m2[r], mx[r] * scale_log2);
      msafe[r] = (mnew == neg_inf()) ? 0.f : mnew;
      alpha[r] = exp2f(m2[r] - msafe[r]);  // m2 == -inf -> 0 (o, l are 0 anyway)
      m2[r] = mnew;
    }
    float ls[2] = {0.f, 0.f};
    uint32_t pa[4][4];
#pragma unroll
    for (int j = 0; j < 8; ++j) {
      const float p0 = exp2f(fmaf(s[j][0], scale_log2, -msafe[0]));
      const float p1 = exp2f(fmaf(s[j][1], scale_log2, -msafe[0]));
      const float p2 = exp2f(fmaf(s[j][2], scale_log2, -msafe[1]));
      const float p3 = exp2f(fmaf(s[j][3], scale_log2, -msafe[1]));
      ls[0] += p0 + p1;
      ls[1] += p2 + p3;
      pa[j >> 1][(j & 1) * 2 + 0] = pack_bf16x2(p0, p1);
      pa[j >> 1][(j & 1) * 2 + 1] = pack_bf16x2(p2, p3);
    }
#pragma unroll
    for (int r = 0; r < 2; ++r) {
      ls[r] += __shfl_xor_sync(0xffffffffu, ls[r], 1);
      ls[r] += __shfl_xor_sync(0xffffffffu, ls[r], 2);
      l[r] = l[r] * alpha[r] + ls[r];
    }
#pragma unroll
    for (int j = 0; j < DO / 8; ++j) {
      o[j][0] *= alpha[0]; o[j][1] *= alpha[0];
      o[j][2] *= alpha[1]; o[j][3] *= alpha[1];
    }
    {
      const int vrow = (lane & 7) + ((lane >> 3) & 1) * 8, vch = lane >> 4;
#pragma unroll
      for (int kk = 0; kk < 4; ++kk) {
        uint32_t bv[DO / 16][4];
#pragma unroll
        for (int jp = 0; jp < DO / 16; ++jp)
          ldsm_x4_trans(bv[jp][0], bv[jp][1], bv[jp][2], bv[jp][3], sV_u + swz(16 * kk + vrow, 2 * jp + vch, ROWBV));
#pragma unroll
        for (int jp = 0; jp < DO / 16; ++jp) {
          mma16816(o[2 * jp], pa[kk], bv[jp][0], bv[jp][1]);
          mma16816(o[2 * jp + 1], pa[kk], bv[jp][2], bv[jp][3]);
        }
      }
    }
  }

  // ---- merge the gathered-block partials (t >= 512 only); this CTA's output half only ----
  if (!dense) {
    constexpr int NCH = DO / 32;  // 16-byte chunks of the partial row per thread for this half
#pragma unroll
    for (int r = 0; r < 2; ++r) {
      const int t = r ? t1 : t0;
      const bool ok = r ? ok1 : ok0;
      const uint4 su = r ? sel1 : sel0;
      if (!ok) continue;
      const size_t prow = ((size_t)bh * S + t) * TOPN;
      bool valid[TOPN];
      float2 ml[TOPN];
#pragma unroll
      for (int slot = 0; slot < TOPN; ++slot) {
        valid[slot] = sel_at(su, slot) <= bt - 2;
        ml[slot] = valid[slot] ? __ldg(pml + prow + slot) : make_float2(neg_inf(), 0.f);
      }
#pragma unroll
      for (int b0 = 0; b0 < TOPN; b0 += MB) {
        uint4 w[MB][NCH];
#pragma unroll
        for (int sIdx = 0; sIdx < MB; ++sIdx) {
          const int slot = b0 + sIdx;
          const uint4* src = reinterpret_cast<const uint4*>(po + (prow + slot) * D) + c;
#pragma unroll
          for (int i = 0; i < NCH; ++i)
            w[sIdx][i] = valid[slot] ? __ldg(src + 4 * (dh * NCH + i)) : make_uint4(0, 0, 0, 0);
        }
#pragma unroll
        for (int sIdx = 0; sIdx < MB; ++sIdx) {
          const int slot = b0 + sIdx;
          const float mnew = fmaxf(m2[r], ml[slot].x);
          const float a = exp2f(m2[r] - mnew);
          const float b = valid[slot] ? exp2f(ml[slot].x - mnew) * ml[slot].y : 0.f;
          l[r] = l[r] * a + b;
          m2[r] = mnew;
#pragma unroll
          for (int i = 0; i < NCH; ++i) {
            const uint32_t ww[4] = {w[sIdx][i].x, w[sIdx][i].y, w[sIdx][i].z, w[sIdx][i].w};
#pragma unroll
            for (int e = 0; e < 4; ++e) {
              const float2 f = unpack_half2(ww[e]);
              const int j = 4 * i + e;
              o[j][2 * r] = o[j][2 * r] * a + f.x * b;
              o[j][2 * r + 1] = o[j][2 * r + 1] * a + f.y * b;
            }
          }
        }
      }
    }
  }
  // ---- reset the dependency flag and pair counter of this tile for the next call ----
  if (tid == 0 && dh == 0 && flags != nullptr) {
    flags[bh * n_blocks + bt] = 0;
    count[bh * n_blocks + bt] = 0;
  }
  // ---- normalize + store this CTA's output half ----
#pragma unroll
  for (int r = 0; r < 2; ++r) {
    const int t = r ? t1 : t0;
    const bool ok = r ? ok1 : ok0;
    if (!ok) continue;
    const float inv = 1.f / l[r];
    uint32_t* dst = reinterpret_cast<uint32_t*>(out + ((size_t)bh * S + t) * D + dh * DO) + c;
#pragma unroll
    for (int j = 0; j < DO / 8; ++j) dst[4 * j] = pack_bf16x2(o[j][2 * r] * inv, o[j][2 * r + 1] * inv);
  }
}

// ------------------------------------------------------------------------------------
// host
// ------------------------------------------------------------------------------------
template <typename... KArgs, typename... AArgs>
static void launch_pdl(void (*kern)(KArgs...), dim3 grid, dim3 block, size_t smem, cudaStream_t stream,
                       AArgs... args) {
  cudaLaunchConfig_t cfg = {};
  cfg.gridDim = grid;
  cfg.blockDim = block;
  cfg.dynamicSmemBytes = smem;
  cfg.stream = stream;
  cudaLaunchAttribute attr;
  attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
  attr.val.programmaticStreamSerializationAllowed = 1;
#ifdef NSA_NO_PDL
  cfg.numAttrs = 0;
#else
  cfg.attrs = &attr;
  cfg.numAttrs = 1;
#endif
  cudaLaunchKernelEx(&cfg, kern, static_cast<KArgs>(args)...);
}
#ifdef NSA_DEBUG
#define DBG_CHECK(name) do { cudaError_t e_ = cudaDeviceSynchronize(); if (e_ == cudaSuccess) e_ = cudaGetLastError(); \
  TORCH_CHECK(e_ == cudaSuccess, "kernel ", name, " failed: ", cudaGetErrorString(e_)); } while (0)
#else
#define DBG_CHECK(name) do {} while (0)
#endif
template <int D>
void launch_all(const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, at::Tensor& out,
                const at::Tensor& kbar3, const at::Tensor& sel, const at::Tensor& count,
                const at::Tensor& entries, const at::Tensor& pml, const at::Tensor& po,
                const at::Tensor& flags,
                int BH, int S, cudaStream_t stream) {
  constexpr int ROWB = D * 2, TILE = BLK * ROWB;
#ifndef NW3_128
#define NW3_128 4
#endif
#ifndef NW3_64
#define NW3_64 8
#endif
  constexpr int NW3 = (D == 64) ? NW3_64 : NW3_128;
  const int n_blocks = (S + BLK - 1) / BLK;
  const float scale = (float)(1.0 / std::sqrt((double)D));
  const float scale_log2 = scale * LOG2E;
  const auto* qp = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>());
  const auto* kp = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>());
  const auto* vp = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<at::BFloat16>());
  auto* op = reinterpret_cast<__nv_bfloat16*>(out.data_ptr<at::BFloat16>());

  constexpr int SPASS = (D == 64) ? SEL_PASS64 : 16, SNBUF = (D == 64) ? SEL_NBUF64 : SEL_NBUF128;
  constexpr int SREGA = (BLK * D * 2 > BLK * (SPASS + 4) * 4) ? BLK * D * 2 : BLK * (SPASS + 4) * 4;
  const int smem2 = SREGA + SNBUF * 3 * SPASS * D * 2 + 64 * 4 + 2 * n_blocks * 4 + 16;
  const int smem3 = 2 * TILE + NW3 * ((D == 64) ? NQB64 : NQB128) * 16 * ROWB;
  constexpr int WDH = (D == 128) ? 2 : 1;
  const int smem4 = ((D == 64) ? WIN_NSTAGE64 : 1) * (TILE + TILE / WDH);
  static bool attr_set = false;
  if (!attr_set) {
    cudaFuncSetAttribute(select_kernel<D>, cudaFuncAttributeMaxDynamicSharedMemorySize, 64 * 1024);
    cudaFuncSetAttribute(gather_attn_kernel<D, NW3>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem3);
    cudaFuncSetAttribute(window_combine_kernel<D>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem4);
    attr_set = true;
  }

  const int16_t* sel_p = nullptr;
  const float2* pml_p = nullptr;
  const __half* po_p = nullptr;
  if (S > BLK * DENSE_TILES) {
    const int cap = S - BLK * DENSE_TILES;
    TORCH_CHECK(kbar3.numel() >= (long)BH * 3 * n_blocks * D, "kbar3 workspace too small");
    TORCH_CHECK(flags.numel() >= (long)BH * n_blocks, "flag workspace too small");
    TORCH_CHECK(sel.numel() >= (long)BH * S * TOPN, "sel workspace too small");
    TORCH_CHECK(count.numel() >= (long)BH * n_blocks, "count workspace too small");
    TORCH_CHECK(entries.numel() >= (long)BH * n_blocks * cap, "entries workspace too small");
    TORCH_CHECK(pml.numel() >= (long)BH * S * TOPN * 2, "pml workspace too small");
    TORCH_CHECK(po.numel() >= (long)BH * S * TOPN * D, "po workspace too small");
    auto* kbar_p = reinterpret_cast<__nv_bfloat16*>(kbar3.data_ptr<at::BFloat16>());
    int16_t* sel_w = sel.data_ptr<int16_t>();
    int* count_p = count.data_ptr<int>();
    uint32_t* ent_p = reinterpret_cast<uint32_t*>(entries.data_ptr<int>());
    float2* pml_w = reinterpret_cast<float2*>(pml.data_ptr<float>());
    __half* po_w = reinterpret_cast<__half*>(po.data_ptr<at::Half>());

    int* flags_p = flags.data_ptr<int>();
    select_kernel<D><<<dim3(n_blocks * BH), 128, smem2, stream>>>(
        qp, kp, kbar_p, flags_p, sel_w, count_p, ent_p, S, n_blocks, BH, cap, scale);
    DBG_CHECK("select");
    sel_p = sel_w;
    pml_p = pml_w;
    po_p = po_w;
    // Partials are written by the gather kernel and read back by the combine kernel;
    // chunk the heads so each chunk's partials stay L2-resident (128 MB L2).
    // The partial buffer is reused by every chunk (same addresses stay dirty in L2 and
    // are never written back inside the timed region).
    const double po_bytes_per_head = (double)S * TOPN * (D * 2 + 8);
    int heads_per_chunk = BH;
    if (po_bytes_per_head * BH > (double)PARTIAL_CHUNK_BYTES) {
      heads_per_chunk = (int)((double)PARTIAL_CHUNK_BYTES / po_bytes_per_head);
      heads_per_chunk = heads_per_chunk < 1 ? 1 : (heads_per_chunk > BH ? BH : heads_per_chunk);
    }
    for (int h0 = 0; h0 < BH; h0 += heads_per_chunk) {
      const int nh = (BH - h0 < heads_per_chunk) ? (BH - h0) : heads_per_chunk;
      const size_t off = (size_t)h0 * S * D;
      launch_pdl(gather_attn_kernel<D, NW3>, dim3(nh, n_blocks - 2), dim3(NW3 * 32), smem3, stream,
          qp + off, kp + off, vp + off, count_p + h0 * n_blocks, ent_p + (size_t)h0 * n_blocks * cap,
          pml_w, po_w, S, n_blocks, cap, scale_log2);
      DBG_CHECK("gather");
      launch_pdl(window_combine_kernel<D>, dim3(nh * WDH, n_blocks), dim3(128), smem4, stream,
          qp + off, kp + off, vp + off, sel_w + (size_t)h0 * S * TOPN, pml_w, po_w, op + off,
          flags_p + h0 * n_blocks, count_p + h0 * n_blocks, S, n_blocks, scale_log2);
      DBG_CHECK("window");
    }
    return;
  }
  window_combine_kernel<D><<<dim3(BH * WDH, n_blocks), 128, smem4, stream>>>(
      qp, kp, vp, sel_p, pml_p, po_p, op, nullptr, nullptr, S, n_blocks, scale_log2);
}

}  // namespace nsa

torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v,
                          torch::Tensor kbar3, torch::Tensor sel, torch::Tensor count,
                          torch::Tensor entries, torch::Tensor pml, torch::Tensor po,
                          torch::Tensor flags) {
  TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "inputs must be CUDA tensors");
  TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 &&
              v.scalar_type() == at::kBFloat16, "inputs must be bf16");
  TORCH_CHECK(q.dim() == 4, "q must be (B,H,S,D)");
  TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "inputs must be contiguous");
  TORCH_CHECK(q.sizes() == k.sizes() && q.sizes() == v.sizes(), "q,k,v shapes must match");
  const int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
  TORCH_CHECK(S >= 1 && S <= 16384, "S out of supported range");
  TORCH_CHECK(D == 64 || D == 128, "head dim must be 64 or 128");
  auto out = at::empty_like(q);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (D == 64) nsa::launch_all<64>(q, k, v, out, kbar3, sel, count, entries, pml, po, flags, B * H, S, stream);
  else nsa::launch_all<128>(q, k, v, out, kbar3, sel, count, entries, pml, po, flags, B * H, S, stream);
  return out;
}
"""

CPP_SRC = (
    "torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, "
    "torch::Tensor kbar3, torch::Tensor sel, torch::Tensor count, torch::Tensor entries, "
    "torch::Tensor pml, torch::Tensor po, torch::Tensor flags);"
)

_ext = None


def _ensure_ninja_on_path():
    import shutil
    import sys

    if shutil.which("ninja"):
        return
    cands = []
    try:
        import ninja  # type: ignore

        cands.append(getattr(ninja, "BIN_DIR", None))
    except Exception:
        pass
    cands.append(os.path.dirname(os.path.realpath(sys.executable)))
    cands.append(os.path.dirname(sys.executable))
    for d in cands:
        if d and os.path.exists(os.path.join(d, "ninja")):
            os.environ["PATH"] = d + os.pathsep + os.environ.get("PATH", "")
            return


def _load():
    global _ext
    if _ext is None:
        _ensure_ninja_on_path()
        _ext = load_inline(
            name="nsa_sparse_attn_final" + os.environ.get("NSA_BUILD_TAG", ""),
            cpp_sources=CPP_SRC,
            cuda_sources=CUDA_SRC,
            functions=["nsa_forward"],
            extra_cuda_cflags=[
                *os.environ.get("NSA_EXTRA_FLAGS", "").split(),
                "-O3",
                "--use_fast_math",
                "-std=c++17",
                "-lineinfo",
                "-gencode=arch=compute_120,code=sm_120",
            ],
            verbose=os.environ.get("NSA_VERBOSE", "0") == "1",
        )
    return _ext


class _Workspace:
    """Per-shape scratch buffers (allocated once, reused across calls)."""

    def __init__(self, B: int, H: int, S: int, D: int, device):
        BH = B * H
        n_blocks = (S + BLOCK_SIZE - 1) // BLOCK_SIZE
        sparse = S > BLOCK_SIZE * 8
        cap = S - BLOCK_SIZE * 8 if sparse else 0
        n = (lambda x: x if sparse else 0)
        self.kbar3 = torch.empty(n(BH * 3 * n_blocks * D), dtype=torch.bfloat16, device=device)
        self.sel = torch.empty(n(BH * S * TOP_N_BLOCKS), dtype=torch.int16, device=device)
        # zero-initialized: reset for the next call by the window kernel, never by a memset
        self.count = torch.zeros(n(BH * n_blocks), dtype=torch.int32, device=device)
        self.entries = torch.empty(n(BH * n_blocks * cap), dtype=torch.int32, device=device)
        self.pml = torch.empty(n(BH * S * TOP_N_BLOCKS * 2), dtype=torch.float32, device=device)
        self.po = torch.empty(n(BH * S * TOP_N_BLOCKS * D), dtype=torch.float16, device=device)
        # dependency flags of the select kernel (binary; reset for the next call by window_combine)
        self.flags = torch.zeros(max(1, BH * n_blocks), dtype=torch.int32, device=device)


_USE_GRAPHS = os.environ.get("NSA_NO_GRAPH", "0") != "1"


class Model(nn.Module):
    def __init__(self, B: int, H: int, S: int, D: int):
        super().__init__()
        self.B, self.H, self.S, self.D = B, H, S, D
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16))
        self._ws = {}
        self._graphs = {}
        _load()

    def _workspace(self, q: torch.Tensor) -> _Workspace:
        key = (tuple(q.shape), q.device.index)
        ws = self._ws.get(key)
        if ws is None:
            ws = _Workspace(*q.shape, q.device)
            self._ws[key] = ws
        return ws

    def _run(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        ws = self._workspace(q)
        return _load().nsa_forward(q, k, v, ws.kbar3, ws.sel, ws.count, ws.entries, ws.pml, ws.po, ws.flags)

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
        if not _USE_GRAPHS or torch.cuda.is_current_stream_capturing():
            return self._run(q, k, v)
        # The whole forward (3 kernel launches) is replayed as a CUDA graph. The graph is
        # keyed on the input buffers; it reads them at replay time, so new contents at the
        # same addresses are handled correctly.
        key = (tuple(q.shape), q.dtype, q.data_ptr(), k.data_ptr(), v.data_ptr(), q.device.index)
        entry = self._graphs.get(key)
        if entry is None:
            try:
                self._run(q, k, v)  # warm-up: extension load, workspace allocation, attributes
                torch.cuda.synchronize()
                g = torch.cuda.CUDAGraph()
                with torch.cuda.graph(g):
                    out = self._run(q, k, v)
                entry = (g, out)
                self._graphs[key] = entry
            except Exception:
                self._graphs[key] = None
                return self._run(q, k, v)
        if entry is None:
            return self._run(q, k, v)
        g, out = entry
        g.replay()
        return out


def get_init_inputs():
    return [B, H, S, D]


def get_inputs():
    q = torch.randn(B, H, S, D, dtype=torch.bfloat16, device="cuda" if torch.cuda.is_available() else "cpu")
    k = torch.randn(B, H, S, D, dtype=torch.bfloat16, device=q.device)
    v = torch.randn(B, H, S, D, dtype=torch.bfloat16, device=q.device)
    return [q, k, v]