From d9d410e64ad12225381b6e92f38b2a96464a4eab Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Mon, 17 Aug 2026 10:55:00 -0700
Subject: [PATCH 1/3] [None][fix] Restore Gemma4 shared-KV draft loading

* Why?

Standalone MTP checkpoint support misclassified Gemma4's complete shared-KV
assistant as an MTP-head-only checkpoint. Startup then rejected the assistant
because it contains model weights rather than `mtp.*` tensors.

* What?

Keep Gemma4's required second draft-weight loading pass while routing it
through the full draft-model loader. Exercise the matching 26B-A4B assistant
in the multimodal MMMU accuracy test and register its MTP baseline.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 tensorrt_llm/llmapi/llm_args.py                          | 7 ++++---
 tests/integration/defs/accuracy/references/mmmu.yaml     | 4 ++++
 .../defs/accuracy/test_llm_api_pytorch_multimodal.py     | 9 +++++++++
 .../unittest/_torch/speculative/hw_agnostic/test_mtp.py  | 4 ++++
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index f1f88a59e953..dc3d70bfc3ac 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -1967,12 +1967,13 @@ def supports_backend(self, backend: str) -> bool:
     def loads_mtp_from_separate_checkpoint(self) -> bool:
         """Whether one-model MTP heads come from ``speculative_model``.
 
-        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.
+        False for shared-KV models, where ``speculative_model`` is a complete
+        draft model, and when it resolves to the target checkpoint, where the
+        heads are loaded from the target weights.
         """
         return (self.spec_dec_mode.is_mtp_one_model()
                 and self.speculative_model is not None
+                and not self._use_shared_kv_cache
                 and not self._mtp_heads_in_target_checkpoint)
 
     @property
diff --git a/tests/integration/defs/accuracy/references/mmmu.yaml b/tests/integration/defs/accuracy/references/mmmu.yaml
index 509272628d86..2643198b7660 100644
--- a/tests/integration/defs/accuracy/references/mmmu.yaml
+++ b/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
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
index 8890f95a995e..9e6c57d0bb9a 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
+++ b/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)
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
index 649da5a18884..af9e79f128e5 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
@@ -17,6 +17,7 @@
 from tensorrt_llm._torch.speculative.utils import (
     get_num_extra_kv_tokens,
     get_num_spec_layers,
+    loads_mtp_from_speculative_model,
     update_spec_config_from_model_config,
 )
 from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig
@@ -1761,6 +1762,9 @@ def test_mtp_shared_kv_config(architecture, expected):
     update_spec_config_from_model_config(spec_config, model_config)
 
     assert spec_config._use_shared_kv_cache is expected
+    assert spec_config.loads_mtp_from_separate_checkpoint is not expected
+    assert loads_mtp_from_speculative_model(spec_config) is not expected
+    assert spec_config.needs_separate_draft_weights
     if expected:
         assert get_num_spec_layers(spec_config) == 0
         assert get_num_extra_kv_tokens(spec_config) == 0

From 343d97c25d8d425139ec24a2270be127a9100d69 Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Mon, 17 Aug 2026 14:55:44 -0700
Subject: [PATCH 2/3] [None][fix] Decouple MTP loading from KV sharing

* Why?

Shared KV cache usage does not describe the contents of an MTP
checkpoint. Gemma4 supplies a complete assistant model, while Nemotron
MTPv2 supplies heads only. Using cache sharing as a weight-loading signal
obscures these distinct contracts and makes future changes fragile.

* What?

