[wave] NSA: CDNA4 wavefront scheduling for sparse gather patterns
- 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
Optimize wavefront scheduling and occupancy for the irregular memory access patterns in NSA's selection attention kernel on MI350 (CDNA4).
### Problem
Selection attention gathers non-contiguous KV blocks based on per-query top-k indices. This creates:
- Irregular global memory access patterns (poor cache utilization)
- Variable work per wavefront (some queries select nearby blocks, others select distant ones)
- Potential wavefront stalls waiting for memory
### Tasks
1. **Gather coalescing analysis**
- Profile actual memory access patterns for representative block_indices distributions
- Measure global memory bandwidth utilization vs theoretical peak
- Identify whether block_size=64 aligns well with MI350 cache lines (64B) and memory channels
2. **Wavefront occupancy tuning**
- Determine optimal number of wavefronts per CU for selection attention
- Trade off register usage (more live state = fewer wavefronts) vs memory latency hiding (more wavefronts)
- MI350 has 64KB VGPR per SIMD, 128 VGPRs per wavefront at full occupancy — find the sweet spot
3. **Prefetch strategies**
- Software prefetch: issue global loads for the next block's KV while computing attention on the current block
- LDS staging: load entire selected KV blocks into LDS before attention computation
- Double-buffering: overlap loads and computation across blocks
4. **Workgroup sizing**
- Current grid: (B, M, G) — one program per query position per GQA group
- Consider coarsening: multiple query positions per workgroup to amortize block index loads
- Consider splitting: if T is large, split blocks across workgroups with a reduction step
### Depends on
- #1248 (selection attention forward — baseline to optimize)
- #1244 (design doc — roofline analysis)
Contributor guide
Assessment
This issue has not been assessed yet.