[wave] NSA: top-k block selection kernel
- 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
Implement the top-k block index selection kernel that identifies which KV blocks are most relevant for each query position. This bridges the compressed attention branch to the selection attention branch.
### Operation
```
Input: Q [B, M, H, D], K_cmp [B, N//bs, G, D], LSE_cmp [B, H, M]
Output: block_indices [B, M, G, block_count] (int32 indices into compressed blocks)
```
For each query position, score all compressed blocks and select the top `block_count` blocks. The scoring uses the attention logits from the compressed branch, adjusted by the LSE for numerical stability.
### Requirements
- Support configurable `block_count` (default 16)
- Produce per-GQA-group indices (all heads in a group share the same selected blocks)
- Causal: only select blocks where `block_start <= query_position`
- Indices must be sorted (ascending) for coalesced memory access in selection attention
- INT32 output indices
### MI350 considerations
- Top-k is inherently sequential — consider approximate top-k (e.g., block-parallel partial sorts with merge) for better wavefront utilization
- The compressed sequence is short (N/block_size), so this kernel is likely latency-bound rather than throughput-bound
- LDS can hold the full score vector for sorting when N/block_size <= 1024
### Depends on
- #1244 (design doc)
- #1246 (compressed attention forward, for LSE)
### References
- `parallel_nsa_topk` in fla.ops.nsa.parallel
- NSA paper Section 3.2
Contributor guide
Assessment
This issue has not been assessed yet.