[CI][Spec Decode] Add MTP placeholder-token regression coverage (#54893)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Sep 2, 2026 at 03:29 UTC
87deddc7a4d673e8e17bfe774a886a55815c3bde
1 file changed
+104
-1
tests/v1/e2e/spec_decode/mtp/test_mtp.py
+104
-1
@@ -7,16 +7,119 @@ import pytest
7
8
from tests.utils import single_gpu_only
9
from vllm import SamplingParams
10
-from vllm.config import CompilationConfig
10
+from vllm.config import CompilationConfig, CUDAGraphMode
11
from vllm.platforms import current_platform
12
13
from ..utils import (
14
_skip_if_insufficient_gpus_for_tp,
15
assert_request_outputs_match,
16
evaluate_llm_for_gsm8k,
17
+ get_spec_decode_metric_value,
18
get_test_prompts,
19
)
20
21
+PLACEHOLDER_REGRESSION_MODEL = "Qwen/Qwen3.5-0.8B-Base"
22
+PLACEHOLDER_REGRESSION_PROMPT = "The capital of France is"
23
+PLACEHOLDER_REGRESSION_MAX_TOKENS = 16
24
+PLACEHOLDER_REGRESSION_ALLOWED_TOKEN_ID = 42
25
+PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS = 3
26
+
27
+
28
+@pytest.mark.parametrize(
29
+ "cudagraph_mode",
30
+ [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE],
31
+ ids=["piecewise", "full-and-piecewise"],
32
+)
33
+@single_gpu_only
34
+def test_mtp_rejected_placeholders_match_greedy_token_ids(
35
+ cudagraph_mode: CUDAGraphMode,
36
+ vllm_runner,
37
+):
38
+ """Rejected MTP padding must not alter the non-speculative token stream."""
39
+ common_kwargs: dict[str, Any] = {
40
+ "block_size": None,
41
+ "max_model_len": 256,
42
+ "max_num_seqs": 1,
43
+ "trust_remote_code": True,
44
+ "enable_chunked_prefill": None,
45
+ "limit_mm_per_prompt": {"image": 0, "video": 0},
46
+ "attention_backend": ("TRITON_ATTN" if current_platform.is_rocm() else "auto"),
47
+ }
48
+ sampling_params = SamplingParams(
49
+ temperature=0.0,
50
+ max_tokens=PLACEHOLDER_REGRESSION_MAX_TOKENS,
51
+ ignore_eos=True,
52
+ allowed_token_ids=[PLACEHOLDER_REGRESSION_ALLOWED_TOKEN_ID],
53
+ logprobs=1,
54
+ )
55
+
56
+ def generate_token_ids(runner) -> tuple[int, ...]:
57
+ outputs = runner.llm.generate([PLACEHOLDER_REGRESSION_PROMPT], sampling_params)
58
+ assert outputs and outputs[0].outputs, "Expected one completion"
59
+ token_ids = tuple(outputs[0].outputs[0].token_ids)
60
+ assert len(token_ids) == PLACEHOLDER_REGRESSION_MAX_TOKENS
61
+ return token_ids
62
+
63
+ def compilation_config() -> CompilationConfig:
64
+ return CompilationConfig(
65
+ cudagraph_mode=cudagraph_mode,
66
+ cudagraph_capture_sizes=[1, PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS + 1],
67
+ )
68
+
69
+ with vllm_runner(
70
+ PLACEHOLDER_REGRESSION_MODEL,
71
+ compilation_config=compilation_config(),
72
+ **common_kwargs,
73
+ ) as ref_runner:
74
+ ref_mode = (
75
+ ref_runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
76
+ )
77
+ assert ref_mode == cudagraph_mode
78
+ ref_token_ids = generate_token_ids(ref_runner)
79
+ expected_token_ids = (PLACEHOLDER_REGRESSION_ALLOWED_TOKEN_ID,) * (
80
+ PLACEHOLDER_REGRESSION_MAX_TOKENS
81
+ )
82
+ assert ref_token_ids == expected_token_ids
83
+
84
+ with vllm_runner(
85
+ PLACEHOLDER_REGRESSION_MODEL,
86
+ speculative_config={
87
+ "method": "mtp",
88
+ "num_speculative_tokens": PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS,
89
+ "max_model_len": 256,
90
+ "rejection_sample_method": "synthetic",
91
+ "synthetic_acceptance_rates": [0.0]
92
+ * PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS,
93
+ },
94
+ disable_log_stats=False,
95
+ compilation_config=compilation_config(),
96
+ **common_kwargs,
97
+ ) as spec_runner:
98
+ spec_mode = (
99
+ spec_runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
100
+ )
101
+ assert spec_mode == cudagraph_mode
102
+ spec_token_ids = generate_token_ids(spec_runner)
103
+ metrics = spec_runner.llm.get_metrics()
104
+ num_drafts = get_spec_decode_metric_value(
105
+ metrics, "vllm:spec_decode_num_drafts"
106
+ )
107
+ num_draft_tokens = get_spec_decode_metric_value(
108
+ metrics, "vllm:spec_decode_num_draft_tokens"
109
+ )
110
+ num_accepted_tokens = get_spec_decode_metric_value(
111
+ metrics, "vllm:spec_decode_num_accepted_tokens"
112
+ )
113
+
114
+ assert num_drafts > 0, "MTP drafter did not run"
115
+ assert num_draft_tokens > 0, "MTP drafter did not produce draft tokens"
116
+ assert num_accepted_tokens == 0, "Synthetic sampler accepted a draft token"
117
+ assert spec_token_ids == ref_token_ids, (
118
+ f"MTP changed greedy token ids under {cudagraph_mode.name}:\n"
119
+ f" reference={ref_token_ids}\n"
120
+ f" speculative={spec_token_ids}"
121
+ )
122
+
123
124
@pytest.mark.parametrize(
125
["model_setup", "mm_enabled", "expected_accuracy_threshold"],