Classify one-model MTP checkpoints as target-embedded, heads-only, or
full draft models. Let target implementations declare when their
assistant must be built from its own config, while keeping shared KV
cache selection as an independent runtime property.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4.py |  2 +
 .../_torch/models/modeling_gemma4mm.py        |  2 +
 .../_torch/models/modeling_nemotron_h.py      |  7 +-
 .../_torch/models/modeling_speculative.py     | 15 ++--
 .../_torch/pyexecutor/model_loader.py         | 17 ++--
 tensorrt_llm/_torch/speculative/utils.py      | 41 +++++++--
 tensorrt_llm/llmapi/llm_args.py               | 47 ++++++----
 .../speculative/hw_agnostic/test_mtp.py       | 28 +++---
 .../test_mtp_separate_checkpoint.py           | 90 ++++++++++++++++---
 9 files changed, 176 insertions(+), 73 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index a9aa54d6029c..2de1e6ab3130 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/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],
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 0300fc2216e4..c31b6dfa5f43 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/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(
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index c6cd70a64ab9..0a6634fa4421 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -912,8 +912,7 @@ 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)
+            has_external_mtp = model_config.spec_config.uses_mtp_head_checkpoint
             assert ckpt_nextn > 0 or has_external_mtp, (
                 "There are not MTP modules in the checkpoint. "
                 "Set speculative_config.speculative_model to a separate MTP "
@@ -987,9 +986,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)
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index 69e69ca8b954..a5943f72c5dc 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -2424,10 +2424,10 @@ 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_full_draft_model_checkpoint:
         if draft_config is None:
             raise ValueError(
-                "Shared-KV speculative decoding requires an external draft "
+                "Full-model MTP speculative decoding requires an external draft "
                 "model config.")
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
@@ -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_full_draft_model_checkpoint:
                     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
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index ff7e105c808c..845a8cfff1fc 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/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
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index d0859c5f23d2..532a10e1dbfd 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/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 external MTP heads only."""
     if spec_config is None:
         return False
-    return spec_config.loads_mtp_from_separate_checkpoint
+    return spec_config.uses_mtp_head_checkpoint
 
 
 def _refers_to_same_checkpoint(lhs, rhs) -> bool:
@@ -211,7 +211,8 @@ 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:
@@ -219,11 +220,12 @@ def resolve_mtp_checkpoint_source(spec_config, checkpoint_dir) -> None:
     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,42 @@ 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
+    # heads-only checkpoint (e.g. NemotronH).
+    # Other implementations advertise a complete 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.DRAFT_MODEL if getattr(
+                target_model_cls, "build_mtp_draft_model_from_config", False)
+                               else _MTPDraftCheckpointType.HEADS)
+        elif checkpoint_type == _MTPDraftCheckpointType.UNRESOLVED:
+            checkpoint_type = _MTPDraftCheckpointType.HEADS
+    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)
 
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index dc3d70bfc3ac..e1e46725b0c3 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/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"
+    HEADS = "heads"
+    DRAFT_MODEL = "draft_model"
+
+
 class DecodingBaseConfig(StrictBaseModel):
     max_draft_len: Optional[NonNegativeInt] = Field(
         default=None, description="The maximum number of draft tokens.")
@@ -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 a heads-only
+    # checkpoint, or supplied as a complete 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,29 +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_mtp_head_checkpoint(self) -> bool:
+        """Whether `speculative_model` contains only external 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.HEADS)
 
