iree-org / iree-org/wave

[wave] NSA: gated output combination kernel

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

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.