diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index 02f473b2b77d..e4890a5140ed 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 not spec_config.mtp_eagle_one_model or spec_config.use_mtp_vanilla: raise ValueError( diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 4ea0222e0c52..51d967fabf7d 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -91,6 +91,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_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 2e53924e7e00..7b322314cfa3 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 @@ -1018,6 +1019,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 @@ -1078,7 +1080,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( @@ -1238,7 +1242,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 @@ -1246,23 +1252,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. @@ -1272,6 +1285,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 @@ -1833,13 +1848,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 @@ -2132,6 +2155,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_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 900c3ef76f50..ced93d59036b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -1665,6 +1665,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 eec349016a63..0a955ae565df 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 @@ -1192,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) @@ -1442,6 +1446,29 @@ def forward(self, ) +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. @@ -1461,7 +1488,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, @@ -1601,6 +1628,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 \ @@ -1608,18 +1637,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, @@ -1637,15 +1676,14 @@ 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, 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 632737785648..226fbd248564 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -696,8 +696,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/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 8efbac572925..452eb8672479 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -902,6 +902,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 1ffb613eb660..777d0fd09e6a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1470,16 +1470,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, @@ -1833,6 +1835,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. 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( default=None, description= @@ -2033,6 +2041,36 @@ 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.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.") + @property def uses_replacement_heads(self) -> bool: """Whether `speculative_model` contains replacement MTP heads.""" @@ -5918,6 +5956,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/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 27ee0153bf3b..d10522aa4659 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -2047,6 +2047,25 @@ "kind": "value", "path": "speculative_config.medusa_choices" }, + { + "allowed_values": [ + "AUTO", + "CUTLASS", + "CUTEDSL", + "TRTLLM", + "DEEPGEMM", + "DENSEGEMM", + "VANILLA", + "TRITON", + "MARLIN", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL" + ], + "annotation": "Optional[Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']]", + "converter": "", + "kind": "categorical", + "path": "speculative_config.moe_backend" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index d6154b6566b3..fd153423ec0b 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 b9156edc34b9..69df0c1b8898 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/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index bb3236abe243..ffb7ee4da3e9 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): @@ -424,3 +427,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/modules/test_mla_registry.py b/tests/unittest/_torch/modules/test_mla_registry.py index 4e40bf419666..27d828f5e65a 100644 --- a/tests/unittest/_torch/modules/test_mla_registry.py +++ b/tests/unittest/_torch/modules/test_mla_registry.py @@ -28,6 +28,7 @@ project_sparse_attn_output, ) from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.fused_moe.interface import MoE from tensorrt_llm._torch.modules.mla import MLA from tensorrt_llm.functional import PositionEmbeddingType @@ -65,7 +66,7 @@ def _make_mla(config: ModelConfig) -> MLA: ) -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) @@ -89,6 +90,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/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 3fdb886854b4..dbe38571a301 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 os import sys import unittest @@ -1769,6 +1784,90 @@ 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", + moe_backend="CUTLASS", + ) + 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_rejected_for_internal_mtp_eagle_one_model(): + spec_config = MTPDecodingConfig( + max_draft_len=1, + 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) + + +@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 + + spec_config = MTPDecodingConfig( + max_draft_len=1, + speculative_model="/tmp/assistant", + moe_backend="CUTLASS", + ) + model_config = SimpleNamespace( + architectures=[architecture], + num_nextn_predict_layers=1, + ) + + 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_rejected_for_shared_kv_replacement_heads(): + class ReplacementHeadTarget: + pass + + spec_config = MTPDecodingConfig( + max_draft_len=1, + speculative_model="/tmp/assistant", + moe_backend="CUTLASS", + ) + model_config = SimpleNamespace( + architectures=["Gemma4ForCausalLM"], + 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, ReplacementHeadTarget) + + assert spec_config.uses_replacement_heads + assert spec_config._use_shared_kv_cache + + 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 0487b403d11e..f5d1c43ecb4b 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -51,7 +51,7 @@ MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, - PeftCacheConfig, + NGramDecodingConfig, PeftCacheConfig, PrefillCudaGraphBackend, PybindMirror, RayPlacementConfig, SkipSoftmaxAttentionConfig, @@ -133,6 +133,85 @@ 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") + + 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_deprecated_two_engine_mtp_is_normalized_to_one_engine(self): + 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): + 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.