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..1d4de8752d6c 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 69e69ca8b954..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._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 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..d41101db7ecd 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 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,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,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) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f1f88a59e953..f408381861c7 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" + 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): 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..3fdb886854b4 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/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) 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..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 @@ -4,24 +4,44 @@ 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 _ExternalDraftModelTarget: + build_mtp_draft_model_from_config = True + + +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, _EmbeddedOrHeadReplacementTarget + ) + + 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 +53,61 @@ 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, _EmbeddedOrHeadReplacementTarget, "embedded"), + ("/path/to/mtp", _EmbeddedOrHeadReplacementTarget, "head_replacement"), + ("/path/to/assistant", _ExternalDraftModelTarget, "external_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) + + external_draft_model = object() + replacement_mtp_heads = object() + monkeypatch.setattr( + modeling_speculative.AutoModelForCausalLM, + "from_config", + lambda draft_config: external_draft_model, + ) + monkeypatch.setattr( + modeling_speculative, + "MTPForCausalLM", + lambda *args: replacement_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 == "external_draft_model" else None + draft_model = modeling_speculative.get_draft_model( + model_config, draft_config, object(), object() + ) + + 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_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): @@ -50,10 +116,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 +132,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 +143,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 +282,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.