A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
[Spec Decode] Fix max_model_len logging in speculative config for draft model (#41571)
Signed-off-by: Lanze Liu <lanzetech@gmail.com>
Lanze Liu committed
May 5, 2026 at 14:56 UTC
79246b5ea6e0ede64a7fe18f4538d59d33dd4792
2 files changed
+30
-1
tests/v1/spec_decode/test_max_len.py
+21
@@ -6,6 +6,7 @@ import pytest
6
7
from tests.utils import get_attn_backend_list_based_on_platform
8
from vllm import LLM, SamplingParams
9
+from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
10
from vllm.platforms import current_platform
11
from vllm.sampling_params import StructuredOutputsParams
12
@@ -77,3 +78,23 @@ def test_eagle_max_len(
78
"is longer than the eagle max length"
79
)
80
assert o.outputs[0].text == "a b c d e " * 15
81
+
82
+
83
+@pytest.mark.parametrize("spec_max_model_len", [80, 150])
84
+def test_mtp_speculative_config_max_model_len(spec_max_model_len: int):
85
+ """Regression test for #41456: max_model_len in speculative config
86
+ should be respected for the draft model."""
87
+ model_config = ModelConfig(
88
+ model="XiaomiMiMo/MiMo-7B-Base",
89
+ runner="generate",
90
+ max_model_len=200,
91
+ trust_remote_code=True,
92
+ )
93
+ spec_config = SpeculativeConfig(
94
+ target_model_config=model_config,
95
+ target_parallel_config=ParallelConfig(),
96
+ method="mtp",
97
+ num_speculative_tokens=1,
98
+ max_model_len=spec_max_model_len,
99
+ )
100
+ assert spec_config.draft_model_config.max_model_len == spec_max_model_len
vllm/config/speculative.py
+9
-1
@@ -626,6 +626,7 @@ class SpeculativeConfig:
626
revision=self.revision,
627
code_revision=self.code_revision,
628
tokenizer_revision=self.target_model_config.tokenizer_revision,
629
+ max_model_len=self.max_model_len, # type: ignore[arg-type]
630
spec_target_max_model_len=self.target_model_config.max_model_len,
631
quantization=self.quantization,
632
enforce_eager=self.target_model_config.enforce_eager,
@@ -837,10 +838,17 @@ class SpeculativeConfig:
838
839
return speculative_max_model_len
840
840
- return min(
841
+ result = min(
842
draft_max_model_len,
843
target_max_model_len,
844
)
845
+ if result != draft_max_model_len:
846
+ logger.info(
847
+ "Overriding draft model max model len from %d to %d",
848
+ draft_max_model_len,
849
+ result,
850
+ )
851
+ return result
852
853
@staticmethod
854
def _verify_and_get_draft_tp(