iree-org / iree-org/wave

[wave] NSA: end-to-end inference pipeline

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

Wire up all NSA inference kernels into a single callable pipeline that replaces dense attention in a transformer layer.

### Pipeline stages (in order)

1. **Mean pool K, V** → K_cmp, V_cmp (#1245)
2. **Compressed attention** → O_cmp, LSE_cmp (#1246)
3. **Top-k block selection** using LSE_cmp → block_indices (#1247)
4. **Selection attention** using block_indices → O_slc (#1248)
5. **Sliding window attention** → O_swa (#1249)
6. **Gated combination** → final O (#1250)

### Requirements

- Single Python-callable entry point: `nsa_forward(q, k, v, g_cmp, g_slc, g_swa, config) -> o`
- Config dataclass with: block_size, block_count, window_size, scale
- Stages 2+4+5 can run concurrently (compressed attn, then selection after top-k; sliding window is independent)
- KV cache integration: support incremental KV append for autoregressive decoding
- Support both prefill (long Q) and decode (Q length 1) modes

### Decode-specific optimizations
- In decode mode (M=1), the compressed attention + top-k path is even more critical since we only have 1 query
- Block indices can potentially be cached/reused across adjacent decode steps if KV hasn't changed much
- Selection attention with M=1 is a single vector-matrix multiply per block — very different perf profile than prefill

### Depends on
- All inference kernel tickets: #1245, #1246, #1247, #1248, #1249, #1250

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.