iree-org / iree-org/wave

[wave] NSA: selection attention backward kernel (dQ, dK, dV)

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

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.