iree-org / iree-org/wave

[wave] NSA: architecture mapping & design doc for MI350

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

Produce a design document that maps the NSA algorithm onto the MI350 (CDNA4) architecture and the wave compiler infrastructure. This is the prerequisite for all implementation tickets.

### Deliverables

1. **Algorithm breakdown** — formalize the three-branch NSA dataflow (compressed, selection, sliding window) with tensor shapes at each stage for representative configs (e.g., DeepSeek-V3: B=1, H=128, G=8 GQA groups, D=128, block_size=64, block_count=16, window_size=512, seqlen=64k)

2. **MI350 hardware mapping**
- Map each branch to CDNA4 compute units (wavefronts, matrix cores, vector ALU)
- Identify memory-bound vs compute-bound phases
- Plan LDS usage for block-gather in selection attention
- Determine optimal wavefront occupancy targets per kernel

3. **Wave compiler integration plan**
- Which ops can lower through existing water dialect paths (e.g., sliding window reusing FA v3 with window param)
- Which ops need new water/wave primitives (block selection gather, top-k, mean pooling)
- Register pressure analysis for selection attention forward/backward kernels

4. **GQA considerations** — how `HEADS_PER_GROUP` tiling interacts with MI350's matrix core shapes (likely 16×16 or 32×32)

5. **Numerical precision plan** — FP16 accumulation strategy, FP32 softmax reduction, atomic add strategy for backward dK/dV

6. **Performance model** — back-of-envelope roofline analysis for each kernel at 64k context on MI350

### References
- arxiv:2502.11089 (NSA paper)
- tilde-research/nsa-impl (Triton reference)
- Existing FA v3 tracker: #142

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.