[wave] DeepSeek Native Sparse Attention (NSA) kernels for MI350
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
# DeepSeek Native Sparse Attention (NSA) for MI350
Implement DeepSeek's Native Sparse Attention mechanism (arxiv:2502.11089) as wave kernels targeting MI350 (CDNA4), covering both inference and training.
## Background
NSA is a hardware-aligned sparse attention mechanism with three parallel branches:
1. **Compressed attention** — coarse-grained token compression via mean pooling of KV into blocks, then standard attention over compressed representations
2. **Selection attention** — fine-grained top-k block selection using compressed attention scores, followed by gather-based sparse attention over selected KV blocks
3. **Sliding window attention** — local context via fixed-size causal window
The three branch outputs are combined via learned gating weights: `o = g_cmp * o_cmp + g_slc * o_slc + g_swa * o_swa`
## Performance targets (64k context, DeepSeek-V3 config)
- Decoding: ~11x speedup over dense attention
- Forward: ~9x speedup over dense attention
- Backward: ~6x speedup over dense attention
## Reference implementations
- DeepSeek paper: https://arxiv.org/abs/2502.11089
- tilde-research/nsa-impl (Triton reference)
- flash-linear-attention (FLA) `fla.ops.nsa`
## Ticket breakdown
### Research & Design
- [ ] #1244 — Architecture mapping & design doc for MI350
### Inference kernels
- [ ] #1245 — Token compression (mean pooling KV blocks)
- [ ] #1246 — Compressed attention forward kernel
- [ ] #1247 — Top-k block selection kernel
- [ ] #1248 — Selection attention forward kernel
- [ ] #1249 — Sliding window attention forward kernel
- [ ] #1250 — Gated output combination kernel
- [ ] #1251 — End-to-end NSA inference pipeline
### Training kernels
- [ ] #1252 — Selection attention backward kernel (dQ, dK, dV)
- [ ] #1253 — Compressed attention backward kernel
- [ ] #1254 — Gating parameter gradients & top-k STE backward
- [ ] #1255 — End-to-end NSA training integration with autograd
### MI350-specific optimizations
- [ ] #1256 — CDNA4 wavefront scheduling for sparse gather patterns
- [ ] #1257 — LDS optimization for block-gather in selection attention
- [ ] #1258 — GQA-aware memory layout & tiling for MI350 cache hierarchy
### Testing & integration
- [ ] #1259 — Numerical correctness tests vs dense attention reference
- [ ] #1260 — Gradient correctness tests (backward pass)
- [ ] #1261 — Performance benchmarks on MI350
- [ ] #1262 — Water dialect lowering for NSA composite ops
- [ ] #1263 — SGLang / inference backend integration
Contributor guide
Assessment
This issue has not been assessed yet.