KernelBench cuda · RTX PRO 6000
DeepSeek NSA Gemini 4 Argon (High)
0.190 msgeomean latency across six shapes · lower is better
manually audited: clean
Gemini 4 Argon wrote exact NSA selection and sparse attention as one tensor-core CUDA kernel, one CTA per 64-query block: the causal diagonal block on mma.sync, fp32 key means published between CTAs with release/acquire flags, a register top-8 over q·mean(k) with ties to the larger block, the 64-token window, and a cp.async-pipelined loop over the CTA's union of selected blocks with online softmax.
harnessagyagent session6h 45mtotal wall6h 52mcheck6mbenchmark1soutput tokens—gpu-lock wait5mgpu-lock held4h 24mregimecompute
Per-shape latencymilliseconds are the measured result; the dense-equivalent roofline is not a useful ceiling
1×16×2048×640.093 ms
1×16×4127×640.236 ms
1×8×8192×640.379 ms
1×8×8191×1280.659 ms
4×8×1024×640.060 ms
2×8×3000×640.144 ms
geomean(0.093 ms · 0.236 ms · 0.379 ms · 0.659 ms · 0.060 ms · 0.144 ms) = 0.190 ms
Kernel source (redacted)
from __future__ import annotations
import os
try:
import ninja
os.environ["PATH"] = f"{ninja.BIN_DIR}:{os.environ.get('PATH', '')}"
except ImportError:
pass
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
OP_TYPE = "deepseek_nsa"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000"]
B, H, S, D = 1, 16, 1024, 64
BLOCK_SIZE = 64
TOP_N_BLOCKS = 8
SLIDING_WINDOW = 64
cuda_source = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <torch/extension.h>
#include <ATen/cuda/EmptyTensor.h>
#include <cstdint>
namespace {
__device__ __forceinline__ uint32_t smem_u32addr(const void* ptr) {
return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
}
__device__ __forceinline__ void cp_async_ca_16(uint32_t smem_addr, const void* gmem_ptr) {
asm volatile(
"cp.async.ca.shared.global [%0], [%1], 16;\n"
:
: "r"(smem_addr), "l"(gmem_ptr)
);
}
__device__ __forceinline__ void cp_async_commit() {
asm volatile("cp.async.commit_group;\n" ::);
}
template <int N>
__device__ __forceinline__ void cp_async_wait() {
asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}
__device__ __forceinline__ void ldmatrix_x4(uint32_t &r0, uint32_t &r1, uint32_t &r2, uint32_t &r3, uint32_t addr) {
asm volatile(
"ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
: "r"(addr)
);
}
__device__ __forceinline__ void ldmatrix_x4_trans(uint32_t &r0, uint32_t &r1, uint32_t &r2, uint32_t &r3, uint32_t addr) {
asm volatile(
"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
: "r"(addr)
);
}
__device__ __forceinline__ void mma_m16n8k16_bf16(
float &d0, float &d1, float &d2, float &d3,
uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
uint32_t b0, uint32_t 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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)
);
}
template <int D>
struct CTASharedStorage {
__nv_bfloat16 k_tile[64][D + 8];
__nv_bfloat16 v_tile[64][D + 8];
uint64_t warp_mask0[4];
uint64_t warp_mask1[4];
};
union SPFrag {
float s[8][4];
uint32_t p[4][4];
};
template <int D, int CTA_LOAD_ITERS, int UINT4_PER_ROW>
__device__ __forceinline__ void issue_cta_tile_cp_async(
__nv_bfloat16 tile[64][D + 8],
const __nv_bfloat16* __restrict__ gmem_base,
int tid_cta
) {
#pragma unroll
for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
int idx = tid_cta + i * 128;
int r = idx / UINT4_PER_ROW;
int c = (idx % UINT4_PER_ROW) * 8;
uint32_t dst = smem_u32addr(&tile[r][c]);
cp_async_ca_16(dst, gmem_base + r * D + c);
}
}
template <int D, int WARPS_PER_CTA = 4, int MIN_BLOCKS = 2>
__global__ __launch_bounds__(WARPS_PER_CTA * 32, MIN_BLOCKS)
void nsa_cta_coop_kernel(
const __nv_bfloat16* __restrict__ Q,
const __nv_bfloat16* __restrict__ K,
const __nv_bfloat16* __restrict__ V,
float* __restrict__ K_mean,
uint32_t* __restrict__ km_flag,
uint32_t gen,
__nv_bfloat16* __restrict__ O,
int S,
int n_full_blocks,
float scale
) {
constexpr int TK = D / 16; // 4 (D=64) or 8 (D=128)
constexpr int TN = D / 8; // 8 (D=64) or 16 (D=128)
constexpr int UINT4_PER_ROW = D / 8; // 8 (D=64) or 16 (D=128)
constexpr int CTA_LOAD_ITERS = (64 * UINT4_PER_ROW) / 128; // 4 (D=64) or 8 (D=128)
__shared__ CTASharedStorage<D> smem;
const int tid_cta = threadIdx.x;
const int w = tid_cta >> 5;
const int lane = tid_cta & 31;
const int groupID = lane >> 2; // 0..7
const int tid = lane & 3; // 0..3
const int bi_curr = blockIdx.x;
const int s0 = bi_curr * 64;
const int t_warp_base = s0 + w * 16;
const int bh = blockIdx.y;
const int64_t bh_offset = static_cast<int64_t>(bh) * S * D;
const __nv_bfloat16* q_head = Q + bh_offset;
const __nv_bfloat16* k_head = K + bh_offset;
const __nv_bfloat16* v_head = V + bh_offset;
__nv_bfloat16* o_head = O + bh_offset;
const int q_top_local = w * 16 + groupID;
const int q_bot_local = w * 16 + groupID + 8;
const int t_top = s0 + q_top_local;
const int t_bot = s0 + q_bot_local;
const bool valid_q_top = (t_top < S);
const bool valid_q_bot = (t_bot < S);
const int sub_r = lane & 15;
const int sub_c = (lane >> 4) << 3;
// 1. Cooperative CTA load of Q[s0 .. s0+63] into smem.k_tile
#pragma unroll
for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
int idx = tid_cta + i * 128;
int r = idx / UINT4_PER_ROW;
int c = (idx % UINT4_PER_ROW) * 8;
int t_r = s0 + r;
uint4 val = (t_r < S)
? __ldg(reinterpret_cast<const uint4*>(q_head + static_cast<int64_t>(t_r) * D + c))
: make_uint4(0, 0, 0, 0);
*reinterpret_cast<uint4*>(&smem.k_tile[r][c]) = val;
}
__syncthreads();
uint32_t q_frag[TK][4];
#pragma unroll
for (int tk = 0; tk < TK; ++tk) {
uint32_t addr = smem_u32addr(&smem.k_tile[w * 16 + sub_r][tk * 16 + sub_c]);
ldmatrix_x4(q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], addr);
}
__syncthreads();
// Online softmax state for q_top and q_bot
float m_top = -1e30f, l_top = 0.0f;
float m_bot = -1e30f, l_bot = 0.0f;
float o_frag[TN][4];
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
o_frag[tn][0] = 0.0f;
o_frag[tn][1] = 0.0f;
o_frag[tn][2] = 0.0f;
o_frag[tn][3] = 0.0f;
}
float sum_curr_top = 0.0f;
float sum_curr_bot = 0.0f;
// Step 1: Cooperative CTA load of K[bi_curr] and V[bi_curr]
if (s0 + 64 <= S) {
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + static_cast<int64_t>(s0) * D, tid_cta);
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + static_cast<int64_t>(s0) * D, tid_cta);
cp_async_commit();
cp_async_wait<0>();
} else {
#pragma unroll
for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
int idx = tid_cta + i * 128;
int r = idx / UINT4_PER_ROW;
int c = (idx % UINT4_PER_ROW) * 8;
int k_idx = s0 + r;
uint4 kv = (k_idx < S) ? __ldg(reinterpret_cast<const uint4*>(k_head + static_cast<int64_t>(k_idx) * D + c)) : make_uint4(0, 0, 0, 0);
uint4 vv = (k_idx < S) ? __ldg(reinterpret_cast<const uint4*>(v_head + static_cast<int64_t>(k_idx) * D + c)) : make_uint4(0, 0, 0, 0);
*reinterpret_cast<uint4*>(&smem.k_tile[r][c]) = kv;
*reinterpret_cast<uint4*>(&smem.v_tile[r][c]) = vv;
}
}
__syncthreads();
if (t_warp_base < S) {
const int max_local_k = min(w * 16 + 15, S - 1 - s0);
for (int k_chunk = 0; k_chunk <= w; ++k_chunk) {
const int k_base = k_chunk * 16;
if (k_base > max_local_k) break;
float s_0[4] = {0.f, 0.f, 0.f, 0.f};
float s_1[4] = {0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int tk = 0; tk < TK; ++tk) {
uint32_t r0, r1, r2, r3;
uint32_t addr = smem_u32addr(&smem.k_tile[k_base + sub_r][tk * 16 + sub_c]);
ldmatrix_x4(r0, r1, r2, r3, addr);
mma_m16n8k16_bf16(s_0[0], s_0[1], s_0[2], s_0[3],
q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r0, r2);
mma_m16n8k16_bf16(s_1[0], s_1[1], s_1[2], s_1[3],
q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r1, r3);
}
const int kc0 = k_base + tid * 2 + 0;
const int kc1 = k_base + tid * 2 + 1;
const int kc2 = k_base + 8 + tid * 2 + 0;
const int kc3 = k_base + 8 + tid * 2 + 1;
bool vt0 = valid_q_top && (kc0 <= q_top_local);
bool vt1 = valid_q_top && (kc1 <= q_top_local);
bool vt2 = valid_q_top && (kc2 <= q_top_local);
bool vt3 = valid_q_top && (kc3 <= q_top_local);
bool vb0 = valid_q_bot && (kc0 <= q_bot_local);
bool vb1 = valid_q_bot && (kc1 <= q_bot_local);
bool vb2 = valid_q_bot && (kc2 <= q_bot_local);
bool vb3 = valid_q_bot && (kc3 <= q_bot_local);
float st0 = vt0 ? (s_0[0] * scale) : -1e30f;
float st1 = vt1 ? (s_0[1] * scale) : -1e30f;
float st2 = vt2 ? (s_1[0] * scale) : -1e30f;
float st3 = vt3 ? (s_1[1] * scale) : -1e30f;
float sb0 = vb0 ? (s_0[2] * scale) : -1e30f;
float sb1 = vb1 ? (s_0[3] * scale) : -1e30f;
float sb2 = vb2 ? (s_1[2] * scale) : -1e30f;
float sb3 = vb3 ? (s_1[3] * scale) : -1e30f;
if (vt0) sum_curr_top += st0;
if (vt1) sum_curr_top += st1;
if (vt2) sum_curr_top += st2;
if (vt3) sum_curr_top += st3;
if (vb0) sum_curr_bot += sb0;
if (vb1) sum_curr_bot += sb1;
if (vb2) sum_curr_bot += sb2;
if (vb3) sum_curr_bot += sb3;
float m_chunk_top = fmaxf(fmaxf(st0, st1), fmaxf(st2, st3));
m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 1));
m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 2));
float m_chunk_bot = fmaxf(fmaxf(sb0, sb1), fmaxf(sb2, sb3));
m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 1));
m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 2));
float m_new_top = fmaxf(m_top, m_chunk_top);
float exp_old_top = (l_top > 0.0f) ? __expf(m_top - m_new_top) : 0.0f;
float pt0 = vt0 ? __expf(st0 - m_new_top) : 0.0f;
float pt1 = vt1 ? __expf(st1 - m_new_top) : 0.0f;
float pt2 = vt2 ? __expf(st2 - m_new_top) : 0.0f;
float pt3 = vt3 ? __expf(st3 - m_new_top) : 0.0f;
float p_sum_top = pt0 + pt1 + pt2 + pt3;
p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 1);
p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 2);
l_top = l_top * exp_old_top + p_sum_top;
m_top = m_new_top;
float m_new_bot = fmaxf(m_bot, m_chunk_bot);
float exp_old_bot = (l_bot > 0.0f) ? __expf(m_bot - m_new_bot) : 0.0f;
float pb0 = vb0 ? __expf(sb0 - m_new_bot) : 0.0f;
float pb1 = vb1 ? __expf(sb1 - m_new_bot) : 0.0f;
float pb2 = vb2 ? __expf(sb2 - m_new_bot) : 0.0f;
float pb3 = vb3 ? __expf(sb3 - m_new_bot) : 0.0f;
float p_sum_bot = pb0 + pb1 + pb2 + pb3;
p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 1);
p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 2);
l_bot = l_bot * exp_old_bot + p_sum_bot;
m_bot = m_new_bot;
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
o_frag[tn][0] *= exp_old_top;
o_frag[tn][1] *= exp_old_top;
o_frag[tn][2] *= exp_old_bot;
o_frag[tn][3] *= exp_old_bot;
}
uint32_t p_frag[4];
*reinterpret_cast<__nv_bfloat162*>(&p_frag[0]) = __floats2bfloat162_rn(pt0, pt1);
*reinterpret_cast<__nv_bfloat162*>(&p_frag[1]) = __floats2bfloat162_rn(pb0, pb1);
*reinterpret_cast<__nv_bfloat162*>(&p_frag[2]) = __floats2bfloat162_rn(pt2, pt3);
*reinterpret_cast<__nv_bfloat162*>(&p_frag[3]) = __floats2bfloat162_rn(pb2, pb3);
#pragma unroll
for (int tn2 = 0; tn2 < TN / 2; ++tn2) {
uint32_t r0, r1, r2, r3;
uint32_t addr = smem_u32addr(&smem.v_tile[k_base + sub_r][tn2 * 16 + sub_c]);
ldmatrix_x4_trans(r0, r1, r2, r3, addr);
mma_m16n8k16_bf16(o_frag[2 * tn2 + 0][0], o_frag[2 * tn2 + 0][1], o_frag[2 * tn2 + 0][2], o_frag[2 * tn2 + 0][3],
p_frag[0], p_frag[1], p_frag[2], p_frag[3], r0, r1);
mma_m16n8k16_bf16(o_frag[2 * tn2 + 1][0], o_frag[2 * tn2 + 1][1], o_frag[2 * tn2 + 1][2], o_frag[2 * tn2 + 1][3],
p_frag[0], p_frag[1], p_frag[2], p_frag[3], r2, r3);
}
}
}
// Fused K_mean computation by Warp 0 (D=64) or Warps 0..1 (D=128) while Warp 3 finishes Step 1!
if (tid_cta < (D / 2) && bi_curr < ((S - 1) >> 6)) {
float* out_ptr = K_mean + (static_cast<int64_t>(bh) * n_full_blocks + bi_curr) * D;
const float scale_inv64 = scale * (1.0f / 64.0f);
const int idx2 = tid_cta;
const int tk_km = idx2 >> 3;
const int sub_tid = (idx2 >> 1) & 3;
const int half_km = idx2 & 1;
const int col_km = tk_km * 16 + half_km * 8 + sub_tid * 2;
float2 s0 = make_float2(0.f, 0.f);
float2 s1 = make_float2(0.f, 0.f);
float2 s2 = make_float2(0.f, 0.f);
float2 s3 = make_float2(0.f, 0.f);
#pragma unroll 4
for (int r = 0; r < 64; r += 4) {
float2 v0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 0][col_km]));
float2 v1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 1][col_km]));
float2 v2 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 2][col_km]));
float2 v3 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 3][col_km]));
s0.x += v0.x; s0.y += v0.y;
s1.x += v1.x; s1.y += v1.y;
s2.x += v2.x; s2.y += v2.y;
s3.x += v3.x; s3.y += v3.y;
}
float2 sum_km = make_float2(
((s0.x + s1.x) + (s2.x + s3.x)) * scale_inv64,
((s0.y + s1.y) + (s2.y + s3.y)) * scale_inv64
);
reinterpret_cast<float2*>(out_ptr)[idx2] = sum_km;
if constexpr (D == 64) {
__syncwarp();
} else {
asm volatile("bar.sync 1, 64;\n" ::: "memory");
}
if (tid_cta == 0) {
uint32_t* flag_ptr = km_flag + bh * n_full_blocks + bi_curr;
asm volatile("st.global.release.gpu.b32 [%0], %1;\n" :: "l"(flag_ptr), "r"(gen) : "memory");
}
}
// If bi_curr >= 1, wait for K_mean[0..bi_curr-1] (if bi_curr >= 8) and prefetch K[bi_curr - 1], V[bi_curr - 1]!
if (bi_curr >= 1) {
if (bi_curr >= 8 && tid_cta < bi_curr) {
const uint32_t* flag_head = km_flag + bh * n_full_blocks;
uint32_t flg;
do {
asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(flg) : "l"(flag_head + tid_cta));
} while (flg != gen);
}
__syncthreads();
const int64_t prev_base = static_cast<int64_t>(bi_curr - 1) * 64;
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + prev_base * D, tid_cta);
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + prev_base * D, tid_cta);
cp_async_commit();
}
// Step 2: Determine selected blocks (mask0_top/bot, mask1_top/bot) using online Top-8 in registers
uint64_t mask0_top = 0ULL, mask1_top = 0ULL;
uint64_t mask0_bot = 0ULL, mask1_bot = 0ULL;
if (t_warp_base < S) {
if (bi_curr < 8) {
uint64_t m = (1ULL << bi_curr) - 1ULL;
mask0_top = valid_q_top ? m : 0ULL;
mask0_bot = valid_q_bot ? m : 0ULL;
} else {
sum_curr_top += __shfl_xor_sync(0xffffffff, sum_curr_top, 1);
sum_curr_top += __shfl_xor_sync(0xffffffff, sum_curr_top, 2);
sum_curr_bot += __shfl_xor_sync(0xffffffff, sum_curr_bot, 1);
sum_curr_bot += __shfl_xor_sync(0xffffffff, sum_curr_bot, 2);
float imp_curr_top = sum_curr_top / static_cast<float>(q_top_local + 1);
float imp_curr_bot = sum_curr_bot / static_cast<float>(q_bot_local + 1);
float ts0 = -1e30f, ts1 = -1e30f;
int tb0 = 2 * tid - 8, tb1 = 2 * tid + 1 - 8;
float min_ts = -1e30f;
int min_tb = -8;
float bs0 = -1e30f, bs1 = -1e30f;
int bb0 = 2 * tid - 8, bb1 = 2 * tid + 1 - 8;
float min_bs = -1e30f;
int min_bb = -8;
const float* km_head = K_mean + static_cast<int64_t>(bh) * n_full_blocks * D;
for (int bi = 0; bi <= bi_curr; ++bi) {
float st, sb;
if (bi < bi_curr) {
const float* km_row = km_head + bi * D;
float dot_top = 0.0f;
float dot_bot = 0.0f;
#pragma unroll
for (int tk = 0; tk < TK; ++tk) {
float2 ft0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][0]));
float2 fb0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][1]));
float2 ft1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][2]));
float2 fb1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][3]));
float4 k01;
asm volatile("ld.global.ca.v4.f32 {%0, %1, %2, %3}, [%4];\n"
: "=f"(k01.x), "=f"(k01.y), "=f"(k01.z), "=f"(k01.w)
: "l"(km_row + tk * 16 + tid * 4));
dot_top = __fmaf_rn(ft0.x, k01.x, dot_top);
dot_top = __fmaf_rn(ft0.y, k01.y, dot_top);
dot_top = __fmaf_rn(ft1.x, k01.z, dot_top);
dot_top = __fmaf_rn(ft1.y, k01.w, dot_top);
dot_bot = __fmaf_rn(fb0.x, k01.x, dot_bot);
dot_bot = __fmaf_rn(fb0.y, k01.y, dot_bot);
dot_bot = __fmaf_rn(fb1.x, k01.z, dot_bot);
dot_bot = __fmaf_rn(fb1.y, k01.w, dot_bot);
}
dot_top += __shfl_xor_sync(0xffffffff, dot_top, 1);
dot_top += __shfl_xor_sync(0xffffffff, dot_top, 2);
dot_bot += __shfl_xor_sync(0xffffffff, dot_bot, 1);
dot_bot += __shfl_xor_sync(0xffffffff, dot_bot, 2);
st = dot_top;
sb = dot_bot;
} else {
st = imp_curr_top;
sb = imp_curr_bot;
}
bool upd_t = (st >= min_ts);
bool upd_b = (sb >= min_bs);
if (__any_sync(0xffffffff, upd_t || upd_b)) {
if (upd_t) {
if (tb0 == min_tb) { ts0 = st; tb0 = bi; }
if (tb1 == min_tb) { ts1 = st; tb1 = bi; }
}
if (upd_b) {
if (bb0 == min_bb) { bs0 = sb; bb0 = bi; }
if (bb1 == min_bb) { bs1 = sb; bb1 = bi; }
}
bool t1_worse = (ts1 < ts0) || (ts1 == ts0 && tb1 < tb0);
float cur_ts = t1_worse ? ts1 : ts0;
int cur_tb = t1_worse ? tb1 : tb0;
bool b1_worse = (bs1 < bs0) || (bs1 == bs0 && bb1 < bb0);
float cur_bs = b1_worse ? bs1 : bs0;
int cur_bb = b1_worse ? bb1 : bb0;
#pragma unroll
for (int mask = 1; mask <= 2; mask <<= 1) {
float oth_ts = __shfl_xor_sync(0xffffffff, cur_ts, mask);
int oth_tb = __shfl_xor_sync(0xffffffff, cur_tb, mask);
if (oth_ts < cur_ts || (oth_ts == cur_ts && oth_tb < cur_tb)) {
cur_ts = oth_ts;
cur_tb = oth_tb;
}
float oth_bs = __shfl_xor_sync(0xffffffff, cur_bs, mask);
int oth_bb = __shfl_xor_sync(0xffffffff, cur_bb, mask);
if (oth_bs < cur_bs || (oth_bs == cur_bs && oth_bb < cur_bb)) {
cur_bs = oth_bs;
cur_bb = oth_bb;
}
}
min_ts = cur_ts;
min_tb = cur_tb;
min_bs = cur_bs;
min_bb = cur_bb;
}
}
if (valid_q_top) {
if (tb0 >= 0 && tb0 < bi_curr) {
if (tb0 < 64) mask0_top |= (1ULL << tb0);
else mask1_top |= (1ULL << (tb0 - 64));
}
if (tb1 >= 0 && tb1 < bi_curr) {
if (tb1 < 64) mask0_top |= (1ULL << tb1);
else mask1_top |= (1ULL << (tb1 - 64));
}
}
if (valid_q_bot) {
if (bb0 >= 0 && bb0 < bi_curr) {
if (bb0 < 64) mask0_bot |= (1ULL << bb0);
else mask1_bot |= (1ULL << (bb0 - 64));
}
if (bb1 >= 0 && bb1 < bi_curr) {
if (bb1 < 64) mask0_bot |= (1ULL << bb1);
else mask1_bot |= (1ULL << (bb1 - 64));
}
}
#pragma unroll
for (int mask = 1; mask <= 2; mask <<= 1) {
mask0_top |= __shfl_xor_sync(0xffffffff, mask0_top, mask);
mask1_top |= __shfl_xor_sync(0xffffffff, mask1_top, mask);
mask0_bot |= __shfl_xor_sync(0xffffffff, mask0_bot, mask);
mask1_bot |= __shfl_xor_sync(0xffffffff, mask1_bot, mask);
}
}
}
// Step 4: Handle block bi_curr - 1 (already prefetched into smem.k_tile and smem.v_tile!)
if (bi_curr >= 1) {
const int prev_bi = bi_curr - 1;
cp_async_wait<0>();
__syncthreads();
if (t_warp_base < S) {
bool has_prev_top = (prev_bi < 64) ? ((mask0_top >> prev_bi) & 1ULL) : ((mask1_top >> (prev_bi - 64)) & 1ULL);
bool has_prev_bot = (prev_bi < 64) ? ((mask0_bot >> prev_bi) & 1ULL) : ((mask1_bot >> (prev_bi - 64)) & 1ULL);
if (prev_bi < 64) {
mask0_top &= ~(1ULL << prev_bi);
mask0_bot &= ~(1ULL << prev_bi);
} else {
mask1_top &= ~(1ULL << (prev_bi - 64));
mask1_bot &= ~(1ULL << (prev_bi - 64));
}
const int min_k_top = has_prev_top ? 0 : (q_top_local + 1);
const int min_k_bot = has_prev_bot ? 0 : (q_bot_local + 1);
int min_k_warp = min(valid_q_top ? min_k_top : 64, valid_q_bot ? min_k_bot : 64);
#pragma unroll
for (int mask = 16; mask >= 4; mask >>= 1) {
min_k_warp = min(min_k_warp, __shfl_xor_sync(0xffffffff, min_k_warp, mask));
}
for (int k_chunk = (min_k_warp >> 4); k_chunk < 4; ++k_chunk) {
const int k_base = k_chunk * 16;
float s_0[4] = {0.f, 0.f, 0.f, 0.f};
float s_1[4] = {0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int tk = 0; tk < TK; ++tk) {
uint32_t r0, r1, r2, r3;
uint32_t addr = smem_u32addr(&smem.k_tile[k_base + sub_r][tk * 16 + sub_c]);
ldmatrix_x4(r0, r1, r2, r3, addr);
mma_m16n8k16_bf16(s_0[0], s_0[1], s_0[2], s_0[3],
q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r0, r2);
mma_m16n8k16_bf16(s_1[0], s_1[1], s_1[2], s_1[3],
q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r1, r3);
}
const int kc0 = k_base + tid * 2 + 0;
const int kc1 = k_base + tid * 2 + 1;
const int kc2 = k_base + 8 + tid * 2 + 0;
const int kc3 = k_base + 8 + tid * 2 + 1;
bool vt0 = valid_q_top && (kc0 >= min_k_top);
bool vt1 = valid_q_top && (kc1 >= min_k_top);
bool vt2 = valid_q_top && (kc2 >= min_k_top);
bool vt3 = valid_q_top && (kc3 >= min_k_top);
bool vb0 = valid_q_bot && (kc0 >= min_k_bot);
bool vb1 = valid_q_bot && (kc1 >= min_k_bot);
bool vb2 = valid_q_bot && (kc2 >= min_k_bot);
bool vb3 = valid_q_bot && (kc3 >= min_k_bot);
float st0 = vt0 ? (s_0[0] * scale) : -1e30f;
float st1 = vt1 ? (s_0[1] * scale) : -1e30f;
float st2 = vt2 ? (s_1[0] * scale) : -1e30f;
float st3 = vt3 ? (s_1[1] * scale) : -1e30f;
float sb0 = vb0 ? (s_0[2] * scale) : -1e30f;
float sb1 = vb1 ? (s_0[3] * scale) : -1e30f;
float sb2 = vb2 ? (s_1[2] * scale) : -1e30f;
float sb3 = vb3 ? (s_1[3] * scale) : -1e30f;
float m_chunk_top = fmaxf(fmaxf(st0, st1), fmaxf(st2, st3));
m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 1));
m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 2));
float m_chunk_bot = fmaxf(fmaxf(sb0, sb1), fmaxf(sb2, sb3));
m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 1));
m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 2));
float m_new_top = fmaxf(m_top, m_chunk_top);
float exp_old_top = (m_new_top > m_top) ? __expf(m_top - m_new_top) : 1.0f;
float pt0 = vt0 ? __expf(st0 - m_new_top) : 0.0f;
float pt1 = vt1 ? __expf(st1 - m_new_top) : 0.0f;
float pt2 = vt2 ? __expf(st2 - m_new_top) : 0.0f;
float pt3 = vt3 ? __expf(st3 - m_new_top) : 0.0f;
float p_sum_top = pt0 + pt1 + pt2 + pt3;
p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 1);
p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 2);
l_top = l_top * exp_old_top + p_sum_top;
m_top = m_new_top;
float m_new_bot = fmaxf(m_bot, m_chunk_bot);
float exp_old_bot = (m_new_bot > m_bot) ? __expf(m_bot - m_new_bot) : 1.0f;
float pb0 = vb0 ? __expf(sb0 - m_new_bot) : 0.0f;
float pb1 = vb1 ? __expf(sb1 - m_new_bot) : 0.0f;
float pb2 = vb2 ? __expf(sb2 - m_new_bot) : 0.0f;
float pb3 = vb3 ? __expf(sb3 - m_new_bot) : 0.0f;
float p_sum_bot = pb0 + pb1 + pb2 + pb3;
p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 1);
p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 2);
l_bot = l_bot * exp_old_bot + p_sum_bot;
m_bot = m_new_bot;
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
o_frag[tn][0] *= exp_old_top;
o_frag[tn][1] *= exp_old_top;
o_frag[tn][2] *= exp_old_bot;
o_frag[tn][3] *= exp_old_bot;
}
uint32_t p_frag[4];
*reinterpret_cast<__nv_bfloat162*>(&p_frag[0]) = __floats2bfloat162_rn(pt0, pt1);
*reinterpret_cast<__nv_bfloat162*>(&p_frag[1]) = __floats2bfloat162_rn(pb0, pb1);
*reinterpret_cast<__nv_bfloat162*>(&p_frag[2]) = __floats2bfloat162_rn(pt2, pt3);
*reinterpret_cast<__nv_bfloat162*>(&p_frag[3]) = __floats2bfloat162_rn(pb2, pb3);
#pragma unroll
for (int tn2 = 0; tn2 < TN / 2; ++tn2) {
uint32_t r0, r1, r2, r3;
uint32_t addr = smem_u32addr(&smem.v_tile[k_base + sub_r][tn2 * 16 + sub_c]);
ldmatrix_x4_trans(r0, r1, r2, r3, addr);
mma_m16n8k16_bf16(o_frag[2 * tn2 + 0][0], o_frag[2 * tn2 + 0][1], o_frag[2 * tn2 + 0][2], o_frag[2 * tn2 + 0][3],
p_frag[0], p_frag[1], p_frag[2], p_frag[3], r0, r1);
mma_m16n8k16_bf16(o_frag[2 * tn2 + 1][0], o_frag[2 * tn2 + 1][1], o_frag[2 * tn2 + 1][2], o_frag[2 * tn2 + 1][3],
p_frag[0], p_frag[1], p_frag[2], p_frag[3], r2, r3);
}
}
}
}
// Step 3: Cooperative CTA union across all 4 warps with software-pipelined cp.async.ca!
uint64_t w_mask0 = mask0_top | mask0_bot;
uint64_t w_mask1 = mask1_top | mask1_bot;
#pragma unroll
for (int mask = 16; mask >= 4; mask >>= 1) {
w_mask0 |= __shfl_xor_sync(0xffffffff, w_mask0, mask);
w_mask1 |= __shfl_xor_sync(0xffffffff, w_mask1, mask);
}
if (lane == 0) {
smem.warp_mask0[w] = w_mask0;
smem.warp_mask1[w] = w_mask1;
}
__syncthreads();
uint64_t cta_mask0 = smem.warp_mask0[0] | smem.warp_mask0[1] | smem.warp_mask0[2] | smem.warp_mask0[3];
uint64_t cta_mask1 = smem.warp_mask1[0] | smem.warp_mask1[1] | smem.warp_mask1[2] | smem.warp_mask1[3];
if (cta_mask0 != 0ULL || cta_mask1 != 0ULL) {
int bi;
bool active_top, active_bot, warp_active;
if (cta_mask0 != 0ULL) {
bi = __ffsll(cta_mask0) - 1;
cta_mask0 &= (cta_mask0 - 1ULL);
active_top = (mask0_top >> bi) & 1ULL;
active_bot = (mask0_bot >> bi) & 1ULL;
warp_active = (w_mask0 >> bi) & 1ULL;
} else {
int bit = __ffsll(cta_mask1) - 1;
bi = 64 + bit;
cta_mask1 &= (cta_mask1 - 1ULL);
active_top = (mask1_top >> bit) & 1ULL;
active_bot = (mask1_bot >> bit) & 1ULL;
warp_active = (w_mask1 >> bit) & 1ULL;
}
// Prologue: issue K[bi_0] and V[bi_0]
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + static_cast<int64_t>(bi) * 64 * D, tid_cta);
cp_async_commit();
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + static_cast<int64_t>(bi) * 64 * D, tid_cta);
cp_async_commit();
// Wait for K[bi_0] (group 1), leaving V[bi_0] (group 0) in flight!
cp_async_wait<1>();
__syncthreads();
while (true) {
SPFrag sp;
if (warp_active) {
// Compute S = Q @ K[bi].T while V[bi] is in flight!
#pragma unroll
for (int kc = 0; kc < 4; ++kc) {
const int k_base = kc * 16;
float s_0[4] = {0.f, 0.f, 0.f, 0.f};
float s_1[4] = {0.f, 0.f, 0.f, 0.f};
#pragma unroll
for (int tk = 0; tk < TK; ++tk) {
uint32_t r0, r1, r2, r3;
uint32_t addr = smem_u32addr(&smem.k_tile[k_base + sub_r][tk * 16 + sub_c]);
ldmatrix_x4(r0, r1, r2, r3, addr);
mma_m16n8k16_bf16(s_0[0], s_0[1], s_0[2], s_0[3],
q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r0, r2);
mma_m16n8k16_bf16(s_1[0], s_1[1], s_1[2], s_1[3],
q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r1, r3);
}
sp.s[2 * kc + 0][0] = s_0[0] * scale;
sp.s[2 * kc + 0][1] = s_0[1] * scale;
sp.s[2 * kc + 0][2] = s_0[2] * scale;
sp.s[2 * kc + 0][3] = s_0[3] * scale;
sp.s[2 * kc + 1][0] = s_1[0] * scale;
sp.s[2 * kc + 1][1] = s_1[1] * scale;
sp.s[2 * kc + 1][2] = s_1[2] * scale;
sp.s[2 * kc + 1][3] = s_1[3] * scale;
}
float m_blk_top = -1e30f;
float m_blk_bot = -1e30f;
#pragma unroll
for (int tn = 0; tn < 8; ++tn) {
m_blk_top = fmaxf(m_blk_top, fmaxf(sp.s[tn][0], sp.s[tn][1]));
m_blk_bot = fmaxf(m_blk_bot, fmaxf(sp.s[tn][2], sp.s[tn][3]));
}
m_blk_top = fmaxf(m_blk_top, __shfl_xor_sync(0xffffffff, m_blk_top, 1));
m_blk_top = fmaxf(m_blk_top, __shfl_xor_sync(0xffffffff, m_blk_top, 2));
m_blk_bot = fmaxf(m_blk_bot, __shfl_xor_sync(0xffffffff, m_blk_bot, 1));
m_blk_bot = fmaxf(m_blk_bot, __shfl_xor_sync(0xffffffff, m_blk_bot, 2));
float m_new_top = active_top ? fmaxf(m_top, m_blk_top) : m_top;
float exp_old_top = active_top ? __expf(m_top - m_new_top) : 1.0f;
float m_new_bot = active_bot ? fmaxf(m_bot, m_blk_bot) : m_bot;
float exp_old_bot = active_bot ? __expf(m_bot - m_new_bot) : 1.0f;
float p_sum_top = 0.0f;
float p_sum_bot = 0.0f;
#pragma unroll
for (int kc = 0; kc < 4; ++kc) {
float pt0 = active_top ? __expf(sp.s[2 * kc + 0][0] - m_new_top) : 0.0f;
float pt1 = active_top ? __expf(sp.s[2 * kc + 0][1] - m_new_top) : 0.0f;
float pt2 = active_top ? __expf(sp.s[2 * kc + 1][0] - m_new_top) : 0.0f;
float pt3 = active_top ? __expf(sp.s[2 * kc + 1][1] - m_new_top) : 0.0f;
p_sum_top += pt0 + pt1 + pt2 + pt3;
float pb0 = active_bot ? __expf(sp.s[2 * kc + 0][2] - m_new_bot) : 0.0f;
float pb1 = active_bot ? __expf(sp.s[2 * kc + 0][3] - m_new_bot) : 0.0f;
float pb2 = active_bot ? __expf(sp.s[2 * kc + 1][2] - m_new_bot) : 0.0f;
float pb3 = active_bot ? __expf(sp.s[2 * kc + 1][3] - m_new_bot) : 0.0f;
p_sum_bot += pb0 + pb1 + pb2 + pb3;
uint32_t pf0, pf1, pf2, pf3;
*reinterpret_cast<__nv_bfloat162*>(&pf0) = __floats2bfloat162_rn(pt0, pt1);
*reinterpret_cast<__nv_bfloat162*>(&pf1) = __floats2bfloat162_rn(pb0, pb1);
*reinterpret_cast<__nv_bfloat162*>(&pf2) = __floats2bfloat162_rn(pt2, pt3);
*reinterpret_cast<__nv_bfloat162*>(&pf3) = __floats2bfloat162_rn(pb2, pb3);
sp.p[kc][0] = pf0;
sp.p[kc][1] = pf1;
sp.p[kc][2] = pf2;
sp.p[kc][3] = pf3;
}
p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 1);
p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 2);
l_top = l_top * exp_old_top + p_sum_top;
m_top = m_new_top;
p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 1);
p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 2);
l_bot = l_bot * exp_old_bot + p_sum_bot;
m_bot = m_new_bot;
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
o_frag[tn][0] *= exp_old_top;
o_frag[tn][1] *= exp_old_top;
o_frag[tn][2] *= exp_old_bot;
o_frag[tn][3] *= exp_old_bot;
}
}
// Wait for V[bi] to arrive and ensure all warps are done reading smem.k_tile
cp_async_wait<0>();
__syncthreads();
const bool has_next = (cta_mask0 != 0ULL || cta_mask1 != 0ULL);
int next_bi = 0;
bool next_active_top = false, next_active_bot = false, next_warp_active = false;
if (has_next) {
if (cta_mask0 != 0ULL) {
next_bi = __ffsll(cta_mask0) - 1;
cta_mask0 &= (cta_mask0 - 1ULL);
next_active_top = (mask0_top >> next_bi) & 1ULL;
next_active_bot = (mask0_bot >> next_bi) & 1ULL;
next_warp_active = (w_mask0 >> next_bi) & 1ULL;
} else {
int bit = __ffsll(cta_mask1) - 1;
next_bi = 64 + bit;
cta_mask1 &= (cta_mask1 - 1ULL);
next_active_top = (mask1_top >> bit) & 1ULL;
next_active_bot = (mask1_bot >> bit) & 1ULL;
next_warp_active = (w_mask1 >> bit) & 1ULL;
}
// Issue K[next_bi] into smem.k_tile WHILE warps compute P @ V[bi]!
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + static_cast<int64_t>(next_bi) * 64 * D, tid_cta);
cp_async_commit();
}
if (warp_active) {
#pragma unroll
for (int kc = 0; kc < 4; ++kc) {
const int k_base = kc * 16;
#pragma unroll
for (int tn2 = 0; tn2 < TN / 2; ++tn2) {
uint32_t r0, r1, r2, r3;
uint32_t addr = smem_u32addr(&smem.v_tile[k_base + sub_r][tn2 * 16 + sub_c]);
ldmatrix_x4_trans(r0, r1, r2, r3, addr);
mma_m16n8k16_bf16(o_frag[2 * tn2 + 0][0], o_frag[2 * tn2 + 0][1], o_frag[2 * tn2 + 0][2], o_frag[2 * tn2 + 0][3],
sp.p[kc][0], sp.p[kc][1], sp.p[kc][2], sp.p[kc][3], r0, r1);
mma_m16n8k16_bf16(o_frag[2 * tn2 + 1][0], o_frag[2 * tn2 + 1][1], o_frag[2 * tn2 + 1][2], o_frag[2 * tn2 + 1][3],
sp.p[kc][0], sp.p[kc][1], sp.p[kc][2], sp.p[kc][3], r2, r3);
}
}
}
if (!has_next) break;
// Wait for K[next_bi] and ensure all warps are done reading smem.v_tile
cp_async_wait<0>();
__syncthreads();
// Issue V[next_bi] into smem.v_tile WHILE warps compute Q @ K[next_bi].T!
issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + static_cast<int64_t>(next_bi) * 64 * D, tid_cta);
cp_async_commit();
bi = next_bi;
active_top = next_active_top;
active_bot = next_active_bot;
warp_active = next_warp_active;
}
}
// Step 5: Normalize o_frag by l_top / l_bot, stage via smem.k_tile for coalesced uint4 global writes
__syncthreads();
if (t_warp_base < S) {
float inv_l_top = (l_top > 0.0f) ? (1.0f / l_top) : 0.0f;
float inv_l_bot = (l_bot > 0.0f) ? (1.0f / l_bot) : 0.0f;
#pragma unroll
for (int tn = 0; tn < TN; ++tn) {
__nv_bfloat162 b_top = __floats2bfloat162_rn(o_frag[tn][0] * inv_l_top, o_frag[tn][1] * inv_l_top);
__nv_bfloat162 b_bot = __floats2bfloat162_rn(o_frag[tn][2] * inv_l_bot, o_frag[tn][3] * inv_l_bot);
*reinterpret_cast<__nv_bfloat162*>(&smem.k_tile[w * 16 + groupID][tn * 8 + tid * 2]) = b_top;
*reinterpret_cast<__nv_bfloat162*>(&smem.k_tile[w * 16 + groupID + 8][tn * 8 + tid * 2]) = b_bot;
}
}
__syncthreads();
#pragma unroll
for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
int idx = tid_cta + i * 128;
int r = idx / UINT4_PER_ROW;
int c = (idx % UINT4_PER_ROW) * 8;
int t_r = s0 + r;
if (t_r < S) {
*reinterpret_cast<uint4*>(o_head + static_cast<int64_t>(t_r) * D + c) =
*reinterpret_cast<const uint4*>(&smem.k_tile[r][c]);
}
}
}
struct KMeanBuffers {
float* km_ptr = nullptr;
uint32_t* flag_ptr = nullptr;
uint32_t gen = 0;
};
KMeanBuffers& get_km_buffers(size_t needed_blocks, int D) {
static KMeanBuffers bufs;
static size_t cap_floats = 0;
static size_t cap_flags = 0;
size_t needed_floats = needed_blocks * D;
if (needed_floats > cap_floats || needed_blocks > cap_flags) {
if (bufs.km_ptr) cudaFree(bufs.km_ptr);
if (bufs.flag_ptr) cudaFree(bufs.flag_ptr);
cap_flags = std::max(needed_blocks, static_cast<size_t>(32 * 128));
cap_floats = std::max(needed_floats, static_cast<size_t>(32 * 128 * 128));
cudaMalloc(&bufs.km_ptr, cap_floats * sizeof(float));
cudaMalloc(&bufs.flag_ptr, cap_flags * sizeof(uint32_t));
cudaMemset(bufs.flag_ptr, 0, cap_flags * sizeof(uint32_t));
}
++bufs.gen;
return bufs;
}
} // namespace
torch::Tensor nsa_forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
const int B = q.size(0);
const int H = q.size(1);
const int S = q.size(2);
const int D = q.size(3);
const int BH = B * H;
const int n_full_blocks = S / 64;
const int n_blocks = (S + 63) / 64;
const float scale = 1.0f / std::sqrt(static_cast<float>(D));
auto o = at::detail::empty_cuda(q.sizes(), at::ScalarType::BFloat16, q.device(), std::nullopt);
auto& bufs = get_km_buffers(static_cast<size_t>(BH) * std::max(n_full_blocks, 1), D);
const auto* q_ptr = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
const auto* k_ptr = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
const auto* v_ptr = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
auto* o_ptr = reinterpret_cast<__nv_bfloat16*>(o.data_ptr());
dim3 grid_attn(n_blocks, BH);
dim3 block_attn(128);
if (D == 64) {
nsa_cta_coop_kernel<64, 4, 3><<<grid_attn, block_attn>>>(
q_ptr, k_ptr, v_ptr, bufs.km_ptr, bufs.flag_ptr, bufs.gen, o_ptr, S, n_full_blocks, scale
);
} else if (D == 128) {
nsa_cta_coop_kernel<128, 4, 2><<<grid_attn, block_attn>>>(
q_ptr, k_ptr, v_ptr, bufs.km_ptr, bufs.flag_ptr, bufs.gen, o_ptr, S, n_full_blocks, scale
);
}
return o;
}
"""
cpp_source = "torch::Tensor nsa_forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);"
_ext = load_inline(
name="nsa_cuda_ext_v10",
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=["nsa_forward_cuda"],
extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_120a"],
verbose=False,
)
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:
return _ext.nsa_forward_cuda(q, k, v)
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]
20261013_175245_agy_gemini-4-argon-high_02_deepseek_nsa