| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The V2 warmup must reserve the KV blocks the scheduler would reserve. |
| 4 | |
| 5 | `KVCacheManager.allocate_slots` sizes every allocation for |
| 6 | `num_computed + num_scheduled + num_lookahead_tokens`, because the speculator |
| 7 | writes KV past the token range the target model was scheduled for. The warmup |
| 8 | hand-builds its `SchedulerOutput`s, so it has to reserve the same blocks. |
| 9 | """ |
| 10 | |
| 11 | from types import SimpleNamespace |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | from tests.v1.core.test_prefix_caching import make_kv_cache_manager |
| 17 | from tests.v1.core.utils import create_requests |
| 18 | from vllm.config.speculative import SpeculativeConfig |
| 19 | from vllm.config.vllm import VllmConfig |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.utils.math_utils import cdiv |
| 22 | from vllm.v1.kv_cache_interface import ( |
| 23 | CircularBufferSpec, |
| 24 | FullAttentionSpec, |
| 25 | KVCacheConfig, |
| 26 | KVCacheGroupSpec, |
| 27 | MambaSpec, |
| 28 | ) |
| 29 | from vllm.v1.worker.gpu.warmup import ( |
| 30 | _reserved_block_count, |
| 31 | run_mixed_prefill_decode_warmup, |
| 32 | warmup_kernels, |
| 33 | ) |
| 34 | |
| 35 | BLOCK_SIZE = 16 |
| 36 | MAX_MODEL_LEN = 1024 |
| 37 | NUM_SPEC_STEPS = 3 |
| 38 | |
| 39 | # `warmup_kernels` ends on `torch.accelerator.synchronize()`. |
| 40 | pytestmark = pytest.mark.skipif( |
| 41 | not current_platform.is_cuda(), reason="warmup synchronizes on the accelerator" |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | def _attention_group() -> KVCacheGroupSpec: |
| 46 | return KVCacheGroupSpec( |
| 47 | ["layer"], |
| 48 | FullAttentionSpec( |
| 49 | block_size=BLOCK_SIZE, |
| 50 | num_kv_heads=1, |
| 51 | head_size=1, |
| 52 | dtype=torch.float32, |
| 53 | ), |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def _circular_group() -> KVCacheGroupSpec: |
| 58 | return KVCacheGroupSpec( |
| 59 | ["circular"], |
| 60 | CircularBufferSpec( |
| 61 | block_size=BLOCK_SIZE, |
| 62 | num_kv_heads=1, |
| 63 | head_size=1, |
| 64 | dtype=torch.float32, |
| 65 | ), |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | def _mamba_group(mamba_cache_mode: str) -> KVCacheGroupSpec: |
| 70 | # Name carries the mode so several groups can coexist in one config. |
| 71 | return KVCacheGroupSpec( |
| 72 | [f"mamba_{mamba_cache_mode}"], |
| 73 | MambaSpec( |
| 74 | block_size=BLOCK_SIZE, |
| 75 | shapes=((1,),), |
| 76 | dtypes=(torch.float32,), |
| 77 | mamba_cache_mode=mamba_cache_mode, |
| 78 | num_speculative_blocks=NUM_SPEC_STEPS, |
| 79 | ), |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | def _make_runner( |
| 84 | kv_cache_groups: list[KVCacheGroupSpec], |
| 85 | num_lookahead_tokens: int, |
| 86 | num_spec_steps: int = NUM_SPEC_STEPS, |
| 87 | ) -> SimpleNamespace: |
| 88 | """Stub model runner exposing only what the warmup entry points read.""" |
| 89 | return SimpleNamespace( |
| 90 | num_speculative_steps=num_spec_steps, |
| 91 | adaptive_verification=None, |
| 92 | rejection_sampler=None, |
| 93 | decode_query_len=num_spec_steps + 1, |
| 94 | is_pooling_model=False, |
| 95 | is_encoder_decoder=False, |
| 96 | is_last_pp_rank=True, |
| 97 | max_num_reqs=4, |
| 98 | max_model_len=MAX_MODEL_LEN, |
| 99 | model_config=SimpleNamespace(get_vocab_size=lambda: 64), |
| 100 | model_state=SimpleNamespace(max_encoder_len=0), |
| 101 | scheduler_config=SimpleNamespace(max_num_seqs=4, max_num_batched_tokens=2048), |
| 102 | kv_cache_config=SimpleNamespace( |
| 103 | kv_cache_groups=kv_cache_groups, num_blocks=1024 |
| 104 | ), |
| 105 | vllm_config=SimpleNamespace( |
| 106 | num_lookahead_tokens=num_lookahead_tokens, is_mm_encoder_only=False |
| 107 | ), |
| 108 | kv_block_zeroer=None, |
| 109 | kv_connector=SimpleNamespace(set_disabled=lambda disabled: None), |
| 110 | ) |
| 111 | |
| 112 | |
| 113 | class _StepRecorder: |
| 114 | """Rebuilds each request's per-group block holdings from the warmup steps.""" |
| 115 | |
| 116 | def __init__(self) -> None: |
| 117 | # (blocks held per group, num_computed_tokens, num_scheduled_tokens) |
| 118 | self.steps: list[tuple[list[int], int, int]] = [] |
| 119 | self._held: dict[str, list[int]] = {} |
| 120 | |
| 121 | def execute_model(self, scheduler_output) -> None: |
| 122 | for new_req in scheduler_output.scheduled_new_reqs: |
| 123 | self._held[new_req.req_id] = [len(ids) for ids in new_req.block_ids] |
| 124 | self._record(new_req.req_id, new_req.num_computed_tokens, scheduler_output) |
| 125 | cached = scheduler_output.scheduled_cached_reqs |
| 126 | for i, req_id in enumerate(cached.req_ids): |
| 127 | new_block_ids = cached.new_block_ids[i] |
| 128 | if new_block_ids is not None: |
| 129 | self._held[req_id] = [ |
| 130 | held + len(ids) |
| 131 | for held, ids in zip(self._held[req_id], new_block_ids) |
| 132 | ] |
| 133 | self._record(req_id, cached.num_computed_tokens[i], scheduler_output) |
| 134 | |
| 135 | def _record(self, req_id: str, num_computed: int, scheduler_output) -> None: |
| 136 | self.steps.append( |
| 137 | ( |
| 138 | list(self._held[req_id]), |
| 139 | num_computed, |
| 140 | scheduler_output.num_scheduled_tokens[req_id], |
| 141 | ) |
| 142 | ) |
| 143 | |
| 144 | def sample_tokens(self, grammar_output=None) -> None: |
| 145 | return None |
| 146 | |
| 147 | |
| 148 | def _assert_covers_lookahead( |
| 149 | steps: list[tuple[list[int], int, int]], num_lookahead_tokens: int |
| 150 | ) -> None: |
| 151 | assert steps, "warmup ran no steps" |
| 152 | for num_blocks, num_computed, num_scheduled in steps: |
| 153 | num_tokens = min( |
| 154 | num_computed + num_scheduled + num_lookahead_tokens, MAX_MODEL_LEN |
| 155 | ) |
| 156 | assert num_blocks[0] >= cdiv(num_tokens, BLOCK_SIZE), ( |
| 157 | f"{num_blocks[0]} blocks for {num_computed}+{num_scheduled} tokens " |
| 158 | f"and {num_lookahead_tokens} lookahead tokens" |
| 159 | ) |
| 160 | |
| 161 | |
| 162 | # 0 covers eagle / MTP / draft models, 1 covers DFlash's extra in-fill query. |
| 163 | @pytest.mark.parametrize("extra_lookahead", [0, 1]) |
| 164 | @pytest.mark.parametrize("num_spec_steps", [2, 3, 5, 7]) |
| 165 | def test_warmup_kernels_reserves_lookahead_blocks(num_spec_steps, extra_lookahead): |
| 166 | num_lookahead_tokens = num_spec_steps + extra_lookahead |
| 167 | recorder = _StepRecorder() |
| 168 | |
| 169 | warmup_kernels( |
| 170 | _make_runner([_attention_group()], num_lookahead_tokens, num_spec_steps), |
| 171 | recorder.execute_model, |
| 172 | recorder.sample_tokens, |
| 173 | ) |
| 174 | |
| 175 | _assert_covers_lookahead(recorder.steps, num_lookahead_tokens) |
| 176 | |
| 177 | |
| 178 | def test_mixed_warmup_reserves_lookahead_blocks(): |
| 179 | num_lookahead_tokens = NUM_SPEC_STEPS + 1 |
| 180 | recorder = _StepRecorder() |
| 181 | |
| 182 | assert run_mixed_prefill_decode_warmup( |
| 183 | _make_runner([_attention_group()], num_lookahead_tokens), |
| 184 | worker_execute_model=recorder.execute_model, |
| 185 | worker_sample_tokens=recorder.sample_tokens, |
| 186 | num_tokens=128, |
| 187 | ) |
| 188 | |
| 189 | _assert_covers_lookahead(recorder.steps, num_lookahead_tokens) |
| 190 | |
| 191 | |
| 192 | @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) |
| 193 | def test_warmup_reserves_mamba_speculative_blocks(mamba_cache_mode): |
| 194 | """Mamba groups hold the running-state block plus the speculative tail. |
| 195 | |
| 196 | `MambaManager` reserves `num_speculative_blocks` blocks past the token |
| 197 | range in every cache mode, and the mamba kernels read all |
| 198 | `1 + num_speculative_blocks` of those block-table columns (see |
| 199 | `mamba_get_block_table_tensor`). |
| 200 | """ |
| 201 | recorder = _StepRecorder() |
| 202 | |
| 203 | warmup_kernels( |
| 204 | _make_runner( |
| 205 | [_attention_group(), _mamba_group(mamba_cache_mode)], NUM_SPEC_STEPS |
| 206 | ), |
| 207 | recorder.execute_model, |
| 208 | recorder.sample_tokens, |
| 209 | ) |
| 210 | |
| 211 | assert recorder.steps, "warmup ran no steps" |
| 212 | for num_blocks, num_computed, num_scheduled in recorder.steps: |
| 213 | # `MambaManager` drops the lookahead tokens in align mode to keep the |
| 214 | # allocation block-aligned, and keeps them otherwise. Pin the token |
| 215 | # range and the speculative tail separately, so an implementation that |
| 216 | # returned a flat `1 + num_speculative_blocks` would fail. |
| 217 | if mamba_cache_mode == "align": |
| 218 | # Align mode sizes from the uncapped main-model range. |
| 219 | lookahead, num_tokens = 0, num_computed + num_scheduled |
| 220 | else: |
| 221 | lookahead = NUM_SPEC_STEPS |
| 222 | num_tokens = min(num_computed + num_scheduled + lookahead, MAX_MODEL_LEN) |
| 223 | assert num_blocks[1] == cdiv(num_tokens, BLOCK_SIZE) + NUM_SPEC_STEPS, ( |
| 224 | f"{mamba_cache_mode}: {num_blocks[1]} mamba blocks for " |
| 225 | f"{num_computed}+{num_scheduled} tokens and {lookahead} lookahead" |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | def _hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: |
| 230 | """Attention, circular, and Mamba groups exercise every reservation branch.""" |
| 231 | return KVCacheConfig( |
| 232 | num_blocks=num_blocks, |
| 233 | kv_cache_tensors=[], |
| 234 | kv_cache_groups=[ |
| 235 | _attention_group(), |
| 236 | _circular_group(), |
| 237 | _mamba_group("none"), |
| 238 | _mamba_group("align"), |
| 239 | ], |
| 240 | ) |
| 241 | |
| 242 | |
| 243 | def test_reserved_block_count_matches_real_kv_cache_manager(): |
| 244 | """`_reserved_block_count` must predict exactly what the real |
| 245 | `KVCacheManager.allocate_slots` consumes, for every group, at every step |
| 246 | of a warmup-realistic trajectory: a prefill followed by decode steps that |
| 247 | mix spec and non-spec shapes, all at the fixed `num_lookahead_tokens` |
| 248 | `_warmup_block_counter` binds once per model runner. |
| 249 | |
| 250 | The existing tests above pin this arithmetic against a hand-built |
| 251 | `_StepRecorder` fake, which only proves the fake and the production |
| 252 | function agree with each other. Driving `allocate_slots` itself is the |
| 253 | only way to catch a real divergence from the allocator warmup has to |
| 254 | match. |
| 255 | |
| 256 | Since `num_computed_tokens` only grows and `num_lookahead_tokens` is held |
| 257 | fixed across the trajectory, the real manager's held-block count is |
| 258 | non-decreasing in lockstep with the prediction, so equality (not just |
| 259 | the reservation lower bound) holds at every step -- exactly the shape of |
| 260 | trajectory `warmup_kernels` itself drives the allocator through. |
| 261 | """ |
| 262 | num_lookahead_tokens = NUM_SPEC_STEPS + 1 |
| 263 | decode_query_len = NUM_SPEC_STEPS + 1 |
| 264 | |
| 265 | kv_cache_config = _hybrid_kv_cache_config(num_blocks=256) |
| 266 | specs = [g.kv_cache_spec for g in kv_cache_config.kv_cache_groups] |
| 267 | manager = make_kv_cache_manager( |
| 268 | kv_cache_config, |
| 269 | max_model_len=MAX_MODEL_LEN, |
| 270 | hash_block_size=BLOCK_SIZE, |
| 271 | enable_caching=False, |
| 272 | ) |
| 273 | (request,) = create_requests( |
| 274 | num_requests=1, num_tokens=decode_query_len + 1, block_size=BLOCK_SIZE |
| 275 | ) |
| 276 | |
| 277 | def _step(num_new_tokens: int) -> None: |
| 278 | computed_blocks, num_new_computed, _ = manager.get_computed_blocks(request) |
| 279 | blocks = manager.allocate_slots( |
| 280 | request, |
| 281 | num_new_tokens, |
| 282 | num_new_computed, |
| 283 | computed_blocks, |
| 284 | num_lookahead_tokens=num_lookahead_tokens, |
| 285 | ) |
| 286 | assert blocks is not None, "block pool exhausted" |
| 287 | held = manager.get_block_ids(request.request_id) |
| 288 | num_tokens = request.num_computed_tokens + num_new_tokens |
| 289 | for spec, group_held in zip(specs, held): |
| 290 | expected = _reserved_block_count( |
| 291 | num_tokens, |
| 292 | spec, |
| 293 | num_lookahead_tokens=num_lookahead_tokens, |
| 294 | max_model_len=MAX_MODEL_LEN, |
| 295 | max_encoder_len=0, |
| 296 | ) |
| 297 | assert len(group_held) == expected, ( |
| 298 | f"{type(spec).__name__} " |
| 299 | f"mode={getattr(spec, 'mamba_cache_mode', None)}: real manager " |
| 300 | f"holds {len(group_held)} blocks for {num_tokens} tokens, " |
| 301 | f"_reserved_block_count predicts {expected}" |
| 302 | ) |
| 303 | request.num_computed_tokens = num_tokens |
| 304 | |
| 305 | # Prefill, a spec decode step, a no-draft decode step, and another spec |
| 306 | # decode step: the same shapes `warmup_kernels` drives its decode steps |
| 307 | # through (see `decode_steps` in `warmup.py`). |
| 308 | _step(decode_query_len + 1) |
| 309 | _step(decode_query_len) |
| 310 | _step(1) |
| 311 | _step(decode_query_len) |
| 312 | |
| 313 | manager.free(request) |
| 314 | |
| 315 | |
| 316 | @pytest.mark.parametrize( |
| 317 | ("method", "expected"), |
| 318 | [ |
| 319 | ("eagle", NUM_SPEC_STEPS), |
| 320 | ("eagle3", NUM_SPEC_STEPS), |
| 321 | ("mtp", NUM_SPEC_STEPS), |
| 322 | ("dspark", NUM_SPEC_STEPS), |
| 323 | ("draft_model", NUM_SPEC_STEPS), |
| 324 | # DFlash's in-fill decoding adds a query for the last sampled token. |
| 325 | ("dflash", NUM_SPEC_STEPS + 1), |
| 326 | ("ngram", 0), |
| 327 | ("ngram_gpu", 0), |
| 328 | ("medusa", 0), |
| 329 | ("mlp_speculator", 0), |
| 330 | ("suffix", 0), |
| 331 | ("extract_hidden_states", 0), |
| 332 | ], |
| 333 | ) |
| 334 | def test_num_lookahead_tokens_per_method(method: str, expected: int): |
| 335 | """`VllmConfig.num_lookahead_tokens` is the single source of the reservation. |
| 336 | |
| 337 | Both the scheduler and the warmup read it, so a wrong answer here silently |
| 338 | changes how many blocks every request holds. Exercises the real property |
| 339 | and the real `SpeculativeConfig` predicates; the config is built without |
| 340 | `__post_init__` because the speculative methods otherwise require a draft |
| 341 | model to be resolvable. |
| 342 | """ |
| 343 | |
| 344 | class _Config: |
| 345 | speculative_config: SpeculativeConfig | None = None |
| 346 | diffusion_config = None |
| 347 | num_speculative_tokens = VllmConfig.num_speculative_tokens |
| 348 | num_lookahead_tokens = VllmConfig.num_lookahead_tokens |
| 349 | |
| 350 | speculative_config = object.__new__(SpeculativeConfig) |
| 351 | object.__setattr__(speculative_config, "method", method) |
| 352 | object.__setattr__(speculative_config, "num_speculative_tokens", NUM_SPEC_STEPS) |
| 353 | |
| 354 | config = _Config() |
| 355 | config.speculative_config = speculative_config |
| 356 | |
| 357 | assert config.num_lookahead_tokens == expected |
| 358 | |
| 359 | |
| 360 | def test_num_lookahead_tokens_without_speculation(): |
| 361 | class _Config: |
| 362 | speculative_config: SpeculativeConfig | None = None |
| 363 | diffusion_config = None |
| 364 | num_speculative_tokens = VllmConfig.num_speculative_tokens |
| 365 | num_lookahead_tokens = VllmConfig.num_lookahead_tokens |
| 366 | |
| 367 | config = _Config() |
| 368 | |
| 369 | assert config.num_lookahead_tokens == 0 |