-        False for shared-KV models, where ``speculative_model`` is a complete
-        draft model, and when it resolves to the target checkpoint, where the
-        heads are loaded from the target weights.
-        """
+    @property
+    def uses_full_draft_model_checkpoint(self) -> bool:
+        """Whether `speculative_model` contains a complete draft model."""
         return (self.spec_dec_mode.is_mtp_one_model()
-                and self.speculative_model is not None
-                and not self._use_shared_kv_cache
-                and not self._mtp_heads_in_target_checkpoint)
+                and self._mtp_draft_checkpoint_type
+                == _MTPDraftCheckpointType.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 complete draft models and heads-only MTP 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_full_draft_model_checkpoint
+                or self.uses_mtp_head_checkpoint)
 
     @property
     def spec_dec_mode(self):
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
index af9e79f128e5..3443ef267a78 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
@@ -17,8 +17,8 @@
 from tensorrt_llm._torch.speculative.utils import (
     get_num_extra_kv_tokens,
     get_num_spec_layers,
-    loads_mtp_from_speculative_model,
     update_spec_config_from_model_config,
+    uses_mtp_head_checkpoint,
 )
 from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig
 
@@ -1742,30 +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_full_draft_model", [True, False])
+def test_mtp_checkpoint_type_config(uses_full_draft_model):
+    class TargetModel:
+        if uses_full_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_full_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
-    assert spec_config.loads_mtp_from_separate_checkpoint is not expected
-    assert loads_mtp_from_speculative_model(spec_config) is not expected
+    assert spec_config._use_shared_kv_cache is uses_full_draft_model
+    assert spec_config.uses_full_draft_model_checkpoint is uses_full_draft_model
+    assert uses_mtp_head_checkpoint(spec_config) is not uses_full_draft_model
     assert spec_config.needs_separate_draft_weights
-    if expected:
+    if uses_full_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)
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
index a5be718ad5af..105cd888451a 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
@@ -4,24 +4,42 @@
 import json
 from types import SimpleNamespace
 
+import pytest
 import torch
+from transformers import PretrainedConfig
 
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models import modeling_speculative
 from tensorrt_llm._torch.models.checkpoints.hf.nemotron_h_weight_mapper import (
     NemotronHHfWeightMapper,
 )
 from tensorrt_llm._torch.speculative.utils import (
     filter_mtp_checkpoint_weights,
-    loads_mtp_from_speculative_model,
     resolve_mtp_checkpoint_source,
     select_mtp_checkpoint_weights,
     skip_modules_for_separate_mtp_checkpoint,
     update_spec_config_from_model_config,
+    uses_mtp_head_checkpoint,
 )
 from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig, MTPDecodingConfig
 
 
+class _FullDraftTarget:
+    build_mtp_draft_model_from_config = True
+
+
+class _EmbeddedOrHeadOnlyTarget:
+    pass
+
+
+def _resolve_as_head_checkpoint(spec_config):
+    pretrained_config = PretrainedConfig(architectures=["TargetModel"], num_nextn_predict_layers=1)
+    update_spec_config_from_model_config(spec_config, pretrained_config, _EmbeddedOrHeadOnlyTarget)
+
+
 def test_needs_separate_draft_weights_for_mtp_with_speculative_model():
     cfg = MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp")
+    _resolve_as_head_checkpoint(cfg)
     assert cfg.needs_separate_draft_weights is True
 
     cfg_no_draft = MTPDecodingConfig(max_draft_len=3)
@@ -33,15 +51,59 @@ def test_needs_separate_draft_weights_still_true_for_eagle3():
     assert cfg.needs_separate_draft_weights is True
 
 
-def test_loads_mtp_from_speculative_model_helper():
-    assert (
-        loads_mtp_from_speculative_model(
-            MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp")
-        )
-        is True
+def test_uses_mtp_head_checkpoint_helper():
+    spec_config = MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp")
+    _resolve_as_head_checkpoint(spec_config)
+    assert uses_mtp_head_checkpoint(spec_config) is True
+    assert uses_mtp_head_checkpoint(MTPDecodingConfig(max_draft_len=3)) is False
+    assert uses_mtp_head_checkpoint(None) is False
+
+
+@pytest.mark.parametrize(
+    ("speculative_model", "target_model_cls", "expected_checkpoint_type"),
+    [
+        (None, _EmbeddedOrHeadOnlyTarget, "embedded"),
+        ("/path/to/mtp", _EmbeddedOrHeadOnlyTarget, "heads"),
+        ("/path/to/assistant", _FullDraftTarget, "draft_model"),
+    ],
+)
+def test_mtp_checkpoint_type_selects_draft_model_constructor(
+    monkeypatch, speculative_model, target_model_cls, expected_checkpoint_type
+):
+    spec_config = MTPDecodingConfig(max_draft_len=3, speculative_model=speculative_model)
+    pretrained_config = PretrainedConfig(architectures=["TargetModel"], num_hidden_layers=52)
+    model_config = ModelConfig(
+        spec_config=spec_config,
+        pretrained_config=pretrained_config,
+    )
+    update_spec_config_from_model_config(spec_config, pretrained_config, target_model_cls)
+
+    full_draft_model = object()
+    mtp_heads = object()
+    monkeypatch.setattr(
+        modeling_speculative.AutoModelForCausalLM,
+        "from_config",
+        lambda draft_config: full_draft_model,
+    )
+    monkeypatch.setattr(
+        modeling_speculative,
+        "MTPForCausalLM",
+        lambda *args: mtp_heads,
     )
-    assert loads_mtp_from_speculative_model(MTPDecodingConfig(max_draft_len=3)) is False
-    assert loads_mtp_from_speculative_model(None) is False
+
+    draft_config = object() if expected_checkpoint_type == "draft_model" else None
+    draft_model = modeling_speculative.get_draft_model(
+        model_config, draft_config, object(), object()
+    )
+
+    if expected_checkpoint_type == "draft_model":
+        assert spec_config.uses_full_draft_model_checkpoint
+        assert not spec_config.uses_mtp_head_checkpoint
+        assert draft_model is full_draft_model
+    else:
+        assert not spec_config.uses_full_draft_model_checkpoint
+        assert spec_config.uses_mtp_head_checkpoint is (expected_checkpoint_type == "heads")
+        assert draft_model is mtp_heads
 
 
 def test_speculative_model_equal_to_target_keeps_embedded_mtp(tmp_path):
@@ -50,10 +112,9 @@ def test_speculative_model_equal_to_target_keeps_embedded_mtp(tmp_path):
     target_dir.mkdir()
 
     cfg = MTPDecodingConfig(max_draft_len=3, speculative_model=str(target_dir))
-    assert loads_mtp_from_speculative_model(cfg) is True
 
     resolve_mtp_checkpoint_source(cfg, str(target_dir))
-    assert loads_mtp_from_speculative_model(cfg) is False
+    assert uses_mtp_head_checkpoint(cfg) is False
     assert cfg.needs_separate_draft_weights is False
     # The user-provided value is left untouched.
     assert cfg.speculative_model == str(target_dir)
@@ -67,7 +128,7 @@ def test_speculative_model_equal_to_target_matches_equivalent_paths(tmp_path):
 
     cfg = MTPDecodingConfig(max_draft_len=3, speculative_model=str(link_dir))
     resolve_mtp_checkpoint_source(cfg, str(target_dir) + "/")
-    assert loads_mtp_from_speculative_model(cfg) is False
+    assert uses_mtp_head_checkpoint(cfg) is False
 
 
 def test_separate_mtp_checkpoint_survives_resolution(tmp_path):
@@ -78,7 +139,8 @@ def test_separate_mtp_checkpoint_survives_resolution(tmp_path):
 
     cfg = MTPDecodingConfig(max_draft_len=3, speculative_model=str(mtp_dir))
     resolve_mtp_checkpoint_source(cfg, str(target_dir))
-    assert loads_mtp_from_speculative_model(cfg) is True
+    _resolve_as_head_checkpoint(cfg)
+    assert uses_mtp_head_checkpoint(cfg) is True
     assert cfg.needs_separate_draft_weights is True
 
 
@@ -216,6 +278,8 @@ def _make_one_engine_stub(spec_config, num_hidden_layers: int = 52):
     """
     from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM
 
+    _resolve_as_head_checkpoint(spec_config)
+
     class _OneEngineStub(SpecDecOneEngineForCausalLM):
         # The real ``config`` is a read-only property over
         # ``model_config.pretrained_config``, which this stub never builds.

From 6f7f93f35dcdaf0e2fb6e49c59226a4f4fe83c35 Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Tue, 18 Aug 2026 12:54:12 -0700
Subject: [PATCH 3/3] Rename fields for clarity

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 .../_torch/models/modeling_nemotron_h.py      | 10 +++--
 .../_torch/models/modeling_speculative.py     |  8 ++--
 tensorrt_llm/_torch/speculative/utils.py      | 17 ++++----
 tensorrt_llm/llmapi/llm_args.py               | 34 ++++++++--------
 .../speculative/hw_agnostic/test_mtp.py       | 16 ++++----
 .../test_mtp_separate_checkpoint.py           | 40 ++++++++++---------
 6 files changed, 66 insertions(+), 59 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index 0a6634fa4421..1d4de8752d6c 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -912,12 +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.uses_mtp_head_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.
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index a5943f72c5dc..0715c39580a3 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/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.uses_full_draft_model_checkpoint:
+    elif model_config.spec_config.uses_external_draft_model:
         if draft_config is None:
             raise ValueError(
-                "Full-model MTP 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.uses_full_draft_model_checkpoint:
+                elif spec_config.uses_external_draft_model:
                     self.draft_config = ModelConfig.from_pretrained(
                         spec_config.speculative_model,
                         trust_remote_code=True,
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 532a10e1dbfd..d41101db7ecd 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -182,10 +182,10 @@ def skip_modules_for_separate_mtp_checkpoint(weights: dict) -> list[str]:
 
 
 def uses_mtp_head_checkpoint(spec_config) -> bool:
-    """True when `speculative_model` contains external MTP heads only."""
+    """True when `speculative_model` contains replacement MTP heads."""
     if spec_config is None:
         return False
-    return spec_config.uses_mtp_head_checkpoint
+    return spec_config.uses_replacement_heads
 
 
 def _refers_to_same_checkpoint(lhs, rhs) -> bool:
@@ -812,19 +812,20 @@ def update_spec_config_from_model_config(spec_config,
 
     # 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
-    # heads-only checkpoint (e.g. NemotronH).
-    # Other implementations advertise a complete assistant architecture, which must be
+    # 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.DRAFT_MODEL if getattr(
-                target_model_cls, "build_mtp_draft_model_from_config", False)
-                               else _MTPDraftCheckpointType.HEADS)
+            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.HEADS
+            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
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index e1e46725b0c3..f408381861c7 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -1755,8 +1755,8 @@ class _MTPDraftCheckpointType(StrEnum):
 
     UNRESOLVED = "unresolved"
     TARGET = "target"
-    HEADS = "heads"
-    DRAFT_MODEL = "draft_model"
+    HEAD_REPLACEMENT = "head_replacement"
+    EXTERNAL_DRAFT_MODEL = "external_draft_model"
 
 
 class DecodingBaseConfig(StrictBaseModel):
@@ -1778,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,
@@ -1859,8 +1859,8 @@ 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)
-    # Describes whether one-model MTP is embedded in the target, supplied as a heads-only
-    # checkpoint, or supplied as a complete draft model.
+    # 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.
@@ -1974,30 +1974,30 @@ def supports_backend(self, backend: str) -> bool:
         return True
 
     @property
