[Test] Add DeepSeek MTP parallel-load tests (#41653)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Stefano Castagnetta committed
Jul 21, 2026 at 16:38 UTC
5812e1a66bc5909508c5e08674c43896f5d3401b
2 files changed
+265
.buildkite/test_areas/spec_decode.yaml
+16
@@ -170,3 +170,19 @@ steps:
170
- tests/v1/e2e/spec_decode/
171
commands:
172
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
173
+
174
+- label: Spec Decode DeepSeek MTP Parallel Load (B200)
175
+ key: spec-decode-deepseek-mtp-parallel-load-b200
176
+ timeout_in_minutes: 30
177
+ device: b200-k8s
178
+ optional: true
179
+ num_devices: 2
180
+ source_file_dependencies:
181
+ - vllm/v1/spec_decode/llm_base_proposer.py
182
+ - vllm/v1/spec_decode/eagle.py
183
+ - vllm/v1/worker/gpu/spec_decode/eagle/
184
+ - vllm/model_executor/models/deepseek_mtp.py
185
+ - vllm/model_executor/models/deepseek_v2.py
186
+ - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
187
+ commands:
188
+ - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
new
+249
@@ -0,0 +1,249 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import asyncio
5
+from contextlib import AsyncExitStack
6
+from dataclasses import dataclass, replace
7
+from typing import Any
8
+
9
+import pytest
10
+import torch
11
+
12
+from tests.utils import multi_gpu_marks
13
+from vllm import LLM, SamplingParams
14
+from vllm.distributed import cleanup_dist_env_and_memory
15
+from vllm.engine.arg_utils import AsyncEngineArgs
16
+from vllm.sampling_params import RequestOutputKind
17
+from vllm.v1.engine.async_llm import AsyncLLM
18
+from vllm.v1.metrics.reader import Metric
19
+
20
+DEEPSEEK_MTP_MAIN_RANDOM = "luccafong/deepseek_mtp_main_random"
21
+DEEPSEEK_MTP_DRAFT_RANDOM = "luccafong/deepseek_mtp_draft_random"
22
+
23
+PROMPT = "The capital of France is"
24
+MAX_TOKENS = 8
25
+MAX_MODEL_LEN = 2048
26
+GPU_MEM_UTIL = 0.85
27
+# Spec / no-spec greedy output should match, but exact match is not guaranteed
28
+# across parallelism layouts (reduction order can flip near-tie argmaxes), so
29
+# gate on a similarity ratio like the spec-decode E2E tests.
30
+MIN_MATCH_RATIO = 0.8
31
+
32
+
33
+def _token_match_ratio(a: tuple[int, ...], b: tuple[int, ...]) -> float:
34
+ n = min(len(a), len(b))
35
+ return sum(x == y for x, y in zip(a, b)) / n if n else 0.0
36
+
37
+
38
+@dataclass(frozen=True)
39
+class InlineConfig:
40
+ id: str
41
+ tp_size: int = 1
42
+ pp_size: int = 1
43
+ enable_expert_parallel: bool = False
44
+ enable_eplb: bool = False
45
+ skip_reason: str | None = None
46
+
47
+
48
+INLINE_CONFIGS = [
49
+ InlineConfig(id="tp2", tp_size=2),
50
+ InlineConfig(id="ep2", tp_size=2, enable_expert_parallel=True),
51
+ InlineConfig(
52
+ id="ep2_eplb",
53
+ tp_size=2,
54
+ enable_expert_parallel=True,
55
+ enable_eplb=True,
56
+ ),
57
+ InlineConfig(
58
+ id="pp2",
59
+ pp_size=2,
60
+ skip_reason=(
61
+ "DeepSeek MTP pipeline-parallel support is in flight upstream; "
62
+ "see https://github.com/vllm-project/vllm/pull/38104"
63
+ ),
64
+ ),
65
+]
66
+
67
+
68
+def _inline_kwargs(config: InlineConfig, *, with_spec: bool) -> dict[str, Any]:
69
+ kwargs: dict[str, Any] = {
70
+ "model": DEEPSEEK_MTP_MAIN_RANDOM,
71
+ "tensor_parallel_size": config.tp_size,
72
+ "pipeline_parallel_size": config.pp_size,
73
+ "enable_expert_parallel": config.enable_expert_parallel,
74
+ "max_model_len": MAX_MODEL_LEN,
75
+ "gpu_memory_utilization": GPU_MEM_UTIL,
76
+ "enforce_eager": True,
77
+ "trust_remote_code": True,
78
+ "disable_log_stats": False,
79
+ # MLA + batch invariance currently disables prefix caching at runtime
80
+ # (see vllm/v1/attention/backends/mla/common.py); set explicitly so the
81
+ # configuration is unambiguous.
82
+ "enable_prefix_caching": False,
83
+ }
84
+ if config.enable_eplb:
85
+ kwargs["enable_eplb"] = True
86
+ # Rearrangement first fires after step_interval // 4 forward steps
87
+ # (the step counter starts at 3/4 of step_interval, see
88
+ # vllm/distributed/eplb/eplb_state.py). Keep these small so the EPLB
89
+ # routine actually runs within the short MAX_TOKENS generation instead
90
+ # of never triggering.
91
+ kwargs["eplb_config"] = {
92
+ "num_redundant_experts": config.tp_size,
93
+ "window_size": 2,
94
+ "step_interval": 4,
95
+ "log_balancedness": False,
96
+ }
97
+ if with_spec:
98
+ kwargs["speculative_config"] = {
99
+ "method": "mtp",
100
+ "model": DEEPSEEK_MTP_DRAFT_RANDOM,
101
+ "num_speculative_tokens": 1,
102
+ }
103
+ return kwargs
104
+
105
+
106
+def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]:
107
+ outputs = llm.generate(
108
+ [PROMPT],
109
+ SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True),
110
+ )
111
+ assert outputs and outputs[0].outputs, "expected one completion"
112
+ return tuple(outputs[0].outputs[0].token_ids)
113
+
114
+
115
+def _spec_decode_num_drafts(metrics: list[Metric]) -> int:
116
+ name2metric = {m.name: m for m in metrics}
117
+ counter = name2metric.get("vllm:spec_decode_num_drafts")
118
+ assert counter is not None, (
119
+ "spec_decode_num_drafts metric missing; check disable_log_stats=False"
120
+ )
121
+ return int(counter.value)
122
+
123
+
124
+@pytest.mark.parametrize(
125
+ "config",
126
+ [
127
+ pytest.param(c, id=c.id, marks=multi_gpu_marks(num_gpus=2))
128
+ for c in INLINE_CONFIGS
129
+ ],
130
+)
131
+def test_deepseek_mtp_load_inline(
132
+ monkeypatch: pytest.MonkeyPatch,
133
+ config: InlineConfig,
134
+):
135
+ """MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy."""
136
+ if config.skip_reason is not None:
137
+ pytest.skip(config.skip_reason)
138
+
139
+ # Reduces run-to-run nondeterminism so spec / no-spec stay close;
140
+ # see tests/v1/distributed/test_eagle_dp.py for the same pattern.
141
+ monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
142
+
143
+ spec_llm = LLM(**_inline_kwargs(config, with_spec=True))
144
+ try:
145
+ spec_tokens = _generate_token_ids_inline(spec_llm)
146
+ n_drafts = _spec_decode_num_drafts(spec_llm.get_metrics())
147
+ finally:
148
+ del spec_llm
149
+ torch.accelerator.empty_cache()
150
+ cleanup_dist_env_and_memory()
151
+
152
+ no_spec_llm = LLM(**_inline_kwargs(config, with_spec=False))
153
+ try:
154
+ no_spec_tokens = _generate_token_ids_inline(no_spec_llm)
155
+ finally:
156
+ del no_spec_llm
157
+ torch.accelerator.empty_cache()
158
+ cleanup_dist_env_and_memory()
159
+
160
+ # Non-vacuity: a silently-broken MTP drafter that falls back to
161
+ # verifier-only decoding would still produce matching output below.
162
+ assert n_drafts > 0, (
163
+ f"MTP drafter never fired under {config.id}: vllm:spec_decode_num_drafts == 0"
164
+ )
165
+ match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens)
166
+ print(f"\n{config.id}: spec/no-spec match_ratio={match_ratio:.3f}")
167
+ assert match_ratio >= MIN_MATCH_RATIO, (
168
+ f"Spec / no-spec output divergence under {config.id}: "
169
+ f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n"
170
+ f" spec_tokens = {spec_tokens}\n"
171
+ f" no_spec_tokens= {no_spec_tokens}"
172
+ )
173
+
174
+
175
+@pytest.mark.asyncio
176
+@pytest.mark.parametrize(
177
+ "dp_size",
178
+ [
179
+ pytest.param(2, marks=pytest.mark.distributed(num_gpus=2), id="dp2"),
180
+ pytest.param(1, id="dp1"),
181
+ ],
182
+)
183
+async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: int):
184
+ """MTP loads under DP (on and off) via AsyncLLM; spec matches no-spec greedy."""
185
+ if torch.accelerator.device_count() < dp_size:
186
+ pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs")
187
+
188
+ monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
189
+
190
+ base_args = AsyncEngineArgs(
191
+ model=DEEPSEEK_MTP_MAIN_RANDOM,
192
+ tensor_parallel_size=1,
193
+ data_parallel_size=dp_size,
194
+ data_parallel_backend="mp",
195
+ max_model_len=MAX_MODEL_LEN,
196
+ gpu_memory_utilization=GPU_MEM_UTIL,
197
+ enforce_eager=True,
198
+ trust_remote_code=True,
199
+ enable_prefix_caching=False,
200
+ )
201
+ spec_args = replace(
202
+ base_args,
203
+ speculative_config={
204
+ "method": "mtp",
205
+ "model": DEEPSEEK_MTP_DRAFT_RANDOM,
206
+ "num_speculative_tokens": 1,
207
+ },
208
+ )
209
+
210
+ sampling_params = SamplingParams(
211
+ max_tokens=MAX_TOKENS,
212
+ ignore_eos=True,
213
+ output_kind=RequestOutputKind.FINAL_ONLY,
214
+ temperature=0.0,
215
+ )
216
+
217
+ async def _generate(args: AsyncEngineArgs, request_id: str) -> tuple[int, ...]:
218
+ try:
219
+ async with AsyncExitStack() as stack:
220
+ engine = AsyncLLM.from_engine_args(args)
221
+ stack.callback(engine.shutdown)
222
+ async for out in engine.generate(
223
+ request_id=request_id,
224
+ prompt=PROMPT,
225
+ sampling_params=sampling_params,
226
+ ):
227
+ token_ids = tuple(out.outputs[0].token_ids)
228
+ assert len(token_ids) == MAX_TOKENS
229
+ return token_ids
230
+ raise AssertionError("AsyncLLM produced no output")
231
+ finally:
232
+ torch.accelerator.empty_cache()
233
+ cleanup_dist_env_and_memory()
234
+
235
+ spec_tokens = await asyncio.wait_for(
236
+ _generate(spec_args, f"deepseek-mtp-dp{dp_size}-spec"), timeout=600
237
+ )
238
+ no_spec_tokens = await asyncio.wait_for(
239
+ _generate(base_args, f"deepseek-mtp-dp{dp_size}-no-spec"), timeout=600
240
+ )
241
+
242
+ match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens)
243
+ print(f"\ndp{dp_size}: spec/no-spec match_ratio={match_ratio:.3f}")
244
+ assert match_ratio >= MIN_MATCH_RATIO, (
245
+ f"Spec / no-spec output divergence under dp{dp_size}: "
246
+ f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n"
247
+ f" spec_tokens = {spec_tokens}\n"
248
+ f" no_spec_tokens= {no_spec_tokens}"
249
+ )