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 5, 11:00 UTC
`tests/v1/e2e/spec_decode`: assert async scheduling is used (#39206)
Signed-off-by: Rishi Puri <riship@nvidia.com> Signed-off-by: Rishi Puri <puririshi98@berkeley.edu> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
Rishi Puri committed
Apr 8, 2026 at 13:30 UTC
ad05edfbcac547cf233cf993b278ce8b34e732c5
2 files changed
+58
-4
tests/v1/e2e/spec_decode/test_async_spec_decode.py
+5
@@ -116,6 +116,11 @@ def test_no_sync_with_spec_decode(
116
async_scheduling=True,
117
)
118
119
+ # Assert async scheduling is actually active before running inference.
120
+ assert llm.llm_engine.vllm_config.scheduler_config.async_scheduling, (
121
+ f"Expected async_scheduling=True for spec decode, got False. method={method}"
122
+ )
123
+
124
outputs = llm.generate(
125
["Hello, my name is"],
126
SamplingParams(temperature=0, max_tokens=10),
tests/v1/e2e/spec_decode/test_spec_decode.py
+53
-4
@@ -30,6 +30,13 @@ from vllm.v1.metrics.reader import Metric
30
MTP_SIMILARITY_RATE = 0.8
31
32
33
+class AsyncSchedulingNotEnabledError(AssertionError):
34
+ """Raised when async_scheduling is expected to be True for draft_model
35
+ spec decode but is False. Tracked in:
36
+ https://github.com/vllm-project/vllm/issues/38929
37
+ """
38
+
39
+
40
def _skip_if_insufficient_gpus_for_tp(tp_size: int):
41
"""Skip test if available GPUs < tp_size on ROCm."""
42
available_gpus = torch.accelerator.device_count()
@@ -206,6 +213,8 @@ def test_ngram_gpu_default_with_async_scheduling(
213
max_model_len=4096,
214
async_scheduling=async_scheduling,
215
)
216
+ # Assert the resolved async_scheduling config matches what was requested.
217
+ assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling == async_scheduling
218
evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
219
del spec_llm
220
cleanup_dist_env_and_memory()
@@ -457,6 +466,8 @@ def _run_eagle_correctness(
466
model_impl=model_impl,
467
attention_config=attention_config,
468
)
469
+ # EAGLE/EAGLE3 supports async scheduling; assert it is active by default.
470
+ assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
471
evaluate_llm_for_gsm8k(
472
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
473
)
@@ -760,6 +771,8 @@ def test_mtp_correctness(
771
max_model_len=2048,
772
attention_backend=attn_backend,
773
)
774
+ # MTP supports async scheduling; assert it is active by default.
775
+ assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
776
evaluate_llm_for_gsm8k(
777
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
778
)
@@ -829,12 +842,22 @@ cases = [
842
@pytest.mark.parametrize("args", cases)
843
@pytest.mark.parametrize("enforce_eager", [True, False])
844
@single_gpu_only
845
+# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
846
+@pytest.mark.xfail(
847
+ raises=AsyncSchedulingNotEnabledError,
848
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
849
+)
850
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
851
args.enforce_eager = enforce_eager
852
assert_draft_model_correctness(args)
853
854
855
@single_gpu_only
856
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
857
+@pytest.mark.xfail(
858
+ raises=AsyncSchedulingNotEnabledError,
859
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
860
+)
861
def test_draft_model_realistic_example():
862
args = ArgsTest(
863
target_model="Qwen/Qwen3-1.7B",
@@ -850,6 +873,11 @@ def test_draft_model_realistic_example():
873
874
875
@single_gpu_only
876
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
877
+@pytest.mark.xfail(
878
+ raises=AsyncSchedulingNotEnabledError,
879
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
880
+)
881
def test_draft_model_parallel_drafting():
882
args = ArgsTest(
883
target_model="Qwen/Qwen3-1.7B",
@@ -876,6 +904,11 @@ def test_draft_model_parallel_drafting():
904
)
905
@pytest.mark.parametrize("enforce_eager", [True, False])
906
@single_gpu_only
907
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
908
+@pytest.mark.xfail(
909
+ raises=AsyncSchedulingNotEnabledError,
910
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
911
+)
912
def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
913
tgt_model, draft_model = models
914
sd_case = ArgsTest(
@@ -888,6 +921,11 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
921
922
923
@multi_gpu_only(num_gpus=2)
924
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
925
+@pytest.mark.xfail(
926
+ raises=AsyncSchedulingNotEnabledError,
927
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
928
+)
929
def test_draft_model_tensor_parallelism():
930
"""Ensure spec decode works when running with TP > 1."""
931
_skip_if_insufficient_gpus_for_tp(2)
@@ -1062,6 +1100,7 @@ def assert_draft_model_correctness(args: ArgsTest):
1100
enforce_eager=args.enforce_eager,
1101
disable_log_stats=False, # enables get_metrics()
1102
)
1103
+
1104
# we don't check the outputs, only check the metrics
1105
spec_llm.chat(test_prompts, args.sampling_config)
1106
metrics = spec_llm.get_metrics()
@@ -1073,10 +1112,6 @@ def assert_draft_model_correctness(args: ArgsTest):
1112
spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
1113
)
1114
1076
- del spec_llm # CLEANUP
1077
- torch.accelerator.empty_cache()
1078
- cleanup_dist_env_and_memory()
1079
-
1115
print(
1116
f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
1117
f"temperature={args.sampling_config.temperature:.2f}, "
@@ -1086,6 +1121,20 @@ def assert_draft_model_correctness(args: ArgsTest):
1121
1122
assert acceptance_rate >= args.expected_acceptance_rate
1123
assert acceptance_len >= args.expected_acceptance_len
1124
+ # draft_model supports async scheduling; assert it is active by default.
1125
+ # Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that
1126
+ # @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this
1127
+ # specific failure — leaving all other assertion failures (e.g. correctness or
1128
+ # acceptance-rate checks above) visible as real test failures.
1129
+ has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
1130
+ del spec_llm # CLEANUP
1131
+ torch.accelerator.empty_cache()
1132
+ cleanup_dist_env_and_memory()
1133
+ if not has_async:
1134
+ raise AsyncSchedulingNotEnabledError(
1135
+ "Expected async_scheduling=True for draft_model spec decode, got False."
1136
+ " See https://github.com/vllm-project/vllm/issues/38929"
1137
+ )
1138
1139
1140
def get_messages(dataset: str, n: int) -> list[Messages]: