iree-org / iree-org/wave

[wave] NSA: selection attention forward kernel

Open
#1,248 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 forward pass of the selection attention branch — the core sparse attention kernel in NSA. Given top-k block indices, this kernel gathers the selected KV blocks and computes attention only over those blocks.

### Operation

```
Input: Q [B, M, H, D], K [B, N, G, D], V [B, N, G, D], block_indices [B, M, G, T]
Output: O_slc [B, M, H, D], LSE_slc [B, H, M]
```

Where T = block_count (number of selected blocks per query).

### Algorithm
For each query position m:
1. Load query vector q[m]
2. For each selected block index t in block_indices[m]:
- Gather K, V for tokens [t*block_size : (t+1)*block_size]
- Compute QK^T, apply causal mask, softmax
- Accumulate weighted V using online softmax (log-sum-exp trick)
3. Store output and LSE

### Requirements

- GQA support: heads within a group share block_indices
- Online softmax across all T blocks (numerically stable via running max + LSE)
- Causal masking within gathered blocks
- FP16 matmuls, FP32 accumulation
- Return LSE (needed for backward pass)
- Support both `block_size=64` and `block_size=32`

### MI350 considerations
- **This is the most performance-critical kernel** — it determines the speedup over dense attention
- Block-gather pattern is irregular — plan for LDS staging of gathered KV blocks
- Each program instance handles one (batch, query_pos, gqa_group) — parallelize across these
- BLOCK_H dimension processes all heads in a GQA group together (register tiling)
- Inner loop over T blocks — unroll for small T (16) if register pressure allows
- The `tl.multiple_of(col, SELECTION_BLOCK_SIZE)` hint in the reference is critical for coalesced loads on MI350

### Performance target
- For B=1, M=64k, T=16, block_size=64, G=8, D=128: should achieve >8x speedup over dense FA at same seqlen

### Depends on
- #1244 (design doc)
- #1247 (top-k selection kernel, for block_indices)

### References
- `_sel_attn_fwd_kernel` in tilde-research/nsa-impl/nsa/selection.py
- NSA paper Section 3.2

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.