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 6, 23:00 UTC
[Test] `test_async_scheduling.py` improvements (#36340)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 10, 2026 at 11:17 UTC
2a68464c5bf1a26821afe76cf49dc53f75b87e98
1 file changed
+57
-48
tests/v1/e2e/test_async_scheduling.py
+57
-48
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
from itertools import repeat
5
from typing import Any
6
@@ -19,6 +20,8 @@ from ...models.utils import check_outputs_equal
20
MODEL = "Qwen/Qwen3-0.6B"
21
MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
22
23
+# Need to enforce eager for MRV2 while we sort out cudagraph issues.
24
+ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1"
25
26
first_prompt = (
27
"The following numbers of the sequence "
@@ -47,10 +50,10 @@ def test_without_spec_decoding(
50
test_sampling_params: list[dict[str, Any]] = [
51
dict(),
52
# dict(min_tokens=20),
50
- dict(presence_penalty=-1.0),
53
+ dict(frequency_penalty=-1.0),
54
dict(bad_words=["the", " the"]),
55
dict(logprobs=2),
53
- dict(logprobs=2, presence_penalty=-1.0),
56
+ dict(logprobs=2, frequency_penalty=-1.0),
57
dict(structured_outputs=struct_outputs),
58
dict(
59
structured_outputs=struct_outputs,
@@ -58,12 +61,12 @@ def test_without_spec_decoding(
61
),
62
dict(
63
structured_outputs=struct_outputs,
61
- presence_penalty=-1.0,
64
+ frequency_penalty=-1.0,
65
),
66
dict(
67
structured_outputs=struct_outputs,
68
logprobs=2,
66
- presence_penalty=-1.0,
69
+ frequency_penalty=-1.0,
70
),
71
]
72
@@ -116,15 +119,15 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
119
120
test_sampling_params = [
121
dict(),
119
- dict(presence_penalty=-1.0),
122
+ dict(frequency_penalty=-1.0),
123
dict(bad_words=["the", " the"]),
124
dict(logprobs=2),
122
- dict(logprobs=2, presence_penalty=-1.0),
125
+ dict(logprobs=2, frequency_penalty=-1.0),
126
dict(structured_outputs=struct_outputs),
127
dict(
128
structured_outputs=struct_outputs,
129
logprobs=2,
127
- presence_penalty=-1.0,
130
+ frequency_penalty=-1.0,
131
),
132
]
133
@@ -144,14 +147,7 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
147
(True, "uni", True, spec_config_short, True),
148
]
149
147
- # On ROCm, use TRITON_ATTN + float32 for better numerical consistency
148
- run_tests(
149
- monkeypatch,
150
- MTP_MODEL,
151
- test_configs,
152
- test_sampling_params,
153
- is_testing_with_spec_decoding=True,
154
- )
150
+ run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
151
152
153
def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
@@ -196,12 +192,11 @@ def run_tests(
192
model: str,
193
test_configs: list[tuple],
194
test_sampling_params: list[dict[str, Any]],
199
- is_testing_with_spec_decoding: bool = False,
195
):
196
"""Test consistency of combos of async scheduling, preemption,
197
uni/multiproc executor with spec decoding."""
198
204
- # Determine attention config based on platform
199
+ # Flex attention supports float32.
200
attention_config = {"backend": "FLEX_ATTENTION"}
201
202
with monkeypatch.context() as m:
@@ -226,7 +221,6 @@ def run_tests(
221
async_scheduling,
222
spec_config,
223
test_prefill_chunking=test_prefill_chunking,
229
- is_testing_with_spec_decoding=is_testing_with_spec_decoding,
224
attention_config=attention_config,
225
)
226
outputs.append(test_results)
@@ -250,6 +244,7 @@ def run_tests(
244
test_acceptance_rates or repeat(None),
245
test_sampling_params,
246
):
247
+ reason = None
248
try:
249
check_outputs_equal(
250
outputs_0_lst=base_outs,
@@ -257,42 +252,57 @@ def run_tests(
252
name_0=f"baseline=[{baseline_config}], params={params}",
253
name_1=f"config=[{test_config}], params={params}",
254
)
260
-
261
- assert _all_logprobs_match(base_logprobs, test_logprobs)
262
-
263
- if (
264
- base_acceptance_rate is not None
265
- and test_acceptance_rate is not None
266
- ):
267
- if "spec_mml=None" in test_config:
268
- # Preemption causes more variance in acceptance rates
269
- if (
270
- current_platform.is_rocm()
271
- and "preemption=True" in test_config
272
- ):
273
- tolerance = 0.10
255
+ except AssertionError as e:
256
+ reason = "outputs ", e
257
+
258
+ if reason is None:
259
+ try:
260
+ assert _all_logprobs_match(base_logprobs, test_logprobs)
261
+ except AssertionError as e:
262
+ reason = "logprobs", e
263
+
264
+ if reason is None:
265
+ try:
266
+ if (
267
+ base_acceptance_rate is not None
268
+ and test_acceptance_rate is not None
269
+ ):
270
+ if "spec_mml=None" in test_config:
271
+ # Preemption causes more variance in acceptance rates
272
+ if (
273
+ current_platform.is_rocm()
274
+ and "preemption=True" in test_config
275
+ ):
276
+ tolerance = 0.10
277
+ else:
278
+ tolerance = 0.05
279
+ assert (
280
+ test_acceptance_rate > base_acceptance_rate
281
+ or test_acceptance_rate
282
+ == pytest.approx(base_acceptance_rate, rel=tolerance)
283
+ )
284
else:
275
- tolerance = 0.05
276
- assert (
277
- test_acceptance_rate > base_acceptance_rate
278
- or test_acceptance_rate
279
- == pytest.approx(base_acceptance_rate, rel=tolerance)
280
- )
281
- else:
282
- # Currently the reported acceptance rate is expected to be
283
- # lower when we sometimes skip drafting altogether.
284
- assert test_acceptance_rate > 0.1
285
+ # Currently the reported acceptance rate is expected to be
286
+ # lower when we sometimes skip drafting altogether.
287
+ assert test_acceptance_rate > 0.1
288
+ except AssertionError as e:
289
+ reason = "accept ", e
290
+
291
+ if reason is None:
292
print(
286
- f"PASSED: config=[{test_config}], params={params}"
293
+ f"\033[32mPASSED\033[0m: "
294
+ f"config=[{test_config}], params={params}"
295
f" accept_rate={test_acceptance_rate}"
296
)
289
- except AssertionError as e:
297
+ else:
298
+ reason_str, _ = reason
299
print(
291
- f"FAILED: config=[{test_config}], params={params}"
300
+ f"\033[31mFAILED\033[0m({reason_str}): "
301
+ f"config=[{test_config}], params={params}"
302
f" accept_rate={test_acceptance_rate}"
303
)
304
if failure is None:
295
- failure = e
305
+ _, failure = reason
306
307
if failure is not None:
308
raise failure
@@ -307,7 +317,6 @@ def run_test(
317
async_scheduling: bool,
318
spec_config: dict[str, Any] | None,
319
test_prefill_chunking: bool,
310
- is_testing_with_spec_decoding: bool = False,
320
attention_config: dict[str, Any] | None = None,
321
):
322
spec_decoding = spec_config is not None
@@ -335,7 +344,7 @@ def run_test(
344
enable_chunked_prefill=test_prefill_chunking,
345
# Force prefill chunking
346
max_num_batched_tokens=48 if test_prefill_chunking else None,
338
- # enforce_eager=True,
347
+ enforce_eager=ENFORCE_EAGER,
348
async_scheduling=async_scheduling,
349
distributed_executor_backend=executor,
350
dtype="float32",