feat: integrate cute-dsl Blackwell GQA decode into BatchDecodeWithPagedKVCacheWrapper (#3360)
<!-- .github/pull_request_template.md -->
## ๐ Description
- Adds a new `cute-dsl` backend to `BatchDecodeWithPagedKVCacheWrapper`
powered by the cute-dsl Blackwell (SM100a) GQA decode kernel, supporting
- NHD + HND paged + non-paged KV layout
- fp8/bf16 QKVO
- headdim multiple of 64 (large headdims supported)
- speculative-decode (q_len_per_req > 1 + causal mask)
- runtime PDL
- skip softmax with threshold scale factor matching trtllm behavior
- per-tensor scale factors (folded into softmax scale and output scale)
- LSE return
- three reduction modes:`kernel` (deterministic split-K + reduction
kernel), `atomic` (cluster reduction + L2 atomic adds), and `none`
(kv_splits == 1, no flash-decoding).
- Wraps the kernel in two PyTorch-facing classes
(`BatchDecodeCuteDSLWrapper`, `BatchDecodePagedCuteDSLWrapper`) modeled
on the prefill cute-dsl pattern. Compile is memoized via
`@functools.cache`; symbolic dims (`cute.sym_int`) cover runtime
variation in batch_size, seq_len, and prediction.
- Bug fix in trtllm-gen `plan()`: `block_id = indptr[0]` was a 0-dim
view, and `block_id += ...` did an in-place add that mutated
`indptr[0]`. Switched to a Python int read from the host-side indptr to
avoid corrupting shared buffers.
- Replace references to deprecated `cute.nvgpu.tcgen05.OperandMajorMode`
with `cute.nvgpu.OperandMajorMode` introduced in cute-dsl 4.5.0
- `BatchDecodeWithPagedKVCacheWrapper` API change: `q_len_per_req` is
now accepted in `plan` instead of `run`, to provide a tile size hint to
the `cute-dsl` compilation. Passing to `run` emits deprecation notice
and is used to validate the expected shape of `q`. Otherwise
`q_len_per_req` is inferred from the shape of `q` at runtime.
## ๐ Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.
### โ
Pre-commit Checks
- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.
> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).
## ๐งช Tests
- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).
- [x] `pytest tests/attention/test_cute_dsl_decode.py` on Blackwell โ
covers
both reduction modes, paged + ragged wrappers, LSE, speculative decode,
v_scale, HND vs NHD layout, and runtime-q_len mismatch.
- [x] `python benchmarks/bench_cute_dsl_decode.py` โ sanity-check
throughput
against fa2 / trtllm-gen.
## Perf results B200
```
b= 1 mtp=0 s= 1024 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 12.243 us 0.3 TB/s
trtllm-gen: 6.301 us 0.7 TB/s
cute-dsl: 4.862 us 0.9 TB/s
b= 1 mtp=0 s= 4096 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 18.184 us 0.9 TB/s
trtllm-gen: 8.128 us 2.1 TB/s
cute-dsl: 11.002 us 1.5 TB/s
b= 1 mtp=0 s=16384 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 44.368 us 1.5 TB/s
trtllm-gen: 17.782 us 3.8 TB/s
cute-dsl: 18.137 us 3.7 TB/s
b= 8 mtp=0 s= 1024 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 28.691 us 1.2 TB/s
trtllm-gen: 10.577 us 3.2 TB/s
cute-dsl: 8.952 us 3.8 TB/s
b= 8 mtp=0 s= 4096 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 85.133 us 1.6 TB/s
trtllm-gen: 24.139 us 5.6 TB/s
cute-dsl: 22.920 us 5.9 TB/s
b= 8 mtp=0 s=16384 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 300.016 us 1.8 TB/s
trtllm-gen: 78.230 us 6.9 TB/s
cute-dsl: 77.682 us 6.9 TB/s
b=64 mtp=0 s= 1024 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 159.955 us 1.7 TB/s
trtllm-gen: 43.420 us 6.2 TB/s
cute-dsl: 45.634 us 5.9 TB/s
b=64 mtp=0 s= 4096 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 589.830 us 1.8 TB/s
trtllm-gen: 156.094 us 6.9 TB/s
cute-dsl: 152.400 us 7.1 TB/s
b=64 mtp=0 s=16384 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
fa2: 2314.402 us 1.9 TB/s
trtllm-gen: 606.548 us 7.1 TB/s
cute-dsl: 598.255 us 7.2 TB/s
b= 1 mtp=3 s= 1024 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 6.703 us 0.6 TB/s
cute-dsl: 6.693 us 0.6 TB/s
b= 1 mtp=3 s= 4096 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 10.593 us 1.6 TB/s
cute-dsl: 11.871 us 1.4 TB/s
b= 1 mtp=3 s=16384 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 19.029 us 3.5 TB/s
cute-dsl: 21.917 us 3.1 TB/s
b= 8 mtp=3 s= 1024 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 11.857 us 2.9 TB/s
cute-dsl: 11.616 us 3.0 TB/s
b= 8 mtp=3 s= 4096 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 25.721 us 5.3 TB/s
cute-dsl: 29.520 us 4.6 TB/s
b= 8 mtp=3 s=16384 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 80.154 us 6.7 TB/s
cute-dsl: 90.475 us 5.9 TB/s
b=64 mtp=3 s= 1024 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 49.531 us 5.6 TB/s
cute-dsl: 63.087 us 4.4 TB/s
b=64 mtp=3 s= 4096 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 162.000 us 6.7 TB/s
cute-dsl: 180.057 us 6.0 TB/s
b=64 mtp=3 s=16384 pg=16 h_q=64 h_kv=8 d=128 dtype=torch.bfloat16
trtllm-gen: 621.074 us 6.9 TB/s
cute-dsl: 690.734 us 6.2 TB/s
```
## Reviewer Notes
<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added a "cute-dsl" decode backend for Blackwell GPUs with
high-performance grouped-query attention decoding, plus two user-facing
wrappers for ragged and paged KV-cache decoding; supports LSE output,
kernel/atomic/none reductions, speculative multi-token decode,
softmax-skip tuning, workspace reuse, and optional user-provided
outputs.
* **Benchmarks**
* Added a CUDA-gated paged-batch GQA decode benchmark exercising
multiple backends, dtypes, batch/seq sizes and reporting
latency/bandwidth.
* **Tests**
* Added comprehensive integration and correctness tests covering decode
paths, LSE, reduction modes, softmax-skip behavior, and workspace
handling.
<!-- review_stack_entry_start -->
[](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3360?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)
<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Mingyang Wang <35635157+saltyminty@users.noreply.github.com> R
Richard Cai committed
bff85f3459707d5d2f1426d1ded4a320ab142078
Parent: 2f372e2
Committed by GitHub <noreply@github.com>
on 5/22/2026, 9:36:34 PM