[wave] NSA: gated output combination kernel
- 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 kernel that combines the three NSA branch outputs using learned gating weights.
### Operation
```
Input:
O_cmp [B, M, H, D] — compressed attention output
O_slc [B, M, H, D] — selection attention output
O_swa [B, M, H, D] — sliding window attention output
g_cmp [B, M, H] — compressed gate (sigmoid output)
g_slc [B, M, H] — selection gate
g_swa [B, M, H] — sliding window gate
Output:
O [B, M, H, D] = g_cmp.unsqueeze(-1) * O_cmp + g_slc.unsqueeze(-1) * O_slc + g_swa.unsqueeze(-1) * O_swa
```
### Requirements
- Element-wise fused multiply-add across three branches
- FP16 input/output, can use FP16 throughout (simple pointwise)
- Consider fusing with the final projection linear layer if applicable
- Backward: compute gradients for all 6 inputs (3 outputs + 3 gates)
### MI350 considerations
- Pure bandwidth-bound pointwise kernel
- Fuse with downstream operations (e.g., residual add, layer norm) if possible to avoid extra memory round-trips
- Consider fusing this into the epilog of the most expensive branch (selection attention) to save a kernel launch
### Depends on
- #1246 (compressed attention forward)
- #1248 (selection attention forward)
- #1249 (sliding window attention forward)
Contributor guide
Assessment
This issue has not been assessed yet.