"""DeepSeek NSA-style block-sparse attention (bench semantics), CUDA implementation for sm_120.

Pipeline (custom CUDA kernels, JIT-built with torch.utils.cpp_extension.load_inline):
  1. prep      : per (b,h,key-block) compensated sum of the 64 keys, kept as three bf16 pieces (fp32-exact) in MMA
                 fragment order
  2. select    : block importance q . mean(k) via tensor cores (diagonal block from the causal score tile), exact top-8
                 (ties -> larger block index), emitted as per-key-block row masks and per-row block bitmaps
  3. attention : key-stationary: one CTA per key block gathers every row that selected it (plus the local-window rows)
                 into dense 16-row tiles (mma.sync) and writes int8 partial outputs + (max, sum, scale) per (row, block);
                 for D = 64 and S <= 1100 selection and attention are fused in one row-stationary kernel (8-row warps,
                 K/V streamed through an mbarrier ring)
  4. merge     : softmax-weighted combination of the partials into the bf16 output
"""
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_CUDA_SRC = r'''
// NSA-style sparse attention v1: prep (block-sum splits) -> select (importance + top-8 bitmaps) -> attend
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cstdint>
#include <type_traits>
#include <map>
#include <tuple>

typedef __nv_bfloat16 bf16;
__device__ int g_pf_ctr;   // dispatch-order ticket counter of the fused kernel (reset by the prep kernel that precedes it)

#define NEG_INF (-INFINITY)

// ---------------------------------------------------------------- PTX helpers
__device__ __forceinline__ uint32_t smem_u32(const void* p) { return (uint32_t)__cvta_generic_to_shared(p); }
__device__ __forceinline__ void cp_async16(uint32_t dst, const void* src, int src_bytes) {
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(dst), "l"(src), "r"(src_bytes));
}
__device__ __forceinline__ void cp_async4(uint32_t dst, const void* src, int src_bytes) {
    asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %2;\n" ::"r"(dst), "l"(src), "r"(src_bytes));
}
__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n" ::); }
template <int N> __device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }

__device__ __forceinline__ 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));
}
__device__ __forceinline__ void ldsm_x4_t(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));
}
__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, 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"(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));
}
__device__ __forceinline__ float ex2f(float x) {
    float y;
    asm("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(y) : "f"(x));
    return y;
}
__device__ __forceinline__ uint32_t movtrans(uint32_t a) {
    uint32_t d;
    asm volatile("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;\n" : "=r"(d) : "r"(a));
    return d;
}
__device__ __forceinline__ uint32_t pack_h2(float lo, float hi) {
    __half2 h = __floats2half2_rn(lo, hi);
    return *reinterpret_cast<uint32_t*>(&h);
}
__device__ __forceinline__ float2 unpack_h2(uint32_t v) {
    __half2 h = *reinterpret_cast<__half2*>(&v);
    return __half22float2(h);
}

__device__ __forceinline__ void mbar_init(uint32_t addr, uint32_t count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(addr), "r"(count));
}
__device__ __forceinline__ void mbar_arrive(uint32_t addr) {
    asm volatile("{\n .reg .b64 st;\n mbarrier.arrive.shared::cta.b64 st, [%0];\n}\n" ::"r"(addr) : "memory");
}
__device__ __forceinline__ void mbar_wait(uint32_t addr, uint32_t parity) {
    asm volatile(
        "{\n"
        " .reg .pred p;\n"
        "WAIT_%=:\n"
        " mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n"
        " @!p bra WAIT_%=;\n"
        "}\n" ::"r"(addr), "r"(parity) : "memory");
}
__device__ __forceinline__ void cp_mbar_arrive_noinc(uint32_t addr) {
    asm volatile("cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n" ::"r"(addr) : "memory");
}
__device__ __forceinline__ uint32_t pack_bf16(float lo, float hi) {
    __nv_bfloat162 t = __floats2bfloat162_rn(lo, hi);
    return *reinterpret_cast<uint32_t*>(&t);
}

// swizzled byte offset of 16B chunk `ch` in row `r` of a [rows][D] bf16 tile
template <int D>
__device__ __forceinline__ uint32_t swz(int r, int ch) {
    return (uint32_t)(r * (D * 2) + ((ch ^ (r & 7)) << 4));
}

// ---------------------------------------------------------------- prep kernel
// For each (bh, key block b): fp64 sum over the 64 keys, split into 3 bf16 pieces, stored in
// mma B-fragment order: word index ((((bh*NT + j)*32 + lane)*3 + s)*KS + ks)*2 + reg
__device__ __forceinline__ void two_sum_acc(float& hi, float& lo, float x) {
    float s = hi + x;
    float bb = s - hi;
    float err = (hi - (s - bb)) + (x - bb);
    hi = s;
    lo += err;
}

template <int D>
__global__ void __launch_bounds__(256) prep_kernel(const bf16* __restrict__ K, uint32_t* __restrict__ kcf, int S, int NT) {
    constexpr int NP = D / 2;
    constexpr int KG = 256 / NP;
    constexpr int KPG = 64 / KG;
    constexpr int KS = D / 16;
    __shared__ float red[KG][NP][4];
    int b = blockIdx.x, bh = blockIdx.y;
    int tid = threadIdx.x;
    if (b == 0 && bh == 0 && tid == 0) g_pf_ctr = 0;
    int p = tid % NP, g = tid / NP;
    const __nv_bfloat162* Kp = reinterpret_cast<const __nv_bfloat162*>(K + (size_t)bh * S * D);
    float ha = 0.f, la = 0.f, hb = 0.f, lb = 0.f;
#pragma unroll
    for (int i = 0; i < KPG; ++i) {
        int key = b * 64 + g * KPG + i;
        if (key < S) {
            __nv_bfloat162 v = Kp[(size_t)key * NP + p];
            two_sum_acc(ha, la, __low2float(v));
            two_sum_acc(hb, lb, __high2float(v));
        }
    }
    red[g][p][0] = ha; red[g][p][1] = la; red[g][p][2] = hb; red[g][p][3] = lb;
    __syncthreads();
    if (g == 0) {
#pragma unroll
        for (int gg = 1; gg < KG; ++gg) {
            two_sum_acc(ha, la, red[gg][p][0]);
            la += red[gg][p][1];
            two_sum_acc(hb, lb, red[gg][p][2]);
            lb += red[gg][p][3];
        }
        // renormalize: value = ha + la
        float sa = ha + la; la = la - (sa - ha); ha = sa;
        float sb = hb + lb; lb = lb - (sb - hb); hb = sb;
        int d0 = 2 * p;
        int ks = d0 / 16, within = d0 % 16, reg = within / 8, c = (within % 8) / 2;
        int j = b / 8, nl = b % 8;
        int lane = nl * 4 + c;
        size_t base = (((size_t)bh * NT + j) * 32 + lane) * 3;
        float ra = ha, rb = hb;
#pragma unroll
        for (int s = 0; s < 3; ++s) {
            bf16 xa = __float2bfloat16_rn(ra);
            bf16 xb = __float2bfloat16_rn(rb);
            ra -= __bfloat162float(xa);
            rb -= __bfloat162float(xb);
            if (s == 0) { ra += la; rb += lb; }
            __nv_bfloat162 pk;
            pk.x = xa;
            pk.y = xb;
            kcf[((base + s) * KS + ks) * 2 + reg] = *reinterpret_cast<uint32_t*>(&pk);
        }
    }
}

__device__ __forceinline__ uint32_t u4get(const uint4& v, int i) { return i == 0 ? v.x : (i == 1 ? v.y : (i == 2 ? v.z : v.w)); }


// ---- 8-element descending sort (19 compare-exchanges) and top-8 merge of two descending lists (bitonic half-cleaners)
__device__ __forceinline__ void ce_desc(float& a, float& b) {
    const float hi = fmaxf(a, b), lo = fminf(a, b);
    a = hi;
    b = lo;
}
__device__ __forceinline__ void sort8_desc(float (&x)[8]) {
    ce_desc(x[0], x[2]); ce_desc(x[1], x[3]); ce_desc(x[4], x[6]); ce_desc(x[5], x[7]);
    ce_desc(x[0], x[4]); ce_desc(x[1], x[5]); ce_desc(x[2], x[6]); ce_desc(x[3], x[7]);
    ce_desc(x[0], x[1]); ce_desc(x[2], x[3]); ce_desc(x[4], x[5]); ce_desc(x[6], x[7]);
    ce_desc(x[2], x[4]); ce_desc(x[3], x[5]);
    ce_desc(x[1], x[4]); ce_desc(x[3], x[6]);
    ce_desc(x[1], x[2]); ce_desc(x[3], x[4]); ce_desc(x[5], x[6]);
}
// lv (descending, running top 8) <- top 8 of lv and xs (descending)
__device__ __forceinline__ void top8_merge(float (&lv)[8], const float (&xs)[8]) {
    float cc[8];
#pragma unroll
    for (int k = 0; k < 8; ++k) cc[k] = fmaxf(lv[k], xs[7 - k]);
    ce_desc(cc[0], cc[4]); ce_desc(cc[1], cc[5]); ce_desc(cc[2], cc[6]); ce_desc(cc[3], cc[7]);
    ce_desc(cc[0], cc[2]); ce_desc(cc[1], cc[3]); ce_desc(cc[4], cc[6]); ce_desc(cc[5], cc[7]);
    ce_desc(cc[0], cc[1]); ce_desc(cc[2], cc[3]); ce_desc(cc[4], cc[5]); ce_desc(cc[6], cc[7]);
#pragma unroll
    for (int k = 0; k < 8; ++k) lv[k] = cc[k];
}

// ---------------------------------------------------------------- select v2
// K diag tile aliases the start of the importance array (dead once the diag pass is done); early prefetch of Q/Kc fragments;
// value-only branch-free top-8; no `sel` output (attend kernels only read the colmask).
template <int D, int DBG = 0>
__global__ void __launch_bounds__(128, 2) select_kernel(const bf16* __restrict__ Q, const bf16* __restrict__ K,
                                                        const uint32_t* __restrict__ kcf, uint2* __restrict__ colmask, uint32_t* __restrict__ sel,
                                                        int NBLK, int S, int NT, int Spad, float scale) {
    constexpr int KS = D / 16;
    constexpr int NBP = 132;
    constexpr int NV = (3 * KS * 2) / 4;
    extern __shared__ __align__(128) uint8_t smem_raw[];
    uint8_t* Kd = smem_raw;                                             // [64][D] swizzled, aliases imp until diag pass done
    uint8_t* Qs = smem_raw + 64 * D * 2;                                // [64][D] swizzled Q tile (also aliases imp)
    static_assert(2 * 64 * D * 2 <= 64 * 132 * 4, "K and Q tiles must fit in the importance array");
    float* imp = reinterpret_cast<float*>(smem_raw);                    // [64][NBP]
    uint32_t* cm = reinterpret_cast<uint32_t*>(smem_raw + 64 * NBP * 4);  // [128][2]
    const int m = blockIdx.x, bh = blockIdx.y;
    const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
    long long T0 = 0, T1 = 0, T2 = 0, T3 = 0, T4 = 0;
    if (DBG) T0 = clock64();
    uint2* cmrow = colmask + ((size_t)bh * NBLK + m) * 128;
    uint32_t* selrow = sel + ((size_t)bh * Spad + (size_t)m * 64) * 4;
    if (m < 8) {
        if (tid < 64) {
            uint4 v4 = make_uint4((1u << (m + 1)) - 1u, 0u, 0u, 0u);
            if (m * 64 + tid >= S) v4 = make_uint4(0u, 0u, 0u, 0u);
            *reinterpret_cast<uint4*>(selrow + tid * 4) = v4;
        }
        int nrows = min(64, S - m * 64);
        unsigned long long vm = (nrows >= 64) ? ~0ull : ((1ull << nrows) - 1ull);
        uint2 v = (tid < m) ? make_uint2((uint32_t)vm, (uint32_t)(vm >> 32)) : make_uint2(0u, 0u);
        cmrow[tid] = v;
        return;
    }
    cm[tid * 2] = 0u;
    cm[tid * 2 + 1] = 0u;
    const bf16* Kh = K + (size_t)bh * S * D;
    const bf16* Qh = Q + (size_t)bh * S * D;
    {
        constexpr int CH = D / 8;
        for (int i = tid; i < 64 * CH; i += 128) {
            int r = i / CH, ch = i % CH;
            int key = m * 64 + r;
            const bf16* src = Kh + (size_t)(key < S ? key : 0) * D + ch * 8;
            cp_async16(smem_u32(Kd) + swz<D>(r, ch), src, key < S ? 16 : 0);
        }
        for (int i = tid; i < 64 * CH; i += 128) {
            int r = i / CH, ch = i % CH;
            int row = m * 64 + r;
            const bf16* src = Qh + (size_t)(row < S ? row : 0) * D + ch * 8;
            cp_async16(smem_u32(Qs) + swz<D>(r, ch), src, row < S ? 16 : 0);
        }
        cp_commit();
    }
    const int c4 = lane & 3;
    const int ntm = (m + 7) >> 3;
    uint32_t qf[(D == 64) ? 4 : 1][KS][4];
    uint4 kvc[NV], kvn[NV];
    int j = (D == 64) ? warp : (warp & 1);
    if (j < ntm) {
        const uint4* pp = reinterpret_cast<const uint4*>(kcf + (((size_t)bh * NT + j) * 32 + lane) * (3 * KS * 2));
#pragma unroll
        for (int i = 0; i < NV; ++i) kvc[i] = pp[i];
    }
    cp_wait<0>();
    __syncthreads();
    uint32_t qd[KS][4];
    uint32_t qp[(D == 128) ? 2 : 1][KS][4];   // D = 128: A fragments of the two row tiles (2P, 2P + 1) of this warp's pair, P = warp >> 1
    if constexpr (D == 128) {
#pragma unroll
        for (int h = 0; h < 2; ++h) {
#pragma unroll
            for (int ks = 0; ks < KS; ++ks) {
                const int qrow = (2 * (warp >> 1) + h) * 16 + (lane & 7) + ((lane >> 3) & 1) * 8;
                ldsm_x4(qp[h][ks][0], qp[h][ks][1], qp[h][ks][2], qp[h][ks][3], smem_u32(Qs) + swz<D>(qrow, 2 * ks + (lane >> 4)));
            }
        }
#pragma unroll
        for (int ks = 0; ks < KS; ++ks) {
#pragma unroll
            for (int i = 0; i < 4; ++i) qd[ks][i] = (warp & 1) ? qp[1][ks][i] : qp[0][ks][i];   // own tile (warp) for the diagonal pass
        }
    } else {
#pragma unroll
        for (int ks = 0; ks < KS; ++ks) {
            const int qrow = warp * 16 + (lane & 7) + ((lane >> 3) & 1) * 8;
            ldsm_x4(qd[ks][0], qd[ks][1], qd[ks][2], qd[ks][3], smem_u32(Qs) + swz<D>(qrow, 2 * ks + (lane >> 4)));
        }
    }
    if constexpr (D == 64) {
#pragma unroll
        for (int t = 0; t < 4; ++t) {
#pragma unroll
            for (int ks = 0; ks < KS; ++ks) {
                const int qrow = t * 16 + (lane & 7) + ((lane >> 3) & 1) * 8;
                ldsm_x4(qf[t][ks][0], qf[t][ks][1], qf[t][ks][2], qf[t][ks][3], smem_u32(Qs) + swz<D>(qrow, 2 * ks + (lane >> 4)));
            }
        }
    }
    // diag scores
    float sc[8][4];
#pragma unroll
    for (int i = 0; i < 8; ++i) sc[i][0] = sc[i][1] = sc[i][2] = sc[i][3] = 0.f;
#pragma unroll
    for (int ks = 0; ks < KS; ++ks) {
#pragma unroll
        for (int np = 0; np < 4; ++np) {
            uint32_t b0, b1, b2, b3;
            int mat = lane >> 3;
            int row = np * 16 + ((mat >> 1) << 3) + (lane & 7);
            int ch = 2 * ks + (mat & 1);
            ldsm_x4(b0, b1, b2, b3, smem_u32(Kd) + swz<D>(row, ch));
            mma16816(sc[2 * np], qd[ks], b0, b1);
            mma16816(sc[2 * np + 1], qd[ks], b2, b3);
        }
    }
    const int o0 = warp * 16 + (lane >> 2), o1 = o0 + 8;
    float sum0 = 0.f, sum1 = 0.f;
#pragma unroll
    for (int nt = 0; nt < 8; ++nt) {
#pragma unroll
        for (int e = 0; e < 2; ++e) {
            int cidx = nt * 8 + c4 * 2 + e;
            if (cidx <= o0) sum0 += sc[nt][e];
            if (cidx <= o1) sum1 += sc[nt][2 + e];
        }
    }
    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();   // every warp is done reading the K tile: its smem becomes the importance array
    if (c4 == 0) {
        imp[o0 * NBP + m] = sum0 * scale / (float)(o0 + 1);
        imp[o1 * NBP + m] = sum1 * scale / (float)(o1 + 1);
    }
    if (DBG) T1 = clock64();
    const float scale64 = scale * (1.0f / 64.0f);
    if constexpr (D == 64) {
        for (; j < ntm; j += 4) {
            int jn = j + 4;
            if (jn < ntm) {
                const uint4* pp = reinterpret_cast<const uint4*>(kcf + (((size_t)bh * NT + jn) * 32 + lane) * (3 * KS * 2));
#pragma unroll
                for (int i = 0; i < NV; ++i) kvn[i] = pp[i];
            }
            int blk = j * 8 + c4 * 2;
#pragma unroll
            for (int t = 0; t < 4; ++t) {
                float acc[3][4];
#pragma unroll
                for (int s3 = 0; s3 < 3; ++s3) {
                    acc[s3][0] = acc[s3][1] = acc[s3][2] = acc[s3][3] = 0.f;
#pragma unroll
                    for (int ks = 0; ks < KS; ++ks) {
                        int w = (s3 * KS + ks) * 2;
                        mma16816(acc[s3], qf[t][ks], u4get(kvc[w >> 2], w & 3), u4get(kvc[(w + 1) >> 2], (w + 1) & 3));
                    }
                }
                int ra = t * 16 + (lane >> 2), rb = ra + 8;
                float v00 = (acc[0][0] + (acc[1][0] + acc[2][0])) * scale64;
                float v01 = (acc[0][1] + (acc[1][1] + acc[2][1])) * scale64;
                float v10 = (acc[0][2] + (acc[1][2] + acc[2][2])) * scale64;
                float v11 = (acc[0][3] + (acc[1][3] + acc[2][3])) * scale64;
                if (blk < m) imp[ra * NBP + blk] = v00;
                if (blk + 1 < m) imp[ra * NBP + blk + 1] = v01;
                if (blk < m) imp[rb * NBP + blk] = v10;
                if (blk + 1 < m) imp[rb * NBP + blk + 1] = v11;
            }
#pragma unroll
            for (int i = 0; i < NV; ++i) kvc[i] = kvn[i];
        }
    } else {
        // D = 128: warp w owns the row-tile pair (2P, 2P + 1) and every other 8-block group (jj = (w & 1) + 2i): each group's pieces are loaded by
        // half as many warps and every loaded group feeds two tiles
        for (int jj = (warp & 1); jj < ntm; jj += 2) {
            if (jj + 2 < ntm) {
                const uint4* pp = reinterpret_cast<const uint4*>(kcf + (((size_t)bh * NT + jj + 2) * 32 + lane) * (3 * KS * 2));
#pragma unroll
                for (int i = 0; i < NV; ++i) kvn[i] = pp[i];
            }
            float acc[2][3][4];
#pragma unroll
            for (int h = 0; h < 2; ++h)
#pragma unroll
                for (int s3 = 0; s3 < 3; ++s3) acc[h][s3][0] = acc[h][s3][1] = acc[h][s3][2] = acc[h][s3][3] = 0.f;
#pragma unroll
            for (int s3 = 0; s3 < 3; ++s3) {
#pragma unroll
                for (int ks = 0; ks < KS; ++ks) {
                    const int w = (s3 * KS + ks) * 2;
                    const uint32_t b0 = u4get(kvc[w >> 2], w & 3), b1 = u4get(kvc[(w + 1) >> 2], (w + 1) & 3);
                    mma16816(acc[0][s3], qp[0][ks], b0, b1);
                    mma16816(acc[1][s3], qp[1][ks], b0, b1);
                }
            }
            const int blk = jj * 8 + c4 * 2;
#pragma unroll
            for (int h = 0; h < 2; ++h) {
                const int ra = (2 * (warp >> 1) + h) * 16 + (lane >> 2), rb = ra + 8;
                float v00 = (acc[h][0][0] + (acc[h][1][0] + acc[h][2][0])) * scale64;
                float v01 = (acc[h][0][1] + (acc[h][1][1] + acc[h][2][1])) * scale64;
                float v10 = (acc[h][0][2] + (acc[h][1][2] + acc[h][2][2])) * scale64;
                float v11 = (acc[h][0][3] + (acc[h][1][3] + acc[h][2][3])) * scale64;
                if (blk < m) imp[ra * NBP + blk] = v00;
                if (blk + 1 < m) imp[ra * NBP + blk + 1] = v01;
                if (blk < m) imp[rb * NBP + blk] = v10;
                if (blk + 1 < m) imp[rb * NBP + blk + 1] = v11;
            }
#pragma unroll
            for (int i = 0; i < NV; ++i) kvc[i] = kvn[i];
        }
    }
    __syncthreads();
    if (DBG) T2 = clock64();
    // top-8: 2 threads per row (lane and lane^16). Value-only branch-free sorted insertion, then T8, then emission.
    {
        const int row = (warp << 4) + (lane & 15);
        const int half = lane >> 4;
        const float* ir = imp + row * NBP;
        const int npairs = (m >> 3) + 1;
        float lv[8];
#pragma unroll
        for (int k = 0; k < 8; ++k) lv[k] = NEG_INF;
        for (int pp = half; pp < npairs; pp += 2) {
            float4 x0 = *reinterpret_cast<const float4*>(ir + 8 * pp);
            float4 x1 = *reinterpret_cast<const float4*>(ir + 8 * pp + 4);
            float xs[8] = {x0.x, x0.y, x0.z, x0.w, x1.x, x1.y, x1.z, x1.w};
            if (pp == (m >> 3)) {
#pragma unroll
                for (int e = 0; e < 8; ++e) if (8 * pp + e > m) xs[e] = NEG_INF;
            }
            sort8_desc(xs);
            top8_merge(lv, xs);
        }
        float T8 = INFINITY;
        {
            float ovv[8];
#pragma unroll
            for (int k = 0; k < 8; ++k) ovv[k] = __shfl_xor_sync(0xffffffffu, lv[k], 16);
#pragma unroll
            for (int k = 0; k < 8; ++k) T8 = fminf(T8, fmaxf(lv[k], ovv[7 - k]));
        }
        uint32_t bm[4] = {0u, 0u, 0u, 0u};
#pragma unroll
        for (int w = 0; w < 4; ++w) {
#pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int pp = 4 * w + half + 2 * i;
                if (pp < npairs) {
                    float4 x0 = *reinterpret_cast<const float4*>(ir + 8 * pp);
                    float4 x1 = *reinterpret_cast<const float4*>(ir + 8 * pp + 4);
                    float xs[8] = {x0.x, x0.y, x0.z, x0.w, x1.x, x1.y, x1.z, x1.w};
                    uint32_t byte = 0u;
#pragma unroll
                    for (int e = 0; e < 8; ++e) {
                        bool ge = (xs[e] >= T8) && (8 * pp + e <= m);
                        byte |= ge ? (1u << e) : 0u;
                    }
                    bm[w] |= byte << (8 * (pp & 3));
                }
            }
        }
#pragma unroll
        for (int w = 0; w < 4; ++w) bm[w] |= __shfl_xor_sync(0xffffffffu, bm[w], 16);
        const int t = m * 64 + row;
        const bool row_ok = t < S;
        int total = __popc(bm[0]) + __popc(bm[1]) + __popc(bm[2]) + __popc(bm[3]);
        if (row_ok && total > 8) {
            // exact ties at the 8th value: larger block index wins
            int cgt = 0;
            for (int b2 = 0; b2 <= m; ++b2) cgt += (ir[b2] > T8) ? 1 : 0;
            int need = 8 - cgt;
            bm[0] = bm[1] = bm[2] = bm[3] = 0u;
            for (int b2 = m; b2 >= 0; --b2) {
                float v2 = ir[b2];
                bool take = (v2 > T8) || (v2 == T8 && need > 0);
                if (v2 == T8 && need > 0) --need;
                if (take) bm[b2 >> 5] |= 1u << (b2 & 31);
            }
        }
        if (!row_ok) { bm[0] = bm[1] = bm[2] = bm[3] = 0u; }
        if (half == 0) {
            uint32_t s0 = bm[0], s1 = bm[1], s2 = bm[2], s3 = bm[3];
            if (row_ok) {   // the stored set S_t also holds the row's own block and its predecessor (the local-window blocks)
                {
                    const int wm = m >> 5;
                    const uint32_t bmk = 1u << (m & 31);
                    s0 |= (wm == 0) ? bmk : 0u; s1 |= (wm == 1) ? bmk : 0u; s2 |= (wm == 2) ? bmk : 0u; s3 |= (wm == 3) ? bmk : 0u;
                }
                if (m > 0) {
                    const int m1 = m - 1;
                    const int wm = m1 >> 5;
                    const uint32_t bmk = 1u << (m1 & 31);
                    s0 |= (wm == 0) ? bmk : 0u; s1 |= (wm == 1) ? bmk : 0u; s2 |= (wm == 2) ? bmk : 0u; s3 |= (wm == 3) ? bmk : 0u;
                }
            }
            *reinterpret_cast<uint4*>(selrow + row * 4) = make_uint4(s0, s1, s2, s3);
        }
#pragma unroll
        for (int w = 0; w < 2; ++w) {
            int ww = half * 2 + w;
            uint32_t x = (ww == 0) ? bm[0] : (ww == 1) ? bm[1] : (ww == 2) ? bm[2] : bm[3];
            while (x) {
                int bpos = __ffs(x) - 1;
                x &= x - 1;
                atomicOr(&cm[(ww * 32 + bpos) * 2 + (row >> 5)], 1u << (row & 31));
            }
        }
    }
    if (DBG) T3 = clock64();
    __syncthreads();
    cmrow[tid] = make_uint2(cm[tid * 2], cm[tid * 2 + 1]);
    if (DBG && tid == 0 && (m == NBLK - 1 || m == NBLK / 2)) {
        T4 = clock64();
        int o = (m == NBLK - 1) ? 0 : 5;
        (void)o; (void)T1; (void)T2; (void)T3; (void)T4;
    }
}


// ---------------------------------------------------------------- B8: static 8-row warps, transposed MMA passes, register state
template <int D, int MODE>
__device__ __forceinline__ void run_pass8(uint32_t kaddr, uint32_t vaddr, const uint32_t* __restrict__ koff,
                                          const uint32_t* __restrict__ voff, const uint32_t (*qb)[2], int lane,
                                          uint32_t keep8, int oadd, int mt_lo, int mt_hi, float cscale, float (*o)[4],
                                          float& mA, float& mB, float& lA, float& lB) {
    constexpr int KS = D / 16;
    constexpr int MTS = 16 * D * 2;
    constexpr bool FULL = (MODE == 0);
    const int j = lane >> 2, c = lane & 3;
    float s[4][4];
#pragma unroll
    for (int mt = 0; mt < 4; ++mt) s[mt][0] = s[mt][1] = s[mt][2] = s[mt][3] = 0.f;
#pragma unroll
    for (int ks = 0; ks < KS; ++ks) {
        const uint32_t kb0 = kaddr + koff[ks];
#pragma unroll
        for (int mt = 0; mt < 4; ++mt) {
            if (FULL || (mt >= mt_lo && mt < mt_hi)) {
                uint32_t a[4];
                ldsm_x4(a[0], a[1], a[2], a[3], kb0 + mt * MTS);
                mma16816(s[mt], a, qb[ks][0], qb[ks][1]);
            }
        }
    }
    // ---- masks. columns: A = row 2c, B = row 2c+1 of this warp's 8-row group
    const bool kA = ((keep8 >> (2 * c)) & 1u) != 0u;
    const bool kB = ((keep8 >> (2 * c + 1)) & 1u) != 0u;
    if (MODE == 0) {
#pragma unroll
        for (int mt = 0; mt < 4; ++mt) {
            if (!kA) { s[mt][0] = NEG_INF; s[mt][2] = NEG_INF; }
            if (!kB) { s[mt][1] = NEG_INF; s[mt][3] = NEG_INF; }
        }
    } else {
        const int oA = oadd + 2 * c, oB = oA + 1;
#pragma unroll
        for (int mt = 0; mt < 4; ++mt) {
#pragma unroll
            for (int r = 0; r < 4; ++r) {
                int key = 16 * mt + j + ((r >> 1) << 3);
                int oo = (r & 1) ? oB : oA;
                bool keep = (r & 1) ? kB : kA;
                bool allowed = (MODE == 1) ? (key <= oo) : (keep || key >= oo + 1);
                if (!allowed) s[mt][r] = NEG_INF;
            }
        }
    }
    // ---- online softmax (state in registers)
    float mxA = fmaxf(fmaxf(s[0][0], s[0][2]), fmaxf(s[1][0], s[1][2]));
    mxA = fmaxf(mxA, fmaxf(fmaxf(s[2][0], s[2][2]), fmaxf(s[3][0], s[3][2])));
    float mxB = fmaxf(fmaxf(s[0][1], s[0][3]), fmaxf(s[1][1], s[1][3]));
    mxB = fmaxf(mxB, fmaxf(fmaxf(s[2][1], s[2][3]), fmaxf(s[3][1], s[3][3])));
    mxA = fmaxf(mxA, __shfl_xor_sync(0xffffffffu, mxA, 4));
    mxB = fmaxf(mxB, __shfl_xor_sync(0xffffffffu, mxB, 4));
    mxA = fmaxf(mxA, __shfl_xor_sync(0xffffffffu, mxA, 8));
    mxB = fmaxf(mxB, __shfl_xor_sync(0xffffffffu, mxB, 8));
    mxA = fmaxf(mxA, __shfl_xor_sync(0xffffffffu, mxA, 16));
    mxB = fmaxf(mxB, __shfl_xor_sync(0xffffffffu, mxB, 16));
    const float nA = fmaxf(mA, mxA * cscale), nB = fmaxf(mB, mxB * cscale);
    const float msA = (nA == NEG_INF) ? 0.f : nA;
    const float msB = (nB == NEG_INF) ? 0.f : nB;
    const float alA = ex2f(mA - msA), alB = ex2f(mB - msB);
    mA = nA;
    mB = nB;
    float la = 0.f, lb = 0.f;
    uint32_t pb[4][2];
#pragma unroll
    for (int mt = 0; mt < 4; ++mt) {
        float pA0 = ex2f(fmaf(s[mt][0], cscale, -msA));
        float pB0 = ex2f(fmaf(s[mt][1], cscale, -msB));
        float pA1 = ex2f(fmaf(s[mt][2], cscale, -msA));
        float pB1 = ex2f(fmaf(s[mt][3], cscale, -msB));
        la += pA0 + pA1;
        lb += pB0 + pB1;
        pb[mt][0] = movtrans(pack_bf16(pA0, pB0));
        pb[mt][1] = movtrans(pack_bf16(pA1, pB1));
    }
    lA = lA * alA + la;
    lB = lB * alB + lb;
#pragma unroll
    for (int dm = 0; dm < KS; ++dm) {
        o[dm][0] *= alA;
        o[dm][2] *= alA;
        o[dm][1] *= alB;
        o[dm][3] *= alB;
    }
#pragma unroll
    for (int dm = 0; dm < KS; ++dm) {
        const uint32_t vb0 = vaddr + voff[dm];
#pragma unroll
        for (int mt = 0; mt < 4; ++mt) {
            if (FULL || (mt >= mt_lo && mt < mt_hi)) {
                uint32_t a[4];
                ldsm_x4_t(a[0], a[1], a[2], a[3], vb0 + mt * MTS);
                mma16816(o[dm], a, pb[mt][0], pb[mt][1]);
            }
        }
    }
}

// ---------------------------------------------------------------- B8D: B8 with a producer warp + mbarrier ring (decoupled warps)
template <int D, int CONS, int STAGES, int MINB = 1>
__global__ void __launch_bounds__((CONS + 1) * 32, MINB) attend_b8d(const bf16* __restrict__ Q, const bf16* __restrict__ K,
                                                                 const bf16* __restrict__ V, const uint2* __restrict__ colmask,
                                                                 bf16* __restrict__ O, int S, int NBLK, int BH, int ntiles,
                                                                 float cscale) {
    constexpr int KS = D / 16;
    constexpr int CH = D / 8;
    constexpr int TILE_BYTES = 64 * D * 2;
    constexpr int STAGE_BYTES = 2 * TILE_BYTES;
    constexpr int BPC = CONS / 8;
    constexpr int OST_STRIDE = (D + 8) * 2;
    constexpr int OST_BYTES = 8 * OST_STRIDE;
    static_assert(CONS % 8 == 0, "CONS");
    extern __shared__ __align__(128) uint8_t smem[];
    const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
    const uint32_t sbase = smem_u32(smem);
    const uint32_t full0 = sbase + STAGES * STAGE_BYTES;
    const uint32_t empty0 = full0 + STAGES * 8;
    uint8_t* ost_base = smem + STAGES * STAGE_BYTES + 2 * STAGES * 8 + 16;

    const int idx = blockIdx.x;
    const int tile = ntiles - 1 - idx / BH;
    const int bh = idx % BH;
    const int blk0 = tile * BPC;
    const int bmax = min(blk0 + BPC - 1, NBLK - 1);
    const int nv = bmax + 1;
    const bf16* Kh = K + (size_t)bh * S * D;
    const bf16* Vh = V + (size_t)bh * S * D;

    if (tid == 0) {
#pragma unroll
        for (int st = 0; st < STAGES; ++st) {
            mbar_init(full0 + st * 8, 32);
            mbar_init(empty0 + st * 8, CONS);
        }
    }
    __syncthreads();

    if (warp == CONS) {
        // ---------------- producer warp
        for (int i = 0; i < nv; ++i) {
            const int st = i % STAGES;
            const int b = bmax - i;
            if (i >= STAGES) mbar_wait(empty0 + st * 8, ((i / STAGES) - 1) & 1);
            const uint32_t kdst = sbase + st * STAGE_BYTES;
            const uint32_t vdst = kdst + TILE_BYTES;
            for (int q = lane; q < 64 * CH; q += 32) {
                int r = q / CH, ch = q % CH;
                int key = b * 64 + r;
                size_t off = (size_t)(key < S ? key : 0) * D + ch * 8;
                int nb = key < S ? 16 : 0;
                cp_async16(kdst + swz<D>(r, ch), Kh + off, nb);
                cp_async16(vdst + swz<D>(r, ch), Vh + off, nb);
            }
            cp_mbar_arrive_noinc(full0 + st * 8);
        }
        return;
    }

    // ---------------- consumer warps: one static 8-row group each
    const int j = lane >> 2, c = lane & 3;
    const int mw = blk0 + (warp >> 3);
    const int sub = warp & 7;
    const int oadd = sub * 8;
    const int wrow0 = mw * 64 + oadd;
    const bool warp_valid = (mw < NBLK) && (wrow0 < S);
    const int rows_valid = min(8, S - wrow0);
    const uint2* cmw = colmask + ((size_t)bh * NBLK + mw) * 128;
    const uint2 zero2 = make_uint2(0u, 0u);

    uint32_t qb[KS][2];
    {
        int row = min(wrow0 + j, S - 1);
        const bf16* qp = Q + ((size_t)bh * S + row) * D;
#pragma unroll
        for (int ks = 0; ks < KS; ++ks) {
            qb[ks][0] = __ldg(reinterpret_cast<const unsigned int*>(qp + 16 * ks + 2 * c));
            qb[ks][1] = __ldg(reinterpret_cast<const unsigned int*>(qp + 16 * ks + 8 + 2 * c));
        }
    }
    uint32_t koff[KS], voff[KS];
#pragma unroll
    for (int q = 0; q < KS; ++q) {
        int r0 = (((lane >> 3) & 1) << 3) + (lane & 7);
        koff[q] = r0 * (D * 2) + ((((2 * q) + (lane >> 4)) ^ (lane & 7)) << 4);
        int r1 = (((lane >> 4) & 1) << 3) + (lane & 7);
        voff[q] = r1 * (D * 2) + ((((2 * q) + ((lane >> 3) & 1)) ^ (lane & 7)) << 4);
    }
    float o[KS][4];
#pragma unroll
    for (int i = 0; i < KS; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f;
    float mA = NEG_INF, mB = NEG_INF, lA = 0.f, lB = 0.f;

    uint2 cm_c = (warp_valid && mw >= bmax) ? __ldg(&cmw[bmax]) : zero2;
    uint2 cm_n1 = (warp_valid && nv > 1) ? __ldg(&cmw[bmax - 1]) : zero2;
    uint2 cm_n2;
    for (int i = 0; i < nv; ++i) {
        const int b = bmax - i;
        const int st = i % STAGES;
        cm_n2 = (warp_valid && i + 2 < nv) ? __ldg(&cmw[b - 2]) : zero2;
        mbar_wait(full0 + st * 8, (i / STAGES) & 1);
        const int d = mw - b;
        if (warp_valid && d >= 0) {
            const uint32_t kaddr = sbase + st * STAGE_BYTES;
            const uint32_t vaddr = kaddr + TILE_BYTES;
            const unsigned long long m64 = ((unsigned long long)cm_c.y << 32) | cm_c.x;
            const uint32_t keep8 = (uint32_t)((m64 >> oadd) & 0xffull);
            if (d == 0) {
                const int mt_hi = min(4, ((oadd + 7) >> 4) + 1);
                run_pass8<D, 1>(kaddr, vaddr, koff, voff, qb, lane, 0u, oadd, 0, mt_hi, cscale, o, mA, mB, lA, lB);
            } else if (d == 1) {
                const int mt_lo = keep8 ? 0 : ((oadd + 1) >> 4);
                run_pass8<D, 2>(kaddr, vaddr, koff, voff, qb, lane, keep8, oadd, mt_lo, 4, cscale, o, mA, mB, lA, lB);
            } else if (keep8 != 0u) {
                run_pass8<D, 0>(kaddr, vaddr, koff, voff, qb, lane, keep8, oadd, 0, 4, cscale, o, mA, mB, lA, lB);
            }
        }
        __syncwarp();
        if (lane == 0) mbar_arrive(empty0 + st * 8);
        cm_c = cm_n1;
        cm_n1 = cm_n2;
    }
    // ---- epilogue (warp-private staging)
    lA += __shfl_xor_sync(0xffffffffu, lA, 4);
    lB += __shfl_xor_sync(0xffffffffu, lB, 4);
    lA += __shfl_xor_sync(0xffffffffu, lA, 8);
    lB += __shfl_xor_sync(0xffffffffu, lB, 8);
    lA += __shfl_xor_sync(0xffffffffu, lA, 16);
    lB += __shfl_xor_sync(0xffffffffu, lB, 16);
    const float invA = 1.f / lA, invB = 1.f / lB;
    uint8_t* ost = ost_base + warp * OST_BYTES;
#pragma unroll
    for (int dm = 0; dm < KS; ++dm) {
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c) * OST_STRIDE + (16 * dm + j) * 2) = __float2bfloat16_rn(o[dm][0] * invA);
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c + 1) * OST_STRIDE + (16 * dm + j) * 2) = __float2bfloat16_rn(o[dm][1] * invB);
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c) * OST_STRIDE + (16 * dm + j + 8) * 2) = __float2bfloat16_rn(o[dm][2] * invA);
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c + 1) * OST_STRIDE + (16 * dm + j + 8) * 2) = __float2bfloat16_rn(o[dm][3] * invB);
    }
    __syncwarp();
    if (warp_valid) {
        bf16* Oh = O + ((size_t)bh * S + wrow0) * D;
#pragma unroll
        for (int it = 0; it < 2 * (D / 64); ++it) {
            int q = lane + it * 32;
            int r = q / CH, ch = q % CH;
            if (r < rows_valid) {
                uint4 v = *reinterpret_cast<const uint4*>(ost + r * OST_STRIDE + ch * 16);
                *reinterpret_cast<uint4*>(Oh + (size_t)r * D + ch * 8) = v;
            }
        }
    }
}



// ---------------------------------------------------------------- FS: fused select + attend for short sequences (D == 64, NBLK <= 33)
// One CTA = one 64-row block. 8 consumer warps (8 rows each, transposed passes, register state) + a producer warp that streams
// K/V blocks through an mbarrier ring from the first cycle. The consumers first compute the block importances (diagonal block from
// the causal score tile, earlier blocks from the block-sum pieces), select the top-8 blocks per row in shared memory, and only then start
// consuming the ring.
template <int WSTAGES>
__global__ void __launch_bounds__(288, 1) attend_fs(const bf16* __restrict__ Q, const bf16* __restrict__ K, const bf16* __restrict__ V,
                                                    const uint32_t* __restrict__ kcf, bf16* __restrict__ O, int S, int NBLK,
                                                    int BH, int NT, float scale, float cscale, int nres) {
    constexpr int D = 64;
    constexpr int KS = D / 16;
    constexpr int CH = D / 8;
    constexpr int CONS = 8;
    constexpr int STAGES = WSTAGES;
    constexpr int TILE_BYTES = 64 * D * 2;
    constexpr int STAGE_BYTES = 2 * TILE_BYTES;
    constexpr int OST_STRIDE = (D + 8) * 2;
    constexpr int OST_BYTES = 8 * OST_STRIDE;
    constexpr int NBP = 40;
    constexpr int NV = (3 * KS * 2) / 4;
    extern __shared__ __align__(128) uint8_t smem[];
    const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
    const uint32_t sbase = smem_u32(smem);
    const uint32_t full0 = sbase + STAGES * STAGE_BYTES;
    const uint32_t empty0 = full0 + STAGES * 8;
    uint8_t* ost_base = smem + STAGES * STAGE_BYTES + 2 * STAGES * 8 + 16;
    float* imp = reinterpret_cast<float*>(ost_base + CONS * OST_BYTES);          // [64][NBP]
    uint32_t* cms = reinterpret_cast<uint32_t*>(imp + 64 * NBP);                  // [34][2] colmask words for this row block

    const int idx = blockIdx.x;
    const int m = NBLK - 1 - idx / BH;
    const int bh = idx % BH;
    const int bmax = m;
    const int nv = bmax + 1;
    const bf16* Kh = K + (size_t)bh * S * D;
    const bf16* Vh = V + (size_t)bh * S * D;
    const bf16* Qh = Q + (size_t)bh * S * D;

    if (tid == 0) {
#pragma unroll
        for (int st = 0; st < STAGES; ++st) {
            mbar_init(full0 + st * 8, 32);
            mbar_init(empty0 + st * 8, CONS);
        }
    }
    __syncthreads();

    if (warp == CONS) {
        auto warm = [&]() {   // L2 warm-up for the CTA that will take an SM slot about `nres` dispatches after this one: its Q rows, block-sum pieces and first K/V blocks
            int tgt = 0;
            if (lane == 0) tgt = atomicAdd(&g_pf_ctr, 1) + nres;
            tgt = __shfl_sync(0xffffffffu, tgt, 0);
            if (tgt < (int)gridDim.x) {
                const int m2 = NBLK - 1 - tgt / BH, bh2 = tgt % BH;
                const size_t hb = (size_t)bh2 * S;
                for (int i = lane; i < 64; i += 32)
                    if (m2 * 64 + i < S) asm volatile("prefetch.global.L2 [%0];" ::"l"(Q + (hb + m2 * 64 + i) * D));
                const int ntm2 = (m2 + 7) >> 3;
                for (int jg = 0; jg < ntm2; ++jg)
                    for (int i = lane; i < 24; i += 32) asm volatile("prefetch.global.L2 [%0];" ::"l"(kcf + (((size_t)bh2 * NT + jg) * 32) * (3 * KS * 2) + i * 32));
                for (int bi = 0; bi < STAGES && bi <= m2; ++bi) {
                    const int b2 = m2 - bi;
                    for (int i = lane; i < 64; i += 32)
                        if (b2 * 64 + i < S) {
                            asm volatile("prefetch.global.L2 [%0];" ::"l"(K + (hb + b2 * 64 + i) * D));
                            asm volatile("prefetch.global.L2 [%0];" ::"l"(V + (hb + b2 * 64 + i) * D));
                        }
                }
            }
        };
        for (int i = 0; i < nv; ++i) {
            const int st = i % STAGES;
            const int b = bmax - i;
            if (i == STAGES) warm();      // first point where the producer has to wait for the consumers anyway
            if (i >= STAGES) mbar_wait(empty0 + st * 8, ((i / STAGES) - 1) & 1);
            const uint32_t kdst = sbase + st * STAGE_BYTES;
            const uint32_t vdst = kdst + TILE_BYTES;
            for (int q = lane; q < 64 * CH; q += 32) {
                int r = q / CH, ch = q % CH;
                int key = b * 64 + r;
                size_t off = (size_t)(key < S ? key : 0) * D + ch * 8;
                int nb = key < S ? 16 : 0;
                cp_async16(kdst + swz<D>(r, ch), Kh + off, nb);
                cp_async16(vdst + swz<D>(r, ch), Vh + off, nb);
            }
            cp_mbar_arrive_noinc(full0 + st * 8);
        }
        if (nv <= STAGES) warm();
        return;
    }

    // ---------------- consumer warps
    const int j = lane >> 2, c = lane & 3;
    const int mw = m;
    const int oadd = warp * 8;
    const int wrow0 = mw * 64 + oadd;
    const bool warp_valid = wrow0 < S;
    const int rows_valid = min(8, S - wrow0);

    // Q fragments of the visit passes (B operand, 8 rows)
    uint32_t qb[KS][2];
    {
        int row = min(wrow0 + j, S - 1);
        const bf16* qp = Qh + (size_t)row * D;
#pragma unroll
        for (int ks = 0; ks < KS; ++ks) {
            qb[ks][0] = __ldg(reinterpret_cast<const unsigned int*>(qp + 16 * ks + 2 * c));
            qb[ks][1] = __ldg(reinterpret_cast<const unsigned int*>(qp + 16 * ks + 8 + 2 * c));
        }
    }
    uint32_t koff[KS], voff[KS];
#pragma unroll
    for (int q = 0; q < KS; ++q) {
        int r0 = (((lane >> 3) & 1) << 3) + (lane & 7);
        koff[q] = r0 * (D * 2) + ((((2 * q) + (lane >> 4)) ^ (lane & 7)) << 4);
        int r1 = (((lane >> 4) & 1) << 3) + (lane & 7);
        voff[q] = r1 * (D * 2) + ((((2 * q) + ((lane >> 3) & 1)) ^ (lane & 7)) << 4);
    }

    // ---------------- selection
    for (int i = tid; i < 34 * 2; i += 256) cms[i] = 0u;
    if (m < 8) {
        // every valid row selects every block below the diagonal
        const int nrows = min(64, S - m * 64);
        const unsigned long long vm = (nrows >= 64) ? ~0ull : ((1ull << nrows) - 1ull);
        for (int b = tid; b < 34; b += 256) {
            const bool on = b < m;
            cms[2 * b] = on ? (uint32_t)vm : 0u;
            cms[2 * b + 1] = on ? (uint32_t)(vm >> 32) : 0u;
        }
        asm volatile("bar.sync 1, 256;\n" ::: "memory");
    } else {
        const int c4 = c;
        const int t = warp & 3;                 // 16-row tile of this warp
        const int jp = warp >> 2;               // block-group parity
        const int ntm = (m + 7) >> 3;
        uint32_t qd[KS][4];
        {
            const int r0 = m * 64 + t * 16 + (lane >> 2), r1 = r0 + 8;
#pragma unroll
            for (int ks = 0; ks < KS; ++ks) {
                uint32_t z = 0;
                qd[ks][0] = r0 < S ? *reinterpret_cast<const uint32_t*>(Qh + (size_t)r0 * D + ks * 16 + c4 * 2) : z;
                qd[ks][1] = r1 < S ? *reinterpret_cast<const uint32_t*>(Qh + (size_t)r1 * D + ks * 16 + c4 * 2) : z;
                qd[ks][2] = r0 < S ? *reinterpret_cast<const uint32_t*>(Qh + (size_t)r0 * D + ks * 16 + 8 + c4 * 2) : z;
                qd[ks][3] = r1 < S ? *reinterpret_cast<const uint32_t*>(Qh + (size_t)r1 * D + ks * 16 + 8 + c4 * 2) : z;
            }
        }
        uint4 kvc[NV], kvn[NV];
        int jj = jp;
        if (jj < ntm) {
            const uint4* pp = reinterpret_cast<const uint4*>(kcf + (((size_t)bh * NT + jj) * 32 + lane) * (3 * KS * 2));
#pragma unroll
            for (int i = 0; i < NV; ++i) kvc[i] = pp[i];
        }
        if (jp == 0) {
            // diagonal block: causal score tile from the ring's first stage (block m)
            mbar_wait(full0, 0);
            float sc[8][4];
#pragma unroll
            for (int i = 0; i < 8; ++i) sc[i][0] = sc[i][1] = sc[i][2] = sc[i][3] = 0.f;
#pragma unroll
            for (int ks = 0; ks < KS; ++ks) {
#pragma unroll
                for (int np = 0; np < 4; ++np) {
                    uint32_t b0, b1, b2, b3;
                    int mat = lane >> 3;
                    int row = np * 16 + ((mat >> 1) << 3) + (lane & 7);
                    int ch = 2 * ks + (mat & 1);
                    ldsm_x4(b0, b1, b2, b3, sbase + swz<D>(row, ch));
                    mma16816(sc[2 * np], qd[ks], b0, b1);
                    mma16816(sc[2 * np + 1], qd[ks], b2, b3);
                }
            }
            const int o0 = t * 16 + (lane >> 2), o1 = o0 + 8;
            float sum0 = 0.f, sum1 = 0.f;
#pragma unroll
            for (int nt = 0; nt < 8; ++nt) {
#pragma unroll
                for (int e = 0; e < 2; ++e) {
                    int cidx = nt * 8 + c4 * 2 + e;
                    if (cidx <= o0) sum0 += sc[nt][e];
                    if (cidx <= o1) sum1 += sc[nt][2 + e];
                }
            }
            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);
            if (c4 == 0) {
                imp[o0 * NBP + m] = sum0 * scale / (float)(o0 + 1);
                imp[o1 * NBP + m] = sum1 * scale / (float)(o1 + 1);
            }
        }
        const float scale64 = scale * (1.0f / 64.0f);
        for (; jj < ntm; jj += 2) {
            const int jn = jj + 2;
            if (jn < ntm) {
                const uint4* pp = reinterpret_cast<const uint4*>(kcf + (((size_t)bh * NT + jn) * 32 + lane) * (3 * KS * 2));
#pragma unroll
                for (int i = 0; i < NV; ++i) kvn[i] = pp[i];
            }
            const int blk = jj * 8 + c4 * 2;
            float acc[3][4];
#pragma unroll
            for (int s3 = 0; s3 < 3; ++s3) {
                acc[s3][0] = acc[s3][1] = acc[s3][2] = acc[s3][3] = 0.f;
#pragma unroll
                for (int ks = 0; ks < KS; ++ks) {
                    int w = (s3 * KS + ks) * 2;
                    mma16816(acc[s3], qd[ks], u4get(kvc[w >> 2], w & 3), u4get(kvc[(w + 1) >> 2], (w + 1) & 3));
                }
            }
            const int ra = t * 16 + (lane >> 2), rb = ra + 8;
            float v00 = (acc[0][0] + (acc[1][0] + acc[2][0])) * scale64;
            float v01 = (acc[0][1] + (acc[1][1] + acc[2][1])) * scale64;
            float v10 = (acc[0][2] + (acc[1][2] + acc[2][2])) * scale64;
            float v11 = (acc[0][3] + (acc[1][3] + acc[2][3])) * scale64;
            if (blk < m) imp[ra * NBP + blk] = v00;
            if (blk + 1 < m) imp[ra * NBP + blk + 1] = v01;
            if (blk < m) imp[rb * NBP + blk] = v10;
            if (blk + 1 < m) imp[rb * NBP + blk + 1] = v11;
#pragma unroll
            for (int i = 0; i < NV; ++i) kvc[i] = kvn[i];
        }
        asm volatile("bar.sync 1, 256;\n" ::: "memory");
        // top-8 per row: 4 consumer threads per row (lanes 4r..4r+3 share a row); value-only insertion, quad merge, emission
        {
            const int row = tid >> 2, part = tid & 3;
            const float* ir = imp + row * NBP;
            const int npairs = (m >> 3) + 1;
            float lv[8];
#pragma unroll
            for (int k = 0; k < 8; ++k) lv[k] = NEG_INF;
            for (int pp = part; pp < npairs; pp += 4) {
                float xs[8];
#pragma unroll
                for (int e = 0; e < 8; ++e) xs[e] = (8 * pp + e <= m) ? ir[8 * pp + e] : NEG_INF;
#pragma unroll
                for (int e = 0; e < 8; ++e) {
                    float prev = INFINITY, cur = xs[e];
                    float nl[8];
#pragma unroll
                    for (int k = 0; k < 8; ++k) {
                        nl[k] = fmaxf(lv[k], fminf(prev, cur));
                        prev = lv[k];
                    }
#pragma unroll
                    for (int k = 0; k < 8; ++k) lv[k] = nl[k];
                }
            }
            // quad merge: two rounds of (top-set of the pair, bitonic sort)
#pragma unroll
            for (int rd = 0; rd < 2; ++rd) {
                float cc[8];
#pragma unroll
                for (int k = 0; k < 8; ++k) {
                    const float o = __shfl_xor_sync(0xffffffffu, lv[7 - k], 1 << rd);
                    cc[k] = fmaxf(lv[k], o);
                }
#pragma unroll
                for (int i = 0; i < 4; ++i) { float a = cc[i], b2 = cc[i + 4]; cc[i] = fmaxf(a, b2); cc[i + 4] = fminf(a, b2); }
#pragma unroll
                for (int i = 0; i < 8; i += 4) {
#pragma unroll
                    for (int u = 0; u < 2; ++u) { float a = cc[i + u], b2 = cc[i + u + 2]; cc[i + u] = fmaxf(a, b2); cc[i + u + 2] = fminf(a, b2); }
                }
#pragma unroll
                for (int i = 0; i < 8; i += 2) { float a = cc[i], b2 = cc[i + 1]; cc[i] = fmaxf(a, b2); cc[i + 1] = fminf(a, b2); }
#pragma unroll
                for (int k = 0; k < 8; ++k) lv[k] = cc[k];
            }
            const float T8 = lv[7];
            uint32_t bm[2] = {0u, 0u};
            for (int pp = part; pp < npairs; pp += 4) {
                uint32_t byte = 0u;
#pragma unroll
                for (int e = 0; e < 8; ++e) {
                    const bool ge = (8 * pp + e <= m) && (ir[8 * pp + e] >= T8);
                    byte |= ge ? (1u << e) : 0u;
                }
                if ((pp >> 2) == 0) bm[0] |= byte << (8 * (pp & 3)); else bm[1] |= byte << (8 * (pp & 3));
            }
#pragma unroll
            for (int w = 0; w < 2; ++w) {
                bm[w] |= __shfl_xor_sync(0xffffffffu, bm[w], 1);
                bm[w] |= __shfl_xor_sync(0xffffffffu, bm[w], 2);
            }
            const int trow = m * 64 + row;
            const bool row_ok = trow < S;
            int total = __popc(bm[0]) + __popc(bm[1]);
            if (row_ok && total > 8) {
                int cgt = 0;
                for (int b2 = 0; b2 <= m; ++b2) cgt += (ir[b2] > T8) ? 1 : 0;
                int need = 8 - cgt;
                bm[0] = bm[1] = 0u;
                for (int b2 = m; b2 >= 0; --b2) {
                    const float v2 = ir[b2];
                    const bool take = (v2 > T8) || (v2 == T8 && need > 0);
                    if (v2 == T8 && need > 0) --need;
                    if (take) { if (b2 >> 5) bm[1] |= 1u << (b2 & 31); else bm[0] |= 1u << (b2 & 31); }
                }
            }
            if (!row_ok) { bm[0] = bm[1] = 0u; }
#pragma unroll
            for (int w = 0; w < 2; ++w) {
                uint32_t x = bm[w];
                while (x) {
                    const int bpos = __ffs(x) - 1;
                    x &= x - 1;
                    if (((w * 32 + bpos) & 3) == part) atomicOr(&cms[(w * 32 + bpos) * 2 + (row >> 5)], 1u << (row & 31));
                }
            }
        }
        asm volatile("bar.sync 1, 256;\n" ::: "memory");
    }

    // ---------------- attend visits (identical to attend_b8d)
    float o[KS][4];
#pragma unroll
    for (int i = 0; i < KS; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f;
    float mA = NEG_INF, mB = NEG_INF, lA = 0.f, lB = 0.f;
    for (int i = 0; i < nv; ++i) {
        const int b = bmax - i;
        const int st = i % STAGES;
        const uint2 cm_c = make_uint2(cms[2 * b], cms[2 * b + 1]);
        mbar_wait(full0 + st * 8, (i / STAGES) & 1);
        const int d = mw - b;
        if (warp_valid && d >= 0) {
            const uint32_t kaddr = sbase + st * STAGE_BYTES;
            const uint32_t vaddr = kaddr + TILE_BYTES;
            const unsigned long long m64 = ((unsigned long long)cm_c.y << 32) | cm_c.x;
            const uint32_t keep8 = (uint32_t)((m64 >> oadd) & 0xffull);
            if (d == 0) {
                const int mt_hi = min(4, ((oadd + 7) >> 4) + 1);
                run_pass8<D, 1>(kaddr, vaddr, koff, voff, qb, lane, 0u, oadd, 0, mt_hi, cscale, o, mA, mB, lA, lB);
            } else if (d == 1) {
                const int mt_lo = keep8 ? 0 : ((oadd + 1) >> 4);
                run_pass8<D, 2>(kaddr, vaddr, koff, voff, qb, lane, keep8, oadd, mt_lo, 4, cscale, o, mA, mB, lA, lB);
            } else if (keep8 != 0u) {
                run_pass8<D, 0>(kaddr, vaddr, koff, voff, qb, lane, keep8, oadd, 0, 4, cscale, o, mA, mB, lA, lB);
            }
        }
        __syncwarp();
        if (lane == 0) mbar_arrive(empty0 + st * 8);
    }
    // ---- epilogue (warp-private staging)
    lA += __shfl_xor_sync(0xffffffffu, lA, 4);
    lB += __shfl_xor_sync(0xffffffffu, lB, 4);
    lA += __shfl_xor_sync(0xffffffffu, lA, 8);
    lB += __shfl_xor_sync(0xffffffffu, lB, 8);
    lA += __shfl_xor_sync(0xffffffffu, lA, 16);
    lB += __shfl_xor_sync(0xffffffffu, lB, 16);
    const float invA = 1.f / lA, invB = 1.f / lB;
    uint8_t* ost = ost_base + warp * OST_BYTES;
#pragma unroll
    for (int dm = 0; dm < KS; ++dm) {
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c) * OST_STRIDE + (16 * dm + j) * 2) = __float2bfloat16_rn(o[dm][0] * invA);
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c + 1) * OST_STRIDE + (16 * dm + j) * 2) = __float2bfloat16_rn(o[dm][1] * invB);
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c) * OST_STRIDE + (16 * dm + j + 8) * 2) = __float2bfloat16_rn(o[dm][2] * invA);
        *reinterpret_cast<__nv_bfloat16*>(ost + (2 * c + 1) * OST_STRIDE + (16 * dm + j + 8) * 2) = __float2bfloat16_rn(o[dm][3] * invB);
    }
    __syncwarp();
    if (warp_valid) {
        bf16* Oh = O + ((size_t)bh * S + wrow0) * D;
#pragma unroll
        for (int it = 0; it < 2 * (D / 64); ++it) {
            int q = lane + it * 32;
            int r = q / CH, ch = q % CH;
            if (r < rows_valid) {
                uint4 v = *reinterpret_cast<const uint4*>(ost + r * OST_STRIDE + ch * 16);
                *reinterpret_cast<uint4*>(Oh + (size_t)r * D + ch * 8) = v;
            }
        }
    }
}

// ---------------------------------------------------------------- KS: key-stationary attention -> per-(row,block) partials, then merge
constexpr int KS_SLOTS = 10;
#ifndef KS_RBPAD
#define KS_RBPAD 32
#endif

// slot of block b in the row's contributing set S_t = sel_t | {m} | {m-1}; also returns the set size
__device__ __forceinline__ int ks_slot(const uint32_t* __restrict__ selp, int t, int b, int& nset) {
    uint4 w = *reinterpret_cast<const uint4*>(selp);
    const int m = t >> 6;
    uint32_t w0 = w.x, w1 = w.y, w2 = w.z, w3 = w.w;
    {
        const int wm = m >> 5;
        const uint32_t bmk = 1u << (m & 31);
        w0 |= (wm == 0) ? bmk : 0u; w1 |= (wm == 1) ? bmk : 0u; w2 |= (wm == 2) ? bmk : 0u; w3 |= (wm == 3) ? bmk : 0u;
    }
    if (m > 0) {
        const int m1 = m - 1;
        const int wm = m1 >> 5;
        const uint32_t bmk = 1u << (m1 & 31);
        w0 |= (wm == 0) ? bmk : 0u; w1 |= (wm == 1) ? bmk : 0u; w2 |= (wm == 2) ? bmk : 0u; w3 |= (wm == 3) ? bmk : 0u;
    }
    nset = __popc(w0) + __popc(w1) + __popc(w2) + __popc(w3);
    const int bw = b >> 5;
    const uint32_t lowmask = (1u << (b & 31)) - 1u;
    int r = (bw > 0 ? __popc(w0) : 0) + (bw > 1 ? __popc(w1) : 0) + (bw > 2 ? __popc(w2) : 0);
    const uint32_t wb = bw == 0 ? w0 : (bw == 1 ? w1 : (bw == 2 ? w2 : w3));
    r += __popc(wb & lowmask);
    return r;
}

// natural orientation: S = Q K^T with 16 rows per tile (M = rows), P feeds the PV A operand straight from the accumulators
template <int D, int MODE>
__device__ __forceinline__ void ksn_pass(uint32_t kaddr, uint32_t vaddr, const uint32_t (&qa)[D / 16][4], int lane, uint32_t keepw,
                                         int oadd, float cscale, float (&o)[D / 8][4], float (&mo)[2], float (&lo)[2]) {
    constexpr int KS = D / 16;
    constexpr int ND = D / 8;
    const int g = lane >> 2, c = lane & 3;
    float s[8][4];
#pragma unroll
    for (int nt = 0; nt < 8; ++nt) s[nt][0] = s[nt][1] = s[nt][2] = s[nt][3] = 0.f;
    const int nk = (MODE == 1) ? ((oadd >> 4) + 1) : 4;     // mode 1 (causal diagonal tile): key groups of 16 above the tile's last row are fully masked
#pragma unroll
    for (int ks = 0; ks < KS; ++ks) {
#pragma unroll
        for (int np = 0; np < 4; ++np) {
            if (MODE != 1 || np < nk) {
                uint32_t b0, b1, b2, b3;
                const int mat = lane >> 3;
                const int row = np * 16 + ((mat >> 1) << 3) + (lane & 7);
                const int ch = 2 * ks + (mat & 1);
                ldsm_x4(b0, b1, b2, b3, kaddr + swz<D>(row, ch));
                mma16816(s[2 * np], qa[ks], b0, b1);
                mma16816(s[2 * np + 1], qa[ks], b2, b3);
            }
        }
    }
    if (MODE == 1 || MODE == 2) {
        const int rA = oadd + g, rB = rA + 8;
        const bool kA = ((keepw >> g) & 1u) != 0u;
        const bool kB = ((keepw >> (g + 8)) & 1u) != 0u;
#pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
#pragma unroll
            for (int e = 0; e < 4; ++e) {
                const int key = 8 * nt + 2 * c + (e & 1);
                const int oo = (e >> 1) ? rB : rA;
                const bool keep = (e >> 1) ? kB : kA;
                const bool allowed = (MODE == 1) ? (key <= oo) : (keep || key >= oo + 1);
                if (!allowed) s[nt][e] = NEG_INF;
            }
        }
    }
    float mxA = fmaxf(fmaxf(s[0][0], s[0][1]), fmaxf(s[1][0], s[1][1]));
    float mxB = fmaxf(fmaxf(s[0][2], s[0][3]), fmaxf(s[1][2], s[1][3]));
#pragma unroll
    for (int nt = 2; nt < 8; ++nt) {
        mxA = fmaxf(mxA, fmaxf(s[nt][0], s[nt][1]));
        mxB = fmaxf(mxB, fmaxf(s[nt][2], s[nt][3]));
    }
    mxA = fmaxf(mxA, __shfl_xor_sync(0xffffffffu, mxA, 1));
    mxB = fmaxf(mxB, __shfl_xor_sync(0xffffffffu, mxB, 1));
    mxA = fmaxf(mxA, __shfl_xor_sync(0xffffffffu, mxA, 2));
    mxB = fmaxf(mxB, __shfl_xor_sync(0xffffffffu, mxB, 2));
    const float mA = mxA * cscale, mB = mxB * cscale;
    const float msA = (mA == NEG_INF) ? 0.f : mA;
    const float msB = (mB == NEG_INF) ? 0.f : mB;
    float la = 0.f, lb = 0.f;
    uint32_t pa[4][4];
#pragma unroll
    for (int kk = 0; kk < 4; ++kk) {
        float p00 = ex2f(fmaf(s[2 * kk][0], cscale, -msA));
        float p01 = ex2f(fmaf(s[2 * kk][1], cscale, -msA));
        float p02 = ex2f(fmaf(s[2 * kk][2], cscale, -msB));
        float p03 = ex2f(fmaf(s[2 * kk][3], cscale, -msB));
        float p10 = ex2f(fmaf(s[2 * kk + 1][0], cscale, -msA));
        float p11 = ex2f(fmaf(s[2 * kk + 1][1], cscale, -msA));
        float p12 = ex2f(fmaf(s[2 * kk + 1][2], cscale, -msB));
        float p13 = ex2f(fmaf(s[2 * kk + 1][3], cscale, -msB));
        la += (p00 + p01) + (p10 + p11);
        lb += (p02 + p03) + (p12 + p13);
        pa[kk][0] = pack_bf16(p00, p01);
        pa[kk][1] = pack_bf16(p02, p03);
        pa[kk][2] = pack_bf16(p10, p11);
        pa[kk][3] = pack_bf16(p12, p13);
    }
#pragma unroll
    for (int n = 0; n < ND; ++n) o[n][0] = o[n][1] = o[n][2] = o[n][3] = 0.f;
#pragma unroll
    for (int kk = 0; kk < 4; ++kk) {
        if (MODE != 1 || kk < nk) {
#pragma unroll
            for (int dp = 0; dp < ND / 2; ++dp) {
                uint32_t b0, b1, b2, b3;
                const int mi = lane >> 3;
                const int key = 16 * kk + (mi & 1) * 8 + (lane & 7);
                const int ch = 2 * dp + (mi >> 1);
                ldsm_x4_t(b0, b1, b2, b3, vaddr + swz<D>(key, ch));
                mma16816(o[2 * dp], pa[kk], b0, b1);
                mma16816(o[2 * dp + 1], pa[kk], b2, b3);
            }
        }
    }
    la += __shfl_xor_sync(0xffffffffu, la, 1);
    lb += __shfl_xor_sync(0xffffffffu, lb, 1);
    la += __shfl_xor_sync(0xffffffffu, la, 2);
    lb += __shfl_xor_sync(0xffffffffu, lb, 2);
    mo[0] = mA; mo[1] = mB; lo[0] = la; lo[1] = lb;   // o is returned unnormalized (sum of p * v); the caller folds 1 / l into the int8 scale
}

template <int D, int WARPS, int MINB>
__global__ void __launch_bounds__(WARPS * 32, MINB) ksn_attend(const bf16* __restrict__ Q, const bf16* __restrict__ K,
                                                               const bf16* __restrict__ V, const uint2* __restrict__ colmask,
                                                               const uint32_t* __restrict__ sel, uint8_t* __restrict__ partO,
                                                               int S, int NBLK, int BH, int Spad,
                                                               float cscale) {
    constexpr int KS = D / 16;
    constexpr int NTH = WARPS * 32;
    constexpr int CHM = (D == 64) ? 64 : 32;      // row blocks per chunk
    constexpr int LCAP = CHM * 64;
    constexpr int RB = D + KS_RBPAD;                  // partial record bytes (int8 data + 16 B trailer + alignment pad)
    constexpr int TILEB = 64 * D * 2;
    constexpr int NCH = D / 8;                    // 16 B chunks per row
    constexpr bool QSM = (D == 64);               // Q rows of a tile are staged in a per-warp shared buffer (cp.async) instead of register prefetch
    constexpr int QWB = 16 * D * 2;               // bytes of the per-warp Q buffer
    extern __shared__ __align__(128) uint8_t smem[];
    uint8_t* Ks = smem;
    uint8_t* Vs = smem + TILEB;
    uint16_t* list = reinterpret_cast<uint16_t*>(smem + 2 * TILEB);
    int* ctr = reinterpret_cast<int*>(smem + 2 * TILEB + LCAP * 2);
    uint8_t* qbuf = smem + 2 * TILEB + LCAP * 2 + 128;
    const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
    const uint32_t qb_addr = smem_u32(qbuf) + warp * QWB;
    const int g = lane >> 2, c = lane & 3;
    const int idx = blockIdx.x;
    const int b = idx / BH, bh = idx % BH;
    const bf16* Qh = Q + (size_t)bh * S * D;
    const bf16* Kh = K + (size_t)bh * S * D;
    const bf16* Vh = V + (size_t)bh * S * D;
    const uint32_t* selh = sel + (size_t)bh * Spad * 4;
    uint8_t* partOh = partO + (size_t)bh * Spad * KS_SLOTS * RB;

    for (int i = tid; i < 64 * NCH; i += NTH) {
        int r = i / NCH, ch = i % NCH;
        int key = b * 64 + r;
        size_t off = (size_t)(key < S ? key : 0) * D + ch * 8;
        int nb = key < S ? 16 : 0;
        cp_async16(smem_u32(Ks) + swz<D>(r, ch), Kh + off, nb);
    }
    {   // V in 4-byte words, transposed inside each row (16 B chunk n, word slot cc <- source word NPT*cc + n): lane c of the PV
        // accumulators then owns the D/4 contiguous output columns [D/4*c, D/4*(c+1)) and can store them straight to global memory
        constexpr int NW = D / 2, NPT = D / 8;
        for (int i = tid; i < 64 * NW; i += NTH) {
            const int r = i / NW, w = i % NW;
            const int key = b * 64 + r;
            const bf16* src = Vh + (size_t)(key < S ? key : 0) * D + 2 * w;
            cp_async4(smem_u32(Vs) + swz<D>(r, w % NPT) + 4 * (w / NPT), src, key < S ? 4 : 0);
        }
    }
    cp_commit();
    if (tid == 0) ctr[0] = 0;
    unsigned long long cm2 = 0ull;
    if (b + 1 < NBLK) {
        uint2 cmv = __ldg(&colmask[((size_t)bh * NBLK + (b + 1)) * 128 + b]);
        cm2 = ((unsigned long long)cmv.y << 32) | cmv.x;
    }
    cp_wait<0>();
    __syncthreads();
    const uint32_t kbase = smem_u32(Ks), vbase = smem_u32(Vs);

    const int nchunks = max(1, (NBLK - (b + 2) + CHM - 1) / CHM);
    for (int chunk = 0; chunk < nchunks; ++chunk) {
        const int mlo = b + 2 + chunk * CHM;
        const int mhi = min(NBLK, mlo + CHM);
        int total = 0;
        __syncthreads();
        {
            const int e = tid;                       // (m, half) entries: 2*CHM = 64
            const int m = mlo + (e >> 1), hf = e & 1;
            uint32_t bits = 0;
            if (e < 2 * CHM && m < mhi) {
                uint2 cmv = __ldg(&colmask[((size_t)bh * NBLK + m) * 128 + b]);
                bits = hf ? cmv.y : cmv.x;
            }
            int cnt = __popc(bits);
            int incl = cnt;
#pragma unroll
            for (int off = 1; off < 32; off <<= 1) {
                int n = __shfl_up_sync(0xffffffffu, incl, off);
                if (lane >= off) incl += n;
            }
            if (lane == 31 && warp < 2 * CHM / 32) ctr[1 + warp] = incl;
            __syncthreads();
            int wbase = 0;
#pragma unroll
            for (int w = 0; w < 2 * CHM / 32; ++w) {
                int wt = ctr[1 + w];
                if (w < warp) wbase += wt;
                total += wt;
            }
            int pos = wbase + incl - cnt;
            while (bits) {
                int p = __ffs(bits) - 1;
                bits &= bits - 1;
                list[pos++] = (uint16_t)(m * 64 + hf * 32 + p);
            }
            if (tid == 0) ctr[0] = 0;
        }
        __syncthreads();
        const int nlist_tiles = (total + 15) >> 4;
        const int nspecial = (chunk == 0) ? (4 + ((b + 1 < NBLK) ? 4 : 0)) : 0;
        const int ntile = nspecial + nlist_tiles;
        auto tile_rows = [&](int tl, int& mode, int& oadd, int& first_row) {
            if (tl < nspecial) {
                if (tl < 4) { mode = 1; oadd = 16 * tl; first_row = b * 64 + 16 * tl; }
                else { mode = 2; oadd = 16 * (tl - 4); first_row = (b + 1) * 64 + 16 * (tl - 4); }
            } else { mode = 0; oadd = 0; first_row = 0; }
        };
        auto row_of = [&](int tl, int mode, int first_row, int n, bool& valid) -> int {
            if (mode == 0) {
                int e = (tl - nspecial) * 16 + n;
                valid = e < total;
                return valid ? (int)list[e] : (int)list[total - 1];
            } else {
                int t = first_row + n;
                valid = (t < S);
                return valid ? t : (S - 1);
            }
        };
        auto fetch = [&]() -> int {
            int t = 0;
            if (lane == 0) t = atomicAdd(&ctr[0], 1);
            return __shfl_sync(0xffffffffu, t, 0);
        };
        struct PreR { uint32_t qa[KS][4]; uint4 selw; int tS; int validS; };
        struct PreS { uint4 selw; int tS; int validS; };
        using Pre = std::conditional_t<QSM, PreS, PreR>;
        auto prefetch = [&](int tl, Pre& p) {
            if (tl >= ntile) return;
            int mode, oadd, fr;
            tile_rows(tl, mode, oadd, fr);
            bool vs;
            const int ts = row_of(tl, mode, fr, lane >> 1, vs);
            p.tS = ts;
            p.validS = vs ? 1 : 0;
            p.selw = vs ? __ldg(reinterpret_cast<const uint4*>(selh + (size_t)ts * 4)) : make_uint4(0u, 0u, 0u, 0u);
            if constexpr (QSM) {
                // Q rows of the tile: 16 rows x 8 chunks of 16 B, four rows per instruction, into the warp's buffer
#pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const int r = 4 * i + (lane >> 3), ch = lane & 7;
                    bool vq;
                    const int tq = row_of(tl, mode, fr, r, vq);
                    cp_async16(qb_addr + swz<D>(r, ch), Qh + (size_t)tq * D + ch * 8, 16);
                }
                cp_commit();
            } else {
                bool v0, v1;
                const int t0 = row_of(tl, mode, fr, g, v0);
                const int t1 = row_of(tl, mode, fr, g + 8, v1);
                const bf16* q0 = Qh + (size_t)t0 * D;
                const bf16* q1 = Qh + (size_t)t1 * D;
#pragma unroll
                for (int ks = 0; ks < KS; ++ks) {
                    p.qa[ks][0] = __ldg(reinterpret_cast<const unsigned int*>(q0 + 16 * ks + 2 * c));
                    p.qa[ks][1] = __ldg(reinterpret_cast<const unsigned int*>(q1 + 16 * ks + 2 * c));
                    p.qa[ks][2] = __ldg(reinterpret_cast<const unsigned int*>(q0 + 16 * ks + 8 + 2 * c));
                    p.qa[ks][3] = __ldg(reinterpret_cast<const unsigned int*>(q1 + 16 * ks + 8 + 2 * c));
                }
            }
        };
        int cur = fetch();
        Pre pcur, pnxt;
        prefetch(cur, pcur);
        if constexpr (QSM) { cp_wait<0>(); __syncwarp(); }
        while (cur < ntile) {
            uint32_t qaS[QSM ? KS : 1][4];
            if constexpr (QSM) {
                // A fragments of the current tile; the buffer is refilled with the next tile's rows only after every lane holds them
#pragma unroll
                for (int ks = 0; ks < KS; ++ks)
                    ldsm_x4(qaS[ks][0], qaS[ks][1], qaS[ks][2], qaS[ks][3], qb_addr + swz<D>((lane & 7) + ((lane >> 3) & 1) * 8, 2 * ks + (lane >> 4)));
#pragma unroll
                for (int ks = 0; ks < KS; ++ks)
#pragma unroll
                    for (int j = 0; j < 4; ++j) asm volatile("" : "+r"(qaS[ks][j]));
                __syncwarp();
            }
            const int nxt = fetch();
            prefetch(nxt, pnxt);
            // record offsets of rows g and g + 8 depend only on the prefetched selection words: computed before the MMA passes
            uint32_t offA, offB;
            int vAB;
            {
                uint32_t roff = 0;
                int rv = 0;
                if (pcur.validS) {
                    const uint32_t w0 = pcur.selw.x, w1 = pcur.selw.y, w2 = pcur.selw.z, w3 = pcur.selw.w;
                    const uint32_t lowmask = (1u << (b & 31)) - 1u;
                    int slot;
                    switch (b >> 5) {                      // b is CTA-uniform
                        case 0: slot = __popc(w0 & lowmask); break;
                        case 1: slot = __popc(w0) + __popc(w1 & lowmask); break;
                        case 2: slot = __popc(w0) + __popc(w1) + __popc(w2 & lowmask); break;
                        default: slot = __popc(w0) + __popc(w1) + __popc(w2) + __popc(w3 & lowmask); break;
                    }
                    roff = ((uint32_t)pcur.tS * KS_SLOTS + (uint32_t)slot) * RB;
                    rv = 1;
                }
                offA = __shfl_sync(0xffffffffu, roff, 2 * g);
                offB = __shfl_sync(0xffffffffu, roff, 2 * g + 16);
                vAB = __shfl_sync(0xffffffffu, rv, 2 * g) | (__shfl_sync(0xffffffffu, rv, 2 * g + 16) << 1);
                asm volatile("" : "+r"(offA), "+r"(offB), "+r"(vAB));   // keep this computation ahead of the MMA passes
            }
            int mode, oadd, fr;
            tile_rows(cur, mode, oadd, fr);
            float o[D / 8][4], mo[2], lo[2];
            if constexpr (QSM) {
                if (mode == 1) {
                    ksn_pass<D, 1>(kbase, vbase, qaS, lane, 0u, oadd, cscale, o, mo, lo);
                } else if (mode == 2) {
                    const uint32_t keepw = (uint32_t)((cm2 >> oadd) & 0xffffull);
                    ksn_pass<D, 2>(kbase, vbase, qaS, lane, keepw, oadd, cscale, o, mo, lo);
                } else {
                    ksn_pass<D, 0>(kbase, vbase, qaS, lane, 0u, 0, cscale, o, mo, lo);
                }
            } else {
                if (mode == 1) {
                    ksn_pass<D, 1>(kbase, vbase, pcur.qa, lane, 0u, oadd, cscale, o, mo, lo);
                } else if (mode == 2) {
                    const uint32_t keepw = (uint32_t)((cm2 >> oadd) & 0xffffull);
                    ksn_pass<D, 2>(kbase, vbase, pcur.qa, lane, keepw, oadd, cscale, o, mo, lo);
                } else {
                    ksn_pass<D, 0>(kbase, vbase, pcur.qa, lane, 0u, 0, cscale, o, mo, lo);
                }
            }
            // ---- quantize the normalized partial to int8 with a per-row scale; every lane stores its own D/4 contiguous columns
            float pA[4] = {0.f, 0.f, 0.f, 0.f}, pB[4] = {0.f, 0.f, 0.f, 0.f};
#pragma unroll
            for (int n = 0; n < D / 8; ++n) {
                pA[n & 3] = fmaxf(pA[n & 3], fmaxf(fabsf(o[n][0]), fabsf(o[n][1])));
                pB[n & 3] = fmaxf(pB[n & 3], fmaxf(fabsf(o[n][2]), fabsf(o[n][3])));
            }
            float amA = fmaxf(fmaxf(pA[0], pA[1]), fmaxf(pA[2], pA[3]));
            float amB = fmaxf(fmaxf(pB[0], pB[1]), fmaxf(pB[2], pB[3]));
            amA = fmaxf(amA, __shfl_xor_sync(0xffffffffu, amA, 1));
            amB = fmaxf(amB, __shfl_xor_sync(0xffffffffu, amB, 1));
            amA = fmaxf(amA, __shfl_xor_sync(0xffffffffu, amA, 2));
            amB = fmaxf(amB, __shfl_xor_sync(0xffffffffu, amB, 2));
            const float qsA = amA > 0.f ? 127.f / amA : 0.f, qsB = amB > 0.f ? 127.f / amB : 0.f;
            const float invA = lo[0] > 0.f ? 1.f / lo[0] : 0.f, invB = lo[1] > 0.f ? 1.f / lo[1] : 0.f;   // amA / amB are maxima of the unnormalized o
            constexpr int NWD = D / 16;               // 32-bit words of int8 data per lane and row (4 or 8)
            uint32_t wA[NWD], wB[NWD];
#pragma unroll
            for (int w = 0; w < NWD; ++w) {
                // round-to-nearest int8 via the 1.5 * 2^23 offset: the low byte of the sum holds the two's complement result
                const uint32_t a0 = __float_as_uint(fmaf(o[2 * w][0], qsA, 12582912.f)), a1 = __float_as_uint(fmaf(o[2 * w][1], qsA, 12582912.f));
                const uint32_t a2 = __float_as_uint(fmaf(o[2 * w + 1][0], qsA, 12582912.f)), a3 = __float_as_uint(fmaf(o[2 * w + 1][1], qsA, 12582912.f));
                const uint32_t b0 = __float_as_uint(fmaf(o[2 * w][2], qsB, 12582912.f)), b1 = __float_as_uint(fmaf(o[2 * w][3], qsB, 12582912.f));
                const uint32_t b2 = __float_as_uint(fmaf(o[2 * w + 1][2], qsB, 12582912.f)), b3 = __float_as_uint(fmaf(o[2 * w + 1][3], qsB, 12582912.f));
                wA[w] = __byte_perm(__byte_perm(a0, a1, 0x0040), __byte_perm(a2, a3, 0x0040), 0x5410);
                wB[w] = __byte_perm(__byte_perm(b0, b1, 0x0040), __byte_perm(b2, b3, 0x0040), 0x5410);
            }
            const int vA = vAB & 1, vB = vAB >> 1;
            constexpr int BPL = D / 4;                // int8 bytes per lane and row (16 or 32)
            if (vA) {
                uint4* p = reinterpret_cast<uint4*>(partOh + offA + BPL * c);
#pragma unroll
                for (int q = 0; q < NWD / 4; ++q) p[q] = make_uint4(wA[4 * q], wA[4 * q + 1], wA[4 * q + 2], wA[4 * q + 3]);
            }
            if (vB) {
                uint4* p = reinterpret_cast<uint4*>(partOh + offB + BPL * c);
#pragma unroll
                for (int q = 0; q < NWD / 4; ++q) p[q] = make_uint4(wB[4 * q], wB[4 * q + 1], wB[4 * q + 2], wB[4 * q + 3]);
            }
            {   // trailer (m, l, scale) of row g (lanes c = 0, 1) and row g + 8 (lanes c = 2, 3), stored twice so every 32 B sector is complete
                const bool second = c >= 2;
                if (second ? vB : vA) {
                    const float4 tv = second ? make_float4(mo[1], lo[1], amB * invB * (1.f / 127.f), 0.f)
                                             : make_float4(mo[0], lo[0], amA * invA * (1.f / 127.f), 0.f);
                    *reinterpret_cast<float4*>(partOh + (second ? offB : offA) + D + (c & 1) * 16) = tv;
                }
            }
            cur = nxt;
            pcur = pnxt;
            if constexpr (QSM) { cp_wait<0>(); __syncwarp(); }
        }
    }
}

// merge: one warp per row (D-generic: each lane owns D/32 consecutive dims); partial records are int8 + (m, l, scale)
template <int D>
__global__ void __launch_bounds__(256) ksn_merge(const uint8_t* __restrict__ partO, const uint32_t* __restrict__ sel,
                                                 bf16* __restrict__ O, int S, int Spad, int BH) {
    constexpr int DL = D / 32;             // dims per lane (2 or 4)
    constexpr int RB = D + KS_RBPAD;
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const long long nrows = (long long)BH * S;
    for (long long R = (long long)blockIdx.x * 8 + warp; R < nrows; R += (long long)gridDim.x * 8) {
        const int bh = (int)(R / S), t = (int)(R % S);
        const size_t rs = (size_t)bh * Spad + t;
        int nset;
        ks_slot(sel + rs * 4, t, 0, nset);
        const uint8_t* pr = partO + rs * KS_SLOTS * RB;
        float4 ml = make_float4(NEG_INF, 0.f, 0.f, 0.f);
        if (lane < nset) ml = *reinterpret_cast<const float4*>(pr + lane * RB + D);
        uint32_t raw[KS_SLOTS];
#pragma unroll
        for (int s2 = 0; s2 < KS_SLOTS; ++s2) {
            raw[s2] = 0u;
            if (s2 < nset) {
                if (DL == 2) raw[s2] = *reinterpret_cast<const uint16_t*>(pr + s2 * RB + lane * 2);
                else raw[s2] = *reinterpret_cast<const uint32_t*>(pr + s2 * RB + lane * 4);
            }
        }
        float M = ml.x;
#pragma unroll
        for (int off = 16; off >= 1; off >>= 1) M = fmaxf(M, __shfl_xor_sync(0xffffffffu, M, off));
        const float Ms = (M == NEG_INF) ? 0.f : M;
        const float wv = (lane < nset) ? ex2f(ml.x - Ms) * ml.y : 0.f;
        float wsum = wv;
#pragma unroll
        for (int off = 16; off >= 1; off >>= 1) wsum += __shfl_xor_sync(0xffffffffu, wsum, off);
        const float wq = wv * ml.z;          // weight folded with the dequantization scale
        float acc[DL];
#pragma unroll
        for (int i = 0; i < DL; ++i) acc[i] = 0.f;
#pragma unroll
        for (int s2 = 0; s2 < KS_SLOTS; ++s2) {
            const float ws2 = __shfl_sync(0xffffffffu, wq, s2);
#pragma unroll
            for (int i = 0; i < DL; ++i) {
                const int qv = (int)(int8_t)((raw[s2] >> (8 * i)) & 0xffu);
                acc[i] = fmaf(ws2, (float)qv, acc[i]);
            }
        }
        const float inv = 1.f / wsum;
#pragma unroll
        for (int i = 0; i < DL / 2; ++i)
            *reinterpret_cast<uint32_t*>(O + (size_t)R * D + lane * DL + 2 * i) = pack_bf16(acc[2 * i] * inv, acc[2 * i + 1] * inv);
    }
}

// merge: LPR lanes per row (32/LPR rows per warp); each lane owns D/LPR dims, loaded with 4/8/16-byte accesses
template <int D, int LPR>
__global__ void __launch_bounds__(256) ksn_merge2(const uint8_t* __restrict__ partO, const uint32_t* __restrict__ sel,
                                                  bf16* __restrict__ O, int S, int Spad, int BH) {
    constexpr int DL = D / LPR;             // dims per lane
    constexpr int RPW = 32 / LPR;           // rows per warp
    constexpr int NW = DL / 4;              // 4-byte words per lane per slot
    constexpr int RB = D + KS_RBPAD;
    constexpr int NSL = (KS_SLOTS + LPR - 1) / LPR;   // slots whose trailer a lane loads (1 for LPR >= 16, 2 for LPR = 8)
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const int grp = lane / LPR, l = lane % LPR;
    const long long nrows = (long long)BH * S;
    for (long long W2 = (long long)blockIdx.x * 8 + warp; W2 * RPW < nrows; W2 += (long long)gridDim.x * 8) {
        const long long R = W2 * RPW + grp;
        const bool valid = R < nrows;
        const long long Rc = valid ? R : (nrows - 1);
        const int bh = (int)(Rc / S), t = (int)(Rc % S);
        const size_t rs = (size_t)bh * Spad + t;
        const uint8_t* pr = partO + rs * KS_SLOTS * RB;
        const uint4 selw = *reinterpret_cast<const uint4*>(sel + rs * 4);
        int nset;
        {
            nset = __popc(selw.x) + __popc(selw.y) + __popc(selw.z) + __popc(selw.w);   // the stored words hold S_t (top-8 plus the local-window blocks)
        }
        float4 ml[NSL];
#pragma unroll
        for (int q = 0; q < NSL; ++q) {
            const int sl = l + q * LPR;
            ml[q] = make_float4(NEG_INF, 0.f, 0.f, 0.f);
            if (sl < nset) ml[q] = *reinterpret_cast<const float4*>(pr + sl * RB + D);
        }
        uint32_t raw[KS_SLOTS][NW];
#pragma unroll
        for (int s2 = 0; s2 < KS_SLOTS; ++s2) {
            if (s2 < nset) {
                const uint8_t* rp = pr + s2 * RB + l * DL;
                if constexpr (NW == 4) { const uint4 w4 = *reinterpret_cast<const uint4*>(rp); raw[s2][0] = w4.x; raw[s2][1] = w4.y; raw[s2][2] = w4.z; raw[s2][3] = w4.w; }
                else if constexpr (NW == 2) { const uint2 w2 = *reinterpret_cast<const uint2*>(rp); raw[s2][0] = w2.x; raw[s2][1] = w2.y; }
                else { raw[s2][0] = *reinterpret_cast<const uint32_t*>(rp); }
            } else {
#pragma unroll
                for (int i = 0; i < NW; ++i) raw[s2][i] = 0u;
            }
        }
        float M = ml[0].x;
#pragma unroll
        for (int q = 1; q < NSL; ++q) M = fmaxf(M, ml[q].x);
#pragma unroll
        for (int off = LPR / 2; off >= 1; off >>= 1) M = fmaxf(M, __shfl_xor_sync(0xffffffffu, M, off));
        const float Ms = (M == NEG_INF) ? 0.f : M;
        float wv[NSL], wqv[NSL];
        float wsum = 0.f;
#pragma unroll
        for (int q = 0; q < NSL; ++q) {
            const int sl = l + q * LPR;
            wv[q] = (sl < nset) ? ex2f(ml[q].x - Ms) * ml[q].y : 0.f;
            wqv[q] = wv[q] * ml[q].z;
            wsum += wv[q];
        }
#pragma unroll
        for (int off = LPR / 2; off >= 1; off >>= 1) wsum += __shfl_xor_sync(0xffffffffu, wsum, off);
        float acc[DL];
#pragma unroll
        for (int i = 0; i < DL; ++i) acc[i] = 0.f;
#pragma unroll
        for (int s2 = 0; s2 < KS_SLOTS; ++s2) {
            const float ws2 = __shfl_sync(0xffffffffu, wqv[s2 / LPR], grp * LPR + (s2 % LPR));
#pragma unroll
            for (int i = 0; i < DL; ++i) {
                const int qv = (int)(int8_t)((raw[s2][i >> 2] >> (8 * (i & 3))) & 0xffu);
                acc[i] = fmaf(ws2, (float)qv, acc[i]);
            }
        }
        const float inv = 1.f / wsum;
        if (valid) {
#pragma unroll
            for (int i = 0; i < DL / 2; ++i)
                *reinterpret_cast<uint32_t*>(O + (size_t)R * D + l * DL + 2 * i) = pack_bf16(acc[2 * i] * inv, acc[2 * i + 1] * inv);
        }
    }
}

// ---------------------------------------------------------------- host
struct Scratch {
    torch::Tensor kcf, sel, cm, po;
};
static std::map<std::tuple<int, int, int, int>, Scratch> g_scratch;

template <int D>
static void launch_fs(const bf16* q, const bf16* k, const bf16* v, bf16* o, Scratch& sc, int BH, int S, cudaStream_t st);

template <int D>
static void launch_front(const bf16* q, const bf16* k, Scratch& sc, int BH, int S, cudaStream_t st) {
    const int NBLK = (S + 63) / 64;
    const int NT = (NBLK + 7) / 8;
    const int Spad = ((S + 127) / 128) * 128;
    const float scale = 1.0f / sqrtf((float)D);
    uint32_t* kcf = reinterpret_cast<uint32_t*>(sc.kcf.data_ptr<int>());
    uint32_t* sel = reinterpret_cast<uint32_t*>(sc.sel.data_ptr<int>());
    uint2* cmp = reinterpret_cast<uint2*>(sc.cm.data_ptr<int>());
    prep_kernel<D><<<dim3(NBLK, BH), 256, 0, st>>>(k, kcf, S, NT);
    constexpr int sel_smem = 64 * 132 * 4 + 128 * 2 * 4;
    static const bool init_sel = (cudaFuncSetAttribute(select_kernel<D, 0>, cudaFuncAttributeMaxDynamicSharedMemorySize, sel_smem) == cudaSuccess);
    (void)init_sel;
    select_kernel<D, 0><<<dim3(NBLK, BH), 128, sel_smem, st>>>(q, k, kcf, cmp, sel, NBLK, S, NT, Spad, scale);
}

template <int D>
static void launch_back(const bf16* q, const bf16* k, const bf16* v, bf16* o, Scratch& sc, int BH, int S, cudaStream_t st, int cfg) {
    const int NBLK = (S + 63) / 64;
    const int Spad = ((S + 127) / 128) * 128;
    const float scale = 1.0f / sqrtf((float)D);
    const float cs = scale * 1.4426950408889634f;
    uint32_t* sel = reinterpret_cast<uint32_t*>(sc.sel.data_ptr<int>());
    uint2* cmp = reinterpret_cast<uint2*>(sc.cm.data_ptr<int>());
    if constexpr (D == 64) {
        if (cfg == 12) {   // small S: static 8-row warps with a decoupled K/V ring, 1 row block per CTA
            constexpr int CONS = 8, STG = 4;
            constexpr int OSTB = 8 * (D + 8) * 2;
            constexpr int smem = STG * 2 * 64 * D * 2 + 2 * STG * 8 + 16 + CONS * OSTB + 64;
            static const bool init_b8 = (cudaFuncSetAttribute(attend_b8d<D, CONS, STG>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess);
            (void)init_b8;
            attend_b8d<D, CONS, STG><<<NBLK * BH, (CONS + 1) * 32, smem, st>>>(q, k, v, cmp, o, S, NBLK, BH, NBLK, cs);
            return;
        }
    }
    uint8_t* po = reinterpret_cast<uint8_t*>(sc.po.data_ptr<int>());
    auto launchn = [&](auto kern, int W) {
        const int smem = 2 * 64 * D * 2 + ((D == 64) ? 4096 : 2048) * 2 + 128 + ((D == 64) ? W * 16 * D * 2 : 0);
        cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
        kern<<<NBLK * BH, W * 32, smem, st>>>(q, k, v, cmp, sel, po, S, NBLK, BH, Spad, cs);
    };
    if (D == 64) {
        if (cfg == 72) launchn(ksn_attend<D, 4, 3>, 4);
        else if (cfg == 73) launchn(ksn_attend<D, 4, 2>, 4);
        else launchn(ksn_attend<D, 8, 2>, 8);
    } else {
        if (cfg == 73) launchn(ksn_attend<D, 2, 2>, 2);
        else if (cfg == 74) launchn(ksn_attend<D, 8, 1>, 8);
        else launchn(ksn_attend<D, 4, 1>, 4);
    }
    const int nrows = BH * S;
    constexpr int LPR = (D == 64) ? 8 : 16;          // lanes per row in the merge (4 or 2 rows per warp)
    const int mb = std::min((int)((nrows + (32 / LPR) - 1) / (32 / LPR)), 188 * 8);
    ksn_merge2<D, LPR><<<mb, 256, 0, st>>>(po, sel, o, S, Spad, BH);
}

template <>
void launch_fs<64>(const bf16* q, const bf16* k, const bf16* v, bf16* o, Scratch& sc, int BH, int S, cudaStream_t st) {
    constexpr int D = 64;
    const int NBLK = (S + 63) / 64;
    const int NT = (NBLK + 7) / 8;
    const float scale = 1.0f / sqrtf((float)D);
    const float cs = scale * 1.4426950408889634f;
    uint32_t* kcf = reinterpret_cast<uint32_t*>(sc.kcf.data_ptr<int>());
    constexpr int STG = 4;
    constexpr int smem = STG * 2 * 64 * D * 2 + 2 * STG * 8 + 16 + 8 * (8 * (D + 8) * 2) + 64 * 40 * 4 + 34 * 2 * 4 + 64;
    static const bool init_fs = (cudaFuncSetAttribute(attend_fs<STG>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem) == cudaSuccess);
    (void)init_fs;
    static const int nsm_fs = [] { int d = 0, n = 0; cudaGetDevice(&d); cudaDeviceGetAttribute(&n, cudaDevAttrMultiProcessorCount, d); return n; }();
    attend_fs<STG><<<NBLK * BH, 288, smem, st>>>(q, k, v, kcf, o, S, NBLK, BH, NT, scale, cs, nsm_fs);
}

static void launch_prep64(const bf16* k, Scratch& sc, int BH, int S, cudaStream_t st) {
    const int NBLK = (S + 63) / 64;
    const int NT = (NBLK + 7) / 8;
    uint32_t* kcf = reinterpret_cast<uint32_t*>(sc.kcf.data_ptr<int>());
    prep_kernel<64><<<dim3(NBLK, BH), 256, 0, st>>>(k, kcf, S, NT);
}

static Scratch& get_scratch(int BH, int S, int D, torch::Device dev) {
    auto key = std::make_tuple(BH, S, D, (int)dev.index());
    auto it = g_scratch.find(key);
    if (it != g_scratch.end()) return it->second;
    const int NBLK = (S + 63) / 64;
    const int NT = (NBLK + 7) / 8;
    const int Spad = ((S + 127) / 128) * 128;
    Scratch sc;
    auto opt = torch::TensorOptions().dtype(torch::kInt32).device(dev);
    sc.kcf = torch::zeros({(long)BH * NT * 32 * 3 * (D / 16) * 2}, opt);
    sc.sel = torch::zeros({(long)BH * Spad * 4}, opt);
    sc.cm = torch::zeros({(long)BH * NBLK * 128 * 2}, opt);
    sc.po = torch::empty({(long)BH * Spad * KS_SLOTS * ((D + KS_RBPAD) / 4)}, opt);
    return g_scratch.emplace(key, sc).first->second;
}

static int auto_cfg(int D, int S) {
    if (D == 64) return (S <= 1100) ? 13 : 74;                         // fused select+attend / key-stationary
    return 72;                                                         // D = 128: key-stationary 4 warps x 1 CTA
}

torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, int cfg = -2) {
    TORCH_CHECK(q.is_cuda() && q.dtype() == torch::kBFloat16, "bf16 cuda");
    TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "contiguous");
    const int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
    TORCH_CHECK((S + 63) / 64 <= 128, "S too long");
    if (cfg == -2) cfg = auto_cfg(D, S);
    Scratch& sc = get_scratch(B * H, S, D, q.device());
    cudaStream_t st = at::cuda::getCurrentCUDAStream();
    const bf16* qp = reinterpret_cast<const bf16*>(q.data_ptr());
    const bf16* kp = reinterpret_cast<const bf16*>(k.data_ptr());
    const bf16* vp = reinterpret_cast<const bf16*>(v.data_ptr());
    if (cfg == 13 && D == 64) {
        launch_prep64(kp, sc, B * H, S, st);
        auto o13 = torch::empty_like(q);   // allocated while the prep kernel runs
        launch_fs<64>(qp, kp, vp, reinterpret_cast<bf16*>(o13.data_ptr()), sc, B * H, S, st);
        return o13;
    }
    if (D == 64) launch_front<64>(qp, kp, sc, B * H, S, st);
    else if (D == 128) launch_front<128>(qp, kp, sc, B * H, S, st);
    else TORCH_CHECK(false, "D must be 64 or 128");
    auto o = torch::empty_like(q);   // allocated while the front kernels run
    bf16* op = reinterpret_cast<bf16*>(o.data_ptr());
    if (D == 64) launch_back<64>(qp, kp, vp, op, sc, B * H, S, st, cfg);
    else launch_back<128>(qp, kp, vp, op, sc, B * H, S, st, cfg);
    return o;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &nsa_forward, "nsa forward", pybind11::arg("q"), pybind11::arg("k"), pybind11::arg("v"), pybind11::arg("cfg") = -2);
}

'''

_ext = load_inline(
    name="nsa_bench_ext_vv93",
    cpp_sources=[""],
    cuda_sources=[_CUDA_SRC],
    functions=None,
    extra_cuda_cflags=["-O3", "-gencode=arch=compute_120,code=sm_120"],
    extra_cflags=["-O3"],
    verbose=False,
)


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))

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        return _ext.forward(q, k, v)

    # skip nn.Module hook bookkeeping: the whole call is a few microseconds on the small shapes
    __call__ = forward