iree-org / iree-org/wave

[wave] DeepSeek Native Sparse Attention (NSA) kernels for MI350

Open
#1,243 0 comments 0 reactions 0 assignees View on GitHub
nsa tracker
Dominant language
Python
Stars
59
Forks
32
PR merge metrics
No merged PRs in 30d

Description

# DeepSeek Native Sparse Attention (NSA) for MI350

Implement DeepSeek's Native Sparse Attention mechanism (arxiv:2502.11089) as wave kernels targeting MI350 (CDNA4), covering both inference and training.

## Background

NSA is a hardware-aligned sparse attention mechanism with three parallel branches:
1. **Compressed attention** — coarse-grained token compression via mean pooling of KV into blocks, then standard attention over compressed representations
2. **Selection attention** — fine-grained top-k block selection using compressed attention scores, followed by gather-based sparse attention over selected KV blocks
3. **Sliding window attention** — local context via fixed-size causal window

The three branch outputs are combined via learned gating weights: `o = g_cmp * o_cmp + g_slc * o_slc + g_swa * o_swa`

## Performance targets (64k context, DeepSeek-V3 config)
- Decoding: ~11x speedup over dense attention
- Forward: ~9x speedup over dense attention
- Backward: ~6x speedup over dense attention

## Reference implementations
- DeepSeek paper: https://arxiv.org/abs/2502.11089
- tilde-research/nsa-impl (Triton reference)
- flash-linear-attention (FLA) `fla.ops.nsa`

## Ticket breakdown

### Research & Design
- [ ] #1244 — Architecture mapping & design doc for MI350

### Inference kernels
- [ ] #1245 — Token compression (mean pooling KV blocks)
- [ ] #1246 — Compressed attention forward kernel
- [ ] #1247 — Top-k block selection kernel
- [ ] #1248 — Selection attention forward kernel
- [ ] #1249 — Sliding window attention forward kernel
- [ ] #1250 — Gated output combination kernel
- [ ] #1251 — End-to-end NSA inference pipeline

### Training kernels
- [ ] #1252 — Selection attention backward kernel (dQ, dK, dV)
- [ ] #1253 — Compressed attention backward kernel
- [ ] #1254 — Gating parameter gradients & top-k STE backward
- [ ] #1255 — End-to-end NSA training integration with autograd

### MI350-specific optimizations
- [ ] #1256 — CDNA4 wavefront scheduling for sparse gather patterns
- [ ] #1257 — LDS optimization for block-gather in selection attention
- [ ] #1258 — GQA-aware memory layout & tiling for MI350 cache hierarchy

### Testing & integration
- [ ] #1259 — Numerical correctness tests vs dense attention reference
- [ ] #1260 — Gradient correctness tests (backward pass)
- [ ] #1261 — Performance benchmarks on MI350
- [ ] #1262 — Water dialect lowering for NSA composite ops
- [ ] #1263 — SGLang / inference backend integration

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.