-    def uses_mtp_head_checkpoint(self) -> bool:
-        """Whether `speculative_model` contains only external MTP heads."""
+    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.HEADS)
+        return (self._mtp_draft_checkpoint_type ==
+                _MTPDraftCheckpointType.HEAD_REPLACEMENT)
 
     @property
-    def uses_full_draft_model_checkpoint(self) -> bool:
-        """Whether `speculative_model` contains a complete draft model."""
+    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._mtp_draft_checkpoint_type
-                == _MTPDraftCheckpointType.DRAFT_MODEL)
+                == _MTPDraftCheckpointType.EXTERNAL_DRAFT_MODEL)
 
     @property
     def needs_separate_draft_weights(self) -> bool:
         """Whether draft weights must be loaded from ``speculative_model``.
 
-        This includes complete draft models and heads-only MTP checkpoints.
+        This includes external draft models and MTP head replacement checkpoints.
         """
         return (self.spec_dec_mode.need_load_draft_weights()
-                or self.uses_full_draft_model_checkpoint
-                or self.uses_mtp_head_checkpoint)
+                or self.uses_external_draft_model
+                or self.uses_replacement_heads)
 
     @property
     def spec_dec_mode(self):
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
index 3443ef267a78..3fdb886854b4 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
@@ -1742,10 +1742,10 @@ def test_prepare_drafter_inputs(
         torch.testing.assert_close(draft_inputs["hidden_states"], ref_previous_hidden_states)
 
 
-@pytest.mark.parametrize("uses_full_draft_model", [True, False])
-def test_mtp_checkpoint_type_config(uses_full_draft_model):
+@pytest.mark.parametrize("uses_external_draft_model", [True, False])
+def test_mtp_checkpoint_type_config(uses_external_draft_model):
     class TargetModel:
-        if uses_full_draft_model:
+        if uses_external_draft_model:
             build_mtp_draft_model_from_config = True
 
     spec_config = MTPDecodingConfig(
@@ -1753,17 +1753,17 @@ class TargetModel:
         speculative_model="/tmp/assistant",
     )
     model_config = SimpleNamespace(
-        architectures=["Gemma4ForCausalLM" if uses_full_draft_model else "TargetModel"],
+        architectures=["Gemma4ForCausalLM" if uses_external_draft_model else "TargetModel"],
         num_nextn_predict_layers=1,
     )
 
     update_spec_config_from_model_config(spec_config, model_config, TargetModel)
 
-    assert spec_config._use_shared_kv_cache is uses_full_draft_model
-    assert spec_config.uses_full_draft_model_checkpoint is uses_full_draft_model
-    assert uses_mtp_head_checkpoint(spec_config) is not uses_full_draft_model
+    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_full_draft_model:
+    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)
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
index 105cd888451a..972b83f542f4 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py
@@ -24,17 +24,19 @@
 from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig, MTPDecodingConfig
 
 
