[None][fix] Restore Gemma4 shared-KV draft loading

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -1225,6 +1225,8 @@ def forward(
 
 @register_auto_model("Gemma4ForCausalLM")
 class Gemma4ForCausalLM(SpecDecOneEngineForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
+    build_mtp_draft_model_from_config = True
+
     def __init__(
         self,
         model_config: ModelConfig[Gemma4TextConfig],

tensorrt_llm/_torch/models/modeling_gemma4mm.py

@@ -963,6 +963,8 @@ class Gemma4ForConditionalGeneration(Gemma4MultimodalModelBase):
     - mm_token_type_ids-based bidirectional masking
     """
 
+    build_mtp_draft_model_from_config = True
+
     def __init__(self, model_config: ModelConfig[Gemma4Config]):
         if _is_mm_disagg():
             raise NotImplementedError(

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -912,13 +912,14 @@ def __init__(
             model_nextn = self.config.num_nextn_predict_layers
             ckpt_nextn = self.config.num_nextn_predict_layers
             self.num_hidden_layers = self.config.num_hidden_layers
-            has_external_mtp = (
-                model_config.spec_config.loads_mtp_from_separate_checkpoint)
-            assert ckpt_nextn > 0 or has_external_mtp, (
+            has_mtp_head_replacement = (
+                model_config.spec_config.uses_replacement_heads)
+            assert ckpt_nextn > 0 or has_mtp_head_replacement, (
                 "There are not MTP modules in the checkpoint. "
                 "Set speculative_config.speculative_model to a separate MTP "
-                "heads checkpoint, or use a target checkpoint that embeds MTP.")
-            if ckpt_nextn == 0 and has_external_mtp:
+                "head replacement checkpoint, or use a target checkpoint that "
+                "embeds MTP.")
+            if ckpt_nextn == 0 and has_mtp_head_replacement:
                 # Neither checkpoint declares a head count: fall back to a
                 # single shared head, matching MTPForCausalLM's MTP-Eagle
                 # default.
@@ -987,9 +988,9 @@ def load_weights(self,
                      weight_mapper: BaseWeightMapper,
                      allow_partial_loading: bool = False):
         from tensorrt_llm._torch.speculative.utils import (
-            filter_mtp_checkpoint_weights, loads_mtp_from_speculative_model)
+            filter_mtp_checkpoint_weights, uses_mtp_head_checkpoint)
 
-        if loads_mtp_from_speculative_model(self.model_config.spec_config):
+        if uses_mtp_head_checkpoint(self.model_config.spec_config):
             # Filter before preprocess: mapper remaps mtp.layers.* ->
             # model.layers.{N}.* and would otherwise load embedded MTP heads.
             weights = filter_mtp_checkpoint_weights(weights)

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -2424,11 +2424,11 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                 f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}"
             )
 
-    elif model_config.spec_config._use_shared_kv_cache:
+    elif model_config.spec_config.uses_external_draft_model:
         if draft_config is None:
             raise ValueError(
-                "Shared-KV speculative decoding requires an external draft "
-                "model config.")
+                "MTP speculative decoding with an external draft model requires "
+                "its model config.")
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
         return MTPForCausalLM(model_config,
@@ -2540,7 +2540,7 @@ def __init__(self,
                     model_config.quant_config.kv_cache_quant_algo
                     self.draft_config.extra_attrs = model_config.extra_attrs
 
-                elif spec_config._use_shared_kv_cache:
+                elif spec_config.uses_external_draft_model:
                     self.draft_config = ModelConfig.from_pretrained(
                         spec_config.speculative_model,
                         trust_remote_code=True,
@@ -2686,10 +2686,10 @@ def load_weights(self,
                      params_map: Optional[Dict[str, str]] = None,
                      allow_partial_loading: bool = False):
         from tensorrt_llm._torch.speculative.utils import (
-            filter_mtp_checkpoint_weights, loads_mtp_from_speculative_model)
+            filter_mtp_checkpoint_weights, uses_mtp_head_checkpoint)
 
         skip_modules = ["draft_model"]
-        if loads_mtp_from_speculative_model(self.spec_config):
+        if uses_mtp_head_checkpoint(self.spec_config):
             # The heads come from speculative_model in a second pass
             # (load_draft_weights), so exclude them here. They must be
             # *skipped* rather than tolerated via allow_partial_loading:
@@ -2711,12 +2711,11 @@ def load_draft_weights(self,
         from tensorrt_llm._torch.models.modeling_utils import \
             _load_weights_impl_v2
         from tensorrt_llm._torch.speculative.utils import (
-            loads_mtp_from_speculative_model,
             remap_preprocessed_mtp_weights_for_draft_model,
             select_mtp_checkpoint_weights,
-            skip_modules_for_separate_mtp_checkpoint)
+            skip_modules_for_separate_mtp_checkpoint, uses_mtp_head_checkpoint)
 
-        if loads_mtp_from_speculative_model(self.spec_config):
+        if uses_mtp_head_checkpoint(self.spec_config):
             # Load MTP heads into draft_model only, and verify every non-shared
             # MTP parameter has a matching tensor. The previous parent-model
             # load used allow_partial_loading=True, which silently left MTP

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -462,11 +462,16 @@ def load_config_and_apply_defaults(
         if llm_args.speculative_config is not None:
             from tensorrt_llm._torch.speculative import \
                 update_spec_config_from_model_config
+            from tensorrt_llm._torch.speculative.utils import \
+                resolve_mtp_checkpoint_source
 
-            # Model defaults reconstruct nested Pydantic configs and drop
-            # init=False runtime fields such as num_nextn_predict_layers.
+            # Model defaults reconstruct nested Pydantic configs and drop private / runtime fields,
+            # so resolve the checkpoint source again before restoring derived MTP state.
+            resolve_mtp_checkpoint_source(llm_args.speculative_config,
+                                          checkpoint_dir)
             update_spec_config_from_model_config(llm_args.speculative_config,
-                                                 config.pretrained_config)
+                                                 config.pretrained_config,
+                                                 preference_cls)
 
         # Resolve "auto" sentinel values after model defaults are applied.
         _resolve_transceiver_runtime_auto(llm_args, preference_cls,
@@ -1377,8 +1382,8 @@ def _load_and_validate_config(
             checkpoint_loader: BaseCheckpointLoader) -> ModelConfig:
         """Loads and validates the model configuration."""
         from tensorrt_llm._torch.speculative.utils import (
-            loads_mtp_from_speculative_model, resolve_mtp_checkpoint_source,
-            update_spec_config_from_model_config)
+            resolve_mtp_checkpoint_source, update_spec_config_from_model_config,
+            uses_mtp_head_checkpoint)
 
         resolve_mtp_checkpoint_source(self.spec_config, checkpoint_dir)
 
@@ -1426,7 +1431,7 @@ def _load_and_validate_config(
 
         config = checkpoint_loader.load_config(**load_config_kwargs)
 
-        if loads_mtp_from_speculative_model(self.spec_config):
+        if uses_mtp_head_checkpoint(self.spec_config):
             # `load_config_and_apply_defaults` already ran this, but against a
             # config object it then discards. The MTP heads' structure fields
             # (head count, block pattern) come from `speculative_model` and

tensorrt_llm/_torch/speculative/utils.py

@@ -181,11 +181,11 @@ def skip_modules_for_separate_mtp_checkpoint(weights: dict) -> list[str]:
     return skip
 
 
-def loads_mtp_from_speculative_model(spec_config) -> bool:
-    """True when one-model MTP should load heads from ``speculative_model``."""
+def uses_mtp_head_checkpoint(spec_config) -> bool:
+    """True when `speculative_model` contains replacement MTP heads."""
     if spec_config is None:
         return False
-    return spec_config.loads_mtp_from_separate_checkpoint
+    return spec_config.uses_replacement_heads
 
 
 def _refers_to_same_checkpoint(lhs, rhs) -> bool:
@@ -211,19 +211,21 @@ def resolve_mtp_checkpoint_source(spec_config, checkpoint_dir) -> None:
     keep that behavior instead of switching to the separate-heads load path,
     which the target checkpoint's key layout may not even satisfy.
     """
-    from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
+    from tensorrt_llm.llmapi.llm_args import (MTPDecodingConfig,
+                                              _MTPDraftCheckpointType)
     if not isinstance(spec_config, MTPDecodingConfig):
         return
     if spec_config.speculative_model is None:
         return
     if not _refers_to_same_checkpoint(spec_config.speculative_model,
                                       checkpoint_dir):
         return
-    if not spec_config._mtp_heads_in_target_checkpoint:
+    if (spec_config._mtp_draft_checkpoint_type
+            != _MTPDraftCheckpointType.TARGET):
         logger.info(
             "speculative_model points at the target checkpoint "
             f"({checkpoint_dir}); loading MTP heads from the target weights.")
-        spec_config._mtp_heads_in_target_checkpoint = True
+        spec_config._mtp_draft_checkpoint_type = _MTPDraftCheckpointType.TARGET
 
 
 def _load_speculative_model_config_dict(spec_config) -> Optional[dict]:
@@ -794,21 +796,43 @@ def get_draft_kv_cache_manager(spec_config, resource_manager):
         ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
 
 
-def update_spec_config_from_model_config(spec_config, model_config):
-    from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
+def update_spec_config_from_model_config(spec_config,
+                                         model_config,
+                                         target_model_cls=None):
+    from tensorrt_llm.llmapi.llm_args import (MTPDecodingConfig,
+                                              _MTPDraftCheckpointType)
     if not isinstance(spec_config, MTPDecodingConfig):
         return
+
     architectures = getattr(model_config, "architectures", None) or ()
     if (architectures
             and architectures[0] in _GEMMA4_SHARED_KV_TARGET_ARCHITECTURES):
         spec_config._use_shared_kv_cache = (
             spec_config.spec_dec_mode.is_mtp_eagle_one_model())
 
+    # The target implementation owns the contract for its MTP drafter. Some one-model MTP
+    # implementations construct `MTPForCausalLM` from the target config, and optionally load a
+    # head replacement checkpoint (e.g. NemotronH).
+    # Other implementations advertise an external assistant architecture, which must be
+    # constructed from the assistant's own config.
+    checkpoint_type = spec_config._mtp_draft_checkpoint_type
+    if spec_config.speculative_model is None:
+        checkpoint_type = _MTPDraftCheckpointType.TARGET
+    elif checkpoint_type != _MTPDraftCheckpointType.TARGET:
+        if target_model_cls is not None:
+            checkpoint_type = (
+                _MTPDraftCheckpointType.EXTERNAL_DRAFT_MODEL if getattr(
+                    target_model_cls, "build_mtp_draft_model_from_config",
+                    False) else _MTPDraftCheckpointType.HEAD_REPLACEMENT)
+        elif checkpoint_type == _MTPDraftCheckpointType.UNRESOLVED:
+            checkpoint_type = _MTPDraftCheckpointType.HEAD_REPLACEMENT
+    spec_config._mtp_draft_checkpoint_type = checkpoint_type
+
     # When MTP heads live in a separate checkpoint, prefer that checkpoint's
     # layer count / pattern over the target model's (which may have no MTP or
     # an older embedded MTP head that will be overridden at weight load).
     draft_nextn = None
-    if loads_mtp_from_speculative_model(spec_config):
+    if uses_mtp_head_checkpoint(spec_config):
         draft_nextn = _merge_mtp_fields_from_speculative_model(
             spec_config, model_config)
 

tensorrt_llm/llmapi/llm_args.py

@@ -1750,6 +1750,15 @@ def skips_top_p(self) -> bool:
                         AdvancedSamplingMode.NO_TOPK_NO_TOPP)
 
 
+class _MTPDraftCheckpointType(StrEnum):
+    """Internal description of where a one-model MTP drafter comes from."""
+
+    UNRESOLVED = "unresolved"
+    TARGET = "target"
+    HEAD_REPLACEMENT = "head_replacement"
+    EXTERNAL_DRAFT_MODEL = "external_draft_model"
+
+
 class DecodingBaseConfig(StrictBaseModel):
     max_draft_len: Optional[NonNegativeInt] = Field(
         default=None, description="The maximum number of draft tokens.")
@@ -1769,9 +1778,9 @@ class DecodingBaseConfig(StrictBaseModel):
         description=
         "The speculative (draft) model. Accepts either (1) a HuggingFace Hub model ID (e.g. 'yuhuili/EAGLE3-LLaMA3.1-Instruct-8B'), "
         "which will be automatically downloaded, or (2) a local filesystem path to a downloaded model directory. "
-        "For MTP, when set to a checkpoint other than the target model, loads MTP heads from it instead of any "
-        "embedded mtp.* weights in the target; pointing it at the target model keeps the embedded heads."
-    )
+        "For one-model MTP, a non-target checkpoint provides either replacement MTP heads or a complete external "
+        "draft model, depending on the target model implementation. Pointing it at the target checkpoint uses the "
+        "target's embedded mtp.* weights.")
 
     max_concurrency: Optional[PositiveInt] = Field(
         default=None,
@@ -1850,9 +1859,10 @@ class DecodingBaseConfig(StrictBaseModel):
     _allow_separate_draft_kv_cache: bool = PrivateAttr(True)
     # If set, the draft model attends directly over the target model KV cache.
     _use_shared_kv_cache: bool = PrivateAttr(False)
-    # If set, speculative_model resolves to the target checkpoint, so one-model
-    # MTP loads its heads from the target weights instead of a separate file.
-    _mtp_heads_in_target_checkpoint: bool = PrivateAttr(False)
+    # Describes whether one-model MTP is embedded in the target, supplied as an MTP head replacement
+    # checkpoint, or supplied as an external draft model.
+    _mtp_draft_checkpoint_type: _MTPDraftCheckpointType = PrivateAttr(
+        default=_MTPDraftCheckpointType.UNRESOLVED)
     # Internal: true when draft_len_schedule was auto-translated from max_concurrency.
     _translated_from_max_concurrency: bool = PrivateAttr(False)
 
@@ -1964,28 +1974,30 @@ def supports_backend(self, backend: str) -> bool:
         return True
 
     @property
-    def loads_mtp_from_separate_checkpoint(self) -> bool:
-        """Whether one-model MTP heads come from ``speculative_model``.
+    def uses_replacement_heads(self) -> bool:
+        """Whether `speculative_model` contains replacement MTP heads."""
+        if (not self.spec_dec_mode.is_mtp_one_model()
+                or self.speculative_model is None):
+            return False
+        return (self._mtp_draft_checkpoint_type ==
+                _MTPDraftCheckpointType.HEAD_REPLACEMENT)
 
-        False when ``speculative_model`` resolves to the target checkpoint:
-        the heads are then loaded from the target weights, as they were
-        before separate MTP checkpoints were supported.
-        """
+    @property
+    def uses_external_draft_model(self) -> bool:
+        """Whether `speculative_model` contains an external draft model."""
         return (self.spec_dec_mode.is_mtp_one_model()
-                and self.speculative_model is not None
-                and not self._mtp_heads_in_target_checkpoint)
+                and self._mtp_draft_checkpoint_type
+                == _MTPDraftCheckpointType.EXTERNAL_DRAFT_MODEL)
 
     @property
     def needs_separate_draft_weights(self) -> bool:
         """Whether draft weights must be loaded from ``speculative_model``.
 
-        True for Eagle3 one-model / external drafters, Gemma4 shared-KV, and
-        one-model MTP when MTP heads live in a separate checkpoint.
+        This includes external draft models and MTP head replacement checkpoints.
         """
