iree-org / iree-org/wave

[wave] NSA: sliding window attention forward kernel

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

Implement or adapt a sliding window causal attention kernel for the NSA sliding window branch. This provides local context awareness and is the simplest of the three NSA branches.

### Operation

```
Input: Q [B, M, H, D], K [B, N, G, D], V [B, N, G, D], window_size
Output: O_swa [B, M, H, D]
```

Standard causal attention but each query only attends to the preceding `window_size` tokens (including itself).

### Requirements

- Configurable window_size (default 512, also support 256, 1024)
- GQA-compatible
- Causal + windowed masking: attend to tokens in range [max(0, m - window_size + 1), m]
- FP16 compute, FP32 softmax accumulation

### Implementation options

1. **Adapt existing FA v3 kernel** (preferred if #142 has window_size support or is close)
- FlashAttention-2/3 already supports `window_size` parameter
- Just need to wire up the parameter in the wave kernel interface

2. **Standalone kernel** (fallback)
- Simpler than full FA since the attention window is small and fixed
- Can tile the window into blocks matching MI350 matrix core size

### MI350 considerations
- Window_size=512 with D=128 means each query touches 512 KV pairs — fits comfortably in LDS
- This kernel should be heavily compute-bound on matrix cores
- Overlap KV loads for next tile with current tile's MFMA

### Depends on
- #1244 (design doc)
- #142 (FA v3 — check if window_size param already works)

### References
- `flash_attn_func(..., window_size=(window_size-1, 0))` in nsa-impl
- NSA paper Section 3.3

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.