From 7f97aa9bd63a95e1c3b53bccd8d9748cceb47985 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Thu, 6 Aug 2026 01:30:37 -0700 Subject: [PATCH 1/6] [None][feat] support draft model MoE backend override Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/llm_args.py | 6 ++ tensorrt_llm/_torch/models/modeling_dspark.py | 56 ++++++++---- .../_torch/models/modeling_speculative.py | 44 +++++++++- tensorrt_llm/_torch/speculative/utils.py | 1 + tensorrt_llm/llmapi/llm_args.py | 54 ++++++++++-- .../modeling/test_modeling_speculative.py | 59 +++++++++++++ .../hw_agnostic/test_dspark_eplb_config.py | 25 ++++++ .../speculative/hw_agnostic/test_mtp.py | 86 +++++++++++++++++++ tests/unittest/llmapi/test_llm_args.py | 86 ++++++++++++++++++- 9 files changed, 388 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index fd56bbb4a13a..509aae61ffa8 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -122,6 +122,12 @@ def validate_supported_speculative_config(self): if spec_config is None: return self + if spec_config.moe_backend is not None: + raise ValueError( + "AutoDeploy does not support speculative_config.moe_backend. " + "This draft-model override is available only with the PyTorch backend." + ) + if isinstance(spec_config, MTPDecodingConfig): if spec_config.use_mtp_vanilla: raise ValueError( diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 58475da2389b..65fbd52b4f1f 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -94,6 +94,7 @@ class or the edge would become a cycle. from ..._utils import is_sm_100f from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..distributed import AllReduceParams +from ..model_config import ModelConfig from ..modules.linear import Linear from ..modules.mhc.hyper_connection import HCHead from ..modules.rms_norm import RMSNorm @@ -1055,6 +1056,7 @@ def __init__( aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], num_stages: Optional[int] = None, block_size: Optional[int] = None, + draft_moe_backend: Optional[str] = None, ): super().__init__() config = model_config.pretrained_config @@ -1115,7 +1117,9 @@ def __init__( # buffers. The draft experts are physically MXFP4 (same as the main # MoE layers), so copy a main MoE layer's experts quant onto the # draft layer keys. - draft_model_config = self._derive_draft_model_config(model_config, base, self.num_stages) + draft_model_config = self._derive_draft_model_config( + model_config, base, self.num_stages, draft_moe_backend=draft_moe_backend + ) self.mtp_layers = nn.ModuleList( [ DSv4DSparkBlock( @@ -1275,7 +1279,9 @@ def _dspark_freqs_table(self, device: torch.device) -> torch.Tensor: return cached @classmethod - def _derive_draft_model_config(cls, model_config, base: int, num_stages: int): + def _derive_draft_model_config( + cls, model_config, base: int, num_stages: int, draft_moe_backend: Optional[str] = None + ): """Return a draft-only ``model_config`` copy with draft-specific fixes. Applies (1) the ``compress_ratios`` draft slice and (2) the @@ -1283,23 +1289,30 @@ def _derive_draft_model_config(cls, model_config, base: int, num_stages: int): experts. A single shallow copy is made (and only when something needs to change) so the shared ``model_config`` and the target model are untouched. - The draft MoE backend is **inherited** from the target's - ``model_config.moe_backend`` (carried by the shallow copy) — not pinned — - matching every other drafter (the MTP module reuses the V4 decoder layer, - whose MoE is built with ``moe_backend=model_config.moe_backend``; separate - Eagle3/DFlash drafts resolve it from their own config the same way). The - draft ``mtp.*`` stages are full V4 blocks, so they share the target's - MXFP4 ``n_routed_experts=384`` / ``n_group=8`` (= 48 experts/group) layout - and therefore the same backend constraints: pick a backend that supports - it (CUTLASS today, DeepGEMM megaMoE once available) on the target and the - draft follows. Note the TRTLLM-Gen ``blockScaleMoe`` routing kernel asserts - ``experts/group <= 32`` (warp size), so it is incompatible with this layout - for both the target and the draft. + The draft MoE backend inherits ``model_config.moe_backend`` unless + ``draft_moe_backend`` is set. AUTO is resolved after the draft-specific + quantization normalization below, so backend selection uses the draft + weights rather than the target's resolved backend. The draft ``mtp.*`` + stages are full V4 blocks, so they share the target's MXFP4 + ``n_routed_experts=384`` / ``n_group=8`` (= 48 experts/group) layout + and therefore the same backend constraints: select a backend that + supports it (CUTLASS today, DeepGEMM megaMoE once available). The + TRTLLM-Gen ``blockScaleMoe`` routing kernel asserts + ``experts/group <= 32`` (warp size), so it is incompatible with this + layout for both the target and the draft. """ new_sa = cls._draft_sparse_config(model_config, base, num_stages) new_qcd = cls._draft_quant_config_dict(model_config, base, num_stages) new_qc = cls._draft_normalized_quant_config(model_config) - if new_sa is None and new_qcd is None and new_qc is None: + resolved_moe_backend = None + if draft_moe_backend is not None: + architectures = getattr(model_config.pretrained_config, "architectures", None) or [] + architecture = architectures[0] if architectures else "" + draft_quant_config = new_qc if new_qc is not None else model_config.quant_config + resolved_moe_backend = ModelConfig.resolve_moe_backend( + draft_moe_backend, architecture, quant_config=draft_quant_config + ) + if new_sa is None and new_qcd is None and new_qc is None and resolved_moe_backend is None: return model_config draft_cfg = copy.copy(model_config) # ModelConfig is a frozen dataclass; bypass the guard for these fields. @@ -1309,6 +1322,8 @@ def _derive_draft_model_config(cls, model_config, base: int, num_stages: int): object.__setattr__(draft_cfg, "quant_config_dict", new_qcd) if new_qc is not None: object.__setattr__(draft_cfg, "quant_config", new_qc) + if resolved_moe_backend is not None: + object.__setattr__(draft_cfg, "moe_backend", resolved_moe_backend) return draft_cfg @staticmethod @@ -1870,13 +1885,21 @@ class DSv4DSparkForCausalLM(nn.Module): attention weights from the in-memory state dict. """ - def __init__(self, draft_config, aux_stream_dict=None, num_stages=None, block_size=None): + def __init__( + self, + draft_config, + aux_stream_dict=None, + num_stages=None, + block_size=None, + draft_moe_backend: Optional[str] = None, + ): super().__init__() self.dspark_model = DSv4DSparkDraftModel( draft_config, aux_stream_dict, num_stages=num_stages, block_size=block_size, + draft_moe_backend=draft_moe_backend, ) # Generic handles expected by the loader / weight mappers. self.model = self.dspark_model @@ -2169,6 +2192,7 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): getattr(model, "aux_stream_dict", None), num_stages=num_stages, block_size=model_config.spec_config.block_size, + draft_moe_backend=getattr(model_config.spec_config, "moe_backend", None), ) # No per-model_type table here. ``DFlashForCausalLM.__init__`` already diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 88827b1b9428..40fcb43a47eb 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # +import copy import inspect from dataclasses import replace from typing import Dict, Generic, List, Optional, Tuple @@ -1307,6 +1308,29 @@ def __init__( self.embed_tokens = model.embed_tokens +def _get_requested_draft_moe_backend(model_config: ModelConfig, + spec_config: object) -> str: + """Return the draft MoE backend request, preserving target inheritance.""" + requested_backend = getattr(spec_config, "moe_backend", None) + return (model_config.moe_backend + if requested_backend is None else requested_backend) + + +def _copy_model_config_with_moe_backend( + model_config: ModelConfig, requested_moe_backend: str) -> ModelConfig: + """Copy a ModelConfig and resolve its MoE backend against its own weights.""" + architectures = getattr(model_config.pretrained_config, "architectures", + None) or [] + architecture = architectures[0] if architectures else "" + resolved_moe_backend = ModelConfig.resolve_moe_backend( + requested_moe_backend, + architecture, + quant_config=model_config.quant_config) + draft_config = copy.copy(model_config) + object.__setattr__(draft_config, "moe_backend", resolved_moe_backend) + return draft_config + + def external_drafter_config_kwargs(model_config, spec_config) -> dict: """`ModelConfig.from_pretrained` kwargs for a one-model external drafter. @@ -1326,7 +1350,7 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: kwargs = dict( trust_remote_code=True, attn_backend=model_config.attn_backend, - moe_backend=model_config.moe_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec_config), mapping=model_config.mapping, spec_config=None, # Avoid recursive spec-dec max_num_tokens=model_config.max_num_tokens, @@ -1460,6 +1484,8 @@ def __init__(self, if spec_config and spec_config.spec_dec_mode.use_one_engine(): # Only create draft_model for modes MTP, Eagle3 (not SA) if not spec_config.spec_dec_mode.is_sa(): + requested_draft_moe_backend = _get_requested_draft_moe_backend( + model_config, spec_config) if spec_config.spec_dec_mode.is_eagle3_one_model(): if spec_config.eagle3_model_arch == "mistral_large3": from tensorrt_llm._torch.models.checkpoints.mistral.config_loader import \ @@ -1467,18 +1493,28 @@ def __init__(self, self.draft_config = MistralConfigLoader().load( spec_config.speculative_model, mapping=model_config.mapping, - moe_backend=model_config.moe_backend, + moe_backend=requested_draft_moe_backend, moe_max_num_tokens=model_config.moe_max_num_tokens, max_num_tokens=model_config.max_num_tokens, moe_load_balancer=model_config.moe_load_balancer, skip_create_weights_in_init=True, ) + if getattr(spec_config, "moe_backend", + None) is not None: + # Unlike ModelConfig.from_pretrained, the Mistral + # loader does not resolve AUTO after loading quant + # metadata. Resolve it against the draft config now, + # before constructing any draft modules. + self.draft_config = \\ + _copy_model_config_with_moe_backend( + self.draft_config, + requested_draft_moe_backend) elif spec_config.eagle3_model_arch == "llama3": self.draft_config = ModelConfig.from_pretrained( model_config.spec_config.speculative_model, trust_remote_code=True, attn_backend=model_config.attn_backend, - moe_backend=model_config.moe_backend, + moe_backend=requested_draft_moe_backend, mapping=model_config.mapping, spec_config=model_config.spec_config, max_num_tokens=model_config.max_num_tokens, @@ -1496,7 +1532,7 @@ def __init__(self, spec_config.speculative_model, trust_remote_code=True, attn_backend=model_config.attn_backend, - moe_backend=model_config.moe_backend, + moe_backend=requested_draft_moe_backend, mapping=model_config.mapping, spec_config=None, max_num_tokens=model_config.max_num_tokens, diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index b4fc992e18da..5bf3c4a2d415 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -840,6 +840,7 @@ def update_spec_config_from_model_config(spec_config, if num_nextn_predict_layers is None: num_nextn_predict_layers = 1 spec_config.num_nextn_predict_layers = num_nextn_predict_layers + spec_config._validate_moe_backend_compatibility(model_config_resolved=True) is_vanilla = spec_config.spec_dec_mode.is_mtp_vanilla() # Resolve max_draft_len when the user didn't set it: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 59c5030de54c..1a830b687168 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1601,16 +1601,18 @@ def get_layer_initial_global_assignments( return assignments +_MoeBackend = Literal["AUTO", "CUTLASS", "CUTEDSL", "TRTLLM", "DEEPGEMM", + "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", + "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"] + + class MoeConfig(StrictBaseModel): """Configuration for MoE.""" - backend: Literal[ - "AUTO", "CUTLASS", "CUTEDSL", "TRTLLM", "DEEPGEMM", "DENSEGEMM", - "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL"] = Field( - default='AUTO', - description="MoE backend to use. " - "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." - ) + backend: _MoeBackend = Field( + default='AUTO', + description="MoE backend to use. " + "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." + ) max_num_tokens: Optional[int] = Field( default=None, @@ -1964,6 +1966,12 @@ class DecodingBaseConfig(StrictBaseModel): "draft model, depending on the target model implementation. Pointing it at the target checkpoint uses the " "target's embedded mtp.* weights.") + moe_backend: Optional[_MoeBackend] = Field( + default=None, + description= + "MoE backend override for the speculative (draft) model on the PyTorch backend. None preserves the existing behavior, AUTO selects a backend from the draft model configuration, and a concrete backend applies only to the draft model. The resolved backend may fall back based on model, quantization, and hardware support. One-engine MTP backed by target-checkpoint draft layers does not support this override because target and draft layers share quantization metadata; shared-KV external assistants are supported." + ) + max_concurrency: Optional[PositiveInt] = Field( default=None, description= @@ -2164,6 +2172,34 @@ def supports_backend(self, backend: str) -> bool: """ return True + def _validate_moe_backend_compatibility(self, + *, + model_config_resolved: bool = False + ) -> None: + if self.moe_backend is None: + return + + spec_mode = self.spec_dec_mode + unsupported_internal_mtp = (spec_mode.is_mtp_vanilla() + or (model_config_resolved + and spec_mode.is_mtp_eagle_one_model() + and not self._use_shared_kv_cache)) + if unsupported_internal_mtp: + raise ValueError( + "speculative_config.moe_backend does not support " + "one-engine MTP backed by target-checkpoint draft layers " + "because target and draft layers share quantization metadata. " + "Leave moe_backend unset to inherit the target backend, or " + "use a separate draft checkpoint.") + + has_neural_drafter = (spec_mode.has_draft_model() + or (spec_mode.use_one_engine() + and not spec_mode.is_sa())) + if not has_neural_drafter: + raise ValueError("speculative_config.moe_backend requires a neural " + "draft model or draft layers, but decoding_type " + f"{self.decoding_type} does not use one.") + @property def uses_replacement_heads(self) -> bool: """Whether `speculative_model` contains replacement MTP heads.""" @@ -6213,6 +6249,8 @@ def validate_speculative_config(self): exclude={"decoding_type"}) self.speculative_config = Eagle3DecodingConfig(**eagle_data) + self.speculative_config._validate_moe_backend_compatibility() + if self.speculative_config.use_rejection_sampling: # Supported paths: Eagle3 one-model, MTP-Eagle one-model, # vanilla MTP, PARD, DFlash, DraftTarget one-model. Classify by diff --git a/tests/unittest/_torch/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index 8aa29ceb6d85..0afd708c77d5 100644 --- a/tests/unittest/_torch/modeling/test_modeling_speculative.py +++ b/tests/unittest/_torch/modeling/test_modeling_speculative.py @@ -29,8 +29,11 @@ from tensorrt_llm._torch.models.modeling_speculative import ( Eagle3ForCausalLM, SpecDecOneEngineForCausalLM, + _copy_model_config_with_moe_backend, + external_drafter_config_kwargs, ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode class _FakeDraftModel(nn.Module): @@ -425,3 +428,59 @@ def test_dflash_trtllm_gen_buffers_reject_capture_time_allocation(): wrapper._dflash_trtllm_gen_counters = torch.empty(16, dtype=torch.uint8, device="meta") with pytest.raises(RuntimeError, match="counter buffer.*before CUDA graph capture"): _prepare_dflash_buffers(wrapper, 2) + + +# --------------------------------------------------------------------------- +# One-engine draft MoE backend selection +# --------------------------------------------------------------------------- + + +def _draft_backend_test_model_config(moe_backend: str = "CUTLASS") -> ModelConfig: + return ModelConfig( + pretrained_config=PretrainedConfig( + architectures=["DraftBackendTestForCausalLM"], + hidden_size=64, + vocab_size=128, + num_hidden_layers=2, + ), + moe_backend=moe_backend, + ) + + +def _external_spec_config(moe_backend: str | None) -> SimpleNamespace: + return SimpleNamespace( + spec_dec_mode=SpeculativeDecodingMode.PARD, + moe_backend=moe_backend, + ) + + +def test_external_draft_moe_backend_none_inherits_target() -> None: + """None preserves the existing target-backend inheritance behavior.""" + model_config = _draft_backend_test_model_config("CUTLASS") + + kwargs = external_drafter_config_kwargs(model_config, _external_spec_config(None)) + + assert kwargs["moe_backend"] == "CUTLASS" + + +def test_external_draft_moe_backend_auto_reaches_draft_loader() -> None: + """AUTO remains unresolved until the draft checkpoint quant config is read.""" + model_config = _draft_backend_test_model_config("TRTLLM") + + kwargs = external_drafter_config_kwargs(model_config, _external_spec_config("AUTO")) + + assert kwargs["moe_backend"] == "AUTO" + + +def test_loaded_draft_moe_backend_uses_isolated_model_config() -> None: + """Resolving a loaded draft config does not modify another config.""" + target_config = _draft_backend_test_model_config("CUTLASS") + with patch.object(ModelConfig, "resolve_moe_backend", return_value="TRTLLM") as resolve_backend: + draft_config = _copy_model_config_with_moe_backend(target_config, "AUTO") + + assert draft_config is not target_config + assert draft_config.moe_backend == "TRTLLM" + assert target_config.moe_backend == "CUTLASS" + resolve_backend.assert_called_once_with( + "AUTO", "DraftBackendTestForCausalLM", quant_config=target_config.quant_config + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py index 4a342e6ee8ef..c4738c0bb9a9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py @@ -149,6 +149,31 @@ def test_external_drafter_kwargs_are_stable_across_modes(): assert set(dspark) - set(common) == {"moe_load_balancer"} +def test_dspark_draft_backend_auto_resolves_on_isolated_copy(): + quant_config = SimpleNamespace(quant_algo=None) + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), + sparse_attention_config=None, + quant_config_dict=None, + quant_config=quant_config, + moe_backend="CUTLASS", + ) + + with patch.object( + modeling_dspark.ModelConfig, "resolve_moe_backend", return_value="TRTLLM" + ) as resolve_backend: + draft_config = modeling_dspark.DSparkDraftModel._derive_draft_model_config( + model_config, NUM_HIDDEN_LAYERS, NUM_STAGES, "AUTO" + ) + + assert draft_config is not model_config + assert draft_config.moe_backend == "TRTLLM" + assert model_config.moe_backend == "CUTLASS" + resolve_backend.assert_called_once_with( + "AUTO", "DeepseekV4ForCausalLM", quant_config=quant_config + ) + + # -------------------------------------------------------------------------- # 2. stage-layer placement coverage # -------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index 525f4cedc8e5..d142ffd49e2b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from types import SimpleNamespace @@ -1768,6 +1783,77 @@ class TargetModel: assert not should_use_separate_draft_kv_cache(spec_config) +@pytest.mark.parametrize("num_nextn_predict_layers", [2, 3]) +def test_mtp_moe_backend_rejected_after_checkpoint_resolves_one_engine( + num_nextn_predict_layers, +): + spec_config = MTPDecodingConfig( + max_draft_len=1, + speculative_model="/tmp/assistant", + mtp_eagle_one_model=False, + moe_backend="CUTLASS", + ) + assert spec_config.spec_dec_mode.is_mtp_eagle() + model_config = SimpleNamespace( + architectures=["LlamaForCausalLM"], + num_nextn_predict_layers=num_nextn_predict_layers, + ) + + with pytest.raises(ValueError, match="does not support one-engine MTP"): + update_spec_config_from_model_config(spec_config, model_config) + + +def test_mtp_moe_backend_allowed_after_checkpoint_keeps_two_engine(): + spec_config = MTPDecodingConfig( + max_draft_len=1, + speculative_model="/tmp/assistant", + mtp_eagle_one_model=False, + moe_backend="CUTLASS", + ) + model_config = SimpleNamespace( + architectures=["LlamaForCausalLM"], + num_nextn_predict_layers=1, + ) + + update_spec_config_from_model_config(spec_config, model_config) + + assert spec_config.spec_dec_mode.is_mtp_eagle() + + +def test_mtp_moe_backend_rejected_for_internal_mtp_eagle_one_model(): + spec_config = MTPDecodingConfig( + max_draft_len=1, + speculative_model="/tmp/assistant", + moe_backend="CUTLASS", + ) + assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() + model_config = SimpleNamespace( + architectures=["LlamaForCausalLM"], + num_nextn_predict_layers=1, + ) + + with pytest.raises(ValueError, match="does not support one-engine MTP"): + update_spec_config_from_model_config(spec_config, model_config) + + +def test_mtp_moe_backend_allowed_for_shared_kv_external_assistant(): + spec_config = MTPDecodingConfig( + max_draft_len=1, + speculative_model="/tmp/assistant", + moe_backend="CUTLASS", + ) + model_config = SimpleNamespace( + architectures=["Gemma4ForCausalLM"], + num_nextn_predict_layers=1, + ) + + update_spec_config_from_model_config(spec_config, model_config) + + assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() + assert spec_config._use_shared_kv_cache + assert spec_config.moe_backend == "CUTLASS" + + def test_mtp_shared_kv_draft_inputs(): spec_config = MTPDecodingConfig( max_draft_len=3, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index e9c9925dfb37..a9d450305a7c 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -52,7 +52,7 @@ MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, - PeftCacheConfig, + NGramDecodingConfig, PeftCacheConfig, PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, @@ -185,6 +185,90 @@ def test_MTPDecodingConfig_default_draft_len_is_not_user_set(): assert "max_draft_len" in explicit_config.model_fields_set +@pytest.mark.cpu_only +class TestDecodingBaseConfigMoeBackend: + + def test_defaults_to_none(self): + config = DecodingBaseConfig() + + assert config.moe_backend is None + assert config.model_dump()["moe_backend"] is None + + @pytest.mark.parametrize( + "moe_backend", + [None, *get_args(MoeConfig.model_fields["backend"].annotation)], + ) + def test_accepts_every_moe_backend(self, moe_backend): + config = DecodingBaseConfig(moe_backend=moe_backend) + + assert config.moe_backend == moe_backend + + @pytest.mark.parametrize("moe_backend", ["INVALID", "cutlass", 0]) + def test_rejects_invalid_moe_backend(self, moe_backend): + with pytest.raises(ValidationError, match="moe_backend"): + DecodingBaseConfig(moe_backend=moe_backend) + + def test_model_dump_and_yaml_parsing(self): + config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") + + assert config.model_dump()["moe_backend"] == "CUTLASS" + + yaml_config = yaml.safe_load(""" +decoding_type: MTP +max_draft_len: 1 +moe_backend: TRTLLM +""") + restored = TypeAdapter(SpeculativeConfig).validate_python(yaml_config) + + assert isinstance(restored, MTPDecodingConfig) + assert restored.moe_backend == "TRTLLM" + assert restored.model_dump()["moe_backend"] == "TRTLLM" + + def test_autodeploy_rejects_override(self): + spec_config = MTPDecodingConfig(max_draft_len=1, + moe_backend="CUTLASS", + mtp_eagle_one_model=False) + + with pytest.raises(ValidationError, + match="available only with the PyTorch backend"): + AutoDeployLlmArgs(model="/target", speculative_config=spec_config) + + def test_rejects_explicit_vanilla_mtp_override(self): + spec_config = MTPDecodingConfig(max_draft_len=1, + moe_backend="CUTLASS", + use_mtp_vanilla=True) + + with pytest.raises(ValidationError, + match="does not support one-engine MTP"): + TorchLlmArgs(model=llama_model_path, speculative_config=spec_config) + + def test_defers_checkpoint_dependent_mtp_eagle_override_validation(self): + spec_config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") + + llm_args = TorchLlmArgs(model=llama_model_path, + speculative_config=spec_config) + + assert llm_args.speculative_config.moe_backend == "CUTLASS" + + def test_accepts_two_engine_mtp_override(self): + spec_config = MTPDecodingConfig(max_draft_len=1, + moe_backend="CUTLASS", + mtp_eagle_one_model=False) + + llm_args = TorchLlmArgs(model=llama_model_path, + speculative_config=spec_config) + + assert llm_args.speculative_config.moe_backend == "CUTLASS" + + def test_rejects_override_without_neural_drafter(self): + spec_config = NGramDecodingConfig(max_draft_len=1, + moe_backend="CUTLASS") + + with pytest.raises(ValidationError, + match="requires a neural draft model"): + TorchLlmArgs(model=llama_model_path, speculative_config=spec_config) + + @pytest.mark.cpu_only def test_rejection_sampling_allows_attention_dp(monkeypatch): """ADP (incl. ADP+LM-head-TP) supports rejection sampling. From 07641d72f742b1fcff99572b3e9fe4bd8142870f Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Tue, 25 Aug 2026 23:09:06 -0700 Subject: [PATCH 2/6] [None][fix] handle draft model MoE backend edge cases Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_dflash.py | 3 ++ tensorrt_llm/_torch/models/modeling_gemma4.py | 3 ++ .../_torch/models/modeling_gemma4mm.py | 4 ++ .../_torch/models/modeling_speculative.py | 6 ++- .../_torch/moe/fused_moe/interface.py | 9 +++- tensorrt_llm/llmapi/llm_args.py | 12 ++--- .../_torch/attention/test_mla_registry.py | 12 ++++- .../_torch/modeling/test_gemma4_multimodal.py | 2 + .../_torch/modeling/test_modeling_gemma4.py | 1 + .../speculative/hw_agnostic/test_mtp.py | 45 ++++++++++++------- tests/unittest/llmapi/test_llm_args.py | 13 ++---- 11 files changed, 75 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 577942639bf6..45a31bccce6a 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -348,6 +348,9 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): # Remove spec_config to prevent recursive spec-dec initialization draft_config_no_spec = replace(draft_config, spec_config=None, lm_head_gather_output=False) + # ModelConfig.extra_attrs is init=False, so dataclasses.replace() does + # not preserve the shared custom-op registries. + draft_config_no_spec.extra_attrs = draft_config.extra_attrs # Weights will be loaded later by ModelLoader.load_draft_weights() self.draft_model_full = DraftModelClass(draft_config_no_spec) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 283e6ef311e0..d5fa9c9c5da2 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -1666,6 +1666,9 @@ def __init__(self, model_config: ModelConfig): pretrained_config=assistant_text_config, spec_config=None, ) + # extra_attrs is init=False and would otherwise be reset by replace(), + # disconnecting the assistant's custom-op registries from the engine. + text_model_config.extra_attrs = model_config.extra_attrs super().__init__( Gemma4TextModel(text_model_config), config=model_config, diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index c52c3e2fd4ba..05476526e185 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -904,6 +904,9 @@ def get_sub_model_config( attn_backend=attn_backend, quant_config=quant_config, ) + # extra_attrs is init=False and would otherwise be reset by replace(). + # All submodels execute under the top-level engine registry. + sub_config.extra_attrs = model_config.extra_attrs if ( hasattr(sub_config.pretrained_config, "torch_dtype") and sub_config.pretrained_config.torch_dtype is None @@ -1072,6 +1075,7 @@ def __init__(self, model_config: ModelConfig[Gemma4Config]): self._mm_token_ids = torch.tensor(_mm_ids, dtype=torch.int32) model_config_cp = copy.deepcopy(model_config) + model_config_cp.extra_attrs = model_config.extra_attrs self.model_config = model_config_cp # --- Language model --- diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 40fcb43a47eb..35368160012a 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1193,6 +1193,9 @@ def __init__(self, draft_config): draft_config_no_spec = replace(draft_config, spec_config=None, lm_head_gather_output=False) + # ModelConfig.extra_attrs is init=False, so dataclasses.replace() does + # not preserve the shared custom-op registries. + draft_config_no_spec.extra_attrs = draft_config.extra_attrs # Weights will be loaded later by ModelLoader.load_draft_weights() self.draft_model_full = DraftModelClass(draft_config_no_spec) @@ -1539,8 +1542,7 @@ def __init__(self, moe_max_num_tokens=model_config.moe_max_num_tokens) self.draft_config.quant_config.kv_cache_quant_algo = \ model_config.quant_config.kv_cache_quant_algo - self.draft_config.extra_attrs = dict( - model_config.extra_attrs) + self.draft_config.extra_attrs = model_config.extra_attrs self.draft_config.extra_attrs[ _SPECULATIVE_POSITION_HEADROOM] = ( 2 * spec_config.tokens_per_gen_step) diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 3ef31b3b1f02..d915c002a700 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -690,8 +690,13 @@ def _register_layer(self, model_config: ModelConfig): if model_config is not None and self.layer_idx_str is not None: if "moe_layers" not in model_config.extra_attrs: model_config.extra_attrs["moe_layers"] = {} - assert self.layer_idx_str not in model_config.extra_attrs["moe_layers"], \ - f"Duplicate MoE layer for layer_idx={self.layer_idx_str}" + suffix = 0 + # ``layer_idx`` is local to a model stack, while one-model + # speculative decoding shares this registry across target and + # draft modules. Preserve every module under a stable unique key. + while self.layer_idx_str in model_config.extra_attrs["moe_layers"]: + self.layer_idx_str = str(self.layer_idx) + f"_{suffix}" + suffix += 1 model_config.extra_attrs["moe_layers"][ self.layer_idx_str] = weakref.ref(self) self.register_to_config = True diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1a830b687168..bf957cae6b97 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1969,7 +1969,7 @@ class DecodingBaseConfig(StrictBaseModel): moe_backend: Optional[_MoeBackend] = Field( default=None, description= - "MoE backend override for the speculative (draft) model on the PyTorch backend. None preserves the existing behavior, AUTO selects a backend from the draft model configuration, and a concrete backend applies only to the draft model. The resolved backend may fall back based on model, quantization, and hardware support. One-engine MTP backed by target-checkpoint draft layers does not support this override because target and draft layers share quantization metadata; shared-KV external assistants are supported." + "MoE backend override for the speculative (draft) model on the PyTorch backend. None preserves the existing behavior, AUTO selects a backend from the draft model configuration, and a concrete backend applies only to the draft model. The resolved backend may fall back based on model, quantization, and hardware support. Vanilla MTP and one-engine MTP-EAGLE backed by target-checkpoint or replacement-head draft layers do not support this override because target and draft layers share quantization metadata; full external MTP-EAGLE draft models are supported." ) max_concurrency: Optional[PositiveInt] = Field( @@ -2183,14 +2183,16 @@ def _validate_moe_backend_compatibility(self, unsupported_internal_mtp = (spec_mode.is_mtp_vanilla() or (model_config_resolved and spec_mode.is_mtp_eagle_one_model() - and not self._use_shared_kv_cache)) + and not self.uses_external_draft_model)) if unsupported_internal_mtp: raise ValueError( - "speculative_config.moe_backend does not support " - "one-engine MTP backed by target-checkpoint draft layers " + "speculative_config.moe_backend does not support one-engine MTP " + "backed by target-checkpoint or replacement-head draft layers; " + "vanilla MTP is also unsupported " "because target and draft layers share quantization metadata. " "Leave moe_backend unset to inherit the target backend, or " - "use a separate draft checkpoint.") + "for one-engine MTP-EAGLE, use a full external draft-model " + "checkpoint.") has_neural_drafter = (spec_mode.has_draft_model() or (spec_mode.use_one_engine() diff --git a/tests/unittest/_torch/attention/test_mla_registry.py b/tests/unittest/_torch/attention/test_mla_registry.py index c3e0beced9f1..9e69b7703de2 100644 --- a/tests/unittest/_torch/attention/test_mla_registry.py +++ b/tests/unittest/_torch/attention/test_mla_registry.py @@ -29,6 +29,7 @@ ) from tensorrt_llm._torch.attention.mla import MLA from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.moe.fused_moe.interface import MoE from tensorrt_llm.functional import PositionEmbeddingType @@ -169,7 +170,7 @@ def _apply_output_gate( _make_mla(config, cls=_GatedMLA) -def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: +def test_duplicate_layer_ids_preserve_all_registrations() -> None: target_config = ModelConfig(skip_create_weights_in_init=True) draft_config = ModelConfig(skip_create_weights_in_init=True) next_config = ModelConfig(skip_create_weights_in_init=True) @@ -193,6 +194,15 @@ def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: assert registry["0_0"]() is draft_mla assert registry["0_1"]() is next_mla + moe_layers = [nn.Module() for _ in range(3)] + for layer in moe_layers: + layer.layer_idx = 0 + layer.layer_idx_str = "0" + MoE._register_layer(layer, target_config) + + assert [layer.layer_idx_str for layer in moe_layers] == ["0", "0_0", "0_1"] + assert [ref() for ref in target_config.extra_attrs["moe_layers"].values()] == moe_layers + def _make_dsv4_epilogue_layer() -> SimpleNamespace: return SimpleNamespace( diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 9c4619413832..682d2e5071ec 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -817,6 +817,8 @@ def test_instantiation_with_vision(self): ) model = Gemma4ForConditionalGeneration(mc) + self.assertIs(model.model_config.extra_attrs, mc.extra_attrs) + self.assertIs(model.llm.model.model_config.extra_attrs, mc.extra_attrs) self.assertIsNotNone(model.llm) self.assertIsNotNone(model.vision_tower) self.assertIsNotNone(model.embed_vision) diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 84bedef0c009..3b6bf453b8f1 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -645,6 +645,7 @@ def test_assistant_uses_target_kv_sources(self): model_config = _make_assistant_model_config() model_config.extra_attrs["_speculative_position_headroom"] = 2 * 4 assistant = Gemma4AssistantForCausalLM(model_config) + self.assertIs(assistant.model.model_config.extra_attrs, model_config.extra_attrs) self.assertEqual(len(assistant.model.layers), 4) self.assertTrue(all(layer.is_kv_shared_layer for layer in assistant.model.layers)) self.assertEqual( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index d142ffd49e2b..78250b9abc3c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -1790,10 +1790,8 @@ def test_mtp_moe_backend_rejected_after_checkpoint_resolves_one_engine( spec_config = MTPDecodingConfig( max_draft_len=1, speculative_model="/tmp/assistant", - mtp_eagle_one_model=False, moe_backend="CUTLASS", ) - assert spec_config.spec_dec_mode.is_mtp_eagle() model_config = SimpleNamespace( architectures=["LlamaForCausalLM"], num_nextn_predict_layers=num_nextn_predict_layers, @@ -1803,40 +1801,55 @@ def test_mtp_moe_backend_rejected_after_checkpoint_resolves_one_engine( update_spec_config_from_model_config(spec_config, model_config) -def test_mtp_moe_backend_allowed_after_checkpoint_keeps_two_engine(): +def test_mtp_moe_backend_rejected_for_internal_mtp_eagle_one_model(): spec_config = MTPDecodingConfig( max_draft_len=1, - speculative_model="/tmp/assistant", - mtp_eagle_one_model=False, moe_backend="CUTLASS", ) + assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() model_config = SimpleNamespace( architectures=["LlamaForCausalLM"], num_nextn_predict_layers=1, ) - update_spec_config_from_model_config(spec_config, model_config) + with pytest.raises(ValueError, match="does not support one-engine MTP"): + update_spec_config_from_model_config(spec_config, model_config) - assert spec_config.spec_dec_mode.is_mtp_eagle() +@pytest.mark.parametrize( + ("architecture", "uses_shared_kv_cache"), + [("Gemma4ForCausalLM", True), ("LlamaForCausalLM", False)], +) +def test_mtp_moe_backend_allowed_for_full_external_assistant( + architecture, + uses_shared_kv_cache, +): + class ExternalDraftModelTarget: + build_mtp_draft_model_from_config = True -def test_mtp_moe_backend_rejected_for_internal_mtp_eagle_one_model(): spec_config = MTPDecodingConfig( max_draft_len=1, speculative_model="/tmp/assistant", moe_backend="CUTLASS", ) - assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() model_config = SimpleNamespace( - architectures=["LlamaForCausalLM"], + architectures=[architecture], num_nextn_predict_layers=1, ) - with pytest.raises(ValueError, match="does not support one-engine MTP"): - update_spec_config_from_model_config(spec_config, model_config) + update_spec_config_from_model_config(spec_config, model_config, ExternalDraftModelTarget) + + assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() + assert spec_config.uses_external_draft_model + assert spec_config._use_shared_kv_cache is uses_shared_kv_cache + assert should_use_separate_draft_kv_cache(spec_config) is not uses_shared_kv_cache + assert spec_config.moe_backend == "CUTLASS" -def test_mtp_moe_backend_allowed_for_shared_kv_external_assistant(): +def test_mtp_moe_backend_rejected_for_shared_kv_replacement_heads(): + class ReplacementHeadTarget: + pass + spec_config = MTPDecodingConfig( max_draft_len=1, speculative_model="/tmp/assistant", @@ -1847,11 +1860,11 @@ def test_mtp_moe_backend_allowed_for_shared_kv_external_assistant(): num_nextn_predict_layers=1, ) - update_spec_config_from_model_config(spec_config, model_config) + with pytest.raises(ValueError, match="does not support one-engine MTP"): + update_spec_config_from_model_config(spec_config, model_config, ReplacementHeadTarget) - assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() + assert spec_config.uses_replacement_heads assert spec_config._use_shared_kv_cache - assert spec_config.moe_backend == "CUTLASS" def test_mtp_shared_kv_draft_inputs(): diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index a9d450305a7c..642af78071d1 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -225,9 +225,7 @@ def test_model_dump_and_yaml_parsing(self): assert restored.model_dump()["moe_backend"] == "TRTLLM" def test_autodeploy_rejects_override(self): - spec_config = MTPDecodingConfig(max_draft_len=1, - moe_backend="CUTLASS", - mtp_eagle_one_model=False) + spec_config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") with pytest.raises(ValidationError, match="available only with the PyTorch backend"): @@ -250,15 +248,12 @@ def test_defers_checkpoint_dependent_mtp_eagle_override_validation(self): assert llm_args.speculative_config.moe_backend == "CUTLASS" - def test_accepts_two_engine_mtp_override(self): + def test_deprecated_two_engine_mtp_is_normalized_to_one_engine(self): spec_config = MTPDecodingConfig(max_draft_len=1, - moe_backend="CUTLASS", mtp_eagle_one_model=False) - llm_args = TorchLlmArgs(model=llama_model_path, - speculative_config=spec_config) - - assert llm_args.speculative_config.moe_backend == "CUTLASS" + assert spec_config.mtp_eagle_one_model + assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() def test_rejects_override_without_neural_drafter(self): spec_config = NGramDecodingConfig(max_draft_len=1, From b813ec6e80287a0e7d5a992087ae7334a3f9217c Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Mon, 7 Sep 2026 23:22:35 -0700 Subject: [PATCH 3/6] [None][feat] support embedded MTP MoE backend override Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- .../_torch/models/modeling_speculative.py | 80 ++++++++++++-- tensorrt_llm/llmapi/llm_args.py | 46 ++++---- .../usage/llm_args_golden_manifest.json | 18 ++++ .../modeling/test_modeling_speculative.py | 102 +++++++++++++++++- .../speculative/hw_agnostic/test_mtp.py | 25 +++-- tests/unittest/llmapi/test_llm_args.py | 34 +++--- 6 files changed, 243 insertions(+), 62 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 35368160012a..5af73e0add7c 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -12,6 +12,7 @@ from transformers import LlamaConfig, PretrainedConfig from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantConfig from ...functional import PositionEmbeddingType from ..attention.attention import Attention @@ -1320,7 +1321,9 @@ def _get_requested_draft_moe_backend(model_config: ModelConfig, def _copy_model_config_with_moe_backend( - model_config: ModelConfig, requested_moe_backend: str) -> ModelConfig: + model_config: ModelConfig, + requested_moe_backend: str, + quant_config: Optional[QuantConfig] = None) -> ModelConfig: """Copy a ModelConfig and resolve its MoE backend against its own weights.""" architectures = getattr(model_config.pretrained_config, "architectures", None) or [] @@ -1328,12 +1331,56 @@ def _copy_model_config_with_moe_backend( resolved_moe_backend = ModelConfig.resolve_moe_backend( requested_moe_backend, architecture, - quant_config=model_config.quant_config) + quant_config=(model_config.quant_config + if quant_config is None else quant_config)) draft_config = copy.copy(model_config) - object.__setattr__(draft_config, "moe_backend", resolved_moe_backend) + # ModelConfig is frozen after loading. Temporarily thaw only the isolated + # draft copy so the target config remains unchanged. + was_frozen = draft_config._frozen + draft_config._frozen = False + try: + draft_config.moe_backend = resolved_moe_backend + finally: + draft_config._frozen = was_frozen return draft_config +def _get_mtp_moe_quant_config(model_config: ModelConfig, + layer_idx: int) -> QuantConfig: + """Return the routed-expert quant config for an embedded MTP layer. + + The first constructed MTP layer represents all MTP layers when resolving + AUTO because they share one backend setting. + """ + layer_quant_configs = model_config.quant_config_dict or {} + prefixes = [f"model.layers.{layer_idx}."] + num_hidden_layers = getattr(model_config.pretrained_config, + "num_hidden_layers", None) + if num_hidden_layers is not None and layer_idx >= num_hidden_layers: + prefixes.append(f"mtp.layers.{layer_idx - num_hidden_layers}.") + + for prefix in prefixes: + quant_config = layer_quant_configs.get(f"{prefix}mlp.experts") + if quant_config is not None: + return quant_config + for prefix in prefixes: + for name, quant_config in layer_quant_configs.items(): + if name.startswith(prefix) and ".experts" in name: + return quant_config + return model_config.quant_config + + +def _enable_trtllm_moe_preload_for_draft(model: nn.Module, + moe_backend: str) -> None: + """Preserve TRTLLM-Gen's serial weight-loading order for draft-only use.""" + preload_weight_modules = getattr(model, "preload_weight_modules", None) + if moe_backend != "TRTLLM" or preload_weight_modules is None: + return + for module_name in ("experts", "routing_method", "all_reduce"): + if module_name not in preload_weight_modules: + preload_weight_modules.append(module_name) + + def external_drafter_config_kwargs(model_config, spec_config) -> dict: """`ModelConfig.from_pretrained` kwargs for a one-model external drafter. @@ -1390,7 +1437,26 @@ def _build_eagle3_one_model_draft(model_config, draft_config, lm_head, model): @register_draft_model(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) def _build_mtp_one_model_draft(model_config, draft_config, lm_head, model): """Build the one-model MTP drafter (vanilla MTP and MTP-Eagle share it).""" - return MTPForCausalLM(model_config, + mtp_model_config = model_config + requested_moe_backend = getattr(model_config.spec_config, "moe_backend", + None) + if requested_moe_backend is not None: + start_layer_idx = model_config.pretrained_config.num_hidden_layers + mtp_model_config = _copy_model_config_with_moe_backend( + model_config, + requested_moe_backend, + quant_config=_get_mtp_moe_quant_config(model_config, + start_layer_idx)) + model_type = model_config.pretrained_config.model_type + if (model_type in {"nemotron_h", "nemotron_h_puzzle"} + and mtp_model_config.moe_backend != model_config.moe_backend): + raise ValueError( + "Nemotron-H embedded MTP layers cannot use a different MoE " + "backend from the target model because their checkpoint " + "weight mapper uses one shared backend-dependent layout.") + _enable_trtllm_moe_preload_for_draft(model, + mtp_model_config.moe_backend) + return MTPForCausalLM(mtp_model_config, model_config.pretrained_config.num_hidden_layers, lm_head, model) @@ -1508,10 +1574,8 @@ def __init__(self, # loader does not resolve AUTO after loading quant # metadata. Resolve it against the draft config now, # before constructing any draft modules. - self.draft_config = \\ - _copy_model_config_with_moe_backend( - self.draft_config, - requested_draft_moe_backend) + self.draft_config = _copy_model_config_with_moe_backend( + self.draft_config, requested_draft_moe_backend) elif spec_config.eagle3_model_arch == "llama3": self.draft_config = ModelConfig.from_pretrained( model_config.spec_config.speculative_model, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index bf957cae6b97..a7764039262c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1968,9 +1968,18 @@ class DecodingBaseConfig(StrictBaseModel): moe_backend: Optional[_MoeBackend] = Field( default=None, - description= - "MoE backend override for the speculative (draft) model on the PyTorch backend. None preserves the existing behavior, AUTO selects a backend from the draft model configuration, and a concrete backend applies only to the draft model. The resolved backend may fall back based on model, quantization, and hardware support. Vanilla MTP and one-engine MTP-EAGLE backed by target-checkpoint or replacement-head draft layers do not support this override because target and draft layers share quantization metadata; full external MTP-EAGLE draft models are supported." - ) + description=( + "MoE backend override for a neural draft model or embedded MTP " + "layers on the PyTorch backend. None inherits the target model's " + "backend. AUTO resolves from the draft checkpoint or embedded MTP " + "layer quantization, and a concrete backend applies only to the " + "draft model or layers. Resolution may fall back based on model, " + "quantization, and hardware support. Replacement-head MTP " + "checkpoints are unsupported because their independent " + "quantization metadata is not loaded. Nemotron-H embedded MTP " + "layers must inherit the target backend because their checkpoint " + "mapper uses a shared backend-dependent layout. Decoding methods " + "without a neural draft model ignore this option.")) max_concurrency: Optional[PositiveInt] = Field( default=None, @@ -2176,31 +2185,14 @@ def _validate_moe_backend_compatibility(self, *, model_config_resolved: bool = False ) -> None: - if self.moe_backend is None: + if (self.moe_backend is None or not model_config_resolved + or not self.uses_replacement_heads): return - - spec_mode = self.spec_dec_mode - unsupported_internal_mtp = (spec_mode.is_mtp_vanilla() - or (model_config_resolved - and spec_mode.is_mtp_eagle_one_model() - and not self.uses_external_draft_model)) - if unsupported_internal_mtp: - raise ValueError( - "speculative_config.moe_backend does not support one-engine MTP " - "backed by target-checkpoint or replacement-head draft layers; " - "vanilla MTP is also unsupported " - "because target and draft layers share quantization metadata. " - "Leave moe_backend unset to inherit the target backend, or " - "for one-engine MTP-EAGLE, use a full external draft-model " - "checkpoint.") - - has_neural_drafter = (spec_mode.has_draft_model() - or (spec_mode.use_one_engine() - and not spec_mode.is_sa())) - if not has_neural_drafter: - raise ValueError("speculative_config.moe_backend requires a neural " - "draft model or draft layers, but decoding_type " - f"{self.decoding_type} does not use one.") + raise ValueError( + "speculative_config.moe_backend does not support replacement-head " + "MTP checkpoints because their independent quantization metadata " + "is not loaded. Leave moe_backend unset to inherit the target " + "backend, or use a full external draft-model checkpoint.") @property def uses_replacement_heads(self) -> bool: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e4f0bcc10f7e..d9c232ceb325 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1634,6 +1634,24 @@ "kind": "value", "path": "speculative_config.max_total_draft_tokens" }, + { + "allowed_values": [ + "AUTO", + "CUTLASS", + "CUTEDSL", + "TRTLLM", + "DEEPGEMM", + "DENSEGEMM", + "VANILLA", + "TRITON", + "MARLIN", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL" + ], + "capture_policy": "literal|none", + "kind": "categorical", + "path": "speculative_config.moe_backend" + }, { "capture_policy": "int|none", "kind": "value", diff --git a/tests/unittest/_torch/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index 0afd708c77d5..9bde9b341114 100644 --- a/tests/unittest/_torch/modeling/test_modeling_speculative.py +++ b/tests/unittest/_torch/modeling/test_modeling_speculative.py @@ -16,7 +16,7 @@ """Unit tests for speculative modeling classes.""" from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, patch, sentinel import pytest import torch @@ -29,11 +29,14 @@ from tensorrt_llm._torch.models.modeling_speculative import ( Eagle3ForCausalLM, SpecDecOneEngineForCausalLM, + _build_mtp_one_model_draft, _copy_model_config_with_moe_backend, external_drafter_config_kwargs, ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo class _FakeDraftModel(nn.Module): @@ -484,3 +487,100 @@ def test_loaded_draft_moe_backend_uses_isolated_model_config() -> None: resolve_backend.assert_called_once_with( "AUTO", "DraftBackendTestForCausalLM", quant_config=target_config.quant_config ) + + +def test_internal_mtp_without_override_reuses_target_model_config() -> None: + target_config = _draft_backend_test_model_config("CUTEDSL") + target_config.spec_config = SimpleNamespace(moe_backend=None) + target_model = SimpleNamespace(aux_stream_dict={}, preload_weight_modules=[]) + + with patch( + "tensorrt_llm._torch.models.modeling_speculative.MTPForCausalLM", + return_value=sentinel.draft_model, + ) as mtp_cls: + draft_model = _build_mtp_one_model_draft( + target_config, None, sentinel.lm_head, target_model + ) + + assert draft_model is sentinel.draft_model + assert mtp_cls.call_args.args[0] is target_config + assert target_model.preload_weight_modules == [] + + +@pytest.mark.parametrize("requested_backend", ["TRTLLM", "AUTO"]) +@pytest.mark.parametrize( + "quant_config_key", + ["model.layers.2.mlp.experts", "mtp.layers.0.mlp.experts"], +) +def test_internal_mtp_moe_backend_uses_isolated_layer_config( + requested_backend: str, + quant_config_key: str, +) -> None: + target_config = _draft_backend_test_model_config("CUTEDSL") + target_config.spec_config = SimpleNamespace(moe_backend=requested_backend) + mtp_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + target_config.quant_config_dict = { + quant_config_key: mtp_quant_config, + } + target_config._frozen = True + target_model = SimpleNamespace(aux_stream_dict={}, preload_weight_modules=[]) + + with ( + patch.object(ModelConfig, "resolve_moe_backend", return_value="TRTLLM") as resolve_backend, + patch( + "tensorrt_llm._torch.models.modeling_speculative.MTPForCausalLM", + return_value=sentinel.draft_model, + ) as mtp_cls, + ): + draft_model = _build_mtp_one_model_draft( + target_config, None, sentinel.lm_head, target_model + ) + + mtp_model_config = mtp_cls.call_args.args[0] + assert draft_model is sentinel.draft_model + assert mtp_model_config is not target_config + assert mtp_model_config.moe_backend == "TRTLLM" + assert mtp_model_config.quant_config_dict is target_config.quant_config_dict + assert mtp_model_config.extra_attrs is target_config.extra_attrs + assert target_config.moe_backend == "CUTEDSL" + assert target_config._frozen + assert target_model.preload_weight_modules == ["experts", "routing_method", "all_reduce"] + resolve_backend.assert_called_once_with( + requested_backend, + "DraftBackendTestForCausalLM", + quant_config=mtp_quant_config, + ) + + +def test_internal_mtp_auto_resolves_from_layer_quantization() -> None: + target_config = _draft_backend_test_model_config("CUTEDSL") + target_config.spec_config = SimpleNamespace(moe_backend="AUTO") + target_config.quant_config = QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION) + target_config.quant_config_dict = { + "model.layers.2.mlp.experts": QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES), + } + target_config._frozen = True + target_model = SimpleNamespace(aux_stream_dict={}, preload_weight_modules=[]) + + with ( + patch("tensorrt_llm._torch.model_config.is_sm_100f", return_value=True), + patch( + "tensorrt_llm._torch.models.modeling_speculative.MTPForCausalLM", + return_value=sentinel.draft_model, + ) as mtp_cls, + ): + _build_mtp_one_model_draft(target_config, None, sentinel.lm_head, target_model) + + mtp_model_config = mtp_cls.call_args.args[0] + assert mtp_model_config.moe_backend == "TRTLLM" + assert target_config.moe_backend == "CUTEDSL" + + +def test_internal_mtp_rejects_nemotron_backend_mismatch() -> None: + target_config = _draft_backend_test_model_config("CUTLASS") + target_config.pretrained_config.model_type = "nemotron_h" + target_config.spec_config = SimpleNamespace(moe_backend="VANILLA") + target_model = SimpleNamespace(aux_stream_dict={}, preload_weight_modules=[]) + + with pytest.raises(ValueError, match="Nemotron-H embedded MTP layers"): + _build_mtp_one_model_draft(target_config, None, sentinel.lm_head, target_model) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index 78250b9abc3c..7304212eeba0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -1784,12 +1784,11 @@ class TargetModel: @pytest.mark.parametrize("num_nextn_predict_layers", [2, 3]) -def test_mtp_moe_backend_rejected_after_checkpoint_resolves_one_engine( - num_nextn_predict_layers, -): +def test_mtp_moe_backend_allowed_after_checkpoint_resolves_vanilla( + num_nextn_predict_layers: int, +) -> None: spec_config = MTPDecodingConfig( max_draft_len=1, - speculative_model="/tmp/assistant", moe_backend="CUTLASS", ) model_config = SimpleNamespace( @@ -1797,11 +1796,13 @@ def test_mtp_moe_backend_rejected_after_checkpoint_resolves_one_engine( num_nextn_predict_layers=num_nextn_predict_layers, ) - with pytest.raises(ValueError, match="does not support one-engine MTP"): - update_spec_config_from_model_config(spec_config, model_config) + update_spec_config_from_model_config(spec_config, model_config) + + assert spec_config.spec_dec_mode.is_mtp_vanilla() + assert spec_config.moe_backend == "CUTLASS" -def test_mtp_moe_backend_rejected_for_internal_mtp_eagle_one_model(): +def test_mtp_moe_backend_allowed_for_internal_mtp_eagle() -> None: spec_config = MTPDecodingConfig( max_draft_len=1, moe_backend="CUTLASS", @@ -1812,8 +1813,10 @@ def test_mtp_moe_backend_rejected_for_internal_mtp_eagle_one_model(): num_nextn_predict_layers=1, ) - with pytest.raises(ValueError, match="does not support one-engine MTP"): - update_spec_config_from_model_config(spec_config, model_config) + update_spec_config_from_model_config(spec_config, model_config) + + assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() + assert spec_config.moe_backend == "CUTLASS" @pytest.mark.parametrize( @@ -1846,7 +1849,7 @@ class ExternalDraftModelTarget: assert spec_config.moe_backend == "CUTLASS" -def test_mtp_moe_backend_rejected_for_shared_kv_replacement_heads(): +def test_mtp_moe_backend_rejected_for_shared_kv_replacement_heads() -> None: class ReplacementHeadTarget: pass @@ -1860,7 +1863,7 @@ class ReplacementHeadTarget: num_nextn_predict_layers=1, ) - with pytest.raises(ValueError, match="does not support one-engine MTP"): + with pytest.raises(ValueError, match="replacement-head MTP"): update_spec_config_from_model_config(spec_config, model_config, ReplacementHeadTarget) assert spec_config.uses_replacement_heads diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 642af78071d1..a1241c34db4c 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -188,7 +188,7 @@ def test_MTPDecodingConfig_default_draft_len_is_not_user_set(): @pytest.mark.cpu_only class TestDecodingBaseConfigMoeBackend: - def test_defaults_to_none(self): + def test_defaults_to_none(self) -> None: config = DecodingBaseConfig() assert config.moe_backend is None @@ -198,17 +198,17 @@ def test_defaults_to_none(self): "moe_backend", [None, *get_args(MoeConfig.model_fields["backend"].annotation)], ) - def test_accepts_every_moe_backend(self, moe_backend): + def test_accepts_every_moe_backend(self, moe_backend: str | None) -> None: config = DecodingBaseConfig(moe_backend=moe_backend) assert config.moe_backend == moe_backend @pytest.mark.parametrize("moe_backend", ["INVALID", "cutlass", 0]) - def test_rejects_invalid_moe_backend(self, moe_backend): + def test_rejects_invalid_moe_backend(self, moe_backend: object) -> None: with pytest.raises(ValidationError, match="moe_backend"): DecodingBaseConfig(moe_backend=moe_backend) - def test_model_dump_and_yaml_parsing(self): + def test_model_dump_and_yaml_parsing(self) -> None: config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") assert config.model_dump()["moe_backend"] == "CUTLASS" @@ -224,23 +224,25 @@ def test_model_dump_and_yaml_parsing(self): assert restored.moe_backend == "TRTLLM" assert restored.model_dump()["moe_backend"] == "TRTLLM" - def test_autodeploy_rejects_override(self): + def test_autodeploy_rejects_override(self) -> None: spec_config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") with pytest.raises(ValidationError, match="available only with the PyTorch backend"): AutoDeployLlmArgs(model="/target", speculative_config=spec_config) - def test_rejects_explicit_vanilla_mtp_override(self): + def test_accepts_explicit_vanilla_mtp_override(self) -> None: spec_config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS", use_mtp_vanilla=True) - with pytest.raises(ValidationError, - match="does not support one-engine MTP"): - TorchLlmArgs(model=llama_model_path, speculative_config=spec_config) + llm_args = TorchLlmArgs(model=llama_model_path, + speculative_config=spec_config) + + assert llm_args.speculative_config.moe_backend == "CUTLASS" - def test_defers_checkpoint_dependent_mtp_eagle_override_validation(self): + def test_defers_checkpoint_dependent_mtp_eagle_override_validation( + self) -> None: spec_config = MTPDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") llm_args = TorchLlmArgs(model=llama_model_path, @@ -248,20 +250,22 @@ def test_defers_checkpoint_dependent_mtp_eagle_override_validation(self): assert llm_args.speculative_config.moe_backend == "CUTLASS" - def test_deprecated_two_engine_mtp_is_normalized_to_one_engine(self): + def test_deprecated_two_engine_mtp_is_normalized_to_one_engine( + self) -> None: spec_config = MTPDecodingConfig(max_draft_len=1, mtp_eagle_one_model=False) assert spec_config.mtp_eagle_one_model assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() - def test_rejects_override_without_neural_drafter(self): + def test_ignores_override_without_neural_drafter(self) -> None: spec_config = NGramDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") - with pytest.raises(ValidationError, - match="requires a neural draft model"): - TorchLlmArgs(model=llama_model_path, speculative_config=spec_config) + llm_args = TorchLlmArgs(model=llama_model_path, + speculative_config=spec_config) + + assert llm_args.speculative_config.moe_backend == "CUTLASS" @pytest.mark.cpu_only From b175fe4eab06c6db518eb6c42b381b7610269506 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:56:23 -0700 Subject: [PATCH 4/6] [None][test] remove obsolete two-model MTP test Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tests/unittest/llmapi/test_llm_args.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index a1241c34db4c..87103f7bfa5e 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -250,14 +250,6 @@ def test_defers_checkpoint_dependent_mtp_eagle_override_validation( assert llm_args.speculative_config.moe_backend == "CUTLASS" - def test_deprecated_two_engine_mtp_is_normalized_to_one_engine( - self) -> None: - spec_config = MTPDecodingConfig(max_draft_len=1, - mtp_eagle_one_model=False) - - assert spec_config.mtp_eagle_one_model - assert spec_config.spec_dec_mode.is_mtp_eagle_one_model() - def test_ignores_override_without_neural_drafter(self) -> None: spec_config = NGramDecodingConfig(max_draft_len=1, moe_backend="CUTLASS") From a8cf31338378a7e542a9ab82fc7455f7e7584242 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:23:41 -0700 Subject: [PATCH 5/6] [None][test] move MoE registry test to MoE suite Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- .../unittest/_torch/attention/test_mla_registry.py | 12 +----------- tests/unittest/_torch/moe/test_moe_module.py | 14 +++++++++++++- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/unittest/_torch/attention/test_mla_registry.py b/tests/unittest/_torch/attention/test_mla_registry.py index 9e69b7703de2..c3e0beced9f1 100644 --- a/tests/unittest/_torch/attention/test_mla_registry.py +++ b/tests/unittest/_torch/attention/test_mla_registry.py @@ -29,7 +29,6 @@ ) from tensorrt_llm._torch.attention.mla import MLA from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.moe.fused_moe.interface import MoE from tensorrt_llm.functional import PositionEmbeddingType @@ -170,7 +169,7 @@ def _apply_output_gate( _make_mla(config, cls=_GatedMLA) -def test_duplicate_layer_ids_preserve_all_registrations() -> None: +def test_duplicate_layer_ids_preserve_all_mla_registrations() -> None: target_config = ModelConfig(skip_create_weights_in_init=True) draft_config = ModelConfig(skip_create_weights_in_init=True) next_config = ModelConfig(skip_create_weights_in_init=True) @@ -194,15 +193,6 @@ def test_duplicate_layer_ids_preserve_all_registrations() -> None: assert registry["0_0"]() is draft_mla assert registry["0_1"]() is next_mla - moe_layers = [nn.Module() for _ in range(3)] - for layer in moe_layers: - layer.layer_idx = 0 - layer.layer_idx_str = "0" - MoE._register_layer(layer, target_config) - - assert [layer.layer_idx_str for layer in moe_layers] == ["0", "0_0", "0_1"] - assert [ref() for ref in target_config.extra_attrs["moe_layers"].values()] == moe_layers - def _make_dsv4_epilogue_layer() -> SimpleNamespace: return SimpleNamespace( diff --git a/tests/unittest/_torch/moe/test_moe_module.py b/tests/unittest/_torch/moe/test_moe_module.py index 489c42b8ab01..e595faa66e5b 100644 --- a/tests/unittest/_torch/moe/test_moe_module.py +++ b/tests/unittest/_torch/moe/test_moe_module.py @@ -86,7 +86,7 @@ create_moe, ) from tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency -from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode +from tensorrt_llm._torch.moe.fused_moe.interface import MoE, MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, @@ -125,6 +125,18 @@ ) +def test_duplicate_layer_ids_preserve_all_moe_registrations() -> None: + model_config = ModelConfig(skip_create_weights_in_init=True) + moe_layers = [torch.nn.Module() for _ in range(3)] + for layer in moe_layers: + layer.layer_idx = 0 + layer.layer_idx_str = "0" + MoE._register_layer(layer, model_config) + + assert [layer.layer_idx_str for layer in moe_layers] == ["0", "0_0", "0_1"] + assert [ref() for ref in model_config.extra_attrs["moe_layers"].values()] == moe_layers + + def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> None: """MegaMoE backends resolve an EP ProcessGroup at construction time. From ff3633682fd4cc01efc5d467b7f3266ea878cd0c Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Tue, 15 Sep 2026 03:44:01 -0700 Subject: [PATCH 6/6] [None][test] update DSpark draft model reference Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- .../_torch/speculative/hw_agnostic/test_dspark_eplb_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py index c4738c0bb9a9..38ca76204105 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py @@ -162,7 +162,7 @@ def test_dspark_draft_backend_auto_resolves_on_isolated_copy(): with patch.object( modeling_dspark.ModelConfig, "resolve_moe_backend", return_value="TRTLLM" ) as resolve_backend: - draft_config = modeling_dspark.DSparkDraftModel._derive_draft_model_config( + draft_config = modeling_dspark.DSv4DSparkDraftModel._derive_draft_model_config( model_config, NUM_HIDDEN_LAYERS, NUM_STAGES, "AUTO" )