@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 924707f1bf

[CI] Split (H200 MIG 35GB) Spec Decode Speculators + MTP into 4 named jobs (#57237)

Signed-off-by: Thang Nguyen <thangnguyenvn647@gmail.com> Co-authored-by: Thang Nguyen <thang.nguyen@inferact.ai> Co-authored-by: Kimi Code <noreply@moonshot.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Thang Nguyen committed Sep 27, 2026 at 08:01 UTC 924707f1bf94ff583d89bff7522ee12ff032c286
12 files changed +499 -258
.buildkite/test_areas/spec_decode.yaml
+98 -6
@@ -98,9 +98,11 @@ steps:
98 commands:
99 - pytest -v -s v1/e2e/spec_decode/eagle/
100
101 -- label: ":nvidia: (H200 MIG 35GB) Spec Decode Speculators + MTP"
102 - key: spec-decode-speculators-mtp
103 - timeout_in_minutes: 50
101 +# MTP correctness is one parametrized test split into per-model folders;
102 +# each job below runs whole directories.
103 +- label: ":nvidia: (H200 MIG 35GB) Spec Decode Speculators"
104 + key: spec-decode-speculators
105 + timeout_in_minutes: 30
106 device: h200_35gb
107 source_file_dependencies:
108 - vllm/v1/spec_decode/
@@ -110,13 +112,103 @@ steps:
112 - tests/v1/e2e/spec_decode/
113 commands:
114 - pytest -v -s v1/e2e/spec_decode/speculators/
113 - - pytest -v -s v1/e2e/spec_decode/mtp/
115 mirror:
116 amd:
116 - label: ":amd: (MI300) Spec Decode Speculators + MTP"
117 + label: ":amd: (MI300) Spec Decode Speculators"
118 dind: false
119 device: mi300_1
119 - timeout_in_minutes: 75
120 + timeout_in_minutes: 40
121 + depends_on:
122 + - image-build-amd
123 + source_file_dependencies:
124 + - vllm/v1/spec_decode/
125 + - vllm/v1/worker/gpu/spec_decode/
126 + - vllm/model_executor/model_loader/
127 + - vllm/v1/sample/
128 + - vllm/model_executor/layers/
129 + - vllm/transformers_utils/configs/speculators/
130 + - tests/v1/e2e/spec_decode/
131 + - vllm/platforms/rocm.py
132 +
133 +- label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP Gemma4"
134 + key: spec-decode-mtp-gemma4
135 + timeout_in_minutes: 30
136 + device: h200_35gb
137 + source_file_dependencies:
138 + - vllm/v1/spec_decode/
139 + - vllm/v1/worker/gpu/spec_decode/
140 + - vllm/v1/attention/backends/
141 + - vllm/transformers_utils/configs/speculators/
142 + - tests/v1/e2e/spec_decode/
143 + commands:
144 + - pytest -v -s v1/e2e/spec_decode/mtp/gemma4/
145 + mirror:
146 + amd:
147 + label: ":amd: (MI300) Spec Decode MTP Gemma4"
148 + dind: false
149 + device: mi300_1
150 + timeout_in_minutes: 40
151 + depends_on:
152 + - image-build-amd
153 + source_file_dependencies:
154 + - vllm/v1/spec_decode/
155 + - vllm/v1/worker/gpu/spec_decode/
156 + - vllm/model_executor/model_loader/
157 + - vllm/v1/sample/
158 + - vllm/model_executor/layers/
159 + - vllm/transformers_utils/configs/speculators/
160 + - tests/v1/e2e/spec_decode/
161 + - vllm/platforms/rocm.py
162 +
163 +- label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP DeepSeek + MiMo"
164 + key: spec-decode-mtp-deepseek-mimo
165 + timeout_in_minutes: 30
166 + device: h200_35gb
167 + source_file_dependencies:
168 + - vllm/v1/spec_decode/
169 + - vllm/v1/worker/gpu/spec_decode/
170 + - vllm/v1/attention/backends/
171 + - vllm/transformers_utils/configs/speculators/
172 + - tests/v1/e2e/spec_decode/
173 + commands:
174 + - pytest -v -s v1/e2e/spec_decode/mtp/deepseek/
175 + - pytest -v -s v1/e2e/spec_decode/mtp/mimo/
176 + mirror:
177 + amd:
178 + label: ":amd: (MI300) Spec Decode MTP DeepSeek + MiMo"
179 + dind: false
180 + device: mi300_1
181 + timeout_in_minutes: 40
182 + depends_on:
183 + - image-build-amd
184 + source_file_dependencies:
185 + - vllm/v1/spec_decode/
186 + - vllm/v1/worker/gpu/spec_decode/
187 + - vllm/model_executor/model_loader/
188 + - vllm/v1/sample/
189 + - vllm/model_executor/layers/
190 + - vllm/transformers_utils/configs/speculators/
191 + - tests/v1/e2e/spec_decode/
192 + - vllm/platforms/rocm.py
193 +
194 +- label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP Qwen3.5"
195 + key: spec-decode-mtp-qwen3-5
196 + timeout_in_minutes: 30
197 + device: h200_35gb
198 + source_file_dependencies:
199 + - vllm/v1/spec_decode/
200 + - vllm/v1/worker/gpu/spec_decode/
201 + - vllm/v1/attention/backends/
202 + - vllm/transformers_utils/configs/speculators/
203 + - tests/v1/e2e/spec_decode/
204 + commands:
205 + - pytest -v -s v1/e2e/spec_decode/mtp/qwen3_5/
206 + mirror:
207 + amd:
208 + label: ":amd: (MI300) Spec Decode MTP Qwen3.5"
209 + dind: false
210 + device: mi300_1
211 + timeout_in_minutes: 40
212 depends_on:
213 - image-build-amd
214 source_file_dependencies:
tests/v1/e2e/spec_decode/mtp/_correctness.py new
+120
@@ -0,0 +1,120 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any
5 +
6 +import pytest
7 +
8 +from vllm import SamplingParams
9 +from vllm.config import CompilationConfig
10 +from vllm.platforms import current_platform
11 +
12 +from ..utils import (
13 + _skip_if_insufficient_gpus_for_tp,
14 + assert_request_outputs_match,
15 + evaluate_llm_for_gsm8k,
16 + get_test_prompts,
17 +)
18 +
19 +
20 +def check_mtp_correctness(
21 + monkeypatch: pytest.MonkeyPatch,
22 + sampling_config: SamplingParams,
23 + model_setup: tuple[str, str, int, str | None],
24 + mm_enabled: bool,
25 + expected_accuracy_threshold: float,
26 + vllm_runner,
27 +):
28 + """Compare the outputs of a original LLM and a speculative LLM
29 + which should be the same when using MTP speculative decoding. Due to some variance
30 + in the engine, it is possible for some outputs to differ, so we expect that at least
31 + 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
32 + reference threshold for each model.
33 + """
34 + method, model_name, tp_size, draft_model = model_setup
35 + _skip_if_insufficient_gpus_for_tp(tp_size)
36 +
37 + # Generate test prompts inside the function instead of using fixture
38 + test_prompts = get_test_prompts(mm_enabled)
39 + with monkeypatch.context() as m:
40 + m.setenv("VLLM_MLA_DISABLE", "1")
41 +
42 + attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
43 +
44 + # Skip multimodal profiling for models that don't need it in this test.
45 + extra_kwargs: dict[str, Any] = {}
46 + if "Qwen3.5" in model_name:
47 + extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
48 + elif "gemma-4" in model_name:
49 + extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0}
50 +
51 + if draft_model is not None and "gemma-4" in draft_model:
52 + import transformers
53 + from packaging.version import Version
54 +
55 + if Version(transformers.__version__) < Version("5.8.0"):
56 + pytest.skip(
57 + "Gemma4 MTP assistant requires transformers>=5.8.0, "
58 + f"got {transformers.__version__}"
59 + )
60 +
61 + with vllm_runner(
62 + model_name,
63 + block_size=None,
64 + max_model_len=2048,
65 + tensor_parallel_size=tp_size,
66 + trust_remote_code=True,
67 + attention_backend=attn_backend,
68 + enable_chunked_prefill=None,
69 + compilation_config=CompilationConfig(),
70 + **extra_kwargs,
71 + ) as ref_runner:
72 + ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
73 + evaluate_llm_for_gsm8k(
74 + ref_runner.llm,
75 + expected_accuracy_threshold=expected_accuracy_threshold,
76 + )
77 +
78 + speculative_config: dict[str, Any] = {
79 + "method": method,
80 + "num_speculative_tokens": 1,
81 + "max_model_len": 2048,
82 + }
83 + if draft_model is not None:
84 + speculative_config["model"] = draft_model
85 + speculative_config["num_speculative_tokens"] = 2
86 +
87 + with vllm_runner(
88 + model_name,
89 + block_size=None,
90 + trust_remote_code=True,
91 + tensor_parallel_size=tp_size,
92 + speculative_config=speculative_config,
93 + max_model_len=2048,
94 + attention_backend=attn_backend,
95 + enable_chunked_prefill=None,
96 + compilation_config=CompilationConfig(),
97 + **extra_kwargs,
98 + ) as spec_runner:
99 + # MTP supports async scheduling by default.
100 + has_async = (
101 + spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling
102 + )
103 + assert has_async, (
104 + f"Expected async scheduling for {method}: target={model_name}, "
105 + f"draft={draft_model}; got {has_async}"
106 + )
107 + evaluate_llm_for_gsm8k(
108 + spec_runner.llm,
109 + expected_accuracy_threshold=expected_accuracy_threshold,
110 + )
111 + spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
112 +
113 + # Heuristic: expect at least 80% of the prompts to match exactly
114 + # Upon failure, inspect the outputs to check for inaccuracy.
115 + assert_request_outputs_match(
116 + ref_outputs,
117 + spec_outputs,
118 + required_matches=int(0.8 * len(ref_outputs)) + 1,
119 + context=f"{method} target={model_name}, draft={draft_model}",
120 + )
tests/v1/e2e/spec_decode/mtp/deepseek/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/mtp/deepseek/test_mtp.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from tests.utils import single_gpu_only
7 +from vllm import SamplingParams
8 +from vllm.platforms import current_platform
9 +
10 +from .._correctness import check_mtp_correctness
11 +
12 +
13 +@pytest.mark.parametrize(
14 + ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
15 + [
16 + pytest.param(
17 + ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1, None),
18 + False,
19 + 0.0,
20 + marks=pytest.mark.skipif(
21 + current_platform.is_device_capability_family(100),
22 + reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
23 + ),
24 + ), # dummy model
25 + ],
26 + ids=["deepseek"],
27 +)
28 +@single_gpu_only
29 +def test_mtp_correctness(
30 + monkeypatch: pytest.MonkeyPatch,
31 + sampling_config: SamplingParams,
32 + model_setup: tuple[str, str, int, str | None],
33 + mm_enabled: bool,
34 + expected_accuracy_threshold: float,
35 + vllm_runner,
36 +):
37 + check_mtp_correctness(
38 + monkeypatch,
39 + sampling_config,
40 + model_setup,
41 + mm_enabled,
42 + expected_accuracy_threshold,
43 + vllm_runner,
44 + )
tests/v1/e2e/spec_decode/mtp/gemma4/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/mtp/gemma4/test_mtp.py new
+39
@@ -0,0 +1,39 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from tests.utils import single_gpu_only
7 +from vllm import SamplingParams
8 +
9 +from .._correctness import check_mtp_correctness
10 +
11 +
12 +@pytest.mark.parametrize(
13 + ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
14 + [
15 + (
16 + ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"),
17 + False,
18 + 0.50,
19 + ), # gemma4 MTP with assistant model, ref: ~62%
20 + ],
21 + ids=["gemma4-e4b"],
22 +)
23 +@single_gpu_only
24 +def test_mtp_correctness(
25 + monkeypatch: pytest.MonkeyPatch,
26 + sampling_config: SamplingParams,
27 + model_setup: tuple[str, str, int, str | None],
28 + mm_enabled: bool,
29 + expected_accuracy_threshold: float,
30 + vllm_runner,
31 +):
32 + check_mtp_correctness(
33 + monkeypatch,
34 + sampling_config,
35 + model_setup,
36 + mm_enabled,
37 + expected_accuracy_threshold,
38 + vllm_runner,
39 + )
tests/v1/e2e/spec_decode/mtp/mimo/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/mtp/mimo/test_mtp.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from tests.utils import single_gpu_only
7 +from vllm import SamplingParams
8 +
9 +from .._correctness import check_mtp_correctness
10 +
11 +
12 +@pytest.mark.parametrize(
13 + ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
14 + [
15 + # Reference accuracy: 65%-70%.
16 + (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1, None), False, 0.5),
17 + ],
18 + ids=["mimo"],
19 +)
20 +@single_gpu_only
21 +def test_mtp_correctness(
22 + monkeypatch: pytest.MonkeyPatch,
23 + sampling_config: SamplingParams,
24 + model_setup: tuple[str, str, int, str | None],
25 + mm_enabled: bool,
26 + expected_accuracy_threshold: float,
27 + vllm_runner,
28 +):
29 + check_mtp_correctness(
30 + monkeypatch,
31 + sampling_config,
32 + model_setup,
33 + mm_enabled,
34 + expected_accuracy_threshold,
35 + vllm_runner,
36 + )
tests/v1/e2e/spec_decode/mtp/qwen3_5/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py new
+39
@@ -0,0 +1,39 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from tests.utils import single_gpu_only
7 +from vllm import SamplingParams
8 +
9 +from .._correctness import check_mtp_correctness
10 +
11 +
12 +@pytest.mark.parametrize(
13 + ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
14 + [
15 + (
16 + ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1, None),
17 + False,
18 + 0.20,
19 + ), # hybrid + MTP, ref: ~34%-35%
20 + ],
21 + ids=["qwen3_5-hybrid"],
22 +)
23 +@single_gpu_only
24 +def test_mtp_correctness(
25 + monkeypatch: pytest.MonkeyPatch,
26 + sampling_config: SamplingParams,
27 + model_setup: tuple[str, str, int, str | None],
28 + mm_enabled: bool,
29 + expected_accuracy_threshold: float,
30 + vllm_runner,
31 +):
32 + check_mtp_correctness(
33 + monkeypatch,
34 + sampling_config,
35 + model_setup,
36 + mm_enabled,
37 + expected_accuracy_threshold,
38 + vllm_runner,
39 + )
tests/v1/e2e/spec_decode/mtp/qwen3_5/test_rejected_placeholders.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any
5 +
6 +import pytest
7 +
8 +from tests.utils import single_gpu_only
9 +from vllm import SamplingParams
10 +from vllm.config import CompilationConfig, CUDAGraphMode
11 +from vllm.platforms import current_platform
12 +
13 +from ...utils import get_spec_decode_metric_value
14 +
15 +PLACEHOLDER_REGRESSION_MODEL = "Qwen/Qwen3.5-0.8B-Base"
16 +PLACEHOLDER_REGRESSION_PROMPT = "The capital of France is"
17 +PLACEHOLDER_REGRESSION_MAX_TOKENS = 16
18 +PLACEHOLDER_REGRESSION_ALLOWED_TOKEN_ID = 42
19 +PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS = 3
20 +
21 +
22 +@pytest.mark.parametrize(
23 + "cudagraph_mode",
24 + [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE],
25 + ids=["piecewise", "full-and-piecewise"],
26 +)
27 +@single_gpu_only
28 +def test_mtp_rejected_placeholders_match_greedy_token_ids(
29 + cudagraph_mode: CUDAGraphMode,
30 + vllm_runner,
31 +):
32 + """Rejected MTP padding must not alter the non-speculative token stream."""
33 + common_kwargs: dict[str, Any] = {
34 + "block_size": None,
35 + "max_model_len": 256,
36 + "max_num_seqs": 1,
37 + "trust_remote_code": True,
38 + "enable_chunked_prefill": None,
39 + "limit_mm_per_prompt": {"image": 0, "video": 0},
40 + "attention_backend": ("TRITON_ATTN" if current_platform.is_rocm() else "auto"),
41 + }
42 + sampling_params = SamplingParams(
43 + temperature=0.0,
44 + max_tokens=PLACEHOLDER_REGRESSION_MAX_TOKENS,
45 + ignore_eos=True,
46 + allowed_token_ids=[PLACEHOLDER_REGRESSION_ALLOWED_TOKEN_ID],
47 + logprobs=1,
48 + )
49 +
50 + def generate_token_ids(runner) -> tuple[int, ...]:
51 + outputs = runner.llm.generate([PLACEHOLDER_REGRESSION_PROMPT], sampling_params)
52 + assert outputs and outputs[0].outputs, "Expected one completion"
53 + token_ids = tuple(outputs[0].outputs[0].token_ids)
54 + assert len(token_ids) == PLACEHOLDER_REGRESSION_MAX_TOKENS
55 + return token_ids
56 +
57 + def compilation_config() -> CompilationConfig:
58 + return CompilationConfig(
59 + cudagraph_mode=cudagraph_mode,
60 + cudagraph_capture_sizes=[1, PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS + 1],
61 + )
62 +
63 + with vllm_runner(
64 + PLACEHOLDER_REGRESSION_MODEL,
65 + compilation_config=compilation_config(),
66 + **common_kwargs,
67 + ) as ref_runner:
68 + ref_mode = (
69 + ref_runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
70 + )
71 + assert ref_mode == cudagraph_mode
72 + ref_token_ids = generate_token_ids(ref_runner)
73 + expected_token_ids = (PLACEHOLDER_REGRESSION_ALLOWED_TOKEN_ID,) * (
74 + PLACEHOLDER_REGRESSION_MAX_TOKENS
75 + )
76 + assert ref_token_ids == expected_token_ids
77 +
78 + with vllm_runner(
79 + PLACEHOLDER_REGRESSION_MODEL,
80 + speculative_config={
81 + "method": "mtp",
82 + "num_speculative_tokens": PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS,
83 + "max_model_len": 256,
84 + "rejection_sample_method": "synthetic",
85 + "synthetic_acceptance_rates": [0.0]
86 + * PLACEHOLDER_REGRESSION_NUM_SPEC_TOKENS,
87 + },
88 + disable_log_stats=False,
89 + compilation_config=compilation_config(),
90 + **common_kwargs,
91 + ) as spec_runner:
92 + spec_mode = (
93 + spec_runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
94 + )
95 + assert spec_mode == cudagraph_mode
96 + spec_token_ids = generate_token_ids(spec_runner)
97 + metrics = spec_runner.llm.get_metrics()
98 + num_drafts = get_spec_decode_metric_value(
99 + metrics, "vllm:spec_decode_num_drafts"
100 + )
101 + num_draft_tokens = get_spec_decode_metric_value(
102 + metrics, "vllm:spec_decode_num_draft_tokens"
103 + )
104 + num_accepted_tokens = get_spec_decode_metric_value(
105 + metrics, "vllm:spec_decode_num_accepted_tokens"
106 + )
107 +
108 + assert num_drafts > 0, "MTP drafter did not run"
109 + assert num_draft_tokens > 0, "MTP drafter did not produce draft tokens"
110 + assert num_accepted_tokens == 0, "Synthetic sampler accepted a draft token"
111 + assert spec_token_ids == ref_token_ids, (
112 + f"MTP changed greedy token ids under {cudagraph_mode.name}:\n"
113 + f" reference={ref_token_ids}\n"
114 + f" speculative={spec_token_ids}"
115 + )
tests/v1/e2e/spec_decode/mtp/test_mtp.py deleted
-252
@@ -1,252 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Any
5 -
6 -import pytest
7 -
8 -from tests.utils import single_gpu_only
9 -from vllm import SamplingParams
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"],
126 - [
127 - # Reference accuracy: 65%-70%.
128 - (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1, None), False, 0.5),
129 - pytest.param(
130 - ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1, None),
131 - False,
132 - 0.0,
133 - marks=pytest.mark.skipif(
134 - current_platform.is_device_capability_family(100),
135 - reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
136 - ),
137 - ), # dummy model
138 - (
139 - ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1, None),
140 - False,
141 - 0.20,
142 - ), # hybrid + MTP, ref: ~34%-35%
143 - (
144 - ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"),
145 - False,
146 - 0.50,
147 - ), # gemma4 MTP with assistant model, ref: ~62%
148 - ],
149 - ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"],
150 -)
151 -@single_gpu_only
152 -def test_mtp_correctness(
153 - monkeypatch: pytest.MonkeyPatch,
154 - sampling_config: SamplingParams,
155 - model_setup: tuple[str, str, int, str | None],
156 - mm_enabled: bool,
157 - expected_accuracy_threshold: float,
158 - vllm_runner,
159 -):
160 - """Compare the outputs of a original LLM and a speculative LLM
161 - which should be the same when using MTP speculative decoding. Due to some variance
162 - in the engine, it is possible for some outputs to differ, so we expect that at least
163 - 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
164 - reference threshold for each model.
165 - """
166 - method, model_name, tp_size, draft_model = model_setup
167 - _skip_if_insufficient_gpus_for_tp(tp_size)
168 -
169 - # Generate test prompts inside the function instead of using fixture
170 - test_prompts = get_test_prompts(mm_enabled)
171 - with monkeypatch.context() as m:
172 - m.setenv("VLLM_MLA_DISABLE", "1")
173 -
174 - attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
175 -
176 - # Skip multimodal profiling for models that don't need it in this test.
177 - extra_kwargs: dict[str, Any] = {}
178 - if "Qwen3.5" in model_name:
179 - extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
180 - elif "gemma-4" in model_name:
181 - extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0}
182 -
183 - if draft_model is not None and "gemma-4" in draft_model:
184 - import transformers
185 - from packaging.version import Version
186 -
187 - if Version(transformers.__version__) < Version("5.8.0"):
188 - pytest.skip(
189 - "Gemma4 MTP assistant requires transformers>=5.8.0, "
190 - f"got {transformers.__version__}"
191 - )
192 -
193 - with vllm_runner(
194 - model_name,
195 - block_size=None,
196 - max_model_len=2048,
197 - tensor_parallel_size=tp_size,
198 - trust_remote_code=True,
199 - attention_backend=attn_backend,
200 - enable_chunked_prefill=None,
201 - compilation_config=CompilationConfig(),
202 - **extra_kwargs,
203 - ) as ref_runner:
204 - ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
205 - evaluate_llm_for_gsm8k(
206 - ref_runner.llm,
207 - expected_accuracy_threshold=expected_accuracy_threshold,
208 - )
209 -
210 - speculative_config: dict[str, Any] = {
211 - "method": method,
212 - "num_speculative_tokens": 1,
213 - "max_model_len": 2048,
214 - }
215 - if draft_model is not None:
216 - speculative_config["model"] = draft_model
217 - speculative_config["num_speculative_tokens"] = 2
218 -
219 - with vllm_runner(
220 - model_name,
221 - block_size=None,
222 - trust_remote_code=True,
223 - tensor_parallel_size=tp_size,
224 - speculative_config=speculative_config,
225 - max_model_len=2048,
226 - attention_backend=attn_backend,
227 - enable_chunked_prefill=None,
228 - compilation_config=CompilationConfig(),
229 - **extra_kwargs,
230 - ) as spec_runner:
231 - # MTP supports async scheduling by default.
232 - has_async = (
233 - spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling
234 - )
235 - assert has_async, (
236 - f"Expected async scheduling for {method}: target={model_name}, "
237 - f"draft={draft_model}; got {has_async}"
238 - )
239 - evaluate_llm_for_gsm8k(
240 - spec_runner.llm,
241 - expected_accuracy_threshold=expected_accuracy_threshold,
242 - )
243 - spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
244 -
245 - # Heuristic: expect at least 80% of the prompts to match exactly
246 - # Upon failure, inspect the outputs to check for inaccuracy.
247 - assert_request_outputs_match(
248 - ref_outputs,
249 - spec_outputs,
250 - required_matches=int(0.8 * len(ref_outputs)) + 1,
251 - context=f"{method} target={model_name}, draft={draft_model}",
252 - )