diff --git a/tests/v1/spec_decode/test_dflash_profile.py b/tests/v1/spec_decode/test_dflash_profile.py new file mode 100644 index 0000000000..3b6f6061eb --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8be523d3de..a5e5ac2497 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 = [ diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 244fa8bc18..eef1704ac1 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/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