[Bugfix] Fix MTP draft model using local cache path instead of S3 URL with runai_streamer (#42079)
Signed-off-by: SoluMilken <ypiheyn.imm02g@g2.nctu.edu.tw>
SoluMilken committed
Aug 22, 2026 at 00:07 UTC
27ec8ac626345498fdac0527a4fcd1451c24bebc
2 files changed
+43
-1
tests/test_config.py
+35
@@ -2065,6 +2065,41 @@ def test_draft_sample_method_gumbel_is_rejected():
2065
)
2066
2067
2068
+@patch("vllm.config.speculative.ModelConfig")
2069
+def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls):
2070
+ """Regression test: MTP + runai_streamer should use model_weights (original
2071
+ S3 URL) for the draft model, not model (local cache dir set by
2072
+ pull_runai_model_from_obj_storage)."""
2073
+ from unittest.mock import MagicMock
2074
+
2075
+ s3_url = "s3://my-bucket/Qwen3-35B-A3B-FP8"
2076
+ local_cache = "/root/.cache/vllm/assets/model_streamer/abcd1234"
2077
+
2078
+ mock_draft = MagicMock()
2079
+ mock_draft.model = local_cache
2080
+ mock_draft.hf_config.model_type = "deepseek_mtp"
2081
+ mock_draft.hf_config.n_predict = None
2082
+ mock_draft.max_model_len = 4096
2083
+ mock_model_config_cls.return_value = mock_draft
2084
+
2085
+ target_config = MagicMock()
2086
+ target_config.model = local_cache
2087
+ target_config.model_weights = s3_url
2088
+ target_config.hf_text_config.model_type = "deepseek_v3"
2089
+ target_config.quantization = None
2090
+ target_config.max_model_len = 4096
2091
+
2092
+ SpeculativeConfig(
2093
+ method="mtp",
2094
+ num_speculative_tokens=1,
2095
+ target_model_config=target_config,
2096
+ target_parallel_config=ParallelConfig(),
2097
+ )
2098
+
2099
+ actual_model = mock_model_config_cls.call_args.kwargs["model"]
2100
+ assert actual_model == s3_url
2101
+
2102
+
2103
def test_ir_op_priority_default():
2104
"""Test that IR op priority defaults are set correctly."""
2105
from vllm.config.kernel import IrOpPriorityConfig
vllm/config/speculative.py
+8
-1
@@ -769,7 +769,14 @@ class SpeculativeConfig:
769
if self.target_model_config is None:
770
raise ValueError("target_model_config must be present for mtp")
771
# use the draft model from the same model:
772
- self.model = self.target_model_config.model
772
+ # Use model_weights if set (e.g. runai_streamer replaces
773
+ # model with a local cache dir but keeps the original path
774
+ # in model_weights), so the draft model can load weights
775
+ # from the same source as the target model.
776
+ self.model = (
777
+ self.target_model_config.model_weights
778
+ or self.target_model_config.model
779
+ )
780
# Align the quantization of draft model for cases such as
781
# --quantization fp8 with a bf16 checkpoint.
782
if not self.quantization: