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_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_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_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 88827b1b9428..5af73e0add7c 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 @@ -11,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 @@ -1192,6 +1194,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) @@ -1307,6 +1312,75 @@ 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, + 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 [] + architecture = architectures[0] if architectures else "" + resolved_moe_backend = ModelConfig.resolve_moe_backend( + requested_moe_backend, + architecture, + quant_config=(model_config.quant_config + if quant_config is None else quant_config)) + draft_config = copy.copy(model_config) + # 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. @@ -1326,7 +1400,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, @@ -1363,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) @@ -1460,6 +1553,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 +1562,26 @@ 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,15 +1599,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 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/_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..a7764039262c 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,21 @@ 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 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, description= @@ -2164,6 +2181,19 @@ 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 or not model_config_resolved + or not self.uses_replacement_heads): + return + 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: """Whether `speculative_model` contains replacement MTP heads.""" @@ -6213,6 +6243,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 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_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/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index 8aa29ceb6d85..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,8 +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): @@ -425,3 +431,156 @@ 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 + ) + + +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/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. 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..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 @@ -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.DSv4DSparkDraftModel._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..7304212eeba0 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,93 @@ 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_allowed_after_checkpoint_resolves_vanilla( + num_nextn_predict_layers: int, +) -> None: + spec_config = MTPDecodingConfig( + max_draft_len=1, + moe_backend="CUTLASS", + ) + model_config = SimpleNamespace( + architectures=["LlamaForCausalLM"], + num_nextn_predict_layers=num_nextn_predict_layers, + ) + + 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_allowed_for_internal_mtp_eagle() -> None: + 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, + ) + + 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( + ("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() -> None: + 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="replacement-head 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 e9c9925dfb37..87103f7bfa5e 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,81 @@ 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) -> None: + 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: 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: object) -> None: + with pytest.raises(ValidationError, match="moe_backend"): + DecodingBaseConfig(moe_backend=moe_backend) + + 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" + + 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) -> 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_accepts_explicit_vanilla_mtp_override(self) -> None: + spec_config = MTPDecodingConfig(max_draft_len=1, + moe_backend="CUTLASS", + use_mtp_vanilla=True) + + 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) -> None: + 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_ignores_override_without_neural_drafter(self) -> None: + spec_config = NGramDecodingConfig(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" + + @pytest.mark.cpu_only def test_rejection_sampling_allows_attention_dp(monkeypatch): """ADP (incl. ADP+LM-head-TP) supports rejection sampling.