[wave] NSA: sliding window 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 or adapt a sliding window causal attention kernel for the NSA sliding window branch. This provides local context awareness and is the simplest of the three NSA branches.
### Operation
```
Input: Q [B, M, H, D], K [B, N, G, D], V [B, N, G, D], window_size
Output: O_swa [B, M, H, D]
```
Standard causal attention but each query only attends to the preceding `window_size` tokens (including itself).
### Requirements
- Configurable window_size (default 512, also support 256, 1024)
- GQA-compatible
- Causal + windowed masking: attend to tokens in range [max(0, m - window_size + 1), m]
- FP16 compute, FP32 softmax accumulation
### Implementation options
1. **Adapt existing FA v3 kernel** (preferred if #142 has window_size support or is close)
- FlashAttention-2/3 already supports `window_size` parameter
- Just need to wire up the parameter in the wave kernel interface
2. **Standalone kernel** (fallback)
- Simpler than full FA since the attention window is small and fixed
- Can tile the window into blocks matching MI350 matrix core size
### MI350 considerations
- Window_size=512 with D=128 means each query touches 512 KV pairs — fits comfortably in LDS
- This kernel should be heavily compute-bound on matrix cores
- Overlap KV loads for next tile with current tile's MFMA
### Depends on
- #1244 (design doc)
- #142 (FA v3 — check if window_size param already works)
### References
- `flash_attn_func(..., window_size=(window_size-1, 0))` in nsa-impl
- NSA paper Section 3.3
Contributor guide
Assessment
This issue has not been assessed yet.