[wave] NSA: end-to-end inference pipeline
- 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
Wire up all NSA inference kernels into a single callable pipeline that replaces dense attention in a transformer layer.
### Pipeline stages (in order)
1. **Mean pool K, V** → K_cmp, V_cmp (#1245)
2. **Compressed attention** → O_cmp, LSE_cmp (#1246)
3. **Top-k block selection** using LSE_cmp → block_indices (#1247)
4. **Selection attention** using block_indices → O_slc (#1248)
5. **Sliding window attention** → O_swa (#1249)
6. **Gated combination** → final O (#1250)
### Requirements
- Single Python-callable entry point: `nsa_forward(q, k, v, g_cmp, g_slc, g_swa, config) -> o`
- Config dataclass with: block_size, block_count, window_size, scale
- Stages 2+4+5 can run concurrently (compressed attn, then selection after top-k; sliding window is independent)
- KV cache integration: support incremental KV append for autoregressive decoding
- Support both prefill (long Q) and decode (Q length 1) modes
### Decode-specific optimizations
- In decode mode (M=1), the compressed attention + top-k path is even more critical since we only have 1 query
- Block indices can potentially be cached/reused across adjacent decode steps if KV hasn't changed much
- Selection attention with M=1 is a single vector-matrix multiply per block — very different perf profile than prefill
### Depends on
- All inference kernel tickets: #1245, #1246, #1247, #1248, #1249, #1250
Contributor guide
Assessment
This issue has not been assessed yet.