SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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

[![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/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