"""DeepSeek NSA-inspired sparse attention in CUDA (bench-faithful semantics).

Two-kernel pipeline, both launched from a single C++ extension entry point:

  nsa_importance : block-importance GEMM.  For each query tile of QTILE queries
      we compute the mean of (q_t . k_j / sqrt(D)) over each 64-key block.  The
      bulk is done with bf16 tensor-core MMA (m16n8k16); only the 1-2 diagonal
      blocks (the ones containing the tile's query positions) are done with
      plain SIMT prefix sums because they are partially causal.
  nsa_attention : top-N block selection + online-softmax sparse attention.
      Each query keeps its top_n blocks (by importance) unioned with a sliding
      window of `window` tokens.  A block-major loop over keys loads each
      k/v block once into shared memory and fans it out to every query in the
      tile that selected it, so selected k/v are not re-read per query.

Correctness gate (check.py) uses S in {256, 384}: fewer than 8 blocks exist,
so top-N selects every block and the output is dense causal attention -- still
handled by the same code path.  The benchmark shapes exercise the real sparse
selection.
"""
from __future__ import annotations

import math

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

_CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>
#include <cfloat>
#include <algorithm>

#define QTILE 128
#define TOP_N 8

using bf16 = __nv_bfloat16;

__device__ __forceinline__ uint32_t ld32u(const void* p) {
    return *(const uint32_t*)p;
}

__device__ __forceinline__ float2 b2_pair(uint32_t u) {
    __nv_bfloat162 p = *reinterpret_cast<const __nv_bfloat162*>(&u);
    return make_float2(__bfloat162float(p.x), __bfloat162float(p.y));
}