-        if (self.spec_dec_mode.need_load_draft_weights()
-                or self._use_shared_kv_cache):
-            return True
-        return self.loads_mtp_from_separate_checkpoint
+        return (self.spec_dec_mode.need_load_draft_weights()
+                or self.uses_external_draft_model
+                or self.uses_replacement_heads)
 
     @property
     def spec_dec_mode(self):

tests/integration/defs/accuracy/references/mmmu.yaml

@@ -19,6 +19,10 @@ google/gemma-4-26B-A4B-it:
   - quant_algo: NVFP4
     kv_cache_quant_algo: FP8
     accuracy: 54.0
+  - quant_algo: NVFP4
+    kv_cache_quant_algo: FP8
+    spec_dec_algo: MTP
+    accuracy: 54.0
 google/gemma-3-12b-it:
   - accuracy: 50.44
   - quant_algo: FP8

tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py

@@ -321,6 +321,7 @@ def test_nvfp4_prequantized(self):
 class TestGemma4_26B_A4B(LlmapiAccuracyTestHarness):
     MODEL_NAME = "google/gemma-4-26B-A4B-it"
     MODEL_PATH = f"{llm_models_root()}/gemma/nvidia-Gemma-4-26B-A4B-NVFP4"
+    MTP_MODEL_PATH = f"{llm_models_root()}/gemma/gemma-4-26B-A4B-it-assistant"
     EXTRA_EVALUATOR_KWARGS = {
         "chat_template_kwargs": {"enable_thinking": False},
     }
