[wave] NSA: end-to-end training integration with autograd
- 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
Wire up all NSA forward and backward kernels into a PyTorch-compatible autograd Function that enables end-to-end training.
### Deliverables
1. **`NSAFunction(torch.autograd.Function)`** with:
- `forward()`: runs the full NSA inference pipeline (#1251), saves tensors for backward
- `backward()`: orchestrates all backward kernels in correct order
2. **Backward orchestration order**:
```
dO → gating backward → dO_cmp, dO_slc, dO_swa, dg_* (#1254)
↓ ↓ ↓
compressed attn bwd selection attn bwd sliding window bwd
→ dQ_cmp, dK_cmp, → dQ_slc, dK_slc, → dQ_swa, dK_swa,
dV_cmp (#1253) dV_slc (#1252) dV_swa (FA bwd)
↓
mean-pool bwd
→ dK_pool, dV_pool
Aggregate: dQ = dQ_cmp + dQ_slc + dQ_swa
dK = dK_pool + dK_slc + dK_swa
dV = dV_pool + dV_slc + dV_swa
dg_cmp, dg_slc, dg_swa (from gating bwd)
```
3. **Memory optimization**:
- Activation checkpointing: recompute compressed KV and block_indices in backward instead of saving
- Gradient accumulation support for pipeline/data parallelism
4. **Convenience wrapper**: `nsa_attention(q, k, v, g_cmp, g_slc, g_swa, config) -> o` that calls NSAFunction.apply
### Testing
- `torch.autograd.gradcheck` with FP64 reference
- Compare training loss curves: NSA vs dense attention on a small model (should match within tolerance)
### Depends on
- All backward kernel tickets: #1252, #1253, #1254
- Inference pipeline: #1251
Contributor guide
Assessment
This issue has not been assessed yet.