[wave] NSA: token compression kernel (mean pooling KV blocks)
- 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 a wave kernel that compresses K and V tensors by mean-pooling consecutive tokens into blocks. This is the first stage of the NSA compressed attention branch.
### Operation
```
Input: K, V of shape [B, N, G, D]
Output: K_cmp, V_cmp of shape [B, N // block_size, G, D]
```
For each block of `block_size` consecutive tokens along the sequence dimension, compute the element-wise mean to produce a single compressed token.
### Requirements
- Support configurable `block_size` (default 64, also support 32, 128)
- Support GQA layout: G KV heads, H query heads where H = G * HEADS_PER_GROUP
- Handle sequence lengths not evenly divisible by block_size (pad or truncate last block)
- FP16 input with FP32 accumulation for numerical stability
- Target MI350 CDNA4 — leverage vector ALU for the reduction
### Performance target
- Should be memory-bandwidth bound; target close to 2× memread bandwidth (reading K + V once)
- For B=1, N=64k, G=8, D=128, block_size=64: ~0.1ms on MI350
### Depends on
- #1244 (design doc)
### References
- `mean_pooling` in fla.ops.utils.pooling
- NSA paper Section 3.1
Contributor guide
Assessment
This issue has not been assessed yet.