| 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 |