[wave] NSA: selection attention backward kernel (dQ, dK, dV)
- 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 the backward pass for the selection attention kernel. This is the most complex backward kernel in NSA due to the irregular gather pattern and the need for atomic accumulations into dK and dV.
### Operation
```
Input:
Q [B, M, H, D], K [B, N, G, D], V [B, N, G, D]
block_indices [B, M, G, T]
O [B, M, H, D], LSE [B, H, M] (saved from forward)
dO [B, M, H, D] (upstream gradient)
Output:
dQ [B, M, H, D]
dK [B, N, G, D]
dV [B, N, G, D]
```
### Algorithm
Two-stage backward (following the reference implementation):
**Stage 1: Preprocess**
- Compute Delta[b,h,m] = sum_d(O[b,m,h,d] * dO[b,m,h,d]) — the row-wise dot product
**Stage 2: Main backward** (parallelized over B × M × G)
For each query position m and each selected block t:
1. Recompute attention weights: `p = softmax(Q[m] @ K[gathered]^T / scale)`
2. `dV[gathered] += p^T @ dO[m]` (atomic add — multiple query positions write to same KV positions)
3. `dS = p * (dO[m] @ V[gathered]^T - Delta[m])` (the softmax backward)
4. `dQ[m] += dS @ K[gathered]`
5. `dK[gathered] += dS^T @ Q[m]` (atomic add)
### Requirements
- **Atomic adds for dK, dV**: multiple query positions select overlapping KV blocks, requiring atomic FP32 accumulation
- Two separate kernels: preprocess (Delta computation) and main backward
- FP32 accumulation throughout backward, cast dK/dV to FP16 at the end
- Must match the forward kernel's numerical behavior exactly (same causal masking, same softmax scale)
- Support both one-pass (atomic, simpler) and two-pass (deterministic, recompute-heavy) variants
### MI350 considerations
- **Atomic FP32 adds are the primary bottleneck** — MI350's global memory atomics are slower than NVIDIA's
- Consider LDS-based partial accumulation within a workgroup before global atomic
- Alternative: deterministic two-pass approach that bins query positions by their selected blocks
- Register pressure is high: need Q, K, V, dO, O, LSE, Delta, dQ, dK, dV all live
- The backward preprocess kernel is simple (pointwise) — fuse with upstream if possible
### Performance target
- Backward should be within 6x of dense attention backward at 64k context (matching paper claims)
### Depends on
- #1248 (selection attention forward — must match numerics exactly)
- #1244 (design doc)
### References
- `_sel_attn_bwd_kernel` and `_sel_attn_bwd_preprocess_kernel` in tilde-research/nsa-impl/nsa/selection.py
- NSA paper Section 4 (training)
Contributor guide
Assessment
This issue has not been assessed yet.