-class _FullDraftTarget:
+class _ExternalDraftModelTarget:
     build_mtp_draft_model_from_config = True
 
 
-class _EmbeddedOrHeadOnlyTarget:
+class _EmbeddedOrHeadReplacementTarget:
     pass
 
 
 def _resolve_as_head_checkpoint(spec_config):
     pretrained_config = PretrainedConfig(architectures=["TargetModel"], num_nextn_predict_layers=1)
-    update_spec_config_from_model_config(spec_config, pretrained_config, _EmbeddedOrHeadOnlyTarget)
+    update_spec_config_from_model_config(
+        spec_config, pretrained_config, _EmbeddedOrHeadReplacementTarget
+    )
 
 
 def test_needs_separate_draft_weights_for_mtp_with_speculative_model():
@@ -62,9 +64,9 @@ def test_uses_mtp_head_checkpoint_helper():
 @pytest.mark.parametrize(
     ("speculative_model", "target_model_cls", "expected_checkpoint_type"),
     [
-        (None, _EmbeddedOrHeadOnlyTarget, "embedded"),
-        ("/path/to/mtp", _EmbeddedOrHeadOnlyTarget, "heads"),
-        ("/path/to/assistant", _FullDraftTarget, "draft_model"),
+        (None, _EmbeddedOrHeadReplacementTarget, "embedded"),
+        ("/path/to/mtp", _EmbeddedOrHeadReplacementTarget, "head_replacement"),
+        ("/path/to/assistant", _ExternalDraftModelTarget, "external_draft_model"),
     ],
 )
 def test_mtp_checkpoint_type_selects_draft_model_constructor(
@@ -78,32 +80,34 @@ def test_mtp_checkpoint_type_selects_draft_model_constructor(
     )
     update_spec_config_from_model_config(spec_config, pretrained_config, target_model_cls)
 
-    full_draft_model = object()
-    mtp_heads = object()
+    external_draft_model = object()
+    replacement_mtp_heads = object()
     monkeypatch.setattr(
         modeling_speculative.AutoModelForCausalLM,
         "from_config",
-        lambda draft_config: full_draft_model,
+        lambda draft_config: external_draft_model,
     )
     monkeypatch.setattr(
         modeling_speculative,
         "MTPForCausalLM",
-        lambda *args: mtp_heads,
+        lambda *args: replacement_mtp_heads,
     )
 
-    draft_config = object() if expected_checkpoint_type == "draft_model" else None
+    draft_config = object() if expected_checkpoint_type == "external_draft_model" else None
     draft_model = modeling_speculative.get_draft_model(
         model_config, draft_config, object(), object()
     )
 
-    if expected_checkpoint_type == "draft_model":
-        assert spec_config.uses_full_draft_model_checkpoint
-        assert not spec_config.uses_mtp_head_checkpoint
-        assert draft_model is full_draft_model
+    if expected_checkpoint_type == "external_draft_model":
+        assert spec_config.uses_external_draft_model
+        assert not spec_config.uses_replacement_heads
+        assert draft_model is external_draft_model
     else:
-        assert not spec_config.uses_full_draft_model_checkpoint
-        assert spec_config.uses_mtp_head_checkpoint is (expected_checkpoint_type == "heads")
-        assert draft_model is mtp_heads
+        assert not spec_config.uses_external_draft_model
+        assert spec_config.uses_replacement_heads is (
+            expected_checkpoint_type == "head_replacement"
+        )
+        assert draft_model is replacement_mtp_heads
 
 
 def test_speculative_model_equal_to_target_keeps_embedded_mtp(tmp_path):