@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 58e61e56b7
Fork 0

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 13, 04:30 UTC

[Test] Rework e2e async scheduling tests (#28744)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Nov 14, 2025 at 16:01 UTC 58e61e56b744da109269586fe45ecc47b10dca5f
1 file changed +268 -90
tests/v1/e2e/test_async_scheduling.py
+268 -90
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from itertools import repeat
4 from typing import Any
5
6 import pytest
@@ -8,126 +9,291 @@ import torch._dynamo.config as dynamo_config
9 from vllm import SamplingParams
10 from vllm.logprobs import Logprob
11 from vllm.sampling_params import StructuredOutputsParams
12 +from vllm.v1.metrics.reader import Metric
13
14 from ...conftest import VllmRunner
15 from ...models.utils import check_outputs_equal
16
17 MODEL = "Qwen/Qwen3-0.6B"
18 +MTP_MODEL = "XiaomiMiMo/MiMo-7B-Base"
19
20
18 -@dynamo_config.patch(cache_size_limit=16)
19 -def test_preempt_and_async_scheduling_e2e(
20 - sample_json_schema, monkeypatch: pytest.MonkeyPatch
21 -):
22 - """Test consistency of combos of async scheduling, preemption,
23 - uni/multiproc executor, and various sampling parameters
24 - including structured outputs."""
21 +first_prompt = (
22 + "The following numbers of the sequence "
23 + + ", ".join(str(i) for i in range(10))
24 + + " are:"
25 +)
26 +example_prompts = [first_prompt, "In one word, the capital of France is "] + [
27 + f"Tell me about the number {i}: " for i in range(32)
28 +]
29
26 - first_prompt = (
27 - "The following numbers of the sequence "
28 - + ", ".join(str(i) for i in range(10))
29 - + " are:"
30 - )
31 - example_prompts = [first_prompt, "In one word, the capital of France is "] + [
32 - f"Tell me about the number {i}: " for i in range(32)
33 - ]
30 +default_params = dict(
31 + temperature=0.0, # greedy
32 + max_tokens=20,
33 +)
34
35 - sampling_param_tests: list[dict[str, Any]] = [
35 +
36 +def test_without_spec_decoding(
37 + sample_json_schema,
38 + monkeypatch: pytest.MonkeyPatch,
39 +):
40 + """Test consistency of combos of async scheduling, preemption,
41 + uni/multiproc executor, prefill chunking."""
42 + struct_outputs = StructuredOutputsParams(json=sample_json_schema)
43 + test_sampling_params: list[dict[str, Any]] = [
44 dict(),
45 # dict(min_tokens=20),
46 dict(presence_penalty=-1.0),
47 dict(bad_words=["the", " the"]),
48 dict(logprobs=2),
49 dict(logprobs=2, presence_penalty=-1.0),
42 - dict(structured_outputs=StructuredOutputsParams(json=sample_json_schema)),
50 + dict(structured_outputs=struct_outputs),
51 dict(
44 - structured_outputs=StructuredOutputsParams(json=sample_json_schema),
52 + structured_outputs=struct_outputs,
53 logprobs=2,
54 presence_penalty=-1.0,
55 ),
56 ]
57
50 - default_params = dict(
51 - temperature=0.0, # greedy
52 - max_tokens=20,
58 + # test_preemption, executor, async_scheduling,
59 + # spec_config, test_prefill_chunking
60 + test_configs = [
61 + (False, "mp", False, None, False),
62 + (True, "mp", False, None, True),
63 + (False, "mp", True, None, False),
64 + (False, "uni", True, None, False),
65 + (True, "mp", True, None, False),
66 + (True, "uni", True, None, False),
67 + (False, "mp", True, None, True),
68 + # Async scheduling + preemption + chunked prefill needs to be fixed (WIP)
69 + # (True, "mp", True, None, True),
70 + # (True, "uni", True, None, True),
71 + ]
72 +
73 + run_tests(
74 + monkeypatch,
75 + MODEL,
76 + test_configs,
77 + test_sampling_params,
78 )
79
80 +
81 +@pytest.mark.skip("MTP model too big to run in fp32 in CI")
82 +def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):
83 + """Test consistency and acceptance rates with some different combos of
84 + preemption, executor, async scheduling, prefill chunking,
85 + spec decoding model length.
86 + """
87 +
88 + spec_config = {
89 + "method": "mtp",
90 + "num_speculative_tokens": 2,
91 + }
92 + spec_config_short = spec_config | {"max_model_len": 50}
93 +
94 + # test_preemption, executor, async_scheduling,
95 + # spec_config, test_prefill_chunking
96 + test_configs = [
97 + (False, "mp", False, None, False),
98 + (False, "mp", False, spec_config, False),
99 + (True, "mp", False, spec_config, True),
100 + (True, "uni", False, spec_config_short, True),
101 + (False, "mp", True, spec_config, False),
102 + (True, "mp", True, spec_config, False),
103 + (False, "mp", True, spec_config_short, True),
104 + (True, "uni", True, spec_config, False),
105 + (True, "uni", True, spec_config_short, False),
106 + # Async scheduling + preemption + chunked prefill needs to be fixed (WIP)
107 + # (True, "mp", True, spec_config, True),
108 + # (True, "uni", True, spec_config_short, True),
109 + ]
110 +
111 + run_tests(
112 + monkeypatch,
113 + MTP_MODEL,
114 + test_configs,
115 + [{}],
116 + )
117 +
118 +
119 +@dynamo_config.patch(cache_size_limit=16)
120 +def run_tests(
121 + monkeypatch: pytest.MonkeyPatch,
122 + model: str,
123 + test_configs: list[tuple],
124 + test_sampling_params: list[dict[str, Any]],
125 +):
126 + """Test consistency of combos of async scheduling, preemption,
127 + uni/multiproc executor with spec decoding."""
128 +
129 with monkeypatch.context() as m:
130 + # avoid precision errors
131 m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
132 # m.setenv("VLLM_BATCH_INVARIANT", "1")
133 + outputs: list[tuple[str, list, list]] = []
134 + for n, (
135 + test_preemption,
136 + executor,
137 + async_scheduling,
138 + spec_config,
139 + test_prefill_chunking,
140 + ) in enumerate(test_configs, 1):
141 + test_str = f"{n}/{len(test_configs)}"
142 + test_results = run_test(
143 + model,
144 + test_str,
145 + test_sampling_params,
146 + test_preemption,
147 + executor,
148 + async_scheduling,
149 + spec_config,
150 + test_prefill_chunking=test_prefill_chunking,
151 + )
152 + outputs.append(test_results)
153 +
154 + baseline_config, baseline_tests, _ = outputs[0]
155 + _, _, baseline_acceptances = next(
156 + (o for o in outputs if o[2] is not None), (None, None, None)
157 + )
158
59 - outputs: list[tuple[str, list]] = []
60 - for test_preemption in [False, True]:
61 - for executor in ["mp", "uni"]:
62 - for async_scheduling in [False, True]:
63 - cache_arg: dict[str, Any] = (
64 - dict(num_gpu_blocks_override=32)
65 - if test_preemption
66 - else dict(gpu_memory_utilization=0.7)
67 - )
68 - test_config = (
69 - f"executor={executor}, preemption={test_preemption},"
70 - f" async_sched={async_scheduling}"
71 - )
72 - print("-" * 80)
73 - print(f"---- TESTING: {test_config}")
74 - print("-" * 80)
75 - with VllmRunner(
76 - MODEL,
77 - max_model_len=512,
78 - enforce_eager=True,
79 - async_scheduling=async_scheduling,
80 - distributed_executor_backend=executor,
81 - dtype="float32", # avoid precision errors
82 - **cache_arg,
83 - ) as vllm_model:
84 - results = []
85 - for override_params in sampling_param_tests:
86 - print(f"----------- RUNNING PARAMS: {override_params}")
87 - results.append(
88 - vllm_model.generate(
89 - example_prompts,
90 - sampling_params=SamplingParams(
91 - **default_params, **override_params
92 - ),
93 - return_logprobs=True,
94 - )
95 - )
96 -
97 - if not outputs:
98 - # First check that the different parameter configs
99 - # actually result in different output.
100 - for (other_test_outs, other_test_logprobs), params in zip(
101 - results[1:], sampling_param_tests[1:]
102 - ):
103 - with pytest.raises(AssertionError):
104 - check_outputs_equal(
105 - outputs_0_lst=results[0][0],
106 - outputs_1_lst=other_test_outs,
107 - name_0=f"baseline params={params}",
108 - name_1=f"other params={params}",
109 - )
110 - assert _all_logprobs_match(
111 - results[0][1], other_test_logprobs
112 - )
113 -
114 - outputs.append((test_config, results))
115 -
116 - baseline_config, baseline_tests = outputs[0]
117 -
118 - for test_config, test_outputs in outputs[1:]:
119 - for (base_outs, base_logprobs), (test_outs, test_logprobs), params in zip(
120 - baseline_tests, test_outputs, sampling_param_tests
159 + print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}")
160 +
161 + failure = None
162 + for test_config, test_outputs, test_acceptance_rates in outputs[1:]:
163 + for (base_outs, base_logprobs), base_acceptance_rate, (
164 + test_outs,
165 + test_logprobs,
166 + ), test_acceptance_rate, params in zip(
167 + baseline_tests,
168 + baseline_acceptances or repeat(None),
169 + test_outputs,
170 + test_acceptance_rates or repeat(None),
171 + test_sampling_params,
172 ):
122 - check_outputs_equal(
123 - outputs_0_lst=base_outs,
124 - outputs_1_lst=test_outs,
125 - name_0=f"baseline=[{baseline_config}], params={params}",
126 - name_1=f"config=[{test_config}], params={params}",
173 + try:
174 + check_outputs_equal(
175 + outputs_0_lst=base_outs,
176 + outputs_1_lst=test_outs,
177 + name_0=f"baseline=[{baseline_config}], params={params}",
178 + name_1=f"config=[{test_config}], params={params}",
179 + )
180 + assert _all_logprobs_match(base_logprobs, test_logprobs)
181 +
182 + if (
183 + base_acceptance_rate is not None
184 + and test_acceptance_rate is not None
185 + ):
186 + if "spec_mml=None" in test_config:
187 + # because the acceptance rate can vary, we use a looser
188 + # tolerance here.
189 + assert (
190 + pytest.approx(test_acceptance_rate, rel=5e-2)
191 + == base_acceptance_rate
192 + )
193 + else:
194 + # Currently the reported acceptance rate is expected to be
195 + # lower when we skip drafting altogether.
196 + assert test_acceptance_rate > 0.05
197 + print(
198 + f"PASSED: config=[{test_config}], params={params}"
199 + f" accept_rate={test_acceptance_rate}"
200 + )
201 + except AssertionError as e:
202 + print(
203 + f"FAILED: config=[{test_config}], params={params}"
204 + f" accept_rate={test_acceptance_rate}"
205 + )
206 + if failure is None:
207 + failure = e
208 +
209 + if failure is not None:
210 + raise failure
211 +
212 +
213 +def run_test(
214 + model: str,
215 + test_str: str,
216 + sampling_param_tests: list[dict[str, Any]],
217 + test_preemption: bool,
218 + executor: str,
219 + async_scheduling: bool,
220 + spec_config: dict[str, Any] | None,
221 + test_prefill_chunking: bool,
222 +):
223 + spec_decoding = spec_config is not None
224 + cache_arg: dict[str, Any] = (
225 + dict(num_gpu_blocks_override=32)
226 + if test_preemption
227 + else dict(gpu_memory_utilization=0.9)
228 + )
229 + spec_mml = (spec_config or {}).get("max_model_len")
230 + test_config = (
231 + f"executor={executor}, preemption={test_preemption}, "
232 + f"async_sched={async_scheduling}, "
233 + f"chunk_prefill={test_prefill_chunking}, "
234 + f"spec_decoding={spec_decoding}, spec_mml={spec_mml}"
235 + )
236 + print("-" * 80)
237 + print(f"---- TESTING {test_str}: {test_config}")
238 + print("-" * 80)
239 + with VllmRunner(
240 + model,
241 + max_model_len=512,
242 + enable_chunked_prefill=test_prefill_chunking,
243 + max_num_batched_tokens=48 if test_prefill_chunking else None,
244 + # enforce_eager=True,
245 + async_scheduling=async_scheduling,
246 + distributed_executor_backend=executor,
247 + dtype="float32", # avoid precision errors
248 + speculative_config=spec_config,
249 + disable_log_stats=False,
250 + **cache_arg,
251 + ) as vllm_model:
252 + results = []
253 + acceptance_rates: list[float] | None = [] if spec_decoding else None
254 + for override_params in sampling_param_tests:
255 + metrics_before = vllm_model.llm.get_metrics()
256 + print(f"----------- RUNNING PARAMS: {override_params}")
257 + results.append(
258 + vllm_model.generate(
259 + example_prompts,
260 + sampling_params=SamplingParams(
261 + **default_params,
262 + **override_params,
263 + ),
264 + return_logprobs=True,
265 + )
266 )
128 - assert _all_logprobs_match(base_logprobs, test_logprobs)
267 + metrics_after = vllm_model.llm.get_metrics()
268 + if acceptance_rates is not None:
269 + acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
270 + acceptance_rates.append(acceptance_rate)
271 + print(f"ACCEPTANCE RATE {acceptance_rate}")
272 +
273 + if test_preemption:
274 + preemptions = _get_count(
275 + metrics_before,
276 + metrics_after,
277 + "vllm:num_preemptions",
278 + )
279 + assert preemptions > 0, "preemption test had no preemptions"
280 +
281 + if len(results) > 1:
282 + # First check that the different parameter configs
283 + # actually result in different output.
284 + for (other_test_outs, other_test_logprobs), params in zip(
285 + results[1:], sampling_param_tests[1:]
286 + ):
287 + with pytest.raises(AssertionError):
288 + check_outputs_equal(
289 + outputs_0_lst=results[0][0],
290 + outputs_1_lst=other_test_outs,
291 + name_0=f"baseline params={params}",
292 + name_1=f"other params={params}",
293 + )
294 + assert _all_logprobs_match(results[0][1], other_test_logprobs)
295
130 - print(f"PASSED: config=[{test_config}], params={params}")
296 + return test_config, results, acceptance_rates
297
298
299 def _all_logprobs_match(req_a, req_b) -> bool:
@@ -149,3 +315,15 @@ def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> boo
315 and a.logprob == pytest.approx(b.logprob, rel=1e-3, abs=1e-6)
316 for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
317 )
318 +
319 +
320 +def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
321 + draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
322 + accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
323 + return accept / draft if draft > 0 else 0.0
324 +
325 +
326 +def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
327 + before_val = next(m.value for m in before if m.name == name)
328 + after_val = next(m.value for m in after if m.name == name)
329 + return after_val - before_val