[Bugfix][Spec Decode] Cap DFlash/DSpark profiling query batch (#56448)
Signed-off-by: wangyicong <wangyicong@bytedance.com>
Yicong Wang committed
Sep 22, 2026 at 12:36 UTC
c4d424c2e27b392da235704ba0edca7ac48bd495
3 files changed
+105
tests/v1/spec_decode/test_dflash_profile.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from contextlib import nullcontext
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+
9
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
10
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
11
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
12
+from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator
13
+
14
+
15
+class _DefaultSpeculator(BaseSpeculator):
16
+ def init_cudagraph_manager(self, cudagraph_mode):
17
+ pass
18
+
19
+ def capture(self):
20
+ pass
21
+
22
+ def propose(self, *args, **kwargs):
23
+ raise NotImplementedError
24
+
25
+
26
+class _StopDummyRun(Exception):
27
+ pass
28
+
29
+
30
+def _run_dummy_batch(
31
+ speculator,
32
+ *,
33
+ max_num_reqs=512,
34
+ is_profile=False,
35
+ skip_attn=False,
36
+):
37
+ captured = {}
38
+
39
+ def execute_model(scheduler_output, **kwargs):
40
+ captured["scheduler_output"] = scheduler_output
41
+ raise _StopDummyRun
42
+
43
+ runner = SimpleNamespace(
44
+ max_num_reqs=max_num_reqs,
45
+ max_num_tokens=2048,
46
+ decode_query_len=1,
47
+ speculator=speculator,
48
+ kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
49
+ is_first_pp_rank=True,
50
+ lora_config=None,
51
+ maybe_dummy_run_with_lora=lambda *args, **kwargs: nullcontext(),
52
+ execute_model=execute_model,
53
+ )
54
+
55
+ with pytest.raises(_StopDummyRun):
56
+ GPUModelRunner._dummy_run(
57
+ runner,
58
+ num_tokens=2048,
59
+ skip_attn=skip_attn,
60
+ is_profile=is_profile,
61
+ )
62
+
63
+ return captured["scheduler_output"]
64
+
65
+
66
+def test_default_speculator_query_width_is_one():
67
+ assert _DefaultSpeculator().num_query_per_req == 1
68
+
69
+
70
+@pytest.mark.parametrize("speculator_cls", [DFlashSpeculator, DSparkSpeculator])
71
+@pytest.mark.parametrize(
72
+ ("is_profile", "skip_attn"),
73
+ [(False, False), (True, False), (True, True)],
74
+)
75
+def test_dummy_run_caps_parallel_draft_requests(speculator_cls, is_profile, skip_attn):
76
+ speculator = speculator_cls.__new__(speculator_cls)
77
+ speculator.num_query_per_req = 5
78
+
79
+ scheduler_output = _run_dummy_batch(
80
+ speculator,
81
+ is_profile=is_profile,
82
+ skip_attn=skip_attn,
83
+ )
84
+ assert len(scheduler_output.num_scheduled_tokens) == 409
85
+ assert scheduler_output.total_num_scheduled_tokens == 2048
86
+
87
+
88
+def test_dummy_run_keeps_already_fitting_parallel_draft_batch():
89
+ speculator = DFlashSpeculator.__new__(DFlashSpeculator)
90
+ speculator.num_query_per_req = 5
91
+
92
+ scheduler_output = _run_dummy_batch(speculator, max_num_reqs=128)
93
+ assert len(scheduler_output.num_scheduled_tokens) == 128
94
+
95
+
96
+def test_dummy_run_keeps_default_speculator_batch():
97
+ scheduler_output = _run_dummy_batch(_DefaultSpeculator())
98
+ assert len(scheduler_output.num_scheduled_tokens) == 512
vllm/v1/worker/gpu/model_runner.py
+5
@@ -790,6 +790,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
790
num_tokens = max(num_tokens, self.decode_query_len)
791
num_reqs = num_tokens // self.decode_query_len
792
assert num_tokens % self.decode_query_len == 0
793
+ elif self.speculator is not None:
794
+ num_reqs = min(
795
+ num_reqs,
796
+ self.max_num_tokens // self.speculator.num_query_per_req,
797
+ )
798
# Distribute the remainder evenly so no dummy request exceeds
799
# ceil(num_tokens / num_reqs) <= max_model_len tokens.
800
num_tokens_per_request = [
vllm/v1/worker/gpu/spec_decode/speculator.py
+2
@@ -57,6 +57,8 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
57
58
59
class BaseSpeculator(ABC):
60
+ num_query_per_req: int = 1
61
+
62
@abstractmethod
63
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
64
pass