[wave] NSA: SGLang / inference backend integration
- Dominant language
- Python
- Stars
- 59
- Forks
- 32
- PR merge metrics
- No merged PRs in 30d
Description
## Parent
Part of #1243 — DeepSeek NSA kernels for MI350
## Description
Integrate the NSA wave kernels into the SGLang inference backend (or equivalent serving framework) to enable end-to-end serving of DeepSeek models with native sparse attention on MI350.
### Integration points
1. **Attention backend registration**
- Register NSA as a new attention backend alongside existing FA v3 (#142)
- Backend selection based on model config: if model uses NSA (e.g., DeepSeek-V3), auto-select NSA backend
- Fallback to dense FA if NSA is not supported for a given config
2. **KV cache integration**
- NSA needs the full KV cache for: compressed attention (mean pool all of KV), selection attention (gather selected blocks), and sliding window
- Integrate with existing paged KV cache: NSA block indices map to KV cache page indices
- Incremental updates: when new tokens are appended, update compressed KV incrementally (only re-pool the last block)
3. **Decode optimization**
- In autoregressive decoding (M=1), NSA reduces to:
a. Pool the new token into the running compressed KV (update last block mean)
b. Score compressed blocks with the new query → top-k indices
c. Gather selected KV blocks + sliding window → sparse attention
- Block indices may have high temporal locality between adjacent decode steps — consider caching
4. **Prefill optimization**
- Long-prompt prefill runs full NSA with M=N
- Chunked prefill: process prompt in chunks, maintaining running compressed KV state
5. **Model support**
- DeepSeek-V3 / V3.2 (primary target)
- Any model using NSA-style attention (configurable parameters)
### Depends on
- #1251 (inference pipeline)
- #142 (FA v3 — sliding window component)
- SGLang wave backend (#97, #477)
Contributor guide
Assessment
This issue has not been assessed yet.