| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Regression tests for the Dynamic SD batch-size schedule helpers.""" |
| 4 | |
| 5 | import logging |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from tests.v1.core.utils import create_requests, create_scheduler |
| 10 | from vllm.v1.core.sched.scheduler import Scheduler |
| 11 | from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup |
| 12 | from vllm.v1.structured_output import StructuredOutputManager |
| 13 | |
| 14 | |
| 15 | def _make_lookup( |
| 16 | num_speculative_tokens_per_batch_size: list[tuple[int, int, int]], |
| 17 | *, |
| 18 | max_batch_size: int = 256, |
| 19 | runtime_num_speculative_tokens: int = 3, |
| 20 | ) -> list[int]: |
| 21 | return build_dynamic_sd_schedule_lookup( |
| 22 | num_speculative_tokens_per_batch_size=num_speculative_tokens_per_batch_size, |
| 23 | vllm_max_batch_size=max_batch_size, |
| 24 | vllm_num_speculative_tokens=runtime_num_speculative_tokens, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | def _make_scheduler_with_dynamic_sd( |
| 29 | schedule: list[tuple[int, int, int]], |
| 30 | *, |
| 31 | max_num_seqs: int = 16, |
| 32 | max_num_batched_tokens: int = 8192, |
| 33 | runtime_num_speculative_tokens: int = 3, |
| 34 | ) -> Scheduler: |
| 35 | base_scheduler = create_scheduler( |
| 36 | max_num_seqs=max_num_seqs, |
| 37 | max_num_batched_tokens=max_num_batched_tokens, |
| 38 | num_speculative_tokens=runtime_num_speculative_tokens, |
| 39 | ) |
| 40 | |
| 41 | speculative_config = base_scheduler.vllm_config.speculative_config |
| 42 | assert speculative_config is not None |
| 43 | speculative_config.num_speculative_tokens_per_batch_size = schedule |
| 44 | |
| 45 | return Scheduler( |
| 46 | vllm_config=base_scheduler.vllm_config, |
| 47 | kv_cache_config=base_scheduler.kv_cache_config, |
| 48 | block_size=base_scheduler.block_size, |
| 49 | log_stats=True, |
| 50 | structured_output_manager=StructuredOutputManager(base_scheduler.vllm_config), |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | def _add_requests_and_schedule( |
| 55 | scheduler: Scheduler, num_requests: int, *, num_tokens: int = 10 |
| 56 | ): |
| 57 | requests = create_requests(num_requests=num_requests, num_tokens=num_tokens) |
| 58 | for request in requests: |
| 59 | scheduler.add_request(request) |
| 60 | return scheduler.schedule() |
| 61 | |
| 62 | |
| 63 | def test_dynamic_sd_uses_batch_size_schedule(): |
| 64 | dynamic_sd_lookup = _make_lookup( |
| 65 | [ |
| 66 | (1, 16, 3), |
| 67 | (32, 128, 2), |
| 68 | (256, 2048, 0), |
| 69 | ] |
| 70 | ) |
| 71 | |
| 72 | assert dynamic_sd_lookup[1] == 3 |
| 73 | assert dynamic_sd_lookup[16] == 3 |
| 74 | assert dynamic_sd_lookup[17] == 3 |
| 75 | assert dynamic_sd_lookup[31] == 3 |
| 76 | assert dynamic_sd_lookup[32] == 2 |
| 77 | assert dynamic_sd_lookup[128] == 2 |
| 78 | assert dynamic_sd_lookup[129] == 2 |
| 79 | assert dynamic_sd_lookup[255] == 2 |
| 80 | assert dynamic_sd_lookup[256] == 0 |
| 81 | |
| 82 | |
| 83 | def test_dynamic_sd_requires_schedule_starting_at_batch_size_one(): |
| 84 | with pytest.raises(ValueError, match="must start at 1"): |
| 85 | _make_lookup([(2, 16, 3)]) |
| 86 | |
| 87 | |
| 88 | def test_dynamic_sd_clamps_k_to_runtime_max(): |
| 89 | dynamic_sd_lookup = _make_lookup( |
| 90 | [(1, 256, 4)], |
| 91 | runtime_num_speculative_tokens=3, |
| 92 | ) |
| 93 | |
| 94 | assert dynamic_sd_lookup[1] == 3 |
| 95 | assert dynamic_sd_lookup[256] == 3 |
| 96 | |
| 97 | |
| 98 | def test_dynamic_sd_rejects_invalid_schedule_entry(): |
| 99 | with pytest.raises(ValueError, match="3-item sequence"): |
| 100 | _make_lookup([(1, 16, 3), (32, 64)]) # type: ignore[list-item] |
| 101 | |
| 102 | |
| 103 | def test_dynamic_sd_rejects_overlapping_ranges(): |
| 104 | with pytest.raises(ValueError, match="non-overlapping and sorted"): |
| 105 | _make_lookup([(1, 16, 3), (16, 32, 2)]) |
| 106 | |
| 107 | |
| 108 | def test_dynamic_sd_rejects_negative_k(): |
| 109 | with pytest.raises(ValueError, match="values must be >= 0"): |
| 110 | _make_lookup([(1, 16, -1)]) |
| 111 | |
| 112 | |
| 113 | def test_dynamic_sd_rejects_empty_schedule(): |
| 114 | with pytest.raises(ValueError, match="must not be empty"): |
| 115 | _make_lookup([]) |
| 116 | |
| 117 | |
| 118 | def test_dynamic_sd_requires_schedule_config(): |
| 119 | with pytest.raises( |
| 120 | ValueError, match="num_speculative_tokens_per_batch_size is required" |
| 121 | ): |
| 122 | build_dynamic_sd_schedule_lookup( |
| 123 | None, |
| 124 | vllm_max_batch_size=256, |
| 125 | vllm_num_speculative_tokens=3, |
| 126 | ) |
| 127 | |
| 128 | |
| 129 | def test_dynamic_sd_lookup_rejects_invalid_batch_size_queries(): |
| 130 | dynamic_sd_lookup = _make_lookup([(1, 256, 3)]) |
| 131 | |
| 132 | assert dynamic_sd_lookup[0] == 0 |
| 133 | with pytest.raises(IndexError): |
| 134 | _ = dynamic_sd_lookup[257] |
| 135 | |
| 136 | |
| 137 | def test_scheduler_initializes_dynamic_sd_lookup_from_speculative_config(): |
| 138 | scheduler = _make_scheduler_with_dynamic_sd( |
| 139 | [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], |
| 140 | runtime_num_speculative_tokens=3, |
| 141 | ) |
| 142 | |
| 143 | assert scheduler.dynamic_sd_lookup is not None |
| 144 | assert scheduler.num_spec_tokens == 3 |
| 145 | |
| 146 | |
| 147 | def test_scheduler_uses_dsd_k_based_on_number_of_scheduled_requests(): |
| 148 | test_cases = [ |
| 149 | (4, 3), |
| 150 | (64, 2), |
| 151 | (256, 0), |
| 152 | ] |
| 153 | |
| 154 | for num_requests, expected_k in test_cases: |
| 155 | scheduler = _make_scheduler_with_dynamic_sd( |
| 156 | [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], |
| 157 | max_num_seqs=num_requests, |
| 158 | max_num_batched_tokens=num_requests * 10, |
| 159 | runtime_num_speculative_tokens=3, |
| 160 | ) |
| 161 | output = _add_requests_and_schedule(scheduler, num_requests) |
| 162 | |
| 163 | assert len(output.num_scheduled_tokens) == num_requests |
| 164 | assert output.num_spec_tokens_to_schedule == expected_k |
| 165 | |
| 166 | |
| 167 | def test_scheduler_clamps_dsd_k_to_runtime_num_speculative_tokens(): |
| 168 | scheduler = _make_scheduler_with_dynamic_sd( |
| 169 | [(1, 256, 5)], |
| 170 | max_num_seqs=16, |
| 171 | max_num_batched_tokens=160, |
| 172 | runtime_num_speculative_tokens=3, |
| 173 | ) |
| 174 | output = _add_requests_and_schedule(scheduler, 16) |
| 175 | |
| 176 | assert len(output.num_scheduled_tokens) == 16 |
| 177 | assert output.num_spec_tokens_to_schedule == 3 |
| 178 | |
| 179 | |
| 180 | def test_scheduler_falls_back_to_static_k_when_dsd_not_configured(): |
| 181 | scheduler = create_scheduler( |
| 182 | max_num_seqs=4, |
| 183 | max_num_batched_tokens=40, |
| 184 | num_speculative_tokens=3, |
| 185 | ) |
| 186 | output = _add_requests_and_schedule(scheduler, 4) |
| 187 | |
| 188 | assert scheduler.dynamic_sd_lookup is None |
| 189 | assert output.num_spec_tokens_to_schedule == 3 |
| 190 | |
| 191 | |
| 192 | def test_dynamic_sd_is_disabled_with_data_parallel(caplog_vllm): |
| 193 | with caplog_vllm.at_level(logging.WARNING, logger="vllm"): |
| 194 | scheduler = create_scheduler( |
| 195 | max_num_seqs=256, |
| 196 | max_num_batched_tokens=2560, |
| 197 | num_speculative_tokens=3, |
| 198 | num_speculative_tokens_per_batch_size=[ |
| 199 | (1, 16, 3), |
| 200 | (64, 128, 2), |
| 201 | (256, 4096, 0), |
| 202 | ], |
| 203 | data_parallel_size=2, |
| 204 | ) |
| 205 | |
| 206 | speculative_config = scheduler.vllm_config.speculative_config |
| 207 | assert speculative_config is not None |
| 208 | assert speculative_config.num_speculative_tokens_per_batch_size is None |
| 209 | assert scheduler.dynamic_sd_lookup is None |
| 210 | assert "Dynamic speculative decoding is not supported with data parallelism" in ( |
| 211 | caplog_vllm.text |
| 212 | ) |
| 213 | |
| 214 | output = _add_requests_and_schedule(scheduler, 256) |
| 215 | assert len(output.num_scheduled_tokens) == 256 |
| 216 | assert output.num_spec_tokens_to_schedule == 3 |
| 217 | |
| 218 | |
| 219 | def test_scheduler_uses_static_k_when_no_requests_are_scheduled(): |
| 220 | scheduler = _make_scheduler_with_dynamic_sd( |
| 221 | [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], |
| 222 | runtime_num_speculative_tokens=3, |
| 223 | ) |
| 224 | output = scheduler.schedule() |
| 225 | |
| 226 | assert len(output.num_scheduled_tokens) == 0 |
| 227 | assert output.num_spec_tokens_to_schedule == 3 |
| 228 | |
| 229 | |
| 230 | def test_scheduler_rejects_bad_dsd_config_at_construction(): |
| 231 | with pytest.raises(ValueError, match="must start at 1"): |
| 232 | _make_scheduler_with_dynamic_sd([(2, 16, 3)]) |
| 233 | |
| 234 | |
| 235 | def test_scheduler_passes_max_num_seqs_as_dsd_runtime_batch_limit(): |
| 236 | scheduler = _make_scheduler_with_dynamic_sd( |
| 237 | [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], |
| 238 | max_num_seqs=16, |
| 239 | max_num_batched_tokens=160, |
| 240 | runtime_num_speculative_tokens=3, |
| 241 | ) |
| 242 | output = _add_requests_and_schedule(scheduler, 16) |
| 243 | |
| 244 | assert scheduler.dynamic_sd_lookup is not None |
| 245 | assert len(scheduler.dynamic_sd_lookup) == 17 |
| 246 | assert len(output.num_scheduled_tokens) == 16 |
| 247 | assert output.num_spec_tokens_to_schedule == 3 |