[None][feat] support draft model MoE backend override

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(

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)

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,31 +1279,40 @@ 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
         ``quant_config_dict`` MXFP4 extension for the draft layers' routed
         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

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,

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 ---

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,25 +1553,35 @@ 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 \
                             MistralConfigLoader
                         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)

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