KernelBench cuda · RTX PRO 6000
DeepSeek NSA Grok 4.5
1.77%geomean peak fraction across shapes
manually audited: clean
Genuine fused CUDA NSA block-select sparse attention kernel via load_inline. Per query it really scores every causal 64-token block (mean of scaled q.k over the block's causal portion), keeps a running top-8 by insertion sort, then runs warp-owned fp32 online-softmax attention over selected blocks unioned with the 64-token sliding window, causal-masked. Empirically proven to recompute on live data (in-place input overwrite changes output, max delta 4.06), to match the reference on off-deck shapes and both scale extremes (max err 0.002 vs 0.1 tol), and to genuinely select a planted high-importance distant block. peak_fraction 0.0177 is low BY DESIGN: the flops_formula is dense-equivalent 4*B*H*S*S*D while a correct kernel exploits sparsity, so this is not evidence of weakness or cheating. Language gate: cuda_raw (global_kernel + cuda_header), Triton/DSL clean.
harnessgrokagent session1h 0mtotal wall1h 10mcheck5mbenchmark5moutput tokens—gpu-lock wait9mgpu-lock held41sregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
1×16×2048×642.304 ms1.5%7 TFLOPS · 1% of 500 TF bf16 peak · also 0.01 TB/s (0% of HBM)
1×16×4127×646.040 ms2.3%12 TFLOPS · 2% of 500 TF bf16 peak · also 0.01 TB/s (0% of HBM)
1×8×8192×649.307 ms2.9%15 TFLOPS · 3% of 500 TF bf16 peak · also 0.00 TB/s (0% of HBM)
1×8×8191×12835.010 ms1.6%8 TFLOPS · 2% of 500 TF bf16 peak · also 0.00 TB/s (0% of HBM)
4×8×1024×641.688 ms1.0%5 TFLOPS · 1% of 500 TF bf16 peak · also 0.01 TB/s (1% of HBM)
2×8×3000×643.848 ms1.9%10 TFLOPS · 2% of 500 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.5% · 2.3% · 2.9% · 1.6% · 1.0% · 1.9%) = 1.8%
Kernel source (redacted)
"""DeepSeek NSA-inspired sparse attention — high-performance CUDA kernel.
Bench semantics matching reference.nsa_attend:
block_size=64, top_n_blocks=8, sliding_window=64, causal, bf16 I/O.
Architecture:
Query-tiled CTAs stream K through shared memory.
Phase 1: parallel block-importance (4 lanes/query) + online top-8.
Phase 2: warp-owned sparse online-softmax over selected ∪ window.
"""
from __future__ import annotations
from typing import Optional
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
B, H, S, D = 1, 16, 1024, 64
CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cmath>
#include <cfloat>
static constexpr int BS = 64;
static constexpr int TOPN = 8;
static constexpr int WIN = 64;
__device__ __forceinline__ float bf2f(__nv_bfloat16 x) {
return __bfloat162float(x);
}
__device__ __forceinline__ __nv_bfloat16 f2bf(float x) {
return __float2bfloat16(x);
}
__device__ __forceinline__ float warp_sum(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v += __shfl_down_sync(0xffffffff, v, off);
return v;
}
__device__ __forceinline__ float warp_sum_bcast(float v) {
return __shfl_sync(0xffffffff, warp_sum(v), 0);
}
__device__ __forceinline__ void topk_insert(float* sc, int* ix, float s, int bi) {
if (s < sc[TOPN - 1]) return;
int pos = TOPN - 1;
while (pos > 0 && s > sc[pos - 1]) {
sc[pos] = sc[pos - 1];
ix[pos] = ix[pos - 1];
--pos;
}
sc[pos] = s;
ix[pos] = bi;
}
template <int TD>
__device__ __forceinline__ float dot_ff(const float* __restrict__ q,
const float* __restrict__ k) {
float acc = 0.f;
#pragma unroll
for (int d = 0; d < TD; d += 4) {
acc += q[d] * k[d] + q[d+1] * k[d+1] + q[d+2] * k[d+2] + q[d+3] * k[d+3];
}
return acc;
}
// =====================================================================
// TD head dim, QT query tile, NT threads per block
// =====================================================================
template <int TD, int QT, int NT>
__global__ void __launch_bounds__(NT, 4)
nsa_kernel(
const __nv_bfloat16* __restrict__ Q,
const __nv_bfloat16* __restrict__ K,
const __nv_bfloat16* __restrict__ V,
__nv_bfloat16* __restrict__ O,
int B_, int H_, int S_, int D_
) {
const int bh = blockIdx.x;
const int tile = blockIdx.y;
const int t0 = tile * QT;
if (t0 >= S_ || bh >= B_ * H_) return;
const int tid = threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
constexpr int NW = NT / 32;
const long long head_off = (static_cast<long long>(bh) * S_) * D_;
const __nv_bfloat16* q_base = Q + head_off;
const __nv_bfloat16* k_base = K + head_off;
const __nv_bfloat16* v_base = V + head_off;
__nv_bfloat16* o_base = O + head_off;
const float scale = rsqrtf(static_cast<float>(D_));
const int n_blocks = (S_ + BS - 1) / BS;
const int t_end = min(t0 + QT, S_);
const int nq = t_end - t0;
// smem: q_f[QT*TD] | k_f[BS*TD] | sel[QT*TOPN] | bsc[QT]
extern __shared__ float smem[];
float* q_f = smem;
float* k_f = q_f + QT * TD;
int* sel = reinterpret_cast<int*>(k_f + BS * TD);
float* bsc = reinterpret_cast<float*>(sel + QT * TOPN);
// Load Q as float
for (int i = tid; i < nq * TD; i += NT) {
int qi = i / TD;
int d = i - qi * TD;
q_f[qi * TD + d] = (d < D_)
? bf2f(q_base[(static_cast<long long>(t0 + qi)) * D_ + d])
: 0.f;
}
for (int i = nq * TD + tid; i < QT * TD; i += NT) q_f[i] = 0.f;
__syncthreads();
float top_sc[TOPN];
int top_bi[TOPN];
#pragma unroll
for (int s = 0; s < TOPN; ++s) { top_sc[s] = -1e30f; top_bi[s] = -1; }
// =================================================================
// Phase 1
// =================================================================
constexpr int LPQ = 4;
constexpr int QPW = 32 / LPQ; // 8 queries per warp
for (int bi = 0; bi < n_blocks; ++bi) {
const int s0 = bi * BS;
if (s0 > t_end - 1) break;
const int nk = min(BS, S_ - s0);
// Load K block → float smem (coalesced)
for (int i = tid; i < nk * D_; i += NT) {
int j = i / D_;
int d = i - j * D_;
k_f[j * TD + d] = bf2f(k_base[(static_cast<long long>(s0 + j)) * D_ + d]);
}
for (int i = nk * TD + tid; i < BS * TD; i += NT) k_f[i] = 0.f;
__syncthreads();
// Score: 4 lanes/query, all lanes participate in shfl
{
const int qi_local = lane / LPQ;
const int sub = lane - qi_local * LPQ;
const int qi = warp * QPW + qi_local;
const bool active = (qi < nq);
float local = 0.f;
int j1 = 0;
bool causal = false;
if (active) {
const int t = t0 + qi;
if (s0 <= t) {
causal = true;
j1 = min(s0 + BS, t + 1) - s0;
for (int j = sub; j < j1; j += LPQ)
local += dot_ff<TD>(q_f + qi * TD, k_f + j * TD);
}
}
local += __shfl_xor_sync(0xffffffff, local, 1);
local += __shfl_xor_sync(0xffffffff, local, 2);
if (active && sub == 0) {
bsc[qi] = causal ? (local / static_cast<float>(j1)) * scale : -1e9f;
}
}
__syncthreads();
if (tid < nq) topk_insert(top_sc, top_bi, bsc[tid], bi);
__syncthreads();
}
if (tid < nq) {
#pragma unroll
for (int s = 0; s < TOPN; ++s) sel[tid * TOPN + s] = top_bi[s];
}
__syncthreads();
// =================================================================
// Phase 2 — sparse online attention
// =================================================================
constexpr int EPT = (TD + 31) / 32;
for (int qi = warp; qi < nq; qi += NW) {
const int t = t0 + qi;
const int w0 = (t + 1 > WIN) ? (t + 1 - WIN) : 0;
int my_sel[TOPN];
#pragma unroll
for (int s = 0; s < TOPN; ++s) my_sel[s] = sel[qi * TOPN + s];
// Selected-block bitset for n_blocks <= 256
unsigned long long sel_bits0 = 0ull, sel_bits1 = 0ull, sel_bits2 = 0ull, sel_bits3 = 0ull;
#pragma unroll
for (int s = 0; s < TOPN; ++s) {
int bi = my_sel[s];
if (bi < 0) continue;
if (bi < 64) sel_bits0 |= (1ull << bi);
else if (bi < 128) sel_bits1 |= (1ull << (bi - 64));
else if (bi < 192) sel_bits2 |= (1ull << (bi - 128));
else if (bi < 256) sel_bits3 |= (1ull << (bi - 192));
}
auto is_sel = [&](int bi) -> bool {
if (bi < 64) return (sel_bits0 >> bi) & 1ull;
if (bi < 128) return (sel_bits1 >> (bi - 64)) & 1ull;
if (bi < 192) return (sel_bits2 >> (bi - 128)) & 1ull;
if (bi < 256) return (sel_bits3 >> (bi - 192)) & 1ull;
return false;
};
float q_reg[EPT];
#pragma unroll
for (int e = 0; e < EPT; ++e) {
int d = lane + e * 32;
q_reg[e] = (d < D_) ? q_f[qi * TD + d] : 0.f;
}
float m_i = -FLT_MAX;
float l_i = 0.f;
float acc[EPT];
#pragma unroll
for (int e = 0; e < EPT; ++e) acc[e] = 0.f;
// Process one key from global memory (all lanes)
#define PROCESS_KEY(j_expr) \
do { \
const int j__ = (j_expr); \
const __nv_bfloat16* kj = k_base + (static_cast<long long>(j__)) * D_; \
const __nv_bfloat16* vj = v_base + (static_cast<long long>(j__)) * D_; \
float partial = 0.f; \
_Pragma("unroll") \
for (int e = 0; e < EPT; ++e) { \
int d = lane + e * 32; \
if (d < D_) partial += q_reg[e] * bf2f(kj[d]); \
} \
float score = warp_sum_bcast(partial) * scale; \
float m_new = fmaxf(m_i, score); \
float ed = __expf(m_i - m_new); \
float es = __expf(score - m_new); \
_Pragma("unroll") \
for (int e = 0; e < EPT; ++e) acc[e] *= ed; \
l_i = l_i * ed + es; \
m_i = m_new; \
_Pragma("unroll") \
for (int e = 0; e < EPT; ++e) { \
int d = lane + e * 32; \
if (d < D_) acc[e] += es * bf2f(vj[d]); \
} \
} while (0)
// Pass A: selected blocks
#pragma unroll
for (int s = 0; s < TOPN; ++s) {
int bi = my_sel[s];
if (bi < 0) continue;
int s0 = bi * BS;
if (s0 > t) continue;
int s1 = min(s0 + BS, t + 1);
for (int j = s0; j < s1; ++j) {
PROCESS_KEY(j);
}
}
// Pass B: window keys not in selected blocks
for (int j = w0; j <= t; ++j) {
if (is_sel(j / BS)) continue;
PROCESS_KEY(j);
}
if (l_i == 0.f) {
PROCESS_KEY(t);
}
#undef PROCESS_KEY
const float inv = 1.f / l_i;
#pragma unroll
for (int e = 0; e < EPT; ++e) {
int d = lane + e * 32;
if (d < D_)
o_base[(static_cast<long long>(t)) * D_ + d] = f2bf(acc[e] * inv);
}
}
}
// Generic fallback for uncommon D
__global__ void nsa_single_kernel(
const __nv_bfloat16* __restrict__ Q,
const __nv_bfloat16* __restrict__ K,
const __nv_bfloat16* __restrict__ V,
__nv_bfloat16* __restrict__ O,
int B_, int H_, int S_, int D_
) {
const int bh = blockIdx.x, t = blockIdx.y;
if (t >= S_ || bh >= B_ * H_) return;
const int tid = threadIdx.x, NT = blockDim.x;
const long long head_off = (static_cast<long long>(bh) * S_) * D_;
const __nv_bfloat16* q_ptr = Q + head_off + (static_cast<long long>(t)) * D_;
const __nv_bfloat16* k_base = K + head_off;
const __nv_bfloat16* v_base = V + head_off;
__nv_bfloat16* o_ptr = O + head_off + (static_cast<long long>(t)) * D_;
const int n_blocks = (S_ + BS - 1) / BS;
const float scale = rsqrtf(static_cast<float>(D_));
extern __shared__ float smem[];
float* q_f = smem;
float* block_imp = q_f + D_;
float* red = block_imp + n_blocks;
int* top_idx = reinterpret_cast<int*>(red + 32);
unsigned char* sel_mask = reinterpret_cast<unsigned char*>(top_idx + TOPN);
for (int d = tid; d < D_; d += NT) q_f[d] = bf2f(q_ptr[d]);
__syncthreads();
auto bsum = [&](float v) -> float {
int lane = tid & 31, wid = tid >> 5;
v = warp_sum(v);
if (lane == 0) red[wid] = v;
__syncthreads();
int nwarps = (NT + 31) >> 5;
v = (tid < nwarps) ? red[tid] : 0.f;
if (wid == 0) v = warp_sum(v);
if (tid == 0) red[0] = v;
__syncthreads();
v = red[0];
__syncthreads();
return v;
};
for (int bi = 0; bi < n_blocks; ++bi) {
int s0 = bi * BS;
if (s0 > t) {
if (tid == 0) block_imp[bi] = -1e9f;
__syncthreads();
continue;
}
int s1 = min(s0 + BS, t + 1);
float local = 0.f;
for (int j = s0 + tid; j < s1; j += NT) {
const __nv_bfloat16* kj = k_base + (static_cast<long long>(j)) * D_;
float dot = 0.f;
for (int d = 0; d < D_; ++d) dot += q_f[d] * bf2f(kj[d]);
local += dot;
}
float total = bsum(local);
if (tid == 0) block_imp[bi] = (total / float(s1 - s0)) * scale;
__syncthreads();
}
if (tid == 0) {
int n_causal = 0;
for (int bi = 0; bi < n_blocks; ++bi)
if (bi * BS <= t) n_causal++; else break;
int k_sel = TOPN < n_causal ? TOPN : n_causal;
for (int s = 0; s < k_sel; ++s) {
float best = -1e30f; int best_i = 0;
for (int bi = 0; bi < n_causal; ++bi) {
bool used = false;
for (int p = 0; p < s; ++p) if (top_idx[p] == bi) used = true;
if (!used && block_imp[bi] > best) { best = block_imp[bi]; best_i = bi; }
}
top_idx[s] = best_i;
}
for (int s = k_sel; s < TOPN; ++s) top_idx[s] = -1;
for (int bi = 0; bi < n_blocks; ++bi) sel_mask[bi] = 0;
for (int s = 0; s < k_sel; ++s) sel_mask[top_idx[s]] = 1;
}
__syncthreads();
float acc[8]; for (int i = 0; i < 8; ++i) acc[i] = 0.f;
float m_i = -FLT_MAX, l_i = 0.f;
int w0 = (t + 1 > WIN) ? (t + 1 - WIN) : 0;
float q_reg[8];
for (int i = 0; i < 8; ++i) {
int d = tid + i * NT;
q_reg[i] = (d < D_) ? q_f[d] : 0.f;
}
for (int bi = 0; bi < n_blocks; ++bi) {
int s0 = bi * BS;
if (s0 > t) break;
int s1 = min(s0 + BS, t + 1);
int j0, j1;
if (sel_mask[bi]) { j0 = s0; j1 = s1; }
else {
j0 = s0 > w0 ? s0 : w0; j1 = s1;
if (j0 >= j1) continue;
}
for (int j = j0; j < j1; ++j) {
const __nv_bfloat16* kj = k_base + (static_cast<long long>(j)) * D_;
const __nv_bfloat16* vj = v_base + (static_cast<long long>(j)) * D_;
float partial = 0.f;
for (int i = 0; i < 8; ++i) {
int d = tid + i * NT;
if (d < D_) partial += q_reg[i] * bf2f(kj[d]);
}
float score = bsum(partial) * scale;
float m_new = fmaxf(m_i, score);
float ed = __expf(m_i - m_new);
float es = __expf(score - m_new);
for (int i = 0; i < 8; ++i) acc[i] *= ed;
l_i = l_i * ed + es;
m_i = m_new;
for (int i = 0; i < 8; ++i) {
int d = tid + i * NT;
if (d < D_) acc[i] += es * bf2f(vj[d]);
}
}
}
if (l_i == 0.f) {
const __nv_bfloat16* kj = k_base + (static_cast<long long>(t)) * D_;
const __nv_bfloat16* vj = v_base + (static_cast<long long>(t)) * D_;
float partial = 0.f;
for (int i = 0; i < 8; ++i) {
int d = tid + i * NT;
if (d < D_) partial += q_reg[i] * bf2f(kj[d]);
}
float score = bsum(partial) * scale;
float m_new = fmaxf(m_i, score);
float ed = __expf(m_i - m_new);
float es = __expf(score - m_new);
for (int i = 0; i < 8; ++i) acc[i] *= ed;
l_i = l_i * ed + es;
for (int i = 0; i < 8; ++i) {
int d = tid + i * NT;
if (d < D_) acc[i] += es * bf2f(vj[d]);
}
}
float inv = 1.f / l_i;
for (int i = 0; i < 8; ++i) {
int d = tid + i * NT;
if (d < D_) o_ptr[d] = f2bf(acc[i] * inv);
}
}
torch::Tensor nsa_forward(
torch::Tensor q, torch::Tensor k, torch::Tensor v,
int64_t, int64_t, int64_t
) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "CUDA required");
TORCH_CHECK(q.scalar_type() == at::kBFloat16, "bf16 required");
TORCH_CHECK(q.dim() == 4, "(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);
auto out = torch::empty_like(qc);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
auto* Qp = reinterpret_cast<const __nv_bfloat16*>(qc.data_ptr<at::BFloat16>());
auto* Kp = reinterpret_cast<const __nv_bfloat16*>(kc.data_ptr<at::BFloat16>());
auto* Vp = reinterpret_cast<const __nv_bfloat16*>(vc.data_ptr<at::BFloat16>());
auto* Op = reinterpret_cast<__nv_bfloat16*>(out.data_ptr<at::BFloat16>());
auto launch = [&](auto kernel, int TD, int QT, int threads) {
dim3 grid(B_ * H_, (S_ + QT - 1) / QT);
size_t smem =
sizeof(float) * ((size_t)QT * TD + (size_t)BS * TD + (size_t)QT) +
sizeof(int) * (size_t)QT * TOPN;
smem = (smem + 15) & ~size_t(15);
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
kernel<<<grid, threads, smem, stream>>>(Qp, Kp, Vp, Op, B_, H_, S_, D_);
};
if (D_ == 64) {
// QT=32, 256 threads: 8 warps → each warp handles 4 queries in phase 2
// smem: q 8KB + k 16KB + sel ~1KB ≈ 25KB
launch(nsa_kernel<64, 32, 256>, 64, 32, 256);
} else if (D_ == 128) {
// QT=16, 256 threads: more phase-2 parallelism
// smem: q 8KB + k 32KB + sel ≈ 41KB
launch(nsa_kernel<128, 16, 256>, 128, 16, 256);
} else {
dim3 grid(B_ * H_, S_);
int threads = 128;
int n_blocks = (S_ + BS - 1) / BS;
size_t smem = sizeof(float) * (D_ + n_blocks + 32) + sizeof(int) * TOPN + n_blocks;
smem = (smem + 15) & ~size_t(15);
nsa_single_kernel<<<grid, threads, smem, stream>>>(Qp, Kp, Vp, Op, B_, H_, S_, D_);
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
return out;
}
"""
CPP_SRC = r"""
torch::Tensor nsa_forward(
torch::Tensor q, torch::Tensor k, torch::Tensor v,
int64_t block_size, int64_t top_n_blocks, int64_t sliding_window);
"""
_nsa_mod: Optional[object] = None
def _get_mod():
global _nsa_mod
if _nsa_mod is None:
_nsa_mod = load_inline(
name="nsa_sparse_attn_v7",
cpp_sources=[CPP_SRC],
cuda_sources=[CUDA_SRC],
functions=["nsa_forward"],
extra_cuda_cflags=[
"-O3",
"--use_fast_math",
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
"-U__CUDA_NO_HALF2_OPERATORS__",
"--expt-relaxed-constexpr",
],
verbose=False,
)
return _nsa_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))
if torch.cuda.is_available():
_get_mod()
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
return _get_mod().nsa_forward(
q, k, v, BLOCK_SIZE, TOP_N_BLOCKS, SLIDING_WINDOW
)
def get_init_inputs():
return [B, H, S, D]
def get_inputs():
q = torch.randn(B, H, S, D, dtype=torch.bfloat16)
k = torch.randn(B, H, S, D, dtype=torch.bfloat16)
v = torch.randn(B, H, S, D, dtype=torch.bfloat16)
return [q, k, v]
20260715_212751_grok_grok-4.5_02_deepseek_nsa