[wave] NSA: compressed attention forward 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 forward pass of the compressed attention branch. This computes standard causal attention over the mean-pooled (compressed) KV representations, producing both the output and log-sum-exp (LSE) values needed by the top-k selection stage.
### Operation
```
Input: Q [B, M, H, D], K_cmp [B, N//bs, G, D], V_cmp [B, N//bs, G, D]
Output: O_cmp [B, M, H, D], LSE_cmp [B, H, M]
```
Causal masking: `q_idx <= (kv_idx + 1) * block_size - 1`
### Requirements
- GQA-compatible (G KV groups, H query heads)
- Return LSE alongside output (needed for top-k scoring in selection branch)
- Causal block mask aligned to compression block boundaries
- FP16 matmuls, FP32 softmax accumulation
- Must produce numerically stable LSE values
### MI350 considerations
- The compressed sequence length is N/block_size (e.g., 1024 for 64k/64), making this a relatively small attention — optimize for this regime
- Consider using matrix cores for the QK^T and PV matmuls even at this reduced size
- May be able to fuse with the mean pooling from the compression kernel
### Depends on
- #1244 (design doc)
- #1245 (token compression kernel)
### References
- `compression_attention` in tilde-research/nsa-impl uses FlexAttention
- NSA paper Section 3.1
Contributor guide
Assessment
This issue has not been assessed yet.