// ---------------------------------------------------------------------------
// Kernel A: block importance = mean of (q . k / sqrt(D)) over causal keys.
// ---------------------------------------------------------------------------
__global__ void nsa_importance(
    const bf16* __restrict__ q,
    const bf16* __restrict__ k,
    float* __restrict__ imp,          // (B*H, S, nb) fp32, -1e9 for non-causal
    int B, int H, int S, int D,
    int block_size, float scale)
{
    const int bh = blockIdx.y;
    const int q0 = blockIdx.x * QTILE;
    const int q1 = min(q0 + QTILE, S);
    const int nq = q1 - q0;
    const int nb = (S + block_size - 1) / block_size;
    const int tid = threadIdx.x;
    const int nthreads = blockDim.x;
    const int warp_id = tid >> 5;
    const int lane = tid & 31;

    const bf16* qb = q + (long)bh * S * D;
    const bf16* kb = k + (long)bh * S * D;
    float* impb = imp + (long)bh * S * nb;

    // Padded row stride avoids 8-way shared bank conflicts on MMA fragments.
    const int LD = D + 2;

    extern __shared__ char smem_raw[];
    bf16* q_sh = (bf16*)smem_raw;
    bf16* k_sh = q_sh + QTILE * LD;

    // Load the query tile (zero-fill rows beyond nq so MMA sees defined data).
    for (int i = tid; i < QTILE * D; i += nthreads) {
        int row = i / D, col = i % D;
        q_sh[row * LD + col] = (row < nq) ? qb[(long)(q0 + row) * D + col] : __float2bfloat16(0.f);
    }
    for (int i = tid; i < nq * nb; i += nthreads)
        impb[(long)(q0 + i / nb) * nb + i % nb] = -1e9f;
    __syncthreads();

    // Blocks strictly before the tile are fully causal: MMA full 64-key sums.
    const int last_full = q0 / block_size;
    for (int bi = 0; bi < last_full; bi++) {
        const int s0 = bi * block_size;
        for (int i = tid; i < block_size * D; i += nthreads)
            k_sh[(i / D) * LD + (i % D)] = kb[(long)s0 * D + i];
        __syncthreads();

        const int qrow0 = warp_id * 16;
        const int g = lane >> 2, l = lane & 3;
        // Separate accumulator per n-tile breaks the accumulate dependency
        // chain across the 8 n-tiles (each is an independent 64-key partial).
        float c0[8], c1[8], c2[8], c3[8];
#pragma unroll
        for (int nt = 0; nt < 8; nt++) { c0[nt] = 0.f; c1[nt] = 0.f; c2[nt] = 0.f; c3[nt] = 0.f; }
#pragma unroll
        for (int ks = 0; ks < D / 16; ks++) {
            const uint32_t a0 = ld32u(&q_sh[(qrow0 + g) * LD + ks * 16 + 2 * l]);
            const uint32_t a1 = ld32u(&q_sh[(qrow0 + g) * LD + ks * 16 + 8 + 2 * l]);
            const uint32_t a2 = ld32u(&q_sh[(qrow0 + g + 8) * LD + ks * 16 + 2 * l]);
            const uint32_t a3 = ld32u(&q_sh[(qrow0 + g + 8) * LD + ks * 16 + 8 + 2 * l]);
#pragma unroll
            for (int nt = 0; nt < 8; nt++) {
                const uint32_t b0 = ld32u(&k_sh[(nt * 8 + g) * LD + ks * 16 + 2 * l]);
                const uint32_t b1 = ld32u(&k_sh[(nt * 8 + g) * LD + ks * 16 + 8 + 2 * l]);
                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"(c0[nt]), "+f"(c1[nt]), "+f"(c2[nt]), "+f"(c3[nt])
                    : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
            }
        }
        // Sum the n-tile partials, then column-sum reduce over lanes 0..3.
        float sum0 = 0.f, sum1 = 0.f;
#pragma unroll
        for (int nt = 0; nt < 8; nt++) { sum0 += c0[nt] + c1[nt]; sum1 += c2[nt] + c3[nt]; }
#pragma unroll
        for (int off = 1; off < 4; off <<= 1) {
            sum0 += __shfl_xor_sync(0xffffffffu, sum0, off);
            sum1 += __shfl_xor_sync(0xffffffffu, sum1, off);
        }
        if ((lane & 3) == 0) {
            const int r0 = qrow0 + g;
            const int r1 = qrow0 + g + 8;
            const float inv = scale / block_size;
            if (r0 < nq) impb[(long)(q0 + r0) * nb + bi] = sum0 * inv;
            if (r1 < nq) impb[(long)(q0 + r1) * nb + bi] = sum1 * inv;
        }
        __syncthreads();
    }

    // Diagonal blocks (those overlapping the query tile): SIMT prefix sums.
    if (tid < nq) {
        const int t = q0 + tid;
        const int b_lo = q0 / block_size;
        const int b_hi = t / block_size;
        const bf16* qrow = &q_sh[tid * LD];
        for (int bi = b_lo; bi <= b_hi; bi++) {
            const int s0 = bi * block_size;
            const int cnt = min(block_size, t + 1 - s0);
            float sum = 0.f;
            for (int j = 0; j < cnt; j++) {
                const bf16* krow = kb + (long)(s0 + j) * D;
                float dot = 0.f;
#pragma unroll
                for (int d = 0; d < D; d += 2) {
                    const float2 qf = b2_pair(ld32u(&qrow[d]));
                    const float2 kf = b2_pair(ld32u(&krow[d]));
                    dot += qf.x * kf.x + qf.y * kf.y;
                }
                sum += dot;
            }
            impb[(long)t * nb + bi] = sum * scale / cnt;
        }
    }
}

// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// Kernel B: top-N select + online-softmax sparse attention.
// Dense block-major with per-query shared state:
//   (1) each query selects top-N blocks unioned with the sliding window;
//   (2) a CSR of (query, block) pairs grouped by block is built in shared;
//   (3) key blocks stream through shared memory once per query tile (so k/v are
//       reused across the tile's queries, not re-read from L2 per query);
//   (4) each block's pairs are processed by one thread pair (D-split) that keeps
//       the block's online-softmax partials in registers and merges once into the
//       query's shared running state.
// ---------------------------------------------------------------------------
template <int D, int QT>
__global__ void nsa_attention(
    const bf16* __restrict__ q,
    const bf16* __restrict__ k,
    const bf16* __restrict__ v,
    const float* __restrict__ imp,
    bf16* __restrict__ o,
    int B, int H, int S, int block_size, int top_n, int window, float scale)
{
    const int bh = blockIdx.y;
    const int q0 = blockIdx.x * QT;
    const int q1 = min(q0 + QT, S);
    const int nq = q1 - q0;
    const int nb = (S + block_size - 1) / block_size;
    const int tid = threadIdx.x;
    const int nthreads = blockDim.x;
    const int qt = tid >> 1;
    const int half = tid & 1;
    const int t = q0 + qt;
    const int d0 = half * (D >> 1);
    const int LD = D + 2;
    const int MAX_PAIR = QT * (TOP_N + 2);

    const bf16* qb = q + (long)bh * S * D;
    const bf16* kb = k + (long)bh * S * D;
    const bf16* vb = v + (long)bh * S * D;
    const float* impb = imp + (long)bh * S * nb;
    bf16* ob = o + (long)bh * S * D;

    extern __shared__ char smem_raw[];
    char* sp = smem_raw;
    bf16* k_sh = (bf16*)sp;                    sp += block_size * LD * sizeof(bf16);
    bf16* v_sh = (bf16*)sp;                    sp += block_size * LD * sizeof(bf16);
    float* acc_sh = (float*)sp;                sp += QT * (D + 2) * sizeof(float);
    float* m_sh = (float*)sp;                  sp += QT * sizeof(float);
    float* l_sh = (float*)sp;                  sp += QT * sizeof(float);
    int* qlist_sh = (int*)sp;                  sp += QT * (TOP_N + 2) * sizeof(int);
    unsigned char* qlist_t8 = (unsigned char*)sp; sp += QT * (TOP_N + 2);
    int* qlist_n = (int*)sp;                   sp += QT * sizeof(int);
    int* block_count = (int*)sp;               sp += nb * sizeof(int);
    int* block_start = (int*)sp;               sp += (nb + 1) * sizeof(int);
    int* cursor = (int*)sp;                    sp += nb * sizeof(int);
    int* sel_query = (int*)sp;                 sp += MAX_PAIR * sizeof(int);
    unsigned char* sel_t8 = (unsigned char*)sp; sp += MAX_PAIR;

    // ---- Selection: build each query's deduped block list ----
    if (qt < nq && half == 0) {
        float best[TOP_N];
        int bidx[TOP_N];
#pragma unroll
        for (int i = 0; i < TOP_N; i++) { best[i] = -1e9f; bidx[i] = -1; }
        const float* impt = impb + (long)t * nb;
        for (int bi = 0; bi < nb; bi++) {
            const float val = impt[bi];
            if (val > best[TOP_N - 1]) {
                int pos = TOP_N - 1;
                while (pos > 0 && val > best[pos - 1]) {
                    best[pos] = best[pos - 1];
                    bidx[pos] = bidx[pos - 1];
                    --pos;
                }
                best[pos] = val;
                bidx[pos] = bi;
            }
        }
        int* ql = &qlist_sh[qt * (TOP_N + 2)];
        unsigned char* qf = &qlist_t8[qt * (TOP_N + 2)];
        int n = 0;
        for (int i = 0; i < TOP_N; i++) {
            if (bidx[i] >= 0) { ql[n] = bidx[i]; qf[n] = 1; ++n; }
        }
        const int w_diag = t / block_size;
        const int w_prev = (t >= block_size && (t % block_size) < window - 1)
                               ? t / block_size - 1 : -1;
        for (int wi = 0; wi < 2; wi++) {
            const int wb = (wi == 0) ? w_diag : w_prev;
            if (wb < 0 || wb >= nb) continue;
            bool dup = false;
            for (int i = 0; i < n; i++)
                if (ql[i] == wb) { dup = true; break; }
            if (!dup) { ql[n] = wb; qf[n] = 0; ++n; }
        }
        qlist_n[qt] = n;
    }

    // ---- Build CSR of (query, block) pairs grouped by block ----
    for (int i = tid; i < nb; i += nthreads) block_count[i] = 0;
    __syncthreads();
    if (qt < nq && half == 0) {
        const int n = qlist_n[qt];
        const int* ql = &qlist_sh[qt * (TOP_N + 2)];
        for (int i = 0; i < n; i++) atomicAdd(&block_count[ql[i]], 1);
    }
    __syncthreads();
    if (tid == 0) {
        block_start[0] = 0;
        for (int bi = 0; bi < nb; bi++) block_start[bi + 1] = block_start[bi] + block_count[bi];
    }
    for (int i = tid; i < nb; i += nthreads) cursor[i] = 0;
    __syncthreads();
    if (qt < nq && half == 0) {
        const int n = qlist_n[qt];
        const int* ql = &qlist_sh[qt * (TOP_N + 2)];
        const unsigned char* qf = &qlist_t8[qt * (TOP_N + 2)];
        for (int i = 0; i < n; i++) {
            const int slot = block_start[ql[i]] + atomicAdd(&cursor[ql[i]], 1);
            sel_query[slot] = qt;
            sel_t8[slot] = qf[i];
        }
    }
    // ---- Init per-query attention state ----
    for (int i = tid; i < nq * (D + 2); i += nthreads) acc_sh[i] = 0.f;
    for (int i = tid; i < nq; i += nthreads) { m_sh[i] = -FLT_MAX; l_sh[i] = 0.f; }
    __syncthreads();

    // ---- Main block stream (thread-pair per (query, block) pair) ----
    const int max_bi = min(nb, (q1 - 1) / block_size + 1);
    const int GROUP = (D == 64) ? 16 : 8;
    for (int bi = 0; bi < max_bi; bi++) {
        const long s0 = (long)bi * block_size;
        for (int i = tid; i < block_size * D; i += nthreads) {
            const int r = i / D, c = i % D;
            k_sh[r * LD + c] = kb[s0 * D + i];
            v_sh[r * LD + c] = vb[s0 * D + i];
        }
        __syncthreads();
        const int count = block_count[bi];
        if (tid < 2 * count) {
            const int p = tid >> 1;
            const int h = tid & 1;
            const int q = sel_query[block_start[bi] + p];
            const bool is_t8 = sel_t8[block_start[bi] + p] != 0;
            const int tt = q0 + q;
            const int w0q = max(0, tt + 1 - window);
            const int lo = is_t8 ? 0 : max(0, w0q - (int)s0);
            const int hi = min(block_size, tt + 1 - (int)s0);
            if (lo < hi) {
                const bf16* qrow = &qb[(long)tt * D + d0];
                uint32_t qr[D / 4];
#pragma unroll
                for (int i = 0; i < D / 4; i++) qr[i] = ld32u(&qrow[2 * i]);
                float pm = -FLT_MAX, pl = 0.f;
                float pacc[D / 2];
#pragma unroll
                for (int i = 0; i < D / 2; i++) pacc[i] = 0.f;
                int j = lo;
                for (; j + GROUP <= hi; j += GROUP) {
                    const unsigned act = __activemask();
                    float sc[GROUP];
#pragma unroll
                    for (int g = 0; g < GROUP; g++) {
                        const bf16* krow = &k_sh[(j + g) * LD + d0];
                        float dot = 0.f;
#pragma unroll
                        for (int i = 0; i < D / 4; i++) {
                            const float2 kf = b2_pair(ld32u(&krow[2 * i]));
                            const float2 qf = b2_pair(qr[i]);
                            dot += qf.x * kf.x + qf.y * kf.y;
                        }
                        dot += __shfl_xor_sync(act, dot, 1);
                        sc[g] = dot * scale;
                    }
#pragma unroll
                    for (int g = 0; g < GROUP; g++) {
                        const float s = sc[g];
                        const float nm = fmaxf(pm, s);
                        const float alpha = expf(pm - nm);
                        const float beta = expf(s - nm);
                        pl = pl * alpha + beta;
                        const bf16* vrow = &v_sh[(j + g) * LD + d0];
#pragma unroll
                        for (int i = 0; i < D / 4; i++) {
                            const float2 vf = b2_pair(ld32u(&vrow[2 * i]));
                            pacc[2 * i] = pacc[2 * i] * alpha + beta * vf.x;
                            pacc[2 * i + 1] = pacc[2 * i + 1] * alpha + beta * vf.y;
                        }
                        pm = nm;
                    }
                }
                for (; j < hi; j++) {
                    const unsigned act = __activemask();
                    const bf16* krow = &k_sh[j * LD + d0];
                    const bf16* vrow = &v_sh[j * LD + d0];
                    float dot = 0.f;
#pragma unroll
                    for (int i = 0; i < D / 4; i++) {
                        const float2 kf = b2_pair(ld32u(&krow[2 * i]));
                        const float2 qf = b2_pair(qr[i]);
                        dot += qf.x * kf.x + qf.y * kf.y;
                    }
                    dot += __shfl_xor_sync(act, dot, 1);
                    const float s = dot * scale;
                    const float nm = fmaxf(pm, s);
                    const float alpha = expf(pm - nm);
                    const float beta = expf(s - nm);
                    pl = pl * alpha + beta;
#pragma unroll
                    for (int i = 0; i < D / 4; i++) {
                        const float2 vf = b2_pair(ld32u(&vrow[2 * i]));
                        pacc[2 * i] = pacc[2 * i] * alpha + beta * vf.x;
                        pacc[2 * i + 1] = pacc[2 * i + 1] * alpha + beta * vf.y;
                    }
                    pm = nm;
                }
                // Merge this block's partial into the query's running state.
                const float mq = m_sh[q];
                const float lq = l_sh[q];
                const float nm = fmaxf(mq, pm);
                const float alpha = expf(mq - nm);
                const float beta = expf(pm - nm);
                m_sh[q] = nm;
                l_sh[q] = lq * alpha + pl * beta;
                float* aq = &acc_sh[q * (D + 2) + d0];
#pragma unroll
                for (int i = 0; i < D / 2; i++) aq[i] = aq[i] * alpha + pacc[i] * beta;
            }
        }
        __syncthreads();
    }

    // ---- Finalize ----
    if (qt < nq) {
        const float lq = l_sh[qt];
        if (lq > 0.f) {
            const float inv = 1.f / lq;
            bf16* orow = &ob[(long)t * D + d0];
            const float* aq = &acc_sh[qt * (D + 2) + d0];
#pragma unroll
            for (int i = 0; i < D / 2; i++)
                orow[i] = __float2bfloat16(aq[i] * inv);
        }
    }
}

