SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

feat: enable glm5 router gemm (#3185)

<!-- .github/pull_request_template.md -->

## ๐Ÿ“Œ Description

It uses hidden_dim = 6144, need to add it to allowlist (ideally could
have router gemm API to avoid duplication), but this is faster for now.

On the performance:

### GLM-5 Router GEMM Benchmark (Flashinfer vs Torch) with CUPTI
benchmark

| M (Tokens) | N (Experts) | K (Dim) | Torch (ms) | Flashinfer (ms) |
Winner | Speedup |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| 1 | 256 | 6144 | 6.048 | 2.976 | Flashinfer | +103.2% |
| 2 | 256 | 6144 | 6.272 | 3.137 | Flashinfer | +99.9% |
| 3 | 256 | 6144 | 6.272 | 3.584 | Flashinfer | +75.0% |
| 4 | 256 | 6144 | 6.304 | 3.872 | Flashinfer | +62.8% |
| 5 | 256 | 6144 | 6.304 | 4.032 | Flashinfer | +56.4% |
| 6 | 256 | 6144 | 6.304 | 4.480 | Flashinfer | +40.7% |
| 7 | 256 | 6144 | 6.304 | 4.480 | Flashinfer | +40.7% |
| 8 | 256 | 6144 | 6.304 | 4.832 | Flashinfer | +30.5% |
| 9 | 256 | 6144 | 7.040 | 5.216 | Flashinfer | +35.0% |
| 10 | 256 | 6144 | 7.040 | 5.408 | Flashinfer | +30.2% |
| 11 | 256 | 6144 | 7.040 | 7.616 | torch | -7.6% |
| 12 | 256 | 6144 | 7.040 | 7.936 | torch | -11.3% |
| 13 | 256 | 6144 | 7.040 | 8.128 | torch | -13.4% |
| 14 | 256 | 6144 | 7.040 | 8.288 | torch | -15.1% |
| 15 | 256 | 6144 | 7.072 | 8.544 | torch | -17.2% |
| 16 | 256 | 6144 | 6.400 | 8.768 | torch | -27.0% |

<!-- What does this PR do? Briefly describe the changes and why theyโ€™re
needed. -->

<!-- Link any related issues here -->

## ๐Ÿš€ 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

Add new hidden dim

```
pytest tests/model_optimizations/test_dsv3_router_gemm.py
============================================================================ test session starts =============================================================================
platform linux -- Python 3.12.3, pytest-9.0.3, pluggy-1.6.0
rootdir: /sgl-workspace/flashinfer
configfile: pytest.ini
plugins: anyio-4.13.0, typeguard-4.5.1
collected 57 items                                                                                                                                                           

tests/model_optimizations/test_dsv3_router_gemm.py .........................................................                                                           [100%]

============================================================================= 57 passed in 2.44s =============================================================================
```

- [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**
* Added support for a new GLM-MoE-DSA router GEMM configuration
(6144ร—256) and its trace/reference implementation.

* **Improvements**
* Generalized router GEMM shape validation and added new dispatch/export
pathway for the 6144ร—256 variant.
* Changed defaults for existing router GEMM launch behavior to enable
optimized dispatch.
* Replaced low-level GPU sync primitives with updated synchronization
calls for CUDA 9.0+.

* **Tests**
* Expanded unit tests to cover the new 6144ร—256 kernel and additional
negative validation cases.

<!-- 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/3185)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
B
Brayden Zhong committed
ee31370cd4927da525e3c157d5baf6a78a091a49
Parent: 885d0cf
Committed by GitHub <noreply@github.com> on 5/14/2026, 10:49:44 PM