KernelBench cuda · RTX PRO 6000
DeepSeek NSA Claude Fable 5
72.7%geomean peak fraction across shapes
manually audited: clean
Genuine, well-engineered NSA kernel. The outlier 0.7192 (vs kinetic 0.4246 board ceiling and fable's own truncated-session 0.2934) is a real algorithmic win, not a hack: the agent inverted the sparse attention from query-centric gather to BLOCK-CENTRIC processing (CTA per KV block stages K/V in smem once, serves all subscriber queries via histogram/scan/ scatter user lists), replaced the O(S^2) importance pass with a tiny Q @ Ksum^T GEMM (mean over block = dot with block K-sum), and moved the block pass onto tensor cores (mma.m16n8k16, 16 users/warp, XOR-swizzled smem, one-shot softmax per block, online-softmax merge kernel). Per-shape fractions >1.0 on long-context shapes (1.64 at S=8192) are the expected dense-equivalent-FLOPs artifact for a fast correct sparse kernel: the deck charges 4*B*H*S^2*D dense FLOPs while NSA touches <=576 keys/query (~14x less work at S=8192); implied REAL throughput is only ~29 TFLOPS of sparse work -- physically modest for this card. The 0.2934 -> 0.7192 jump is the same design family finished: the 07-19 session (api_error at 96 min) had already landed the block-centric inversion but with a SIMT warp-per-subscriber block pass; this 2.3 h session added the tensor-core block pass, one-global-atomic-per-CTA selection, stripe splitting for occupancy, and a CUDA-graph capture of the whole pipeline, iterating through 20+ ncu profile rounds (0.1222 -> 0.3858 -> 0.7296 -> 0.7349 in-session; 0.7192 at grading).
harnessor-fableagent session2h 19mtotal wall2h 20mcheck50sbenchmark2soutput tokens—gpu-lock wait0sgpu-lock held52sregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
1×16×2048×640.083 ms41.4%207 TFLOPS · 41% of 500 TF bf16 peak · also 0.20 TB/s (11% of HBM)
1×16×4127×640.146 ms95.3%477 TFLOPS · 95% of 500 TF bf16 peak · also 0.23 TB/s (13% of HBM)
1×8×8192×640.162 ms169.3%846 TFLOPS · 169% of 500 TF bf16 peak · also 0.21 TB/s (11% of HBM)
1×8×8191×1280.360 ms152.7%763 TFLOPS · 153% of 500 TF bf16 peak · also 0.19 TB/s (10% of HBM)
4×8×1024×640.076 ms22.7%114 TFLOPS · 23% of 500 TF bf16 peak · also 0.22 TB/s (12% of HBM)
2×8×3000×640.111 ms66.4%332 TFLOPS · 66% of 500 TF bf16 peak · also 0.22 TB/s (12% of HBM)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(41.4% · 95.3% · 169.3% · 152.7% · 22.7% · 66.4%) = 73.2% · published 72.7% (lower of repeated isolated re-benchmark passes)
Kernel source (redacted)
"""DeepSeek NSA-inspired sparse attention — CUDA solution.
Semantics (bench-simplified NSA, matches reference.nsa_attend):
per query t: block importance = mean_j(q·k_j)/sqrt(D) over causal keys of each
64-block; top-8 blocks by importance (ties -> higher block id) union sliding
window of last 64 tokens (causal); softmax attention over selected keys.
Key algorithmic ideas:
* mean_j(q·k_j) = q·mean_j(k_j): full-block scoring collapses to a tiny fp32
GEMM Q @ Ksum^T instead of an S×S pass; the partial "current" block comes
from a within-block inclusive cumsum of K.
* Attention over a union of blocks = softmax-merge of independent per-block
partials. So instead of a per-query random gather (which is L2 latency
bound), invert the selection: one CTA per KV block stages the block in
shared memory ONCE and computes partials for every query that selected it
(user lists built by histogram+scan+scatter). A sliding-window pass
handles the window remainder, and a merge kernel combines each query's
partials with online-softmax algebra.
Pipeline (all CUDA, no torch ops in the hot path; whole forward is CUDA-graph
captured when the benchmark reuses the same input buffers):
1. nsa_prep per (head, block): K column prefix sums -> block K sums +
partial current-block importance (4 row-chunks in parallel).
2. nsa_imp_gemm imp = coef * Q @ Ksum^T (fp32 SIMT GEMM, causal tiles only,
dense region t<512 skipped — it never reads imp).
3. nsa_select 2 queries/warp: top-8 blocks (reference tie-break: higher
block id wins ties) + window-coverage flags; per-block user
counts aggregated in shared memory (one global atomic per
CTA per block); window remainders emitted as slots 8/9.
4. nsa_scan exclusive scan of the per-(head,block) user counts.
5. nsa_scatter deterministic invert to per-block user lists (no atomics).
6. nsa_block tensor-core pass, CTA per (head, block, stripe): K/V staged
once in XOR-swizzled smem; each warp handles 16 users:
S = Q@K^T (mma.m16n8k16), causal/window mask, one-shot
softmax (no online rescale needed within one block),
P@V via C->A fragment reuse; acc partial stored bf16
(staged through smem for coalesced writes), m/l fp32.
7. nsa_merge warp per query: softmax-merge the <=10 partials, write o.
"""
from __future__ import annotations
import math
import os
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
_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#define FULL_MASK 0xffffffffu
#define SLOTS 10 // 8 selected blocks + window parts in blocks cb-1 / cb
__device__ __forceinline__ void bf16x8_to_f32(const uint4 u, float f[8]) {
const __nv_bfloat162* p = reinterpret_cast<const __nv_bfloat162*>(&u);
#pragma unroll
for (int i = 0; i < 4; ++i) {
float2 t = __bfloat1622float2(p[i]);
f[2 * i] = t.x;
f[2 * i + 1] = t.y;
}
}
// ---------------------------------------------------------------------------
// Kernel 1: per (head, block) — column prefix sums of K.
// ---------------------------------------------------------------------------
template <int D>
__global__ void __launch_bounds__(D * 4) nsa_prep(
const __nv_bfloat16* __restrict__ q,
const __nv_bfloat16* __restrict__ k,
float* __restrict__ ksum,
float* __restrict__ cur_imp,
int S, int NB, float scale) {
// 4 chunks of 16 rows in parallel; chunk offsets fixed up via an
// exclusive prefix over chunk sums (applied lazily in phase 2).
__shared__ float P[64][D];
__shared__ float coff[4][D];
int bi = blockIdx.x % NB;
int bh = blockIdx.x / NB;
int s0 = bi << 6;
int len = S - s0;
if (len > 64) len = 64;
int tid = threadIdx.x;
int ch = tid / D, d = tid % D;
const __nv_bfloat16* kb = k + ((long long)bh * S + s0) * D;
{
float run = 0.f;
int j0 = ch * 16;
int j1 = min(j0 + 16, len);
for (int j = j0; j < j1; ++j) {
run += __bfloat162float(kb[j * D + d]);
P[j][d] = run;
}
coff[ch][d] = run;
}
__syncthreads();
if (ch == 0) { // exclusive prefix of chunk sums; total -> ksum
float s = 0.f;
#pragma unroll
for (int c = 0; c < 4; ++c) {
float t = coff[c][d];
coff[c][d] = s;
s += t;
}
ksum[((long long)bh * NB + bi) * D + d] = s;
}
if (bi < 8) return; // cur_imp only read for t >= 512
__syncthreads();
// warp per row: cur_imp[t] = q_t . (P[t] + coff[t/16]) * scale / (t%64+1)
int warp = tid >> 5, lane = tid & 31;
constexpr int NW = D * 4 / 32;
const __nv_bfloat16* qb = q + ((long long)bh * S + s0) * D;
for (int tl = warp; tl < len; tl += NW) {
float dsum = 0.f;
int c = tl >> 4;
#pragma unroll
for (int dd = lane, x = 0; x < D / 32; dd += 32, ++x)
dsum += __bfloat162float(qb[tl * D + dd]) * (P[tl][dd] + coff[c][dd]);
#pragma unroll
for (int o = 16; o > 0; o >>= 1) dsum += __shfl_xor_sync(FULL_MASK, dsum, o);
if (lane == 0)
cur_imp[(long long)bh * S + s0 + tl] = dsum * scale / (float)(tl + 1);
}
}
// ---------------------------------------------------------------------------
// Kernel 2: importance GEMM imp[bh, t, bi] = coef * q_t . ksum[bi]
// ---------------------------------------------------------------------------
template <int D>
__global__ void __launch_bounds__(256) nsa_imp_gemm(
const __nv_bfloat16* __restrict__ q,
const float* __restrict__ ksum,
float* __restrict__ imp,
int S, int NB, int NBp, float coef) {
constexpr int QT = 64, BT = 32;
int qt0 = blockIdx.x * QT;
int bt0 = blockIdx.y * BT;
int bh = blockIdx.z;
if ((bt0 << 6) > qt0 + QT - 1) return; // wholly non-causal
if (qt0 + QT - 1 < 512) return; // dense region never reads imp
extern __shared__ __align__(16) float sm[];
float* Qs = sm; // [QT][D+2]
float* Bs = sm + QT * (D + 2); // [BT][D+1]
int tid = threadIdx.x;
const __nv_bfloat16* qb = q + ((long long)bh * S + qt0) * D;
for (int i = tid; i < QT * D / 2; i += 256) {
int row = i / (D / 2), col = i - row * (D / 2);
float2 f = make_float2(0.f, 0.f);
if (qt0 + row < S) {
__nv_bfloat162 h =
*reinterpret_cast<const __nv_bfloat162*>(qb + row * D + col * 2);
f = __bfloat1622float2(h);
}
Qs[row * (D + 2) + col * 2] = f.x;
Qs[row * (D + 2) + col * 2 + 1] = f.y;
}
const float* bb = ksum + ((long long)bh * NB + bt0) * D;
for (int i = tid; i < BT * D; i += 256) {
int row = i / D, col = i - row * D;
Bs[row * (D + 1) + col] = (bt0 + row < NB) ? bb[row * D + col] : 0.f;
}
__syncthreads();
int qa = (tid >> 3) * 2, bx = (tid & 7) * 4;
float a00 = 0.f, a01 = 0.f, a02 = 0.f, a03 = 0.f;
float a10 = 0.f, a11 = 0.f, a12 = 0.f, a13 = 0.f;
const float* q0p = Qs + qa * (D + 2);
const float* q1p = q0p + (D + 2);
const float* b0 = Bs + bx * (D + 1);
const float* b1 = b0 + (D + 1);
const float* b2 = b1 + (D + 1);
const float* b3 = b2 + (D + 1);
#pragma unroll 8
for (int d = 0; d < D; ++d) {
float q0 = q0p[d], q1 = q1p[d];
float v0 = b0[d], v1 = b1[d], v2 = b2[d], v3 = b3[d];
a00 = fmaf(q0, v0, a00); a01 = fmaf(q0, v1, a01);
a02 = fmaf(q0, v2, a02); a03 = fmaf(q0, v3, a03);
a10 = fmaf(q1, v0, a10); a11 = fmaf(q1, v1, a11);
a12 = fmaf(q1, v2, a12); a13 = fmaf(q1, v3, a13);
}
float4 o0 = make_float4(a00 * coef, a01 * coef, a02 * coef, a03 * coef);
float4 o1 = make_float4(a10 * coef, a11 * coef, a12 * coef, a13 * coef);
int t0 = qt0 + qa;
if (t0 < S)
*reinterpret_cast<float4*>(imp + ((long long)bh * S + t0) * NBp + bt0 + bx) = o0;
if (t0 + 1 < S)
*reinterpret_cast<float4*>(imp + ((long long)bh * S + t0 + 1) * NBp + bt0 + bx) = o1;
}
// ---------------------------------------------------------------------------
// Kernel 3: top-8 selection, 2 queries per warp (16-lane groups), 32 queries
// per CTA. Per-block user counts are aggregated in shared memory first and
// published with ONE global atomic per (CTA, block) — same-address global
// atomic contention was the previous bottleneck.
// sel8[q][slot], seq8[q][slot] (slot 8/9 = window parts), flags[q]
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(512) nsa_select(
const float* __restrict__ imp,
const float* __restrict__ cur_imp,
unsigned char* __restrict__ sel8,
short* __restrict__ seq8,
unsigned char* __restrict__ flags,
int* __restrict__ hist,
int S, int NB, int NBp) {
__shared__ int lhist[128];
__shared__ int lbase[128];
__shared__ unsigned char ssel[32][SLOTS];
__shared__ short sseq[32][SLOTS];
__shared__ unsigned char sflag[32];
__shared__ unsigned short svalid[32]; // bitmask of live slots
int tid = threadIdx.x;
int bh = blockIdx.y;
int t0 = blockIdx.x * 32;
for (int i = tid; i < 128; i += 512) lhist[i] = 0;
if (tid < 32) svalid[tid] = 0;
__syncthreads();
int lane = tid & 31;
int half = lane >> 4, l16 = lane & 15;
int ql = (tid >> 5) * 2 + half; // query index within CTA
int t = t0 + ql;
long long bhS = (long long)bh * S;
long long qi = bhS + t;
unsigned gmask = 0xffffu << (half << 4);
if (t < S) {
int cb = t >> 6;
if (cb < 8) {
if (l16 <= cb) {
ssel[ql][l16] = (unsigned char)l16;
int sq = atomicAdd(&lhist[l16], 1);
sseq[ql][l16] = (short)sq;
}
if (l16 == 0) {
sflag[ql] = 3;
svalid[ql] = (unsigned short)((1u << (cb + 1)) - 1u);
}
} else {
float ci = cur_imp[qi];
const float4* irow4 = reinterpret_cast<const float4*>(imp + qi * NBp);
float cv[8];
{
float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f);
float4 f0 = (l16 * 4 < NBp) ? irow4[l16] : z4;
float4 f1 = (64 + l16 * 4 < NBp) ? irow4[16 + l16] : z4;
cv[0] = f0.x; cv[1] = f0.y; cv[2] = f0.z; cv[3] = f0.w;
cv[4] = f1.x; cv[5] = f1.y; cv[6] = f1.z; cv[7] = f1.w;
#pragma unroll
for (int r = 0; r < 8; ++r) {
int bi = (r < 4) ? (l16 * 4 + r) : (64 + l16 * 4 + r - 4);
if (bi > cb) cv[r] = -INFINITY;
else if (bi == cb) cv[r] = ci;
}
}
int w0 = t - 63;
int wb0 = w0 >> 6;
bool have_w0 = false, have_cb = false;
#pragma unroll
for (int i = 0; i < 8; ++i) {
float bv = cv[0];
int bb = l16 * 4;
#pragma unroll
for (int r = 1; r < 8; ++r) {
int bi = (r < 4) ? (l16 * 4 + r) : (64 + l16 * 4 + r - 4);
if (cv[r] > bv || (cv[r] == bv && bi > bb)) { bv = cv[r]; bb = bi; }
}
#pragma unroll
for (int off = 8; off > 0; off >>= 1) {
float ov = __shfl_xor_sync(gmask, bv, off);
int obb = __shfl_xor_sync(gmask, bb, off);
if (ov > bv || (ov == bv && obb > bb)) { bv = ov; bb = obb; }
}
if (l16 == i) {
ssel[ql][i] = (unsigned char)bb;
int sq = atomicAdd(&lhist[bb], 1);
sseq[ql][i] = (short)sq;
}
have_w0 |= (bb == wb0);
have_cb |= (bb == cb);
int owner = (bb < 64) ? (bb >> 2) : ((bb - 64) >> 2);
if (owner == l16) {
int r = (bb < 64) ? (bb & 3) : (4 + (bb & 3));
#pragma unroll
for (int x = 0; x < 8; ++x)
if (x == r) cv[x] = -INFINITY;
}
}
bool lo_ex = (wb0 < cb) && !have_w0;
bool hi_ex = !have_cb;
if (l16 == 0) {
sflag[ql] = (have_w0 ? 1 : 0) | (have_cb ? 2 : 0);
svalid[ql] = (unsigned short)(0xffu | (lo_ex ? 0x100u : 0u) |
(hi_ex ? 0x200u : 0u));
}
if (l16 == 8 && lo_ex) {
ssel[ql][8] = (unsigned char)wb0;
int sq = atomicAdd(&lhist[wb0], 1);
sseq[ql][8] = (short)sq;
}
if (l16 == 9 && hi_ex) {
ssel[ql][9] = (unsigned char)cb;
int sq = atomicAdd(&lhist[cb], 1);
sseq[ql][9] = (short)sq;
}
}
}
__syncthreads();
// publish CTA counts with one global atomic per block id
for (int i = tid; i < NB && i < 128; i += 512) {
int cnt = lhist[i];
lbase[i] = cnt ? atomicAdd(&hist[bh * NB + i], cnt) : 0;
}
__syncthreads();
// write out with CTA base applied
for (int e = tid; e < 32 * SLOTS; e += 512) {
int ql2 = e / SLOTS, slot = e - ql2 * SLOTS;
int t2 = t0 + ql2;
if (t2 >= S) continue;
long long q2 = bhS + t2;
if (slot == 0) flags[q2] = sflag[ql2];
if (svalid[ql2] & (1u << slot)) {
int bi = ssel[ql2][slot];
sel8[q2 * SLOTS + slot] = (unsigned char)bi;
seq8[q2 * SLOTS + slot] = (short)(lbase[bi] + sseq[ql2][slot]);
}
}
}
// ---------------------------------------------------------------------------
// Kernel 4: exclusive scan of hist (n <= 2048), one CTA.
// off[0..n] = exclusive scan (off[n] = total), cur = copy of off[0..n-1]
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(1024) nsa_scan(
const int* __restrict__ hist, int* __restrict__ off, int n) {
__shared__ int smA[2048], smB[2048];
int tid = threadIdx.x;
for (int i = tid; i < 2048; i += 1024) smA[i] = (i < n) ? hist[i] : 0;
__syncthreads();
int* a = smA;
int* b = smB;
for (int d = 1; d < 2048; d <<= 1) {
for (int i = tid; i < 2048; i += 1024)
b[i] = a[i] + ((i >= d) ? a[i - d] : 0);
__syncthreads();
int* tmp = a; a = b; b = tmp;
}
// a = inclusive scan
for (int i = tid; i < n; i += 1024) off[i] = (i == 0) ? 0 : a[i - 1];
if (tid == 0) off[n] = a[n - 1];
}
// ---------------------------------------------------------------------------
// Kernel 5: scatter (query, slot) pairs into per-block user lists.
// deterministic: position = off[block] + seq (no atomics)
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(256) nsa_scatter(
const unsigned char* __restrict__ sel8,
const short* __restrict__ seq8,
const unsigned char* __restrict__ flags,
const int* __restrict__ off,
int* __restrict__ ulist,
int BH, int S, int NB) {
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
long long total = (long long)BH * S * SLOTS;
if (idx >= total) return;
long long qi = idx / SLOTS;
int slot = (int)(idx - qi * SLOTS);
int t = (int)(qi % S);
int cb = t >> 6;
if (slot < 8) {
int count = cb + 1 < 8 ? cb + 1 : 8;
if (slot >= count) return;
} else {
if (cb < 8) return;
int f = flags[qi];
int wb0 = (t - 63) >> 6;
if (slot == 8 && (wb0 >= cb || (f & 1))) return;
if (slot == 9 && (f & 2)) return;
}
int bh = (int)(qi / S);
int bi = sel8[idx];
int pos = off[bh * NB + bi] + seq8[idx];
ulist[pos] = (t << 4) | slot;
}
// ---------------------------------------------------------------------------
// Kernel 6: block pass (tensor cores) — CTA per (head, block); K/V staged in
// smem once; each warp computes 16 users' partials per pass:
// S = Q(16xD) @ K^T(Dx64) via mma.m16n8k16, causal mask, one-shot softmax
// (a block partial needs no online rescale), P @ V via C->A fragment reuse.
// acc partials stored bf16 (staged through Qsm for coalesced writes), m/l fp32.
// ---------------------------------------------------------------------------
__device__ __forceinline__ void mma16816(float c[4], unsigned a0, unsigned a1,
unsigned a2, unsigned a3,
unsigned b0, unsigned 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"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
template <int D, int NW>
__global__ void __launch_bounds__(NW * 32) nsa_block(
const __nv_bfloat16* __restrict__ q,
const __nv_bfloat16* __restrict__ k,
const __nv_bfloat16* __restrict__ v,
const int* __restrict__ off,
const int* __restrict__ ulist,
float* __restrict__ pml,
__nv_bfloat16* __restrict__ pacc,
int S, int NB, int split, float pres) {
constexpr int KT = D / 16; // k-tiles for QK^T
extern __shared__ __align__(16) char bsm[];
__nv_bfloat16* Kt = reinterpret_cast<__nv_bfloat16*>(bsm);
__nv_bfloat16* Vt = Kt + 64 * D;
__nv_bfloat16* Qsm = Vt + 64 * D; // NW warps x 16 x D
int sp = blockIdx.x % split;
int gb = blockIdx.x / split;
int bi = gb % NB;
int bh = gb / NB;
int base = bh * NB + bi;
int u0 = off[base], u1 = off[base + 1];
if (u0 == u1) return;
if (sp * NW * 16 >= u1 - u0) return; // this stripe has no tiles
int s0 = bi << 6;
int len = S - s0;
if (len > 64) len = 64;
{
const uint4* gk = reinterpret_cast<const uint4*>(k + ((long long)bh * S + s0) * D);
const uint4* gv = reinterpret_cast<const uint4*>(v + ((long long)bh * S + s0) * D);
uint4* sk = reinterpret_cast<uint4*>(Kt);
uint4* sv = reinterpret_cast<uint4*>(Vt);
constexpr int CPR = D / 8;
int nvec = len * CPR;
uint4 z = make_uint4(0, 0, 0, 0);
for (int i = threadIdx.x; i < 64 * CPR; i += NW * 32) {
int row = i / CPR, c = i - row * CPR;
int j = row * CPR + (c ^ (row & 7)); // XOR swizzle vs ldmatrix
sk[j] = (i < nvec) ? gk[i] : z;
sv[j] = (i < nvec) ? gv[i] : z;
}
}
__syncthreads();
int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
__nv_bfloat16* Qw = Qsm + warp * 16 * D;
long long bhS = (long long)bh * S;
int nusers = u1 - u0;
int ntiles = (nusers + 15) >> 4;
int quad = lane >> 2, qsub = lane & 3; // C-frag row group / col pair
for (int tile = sp * NW + warp; tile < ntiles; tile += NW * split) {
int ubase = u0 + (tile << 4);
int nact = u1 - ubase;
if (nact > 16) nact = 16;
// lane r < 16 owns user row r
int e_r = ulist[ubase + ((lane < nact) ? lane : nact - 1)];
// gather Q rows into Qw (16 x D)
{
constexpr int CPR = D / 8; // 16B chunks per row
#pragma unroll
for (int x = 0; x < 16 * CPR / 32; ++x) {
int c = lane + 32 * x;
int row = c / CPR, ch = c - row * CPR;
int er = __shfl_sync(FULL_MASK, e_r, row);
int tr = er >> 4;
*reinterpret_cast<uint4*>(Qw + row * D + ((ch ^ (row & 7)) * 8)) =
*reinterpret_cast<const uint4*>(q + (bhS + tr) * D + ch * 8);
}
}
__syncwarp();
// A-frags for Q (KT k-tiles)
unsigned A[KT][4];
{
int row = lane % 16;
unsigned abase = (unsigned)__cvta_generic_to_shared(Qw + row * D);
#pragma unroll
for (int kt = 0; kt < KT; ++kt) {
int ch = (kt * 2 + lane / 16) ^ (row & 7);
asm volatile(
"ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(A[kt][0]), "=r"(A[kt][1]), "=r"(A[kt][2]), "=r"(A[kt][3])
: "r"(abase + ch * 16));
}
}
// scores C[8][4] = Q @ K^T (16 x 64)
float C[8][4];
#pragma unroll
for (int n = 0; n < 8; ++n)
#pragma unroll
for (int x = 0; x < 4; ++x) C[n][x] = 0.f;
{
int krow = lane % 16;
#pragma unroll
for (int np = 0; np < 4; ++np) { // key 16-pairs
unsigned kbase = (unsigned)__cvta_generic_to_shared(
Kt + (np * 16 + krow) * D);
#pragma unroll
for (int kt = 0; kt < KT; ++kt) {
int ch = (kt * 2 + lane / 16) ^ (krow & 7);
unsigned r0, r1, r2, r3;
asm volatile(
"ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
: "r"(kbase + ch * 16));
mma16816(C[2 * np], A[kt][0], A[kt][1], A[kt][2], A[kt][3], r0, r2);
mma16816(C[2 * np + 1], A[kt][0], A[kt][1], A[kt][2], A[kt][3], r1, r3);
}
}
}
// mask (causal + block length) & scale into base-2 domain
int e_a = __shfl_sync(FULL_MASK, e_r, quad);
int e_b = __shfl_sync(FULL_MASK, e_r, quad + 8);
int t_a = e_a >> 4, t_b = e_b >> 4;
int lim_a = t_a - s0; // lo <= j <= lim valid
int lim_b = t_b - s0;
if (lim_a > len - 1) lim_a = len - 1;
if (lim_b > len - 1) lim_b = len - 1;
int lo_a = ((e_a & 15) >= 8) ? max(t_a - 63 - s0, 0) : 0;
int lo_b = ((e_b & 15) >= 8) ? max(t_b - 63 - s0, 0) : 0;
#pragma unroll
for (int n = 0; n < 8; ++n) {
int j0 = n * 8 + qsub * 2;
C[n][0] = (j0 >= lo_a && j0 <= lim_a) ? C[n][0] * pres : -1e30f;
C[n][1] = (j0 + 1 >= lo_a && j0 + 1 <= lim_a) ? C[n][1] * pres : -1e30f;
C[n][2] = (j0 >= lo_b && j0 <= lim_b) ? C[n][2] * pres : -1e30f;
C[n][3] = (j0 + 1 >= lo_b && j0 + 1 <= lim_b) ? C[n][3] * pres : -1e30f;
}
// row max / exp / row sum (one-shot softmax; no online rescale needed)
float ma = -1e30f, mb = -1e30f;
#pragma unroll
for (int n = 0; n < 8; ++n) {
ma = fmaxf(ma, fmaxf(C[n][0], C[n][1]));
mb = fmaxf(mb, fmaxf(C[n][2], C[n][3]));
}
#pragma unroll
for (int o = 1; o <= 2; o <<= 1) {
ma = fmaxf(ma, __shfl_xor_sync(FULL_MASK, ma, o));
mb = fmaxf(mb, __shfl_xor_sync(FULL_MASK, mb, o));
}
float la = 0.f, lb = 0.f;
unsigned P[4][4]; // P a-frags: 4 key-16 tiles x 4 regs
#pragma unroll
for (int np = 0; np < 4; ++np) {
float p00 = exp2f(C[2 * np][0] - ma);
float p01 = exp2f(C[2 * np][1] - ma);
float p02 = exp2f(C[2 * np][2] - mb);
float p03 = exp2f(C[2 * np][3] - mb);
float p10 = exp2f(C[2 * np + 1][0] - ma);
float p11 = exp2f(C[2 * np + 1][1] - ma);
float p12 = exp2f(C[2 * np + 1][2] - mb);
float p13 = exp2f(C[2 * np + 1][3] - mb);
la += p00 + p01 + p10 + p11;
lb += p02 + p03 + p12 + p13;
__nv_bfloat162 x0 = __floats2bfloat162_rn(p00, p01);
__nv_bfloat162 x1 = __floats2bfloat162_rn(p02, p03);
__nv_bfloat162 x2 = __floats2bfloat162_rn(p10, p11);
__nv_bfloat162 x3 = __floats2bfloat162_rn(p12, p13);
P[np][0] = *reinterpret_cast<unsigned*>(&x0);
P[np][1] = *reinterpret_cast<unsigned*>(&x1);
P[np][2] = *reinterpret_cast<unsigned*>(&x2);
P[np][3] = *reinterpret_cast<unsigned*>(&x3);
}
#pragma unroll
for (int o = 1; o <= 2; o <<= 1) {
la += __shfl_xor_sync(FULL_MASK, la, o);
lb += __shfl_xor_sync(FULL_MASK, lb, o);
}
// O = P @ V (16 x D), in 64-dim halves to bound registers;
// acc (bf16) staged into Qw for a coalesced global write
__syncwarp();
int vrow16 = lane % 16;
#pragma unroll
for (int half = 0; half < D / 64; ++half) {
float O[8][4];
#pragma unroll
for (int n = 0; n < 8; ++n)
#pragma unroll
for (int x = 0; x < 4; ++x) O[n][x] = 0.f;
#pragma unroll
for (int kt = 0; kt < 4; ++kt) { // key 16-tiles
#pragma unroll
for (int np = 0; np < 4; ++np) { // dim 16-pairs within half
int vrow = kt * 16 + vrow16;
int ch = (half * 8 + np * 2 + lane / 16) ^ (vrow & 7);
unsigned vad = (unsigned)__cvta_generic_to_shared(Vt + vrow * D) + ch * 16;
unsigned v0, v1, v2, v3;
asm volatile(
"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(v0), "=r"(v1), "=r"(v2), "=r"(v3)
: "r"(vad));
mma16816(O[2 * np], P[kt][0], P[kt][1], P[kt][2], P[kt][3], v0, v1);
mma16816(O[2 * np + 1], P[kt][0], P[kt][1], P[kt][2], P[kt][3], v2, v3);
}
}
#pragma unroll
for (int n = 0; n < 8; ++n) {
int ch = half * 8 + n;
int ja = ((ch ^ (quad & 7)) * 8) + qsub * 2;
__nv_bfloat162 xa = __floats2bfloat162_rn(O[n][0], O[n][1]);
__nv_bfloat162 xb = __floats2bfloat162_rn(O[n][2], O[n][3]);
*reinterpret_cast<__nv_bfloat162*>(Qw + quad * D + ja) = xa;
*reinterpret_cast<__nv_bfloat162*>(Qw + (quad + 8) * D + ja) = xb;
}
}
if (qsub == 0) {
if (quad < nact) {
long long pi = (bhS + (e_a >> 4)) * SLOTS + (e_a & 15);
*reinterpret_cast<float2*>(pml + pi * 2) = make_float2(ma, la);
}
if (quad + 8 < nact) {
long long pi = (bhS + (e_b >> 4)) * SLOTS + (e_b & 15);
*reinterpret_cast<float2*>(pml + pi * 2) = make_float2(mb, lb);
}
}
__syncwarp();
{
constexpr int CPR = D / 8;
#pragma unroll
for (int x = 0; x < 16 * CPR / 32; ++x) {
int c = lane + 32 * x;
int row = c / CPR, ch = c - row * CPR;
int er = __shfl_sync(FULL_MASK, e_r, row);
if (row < nact) {
long long pi = (bhS + (er >> 4)) * SLOTS + (er & 15);
*reinterpret_cast<uint4*>(pacc + pi * D + ch * 8) =
*reinterpret_cast<const uint4*>(Qw + row * D + ((ch ^ (row & 7)) * 8));
}
}
}
__syncwarp();
}
}
// ---------------------------------------------------------------------------
// Kernel 8: merge partials per query and write bf16 output. Warp per query.
// ---------------------------------------------------------------------------
template <int D>
__global__ void __launch_bounds__(256) nsa_merge(
const float* __restrict__ pml,
const __nv_bfloat16* __restrict__ pacc,
const unsigned char* __restrict__ flags,
__nv_bfloat16* __restrict__ o,
int BH, int S) {
constexpr int EL = D / 32; // dims per lane
long long warp = (long long)(blockIdx.x * blockDim.x + threadIdx.x) >> 5;
int lane = threadIdx.x & 31;
long long nq = (long long)BH * S;
if (warp >= nq) return;
int t = (int)(warp % S);
int cb = t >> 6;
int count = cb + 1 < 8 ? cb + 1 : 8;
int nslot = count;
int slist89 = 0; // bit0: slot8, bit1: slot9
if (cb >= 8) {
int f = flags[warp];
int wb0 = (t - 63) >> 6;
if (wb0 < cb && !(f & 1)) slist89 |= 1;
if (!(f & 2)) slist89 |= 2;
}
int total = count + ((slist89 & 1) ? 1 : 0) + ((slist89 & 2) ? 1 : 0);
float m = -1e30f, l = 0.f;
float a[EL];
#pragma unroll
for (int x = 0; x < EL; ++x) a[x] = 0.f;
for (int s = 0; s < total; ++s) {
int sl = s;
if (s >= count) sl = (s == count && (slist89 & 1)) ? 8 : 9;
long long pi = warp * SLOTS + sl;
float2 ml = *reinterpret_cast<const float2*>(pml + pi * 2);
float pm = ml.x;
float pl = ml.y;
float pa[EL];
{
const __nv_bfloat16* pp = pacc + pi * D + lane * EL;
#pragma unroll
for (int x = 0; x < EL; x += 2) {
float2 f = __bfloat1622float2(
*reinterpret_cast<const __nv_bfloat162*>(pp + x));
pa[x] = f.x;
pa[x + 1] = f.y;
}
}
float mn = fmaxf(m, pm);
float c1 = exp2f(m - mn);
float c2 = exp2f(pm - mn);
l = l * c1 + pl * c2;
#pragma unroll
for (int x = 0; x < EL; ++x) a[x] = a[x] * c1 + pa[x] * c2;
m = mn;
}
float inv = 1.f / l;
__nv_bfloat16 ob[EL];
#pragma unroll
for (int x = 0; x < EL; ++x) ob[x] = __float2bfloat16(a[x] * inv);
__nv_bfloat16* dst = o + warp * D + lane * EL;
if (EL == 2) {
*reinterpret_cast<unsigned int*>(dst) = *reinterpret_cast<unsigned int*>(ob);
} else {
*reinterpret_cast<uint2*>(dst) = *reinterpret_cast<uint2*>(ob);
}
}
torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
TORCH_CHECK(q.is_cuda() && q.dtype() == torch::kBFloat16);
TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
int NB = (S + 63) / 64;
int BH = B * H;
long long nq = (long long)BH * S;
auto o = torch::empty_like(q);
auto fopts = q.options().dtype(torch::kFloat32);
auto iopts = q.options().dtype(torch::kInt32);
auto bopts = q.options().dtype(torch::kUInt8);
auto ksum = torch::empty({(long long)BH * NB * D}, fopts);
auto cur_imp = torch::empty({nq}, fopts);
auto imp = torch::empty({nq * ((NB + 31) & ~31)}, fopts);
int NBp = (NB + 31) & ~31;
const int SLOTS_H = 10;
auto sel8 = torch::empty({nq * SLOTS_H}, bopts);
auto seq8 = torch::empty({nq * SLOTS_H}, q.options().dtype(torch::kInt16));
auto flags = torch::empty({nq}, bopts);
auto hist = torch::zeros({BH * NB}, iopts);
auto offs = torch::empty({BH * NB + 1}, iopts);
auto ulist = torch::empty({nq * SLOTS_H}, iopts);
float scale = 1.0f / sqrtf((float)D);
float pres = scale * 1.4426950408889634f;
auto stream = at::cuda::getCurrentCUDAStream();
auto qp = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
auto kp = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
auto vp = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
auto op = reinterpret_cast<__nv_bfloat16*>(o.data_ptr());
#define NSA_DISPATCH(DD) \
do { \
auto pml = torch::empty({nq * SLOTS * 2}, fopts); \
auto pacc = torch::empty({nq * SLOTS * DD}, q.options()); \
auto pmlp = pml.data_ptr<float>(); \
auto paccp = reinterpret_cast<__nv_bfloat16*>(pacc.data_ptr()); \
nsa_prep<DD><<<BH * NB, DD * 4, 0, stream>>>( \
qp, kp, ksum.data_ptr<float>(), cur_imp.data_ptr<float>(), \
S, NB, scale); \
dim3 gg((S + 63) / 64, (NB + 31) / 32, BH); \
size_t gsh = (64 * (DD + 2) + 32 * (DD + 1)) * sizeof(float); \
static bool set_##DD = false; \
if (!set_##DD) { \
cudaFuncSetAttribute(nsa_imp_gemm<DD>, \
cudaFuncAttributeMaxDynamicSharedMemorySize, (int)gsh); \
cudaFuncSetAttribute(nsa_block<DD, DD == 64 ? 8 : 16>, \
cudaFuncAttributeMaxDynamicSharedMemorySize, \
64 * DD * 2 * 2 + (DD == 64 ? 8 : 16) * 16 * DD * 2); \
set_##DD = true; \
} \
nsa_imp_gemm<DD><<<gg, 256, gsh, stream>>>( \
qp, ksum.data_ptr<float>(), imp.data_ptr<float>(), \
S, NB, NBp, scale / 64.0f); \
dim3 sg((S + 31) / 32, BH); \
nsa_select<<<sg, 512, 0, stream>>>( \
imp.data_ptr<float>(), cur_imp.data_ptr<float>(), \
sel8.data_ptr<unsigned char>(), seq8.data_ptr<short>(), \
flags.data_ptr<unsigned char>(), \
hist.data_ptr<int>(), S, NB, NBp); \
nsa_scan<<<1, 1024, 0, stream>>>( \
hist.data_ptr<int>(), offs.data_ptr<int>(), BH * NB); \
nsa_scatter<<<(int)((nq * SLOTS_H + 255) / 256), 256, 0, stream>>>( \
sel8.data_ptr<unsigned char>(), seq8.data_ptr<short>(), \
flags.data_ptr<unsigned char>(), \
offs.data_ptr<int>(), ulist.data_ptr<int>(), BH, S, NB); \
int split = (DD == 64 ? 6016 : 3008) / (BH * NB); \
if (split < 1) split = 1; \
if (split > 16) split = 16; \
nsa_block<DD, DD == 64 ? 8 : 16> \
<<<BH * NB * split, DD == 64 ? 256 : 512, \
64 * DD * 2 * 2 + (DD == 64 ? 8 : 16) * 16 * DD * 2, stream>>>( \
qp, kp, vp, offs.data_ptr<int>(), ulist.data_ptr<int>(), \
pmlp, paccp, S, NB, split, pres); \
nsa_merge<DD><<<(int)((nq * 32 + 255) / 256), 256, 0, stream>>>( \
pmlp, paccp, flags.data_ptr<unsigned char>(), op, BH, S); \
} while (0)
if (D == 64) {
NSA_DISPATCH(64);
} else if (D == 128) {
NSA_DISPATCH(128);
} else {
TORCH_CHECK(false, "unsupported D");
}
#undef NSA_DISPATCH
return o;
}
"""
_CPP_SRC = "torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v);"
os.environ["TORCH_CUDA_ARCH_LIST"] = "12.0"
_ext = None
def _get_ext():
global _ext
if _ext is None:
_ext = load_inline(
name="nsa_sparse_attn_v20",
cpp_sources=[_CPP_SRC],
cuda_sources=[_CUDA_SRC],
functions=["nsa_forward"],
extra_cuda_cflags=["-O3"],
verbose=os.environ.get("NSA_VERBOSE_BUILD", "0") == "1",
)
return _ext
def _nsa_torch_fallback(q, k, v, block_size=BLOCK_SIZE, top_n=TOP_N_BLOCKS,
window=SLIDING_WINDOW):
"""Vectorized reference-equivalent (fp32), used for non-CUDA / odd-D inputs."""
prev = torch.get_float32_matmul_precision()
torch.set_float32_matmul_precision("highest")
try:
return _nsa_torch_fallback_impl(q, k, v, block_size, top_n, window)
finally:
torch.set_float32_matmul_precision(prev)
def _nsa_torch_fallback_impl(q, k, v, block_size, top_n, window):
qf, kf, vf = q.float(), k.float(), v.float()
B, H, S, D = qf.shape
scale = 1.0 / math.sqrt(D)
nb = (S + block_size - 1) // block_size
dev = qf.device
pad = nb * block_size - S
kp = torch.nn.functional.pad(kf, (0, 0, 0, pad))
kb = kp.view(B, H, nb, block_size, D)
ksum = kb.sum(3) # (B,H,nb,D)
t_idx = torch.arange(S, device=dev)
cb = t_idx // block_size
# full-block importance
imp = torch.matmul(qf, ksum.transpose(-1, -2)) * (scale / block_size)
# partial current-block importance
cum = kb.cumsum(3).view(B, H, nb * block_size, D)[:, :, :S, :]
lens = (t_idx % block_size + 1).float()
cur = torch.einsum("bhsd,bhsd->bhs", qf, cum) * scale / lens
imp = imp.clone()
imp[..., :] = torch.where(
torch.arange(nb, device=dev)[None, None, None, :] == cb[None, None, :, None],
cur[..., None], imp)
causal_b = torch.arange(nb, device=dev)[None, None, None, :] <= cb[None, None, :, None]
imp = torch.where(causal_b, imp, torch.full_like(imp, -float("inf")))
# top-n with ties broken toward higher block id: flip, topk, map back
if nb <= top_n:
blk_mask = causal_b
else:
impf = imp.flip(-1)
topi = impf.topk(top_n, dim=-1).indices
topi = nb - 1 - topi
blk_mask = torch.zeros_like(imp, dtype=torch.bool)
blk_mask.scatter_(-1, topi, True)
blk_mask &= causal_b
key_mask = blk_mask[..., :, torch.arange(S, device=dev) // block_size]
j_idx = torch.arange(S, device=dev)
causal = j_idx[None, :] <= t_idx[:, None]
win = j_idx[None, :] >= (t_idx[:, None] - (window - 1))
key_mask = (key_mask | (win[None, None])) & causal[None, None]
sc = torch.matmul(qf, kf.transpose(-1, -2)) * scale
sc = sc.masked_fill(~key_mask, -float("inf"))
att = torch.softmax(sc, dim=-1)
return torch.matmul(att, vf)
class Model(nn.Module):
def __init__(self, B: int, H: int, S: int, D: int):
super().__init__()
self.B, self.H, self.S, self.D = B, H, S, D
self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16))
self._graphs = {}
if torch.cuda.is_available() and D in (64, 128):
_get_ext()
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
B, H, S, D = q.shape
if not (q.is_cuda and q.dtype == torch.bfloat16 and D in (64, 128)
and 64 <= S <= 8192):
return _nsa_torch_fallback(q, k, v).to(torch.bfloat16)
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
ext = _get_ext()
# CUDA-graph the whole pipeline when the same input buffers are used
# repeatedly (benchmark steady state); fresh buffers stay eager.
key = (q.data_ptr(), k.data_ptr(), v.data_ptr(), B, H, S, D)
ent = self._graphs.get(key)
if ent is None:
if len(self._graphs) > 16:
self._graphs.clear()
self._graphs[key] = [1, None, None]
return ext.nsa_forward(q, k, v)
if ent[1] is not None:
ent[1].replay()
return ent[2]
ent[0] += 1
if ent[0] < 3:
return ext.nsa_forward(q, k, v)
try:
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
out = ext.nsa_forward(q, k, v)
ent[1] = g
ent[2] = out
g.replay()
return out
except Exception:
self._graphs[key] = [1, None, None]
return ext.nsa_forward(q, k, v)
def get_init_inputs():
return [1, 16, 1024, 64]
def get_inputs():
q = torch.randn(1, 16, 1024, 64, dtype=torch.bfloat16)
k = torch.randn(1, 16, 1024, 64, dtype=torch.bfloat16)
v = torch.randn(1, 16, 1024, 64, dtype=torch.bfloat16)
return [q, k, v]
20260721_153000_or-fable_anthropic_claude-fable-5_02_deepseek_nsa