feat(moe): support shared expert fusion for trtllm-gen fp4 moe - #4239
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
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:
📝 WalkthroughWalkthroughFP4 and FP8 TensorRT-LLM MoE paths now propagate fused shared-expert counts, expand routing and workspace sizing, enforce DeepSeek-V3 and replay constraints, and add trace, autotuning, integration, and validation coverage. ChangesFused shared-expert MoE support
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant MoERunner
participant FP4Launcher
participant DSReference
Caller->>MoERunner: provide num_fused_shared_experts
MoERunner->>FP4Launcher: dispatch fused routing parameters
DSReference->>FP4Launcher: append shared expert ids and unit weights
FP4Launcher-->>Caller: execute fused FP4 MoE
🚥 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 |
![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
🧹 Nitpick comments (2)
csrc/trtllm_fused_moe_kernel_launcher.cu (1)
1924-1944: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider mirroring the EP defense-in-depth pattern for
UnpackedPrecomputed+ fused shared experts.The EP restriction (
num_fused_shared_experts > 0requires no sharding) is enforced at both the Python wrapper and here in C++ (init_common, Lines 706-725). TheUnpackedPrecomputed+num_fused_shared_experts > 0restriction, however, is only asserted in the Python op (core.pyLines 2761-2763) —check_routing()here has no equivalent guard, relying solely on the generic numel mismatch to eventually fail if this combination is ever reached directly.🛡️ Optional defensive check
if (routing_input_mode_ == RoutingInputMode::UnpackedPrecomputed) { TVM_FFI_ICHECK_EQ(topk_ids.dtype(), dl_int32) << "topk_ids must be int32 for unpacked precomputed routing."; TVM_FFI_ICHECK(topk_weights.dtype() == dl_bfloat16 || topk_weights.dtype() == dl_float32) << "topk_weights must be bfloat16 or float32 for unpacked precomputed routing."; + TVM_FFI_ICHECK(args->num_fused_shared_experts == 0) + << "num_fused_shared_experts > 0 is not supported with UnpackedPrecomputed routing."; }🤖 Prompt for 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. In `@csrc/trtllm_fused_moe_kernel_launcher.cu` around lines 1924 - 1944, Update check_routing() to explicitly reject the combination of RoutingInputMode::UnpackedPrecomputed and args->num_fused_shared_experts > 0, mirroring the existing Python validation. Add the guard alongside the routing-mode and fused-shared-expert checks, with a clear validation message; preserve the existing dtype and numel validations for supported combinations.flashinfer/fused_moe/core.py (1)
4120-4135: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated EP/DeepSeekV3-only validation between FP8 and FP4 wrappers.
trtllm_fp8_block_scale_moe(Lines 4120-4135) andtrtllm_fp4_block_scale_moe(Lines 4574-4594) implement identical EP-rejection and DeepSeekV3-only enforcement logic, differing only in local variable naming (nfsevsnsfe) and message wording. This duplication has already begun to drift; extracting a shared helper (e.g._validate_fused_shared_experts_config(num_fused_shared_experts, local_expert_offset, local_num_experts, num_experts, routing_method_type)) would prevent the two paths from silently diverging further.♻️ Proposed shared helper
+def _validate_fused_shared_experts_config( + num_fused_shared_experts: int, + local_expert_offset: int, + local_num_experts: int, + num_experts: int, + routing_method_type: int, +) -> None: + if num_fused_shared_experts > 0 and ( + local_expert_offset != 0 or local_num_experts != num_experts + ): + raise ValueError( + "Fused shared experts (num_fused_shared_experts > 0) do not yet support " + "expert parallelism: require local_expert_offset == 0 and " + "local_num_experts == num_experts. Got " + f"num_fused_shared_experts={num_fused_shared_experts}, " + f"local_expert_offset={local_expert_offset}, " + f"local_num_experts={local_num_experts}, num_experts={num_experts}." + ) + if num_fused_shared_experts > 0 and routing_method_type != RoutingMethodType.DeepSeekV3: + raise ValueError( + "Fused shared experts (num_fused_shared_experts > 0) are only supported " + f"with DeepSeekV3 routing; got routing_method_type={routing_method_type}." + )Then both
trtllm_fp8_block_scale_moeandtrtllm_fp4_block_scale_moecall_validate_fused_shared_experts_config(nfse, local_expert_offset, local_num_experts, num_experts, routing_method_type)followed by_validate_routing_replay_out(...).Also applies to: 4574-4594
🤖 Prompt for 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. In `@flashinfer/fused_moe/core.py` around lines 4120 - 4135, Extract the duplicated fused shared-expert validation from trtllm_fp8_block_scale_moe and trtllm_fp4_block_scale_moe into a shared _validate_fused_shared_experts_config helper accepting num_fused_shared_experts, local_expert_offset, local_num_experts, num_experts, and routing_method_type. Preserve both existing checks and their error behavior in the helper, then have both wrappers call it before _validate_routing_replay_out.
🤖 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
`@tests/trace/fi_trace_out/moe_fp4_block_scale_ds_routing_topk8_e32_h7168_i2048_ng8_kg4.json`:
- Line 239: Regenerate the trace JSON artifact from the current moe.py template
so the embedded _trtllm_fp4_block_scale_moe_ds_routing_reference matches it
exactly, including the **activation_kwargs parameter, nsfe naming, and current
docstring. Preserve the activation-related inputs and do not manually edit the
embedded reference.
---
Nitpick comments:
In `@csrc/trtllm_fused_moe_kernel_launcher.cu`:
- Around line 1924-1944: Update check_routing() to explicitly reject the
combination of RoutingInputMode::UnpackedPrecomputed and
args->num_fused_shared_experts > 0, mirroring the existing Python validation.
Add the guard alongside the routing-mode and fused-shared-expert checks, with a
clear validation message; preserve the existing dtype and numel validations for
supported combinations.
In `@flashinfer/fused_moe/core.py`:
- Around line 4120-4135: Extract the duplicated fused shared-expert validation
from trtllm_fp8_block_scale_moe and trtllm_fp4_block_scale_moe into a shared
_validate_fused_shared_experts_config helper accepting num_fused_shared_experts,
local_expert_offset, local_num_experts, num_experts, and routing_method_type.
Preserve both existing checks and their error behavior in the helper, then have
both wrappers call it before _validate_routing_replay_out.
🪄 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 Plus
Run ID: 2c34c898-0d05-49fc-a201-824bd92775a1
📥 Commits
Reviewing files that changed from the base of the PR and between 92274ba and 7c8056205a7485054c30a728f4e5b1fdc5b3dc87.
📒 Files selected for processing (8)
csrc/trtllm_fused_moe_kernel_launcher.cuflashinfer/fused_moe/core.pyflashinfer/fused_moe/runners.pyflashinfer/trace/templates/moe.pytests/autotuner/test_trtllm_fused_moe_autotuner_integration.pytests/moe/test_trtllm_gen_fused_moe.pytests/moe/trtllm_gen_fused_moe_utils.pytests/trace/fi_trace_out/moe_fp4_block_scale_ds_routing_topk8_e32_h7168_i2048_ng8_kg4.json
7c80562 to
2874b4c
Compare
|
/bot run tests/moe |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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`:
- Line 2808: Before assigning args->num_fused_shared_experts in init_common(),
validate that every expert-major weight, scale, and optional per-expert tensor
has an expert dimension of at least local_num_experts + nFusedShared, matching
the fused-ID mapping; reject undersized routed-only inputs with a host-side
error before launch.
- Around line 2702-2725: Validate that the optional num_fused_shared_experts
value is non-negative before using it in the dimension and scale-count
calculations in the exported launcher. Reject negative values with a clear
validation error, while preserving value_or(0) behavior when the option is
absent.
- Around line 1935-1942: Scope the expanded topk_ids and topk_weights size
checks in the num_fused_shared_experts block to routing_input_mode_ ==
RoutingInputMode::UnpackedPrecomputed. Preserve the existing validation for
unpacked precomputed inputs while allowing FromLogits and PackedPrecomputed to
use their respective buffer contracts.
In `@tests/moe/test_trtllm_gen_fused_moe.py`:
- Line 1964: Update the pytest.raises call in the relevant test to use a raw
regex string for the match pattern, preserving the existing ValueError assertion
and regex behavior while resolving Ruff RUF043.
🪄 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 Plus
Run ID: 29ed239c-e566-42e8-9182-8bca19246682
📥 Commits
Reviewing files that changed from the base of the PR and between 7c8056205a7485054c30a728f4e5b1fdc5b3dc87 and 2874b4cf77755e3394a87e95ff327d3f37d0491b.
📒 Files selected for processing (8)
csrc/trtllm_fused_moe_kernel_launcher.cuflashinfer/fused_moe/core.pyflashinfer/fused_moe/runners.pyflashinfer/trace/templates/moe.pytests/autotuner/test_trtllm_fused_moe_autotuner_integration.pytests/moe/test_trtllm_gen_fused_moe.pytests/moe/trtllm_gen_fused_moe_utils.pytests/trace/fi_trace_out/moe_fp4_block_scale_ds_routing_topk8_e32_h7168_i2048_ng8_kg4.json
🚧 Files skipped from review as they are similar to previous changes (6)
- tests/moe/trtllm_gen_fused_moe_utils.py
- tests/autotuner/test_trtllm_fused_moe_autotuner_integration.py
- flashinfer/fused_moe/runners.py
- flashinfer/trace/templates/moe.py
- tests/trace/fi_trace_out/moe_fp4_block_scale_ds_routing_topk8_e32_h7168_i2048_ng8_kg4.json
- flashinfer/fused_moe/core.py
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 2046-2066: Extend the fused shared-expert validation block around
check_per_expert_scale to validate non-null gemm1_bias and gemm2_bias tensors
before constructing MoERunnerArgs. Require each bias to have totalLocalExperts
rows and preserve the existing expected trailing-dimension checks, then include
routed-only and fused-expert bias cases in the regression matrix.
🪄 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 Plus
Run ID: 4007a634-9773-4323-8538-13dd6979430c
📥 Commits
Reviewing files that changed from the base of the PR and between 2874b4cf77755e3394a87e95ff327d3f37d0491b and 43df9633c007334be61803f02c665469f65b0f5b.
📒 Files selected for processing (2)
csrc/trtllm_fused_moe_kernel_launcher.cutests/moe/test_trtllm_gen_fused_moe.py
b313efa to
6c2fa53
Compare
|
/bot run tests/moe |
aleozlx
left a comment
aleozlx
left a comment
There was a problem hiding this comment.
one comment about testing organization. rest looks good
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Signed-off-by: Haobin Guo <haobing@nvidia.com>
Signed-off-by: Haobin Guo <haobing@nvidia.com>
6c2fa53 to
b31860e
Compare
|
/bot run tests/moe |
|
[FAILED] Pipeline #60237310 — 17/18 executed test jobs passed Compared with nightly #59923506. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
3 participants
📌 Description
Follow-up work for #2625 to add shared-expert fusion for trtllm-gen FP4 MoE.
🔍 Related Issues
#2551
🚀 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
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
num_fused_shared_expertssupport for FP4 block-scale fused-shared-expert MoE, including DeepSeek-V3 routing, dispatcher plumbing, and updated trace/template inputs.