[Feat] Add num_heads < 128 support for mla decode kernel (#3309)
<!-- .github/pull_request_template.md -->
## ๐ Description
Features addition:
1. Add num_heads < 128 support, fold seqlen_q into num_heads to fill the
M tile of MLA.
2. return lse support.
Performance enhancement:
1. mainly improve the fp8 mla decode.
Performance collection **before & after** this commit with MTP==4.
Blackwell B200, no freq lock. seqlen_k = 80k.
```
num_heads=16 (F=4 fold)
bs Before This MR speedup
1 44 30 1.46x
2 72 36 1.99x
4 138 53 2.60x
8 263 85 3.09x
16 524 152 3.44x
32 1084 300 3.61x
64 2256 623 3.62x
96 3507 1161 3.02x
128 4624 1403 3.30x
192 7789 1999 3.90x
256 9386 2915 3.22x
384 15215 4440 3.43x
512 21023 6097 3.45x
num_heads=32 (F=4 fold)
bs Before This MR speedup
1 46 31 1.50x
4 141 55 2.59x
8 268 90 2.99x
16 532 157 3.38x
32 1111 316 3.52x
64 2309 655 3.53x
96 3702 1196 3.10x
128 4493 1345 3.34x
192 7524 2137 3.52x
256 10628 3028 3.51x
num_heads=64 (F=2 fold)
bs Before This MR speedup
1 79 33 2.41x
4 167 76 2.20x
8 278 146 1.90x
16 545 290 1.88x
32 1144 581 1.97x
64 2354 1248 1.89x
96 3766 1973 1.91x
128 5224 2652 1.97x
192 8211 3947 2.08x
256 10515 5486 1.92x
num_heads=128 (F=1, no fold โ sanity check)
bs Before This MR speedup
1 48 46 1.05x
4 152 143 1.06x
8 296 280 1.06x
16 580 545 1.07x
32 1357 1157 1.17x
64 2729 2496 1.09x
96 4137 3791 1.09x
128 5191 4863 1.07x
192 8186 7474 1.10x
256 11432 10167 1.12x
```
## ๐ Related Issues
https://github.com/flashinfer-ai/flashinfer/issues/3287
## ๐ 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.).
## 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**
* Optional log-sum-exp (LSE) output for attention APIs.
* Spec-decoding (MTP) causal masking support in applicable attention
paths.
* Token-folding optimization for Blackwell GPUs to improve throughput.
* **Tests**
* Expanded coverage across both attention implementations, LSE behavior,
masking, sinks/dispatch, and end-to-end API scenarios (including
expected error cases).
* **Chores**
* Standardized benchmark q_len grid across backends.
<!-- review_stack_entry_start -->
[](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3309?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: Prasun Gera <prasun.gera@gmail.com> O
Observer007 committed
2cf8f4cc7aacd49a442b072ed2d60c73e32007f1
Parent: 2bcb37a
Committed by GitHub <noreply@github.com>
on 5/29/2026, 12:26:08 AM