[TRTLLM-14956][refactor] make MoE implementation selection reproducible - #17532
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:
WalkthroughThis change centralizes MoE backend resolution. It adds structured contracts for selection inputs, deployment state, environment data, eligibility, implementation identity, and resolution reports. Backends, model integrations, construction paths, tests, and documentation use the new resolver. ChangesMoE resolution architecture
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLayer
participant create_moe
participant moe_resolution
participant MoEBackend
participant MoEResolutionReport
ModelLayer->>create_moe: provide model and layer metadata
create_moe->>moe_resolution: resolve_moe_impl(...)
moe_resolution->>MoEBackend: can_implement(problem, deployment)
MoEBackend-->>moe_resolution: return eligibility verdict
moe_resolution->>MoEResolutionReport: record selection or degradation
MoEResolutionReport-->>create_moe: return implementation class
create_moe-->>ModelLayer: construct MoE layer
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
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: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/modules/fused_moe/impl_identity.py (1)
171-198: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winCheck token disjointness within one identity too.
_check_tokens_disjointcompares each token only against_token_to_field, which does not yet contain the identity being registered. An identity that reuses one token across two of its own fields passes the check. Line 196-197 then records that token under the last field in_ID_FIELDSorder, sofield_ofandparse_queryresolve it to one field only. The error message promises "value sets must stay disjoint across fields", but that invariant is not enforced for self-collisions.No registered identity collides today. The gap becomes reachable when a new impl reuses a token, for example
provider="triton"withtechnique="triton".🔧 Proposed fix
def _check_tokens_disjoint(self, identity: MoEImplId) -> None: + seen: Dict[str, str] = {} for name in _ID_FIELDS: token = getattr(identity, name) - owner = self._token_to_field.get(token) + owner = self._token_to_field.get(token) or seen.get(token) if owner is not None and owner != name: raise ValueError( f"cannot register {identity.canonical()}: token {token!r} is already a " f"value of field {owner!r}, so a user writing {token!r} could mean either " f"field. Rename one of them -- value sets must stay disjoint across fields." ) + seen[token] = name🤖 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_identity.py` around lines 171 - 198, Update _check_tokens_disjoint to track tokens encountered while iterating the current identity as well as tokens in _token_to_field. Raise the same collision error when two fields within one identity reuse a token, before register stores the identity or updates _token_to_field, preserving the disjoint-value invariant for both intra- and inter-identity collisions.
🧹 Nitpick comments (14)
tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py (1)
716-719: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe failure message names the wrong function.
The assertion compares
canonical_activation(default)againstcanonical_activation(None), but the message attributes the folding tobuild_moe_problem. A developer reading this failure would inspect the wrong function.♻️ Proposed message fix
assert canonical_activation(default) == folded, ( f"{owner.__qualname__} defaults activation_type to {default}, but " - f"build_moe_problem folds an absent one to {folded}" + f"canonical_activation folds an absent one to {folded}" )🤖 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 716 - 719, Update the assertion failure message in the test around canonical_activation to identify the function that actually folds an absent activation value, rather than naming build_moe_problem; keep the assertion and compared values unchanged.tests/unittest/_torch/modules/moe/test_moe_impl_contracts.py (2)
17-17: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
from __future__ import annotationsis not needed.TensorRT-LLM requires Python >=3.10, so
str | Noneat Line 652 andtuple[str, ...]at Line 654 evaluate at runtime without this import.Based on learnings: "In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features (e.g., PEP 585 generics like dict[str, int], list[str], etc.) throughout the codebase, and you do not need to add
from __future__ import annotations."🤖 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` at line 17, Remove the unnecessary from __future__ import annotations statement from test_moe_impl_contracts.py; the annotations in the affected test code are supported directly by the project’s Python 3.10+ requirement.Source: Learnings
396-401: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe cleanup reaches into two private registry indexes.
The
finallyblock pops fromMOE_IMPL_REGISTRY._storeandMOE_IMPL_REGISTRY._token_to_fielddirectly. IfMoEImplRegistrygains a third index, this cleanup leaks state into the module-level singleton andtest_global_registry_has_no_implementations_yetat Line 375 fails for an unrelated reason.Consider adding a public
unregisterorclearmethod onMoEImplRegistryand calling it here, so the cleanup stays correct as the registry grows.🤖 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 396 - 401, Replace the direct `_store` and `_token_to_field` mutations in the test’s `finally` cleanup with a public `MoEImplRegistry` cleanup API, preferably `unregister(identity)` or an equivalent clear operation. Implement that method to remove the implementation and all associated indexes, then call it on `MOE_IMPL_REGISTRY` so future registry indexes are cleaned consistently.tests/unittest/_torch/modules/moe/test_moe_backend.py (2)
1222-1225: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe function now ignores its own
num_expertsandtop_kparameters.Lines 1222-1223 reassign
num_expertsandtop_kto the module constants. The parameters declared at Lines 1215-1216 already default to the same constants, so any value a caller passes is silently discarded.No current caller passes them, so behavior is unchanged today. Remove the parameters to prevent a future parametrization from being ignored without warning.
♻️ Proposed cleanup
def test_trtllm_bf16_unquantized_moe( routing_kind, activation_type, seq_len, trtllm_use_router_logits, - num_experts=_BF16_UNQUANT_NUM_EXPERTS, - top_k=_BF16_UNQUANT_TOP_K, ): """TRTLLM-Gen BF16 (unquantized) MoE accuracy vs the reference impl.""" backend_type = MoeBackendType.TRTLLM dtype = torch.bfloat16 num_experts = _BF16_UNQUANT_NUM_EXPERTS top_k = _BF16_UNQUANT_TOP_K🤖 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.py` around lines 1222 - 1225, Remove the num_experts and top_k parameters from the affected test function’s signature and delete their local reassignments, while retaining the module constants for the fixed test configuration. Update any affected calls to match the simplified signature.
405-414: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAssert
report.degradedbefore readingdegraded_from.
MoEResolutionReport.degraded_fromreturnsNonewhenselected_byis not"heuristic". If the resolver ever pins Marlin here, Line 414 raisesAttributeError: 'NoneType' object has no attribute 'reason'instead of reporting which implementation was selected.Add the same
assert report.degradedthattest_marlin_degrades_to_cutlass_on_non_nvfp4uses at Line 401.♻️ Proposed assertion
assert impl_class_for(report) is CutlassFusedMoE + assert report.degraded assert report.degraded_from.reason is MoERejectReason.QUANT_UNSUPPORTED🤖 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.py` around lines 405 - 414, Add assert report.degraded immediately before accessing report.degraded_from in test_marlin_override_quant_config_degrades_per_layer, matching test_marlin_degrades_to_cutlass_on_non_nvfp4, while preserving the existing implementation and rejection-reason assertions.tests/unittest/_torch/modules/moe/moe_test_utils.py (1)
1165-1182: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the
model_config is Noneguards or make the parameter Optional.The new code treats
model_configas possiblyNonein six places. The signature at Line 1150 declares it as"MoeModelConfig", and Line 1240 readsmodel_config.hidden_sizewithout a guard. Ifmodel_configwere everNone, Line 1240 would raiseAttributeErrorfor every quantized configuration, so the new guards cannot make the function safe.Pick one contract. Either declare
model_config: Optional["MoeModelConfig"]and guard Line 1240 as well, or remove the guards and rely on the declared non-optional type.🤖 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/moe_test_utils.py` around lines 1165 - 1182, The model_config contract is inconsistent in the affected test helper. Either change the parameter annotation to Optional["MoeModelConfig"] and guard every access, including the model_config.hidden_size use near the function’s later validation, or keep the non-optional annotation and remove all model_config is None fallbacks in MoEProblem and MoEDeployment construction.tests/microbenchmarks/bench_moe/search.py (1)
90-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the exception type and keep the cause.
except Exceptioncatches every failure, includingKeyboardInterruptsubclasses ofExceptionsuch as programming errors in the sweep itself. The coding guidelines require the narrowest exception. Ruff also reports BLE001 here.The realistic failure modes of
can_implementareAttributeError,KeyError,TypeError,ValueError, andRuntimeErrorfrom a backend that has not fully migrated to the contract. Catch those and let anything else propagate.As per coding guidelines: "Catch the narrowest exception possible" and "Catch specific exceptions instead of using broad or bare
except:handlers."♻️ Proposed narrowing
try: verdict = backend_cls.can_implement(problem, deployment) - except Exception as exc: + except (AttributeError, KeyError, TypeError, ValueError, RuntimeError) as exc: return False, (f"{backend_cls.__name__}.can_implement raised {type(exc).__name__}: {exc}")🤖 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/microbenchmarks/bench_moe/search.py` around lines 90 - 93, Update the exception handler around backend_cls.can_implement in the benchmark sweep to catch only AttributeError, KeyError, TypeError, ValueError, and RuntimeError. Preserve the existing false verdict and diagnostic message for those failures, and chain the original exception as the cause; allow all other exceptions to propagate.Sources: Coding guidelines, Linters/SAST tools
tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py (2)
213-214: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSelect the quantization config with an explicit
Nonecheck.
override_quant_config or model_config.quant_configdiscards the override whenever the override object is falsy.QuantConfigis truthy today, so behavior is correct. Theorform makes the selection depend onQuantConfig.__bool__, and a future__bool__or__len__on that class would silently route to the model-level config and change the resolved backend without any error.♻️ Proposed refactor
- quant_config = override_quant_config or model_config.quant_config + quant_config = ( + override_quant_config if override_quant_config is not None else model_config.quant_config + )🤖 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/moe_resolution.py` around lines 213 - 214, Update the quant_config selection in the surrounding resolution logic to use an explicit None check: choose override_quant_config whenever it is not None, otherwise use model_config.quant_config. Keep the subsequent quant_algo derivation unchanged.
464-478: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTighten the return type to
Type[MoE].
impl_class_forandresolve_moe_clsreturn bareType. Every element ofIMPL_PRIORITYis anMoEsubclass, andcreate_moe_backenddeclaresmoe_cls: Type[MoE].Type[MoE]states the contract and lets a type checker validate thecreate_moeandConfigurableMoEcall sites. Both functions are re-exported fromtensorrt_llm/_torch/modules/fused_moe/__init__.py, so the annotation is part of the public surface.♻️ Proposed refactor
-def impl_class_for(report: MoEResolutionReport) -> Type: +def impl_class_for(report: MoEResolutionReport) -> Type["MoE"]: """The class a report's winner names, or raise with the whole trail."""-def resolve_moe_cls(model_config: ModelConfig, **kwargs) -> Type: +def resolve_moe_cls(model_config: ModelConfig, **kwargs) -> Type["MoE"]: """Resolve and return only the implementation class."""Add
MoEto theTYPE_CHECKINGblock:if TYPE_CHECKING: + from .interface import MoE from .routing import BaseMoeRoutingMethod, RoutingMethodTypeAs per coding guidelines: "Annotate every function ... use precise
Callablearguments, use@overloadorTypeVarwhen return types depend on inputs."🤖 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/moe_resolution.py` around lines 464 - 478, Update impl_class_for and resolve_moe_cls to return Type[MoE] instead of bare Type, importing MoE under the existing TYPE_CHECKING guard as needed. Preserve the current resolution behavior and ensure the annotations match the MoE subclass contract used by create_moe_backend and the re-exported public API.Source: Coding guidelines
tensorrt_llm/_torch/modules/fused_moe/create_moe.py (3)
30-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__.Ruff reports RUF022 on this list. Sorting it also matches the ordering used elsewhere in the package.
♻️ Proposed refactor
__all__ = [ + "WIDEEP_DEPRECATION_MESSAGE", "create_moe", "create_moe_backend", "infer_swiglu_gptoss_style", "resolve_moe_cls", "resolve_moe_impl", - "WIDEEP_DEPRECATION_MESSAGE", ]🤖 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/create_moe.py` around lines 30 - 37, Sort the exported names in __all__ alphabetically to resolve Ruff RUF022, keeping the same symbols and export behavior unchanged.Source: Linters/SAST tools
98-118: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOne derive-and-validate block is duplicated across both MoE entry points.
create_moe_backendandcreate_moeeach callderive_moe_layer_shapes, unpack the same four fields, and repeat the same three assertion messages verbatim. Both copies must change together whenever the inference rules or the message text change.
tensorrt_llm/_torch/modules/fused_moe/create_moe.py#L98-L118: replace the block with a call to a sharedrequire_moe_layer_shapeshelper added next toderive_moe_layer_shapesinmoe_resolution.py.tensorrt_llm/_torch/modules/fused_moe/create_moe.py#L390-L410: replace the identical block with the same helper call.🤖 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/create_moe.py` around lines 98 - 118, Add a shared require_moe_layer_shapes helper next to derive_moe_layer_shapes in moe_resolution.py that performs shape derivation, field unpacking, and the existing validations, then replace both duplicated blocks in create_moe_backend at tensorrt_llm/_torch/modules/fused_moe/create_moe.py#L98-L118 and create_moe at tensorrt_llm/_torch/modules/fused_moe/create_moe.py#L390-L410 with calls to that helper, preserving the current arguments and returned shape values.
431-433: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a drift guard for the ConfigurableMoE membership tuple.
This inline tuple must stay in sync with
IMPL_PRIORITYminusTritonFusedMoE,VanillaMoE, andWideEPMoE. A new backend added toIMPL_PRIORITYandBACKEND_FAMILYbut omitted here silently takes thecreate_moe_backendpath. It then loses the communication strategy and the scheduler thatConfigurableMoEbuilds, and the omission produces no error.
moe_resolution.pyalready guards theBACKEND_FAMILY/IMPL_PRIORITYpairing at import time with_UNRANKED. Apply the same treatment here.♻️ Proposed refactor
Define the set once at module scope with a drift check:
# Backends that ConfigurableMoE wraps. The rest own their comm and forward paths. _UNWRAPPED_IMPLS = frozenset({TritonFusedMoE, VanillaMoE, WideEPMoE}) _CONFIGURABLE_IMPLS = frozenset(IMPL_PRIORITY) - _UNWRAPPED_IMPLS _UNCLASSIFIED = set(IMPL_PRIORITY) - _CONFIGURABLE_IMPLS - _UNWRAPPED_IMPLS if _UNCLASSIFIED: raise RuntimeError( f"MoE impls in IMPL_PRIORITY with no ConfigurableMoE decision: " f"{sorted(cls.__name__ for cls in _UNCLASSIFIED)}" )Then:
- if moe_cls in (DeepGemmFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, - CuteDslB12xFusedMoE, CutlassFusedMoE, DenseGEMMFusedMoE, - MegaMoEDeepGemm, MegaMoECuteDsl, MarlinFusedMoE): + if moe_cls in _CONFIGURABLE_IMPLS:🤖 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/create_moe.py` around lines 431 - 433, Replace the inline membership tuple in the ConfigurableMoE selection logic with module-level _UNWRAPPED_IMPLS and _CONFIGURABLE_IMPLS derived from IMPL_PRIORITY. Add an import-time _UNCLASSIFIED drift check, matching moe_resolution.py, that raises a RuntimeError listing any IMPL_PRIORITY implementation absent from both sets; use _CONFIGURABLE_IMPLS for the membership test.tensorrt_llm/_torch/modules/fused_moe/impl_identity.py (1)
99-99: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the missing return annotations.
The coding guidelines require an annotation on every function.
MoEImplId.__post_init__(Line 66),MoEImplQuery.__post_init__(Line 99),MoEImplQuery.__str__(Line 140), andMoEImplRegistry.__len__(Line 269) have none.🔧 Proposed fix
- def __post_init__(self): + def __post_init__(self) -> None: for name in _ID_FIELDS: value = getattr(self, name) if value is not None: object.__setattr__(self, name, _normalize(name, value))- def __str__(self): + def __str__(self) -> str: return self.describe()As per coding guidelines: "Annotate every function, use
Nonefor procedures, avoid unnecessaryAnyandtype: ignore."Also applies to: 140-141
🤖 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_identity.py` at line 99, Add explicit return annotations to MoEImplId.__post_init__, MoEImplQuery.__post_init__, MoEImplQuery.__str__, and MoEImplRegistry.__len__; use None for the __post_init__ procedures and the appropriate string and integer return types for __str__ and __len__.Source: Coding guidelines
tensorrt_llm/_torch/modules/fused_moe/impl_environment.py (1)
172-181: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the return type of
override_moe_environment.The coding guidelines require an annotation on every function. This contextmanager generator has none.
🔧 Proposed fix
+from collections.abc import Iterator + `@contextmanager` -def override_moe_environment(environment: MoEEnvironment): +def override_moe_environment(environment: MoEEnvironment) -> Iterator[MoEEnvironment]: """Temporarily override the collected MoE selection environment."""As per coding guidelines: "Annotate every function, use
Nonefor procedures, avoid unnecessaryAnyandtype: ignore, prefer built-in generic types and|."🤖 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_environment.py` around lines 172 - 181, Annotate the return type of override_moe_environment with the appropriate context-manager type for its yielded MoEEnvironment, while preserving the existing temporary override and restoration behavior.Source: Coding guidelines
🤖 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/models/modeling_deepseekv4.py`:
- Around line 1507-1517: Ensure the preflight resolver uses the same
swiglu_gptoss_style value as create_moe, resolving the default consistently as
False so TritonFusedMoE selection cannot diverge from creation; update the
resolver calls in tensorrt_llm/_torch/models/modeling_deepseekv4.py:1507-1517,
tensorrt_llm/_torch/models/modeling_laguna.py:135-138, and
tensorrt_llm/_torch/models/modeling_qwen3_moe.py:114-118, or reuse a shared
MoEResolutionReport where appropriate.
In `@tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py`:
- Around line 319-325: Update infer_swiglu_gptoss_style to normalize its
activation_type argument through ActivationType before checking for
ActivationType.SwigluBias, so both enum and integer inputs detect SwigluBias
correctly. This fixes the direct-construction call in MoE.__init__ while
preserving create_moe behavior and existing non-Swiglu results.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py`:
- Around line 98-100: Update the FC2 tile-alignment gate and its associated
comment: remove the reference to the deleted constructor check, and validate the
per-partition intermediate size by dividing the global intermediate size by
d.tp_size before applying the _FC2_MMA_TILE_K modulo check. Preserve the
existing rejection behavior for non-aligned shards in the gate around
run_moe_nvfp4.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py`:
- Around line 158-169: Update the unquantized branch of can_implement around
quant_algo and p.swiglu_gptoss_style to reject activation types outside
_BF16_SUPPORTED_ACTIVATIONS, matching _check_configs. Preserve the existing
custom-parameter and FlashInfer dependency checks, and return the appropriate
unsupported-activation rejection before reporting eligibility.
- Around line 19-35: Update the imports in the fused MoE module to bind
get_sm_version and logger from their existing utility or logging modules,
matching the project’s established import locations. Ensure the symbols used by
the affected runtime paths resolve without changing their behavior.
In `@tensorrt_llm/_torch/modules/fused_moe/impl_contract.py`:
- Around line 348-376: Extend the selection-report serialization in to_dict to
include problem.activation, problem.routing, deployment.parallel_size, and
deployment.cluster_size, preserving their original values and the existing
artifact field structure so replay consumers can reproduce can_implement
verdicts.
In `@tensorrt_llm/_torch/modules/fused_moe/impl_environment.py`:
- Around line 123-132: Replace printf-style logging arguments with preformatted
f-string messages at all three sites: in impl_environment.py lines 123-132,
update _run_probe’s warning and debug calls to include dep.value,
type(exc).__name__, exc, and detail directly; in impl_environment.py lines
155-161, format environment.sm, available, env_flags, and
environment.fingerprint() directly in the debug message; in moe_resolution.py
lines 449-459, format requested, location, cause.reason.value, cause.detail,
winner_cls.__name__, and report.describe() in the warning, and pass
report.describe() directly to the debug call.
In `@tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py`:
- Around line 268-274: Replace the `assert not apply_router_weight_on_input`
guard in `MegaMoEDeepGemm` with an unconditional `ValueError` using the existing
explanatory message, matching the `MegaMoECuteDsl` pattern so the unsupported
flag is rejected even under `python -O`.
In `@tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py`:
- Around line 254-269: Update the num_slots initialization in the MoE deployment
construction to use balancer_config.num_slots only when get_moe_load_balancer()
returns an active balancer; otherwise fall back to num_experts, preserving 0
when num_experts is unavailable. Keep the eplb_enabled behavior unchanged.
---
Outside diff comments:
In `@tensorrt_llm/_torch/modules/fused_moe/impl_identity.py`:
- Around line 171-198: Update _check_tokens_disjoint to track tokens encountered
while iterating the current identity as well as tokens in _token_to_field. Raise
the same collision error when two fields within one identity reuse a token,
before register stores the identity or updates _token_to_field, preserving the
disjoint-value invariant for both intra- and inter-identity collisions.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py`:
- Around line 30-37: Sort the exported names in __all__ alphabetically to
resolve Ruff RUF022, keeping the same symbols and export behavior unchanged.
- Around line 98-118: Add a shared require_moe_layer_shapes helper next to
derive_moe_layer_shapes in moe_resolution.py that performs shape derivation,
field unpacking, and the existing validations, then replace both duplicated
blocks in create_moe_backend at
tensorrt_llm/_torch/modules/fused_moe/create_moe.py#L98-L118 and create_moe at
tensorrt_llm/_torch/modules/fused_moe/create_moe.py#L390-L410 with calls to that
helper, preserving the current arguments and returned shape values.
- Around line 431-433: Replace the inline membership tuple in the
ConfigurableMoE selection logic with module-level _UNWRAPPED_IMPLS and
_CONFIGURABLE_IMPLS derived from IMPL_PRIORITY. Add an import-time _UNCLASSIFIED
drift check, matching moe_resolution.py, that raises a RuntimeError listing any
IMPL_PRIORITY implementation absent from both sets; use _CONFIGURABLE_IMPLS for
the membership test.
In `@tensorrt_llm/_torch/modules/fused_moe/impl_environment.py`:
- Around line 172-181: Annotate the return type of override_moe_environment with
the appropriate context-manager type for its yielded MoEEnvironment, while
preserving the existing temporary override and restoration behavior.
In `@tensorrt_llm/_torch/modules/fused_moe/impl_identity.py`:
- Line 99: Add explicit return annotations to MoEImplId.__post_init__,
MoEImplQuery.__post_init__, MoEImplQuery.__str__, and MoEImplRegistry.__len__;
use None for the __post_init__ procedures and the appropriate string and integer
return types for __str__ and __len__.
In `@tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py`:
- Around line 213-214: Update the quant_config selection in the surrounding
resolution logic to use an explicit None check: choose override_quant_config
whenever it is not None, otherwise use model_config.quant_config. Keep the
subsequent quant_algo derivation unchanged.
- Around line 464-478: Update impl_class_for and resolve_moe_cls to return
Type[MoE] instead of bare Type, importing MoE under the existing TYPE_CHECKING
guard as needed. Preserve the current resolution behavior and ensure the
annotations match the MoE subclass contract used by create_moe_backend and the
re-exported public API.
In `@tests/microbenchmarks/bench_moe/search.py`:
- Around line 90-93: Update the exception handler around
backend_cls.can_implement in the benchmark sweep to catch only AttributeError,
KeyError, TypeError, ValueError, and RuntimeError. Preserve the existing false
verdict and diagnostic message for those failures, and chain the original
exception as the cause; allow all other exceptions to propagate.
In `@tests/unittest/_torch/modules/moe/moe_test_utils.py`:
- Around line 1165-1182: The model_config contract is inconsistent in the
affected test helper. Either change the parameter annotation to
Optional["MoeModelConfig"] and guard every access, including the
model_config.hidden_size use near the function’s later validation, or keep the
non-optional annotation and remove all model_config is None fallbacks in
MoEProblem and MoEDeployment construction.
In `@tests/unittest/_torch/modules/moe/test_moe_backend_selection_consistency.py`:
- Around line 716-719: Update the assertion failure message in the test around
canonical_activation to identify the function that actually folds an absent
activation value, rather than naming build_moe_problem; keep the assertion and
compared values unchanged.
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 1222-1225: Remove the num_experts and top_k parameters from the
affected test function’s signature and delete their local reassignments, while
retaining the module constants for the fixed test configuration. Update any
affected calls to match the simplified signature.
- Around line 405-414: Add assert report.degraded immediately before accessing
report.degraded_from in test_marlin_override_quant_config_degrades_per_layer,
matching test_marlin_degrades_to_cutlass_on_non_nvfp4, while preserving the
existing implementation and rejection-reason assertions.
In `@tests/unittest/_torch/modules/moe/test_moe_impl_contracts.py`:
- Line 17: Remove the unnecessary from __future__ import annotations statement
from test_moe_impl_contracts.py; the annotations in the affected test code are
supported directly by the project’s Python 3.10+ requirement.
- Around line 396-401: Replace the direct `_store` and `_token_to_field`
mutations in the test’s `finally` cleanup with a public `MoEImplRegistry`
cleanup API, preferably `unregister(identity)` or an equivalent clear operation.
Implement that method to remove the implementation and all associated indexes,
then call it on `MOE_IMPL_REGISTRY` so future registry indexes are cleaned
consistently.
🪄 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: ba90f403-5e5f-4f67-b918-27993394f1a6
📒 Files selected for processing (35)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_laguna.pytensorrt_llm/_torch/models/modeling_qwen3_moe.pytensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/__init__.pytensorrt_llm/_torch/modules/fused_moe/configurable_moe.pytensorrt_llm/_torch/modules/fused_moe/create_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_triton.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.pytensorrt_llm/_torch/modules/fused_moe/impl_contract.pytensorrt_llm/_torch/modules/fused_moe/impl_environment.pytensorrt_llm/_torch/modules/fused_moe/impl_identity.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_resolution.pytensorrt_llm/_torch/peft/lora/validation.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/microbenchmarks/bench_moe/backend.pytests/microbenchmarks/bench_moe/search.pytests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.pytests/unittest/_torch/modules/moe/moe_test_utils.pytests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.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_impl_contracts.pytests/unittest/_torch/modules/moe/test_moe_module.py
ec8c5c7 to
f5f58a4
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (4)
tensorrt_llm/_torch/modules/fused_moe/impl_environment.py (4)
27-54: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare the public module interface.
This module exposes
MoEDep,MoEEnvFlag,DepProbe, and the public environment helpers without__all__. Define__all__. IfDepProbeis internal, rename it to_DepProbeinstead of exporting it.As per coding guidelines: “keep
__all__updated for public interfaces.”🤖 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_environment.py` around lines 27 - 54, Declare the module’s public interface with an __all__ containing MoEDep, MoEEnvFlag, and the public environment helper symbols defined in the module. Treat DepProbe as internal by renaming it to _DepProbe and update all references, rather than exporting the type alias.Source: Coding guidelines
135-136: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument public parameters and context-manager behavior.
Document
force, cache behavior,environment, and the yielded override value in Google-style docstrings. These functions define the reproducibility boundary for MoE selection.As per coding guidelines: “Google-style docstrings for classes and functions.”
Also applies to: 169-171
🤖 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_environment.py` around lines 135 - 136, Update the Google-style docstrings for collect_moe_environment and the related context-manager function to document the force parameter, cache behavior, environment value, and yielded override value. Describe how force affects collection and clarify the context manager’s temporary override and restoration behavior, preserving the existing implementation.Source: Coding guidelines
20-20: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse Python 3.10 annotations consistently.
Replace
Dict,Optional, andTuplewithdict,MoEEnvironment | None, andtuple. Annotateoverride_moe_environmentwith its yielded type, such asIterator[MoEEnvironment].As per coding guidelines: “Annotate every function” and “prefer built-in generic types and
|.”Also applies to: 53-57, 65-65, 80-80, 90-90, 97-97, 106-106, 115-120, 169-170
🤖 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_environment.py` at line 20, Update all annotations in impl_environment.py to use Python 3.10 built-in generics and union syntax: replace Dict with dict, Optional[MoEEnvironment] with MoEEnvironment | None, and Tuple with tuple. Add the yielded return annotation to override_moe_environment, such as Iterator[MoEEnvironment], and ensure every function in the affected ranges is annotated consistently.Source: Coding guidelines
119-120: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse snake_case for mutable module state.
_CACHED_ENVIRONMENTand_OVERRIDE_ENVIRONMENTare reassigned. Rename them to_cached_environmentand_override_environment.As per coding guidelines: “Use snake_case for ... mutable globals.”
🤖 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_environment.py` around lines 119 - 120, Rename the mutable module-level state variables _CACHED_ENVIRONMENT and _OVERRIDE_ENVIRONMENT to _cached_environment and _override_environment, and update every reference to them throughout the module while preserving their existing behavior.Source: Coding guidelines
🤖 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/impl_environment.py`:
- Around line 60-61: Update _run_probe in
tensorrt_llm/_torch/modules/fused_moe/impl_environment.py at lines 60-61, 68-69,
and 126-129 to catch only expected import or load exceptions for the FlashInfer
probes; remove the broad Exception handling so unexpected probe defects
propagate instead of returning False and selecting a fallback backend.
In `@tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md`:
- Line 160: Update the WideEP (deprecated) column in the capability matrix,
including the entries around the referenced additional lines, to mark every
combination as reference-only or unavailable for selection. Ensure the matrix
consistently reflects that moe_resolution.py rejects the WIDEEP backend literal
while preserving the existing supported-backend guidance.
- Around line 245-268: Clarify the resolution-report contract around
resolve_moe_impl: callers must catch and handle ValueError for unknown or
deprecated backend literals, and no MoEResolutionReport is returned for these
configuration errors. Preserve report generation for valid backend requests,
including cases where winner is None.
- Around line 212-237: Update the “Unknown is not false” guidance to explicitly
state that gates must skip optional shape checks when the field is None and
return MoEEligibility.ok(). Note that resolve_moe_impl treats this result as
eligible and may select the backend without proving the shape constraint.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/impl_environment.py`:
- Around line 27-54: Declare the module’s public interface with an __all__
containing MoEDep, MoEEnvFlag, and the public environment helper symbols defined
in the module. Treat DepProbe as internal by renaming it to _DepProbe and update
all references, rather than exporting the type alias.
- Around line 135-136: Update the Google-style docstrings for
collect_moe_environment and the related context-manager function to document the
force parameter, cache behavior, environment value, and yielded override value.
Describe how force affects collection and clarify the context manager’s
temporary override and restoration behavior, preserving the existing
implementation.
- Line 20: Update all annotations in impl_environment.py to use Python 3.10
built-in generics and union syntax: replace Dict with dict,
Optional[MoEEnvironment] with MoEEnvironment | None, and Tuple with tuple. Add
the yielded return annotation to override_moe_environment, such as
Iterator[MoEEnvironment], and ensure every function in the affected ranges is
annotated consistently.
- Around line 119-120: Rename the mutable module-level state variables
_CACHED_ENVIRONMENT and _OVERRIDE_ENVIRONMENT to _cached_environment and
_override_environment, and update every reference to them throughout the module
while preserving their existing behavior.
🪄 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: 739febf3-42fb-41c5-b584-01572b359c61
📒 Files selected for processing (10)
tensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_laguna.pytensorrt_llm/_torch/models/modeling_qwen3_moe.pytensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/impl_contract.pytensorrt_llm/_torch/modules/fused_moe/impl_environment.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.pytensorrt_llm/_torch/modules/fused_moe/moe_resolution.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tensorrt_llm/_torch/models/modeling_qwen3_moe.py
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
- tensorrt_llm/_torch/models/modeling_laguna.py
- tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py
- tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
- tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
- tensorrt_llm/_torch/modules/fused_moe/impl_contract.py
- tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py
f5f58a4 to
7d0f4d2
Compare
7d0f4d2 to
9a76eed
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py (2)
21-21: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winModernize and complete the annotations.
Line 21 imports legacy generic aliases. Replace
Dict,FrozenSet,List,Optional,Tuple,Type, andUnionusages with built-in generic types and|.Line 430 leaves
kwargsunannotated. Define a precise typed keyword surface that matchesresolve_moe_impl, then usetypeas the return annotation.As per coding guidelines: “Annotate every function” and “prefer built-in generic types and
|.”Also applies to: 430-430
🤖 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/moe_resolution.py` at line 21, Modernize annotations throughout moe_resolution.py by replacing Dict, FrozenSet, List, Optional, Tuple, Type, and Union usages with built-in generics and | syntax, removing obsolete imports. In resolve_moe_impl, annotate kwargs with a precise typed keyword surface matching the function’s accepted parameters, and change its return annotation to type; ensure every function in the touched code is annotated.Source: Coding guidelines
106-108: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd complete Google-style interface documentation.
These public contracts omit
Args,Returns, andRaisessections. Line 229 also accepts Tensor arguments without documented dimensions or dtype constraints. Add complete Google-style docstrings for these classes and functions.As per coding guidelines: “Use docstrings rather than comments for externally usable interfaces, Google-style docstrings for classes and functions, and document public Tensor-like argument dimensions and constrained dtypes.”
Also applies to: 121-141, 186-198, 229-235, 249-254, 287-289, 301-304, 315-329, 418-430
🤖 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/moe_resolution.py` around lines 106 - 108, Add complete Google-style docstrings to the externally usable classes and functions in this module, including _legacy_backend_name and the referenced ranges. Document every argument, return value, and raised exception; for the Tensor arguments in the function around line 229, explicitly state required dimensions and constrained dtypes. Preserve existing behavior and use the actual symbols and contracts rather than generic descriptions.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py`:
- Line 21: Modernize annotations throughout moe_resolution.py by replacing Dict,
FrozenSet, List, Optional, Tuple, Type, and Union usages with built-in generics
and | syntax, removing obsolete imports. In resolve_moe_impl, annotate kwargs
with a precise typed keyword surface matching the function’s accepted
parameters, and change its return annotation to type; ensure every function in
the touched code is annotated.
- Around line 106-108: Add complete Google-style docstrings to the externally
usable classes and functions in this module, including _legacy_backend_name and
the referenced ranges. Document every argument, return value, and raised
exception; for the Tensor arguments in the function around line 229, explicitly
state required dimensions and constrained dtypes. Preserve existing behavior and
use the actual symbols and contracts rather than generic descriptions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 62478ed2-c1d9-4ef3-8dde-8091d349ec7c
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/moe_resolution.py
654a895 to
ca4ac61
Compare
can_implement was not a function of its arguments: it read the SM version, probed optional imports, and in three places needed a constructed instance. The same request could therefore resolve differently on two ranks of one job, and a selection could not be replayed from a record. Machine facts are now collected once into MoEEnvironment and carried in MoEDeployment, so can_implement methods are pure classmethods over (problem, deployment). resolve_moe_impl is the single impl-selection entry point and returns a MoEResolutionReport naming the winner, every rejected candidate with a reason code, the ordered eligible list, and the environment fingerprint. ModelConfig.resolve_moe_backend stays a separate earlier phase (AUTO to backend literal during checkpoint load); merging it is blocked by a real quant_config cycle and is documented as such. Identity is revised for the S4 pin path: MoEImplId is provider.technique.kernel_name.quant, MoEImplQuery parses partial names by value with ordered segments, IMPL_PRIORITY replaces per-backend candidate lists, and selected_by is pinned/heuristic/failed (no auto while requests are still legacy backend literals). Routed-expert MoE LoRA quant constraints move into CutlassFusedMoE.can_implement; moe_lora_enabled is narrowed to has_moe_lora_targets so attention-only LoRA does not force Cutlass. Duplicate eligibility asserts that already live in can_implement are removed from backend __init__/validate paths. Signed-off-by: xxi <xxi@nvidia.com>
The fused_moe kernel already implements SwigluBias for BF16/FP16; rejecting quant_algo=None was a selection-layer overreach that broke dummy gpt-oss CUTLASS. Signed-off-by: xxi <xxi@nvidia.com>
…mits Allow MiniMax NVFP4 Cutlass SwiGLU without expert bias, and reject FlashInfer BF16 when the per-rank intermediate size is not 128-aligned. Signed-off-by: xxi <xxi@nvidia.com>
ca4ac61 to
c478216
Compare
|
/bot run --disable-fail-fast/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #66473 [ run ] triggered by Bot. Commit: |
|
PR_Github #66474 [ ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #66475 [ run ] triggered by Bot. Commit: |
|
PR_Github #66473 [ run ] completed with state |
|
PR_Github #66475 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66500 [ run ] triggered by Bot. Commit: |
|
PR_Github #66500 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66513 [ run ] triggered by Bot. Commit: |
|
PR_Github #66513 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66518 [ run ] triggered by Bot. Commit: |
|
PR_Github #66518 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66520 [ run ] triggered by Bot. Commit: |
|
PR_Github #66520 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66529 [ run ] triggered by Bot. Commit: |
|
PR_Github #66529 [ run ] completed with state |
9 participants
Summary
can_implementa pure classmethod over(problem, deployment)by freezing machine facts intoMoEEnvironmentonMoEDeployment.resolve_moe_implthat returnsMoEResolutionReport(winner, rejected trail with reason codes, eligible order, env fingerprint), replacing the old dualget_moe_cls/resolve_moe_clspaths.model_configviaderive_moe_layer_shapesso call sites only pass what config cannot supply; move MoE LoRA quant gates into Cutlasscan_implement.Test plan
test_moe_backend_selection_consistency.py+test_moe_impl_contracts.py(1579 passed)test_moe_module.py(666 passed, 1908 skipped)test_moe_backend.py(289 passed, 24 skipped, 5 failed — 4 knowntest_trtllm_bf16_unquantized_moe[*-fused_routing]+ 1test_megamoe_init_rejects_uneven_num_slots_with_value_errorto follow up)Dev Engineer Review
MoEEnvironmentandMoEDeployment.MoEProblem/MoEDeployment/MoEEligibilitycontract.top_kfrommodel_config.CutlassFusedMoE.can_implement.QA Engineer Review
tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.pytests/unittest/_torch/modules/moe/test_moe_backend.pytests/unittest/_torch/modules/moe/test_moe_module.pytests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.pytests/unittest/_torch/modules/moe/moe_test_utils.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/microbenchmarks/bench_moe/backend.pytests/microbenchmarks/bench_moe/search.pytests/integration/test_lists/,test-db/,qa/, orwaives.txtchanges were identified.