[TRTLLM-14609][chore] Remove legacy MoE path in TRTLLMGenFusedMoE - #16862
Conversation
…om TRTLLMGenFusedMoE ConfigurableMoE has been the default MoE path for months (ENABLE_CONFIGURABLE_MOE defaults to 1); it owns communication via CommunicationFactory and drives compute through backend.run_moe. The TRTLLMGenFusedMoE self-built alltoall workspace/MoeAlltoAll construction and the legacy forward_impl are dead code on that path. Remove them and the now-unused select_alltoall_method_type. The without_comm constructor param is kept (vestigial) so create_moe_backend is untouched; the backend now unconditionally records the not-enabled alltoall state so run_moe internal-alltoall stays off, matching prior without_comm=True behavior. run_moe and the live ConfigurableMoE contract (quantize_input, supports_moe_output_in_alltoall_workspace, forward_fake, enable_alltoall) are unchanged. Signed-off-by: xxi <xxi@nvidia.com>
d0214ff to
b9c5e8b
Compare
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
WalkthroughThe TRTLLMGen fused MoE module now uses disabled all-to-all state, removes its native forward implementation, and retains fake forwarding. Layer-wise benchmarks simplify MoE patching to the WideEPMoE selector. ChangesTRTLLMGen fused MoE changes
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #61757 [ run ] triggered by Bot. Commit: |
|
PR_Github #61757 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61817 [ run ] triggered by Bot. Commit: |
|
PR_Github #61817 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61887 [ run ] triggered by Bot. Commit: |
|
PR_Github #61887 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61895 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #61905 [ run ] triggered by Bot. Commit: |
|
PR_Github #61895 [ run ] completed with state |
|
PR_Github #61905 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61930 [ run ] triggered by Bot. Commit: |
|
PR_Github #61930 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61984 [ run ] triggered by Bot. Commit: |
|
PR_Github #62124 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62191 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #62245 [ run ] triggered by Bot. Commit: |
|
PR_Github #62191 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62260 [ run ] triggered by Bot. Commit: |
|
PR_Github #62245 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62269 [ run ] triggered by Bot. Commit: |
|
PR_Github #62260 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62281 [ run ] triggered by Bot. Commit: |
|
PR_Github #62269 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62296 [ run ] triggered by Bot. Commit: |
|
PR_Github #62281 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62307 [ run ] triggered by Bot. Commit: |
|
PR_Github #62296 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #62307 [ run ] completed with state
|
|
PR_Github #62316 [ run ] triggered by Bot. Commit: |
|
PR_Github #62316 [ run ] completed with state |
5 participants
Description
For the sake of MoE code health, this PR removes the legacy (pre-
ConfigurableMoE) communication construction and forward path fromTRTLLMGenFusedMoE.ConfigurableMoEhas been the default MoE path for months (ENABLE_CONFIGURABLE_MOEdefaults to"1"), which means the backend-internal communication setup and the backend-internal forward pipeline are no longer reachable in the default flow. Keeping them around costs review time, obscures the real control flow, and makes the upcoming MoE refactor harder.What is removed (
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py)if not without_comm:communication-construction block in__init__. It is replaced by an unconditional minimal initialization, mirroring theCutlassFusedMoEchange.select_alltoall_method_type()forward_impl()Net effect: 5 insertions, 375 deletions in a single file.
What is deliberately preserved
ConfigurableMoE/MoEScheduleris untouched:run_moe()andsupports_moe_output_in_alltoall_workspace().without_commconstructor parameter is kept as a vestigial argument for API compatibility.Test Coverage
Verified on OCI GB200 (1 node, 4 GPUs) with the container pinned by
jenkins/current_image_tags.properties, runningtests/unittest/_torch/modules/moe/test_moe_module.pyfiltered to this backend:test_configurable_moe_single_gputest_configurable_moe_multi_gputest_configurable_moe_multi_gpu_eplb138 passed, 0 failed.
PR Checklist
[JIRA/NVBUG][type] descriptionDev Engineer Review
TRTLLMGenFusedMoEto remove legacy all-to-all communication construction,select_alltoall_method_type(), andforward_impl(), leaving a minimal disabled all-to-all state in__init__(NotEnabled, workspaces set toNone, low-precision combine disabled, andmoe_a2a = None).without_commparameter and retained therun_moe()/supports_moe_output_in_alltoall_workspace()contracts.forward_fake()remains in place with its prior bfloat16/fp4 fake-output behavior fordo_finalizemodes.tensorrt_llm/tools/layer_wise_benchmarks/runner.pyto simplify MoE patching duringRunner.scaled_from_ctx: removedunittest.mockusage and theCutlassFusedMoEdependency, and now patches onlyWideEPMoE.select_alltoall_method_type(with proper restoration ofload_pretrained_config).QA Engineer Review
No test changes.