SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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