KernelBench cuda · H100
DeepSeek NSA DeepSeek V4 Flash (0731)
1.61%geomean peak fraction across shapes
manually audited: clean
Manual static audit covered all 540 lines of solution.py and all 362358 transcript records, including 369 agent tool calls, plus result.json, check.log, benchmark.log, the CUDA-language sidecar, and the frozen grader/template files. solution.py computes live block means from k, live q-to-block importance scores, per-query top-8 and sliding-window selection, and online-softmax attention over the selected live k/v data. Each forward makes fresh contiguous inputs as needed, a fresh output, and a fresh k-mean workspace, then launches both CUDA kernels. The only retained object is the compiled extension handle. There is no input identity/data_ptr cache, CUDA graph, cached or constant output, result table, fake computation, stack/caller or check.py sniff, reference import, or forbidden library, so no same-buffer overwrite cache test is required. The trace writes only this cell's solution and local validation/development helpers, temporary profiling artifacts, and project-scoped agent memory. A ps listing and one broad find passively exposed sibling prompts and foreign cuda.cu pathnames, but no foreign artifact was opened, copied, or used. The agent read the unmodified checker and shared eval modules but did not edit a grader, alter tolerances, or set numeric-stress controls. All seven archived grader files are byte-for-byte equal to template_files, consistent with result.json template_mutated=false. cuda_language.json reports framework=ptx, triton_cheat=false, dsl_cheat=false, forbidden_hits=[], and ok=true. The archived pre-regrade check failed during extension compilation because the run-local nvcc wrapper reported "nvcc is unavailable"; no numeric comparison was reached, so that was infrastructure rather than model incorrectness. Publication metrics supersede that status and come from the sequential isolated regrade on 2026-08-03 using NVIDIA H100 PCIe: check.log records PASS and benchmark.log records RESULT: OK, correct=true, and peak_fraction=0.0161, with per-shape fractions 0.0134, 0.0169, 0.0178, 0.0240, 0.0118, and 0.0151. The isolated check.log names no individual numeric-stress case or magnitude, so none is claimed.
harnessor-fableagent session2h 38mtotal wall2h 38mcheck27sbenchmark4soutput tokens—gpu-lock wait12mgpu-lock held3sregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
1×16×2048×641.698 ms1.3%10 TFLOPS · 1% of 756 TF bf16 peak · also 0.01 TB/s (0% of HBM)
1×16×4127×645.454 ms1.7%13 TFLOPS · 2% of 756 TF bf16 peak · also 0.01 TB/s (0% of HBM)
1×8×8192×6410.208 ms1.8%13 TFLOPS · 2% of 756 TF bf16 peak · also 0.00 TB/s (0% of HBM)
1×8×8191×12815.163 ms2.4%18 TFLOPS · 2% of 756 TF bf16 peak · also 0.00 TB/s (0% of HBM)
4×8×1024×640.964 ms1.2%9 TFLOPS · 1% of 756 TF bf16 peak · also 0.02 TB/s (1% of HBM)
2×8×3000×643.224 ms1.5%11 TFLOPS · 2% of 756 TF bf16 peak · also 0.01 TB/s (0% of HBM)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(1.3% · 1.7% · 1.8% · 2.4% · 1.2% · 1.5%) = 1.6%
Kernel source (redacted)
"""DeepSeek NSA-inspired sparse attention, CUDA kernel.
Implements the bench-simplified NSA semantics of the reference oracle:
1. Block importance = mean of q·k/sqrt(D) over causal keys in each 64-key block.
2. top-8 blocks by importance, unioned with the last 64-token sliding window.
3. Online-softmax attention over the selected keys only.
The block importance is computed exactly as q·(mean of k over the block) which
equals the mean of q·k (the reference computation) in real arithmetic, up to
fp accumulation order. The sliding window always covers the whole causal part
of the block that contains the query, so only that block's importance needs a
per-token (prefix-sum) correction; everything else uses precomputed block means.
CUDA-only: two kernels, plain mma.sync tensor-core GEMMs, no Triton/DSL.
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
BLOCK_SIZE = 64
TOP_N_BLOCKS = 8
SLIDING_WINDOW = 64
_Q_TILE = 64
_THREADS = 256
_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <mma.h>
#include <math.h>
using namespace nvcuda;
#define Q_TILE 64
#define BLOCK_K 64
#define THREADS 256
#define TOP_N_BLOCKS 8
__device__ __forceinline__ float b2f(__nv_bfloat16 x) { return __bfloat162float(x); }
__device__ __forceinline__ __nv_bfloat16 f2b(float x) { return __float2bfloat16(x); }
// ---------------------------------------------------------------------------
// Kernel 1: block means of K -> KMEAN (B,H,N,D) fp32 (unscaled).
// ---------------------------------------------------------------------------
__global__ void kmean_kernel(const __nv_bfloat16* __restrict__ K,
float* __restrict__ KMEAN,
int S, int D, int N) {
const int bh = blockIdx.x;
const __nv_bfloat16* Kb = K + (size_t)bh * S * D;
float* Km = KMEAN + (size_t)bh * N * D;
for (int j = blockIdx.y; j < N; j += gridDim.y) {
const int s0 = j * BLOCK_K;
const int s1 = min(s0 + BLOCK_K, S);
const int cnt = s1 - s0;
for (int d = threadIdx.x; d < D; d += blockDim.x) {
float sum = 0.f;
for (int c = s0; c < s1; ++c) sum += b2f(Kb[(size_t)c * D + d]);
Km[(size_t)j * D + d] = sum / (float)cnt;
}
}
}
// ---------------------------------------------------------------------------
// Kernel 2: fused block-score + select + sparse attention.
// Grid: (ceil(S/64), B*H). Block: 256 threads.
// Template on DIM (head dim, 64 or 128).
// ---------------------------------------------------------------------------
template <int DIM>
__global__ void __launch_bounds__(THREADS, 1)
nsa_kernel(const __nv_bfloat16* __restrict__ Q,
const __nv_bfloat16* __restrict__ K,
const __nv_bfloat16* __restrict__ V,
__nv_bfloat16* __restrict__ O,
const float* __restrict__ KMEAN,
int S, int N, float scale) {
constexpr int DSTR = DIM + (DIM == 64 ? 6 : 2); // padded row stride
constexpr int BSTR = BLOCK_K + (DIM == 64 ? 6 : 2);
constexpr int VSTR = BLOCK_K + (DIM == 64 ? 8 : 2); // transposed V stride
constexpr int KSTEP = DIM / 16; // mma k-steps for score gemm
constexpr int KPSTEP = BLOCK_K / 16; // mma k-steps for P@V
const int bh = blockIdx.y;
const int G = blockIdx.x; // tile index == key block containing queries
const int t0 = G * Q_TILE;
const size_t off = (size_t)bh * S * DIM;
extern __shared__ char smem_raw[];
float* s_kmean = (float*)smem_raw; // N*DSTR
__nv_bfloat16* s_q = (__nv_bfloat16*)(s_kmean + (size_t)N * DSTR); // Q_TILE*DSTR
float* s_bscore = (float*)(s_q + Q_TILE * DSTR); // Q_TILE*N fp32
float* s_acc = (float*)(s_bscore + (size_t)Q_TILE * N); // Q_TILE*DSTR
float* s_m = s_acc + Q_TILE * DSTR; // Q_TILE
float* s_l = s_m + Q_TILE; // Q_TILE
__nv_bfloat16* s_kblk = (__nv_bfloat16*)(s_l + Q_TILE); // BLOCK_K*DSTR
__nv_bfloat16* s_vblk = (__nv_bfloat16*)(s_kblk + BLOCK_K * DSTR);
__nv_bfloat16* s_ptile = (__nv_bfloat16*)(s_vblk + DIM * VSTR); // Q_TILE*BSTR
float* s_stile = (float*)(s_ptile + Q_TILE * BSTR); // Q_TILE*BSTR
uint32_t* s_sel = (uint32_t*)(s_stile + Q_TILE * BSTR); // Q_TILE*ceil(N/32)
__nv_bfloat16* s_kG = (__nv_bfloat16*)(s_sel + Q_TILE * ((N + 31) / 32)); // BLOCK_K*DSTR
uint8_t* s_block_active = (uint8_t*)(s_kG + BLOCK_K * DSTR); // N
const int tid = threadIdx.x;
const int warp = tid / 32;
const int lane = tid % 32;
// ---- load kmean (fp32) ----
for (int i = tid; i < N * DIM; i += THREADS) {
int j = i / DIM, d = i % DIM;
s_kmean[j * DSTR + d] = KMEAN[(size_t)bh * N * DIM + i];
}
// ---- load Q tile (raw bf16) ----
for (int i = tid; i < Q_TILE * DIM; i += THREADS) {
int q = i / DIM, d = i % DIM;
int t = t0 + q;
s_q[q * DSTR + d] = (t < S) ? Q[off + (size_t)t * DIM + d] : __float2bfloat16(0.f);
}
// ---- load block G's K (raw) ----
{
const int L = min(BLOCK_K, S - t0);
for (int i = tid; i < BLOCK_K * DIM; i += THREADS) {
int c = i / DIM, d = i % DIM;
__nv_bfloat16 val = __float2bfloat16(0.f);
if (c < L) val = K[off + (size_t)(t0 + c) * DIM + d];
s_kG[c * DSTR + d] = val;
}
}
__syncthreads();
// ---- block scores: bscore[q][j] = q · kmean[j] (SIMT fp32) ----
{
const int q = tid / 4;
const int sub = tid % 4;
if (q < Q_TILE) {
const __nv_bfloat16* qrow = s_q + q * DSTR;
for (int j = sub; j < N; j += 4) {
float acc = 0.f;
const float* kr = s_kmean + j * DSTR;
#pragma unroll
for (int d = 0; d < DIM; ++d) acc += b2f(qrow[d]) * kr[d];
s_bscore[q * N + j] = acc;
}
}
}
__syncthreads();
// ---- fix-up partial block + non-causal masking ----
if (tid < Q_TILE) {
const int q = tid;
const int t = t0 + q;
const bool real = (t < S);
if (real) {
for (int j = G + 1; j < N; ++j) s_bscore[q * N + j] = -INFINITY;
const int cnt = q + 1; // causal keys in block G (all in sliding window)
float acc = 0.f;
// partial mean = q · (sum of the first cnt keys of block G) / cnt
for (int i = 0; i < cnt; ++i)
for (int d = 0; d < DIM; ++d)
acc += b2f(s_q[q * DSTR + d]) * b2f(s_kG[i * DSTR + d]);
s_bscore[q * N + G] = acc / (float)cnt;
} else {
for (int j = 0; j < N; ++j) s_bscore[q * N + j] = -INFINITY;
}
}
__syncthreads();
// ---- top-8 selection per query ----
for (int i = tid; i < Q_TILE * ((N + 31) / 32); ++i) ((uint32_t*)s_sel)[i] = 0u;
for (int i = tid; i < N; i += THREADS) s_block_active[i] = 0;
__syncthreads();
if (tid < Q_TILE) {
const int q = tid;
const int t = t0 + q;
if (t < S) {
int top[8];
int ntop = 0;
for (int pass = 0; pass < TOP_N_BLOCKS; ++pass) {
float best = -INFINITY;
int bestj = -1;
for (int j = 0; j <= G; ++j) {
bool used = false;
for (int k = 0; k < ntop; ++k) if (top[k] == j) { used = true; break; }
if (used) continue;
float v = s_bscore[q * N + j];
if (v > best) { best = v; bestj = j; }
}
if (bestj >= 0 && best > -1e29f) top[ntop++] = bestj;
}
uint32_t* selrow = s_sel + q * ((N + 31) / 32);
for (int k = 0; k < ntop; ++k) {
int j = top[k];
selrow[j >> 5] |= (1u << (j & 31));
s_block_active[j] = 1;
}
// sliding window blocks
s_block_active[G] = 1;
if (G >= 1) s_block_active[G - 1] = 1;
}
}
__syncthreads();
// ---- attention state init ----
for (int i = tid; i < Q_TILE * DIM; i += THREADS) s_acc[i] = 0.f;
if (tid < Q_TILE) { s_m[tid] = -INFINITY; s_l[tid] = 0.f; }
__syncthreads();
// ---- sparse attention loop over key blocks ----
const int gid = lane / 4;
const int tl = lane % 4;
for (int j = 0; j <= G; ++j) {
__syncthreads();
if (!s_block_active[j]) continue;
const __nv_bfloat16* kblk;
if (j == G) {
kblk = s_kG;
} else {
for (int i = tid; i < BLOCK_K * DIM; i += THREADS) {
int c = i / DIM, d = i % DIM;
s_kblk[c * DSTR + d] = K[off + (size_t)(j * BLOCK_K + c) * DIM + d];
}
kblk = s_kblk;
}
for (int i = tid; i < BLOCK_K * DIM; i += THREADS) {
int c = i / DIM, d = i % DIM;
int key = j * BLOCK_K + c;
s_vblk[d * VSTR + c] = (key < S)
? V[off + (size_t)key * DIM + d]
: __float2bfloat16(0.f);
}
__syncthreads();
// ---- score tile: 8 warps, each handles mt=warp/2, 4 n-tiles ----
{
const int mt = warp / 2;
const int ntbase = (warp % 2) * 4;
float acc4[4][4];
#pragma unroll
for (int nt = 0; nt < 4; ++nt)
acc4[nt][0] = acc4[nt][1] = acc4[nt][2] = acc4[nt][3] = 0.f;
#pragma unroll
for (int ks = 0; ks < KSTEP; ++ks) {
uint32_t a0 = *((const uint32_t*)&s_q[(mt * 16 + gid) * DSTR + ks * 16 + 2 * tl]);
uint32_t a1 = *((const uint32_t*)&s_q[(mt * 16 + gid + 8) * DSTR + ks * 16 + 2 * tl]);
uint32_t a2 = *((const uint32_t*)&s_q[(mt * 16 + gid) * DSTR + ks * 16 + 2 * tl + 8]);
uint32_t a3 = *((const uint32_t*)&s_q[(mt * 16 + gid + 8) * DSTR + ks * 16 + 2 * tl + 8]);
#pragma unroll
for (int nt = 0; nt < 4; ++nt) {
const __nv_bfloat16* krow = &kblk[((ntbase + nt) * 8 + gid) * DSTR + ks * 16 + 2 * tl];
uint32_t b0 = *((const uint32_t*)krow);
uint32_t b1 = *((const uint32_t*)(krow + 8));
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"(acc4[nt][0]), "+f"(acc4[nt][1]), "+f"(acc4[nt][2]), "+f"(acc4[nt][3])
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
}
#pragma unroll
for (int nt = 0; nt < 4; ++nt) {
int n = (ntbase + nt) * 8;
s_stile[(mt * 16 + gid) * BSTR + n + 2 * tl] = acc4[nt][0];
s_stile[(mt * 16 + gid) * BSTR + n + 2 * tl + 1] = acc4[nt][1];
s_stile[(mt * 16 + gid + 8) * BSTR + n + 2 * tl] = acc4[nt][2];
s_stile[(mt * 16 + gid + 8) * BSTR + n + 2 * tl + 1] = acc4[nt][3];
}
}
__syncthreads();
// ---- online softmax (one thread per query) ----
if (tid < Q_TILE) {
const int q = tid;
const int t = t0 + q;
const bool real = (t < S);
const uint32_t* selrow = s_sel + q * ((N + 31) / 32);
auto selbit = [&](int jj) -> bool {
return (selrow[jj >> 5] >> (jj & 31)) & 1u;
};
float rowmax = -INFINITY;
#pragma unroll
for (int c = 0; c < BLOCK_K; ++c) {
if (!real) continue;
if (j * BLOCK_K + c > t) continue;
bool sel = false;
if (j == G && c <= q) sel = true;
else if (j == G - 1 && c >= q + 1) sel = true;
else sel = selbit(j);
if (!sel) continue;
float v = s_stile[q * BSTR + c] * scale;
if (v > rowmax) rowmax = v;
}
if (rowmax == -INFINITY) {
#pragma unroll
for (int c = 0; c < BLOCK_K; ++c) s_ptile[q * BSTR + c] = __float2bfloat16(0.f);
} else {
const float m_old = s_m[q];
const float m_new = fmaxf(m_old, rowmax);
const float rescale = __expf(m_old - m_new);
s_m[q] = m_new;
float lsum = 0.f;
#pragma unroll
for (int c = 0; c < BLOCK_K; ++c) {
float p = 0.f;
if (real && j * BLOCK_K + c <= t) {
bool sel = false;
if (j == G && c <= q) sel = true;
else if (j == G - 1 && c >= q + 1) sel = true;
else sel = selbit(j);
if (sel) p = __expf(s_stile[q * BSTR + c] * scale - m_new);
}
s_ptile[q * BSTR + c] = f2b(p);
lsum += p;
}
s_l[q] = s_l[q] * rescale + lsum;
if (rescale != 1.f) {
for (int d = 0; d < DIM; ++d) s_acc[q * DSTR + d] *= rescale;
}
}
}
__syncthreads();
// ---- P @ V accumulate into s_acc (8 warps) ----
{
const int mt = warp / 2;
const int ntbase = (warp % 2) * (DIM / 16); // 4 for D=64, 8 for D=128
constexpr int NT = DIM / 8;
const int HNT = NT / 2;
float acc4[HNT][4];
#pragma unroll
for (int nt = 0; nt < HNT; ++nt) {
int n = (ntbase + nt) * 8;
acc4[nt][0] = s_acc[(mt * 16 + gid) * DSTR + n + 2 * tl];
acc4[nt][1] = s_acc[(mt * 16 + gid) * DSTR + n + 2 * tl + 1];
acc4[nt][2] = s_acc[(mt * 16 + gid + 8) * DSTR + n + 2 * tl];
acc4[nt][3] = s_acc[(mt * 16 + gid + 8) * DSTR + n + 2 * tl + 1];
}
#pragma unroll
for (int ks = 0; ks < KPSTEP; ++ks) {
uint32_t a0 = *((const uint32_t*)&s_ptile[(mt * 16 + gid) * BSTR + ks * 16 + 2 * tl]);
uint32_t a1 = *((const uint32_t*)&s_ptile[(mt * 16 + gid + 8) * BSTR + ks * 16 + 2 * tl]);
uint32_t a2 = *((const uint32_t*)&s_ptile[(mt * 16 + gid) * BSTR + ks * 16 + 2 * tl + 8]);
uint32_t a3 = *((const uint32_t*)&s_ptile[(mt * 16 + gid + 8) * BSTR + ks * 16 + 2 * tl + 8]);
#pragma unroll
for (int nt = 0; nt < HNT; ++nt) {
const int dim = (ntbase + nt) * 8 + gid;
const __nv_bfloat16* vrow = &s_vblk[dim * VSTR + ks * 16 + 2 * tl];
uint32_t b0 = *((const uint32_t*)vrow);
uint32_t b1 = *((const uint32_t*)(vrow + 8));
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"(acc4[nt][0]), "+f"(acc4[nt][1]), "+f"(acc4[nt][2]), "+f"(acc4[nt][3])
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
}
#pragma unroll
for (int nt = 0; nt < HNT; ++nt) {
int n = (ntbase + nt) * 8;
s_acc[(mt * 16 + gid) * DSTR + n + 2 * tl] = acc4[nt][0];
s_acc[(mt * 16 + gid) * DSTR + n + 2 * tl + 1] = acc4[nt][1];
s_acc[(mt * 16 + gid + 8) * DSTR + n + 2 * tl] = acc4[nt][2];
s_acc[(mt * 16 + gid + 8) * DSTR + n + 2 * tl + 1] = acc4[nt][3];
}
}
}
// ---- write output ----
__syncthreads();
if (tid < Q_TILE) {
const int q = tid;
const int t = t0 + q;
if (t < S) {
const float l = s_l[q];
__nv_bfloat16* orow = O + off + (size_t)t * DIM;
for (int d = 0; d < DIM; ++d) orow[d] = f2b(s_acc[q * DSTR + d] / l);
}
}
}
// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
static int g_S = 0, g_D = 0, g_N = 0;
static int g_ntiles = 0, g_bh = 0;
static size_t g_smem = 0;
static float g_scale = 1.f;
void nsa_configure(int S, int D, int N, int ntiles, int bh, size_t smem, float scale) {
g_S = S; g_D = D; g_N = N; g_ntiles = ntiles; g_bh = bh; g_smem = smem; g_scale = scale;
}
void nsa_launch(const __nv_bfloat16* Q, const __nv_bfloat16* K, const __nv_bfloat16* V,
__nv_bfloat16* O, const float* KMEAN, cudaStream_t stream) {
dim3 grid(g_ntiles, g_bh);
size_t smem = g_smem;
if (g_D == 64) {
cudaFuncSetAttribute(nsa_kernel<64>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
nsa_kernel<64><<<grid, THREADS, smem, stream>>>(Q, K, V, O, KMEAN, g_S, g_N, g_scale);
} else {
cudaFuncSetAttribute(nsa_kernel<128>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
nsa_kernel<128><<<grid, THREADS, smem, stream>>>(Q, K, V, O, KMEAN, g_S, g_N, g_scale);
}
}
void kmean_launch(const __nv_bfloat16* K, float* KMEAN, int S, int D, int N, int bh, cudaStream_t stream) {
dim3 grid(bh, min(N, 256));
kmean_kernel<<<grid, THREADS, 0, stream>>>(K, KMEAN, S, D, N);
}
"""
_CPP_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>
void nsa_configure(int S, int D, int N, int ntiles, int bh, size_t smem, float scale);
void nsa_launch(const __nv_bfloat16* Q, const __nv_bfloat16* K, const __nv_bfloat16* V,
__nv_bfloat16* O, const float* KMEAN, cudaStream_t stream);
void kmean_launch(const __nv_bfloat16* K, float* KMEAN, int S, int D, int N, int bh, cudaStream_t stream);
torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "cuda tensors required");
TORCH_CHECK(q.scalar_type() == torch::kBFloat16, "bf16 required");
TORCH_CHECK(q.dim() == 4, "expected (B,H,S,D)");
auto qc = q.contiguous();
auto kc = k.contiguous();
auto vc = v.contiguous();
const int B = qc.size(0), H = qc.size(1), S = qc.size(2), D = qc.size(3);
const int N = (S + 63) / 64;
const int ntiles = (S + 63) / 64;
const int bh = B * H;
auto o = torch::empty_like(qc);
auto kmean = torch::empty({B, H, N, D}, qc.options().dtype(torch::kFloat32));
// shared size computation (must match kernel layout)
constexpr int Q_TILE = 64, BLOCK_K = 64, THREADS = 256;
const int DSTR = D + (D == 64 ? 6 : 2), BSTR = BLOCK_K + (D == 64 ? 6 : 2);
size_t smem = 0;
smem += (size_t)N * DSTR * 4; // s_kmean fp32
smem += Q_TILE * DSTR * 2; // s_q
smem += Q_TILE * N * 4; // s_bscore fp32
smem += Q_TILE * DSTR * 4; // s_acc
smem += 2 * Q_TILE * 4; // s_m, s_l
smem += BLOCK_K * DSTR * 2; // s_kblk
smem += D * (BLOCK_K + (D == 64 ? 8 : 2)) * 2; // s_vblk (transposed)
smem += Q_TILE * BSTR * 2; // s_ptile
smem += Q_TILE * BSTR * 4; // s_stile
smem += Q_TILE * ((N + 31) / 32) * 4; // s_sel
smem += BLOCK_K * DSTR * 2; // s_kG
smem += N; // s_block_active
smem = (smem + 127) & ~(size_t)127;
// Force 1 block/SM: co-resident blocks on the same SM trigger a timing race
// in the mma/softmax pipeline on some shapes, so keep the shared request
// above half the SM's shared capacity.
if (smem < 118 * 1024) smem = 118 * 1024;
float scale = (float)(1.0 / std::sqrt((double)D));
auto stream = at::cuda::getCurrentCUDAStream();
nsa_configure(S, D, N, ntiles, bh, smem, scale);
kmean_launch((const __nv_bfloat16*)kc.data_ptr(), (float*)kmean.data_ptr(), S, D, N, bh, stream);
nsa_launch((const __nv_bfloat16*)qc.data_ptr(), (const __nv_bfloat16*)kc.data_ptr(),
(const __nv_bfloat16*)vc.data_ptr(), (__nv_bfloat16*)o.data_ptr(),
(const float*)kmean.data_ptr(), stream);
return o;
}
"""
_load_attempted = False
_MOD = None
def _load():
global _load_attempted, _MOD
if _load_attempted:
return _MOD
_load_attempted = True
# Make sure a real nvcc is findable even if the outer environment shadows it
# with a broken wrapper (the run-directory bin/ shims).
import os
_candidates = ["/usr/local/cuda-12.8/bin", "/usr/local/cuda/bin", "/usr/local/cuda-12/bin"]
_need = True
for _d in _candidates:
_nvcc = os.path.join(_d, "nvcc")
if os.path.exists(_nvcc):
os.environ.setdefault("CUDA_HOME", os.path.dirname(_d))
if not os.environ["PATH"].startswith(_d + os.pathsep):
os.environ["PATH"] = _d + os.pathsep + os.environ["PATH"]
_need = False
break
if _need:
raise RuntimeError("no real nvcc found; cannot build CUDA kernel")
# The cu12x torch wheels ship CUDA-library headers (cublas/cusparse/...)
# inside the nvidia pip packages; the system CUDA install here is a
# reduced toolkit that lacks them. Add those header dirs explicitly.
import glob as _glob
_extra = []
_nv_root = os.path.join(os.path.dirname(os.path.dirname(torch.__file__)), "nvidia")
if os.path.isdir(_nv_root):
for _pkg in sorted(os.listdir(_nv_root)):
_inc = os.path.join(_nv_root, _pkg, "include")
if os.path.isdir(_inc):
_extra.append(_inc)
_MOD = load_inline(
name="nsa_kernel",
cpp_sources=_CPP_SRC,
cuda_sources=_CUDA_SRC,
functions=["nsa_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
extra_include_paths=_extra,
verbose=False,
)
return _MOD
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:
mod = _load()
return mod.nsa_forward(q, k, v)
20260802_204049_or-fable_deepseek_deepseek-v4-flash-0731_02_deepseek_nsa