NSA kernel rung: flat VRAM to 32k tokens

The native sparse attention rung is in. VRAM now stays flat out to 32k tokens instead of climbing with sequence length, a 45% reduction at the long end, and gradient parity against the dense path is verified.

What changed

The attention path now selects a sparse block set per query instead of materializing the full score matrix. Memory is bounded by the block budget, not the context length.

import torch
from openbox_llm import sparse_attention

q, k, v = (torch.randn(1, 32, 32768, 128, device="cuda") for _ in range(3))
out = sparse_attention(q, k, v, block_size=64, budget=512)
print(out.shape)  # torch.Size([1, 32, 32768, 128])

Next

← all notes