[Bugfix][Spec Decode] Cap DFlash/DSpark profiling query batch (#56448)

Signed-off-by: wangyicong <wangyicong@bytedance.com>

tests/v1/spec_decode/test_dflash_profile.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator
+
+
+class _DefaultSpeculator(BaseSpeculator):
+    def init_cudagraph_manager(self, cudagraph_mode):
+        pass
+
+    def capture(self):
+        pass
+
+    def propose(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class _StopDummyRun(Exception):
+    pass
+
+
+def _run_dummy_batch(
+    speculator,
+    *,
+    max_num_reqs=512,
+    is_profile=False,
+    skip_attn=False,
+):
+    captured = {}
+
+    def execute_model(scheduler_output, **kwargs):
+        captured["scheduler_output"] = scheduler_output
+        raise _StopDummyRun
+
+    runner = SimpleNamespace(
+        max_num_reqs=max_num_reqs,
+        max_num_tokens=2048,
+        decode_query_len=1,
+        speculator=speculator,
+        kv_connector=SimpleNamespace(set_disabled=lambda disabled: None),
+        is_first_pp_rank=True,
+        lora_config=None,
+        maybe_dummy_run_with_lora=lambda *args, **kwargs: nullcontext(),
+        execute_model=execute_model,
+    )
+
+    with pytest.raises(_StopDummyRun):
+        GPUModelRunner._dummy_run(
+            runner,
+            num_tokens=2048,
+            skip_attn=skip_attn,
+            is_profile=is_profile,
+        )
+
+    return captured["scheduler_output"]
+
+
+def test_default_speculator_query_width_is_one():
+    assert _DefaultSpeculator().num_query_per_req == 1
+
+
+@pytest.mark.parametrize("speculator_cls", [DFlashSpeculator, DSparkSpeculator])
+@pytest.mark.parametrize(
+    ("is_profile", "skip_attn"),
+    [(False, False), (True, False), (True, True)],
+)
+def test_dummy_run_caps_parallel_draft_requests(speculator_cls, is_profile, skip_attn):
+    speculator = speculator_cls.__new__(speculator_cls)
+    speculator.num_query_per_req = 5
+
+    scheduler_output = _run_dummy_batch(
+        speculator,
+        is_profile=is_profile,
+        skip_attn=skip_attn,
+    )
+    assert len(scheduler_output.num_scheduled_tokens) == 409
+    assert scheduler_output.total_num_scheduled_tokens == 2048
+
+
+def test_dummy_run_keeps_already_fitting_parallel_draft_batch():
+    speculator = DFlashSpeculator.__new__(DFlashSpeculator)
+    speculator.num_query_per_req = 5
+
+    scheduler_output = _run_dummy_batch(speculator, max_num_reqs=128)
+    assert len(scheduler_output.num_scheduled_tokens) == 128
+
+
+def test_dummy_run_keeps_default_speculator_batch():
+    scheduler_output = _run_dummy_batch(_DefaultSpeculator())
+    assert len(scheduler_output.num_scheduled_tokens) == 512

vllm/v1/worker/gpu/model_runner.py

@@ -790,6 +790,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             num_tokens = max(num_tokens, self.decode_query_len)
             num_reqs = num_tokens // self.decode_query_len
             assert num_tokens % self.decode_query_len == 0
+        elif self.speculator is not None:
+            num_reqs = min(
+                num_reqs,
+                self.max_num_tokens // self.speculator.num_query_per_req,
+            )
         # Distribute the remainder evenly so no dummy request exceeds
         # ceil(num_tokens / num_reqs) <= max_model_len tokens.
         num_tokens_per_request = [

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -57,6 +57,8 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
 
 
 class BaseSpeculator(ABC):
+    num_query_per_req: int = 1
+
     @abstractmethod
     def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
         pass