// ---------------------------------------------------------------------------
void nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v,
                 torch::Tensor imp, torch::Tensor o,
                 int64_t block_size, int64_t top_n, int64_t window, double scale) {
    const int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
    const int nb = (S + block_size - 1) / block_size;
    const int nqt = (S + QTILE - 1) / QTILE;
    const int nthreads = 256;
    dim3 grid(nqt, B * H);

    static bool configured = false;
    if (!configured) {
        const int max_smem = 101376;  // device max dynamic shared memory (bytes)
        cudaFuncSetAttribute(nsa_importance,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, max_smem);
        cudaFuncSetAttribute(nsa_attention<64, 128>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, max_smem);
        cudaFuncSetAttribute(nsa_attention<128, 64>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, max_smem);
        configured = true;
    }

    const int LD = D + 2;
    const size_t smemA = ((size_t)QTILE * LD + (size_t)block_size * LD) * sizeof(bf16);
    nsa_importance<<<grid, nthreads, smemA>>>(
        (const bf16*)q.data_ptr(), (const bf16*)k.data_ptr(),
        (float*)imp.data_ptr(), B, H, S, D, (int)block_size, (float)scale);

    // Kernel B: dense block-major.  QT=128 for D=64, 64 for D=128 (shared budget).
    const int QTB = (D == 64) ? 128 : 64;
    const int MAX_PAIR = QTB * (TOP_N + 2);
    size_t smemB = 2 * (size_t)block_size * LD * sizeof(bf16)      // k_sh, v_sh
                 + (size_t)QTB * (D + 2) * sizeof(float)           // acc_sh (padded)
                 + (size_t)QTB * sizeof(float) * 2                 // m_sh, l_sh
                 + (size_t)QTB * (TOP_N + 2) * sizeof(int)         // qlist_sh
                 + (size_t)QTB * (TOP_N + 2)                       // qlist_t8
                 + (size_t)QTB * sizeof(int)                       // qlist_n
                 + ((size_t)nb * 2 + (nb + 1)) * sizeof(int)       // block_count, cursor, block_start
                 + (size_t)MAX_PAIR * (sizeof(int) + 1);           // sel_query, sel_t8
    const dim3 gridB((S + QTB - 1) / QTB, B * H);
    const int nthreadsB = 256;
    if (D == 64) {
        nsa_attention<64, 128><<<gridB, nthreadsB, smemB>>>(
            (const bf16*)q.data_ptr(), (const bf16*)k.data_ptr(),
            (const bf16*)v.data_ptr(), (const float*)imp.data_ptr(),
            (bf16*)o.data_ptr(), B, H, S, (int)block_size, (int)top_n, (int)window,
            (float)scale);
    } else {
        nsa_attention<128, 64><<<gridB, nthreadsB, smemB>>>(
            (const bf16*)q.data_ptr(), (const bf16*)k.data_ptr(),
            (const bf16*)v.data_ptr(), (const float*)imp.data_ptr(),
            (bf16*)o.data_ptr(), B, H, S, (int)block_size, (int)top_n, (int)window,
            (float)scale);
    }
}
"""

_CPP_SOURCE = r"""
#include <torch/extension.h>
void nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v,
                 torch::Tensor imp, torch::Tensor o,
                 int64_t block_size, int64_t top_n, int64_t window, double scale);
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is not None:
        return _extension
    import importlib.util
    import os

    cache_dir = os.path.join(
        os.environ.get("TORCH_EXTENSIONS_DIR", os.path.expanduser("~/.cache/torch_extensions")),
        "nsa_sparse_attn",
    )
    so_path = os.path.join(cache_dir, "nsa_sparse_attn.so")
    if os.path.exists(so_path):
        try:
            spec = importlib.util.spec_from_file_location("nsa_sparse_attn", so_path)
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
            _extension = mod
            return _extension
        except Exception:
            pass
    _extension = load_inline(
        name="nsa_sparse_attn",
        cpp_sources=_CPP_SOURCE,
        cuda_sources=_CUDA_SOURCE,
        functions=["nsa_forward"],
        extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
        verbose=False,
    )
    return _extension


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.block_size = 64
        self.top_n = 8
        self.sliding_window = 64
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16))
        self._ext = None

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        if self._ext is None:
            self._ext = _get_extension()
        B, H, S, D = q.shape
        if q.dtype != torch.bfloat16:
            q = q.to(torch.bfloat16)
        if k.dtype != torch.bfloat16:
            k = k.to(torch.bfloat16)
        if v.dtype != torch.bfloat16:
            v = v.to(torch.bfloat16)
        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
        nb = (S + self.block_size - 1) // self.block_size
        imp = torch.empty(B * H, S, nb, dtype=torch.float32, device=q.device)
        o = torch.empty_like(q)
        scale = 1.0 / math.sqrt(D)
        self._ext.nsa_forward(
            q, k, v, imp, o,
            self.block_size, self.top_n, self.sliding_window, scale,
        )
        return o


def get_init_inputs():
    return [1, 16, 1024, 64]


def get_inputs():
    return [torch.randn(1, 16, 1024, 64, dtype=torch.bfloat16) for _ in range(3)]