[Bugfix][Spec Decode] Auto-enable async scheduling for draft models (#48341)
Signed-off-by: BWAAEEEK <jooho414@gmail.com>
JooHo Lee committed
Aug 6, 2026 at 14:51 UTC
c0202c560333e8e475819a93a7cc00d9b2a115db
3 files changed
+25
-39
tests/test_config.py
+23
@@ -407,6 +407,29 @@ def test_reconfigure_for_independent_dp_rank_on_multinode_dense_model():
407
assert parallel_config.world_size == 8
408
409
410
+def test_draft_model_enables_async_scheduling_by_default():
411
+ parallel_config = ParallelConfig(distributed_executor_backend="uni")
412
+ model_config = ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048)
413
+ speculative_config = SpeculativeConfig(
414
+ method="draft_model",
415
+ model="Qwen/Qwen3-0.6B",
416
+ num_speculative_tokens=3,
417
+ target_model_config=model_config,
418
+ target_parallel_config=parallel_config,
419
+ )
420
+ cfg = VllmConfig(
421
+ model_config=model_config,
422
+ scheduler_config=SchedulerConfig(
423
+ max_model_len=2048,
424
+ is_encoder_decoder=False,
425
+ ),
426
+ parallel_config=parallel_config,
427
+ speculative_config=speculative_config,
428
+ )
429
+
430
+ assert cfg.scheduler_config.async_scheduling is True
431
+
432
+
433
@dataclass
434
class _TestConfigFields:
435
a: int
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
+1
-39
@@ -25,10 +25,6 @@ from ..utils import (
25
)
26
27
28
-class AsyncSchedulingNotEnabledError(AssertionError):
29
- """Raised when draft-model spec decode does not enable async scheduling."""
30
-
31
-
28
@dataclass
29
class ArgsTest:
30
target_model: str
@@ -93,22 +89,12 @@ cases = [
89
@pytest.mark.parametrize("args", cases)
90
@pytest.mark.parametrize("enforce_eager", [True, False])
91
@single_gpu_only
96
-# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
97
-@pytest.mark.xfail(
98
- raises=AsyncSchedulingNotEnabledError,
99
- reason="draft_model does not yet enable async_scheduling: issue #38929",
100
-)
92
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
93
args.enforce_eager = enforce_eager
94
assert_draft_model_correctness(args)
95
96
97
@single_gpu_only
107
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
108
-@pytest.mark.xfail(
109
- raises=AsyncSchedulingNotEnabledError,
110
- reason="draft_model does not yet enable async_scheduling: issue #38929",
111
-)
98
def test_draft_model_realistic_example():
99
args = ArgsTest(
100
target_model="Qwen/Qwen3-1.7B",
@@ -124,11 +110,6 @@ def test_draft_model_realistic_example():
110
111
112
@single_gpu_only
127
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
128
-@pytest.mark.xfail(
129
- raises=AsyncSchedulingNotEnabledError,
130
- reason="draft_model does not yet enable async_scheduling: issue #38929",
131
-)
113
def test_draft_model_parallel_drafting():
114
args = ArgsTest(
115
target_model="Qwen/Qwen3-1.7B",
@@ -155,11 +136,6 @@ def test_draft_model_parallel_drafting():
136
)
137
@pytest.mark.parametrize("enforce_eager", [True, False])
138
@single_gpu_only
158
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
159
-@pytest.mark.xfail(
160
- raises=AsyncSchedulingNotEnabledError,
161
- reason="draft_model does not yet enable async_scheduling: issue #38929",
162
-)
139
def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
140
tgt_model, draft_model = models
141
sd_case = ArgsTest(
@@ -172,11 +148,6 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
148
149
150
@multi_gpu_only(num_gpus=2)
175
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
176
-@pytest.mark.xfail(
177
- raises=AsyncSchedulingNotEnabledError,
178
- reason="draft_model does not yet enable async_scheduling: issue #38929",
179
-)
151
def test_draft_model_tensor_parallelism():
152
"""Ensure spec decode works when running with TP > 1."""
153
_skip_if_insufficient_gpus_for_tp(2)
@@ -372,17 +343,8 @@ def assert_draft_model_correctness(args: ArgsTest):
343
344
assert acceptance_rate >= args.expected_acceptance_rate
345
assert acceptance_len >= args.expected_acceptance_len
375
- # draft_model supports async scheduling; assert it is active by default.
376
- # Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that
377
- # @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this
378
- # specific failure — leaving all other assertion failures (e.g. correctness or
379
- # acceptance-rate checks above) visible as real test failures.
346
has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
347
del spec_llm # CLEANUP
348
torch.accelerator.empty_cache()
349
cleanup_dist_env_and_memory()
384
- if not has_async:
385
- raise AsyncSchedulingNotEnabledError(
386
- "Expected async_scheduling=True for draft_model spec decode, got False."
387
- " See https://github.com/vllm-project/vllm/issues/38929"
388
- )
350
+ assert has_async, "Expected async_scheduling=True for draft_model spec decode"
vllm/config/vllm.py
+1
@@ -1126,6 +1126,7 @@ class VllmConfig:
1126
self.speculative_config is not None
1127
and self.speculative_config.method not in get_args(EagleModelTypes)
1128
and self.speculative_config.method not in get_args(NgramGPUTypes)
1129
+ and self.speculative_config.method != "draft_model"
1130
and self.speculative_config.method != "dspark"
1131
):
1132
logger.warning_once(