[wave] NSA: LDS optimization for block-gather in selection attention
- 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
Design and implement LDS (Local Data Share) staging strategies for the selection attention kernel's KV block gathering on MI350.
### Problem
Selection attention gathers `T * block_size` KV pairs per query, where blocks are non-contiguous in global memory. Direct global memory access for each QK^T and PV multiply leads to redundant loads and poor bandwidth utilization.
### LDS staging strategies
1. **Per-block LDS staging**
- For each block index t, cooperatively load K[t*bs:(t+1)*bs, :] and V[t*bs:(t+1)*bs, :] into LDS
- All heads in the GQA group then read from LDS for QK^T and PV
- LDS requirement per block: 2 * block_size * D * sizeof(FP16) = 2 * 64 * 128 * 2 = 32KB
- MI350 has 64KB LDS per workgroup — fits 2 blocks simultaneously (double-buffer)
2. **Shared block deduplication**
- Within a workgroup processing multiple query positions, some may select the same blocks
- Detect overlapping block indices and load each unique block only once
- Requires coordination between threads — use LDS for the dedup index
3. **K transposed layout in LDS**
- Store K in LDS as [D, block_size] (transposed) for efficient QK^T = Q @ K^T without runtime transpose
- V stored as [block_size, D] (natural layout) for PV = P @ V
### MI350 LDS specifics
- 64KB LDS per workgroup, 32 banks, 4 bytes per bank
- Avoid bank conflicts: pad K/V in LDS if D is a multiple of 32 (it often is — D=128)
- Add 1 float16 padding per row: stride = D + 1
- LDS read bandwidth: 32 × 4B × frequency per cycle per CU
### Backward pass considerations
- Backward kernel also needs gathered K, V, and additionally stores dK, dV atomically
- LDS can stage partial dK/dV accumulations before global atomic add (reduce atomic contention)
### Depends on
- #1248 (selection attention forward)
- #1252 (selection attention backward)
- #1256 (wavefront scheduling — determines workgroup structure)
Contributor guide
Assessment
This issue has not been assessed yet.