feat: Expose unpacked topk weights for routed moe (fp4) (#2425)
<!-- .github/pull_request_template.md -->
## ๐ Description
Summary
Add support for pre-computed routing with unpacked format in
trtllm_fp4_block_scale_routed_moe, and
improve code clarity with explicit routing mode enum.
Changes
New Feature: Unpacked Pre-computed Routing (Mode 3)
trtllm_fp4_block_scale_routed_moe now accepts routing input in two
formats:
# Packed format (existing behavior)
trtllm_fp4_block_scale_routed_moe(packed_tensor, ...)
# Unpacked format (new - pass tuple)
trtllm_fp4_block_scale_routed_moe((topk_ids, topk_weights), ...)
This is a backwards-compatible enhancement - existing code continues to
work unchanged.
Code Clarity Improvements
- Added RoutingInputMode enum in both C++ and Python to explicitly
define three routing modes:
- FromLogits: Compute routing from logits (Mode 1)
- PackedPrecomputed: Pre-computed with packed (score << 16 | id) format
(Mode 2)
- UnpackedPrecomputed: Pre-computed with separate tensors (Mode 3)
- Renamed expert_weights โ topk_weights in public APIs for consistency
(MLP weights like
fc1_expert_weights unchanged)
- Added routing_input_mode as first parameter to internal C++ function
for explicit mode selection
Files Modified
- flashinfer/fused_moe/core.py - Python API with RoutingInputMode enum
and tuple support
- csrc/trtllm_fused_moe_kernel_launcher.cu - C++ launcher with
RoutingInputMode enum
- tests/moe/test_trtllm_gen_routed_fused_moe.py - Extended tests for
both formats
- INTEGRATION_STATUS.md - Documentation (to be removed after merge)
Test plan
- test_trtllm_gen_fused_moe.py - All passed (Mode 1: FromLogits)
- test_trtllm_gen_routed_fused_moe.py with routing_format=packed - All
passed (Mode 2)
- test_trtllm_gen_routed_fused_moe.py with routing_format=unpacked - All
passed (Mode 3)
pytest tests/moe/test_trtllm_gen_routed_fused_moe.py -v
๐ค Generated with https://claude.ai/code (fully reviewed)
## ๐ Related Issues
#2373
## ๐ 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.).
```
$ pytest tests/moe/test_trtllm_gen_routed_fused_moe.py
3456 passed in 336.66s (0:05:36)
```
## 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 RoutingInputMode with three modes and end-to-end support for
packed or unpacked precomputed routing inputs; intermediate routed
outputs now use topk_weights.
* **Public API**
* Core APIs now accept precomputed routing IDs/weights and an explicit
routing mode to enable mode-aware routing paths.
* **Tests**
* Parameterized tests added for "packed" and "unpacked" routing formats.
* **Documentation**
* Docstrings and return descriptions updated to reflect new routing
formats and topk_weights semantics.
<!-- end of auto-generated comment: release notes by coderabbit.ai --> A
Alex Yang committed
719ee23fd82cb220d51ad118ca60198718f6c9d1
Parent: a1b8a60
Committed by GitHub <noreply@github.com>
on 5/16/2026, 12:24:00 AM