SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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 -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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