iree-org / iree-org/wave

[wave] NSA: gating parameter gradients & top-k STE backward

Open
#1,254 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

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

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.