iree-org / iree-org/wave

[wave] NSA: end-to-end training integration with autograd

Open
#1,255 0 comments 0 reactions 0 assignees View on GitHub
enhancement nsa
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.