Sparse attention on Blackwell
A CUDA kernel that makes DeepSeek-style sparse attention fly on Blackwell. Passed every workload in NVIDIA's competition, dozens of times faster than the reference.
NVIDIA MLSys 2026 competitioncudallmcode private
The problem
DeepSeek-style sparse attention lets each query read only a small set of a long context’s keys instead of all of them. That removes most of the arithmetic, and replaces it with scattered, irregular memory reads.
How it got fast
Four versions, each answering the bottleneck the previous one exposed:
- Tensor Cores — make the math fast.
- Overlap the gather — stop waiting on scattered loads.
- Split the work — the change that mattered. The first versions gave the GPU only 8 blocks of work for its 148 SMs, so most of the chip sat idle. Splitting along the keys filled the whole thing.
- Autotune — measure each shape instead of guessing.
Result
It passed every workload in the competition, dozens of times faster than the PyTorch reference.
The full walkthrough, with a diagram for each version, is in my ML notes.