[TRTLLM-14955][refactor] Declare MoE backend behaviour instead of comparing classes - #17411
Open
xxi-nv wants to merge 1 commit into
Open
[TRTLLM-14955][refactor] Declare MoE backend behaviour instead of comparing classes#17411xxi-nv wants to merge 1 commit into
xxi-nv wants to merge 1 commit into
Conversation
xxi-nv
requested review from
Barry-Delaney,
allisonlim-nv,
dhansen-nvidia,
hyukn,
leslie-fang25 and
zhaoyangwang-nvidia
August 7, 2026 08:04
Contributor
|
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 (17)
🚧 Files skipped from review as they are similar to previous changes (17)
WalkthroughMoE execution now uses ChangesMoE execution contract
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant MoEScheduler
participant MoECommPlan
participant MoERunContext
participant MoEBackend
MoEScheduler->>MoECommPlan: build routing and communication state
MoEScheduler->>MoERunContext: package inputs and plan
MoERunContext->>MoEBackend: provide execution context
MoEBackend->>MoEBackend: execute with context and optional workspace
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py (1)
329-336: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the exception type in
query_factory.The coding guidelines require catching specific exceptions.
except Exceptionhere also swallows unrelated failures, for example aTypeErrorfrom a signature change inresolve_moe_cls, and records them as a normal refusal in the golden matrix. Catch the types the factory actually raises, so an unexpected failure surfaces as a test error.♻️ Proposed narrowing
- except Exception as exc: # noqa: BLE001 - the exception type is the result + except (ValueError, NotImplementedError, ImportError, RuntimeError) as exc: + # The exception type is part of the characterized result. return None, f"{type(exc).__name__}: {exc}"As per coding guidelines: "Catch specific exceptions instead of using broad or bare
except:handlers."🤖 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 `@tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py` around lines 329 - 336, Update query_factory’s exception handling around resolve_moe_cls and get_moe_cls to catch only the specific exception types those factories use to indicate an unsupported or unavailable MoE implementation. Preserve the existing refusal tuple for those expected failures, while allowing unexpected errors such as TypeError from signature changes to propagate as test errors.Source: Coding guidelines
tests/unittest/_torch/modules/moe/test_moe_impl_contracts.py (1)
243-263: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe global-registry emptiness assertion couples this module to import order.
test_global_registry_has_no_implementations_yetassertslen(MOE_IMPL_REGISTRY) == 0.test_register_moe_impl_decorator_uses_the_global_registryrestores the state in afinally, so the module is self-consistent. The assertion still breaks as soon as any imported module registers an implementation at import time, which the descriptor comment says will happen as backends migrate. Consider asserting the absence of the specific_id()only, and dropping the length check.🤖 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 `@tests/unittest/_torch/modules/moe/test_moe_impl_contracts.py` around lines 243 - 263, Update test_global_registry_has_no_implementations_yet to remove the global len(MOE_IMPL_REGISTRY) == 0 assertion and retain only the check that MOE_IMPL_REGISTRY.lookup(_id()) is None, so the test remains valid when other implementations register during import.tensorrt_llm/_torch/modules/fused_moe/impl_contract.py (1)
287-306: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRaise an explicit exception instead of
assertinrequire_comm_plan.Python
-Ostripsassert. Under-Othis helper returnsNone, and the caller then fails with an opaqueAttributeError: 'NoneType' object has no attribute 'moe_output'instead of the actionable message written here.Two backends in this same change already use the explicit-raise pattern for the same reason. See
mega_moe_cute_dsl.py("Constructor-time invariant checks raise ValueError so that Python-O(which stripsassert) does not silently let an invalid topology through") andmega_moe_deepgemm.pyload_weights.♻️ Proposed change to keep the diagnostic under `-O`
- assert ctx.comm_plan is not None, ( - f"{type(impl).__name__}.run_moe needs ctx.comm_plan, and the scheduler " - "that drives it always supplies one. A missing plan means run_moe was " - "called without going through ExternalCommMoEScheduler." - ) + if ctx.comm_plan is None: + raise ValueError( + f"{type(impl).__name__}.run_moe needs ctx.comm_plan, and the scheduler " + "that drives it always supplies one. A missing plan means run_moe was " + "called without going through ExternalCommMoEScheduler." + ) return ctx.comm_plan🤖 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 `@tensorrt_llm/_torch/modules/fused_moe/impl_contract.py` around lines 287 - 306, Replace the assert in require_comm_plan with an explicit exception when ctx.comm_plan is None, preserving the existing actionable diagnostic message and returning ctx.comm_plan only after validation. Ensure the check remains active under Python -O.
🤖 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 `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py`:
- Around line 157-168: Update MarlinFusedMoE.run_moe to inspect
require_comm_plan(self, ctx).enable_alltoall before remapping expert IDs.
Subtract slot_start only on the non-all-to-all path; preserve the already-local
expert IDs when all-to-all is enabled, matching CutlassFusedMoE behavior.
In `@tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py`:
- Around line 15-41: Add all three CPU-only modules to the appropriate
CPU-capable test-db list:
tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py
(lines 15-41), tests/unittest/_torch/modules/moe/test_moe_comm_plan.py (lines
15-27), and tests/unittest/_torch/modules/moe/test_moe_impl_contracts.py (lines
15-19). Also commit moe_backend_selection_golden.json alongside the
backend-selection module so its golden tests run successfully.
- Around line 129-130: Commit the missing moe_backend_selection_golden.json
fixture and update the assertion near the golden-file check to include the
existing regeneration hint used around line 459, referencing UPDATE_GOLDEN_ENV
so failures explain how to recreate the fixture.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/impl_contract.py`:
- Around line 287-306: Replace the assert in require_comm_plan with an explicit
exception when ctx.comm_plan is None, preserving the existing actionable
diagnostic message and returning ctx.comm_plan only after validation. Ensure the
check remains active under Python -O.
In `@tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py`:
- Around line 329-336: Update query_factory’s exception handling around
resolve_moe_cls and get_moe_cls to catch only the specific exception types those
factories use to indicate an unsupported or unavailable MoE implementation.
Preserve the existing refusal tuple for those expected failures, while allowing
unexpected errors such as TypeError from signature changes to propagate as test
errors.
In `@tests/unittest/_torch/modules/moe/test_moe_impl_contracts.py`:
- Around line 243-263: Update test_global_registry_has_no_implementations_yet to
remove the global len(MOE_IMPL_REGISTRY) == 0 assertion and retain only the
check that MOE_IMPL_REGISTRY.lookup(_id()) is None, so the test remains valid
when other implementations register during import.
🪄 Autofix
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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d5c30dc3-0f62-4d4b-8505-5ec10c36995a
📒 Files selected for processing (20)
tensorrt_llm/_torch/modules/fused_moe/configurable_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/impl_contract.pytensorrt_llm/_torch/modules/fused_moe/interface.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.pytensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/tools/layer_wise_benchmarks/runner.pytests/microbenchmarks/bench_moe/routing/native_logits.pytests/unittest/_torch/lora/test_moe_lora_model_path.pytests/unittest/_torch/modules/moe/test_moe_backend.pytests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.pytests/unittest/_torch/modules/moe/test_moe_comm_plan.pytests/unittest/_torch/modules/moe/test_moe_impl_contracts.py
…paring classes ExternalCommMoEScheduler and ConfigurableMoE decided what to prepare for a backend by naming its class: 13 sites spread over __class__ ==, isinstance, and a five-way chain in _get_backend_kwargs. Adding a backend therefore meant editing the scheduler, and the checks disagreed with each other about subclasses. Each site now reads a declaration off the backend class: - MoEStaticCapability for the multi-chunk LoRA guard and the DWDP gate. - MoEInputRequirement for routing-scale precision, run_moe workspace allocation, DeepEP expert-id sanitization, and the NVLink one-sided combine workspace dtype. - MoERunContext plus MoECommPlan for run_moe, which took 14 loosely-typed arguments assembled by the class chain. Backends deriving from another backend restate every field rather than inherit it, because the old predicates were inconsistent: the LoRA check was exact-class while the DWDP one was isinstance. Verified all 11 impl classes reproduce the old predicates for every declared field. That restatement rule is transitional and only has to hold while backends still derive from backends; TRTLLM-14960..14969 cut those inheritance edges. Two collapses are intentional. The float32 assertion and the bfloat16 conversion become one cast, which is a no-op for the backends that used to assert. The sanitize request drops its isinstance(moe.comm, DeepEP) half because only DeepEP.dispatch reads the kwarg and every other strategy absorbs it through **kwargs. Removing the class chain also fixes latent defects, because it had no else branch: any impl outside it silently received empty kwargs and fell back to signature defaults. On CuteDslB12xFusedMoE the CUTLASS prefill path was told is_sf_swizzled=True while quantize_input had produced unswizzled scale factors under post-quant dispatch, and output_dtype was forced to bfloat16 for fp16 models. payload_in_workspace is now assigned on every path, so it can no longer be inherited from a previous forward. The five external-comm impls take the plan through require_comm_plan() rather than defaulting each field when it is absent. comm_plan stays optional on the context because a fused-comm impl owns the exchange and nothing outside its kernel decides anything, but an external-comm impl is only reachable through ExternalCommMoEScheduler, which builds a plan on every path. Defaulting there reintroduced the defect above in a new place: a wrong moe_output or enable_alltoall fails loudly, while a wrong input_sf_swizzled just makes the kernel read scale factors at the wrong stride. The three impls that had such a default disagreed on its value, which is what a guess rather than a decision looks like. Also drops the invalid_token_expert_id = -1 write, which every Communication.__init__ already performs. MoECommPlan is still produced by a scheduler helper rather than by the comm strategy. Relocating it needs combine() to read the flag off the plan instead of off the strategy, which changes the dispatch and combine signatures and stays with TRTLLM-14972. The router_logits filter that sat in the TRTLLMGen arm of the class chain moves into TRTLLMGenFusedMoE._routes_outside_the_kernel, beside the kernel that cares about it. Its three triggers are carried over unchanged. FORCE_SEPARATED_ROUTING moves to interface.py because both sides read it -- the scheduler to decide whether to precompute top-k at all, the backend to decide whether its kernel may route again -- and interface.py is the only module both already import, so neither gains a dependency. Reaching run_moe with router_logits set but no precomputed top-k now raises rather than dereferencing None. The scheduler cannot produce that pairing, since whatever forces the filter also forces the precompute; test_moe_backend calls run_moe directly and can, so it skips the combination. Signed-off-by: xxi <xxi@nvidia.com>
xxi-nv
force-pushed
the
feat/trtllm-14955-moe-eliminate-class-dispatch
branch
from
August 7, 2026 08:38
843b87b to
eb5b8f4
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #64570 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64570 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
ExternalCommMoESchedulerandConfigurableMoEdecided what to prepare for a backend by naming its class: 13 sites spread over__class__ ==,isinstance, and a five-way chain in_get_backend_kwargs. Adding a backend therefore meant editing the scheduler, and the checks disagreed with each other about subclasses — the multi-chunk LoRA guard was exact-class while the DWDP gate wasisinstance.Each site now reads a declaration off the backend class instead:
MoEStaticCapabilityMoEInputRequirementrun_moeworkspace allocation, DeepEP expert-id sanitization, NVLink one-sided combine workspace dtypeMoERunContext+MoECommPlanrun_moe, which took 14 loosely-typed arguments assembled by the class chainBackends deriving from another backend restate every field rather than inherit it, because the old predicates were inconsistent. All 11 impl classes were verified to reproduce the old predicate for every declared field. That restatement rule is transitional and only has to hold while backends still derive from backends; TRTLLM-14960..14969 cut those inheritance edges.
Two collapses are intentional. The float32 assertion and the bfloat16 conversion become one cast, a no-op for the backends that used to assert. The sanitize request drops its
isinstance(moe.comm, DeepEP)half because onlyDeepEP.dispatchreads the kwarg and every other strategy absorbs it through**kwargs.Removing the class chain also fixes latent defects, because it had no
elsebranch: any impl outside it silently received empty kwargs and fell back to signature defaults. OnCuteDslB12xFusedMoEthe CUTLASS prefill path was toldis_sf_swizzled=Truewhilequantize_inputhad produced unswizzled scale factors under post-quant dispatch, andoutput_dtypewas forced to bfloat16 for fp16 models.payload_in_workspaceis now assigned on every path, so it can no longer be inherited from a previous forward.The five external-comm impls take the plan through
require_comm_plan()rather than defaulting each field when it is absent.comm_planstays optional on the context because a fused-comm impl owns the exchange, but an external-comm impl is only reachable throughExternalCommMoEScheduler, which builds a plan on every path. Defaulting there reintroduced the same class of defect in a new place: a wrongmoe_outputorenable_alltoallfails loudly, while a wronginput_sf_swizzledjust makes the kernel read scale factors at the wrong stride. The three impls that had such a default disagreed on its value.The
router_logitsfilter that sat in the TRTLLMGen arm of the chain moves intoTRTLLMGenFusedMoE._routes_outside_the_kernel, beside the kernel that cares about it; its three triggers are carried over unchanged.FORCE_SEPARATED_ROUTINGmoves tointerface.pybecause both sides read it and that is the only module both already import.No functional change is intended beyond the latent-defect fixes called out above.
Test Coverage
This is a behaviour-preserving refactor, so the existing MoE suites are the safeguard. Both were exercised on GB300 (SM103, 4 GPU):
tests/unittest/_torch/modules/moe/test_moe_backend.py— updated to build aMoERunContextand attach aMoECommPlanfor external-comm backends, so every backend'srun_moeis driven through the new contract. Covers Cutlass, CuteDSL, DeepGemm, DenseGEMM and TRTLLMGen across the quantization modes.tests/unittest/_torch/modules/moe/test_moe_module.py— end-to-end single-GPU and multi-GPU (EP/DP, EPLB) paths throughConfigurableMoEand both schedulers.tests/unittest/_torch/lora/test_moe_lora_model_path.py— updated for theMoEStaticCapability.supports_moe_loradeclaration that replaces the exact-class LoRA check.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
MoERunContextandMoECommPlanfor typed execution and communication state.run_moeAPI.QA Engineer Review
MoERunContext,MoECommPlan, workspace handling, and fused communication.tests/integration/test_lists/,test-db/, orqa/was not identified.