[wave] NSA: gating parameter gradients & top-k STE backward
- 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 backward pass for the gating combination and handle the gradient flow through the top-k block selection (which is non-differentiable and requires a straight-through estimator or similar).
### Part 1: Gating gradients
Forward: `O = g_cmp * O_cmp + g_slc * O_slc + g_swa * O_swa`
Backward from dO:
```
dO_cmp = dO * g_cmp.unsqueeze(-1)
dO_slc = dO * g_slc.unsqueeze(-1)
dO_swa = dO * g_swa.unsqueeze(-1)
dg_cmp = (dO * O_cmp).sum(dim=-1)
dg_slc = (dO * O_slc).sum(dim=-1)
dg_swa = (dO * O_swa).sum(dim=-1)
```
This is straightforward pointwise + reduction.
### Part 2: Top-k selection gradient
The `top-k` operation in the block selection stage is non-differentiable. Options:
1. **Straight-Through Estimator (STE)**: Pass gradients through as if top-k were identity on selected blocks, zero on unselected. This is what the paper uses.
2. **Gumbel-Softmax relaxation**: Differentiable approximation of discrete selection (more complex, potentially better gradients)
3. **Detach selection indices**: Treat block_indices as fixed during backward (simplest, used in reference impl). Gradients only flow through the attention computation given fixed indices, not through the selection itself.
### Requirements
- Gating backward should be fused into a single kernel (6 outputs from 7 inputs)
- For top-k: implement option 3 (detached indices) first, with option 1 (STE) as a follow-up
- Gate gradients feed into the gating MLP's backward pass (upstream)
### Depends on
- #1250 (gated output combination forward)
- #1247 (top-k selection)
Contributor guide
Assessment
This issue has not been assessed yet.