[gg-rebased] fix: pin same-repository MTP to the target revision

tests/config/test_speculative_draft_hf_overrides.py

@@ -12,6 +12,7 @@
 """
 
 import functools
+from types import SimpleNamespace
 
 import pytest
 from transformers import PretrainedConfig
@@ -132,3 +133,43 @@ def test_composed_override_is_picklable():
 
     out = composed(_make_hf_config())
     assert out.num_hidden_layers == 1
+
+
+@pytest.mark.cpu_test
+def test_mtp_same_model_inherits_target_revisions():
+    spec = SimpleNamespace(
+        method="mtp",
+        model="org/model",
+        revision=None,
+        code_revision=None,
+        target_model_config=SimpleNamespace(
+            model="org/model",
+            revision="weights-commit",
+            code_revision="code-commit",
+        ),
+    )
+
+    SpeculativeConfig._inherit_target_revision_for_mtp(spec)
+
+    assert spec.revision == "weights-commit"
+    assert spec.code_revision == "code-commit"
+
+
+@pytest.mark.cpu_test
+def test_mtp_explicit_draft_revisions_are_preserved():
+    spec = SimpleNamespace(
+        method="mtp",
+        model="org/model",
+        revision="draft-weights",
+        code_revision="draft-code",
+        target_model_config=SimpleNamespace(
+            model="org/model",
+            revision="target-weights",
+            code_revision="target-code",
+        ),
+    )
+
+    SpeculativeConfig._inherit_target_revision_for_mtp(spec)
+
+    assert spec.revision == "draft-weights"
+    assert spec.code_revision == "draft-code"

tests/models/test_deepseek_mtp_mxfp8_loader.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
+from types import SimpleNamespace
 
 import torch
 from transformers import PretrainedConfig
@@ -10,7 +11,10 @@
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     dequant_mxfp8_to_bf16,
 )
-from vllm.model_executor.models.deepseek_mtp import _try_load_fp8_linear_as_bf16
+from vllm.model_executor.models.deepseek_mtp import (
+    _get_local_model_path,
+    _try_load_fp8_linear_as_bf16,
+)
 
 
 def _write_serialized_nextn_index(model_dir, layer: int) -> None:
@@ -81,3 +85,75 @@ def test_mtp_fallback_loader_accepts_mxfp8_weight_scale():
     assert torch.equal(param.data, expected)
     assert "model.layers.78.self_attn.fused_qkv_a_proj.weight" in loaded
     assert pending == {}
+
+
+def test_mtp_serialized_probe_uses_target_revision_for_same_model(monkeypatch):
+    calls: list[tuple[str | None, str | None]] = []
+
+    def fake_resolve(model_path, revision=None):
+        calls.append((model_path, revision))
+        return "/cache/pinned-snapshot"
+
+    monkeypatch.setattr(
+        "vllm.model_executor.models.deepseek_mtp._resolve_cached_hf_model_path",
+        fake_resolve,
+    )
+    model = "org/model"
+    config = PretrainedConfig()
+    config._name_or_path = model
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            revision=None,
+            draft_model_config=SimpleNamespace(
+                model=model,
+                model_path=None,
+                model_weights=None,
+                revision=None,
+            ),
+        ),
+        model_config=SimpleNamespace(
+            model=model,
+            model_path=None,
+            model_weights=None,
+            revision="target-commit",
+        ),
+    )
+
+    assert _get_local_model_path(config, vllm_config) == "/cache/pinned-snapshot"
+    assert calls == [(model, "target-commit")]
+
+
+def test_mtp_serialized_probe_prefers_explicit_draft_revision(monkeypatch):
+    calls: list[tuple[str | None, str | None]] = []
+
+    def fake_resolve(model_path, revision=None):
+        calls.append((model_path, revision))
+        return "/cache/draft-snapshot"
+
+    monkeypatch.setattr(
+        "vllm.model_executor.models.deepseek_mtp._resolve_cached_hf_model_path",
+        fake_resolve,
+    )
+    model = "org/model"
+    config = PretrainedConfig()
+    config._name_or_path = model
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            revision="draft-commit",
+            draft_model_config=SimpleNamespace(
+                model=model,
+                model_path=None,
+                model_weights=None,
+                revision="draft-commit",
+            ),
+        ),
+        model_config=SimpleNamespace(
+            model=model,
+            model_path=None,
+            model_weights=None,
+            revision="target-commit",
+        ),
+    )
+
+    assert _get_local_model_path(config, vllm_config) == "/cache/draft-snapshot"
+    assert calls == [(model, "draft-commit")]

vllm/config/speculative.py

@@ -680,6 +680,16 @@ def _is_custom_proposer_path(model: str | None) -> bool:
         parts = model.split(".")
         return len(parts) >= 2 and all(part.isidentifier() for part in parts)
 
+    def _inherit_target_revision_for_mtp(self) -> None:
+        """Pin an in-checkpoint MTP draft to the target model revision."""
+        target = self.target_model_config
+        if self.method != "mtp" or target is None or self.model != target.model:
+            return
+        if self.revision is None:
+            self.revision = target.revision
+        if self.code_revision is None:
+            self.code_revision = target.code_revision
+
     def __post_init__(self):
         # Note: "method" is a new parameter that helps to extend the
         # configuration of non-model-based proposers, and the "model" parameter
@@ -852,6 +862,7 @@ def __post_init__(self):
             self.prompt_lookup_min = 0
 
             if self.model is not None:
+                self._inherit_target_revision_for_mtp()
                 # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*)
                 # lack a model_type key in config.json, so AutoConfig cannot
                 # detect them. When the method is explicitly "medusa", inject

vllm/model_executor/models/deepseek_mtp.py

@@ -68,14 +68,18 @@ def _restore_full_token_layout_if_needed(
     return hidden_states, residual
 
 
-def _resolve_cached_hf_model_path(model_path: str | None) -> str | None:
+def _resolve_cached_hf_model_path(
+    model_path: str | None, revision: str | None = None
+) -> str | None:
     if not model_path or os.path.isdir(model_path):
         return model_path
     try:
         from huggingface_hub import try_to_load_from_cache
 
         cached_index = try_to_load_from_cache(
-            model_path, "model.safetensors.index.json"
+            model_path,
+            "model.safetensors.index.json",
+            revision=revision or "main",
         )
         if isinstance(cached_index, str):
             return os.path.dirname(cached_index)
@@ -87,22 +91,49 @@ def _resolve_cached_hf_model_path(model_path: str | None) -> str | None:
 def _get_local_model_path(
     config: PretrainedConfig, vllm_config: VllmConfig
 ) -> str | None:
+    speculative_config = getattr(vllm_config, "speculative_config", None)
+    draft_model_config = getattr(speculative_config, "draft_model_config", None)
+    model_config = getattr(vllm_config, "model_config", None)
+
+    def _paths(model_config) -> tuple[str, ...]:
+        if model_config is None:
+            return ()
+        return tuple(
+            path
+            for path in (
+                getattr(model_config, "model", None),
+                getattr(model_config, "model_path", None),
+                getattr(model_config, "model_weights", None),
+            )
+            if path
+        )
+
+    draft_paths = _paths(draft_model_config)
+    target_paths = _paths(model_config)
+    draft_revision = getattr(draft_model_config, "revision", None) or getattr(
+        speculative_config, "revision", None
+    )
+    target_revision = getattr(model_config, "revision", None)
+
+    def _revision_for(model_path: str | None) -> str | None:
+        if model_path in draft_paths and draft_revision is not None:
+            return draft_revision
+        if model_path in target_paths:
+            return target_revision
+        return None
+
     for attr in ("_name_or_path", "name_or_path"):
         model_path = getattr(config, attr, None)
-        resolved_model_path = _resolve_cached_hf_model_path(model_path)
+        resolved_model_path = _resolve_cached_hf_model_path(
+            model_path, _revision_for(model_path)
+        )
         if resolved_model_path:
             return resolved_model_path
 
-    speculative_config = getattr(vllm_config, "speculative_config", None)
-    draft_model_config = getattr(speculative_config, "draft_model_config", None)
-    model_config = getattr(vllm_config, "model_config", None)
-    for model_path in (
-        getattr(draft_model_config, "model", None),
-        getattr(draft_model_config, "model_path", None),
-        getattr(model_config, "model", None),
-        getattr(model_config, "model_path", None),
-    ):
-        resolved_model_path = _resolve_cached_hf_model_path(model_path)
+    for model_path in (*draft_paths, *target_paths):
+        resolved_model_path = _resolve_cached_hf_model_path(
+            model_path, _revision_for(model_path)
+        )
         if resolved_model_path:
             return resolved_model_path
     return None