iree-org / iree-org/wave

[wave] NSA: SGLang / inference backend integration

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

Integrate the NSA wave kernels into the SGLang inference backend (or equivalent serving framework) to enable end-to-end serving of DeepSeek models with native sparse attention on MI350.

### Integration points

1. **Attention backend registration**
- Register NSA as a new attention backend alongside existing FA v3 (#142)
- Backend selection based on model config: if model uses NSA (e.g., DeepSeek-V3), auto-select NSA backend
- Fallback to dense FA if NSA is not supported for a given config

2. **KV cache integration**
- NSA needs the full KV cache for: compressed attention (mean pool all of KV), selection attention (gather selected blocks), and sliding window
- Integrate with existing paged KV cache: NSA block indices map to KV cache page indices
- Incremental updates: when new tokens are appended, update compressed KV incrementally (only re-pool the last block)

3. **Decode optimization**
- In autoregressive decoding (M=1), NSA reduces to:
a. Pool the new token into the running compressed KV (update last block mean)
b. Score compressed blocks with the new query → top-k indices
c. Gather selected KV blocks + sliding window → sparse attention
- Block indices may have high temporal locality between adjacent decode steps — consider caching

4. **Prefill optimization**
- Long-prompt prefill runs full NSA with M=N
- Chunked prefill: process prompt in chunks, maintaining running compressed KV state

5. **Model support**
- DeepSeek-V3 / V3.2 (primary target)
- Any model using NSA-style attention (configurable parameters)

### Depends on
- #1251 (inference pipeline)
- #142 (FA v3 — sliding window component)
- SGLang wave backend (#97, #477)

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.