@@ -345,6 +346,14 @@ def test_nvfp4(self):
             max_batch_size=16,
             kv_cache_config=self.kv_cache_config,
             enable_chunked_prefill=True,
+            # Shared-KV MTP overlap can expose too few FlashInfer pages and cause an illegal access
+            # in `AppendPagedKVCache`. Re-enable this after the overlap-MTP KV accounting fix lands.
+            disable_overlap_scheduler=True,
+            speculative_config=MTPDecodingConfig(
+                max_draft_len=3,
+                mtp_eagle_one_model=True,
+                speculative_model=self.MTP_MODEL_PATH,
+            ),
         ) as llm:
             assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4
             task = MMMU(self.MODEL_NAME)

tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py

@@ -18,6 +18,7 @@
     get_num_extra_kv_tokens,
     get_num_spec_layers,
     update_spec_config_from_model_config,
+    uses_mtp_head_checkpoint,
 )
 from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig
 
@@ -1741,27 +1742,28 @@ def test_prepare_drafter_inputs(
         torch.testing.assert_close(draft_inputs["hidden_states"], ref_previous_hidden_states)
 
 
-@pytest.mark.parametrize(
-    ("architecture", "expected"),
-    [
-        ("Gemma4ForCausalLM", True),
-        ("LlamaForCausalLM", False),
-    ],
-)
-def test_mtp_shared_kv_config(architecture, expected):
+@pytest.mark.parametrize("uses_external_draft_model", [True, False])
+def test_mtp_checkpoint_type_config(uses_external_draft_model):
+    class TargetModel:
+        if uses_external_draft_model:
+            build_mtp_draft_model_from_config = True
+
     spec_config = MTPDecodingConfig(
         max_draft_len=3,
         speculative_model="/tmp/assistant",
     )
     model_config = SimpleNamespace(
-        architectures=[architecture],
+        architectures=["Gemma4ForCausalLM" if uses_external_draft_model else "TargetModel"],
         num_nextn_predict_layers=1,
     )
 
-    update_spec_config_from_model_config(spec_config, model_config)
+    update_spec_config_from_model_config(spec_config, model_config, TargetModel)
 
-    assert spec_config._use_shared_kv_cache is expected
-    if expected:
+    assert spec_config._use_shared_kv_cache is uses_external_draft_model
+    assert spec_config.uses_external_draft_model is uses_external_draft_model
+    assert uses_mtp_head_checkpoint(spec_config) is not uses_external_draft_model
+    assert spec_config.needs_separate_draft_weights
+    if uses_external_draft_model:
         assert get_num_spec_layers(spec_config) == 0
         assert get_num_extra_kv_tokens(spec_config) == 0
         assert not should_use_separate_draft_kv_cache(spec_config)