iree-org / iree-org/wave

[wave] NSA: CDNA4 wavefront scheduling for sparse gather patterns

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

Optimize wavefront scheduling and occupancy for the irregular memory access patterns in NSA's selection attention kernel on MI350 (CDNA4).

### Problem

Selection attention gathers non-contiguous KV blocks based on per-query top-k indices. This creates:
- Irregular global memory access patterns (poor cache utilization)
- Variable work per wavefront (some queries select nearby blocks, others select distant ones)
- Potential wavefront stalls waiting for memory

### Tasks

1. **Gather coalescing analysis**
- Profile actual memory access patterns for representative block_indices distributions
- Measure global memory bandwidth utilization vs theoretical peak
- Identify whether block_size=64 aligns well with MI350 cache lines (64B) and memory channels

2. **Wavefront occupancy tuning**
- Determine optimal number of wavefronts per CU for selection attention
- Trade off register usage (more live state = fewer wavefronts) vs memory latency hiding (more wavefronts)
- MI350 has 64KB VGPR per SIMD, 128 VGPRs per wavefront at full occupancy — find the sweet spot

3. **Prefetch strategies**
- Software prefetch: issue global loads for the next block's KV while computing attention on the current block
- LDS staging: load entire selected KV blocks into LDS before attention computation
- Double-buffering: overlap loads and computation across blocks

4. **Workgroup sizing**
- Current grid: (B, M, G) — one program per query position per GQA group
- Consider coarsening: multiple query positions per workgroup to amortize block index loads
- Consider splitting: if T is large, split blocks across workgroups with a reduction step

### Depends on
- #1248 (selection attention forward — baseline to optimize)
- #1244 (design doc — roofline analysis)

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.