feat: trtllm FP4 routed-MoE: accept fp32 topk_weights (copy-free) - #3763
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds float32 ChangesFP4 MoE UnpackedPrecomputed float32 weights
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for float32 routing weights in the TRT-LLM FP4 block scale MoE kernel when using the unpacked precomputed routing mode. This includes adding validation checks, updating the Python API documentation and assertions, and expanding test coverage to include unpacked_fp32 routing format. The reviewer points out that setting args->mDtypeExpW conditionally only for UnpackedPrecomputed mode could cause correctness issues in other routing modes, and suggests setting it unconditionally based on topk_weights.dtype().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/trtllm_fused_moe_kernel_launcher.cu`:
- Around line 1820-1823: The `mDtypeExpW` update in
`trtllm_fused_moe_kernel_launcher.cu` only covers
`RoutingInputMode::UnpackedPrecomputed`, but `trtllm_fp4_block_scale_moe_op()`
also builds `topk_weights` from `routing_logits` in the `FromLogits` path.
Update the launcher logic around `args->mDtypeExpW` to also set the
exponent-weight dtype based on `topk_weights.dtype()` when routing comes from
logits, so fp32 logits correctly propagate `btg::Dtype::Fp32` instead of falling
back to the bf16 default.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: be72cfa3-fae3-4174-8b90-a542417d984e
📥 Commits
Reviewing files that changed from the base of the PR and between eb79c53 and 814a3a8ea4b39dc88d3fdd02d048b16565462c9f.
📒 Files selected for processing (3)
csrc/trtllm_fused_moe_kernel_launcher.cuflashinfer/fused_moe/core.pytests/moe/test_trtllm_gen_routed_fused_moe.py
bf6f828 to
cf51af3
Compare
|
/bot run tests/moe |
|
[FAILED] Pipeline #56703724: 11/20 passed |
cf51af3 to
84c9623
Compare
84c9623 to
a025362
Compare
|
/bot run tests/moe |
|
[FAILED] Pipeline #58082989: 11/20 passed |
a025362 to
e8f80b3
Compare
|
/bot run tests/moe |
|
[FAILED] Pipeline #58251976: 1/20 passed |
e8f80b3 to
79eba7c
Compare
|
/bot run tests/moe |
|
[FAILED] Pipeline #58416904: 13/20 passed |
…expert ids Signed-off-by: jdebache <jdebache@nvidia.com>
79eba7c to
f6f81cf
Compare
|
/bot run tests/moe |
|
[FAILED] Pipeline #58742486: 11/20 passed |
aleozlx
left a comment
aleozlx
left a comment
There was a problem hiding this comment.
change looks to me
adds one test, which sounds fine
The packed `(expert_id << 16) | bf16(weight)` topk_ids format truncates the fp32 routing weights callers already hold. Take a `(topk_ids, topk_weights)` pair instead, as the FP4 runner has since flashinfer-ai#3763, and pick `mDtypeExpW` from the caller's weights dtype so fp32 reaches the combine. `routing_input_mode` now moves to `FusedMoeLauncher`, which both entry points pass through explicitly; FP4 drops its own copy of the field.
`trtllm_fp8_block_scale_routed_moe` and `trtllm_bf16_routed_moe` only take the packed `(expert_id << 16) | bf16(weight)` form, which truncates the fp32 routing weights callers already have. However FP4 already handles this. #3763 gave the FP4 path`(topk_ids, topk_weights)` form; this does the same for the FP8 block-scale and BF16 runners. And unified support in #4104 Now fp32 weights reach the combine as fp32. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added logits-based, packed precomputed, and unpacked precomputed routing for BF16, FP8, and FP4 Mixture-of-Experts operations. - Routed APIs now accept packed expert IDs or separate expert IDs and weights. - Precomputed routing supports FP32 and BF16 expert weights where applicable. - **Bug Fixes** - Improved validation and handling of unpacked routing inputs. - **Tests** - Added FP8 routing coverage for multiple expert-weight formats. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Brayden Zhong <brayden@radixark.ai>
…4448) ## 📌 Description Add unpacked FP32 routing coverage for BF16, FP8, FP4, and MXFP8 routed MoE paths. ## 🔍 Related Issues Testing to address this comment: vllm-project/vllm#46872 (review). ## 🚀 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. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes None. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Expanded routed mixture-of-experts coverage to support both packed and unpacked FP32 routing inputs. * Added coverage for BF16, FP8 block-scale, and FP4 block-scale implementations. * Improved activation-parity checks to report the routing format when mismatches occur. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: jdebache <jdebache@nvidia.com>
5 participants
📌 Description
trtllm_fp4_block_scale_routed_moe's unpacked routing path(
topk_ids=(ids, weights)) silently assumed bf16topk_weights. vLLMrouters emit fp32 by design, so the finalize kernel reinterpreted the fp32
bytes as bf16 → ~1e36/NaN per-expert scales → garbage output).
Fix
FP4BlockScaleLauncher::prepare_routing()— setargs->mDtypeExpWfrom theactual
topk_weightsdtype inUnpackedPrecomputedmode. The finalize kernelalready has a compiled fp32 path; this just selects it. No new kernel, no copy.
FP4BlockScaleLauncher::check_routing()— asserttopk_idsis int32 andtopk_weights ∈ {bf16, fp32}, so a wrong dtype fails loudly instead of beingmisread.
flashinfer/fused_moe/core.py— matching Python-side dtype guard + docstringnow documents bf16-or-fp32 (consumed natively, no cast).
unpacked_fp32variant intest_trtllm_gen_routed_fused_moe(fails with garbage before the fix, matches reference after).
🔍 Related Issues
Enables merging vllm-project/vllm#46872, saving 1.5 to 2 us per layer, which matters in decode.
🚀 Pull Request Checklist
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Summary by CodeRabbit
Bug Fixes
topkIDs useint32andtopkweights usebfloat16orfloat32.float32weights so they’re consumed in their native precision without unintended reinterpretation.Documentation
topk_weightsare consumed at native dtype (commonly copy-free forfloat32routers).Tests
unpacked_fp32routing format case.