KernelBench cuda · RTX PRO 6000
DeepSeek NSA DeepSeek V4 Flash (0731)
4.55%geomean peak fraction across shapes
manually audited: clean
Manual static audit covered all 553 lines of solution.py, all 319427 transcript records and 321 agent tool calls, result.json, check.log, benchmark.log, the CUDA-language sidecar, and the frozen grader/template files. solution.py:54-166 computes causal block importance from the live q and k tensors, using bf16 mma.sync for full blocks and SIMT for partial causal blocks. solution.py:180-412 performs input-dependent top-8 block selection, sliding-window union, online softmax, and v accumulation. solution.py:523-544 makes fresh importance and output tensors on every forward and launches both CUDA kernels on the current q, k, and v buffers. The only retained Python object is the compiled extension handle. No output cache, input identity/data_ptr key, CUDA graph, result table, constant answer, stack/caller or check.py sniff, reference import, or forbidden library exists in the final computation, so no empirical same-buffer overwrite cache test is required. The fixed nsa_sparse_attn.so path is under this run's TORCH_EXTENSIONS_DIR and caches only compiled code, not inputs or outputs. The trace edits this cell's solution.py and local runit.sh/rebuild.sh development helpers, plus project-scoped agent memory notes and disposable profiling artifacts. Foreign run IDs and prompts appear only in passive ps output while diagnosing GPU-lock contention; no tool input names a foreign artifact and none is opened or copied. The agent read the unmodified check.py and shared eval modules but did not edit a grader, alter tolerances, set KBH_NUMERIC_STRESS, or branch submitted code on grader behavior. All seven archived grader files are byte-for-byte equal to template_files, consistent with result.json template_mutated=false. check.log records framework=ptx and PASS from that unmodified checker. check.py loops over both S=256/384 shapes, seeds 42/123, and the configured numeric_stress_cases before PASS, but check.log gives no per-case names or magnitudes, so none are claimed. cuda_language.json reports triton_cheat=false, dsl_cheat=false, forbidden_hits=[], and ok=true. The archived pre-regrade benchmark recorded per-shape fractions 0.0340, 0.0636, 0.0701, 0.0515, 0.0231, and 0.0510, with peak_fraction 0.0457 and RESULT: OK. Publication metrics come from the sequential isolated regrade on 2026-08-03 using NVIDIA RTX PRO 6000 Blackwell Server Edition: correct=true and peak_fraction=0.0455. Its per-shape fractions are 0.0336, 0.0636, 0.0699, 0.0515, 0.0227, and 0.0507. The archived 0.0457 value is retained here only as pre-regrade provenance and is superseded by this isolated publish-grade result.
harnessor-fableagent session3h 13mtotal wall3h 13mcheck4sbenchmark2soutput tokens—gpu-lock wait53mgpu-lock held6mregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
1×16×2048×641.022 ms3.4%17 TFLOPS · 3% of 500 TF bf16 peak · also 0.02 TB/s (1% of HBM)
1×16×4127×642.192 ms6.4%32 TFLOPS · 6% of 500 TF bf16 peak · also 0.02 TB/s (1% of HBM)
1×8×8192×643.932 ms7.0%35 TFLOPS · 7% of 500 TF bf16 peak · also 0.01 TB/s (0% of HBM)
1×8×8191×12810.676 ms5.1%26 TFLOPS · 5% of 500 TF bf16 peak · also 0.01 TB/s (0% of HBM)
4×8×1024×640.756 ms2.3%11 TFLOPS · 2% of 500 TF bf16 peak · also 0.02 TB/s (1% of HBM)
2×8×3000×641.453 ms5.1%25 TFLOPS · 5% of 500 TF bf16 peak · also 0.02 TB/s (1% of HBM)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(3.4% · 6.4% · 7.0% · 5.1% · 2.3% · 5.1%) = 4.5%
Kernel source (redacted)
"""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)]
20260802_203913_or-fable_deepseek_deepseek-v4-flash-0731_02_deepseek_nsa