iree-org / iree-org/wave

[wave] NSA: token compression kernel (mean pooling KV blocks)

Open
#1,245 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 a wave kernel that compresses K and V tensors by mean-pooling consecutive tokens into blocks. This is the first stage of the NSA compressed attention branch.

### Operation

```
Input: K, V of shape [B, N, G, D]
Output: K_cmp, V_cmp of shape [B, N // block_size, G, D]
```

For each block of `block_size` consecutive tokens along the sequence dimension, compute the element-wise mean to produce a single compressed token.

### Requirements

- Support configurable `block_size` (default 64, also support 32, 128)
- Support GQA layout: G KV heads, H query heads where H = G * HEADS_PER_GROUP
- Handle sequence lengths not evenly divisible by block_size (pad or truncate last block)
- FP16 input with FP32 accumulation for numerical stability
- Target MI350 CDNA4 — leverage vector ALU for the reduction

### Performance target
- Should be memory-bandwidth bound; target close to 2× memread bandwidth (reading K + V once)
- For B=1, N=64k, G=8, D=128, block_size=64: ~0.1ms on MI350

### Depends on
- #1244 (design doc)

### References
- `mean_pooling` in fla.ops.utils.pooling
- NSA paper Section 3.1

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.