iree-org / iree-org/wave

[wave] NSA: compressed attention backward kernel

Open
#1,253 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 compressed attention branch. This includes gradients through both the attention over compressed KV and the mean-pooling compression itself.

### Operation

```
Forward recap:
K_cmp, V_cmp = mean_pool(K, V, block_size)
O_cmp, LSE_cmp = attention(Q, K_cmp, V_cmp, causal_block_mask)

Backward:
dO_cmp [B, M, H, D] (from gating backward)
→ dQ_cmp [B, M, H, D]
→ dK_cmp [B, N//bs, G, D], dV_cmp [B, N//bs, G, D]
→ dK [B, N, G, D], dV [B, N, G, D] (scatter mean-pool gradients back)
```

### Two sub-stages

1. **Attention backward** (dQ_cmp, dK_cmp, dV_cmp from dO_cmp)
- Standard attention backward over compressed sequence
- Compressed sequence is short (N/block_size), so this is a small attention backward

2. **Mean-pool backward** (dK, dV from dK_cmp, dV_cmp)
- Gradient of mean pooling: `dK[n] = dK_cmp[n // block_size] / block_size` for each n
- Simple scatter/broadcast operation

### Requirements

- Fuse the two stages if beneficial (attention backward + mean-pool scatter)
- Must also produce dLSE_cmp for the top-k selection gradient path
- FP32 accumulation in attention backward

### MI350 considerations
- The compressed attention backward operates on a small sequence (N/block_size)
- The mean-pool backward is pure bandwidth — good candidate for fusion
- Consider whether the small attention backward is better done via a general-purpose kernel or a specialized one

### Depends on
- #1246 (compressed attention forward)
- #1245 (token compression kernel)
- #1244 (design doc)

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.