| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for SimpleCPUOffloadScheduler.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import logging |
| 8 | from dataclasses import dataclass |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm import SamplingParams |
| 14 | from vllm.config import ( |
| 15 | CacheConfig, |
| 16 | DeviceConfig, |
| 17 | KVTransferConfig, |
| 18 | ModelConfig, |
| 19 | SchedulerConfig, |
| 20 | VllmConfig, |
| 21 | ) |
| 22 | from vllm.config.cache import MambaCacheMode |
| 23 | from vllm.config.kv_events import KVEventsConfig |
| 24 | from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( |
| 25 | SimpleCPUOffloadConnector, |
| 26 | ) |
| 27 | from vllm.utils.hashing import sha256 |
| 28 | from vllm.v1.core.block_pool import BlockPool |
| 29 | from vllm.v1.core.kv_cache_manager import KVCacheBlocks |
| 30 | from vllm.v1.core.kv_cache_utils import ( |
| 31 | get_request_block_hasher, |
| 32 | init_none_hash, |
| 33 | make_block_hash_with_group_id, |
| 34 | ) |
| 35 | from vllm.v1.core.sched.output import ( |
| 36 | CachedRequestData, |
| 37 | KVConnectorBlockState, |
| 38 | NewRequestData, |
| 39 | SchedulerOutput, |
| 40 | ) |
| 41 | from vllm.v1.core.single_type_kv_cache_manager import ( |
| 42 | register_all_kvcache_specs, |
| 43 | ) |
| 44 | from vllm.v1.kv_cache_interface import ( |
| 45 | CircularBufferSpec, |
| 46 | FullAttentionSpec, |
| 47 | KpoolTailSpec, |
| 48 | KVCacheConfig, |
| 49 | KVCacheGroupSpec, |
| 50 | KVCacheTensor, |
| 51 | MambaSpec, |
| 52 | SlidingWindowSpec, |
| 53 | ) |
| 54 | from vllm.v1.outputs import KVConnectorOutput |
| 55 | from vllm.v1.request import Request |
| 56 | from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler |
| 57 | from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadWorkerMetadata |
| 58 | |
| 59 | pytestmark = pytest.mark.skip_global_cleanup |
| 60 | |
| 61 | |
| 62 | # --------------------------------------------------------------------------- |
| 63 | # Constants |
| 64 | # --------------------------------------------------------------------------- |
| 65 | BLOCK_SIZE = 16 |
| 66 | HEAD_SIZE = 16 |
| 67 | NUM_KV_HEADS = 1 |
| 68 | DTYPE = torch.float16 |
| 69 | # bytes per block per tensor: |
| 70 | # block_size * num_kv_heads * head_size * 2 (K+V) * element_size |
| 71 | _BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize |
| 72 | |
| 73 | # Ensure none_hash is initialized once |
| 74 | init_none_hash(sha256) |
| 75 | |
| 76 | |
| 77 | # --------------------------------------------------------------------------- |
| 78 | # Helpers |
| 79 | # --------------------------------------------------------------------------- |
| 80 | |
| 81 | |
| 82 | def _make_kv_cache_config( |
| 83 | num_blocks: int, |
| 84 | num_groups: int = 1, |
| 85 | ) -> KVCacheConfig: |
| 86 | """Build a KVCacheConfig with non-empty kv_cache_tensors.""" |
| 87 | groups = [] |
| 88 | tensors = [] |
| 89 | register_all_kvcache_specs( |
| 90 | vllm_config=None |
| 91 | ) # Ensure specs are registered for tests |
| 92 | for g in range(num_groups): |
| 93 | layer_names = [f"layer_{g}"] |
| 94 | spec = ( |
| 95 | FullAttentionSpec( |
| 96 | block_size=BLOCK_SIZE, |
| 97 | num_kv_heads=NUM_KV_HEADS, |
| 98 | head_size=HEAD_SIZE, |
| 99 | dtype=DTYPE, |
| 100 | ) |
| 101 | if g == 0 |
| 102 | else SlidingWindowSpec( |
| 103 | block_size=BLOCK_SIZE, |
| 104 | num_kv_heads=NUM_KV_HEADS, |
| 105 | head_size=HEAD_SIZE, |
| 106 | dtype=DTYPE, |
| 107 | sliding_window=BLOCK_SIZE * 4, |
| 108 | ) |
| 109 | ) |
| 110 | groups.append( |
| 111 | KVCacheGroupSpec( |
| 112 | layer_names, |
| 113 | spec, |
| 114 | ) |
| 115 | ) |
| 116 | tensors.append( |
| 117 | KVCacheTensor( |
| 118 | size=_BYTES_PER_BLOCK * num_blocks, |
| 119 | layers=layer_names, |
| 120 | layer_stride=_BYTES_PER_BLOCK * num_blocks, |
| 121 | block_stride=_BYTES_PER_BLOCK, |
| 122 | ) |
| 123 | ) |
| 124 | return KVCacheConfig( |
| 125 | num_blocks=num_blocks, |
| 126 | kv_cache_tensors=tensors, |
| 127 | kv_cache_groups=groups, |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | def _make_scratch_kv_cache_config( |
| 132 | num_blocks: int, scratch_block_size: int = 4 |
| 133 | ) -> KVCacheConfig: |
| 134 | """FullAttention group plus a non-prefix-cacheable scratch group. |
| 135 | |
| 136 | Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request |
| 137 | scratch block of ``index_kpool`` tokens: it is excluded from prefix |
| 138 | caching, so its block size neither divides nor is divided by the |
| 139 | hash block size in general. |
| 140 | """ |
| 141 | register_all_kvcache_specs(vllm_config=None) |
| 142 | fa_config = _make_kv_cache_config(num_blocks, num_groups=1) |
| 143 | scratch_layers = ["layer_scratch"] |
| 144 | scratch_spec = KpoolTailSpec( |
| 145 | block_size=scratch_block_size, |
| 146 | num_kv_heads=2, |
| 147 | head_size=HEAD_SIZE, |
| 148 | head_size_v=0, |
| 149 | dtype=DTYPE, |
| 150 | sliding_window=scratch_block_size, |
| 151 | ) |
| 152 | assert not scratch_spec.prefix_cacheable |
| 153 | scratch_bytes = scratch_spec.page_size_bytes * num_blocks |
| 154 | return KVCacheConfig( |
| 155 | num_blocks=num_blocks, |
| 156 | kv_cache_tensors=fa_config.kv_cache_tensors |
| 157 | + [ |
| 158 | KVCacheTensor( |
| 159 | size=scratch_bytes, |
| 160 | layers=scratch_layers, |
| 161 | layer_stride=scratch_bytes, |
| 162 | block_stride=scratch_spec.page_size_bytes, |
| 163 | ) |
| 164 | ], |
| 165 | kv_cache_groups=fa_config.kv_cache_groups |
| 166 | + [KVCacheGroupSpec(scratch_layers, scratch_spec)], |
| 167 | ) |
| 168 | |
| 169 | |
| 170 | def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig: |
| 171 | """Minimal VllmConfig for scheduler tests (no GPU).""" |
| 172 | model_config = ModelConfig( |
| 173 | model="facebook/opt-125m", |
| 174 | trust_remote_code=True, |
| 175 | dtype="float16", |
| 176 | seed=42, |
| 177 | ) |
| 178 | scheduler_config = SchedulerConfig( |
| 179 | max_num_seqs=16, |
| 180 | max_num_batched_tokens=64, |
| 181 | max_model_len=10000, |
| 182 | enable_chunked_prefill=True, |
| 183 | is_encoder_decoder=False, |
| 184 | ) |
| 185 | cache_config = CacheConfig( |
| 186 | block_size=block_size, |
| 187 | gpu_memory_utilization=0.9, |
| 188 | enable_prefix_caching=True, |
| 189 | ) |
| 190 | kv_transfer_config = KVTransferConfig( |
| 191 | kv_connector="SimpleCPUOffloadConnector", |
| 192 | kv_role="kv_both", |
| 193 | ) |
| 194 | return VllmConfig( |
| 195 | scheduler_config=scheduler_config, |
| 196 | model_config=model_config, |
| 197 | cache_config=cache_config, |
| 198 | kv_transfer_config=kv_transfer_config, |
| 199 | device_config=DeviceConfig("cpu"), |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | @dataclass |
| 204 | class SchedulerFixture: |
| 205 | """Bundle returned by make_scheduler for convenient access.""" |
| 206 | |
| 207 | scheduler: SimpleCPUOffloadScheduler |
| 208 | gpu_block_pool: BlockPool |
| 209 | vllm_config: VllmConfig |
| 210 | kv_cache_config: KVCacheConfig |
| 211 | num_groups: int = 1 |
| 212 | |
| 213 | |
| 214 | def make_scheduler( |
| 215 | num_cpu_blocks: int = 8, |
| 216 | num_gpu_blocks: int = 16, |
| 217 | num_groups: int = 1, |
| 218 | lazy: bool = False, |
| 219 | kv_cache_config: KVCacheConfig | None = None, |
| 220 | ) -> SchedulerFixture: |
| 221 | """Build a SimpleCPUOffloadScheduler with small block pools.""" |
| 222 | if kv_cache_config is None: |
| 223 | kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) |
| 224 | else: |
| 225 | num_groups = len(kv_cache_config.kv_cache_groups) |
| 226 | vllm_config = _make_vllm_config() |
| 227 | cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups |
| 228 | |
| 229 | sched = SimpleCPUOffloadScheduler( |
| 230 | vllm_config=vllm_config, |
| 231 | kv_cache_config=kv_cache_config, |
| 232 | cpu_capacity_bytes=cpu_capacity_bytes, |
| 233 | scheduler_block_size=BLOCK_SIZE, |
| 234 | hash_block_size=BLOCK_SIZE, |
| 235 | lazy_offload=lazy, |
| 236 | ) |
| 237 | |
| 238 | # Build a real GPU block pool and bind it |
| 239 | gpu_block_pool = BlockPool( |
| 240 | num_gpu_blocks=num_gpu_blocks, |
| 241 | enable_caching=True, |
| 242 | hash_block_size=BLOCK_SIZE, |
| 243 | ) |
| 244 | sched.bind_gpu_block_pool(gpu_block_pool) |
| 245 | |
| 246 | return SchedulerFixture( |
| 247 | scheduler=sched, |
| 248 | gpu_block_pool=gpu_block_pool, |
| 249 | vllm_config=vllm_config, |
| 250 | kv_cache_config=kv_cache_config, |
| 251 | num_groups=num_groups, |
| 252 | ) |
| 253 | |
| 254 | |
| 255 | _req_counter = 0 |
| 256 | |
| 257 | |
| 258 | def make_request( |
| 259 | num_blocks: int = 2, |
| 260 | request_id: str | None = None, |
| 261 | extra_tokens: int = 1, |
| 262 | ) -> Request: |
| 263 | """Create a Request with deterministic block hashes.""" |
| 264 | global _req_counter |
| 265 | _req_counter += 1 |
| 266 | if request_id is None: |
| 267 | request_id = f"req-{_req_counter}" |
| 268 | |
| 269 | num_tokens = num_blocks * BLOCK_SIZE + extra_tokens |
| 270 | start = _req_counter * 10000 |
| 271 | prompt_token_ids = list(range(start, start + num_tokens)) |
| 272 | sampling_params = SamplingParams(max_tokens=1) |
| 273 | |
| 274 | req = Request( |
| 275 | request_id=request_id, |
| 276 | prompt_token_ids=prompt_token_ids, |
| 277 | sampling_params=sampling_params, |
| 278 | pooling_params=None, |
| 279 | mm_features=None, |
| 280 | block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256), |
| 281 | ) |
| 282 | return req |
| 283 | |
| 284 | |
| 285 | def make_scheduler_output( |
| 286 | req_id_to_num_tokens: dict[str, int], |
| 287 | *, |
| 288 | new_reqs: dict[str, tuple[list[int], ...]] | None = None, |
| 289 | cached_req_new_blocks: dict[str, tuple[list[int], ...] | None] | None = None, |
| 290 | ) -> SchedulerOutput: |
| 291 | """Build a minimal SchedulerOutput with num_scheduled_tokens. |
| 292 | |
| 293 | Args: |
| 294 | new_reqs: For first-time requests, maps req_id -> block_ids tuple. |
| 295 | These are placed into ``scheduled_new_reqs`` as ``NewRequestData``. |
| 296 | cached_req_new_blocks: For returning (cached) requests, maps |
| 297 | req_id -> new_block_ids (incremental) or None. |
| 298 | These are placed into ``scheduled_cached_reqs``. |
| 299 | |
| 300 | """ |
| 301 | scheduled_new_reqs: list[NewRequestData] = [] |
| 302 | if new_reqs: |
| 303 | for req_id, block_ids in new_reqs.items(): |
| 304 | scheduled_new_reqs.append( |
| 305 | NewRequestData( |
| 306 | req_id=req_id, |
| 307 | prompt_token_ids=None, |
| 308 | mm_features=[], |
| 309 | sampling_params=None, |
| 310 | pooling_params=None, |
| 311 | block_ids=block_ids, |
| 312 | num_computed_tokens=0, |
| 313 | lora_request=None, |
| 314 | ) |
| 315 | ) |
| 316 | |
| 317 | if cached_req_new_blocks: |
| 318 | cached_req_ids = list(cached_req_new_blocks.keys()) |
| 319 | cached_new_block_ids = [cached_req_new_blocks[rid] for rid in cached_req_ids] |
| 320 | cached_reqs = CachedRequestData( |
| 321 | req_ids=cached_req_ids, |
| 322 | resumed_req_ids=set(), |
| 323 | new_token_ids=[[] for _ in cached_req_ids], |
| 324 | all_token_ids={}, |
| 325 | new_block_ids=cached_new_block_ids, |
| 326 | num_computed_tokens=[0] * len(cached_req_ids), |
| 327 | num_output_tokens=[0] * len(cached_req_ids), |
| 328 | ) |
| 329 | else: |
| 330 | cached_reqs = CachedRequestData.make_empty() |
| 331 | |
| 332 | return SchedulerOutput( |
| 333 | scheduled_new_reqs=scheduled_new_reqs, |
| 334 | scheduled_cached_reqs=cached_reqs, |
| 335 | num_scheduled_tokens=req_id_to_num_tokens, |
| 336 | total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()), |
| 337 | scheduled_spec_decode_tokens={}, |
| 338 | scheduled_encoder_inputs={}, |
| 339 | num_common_prefix_blocks=[], |
| 340 | preempted_req_ids=set(), |
| 341 | finished_req_ids=set(), |
| 342 | free_encoder_mm_hashes=[], |
| 343 | ) |
| 344 | |
| 345 | |
| 346 | def simulate_store_completion( |
| 347 | scheduler: SimpleCPUOffloadScheduler, |
| 348 | event_idx: int, |
| 349 | ) -> None: |
| 350 | """Simulate worker reporting a store event completion.""" |
| 351 | output = KVConnectorOutput( |
| 352 | finished_recving=set(), |
| 353 | kv_connector_worker_meta=SimpleCPUOffloadWorkerMetadata( |
| 354 | completed_store_events={event_idx: scheduler._expected_worker_count}, |
| 355 | ), |
| 356 | ) |
| 357 | scheduler.update_connector_output(output) |
| 358 | |
| 359 | |
| 360 | def simulate_load_completion( |
| 361 | scheduler: SimpleCPUOffloadScheduler, |
| 362 | req_ids: set[str], |
| 363 | ) -> None: |
| 364 | """Simulate worker reporting load completions for requests.""" |
| 365 | output = KVConnectorOutput( |
| 366 | finished_sending=set(), |
| 367 | finished_recving=req_ids, |
| 368 | ) |
| 369 | scheduler.update_connector_output(output) |
| 370 | |
| 371 | |
| 372 | def get_cpu_free_blocks(scheduler: SimpleCPUOffloadScheduler) -> int: |
| 373 | """Return number of free CPU blocks.""" |
| 374 | return scheduler.cpu_block_pool.get_num_free_blocks() |
| 375 | |
| 376 | |
| 377 | def _allocate_gpu_blocks( |
| 378 | gpu_block_pool: BlockPool, |
| 379 | request: Request, |
| 380 | num_blocks: int, |
| 381 | group_id: int = 0, |
| 382 | ) -> list: |
| 383 | """Allocate GPU blocks, cache them with hashes, return block list. |
| 384 | |
| 385 | Mimics what KVCacheManager does: allocate blocks from pool, then |
| 386 | register them in the prefix cache via cache_full_blocks so that |
| 387 | re-allocation properly evicts stale hashes. |
| 388 | """ |
| 389 | blocks = gpu_block_pool.get_new_blocks(num_blocks) |
| 390 | num_full = min(num_blocks, len(request.block_hashes)) |
| 391 | if num_full > 0: |
| 392 | gpu_block_pool.cache_full_blocks( |
| 393 | request=request, |
| 394 | blocks=blocks, |
| 395 | num_cached_blocks=0, |
| 396 | num_full_blocks=num_full, |
| 397 | block_size=BLOCK_SIZE, |
| 398 | kv_cache_group_id=group_id, |
| 399 | ) |
| 400 | return blocks |
| 401 | |
| 402 | |
| 403 | def _alloc_and_register( |
| 404 | fix: SchedulerFixture, |
| 405 | request: Request, |
| 406 | num_blocks: int, |
| 407 | *, |
| 408 | confirmed: bool = True, |
| 409 | group_id: int = 0, |
| 410 | ) -> KVCacheBlocks: |
| 411 | """Allocate GPU blocks and return KVCacheBlocks. |
| 412 | |
| 413 | Block IDs are no longer registered in a mock KVCacheManager; instead |
| 414 | tests pass them through ``make_scheduler_output`` so that |
| 415 | ``yield_req_data`` can pick them up. |
| 416 | |
| 417 | If ``confirmed`` is True, advance ``request.num_computed_tokens`` to simulate |
| 418 | the scheduler's ``_update_after_schedule`` from a prior step. |
| 419 | """ |
| 420 | gpu_blocks = _allocate_gpu_blocks( |
| 421 | fix.gpu_block_pool, request, num_blocks, group_id=group_id |
| 422 | ) |
| 423 | kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,)) |
| 424 | if confirmed: |
| 425 | request.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 426 | return kv_blocks |
| 427 | |
| 428 | |
| 429 | # --------------------------------------------------------------------------- |
| 430 | # Test 1a: Eager store-and-load roundtrip |
| 431 | # --------------------------------------------------------------------------- |
| 432 | def test_eager_store_and_load_roundtrip() -> None: |
| 433 | """Eager mode: store blocks on compute, complete store, verify cache hit.""" |
| 434 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 435 | sched = fix.scheduler |
| 436 | |
| 437 | num_blocks = 2 |
| 438 | req = make_request(num_blocks=num_blocks) |
| 439 | |
| 440 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 441 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 442 | block_ids = kv_blocks.get_block_ids() |
| 443 | sched_out = make_scheduler_output( |
| 444 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 445 | new_reqs={req.request_id: block_ids}, |
| 446 | ) |
| 447 | |
| 448 | meta = sched.build_connector_meta(sched_out) |
| 449 | assert meta.store_event >= 0, "Expected a store event to be scheduled" |
| 450 | assert len(meta.store_gpu_blocks) > 0 |
| 451 | assert len(meta.store_cpu_blocks) == len(meta.store_gpu_blocks) |
| 452 | simulate_store_completion(sched, meta.store_event) |
| 453 | |
| 454 | # New request with same tokens should get CPU cache hit |
| 455 | req2 = Request( |
| 456 | request_id="req-eager-load", |
| 457 | prompt_token_ids=req.prompt_token_ids, |
| 458 | sampling_params=req.sampling_params, |
| 459 | pooling_params=None, |
| 460 | mm_features=None, |
| 461 | block_hasher=req._block_hasher, |
| 462 | ) |
| 463 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 464 | # make_request pads num_tokens by +1 beyond the last full block, so the |
| 465 | # manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact. |
| 466 | assert hit_tokens == num_blocks * BLOCK_SIZE |
| 467 | assert is_async is True |
| 468 | |
| 469 | gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) |
| 470 | kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) |
| 471 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 472 | |
| 473 | block_ids2 = kv_blocks2.get_block_ids() |
| 474 | sched_out2 = make_scheduler_output( |
| 475 | {req2.request_id: 1}, |
| 476 | new_reqs={req2.request_id: block_ids2}, |
| 477 | ) |
| 478 | meta2 = sched.build_connector_meta(sched_out2) |
| 479 | assert meta2.load_event >= 0, "Expected a load event to be assigned" |
| 480 | assert len(meta2.load_gpu_blocks) > 0 |
| 481 | assert len(meta2.load_cpu_blocks) == len(meta2.load_gpu_blocks) |
| 482 | |
| 483 | |
| 484 | def test_prompt_logprobs_skip_cpu_cache_lookup() -> None: |
| 485 | """Prompt logprobs require every prompt token to be recomputed.""" |
| 486 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 487 | sched = fix.scheduler |
| 488 | |
| 489 | num_blocks = 2 |
| 490 | req = make_request(num_blocks=num_blocks) |
| 491 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 492 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 493 | sched_out = make_scheduler_output( |
| 494 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 495 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 496 | ) |
| 497 | meta = sched.build_connector_meta(sched_out) |
| 498 | simulate_store_completion(sched, meta.store_event) |
| 499 | |
| 500 | retry_request_id = "req-prompt-logprobs" |
| 501 | cached_request = Request( |
| 502 | request_id=retry_request_id, |
| 503 | prompt_token_ids=req.prompt_token_ids, |
| 504 | sampling_params=req.sampling_params, |
| 505 | pooling_params=None, |
| 506 | mm_features=None, |
| 507 | block_hasher=req._block_hasher, |
| 508 | ) |
| 509 | assert sched.get_num_new_matched_tokens(cached_request, num_computed_tokens=0) == ( |
| 510 | num_blocks * BLOCK_SIZE, |
| 511 | True, |
| 512 | ) |
| 513 | pending_hit = sched._pending_cpu_hits[retry_request_id] |
| 514 | pinned_blocks = [ |
| 515 | block for group in pending_hit[0] for block in group if not block.is_null |
| 516 | ] |
| 517 | assert pinned_blocks |
| 518 | assert all(block.ref_cnt == 1 for block in pinned_blocks) |
| 519 | |
| 520 | prompt_logprobs_request = Request( |
| 521 | request_id=retry_request_id, |
| 522 | prompt_token_ids=req.prompt_token_ids, |
| 523 | sampling_params=SamplingParams(max_tokens=1, prompt_logprobs=1), |
| 524 | pooling_params=None, |
| 525 | mm_features=None, |
| 526 | block_hasher=req._block_hasher, |
| 527 | ) |
| 528 | assert prompt_logprobs_request.skip_reading_prefix_cache |
| 529 | |
| 530 | assert sched.get_num_new_matched_tokens( |
| 531 | prompt_logprobs_request, num_computed_tokens=0 |
| 532 | ) == (0, False) |
| 533 | assert retry_request_id not in sched._pending_cpu_hits |
| 534 | assert all(block.ref_cnt == 0 for block in pinned_blocks) |
| 535 | |
| 536 | |
| 537 | def test_eager_store_preserves_secondary_block_hashes() -> None: |
| 538 | """CPU copies retain fine-grained hashes owned by the GPU block.""" |
| 539 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 540 | sched = fix.scheduler |
| 541 | req = make_request(num_blocks=1) |
| 542 | kv_blocks = _alloc_and_register(fix, req, num_blocks=1) |
| 543 | gpu_block = kv_blocks.blocks[0][0] |
| 544 | primary_hash = gpu_block.block_hash |
| 545 | primary_num_tokens = gpu_block.block_hash_num_tokens |
| 546 | assert primary_hash is not None |
| 547 | |
| 548 | fine_grained_req = Request( |
| 549 | request_id="req-fine-grained-hash", |
| 550 | prompt_token_ids=req.prompt_token_ids, |
| 551 | sampling_params=req.sampling_params, |
| 552 | pooling_params=None, |
| 553 | mm_features=None, |
| 554 | block_hasher=get_request_block_hasher(BLOCK_SIZE // 2, sha256), |
| 555 | ) |
| 556 | secondary_hash = make_block_hash_with_group_id(fine_grained_req.block_hashes[0], 0) |
| 557 | fix.gpu_block_pool._insert_block_hash( |
| 558 | secondary_hash, gpu_block, num_tokens=BLOCK_SIZE // 2 |
| 559 | ) |
| 560 | assert ( |
| 561 | secondary_hash |
| 562 | in fix.gpu_block_pool.cached_block_hashes_by_block[gpu_block.block_id] |
| 563 | ) |
| 564 | |
| 565 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 566 | block_ids = kv_blocks.get_block_ids() |
| 567 | meta = sched.build_connector_meta( |
| 568 | make_scheduler_output( |
| 569 | {req.request_id: BLOCK_SIZE}, |
| 570 | new_reqs={req.request_id: block_ids}, |
| 571 | ) |
| 572 | ) |
| 573 | simulate_store_completion(sched, meta.store_event) |
| 574 | |
| 575 | cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 576 | secondary_hash |
| 577 | ) |
| 578 | assert cpu_block is not None |
| 579 | assert cpu_block.block_hash == primary_hash |
| 580 | assert cpu_block.block_hash_num_tokens == primary_num_tokens |
| 581 | assert ( |
| 582 | secondary_hash |
| 583 | in sched.cpu_block_pool.cached_block_hashes_by_block[cpu_block.block_id] |
| 584 | ) |
| 585 | |
| 586 | |
| 587 | # --------------------------------------------------------------------------- |
| 588 | # Test 1b: Finished eager stores |
| 589 | # --------------------------------------------------------------------------- |
| 590 | def test_finished_eager_stores_are_batched() -> None: |
| 591 | """Finished requests in one step retain, complete, and cache every store.""" |
| 592 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 593 | sched = fix.scheduler |
| 594 | |
| 595 | requests = [make_request(num_blocks=2) for _ in range(2)] |
| 596 | block_ids = [] |
| 597 | for request in requests: |
| 598 | kv_blocks = _alloc_and_register(fix, request, num_blocks=2) |
| 599 | sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0) |
| 600 | block_ids.append(kv_blocks.get_block_ids()) |
| 601 | |
| 602 | for request, request_block_ids in zip(requests, block_ids): |
| 603 | sched.request_finished_all_groups(request, request_block_ids) |
| 604 | |
| 605 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 606 | assert len(meta.store_gpu_blocks) == 4 |
| 607 | assert len(meta.store_cpu_blocks) == 4 |
| 608 | simulate_store_completion(sched, meta.store_event) |
| 609 | |
| 610 | for request in requests: |
| 611 | for block_hash in request.block_hashes[:2]: |
| 612 | cached = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 613 | make_block_hash_with_group_id(block_hash, 0) |
| 614 | ) |
| 615 | assert cached is not None |
| 616 | |
| 617 | request = requests[0] |
| 618 | cached_request = Request( |
| 619 | request_id="req-finished-eager-load", |
| 620 | prompt_token_ids=request.prompt_token_ids, |
| 621 | sampling_params=request.sampling_params, |
| 622 | pooling_params=None, |
| 623 | mm_features=None, |
| 624 | block_hasher=request._block_hasher, |
| 625 | ) |
| 626 | hit_tokens, is_async = sched.get_num_new_matched_tokens( |
| 627 | cached_request, num_computed_tokens=0 |
| 628 | ) |
| 629 | assert hit_tokens == 2 * BLOCK_SIZE |
| 630 | assert is_async is True |
| 631 | |
| 632 | cached_gpu_blocks = fix.gpu_block_pool.get_new_blocks(2) |
| 633 | sched.update_state_after_alloc( |
| 634 | cached_request, |
| 635 | KVCacheBlocks(blocks=(cached_gpu_blocks,)), |
| 636 | num_external_tokens=hit_tokens, |
| 637 | ) |
| 638 | load_meta = sched.build_connector_meta( |
| 639 | make_scheduler_output( |
| 640 | {cached_request.request_id: 1}, |
| 641 | new_reqs={ |
| 642 | cached_request.request_id: ( |
| 643 | [block.block_id for block in cached_gpu_blocks], |
| 644 | ) |
| 645 | }, |
| 646 | ) |
| 647 | ) |
| 648 | assert load_meta.load_event >= 0 |
| 649 | assert len(load_meta.load_gpu_blocks) == 2 |
| 650 | assert len(load_meta.load_cpu_blocks) == 2 |
| 651 | |
| 652 | |
| 653 | def test_finished_eager_store_skips_existing_cpu_cache_entries() -> None: |
| 654 | """A final flush must not consume CPU capacity for an existing hash.""" |
| 655 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 656 | sched = fix.scheduler |
| 657 | |
| 658 | first = make_request(num_blocks=2) |
| 659 | first_blocks = _alloc_and_register(fix, first, num_blocks=2) |
| 660 | sched.update_state_after_alloc(first, first_blocks, num_external_tokens=0) |
| 661 | sched.request_finished_all_groups(first, first_blocks.get_block_ids()) |
| 662 | first_meta = sched.build_connector_meta(make_scheduler_output({})) |
| 663 | simulate_store_completion(sched, first_meta.store_event) |
| 664 | free_blocks = get_cpu_free_blocks(sched) |
| 665 | |
| 666 | duplicate = Request( |
| 667 | request_id="req-finished-eager-duplicate", |
| 668 | prompt_token_ids=first.prompt_token_ids, |
| 669 | sampling_params=first.sampling_params, |
| 670 | pooling_params=None, |
| 671 | mm_features=None, |
| 672 | block_hasher=first._block_hasher, |
| 673 | ) |
| 674 | duplicate_blocks = _alloc_and_register(fix, duplicate, num_blocks=2) |
| 675 | sched.update_state_after_alloc(duplicate, duplicate_blocks, num_external_tokens=0) |
| 676 | sched.request_finished_all_groups(duplicate, duplicate_blocks.get_block_ids()) |
| 677 | |
| 678 | duplicate_meta = sched.build_connector_meta(make_scheduler_output({})) |
| 679 | assert duplicate_meta.store_event == -1 |
| 680 | assert get_cpu_free_blocks(sched) == free_blocks |
| 681 | |
| 682 | |
| 683 | def test_finished_eager_store_is_reported_pending_before_metadata_build() -> None: |
| 684 | """Queued finished stores must keep the connector pending until submitted.""" |
| 685 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 686 | sched = fix.scheduler |
| 687 | request = make_request(num_blocks=2) |
| 688 | kv_blocks = _alloc_and_register(fix, request, num_blocks=2) |
| 689 | connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector) |
| 690 | connector.scheduler_manager = sched |
| 691 | sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0) |
| 692 | |
| 693 | sched.request_finished_all_groups(request, kv_blocks.get_block_ids()) |
| 694 | |
| 695 | assert sched.has_pending_stores() |
| 696 | assert connector.has_pending_push_work() |
| 697 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 698 | assert meta.store_event >= 0 |
| 699 | assert not sched._pending_finished_stores |
| 700 | assert connector.has_pending_push_work() |
| 701 | simulate_store_completion(sched, meta.store_event) |
| 702 | assert not sched.has_pending_stores() |
| 703 | assert not connector.has_pending_push_work() |
| 704 | |
| 705 | |
| 706 | def test_finished_eager_store_caches_all_groups() -> None: |
| 707 | """The final flush retains completed blocks from every KV cache group.""" |
| 708 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=2, lazy=False) |
| 709 | sched = fix.scheduler |
| 710 | request = make_request(num_blocks=2) |
| 711 | group_blocks = tuple( |
| 712 | _allocate_gpu_blocks(fix.gpu_block_pool, request, 2, group_id=group_id) |
| 713 | for group_id in range(2) |
| 714 | ) |
| 715 | request.num_computed_tokens = 2 * BLOCK_SIZE |
| 716 | kv_blocks = KVCacheBlocks(blocks=group_blocks) |
| 717 | sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0) |
| 718 | |
| 719 | sched.request_finished_all_groups(request, kv_blocks.get_block_ids()) |
| 720 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 721 | assert len(meta.store_gpu_blocks) == 4 |
| 722 | simulate_store_completion(sched, meta.store_event) |
| 723 | |
| 724 | for group_id in range(2): |
| 725 | for block_hash in request.block_hashes[:2]: |
| 726 | cached = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 727 | make_block_hash_with_group_id(block_hash, group_id) |
| 728 | ) |
| 729 | assert cached is not None |
| 730 | |
| 731 | |
| 732 | def test_finished_eager_store_respects_cpu_capacity() -> None: |
| 733 | """A finished request skips excess blocks when the CPU cache is full.""" |
| 734 | fix = make_scheduler(num_cpu_blocks=3, num_gpu_blocks=16, lazy=False) |
| 735 | sched = fix.scheduler |
| 736 | request = make_request(num_blocks=3) |
| 737 | kv_blocks = _alloc_and_register(fix, request, num_blocks=3) |
| 738 | sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0) |
| 739 | available = get_cpu_free_blocks(sched) |
| 740 | |
| 741 | sched.request_finished_all_groups(request, kv_blocks.get_block_ids()) |
| 742 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 743 | |
| 744 | assert len(meta.store_gpu_blocks) == available |
| 745 | assert len(meta.store_cpu_blocks) == available |
| 746 | |
| 747 | |
| 748 | def test_reset_releases_unsubmitted_finished_stores() -> None: |
| 749 | """Reset releases finished-store refs that were not sent to a worker.""" |
| 750 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 751 | sched = fix.scheduler |
| 752 | request = make_request(num_blocks=2) |
| 753 | kv_blocks = _alloc_and_register(fix, request, num_blocks=2) |
| 754 | sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0) |
| 755 | |
| 756 | sched.request_finished_all_groups(request, kv_blocks.get_block_ids()) |
| 757 | |
| 758 | assert sched.reset() |
| 759 | |
| 760 | |
| 761 | # --------------------------------------------------------------------------- |
| 762 | # Test 1b: Boundary — max_hit_len cap drops the last full block when the |
| 763 | # prompt is an exact multiple of BLOCK_SIZE. |
| 764 | # --------------------------------------------------------------------------- |
| 765 | def test_max_hit_len_cap_drops_last_full_block() -> None: |
| 766 | """When num_tokens is an exact multiple of BLOCK_SIZE, the manager's |
| 767 | ``max_hit_len = num_tokens - 1`` cap forces ``find_longest_cache_hit`` to |
| 768 | drop the final block (since ``max_length // block_size`` rounds down). |
| 769 | """ |
| 770 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 771 | sched = fix.scheduler |
| 772 | |
| 773 | num_blocks = 2 |
| 774 | req = make_request(num_blocks=num_blocks, extra_tokens=0) |
| 775 | assert req.num_tokens == num_blocks * BLOCK_SIZE |
| 776 | |
| 777 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 778 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 779 | sched_out = make_scheduler_output( |
| 780 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 781 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 782 | ) |
| 783 | meta = sched.build_connector_meta(sched_out) |
| 784 | simulate_store_completion(sched, meta.store_event) |
| 785 | |
| 786 | req2 = Request( |
| 787 | request_id="req-cap-boundary", |
| 788 | prompt_token_ids=req.prompt_token_ids, |
| 789 | sampling_params=req.sampling_params, |
| 790 | pooling_params=None, |
| 791 | mm_features=None, |
| 792 | block_hasher=req._block_hasher, |
| 793 | ) |
| 794 | hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 795 | assert hit_tokens == (num_blocks - 1) * BLOCK_SIZE |
| 796 | |
| 797 | |
| 798 | # --------------------------------------------------------------------------- |
| 799 | # Test 1c: Lazy store-and-load roundtrip |
| 800 | # --------------------------------------------------------------------------- |
| 801 | def _flush_old_blocks_to_lru_head( |
| 802 | gpu_pool: BlockPool, |
| 803 | num_filler_blocks: int, |
| 804 | ) -> list: |
| 805 | """Allocate filler blocks so that previously-freed (hashed) blocks migrate |
| 806 | to the LRU head of the free queue. Returns the filler blocks (caller must |
| 807 | free them later to restore pool capacity). |
| 808 | |
| 809 | In a real engine the same thing happens naturally: after one request |
| 810 | finishes and frees its blocks, subsequent requests allocate from the LRU |
| 811 | head, consuming the unhashed blocks and leaving the old hashed blocks at |
| 812 | the front of the queue. |
| 813 | """ |
| 814 | fillers = gpu_pool.get_new_blocks(num_filler_blocks) |
| 815 | return fillers |
| 816 | |
| 817 | |
| 818 | def test_lazy_store_and_load_roundtrip() -> None: |
| 819 | """Lazy mode: schedule a request, finish it so its hashed blocks are freed, |
| 820 | then schedule new requests so the old blocks migrate to the LRU head. |
| 821 | The lazy scanner offloads them to CPU. Re-scheduling the old request |
| 822 | triggers a CPU cache hit + load. |
| 823 | |
| 824 | GPU pool: 8 blocks (7 usable). _target_free = ceil(64/16) = 4. |
| 825 | """ |
| 826 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True) |
| 827 | sched = fix.scheduler |
| 828 | gpu_pool = fix.gpu_block_pool |
| 829 | |
| 830 | num_blocks = 2 |
| 831 | |
| 832 | # --- Step 1: Schedule req_old, compute, and finish --- |
| 833 | req_old = make_request(num_blocks=num_blocks) |
| 834 | gpu_blocks_old = _allocate_gpu_blocks(gpu_pool, req_old, num_blocks, group_id=0) |
| 835 | gpu_pool.free_blocks(gpu_blocks_old) |
| 836 | |
| 837 | # Allocate filler blocks so req_old's hashed blocks move to LRU head. |
| 838 | # 7 usable - 2 (req_old freed) = 5 other free blocks to consume. |
| 839 | fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) |
| 840 | |
| 841 | # --- Step 2: Lazy scanner should offload req_old's blocks --- |
| 842 | sched_out = make_scheduler_output({}) |
| 843 | meta = sched.build_connector_meta(sched_out) |
| 844 | assert meta.store_event >= 0, "Expected lazy store to offload old blocks" |
| 845 | assert len(meta.store_gpu_blocks) == num_blocks |
| 846 | simulate_store_completion(sched, meta.store_event) |
| 847 | |
| 848 | # Free fillers to restore pool capacity. |
| 849 | gpu_pool.free_blocks(fillers) |
| 850 | |
| 851 | # --- Step 3: Re-schedule req_old — should get CPU cache hit --- |
| 852 | req_old2 = Request( |
| 853 | request_id="req-old-reload", |
| 854 | prompt_token_ids=req_old.prompt_token_ids, |
| 855 | sampling_params=req_old.sampling_params, |
| 856 | pooling_params=None, |
| 857 | mm_features=None, |
| 858 | block_hasher=req_old._block_hasher, |
| 859 | ) |
| 860 | hit_tokens, is_async = sched.get_num_new_matched_tokens( |
| 861 | req_old2, num_computed_tokens=0 |
| 862 | ) |
| 863 | # make_request pads num_tokens by +1 beyond the last full block, so the |
| 864 | # manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact. |
| 865 | expected_hit = num_blocks * BLOCK_SIZE |
| 866 | assert hit_tokens == expected_hit, ( |
| 867 | f"Expected {expected_hit} hit tokens, got {hit_tokens}" |
| 868 | ) |
| 869 | assert is_async is True |
| 870 | |
| 871 | # Allocate fresh GPU blocks for the load. |
| 872 | gpu_blocks_load = gpu_pool.get_new_blocks(num_blocks) |
| 873 | kv_blocks_load = KVCacheBlocks(blocks=(gpu_blocks_load,)) |
| 874 | sched.update_state_after_alloc( |
| 875 | req_old2, kv_blocks_load, num_external_tokens=hit_tokens |
| 876 | ) |
| 877 | |
| 878 | sched_out2 = make_scheduler_output({req_old2.request_id: 1}) |
| 879 | meta2 = sched.build_connector_meta(sched_out2) |
| 880 | assert meta2.load_event >= 0, "Expected a load event to be assigned" |
| 881 | assert len(meta2.load_gpu_blocks) > 0 |
| 882 | |
| 883 | |
| 884 | # --------------------------------------------------------------------------- |
| 885 | # Test 2a: Eager duplicate store is skipped |
| 886 | # --------------------------------------------------------------------------- |
| 887 | def test_eager_duplicate_store_skipped() -> None: |
| 888 | """Eager: storing the same block hashes twice should not allocate new CPU blocks.""" |
| 889 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 890 | sched = fix.scheduler |
| 891 | |
| 892 | num_blocks = 2 |
| 893 | req = make_request(num_blocks=num_blocks) |
| 894 | |
| 895 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 896 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 897 | block_ids = kv_blocks.get_block_ids() |
| 898 | sched_out = make_scheduler_output( |
| 899 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 900 | new_reqs={req.request_id: block_ids}, |
| 901 | ) |
| 902 | |
| 903 | meta1 = sched.build_connector_meta(sched_out) |
| 904 | assert meta1.store_event >= 0 |
| 905 | simulate_store_completion(sched, meta1.store_event) |
| 906 | cpu_free_after_first = get_cpu_free_blocks(sched) |
| 907 | |
| 908 | # Second request with identical hashes — should skip store |
| 909 | req2 = Request( |
| 910 | request_id="req-dup-eager", |
| 911 | prompt_token_ids=req.prompt_token_ids, |
| 912 | sampling_params=req.sampling_params, |
| 913 | pooling_params=None, |
| 914 | mm_features=None, |
| 915 | block_hasher=req._block_hasher, |
| 916 | ) |
| 917 | kv_blocks2 = _alloc_and_register(fix, req2, num_blocks) |
| 918 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=0) |
| 919 | block_ids2 = kv_blocks2.get_block_ids() |
| 920 | sched_out2 = make_scheduler_output( |
| 921 | {req2.request_id: num_blocks * BLOCK_SIZE}, |
| 922 | new_reqs={req2.request_id: block_ids2}, |
| 923 | ) |
| 924 | |
| 925 | meta2 = sched.build_connector_meta(sched_out2) |
| 926 | if meta2.store_event >= 0: |
| 927 | assert len(meta2.store_cpu_blocks) == 0, ( |
| 928 | "Expected no new CPU blocks for duplicate hashes" |
| 929 | ) |
| 930 | assert get_cpu_free_blocks(sched) == cpu_free_after_first |
| 931 | |
| 932 | |
| 933 | # --------------------------------------------------------------------------- |
| 934 | # Test 2b: Eager dedup of in-flight stores across consecutive steps |
| 935 | # --------------------------------------------------------------------------- |
| 936 | def test_eager_in_flight_store_dedup_across_steps() -> None: |
| 937 | """Eager: a second request sharing a prefix with an in-flight store |
| 938 | must not re-offload the same GPU blocks before completion lands. |
| 939 | |
| 940 | Simulates a GPU prefix-cache hit by reusing the first request's |
| 941 | GPU block IDs in the second scheduler step, which is the path the |
| 942 | real scheduler takes when two requests share a prefix. |
| 943 | """ |
| 944 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 945 | sched = fix.scheduler |
| 946 | |
| 947 | num_blocks = 2 |
| 948 | req = make_request(num_blocks=num_blocks) |
| 949 | |
| 950 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 951 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 952 | block_ids = kv_blocks.get_block_ids() |
| 953 | sched_out = make_scheduler_output( |
| 954 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 955 | new_reqs={req.request_id: block_ids}, |
| 956 | ) |
| 957 | |
| 958 | meta1 = sched.build_connector_meta(sched_out) |
| 959 | assert meta1.store_event >= 0 |
| 960 | assert len(meta1.store_cpu_blocks) == num_blocks |
| 961 | # In-flight set tracks the scheduled GPU blocks until completion. |
| 962 | assert sched._in_flight_store_gpu_blocks == set(meta1.store_gpu_blocks) |
| 963 | cpu_free_after_first = get_cpu_free_blocks(sched) |
| 964 | |
| 965 | # Second request shares the prefix and reuses the same GPU block IDs |
| 966 | # (the real scheduler path: GPU prefix cache returns the same blocks). |
| 967 | # Do NOT simulate completion — the first store is still in-flight. |
| 968 | req2 = Request( |
| 969 | request_id="req-dup-eager-inflight", |
| 970 | prompt_token_ids=req.prompt_token_ids, |
| 971 | sampling_params=req.sampling_params, |
| 972 | pooling_params=None, |
| 973 | mm_features=None, |
| 974 | block_hasher=req._block_hasher, |
| 975 | ) |
| 976 | req2.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 977 | sched.update_state_after_alloc(req2, kv_blocks, num_external_tokens=0) |
| 978 | sched_out2 = make_scheduler_output( |
| 979 | {req2.request_id: num_blocks * BLOCK_SIZE}, |
| 980 | new_reqs={req2.request_id: block_ids}, |
| 981 | ) |
| 982 | |
| 983 | meta2 = sched.build_connector_meta(sched_out2) |
| 984 | if meta2.store_event >= 0: |
| 985 | assert len(meta2.store_cpu_blocks) == 0, ( |
| 986 | "Expected no new CPU blocks for in-flight duplicate hashes" |
| 987 | ) |
| 988 | assert get_cpu_free_blocks(sched) == cpu_free_after_first, ( |
| 989 | "Second request should not consume CPU blocks while the first " |
| 990 | "store is still in-flight" |
| 991 | ) |
| 992 | |
| 993 | # After completion, the in-flight set is cleared. |
| 994 | simulate_store_completion(sched, meta1.store_event) |
| 995 | assert sched._in_flight_store_gpu_blocks == set() |
| 996 | |
| 997 | |
| 998 | # --------------------------------------------------------------------------- |
| 999 | # Test 2c: Lazy duplicate store is skipped |
| 1000 | # --------------------------------------------------------------------------- |
| 1001 | def test_lazy_duplicate_store_skipped() -> None: |
| 1002 | """Lazy: blocks already offloaded to CPU should not be offloaded again. |
| 1003 | |
| 1004 | Same pattern as the lazy roundtrip: flush old blocks to LRU head, offload, |
| 1005 | then repeat with the same hashes and verify no new CPU allocation. |
| 1006 | """ |
| 1007 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True) |
| 1008 | sched = fix.scheduler |
| 1009 | gpu_pool = fix.gpu_block_pool |
| 1010 | |
| 1011 | num_blocks = 2 |
| 1012 | req = make_request(num_blocks=num_blocks) |
| 1013 | |
| 1014 | # Schedule + finish → hashed blocks in free queue |
| 1015 | gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 1016 | gpu_pool.free_blocks(gpu_blocks) |
| 1017 | |
| 1018 | # Flush old blocks to LRU head, then trigger lazy offload. |
| 1019 | fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) |
| 1020 | meta1 = sched.build_connector_meta(make_scheduler_output({})) |
| 1021 | assert meta1.store_event >= 0 |
| 1022 | simulate_store_completion(sched, meta1.store_event) |
| 1023 | gpu_pool.free_blocks(fillers) |
| 1024 | cpu_free_after_first = get_cpu_free_blocks(sched) |
| 1025 | |
| 1026 | # Allocate blocks with the same hashes and free them again. |
| 1027 | # The scanner should see they are already in CPU cache and skip them. |
| 1028 | req2 = Request( |
| 1029 | request_id="req-dup-lazy", |
| 1030 | prompt_token_ids=req.prompt_token_ids, |
| 1031 | sampling_params=req.sampling_params, |
| 1032 | pooling_params=None, |
| 1033 | mm_features=None, |
| 1034 | block_hasher=req._block_hasher, |
| 1035 | ) |
| 1036 | gpu_blocks2 = _allocate_gpu_blocks(gpu_pool, req2, num_blocks, group_id=0) |
| 1037 | gpu_pool.free_blocks(gpu_blocks2) |
| 1038 | |
| 1039 | # Flush again so the hashed blocks are at LRU head for the scanner. |
| 1040 | fillers2 = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) |
| 1041 | meta2 = sched.build_connector_meta(make_scheduler_output({})) |
| 1042 | gpu_pool.free_blocks(fillers2) |
| 1043 | |
| 1044 | # Either no store event, or zero new CPU blocks (already cached). |
| 1045 | if meta2.store_event >= 0: |
| 1046 | assert len(meta2.store_cpu_blocks) == 0, ( |
| 1047 | "Expected no new CPU blocks for duplicate hashes" |
| 1048 | ) |
| 1049 | assert get_cpu_free_blocks(sched) == cpu_free_after_first |
| 1050 | |
| 1051 | |
| 1052 | # --------------------------------------------------------------------------- |
| 1053 | # Test 3: LRU eviction order |
| 1054 | # --------------------------------------------------------------------------- |
| 1055 | def test_lru_eviction_order() -> None: |
| 1056 | """With limited CPU space, oldest blocks should be evicted first. |
| 1057 | |
| 1058 | CPU block pool: num_cpu_blocks=5 -> 4 free usable blocks (1 taken by null_block). |
| 1059 | After storing 4 blocks (2 req_a + 2 req_b), all free slots are occupied by |
| 1060 | cached blocks (ref_cnt=0, in hash map). When 2 more are stored (req_c), |
| 1061 | 2 LRU blocks from req_a get evicted from the cache to make room. |
| 1062 | """ |
| 1063 | # 5 total = 4 usable (null_block takes 1), filling exactly with 4 blocks |
| 1064 | fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False) |
| 1065 | sched = fix.scheduler |
| 1066 | |
| 1067 | # Fill CPU with 4 blocks: 2 requests x 2 blocks (in LRU insertion order) |
| 1068 | req_a = make_request(num_blocks=2) |
| 1069 | req_b = make_request(num_blocks=2) |
| 1070 | |
| 1071 | kv_a = _alloc_and_register(fix, req_a, 2) |
| 1072 | kv_b = _alloc_and_register(fix, req_b, 2) |
| 1073 | sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0) |
| 1074 | sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0) |
| 1075 | |
| 1076 | ids_a = kv_a.get_block_ids() |
| 1077 | ids_b = kv_b.get_block_ids() |
| 1078 | sched_out = make_scheduler_output( |
| 1079 | { |
| 1080 | req_a.request_id: 2 * BLOCK_SIZE, |
| 1081 | req_b.request_id: 2 * BLOCK_SIZE, |
| 1082 | }, |
| 1083 | new_reqs={ |
| 1084 | req_a.request_id: ids_a, |
| 1085 | req_b.request_id: ids_b, |
| 1086 | }, |
| 1087 | ) |
| 1088 | meta = sched.build_connector_meta(sched_out) |
| 1089 | assert meta.store_event >= 0 |
| 1090 | simulate_store_completion(sched, meta.store_event) |
| 1091 | |
| 1092 | # Verify all 4 blocks are cached in CPU hash map |
| 1093 | for i, bhash in enumerate(req_a.block_hashes[:2]): |
| 1094 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1095 | assert ( |
| 1096 | sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 1097 | bhash_with_group |
| 1098 | ) |
| 1099 | is not None |
| 1100 | ), f"req_a block {i} should be cached after store" |
| 1101 | for i, bhash in enumerate(req_b.block_hashes[:2]): |
| 1102 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1103 | assert ( |
| 1104 | sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 1105 | bhash_with_group |
| 1106 | ) |
| 1107 | is not None |
| 1108 | ), f"req_b block {i} should be cached after store" |
| 1109 | |
| 1110 | # Store 2 more blocks from a new request - must evict 2 LRU blocks (req_a) |
| 1111 | req_c = make_request(num_blocks=2) |
| 1112 | kv_c = _alloc_and_register(fix, req_c, 2) |
| 1113 | sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0) |
| 1114 | |
| 1115 | ids_c = kv_c.get_block_ids() |
| 1116 | sched_out2 = make_scheduler_output( |
| 1117 | {req_c.request_id: 2 * BLOCK_SIZE}, |
| 1118 | new_reqs={req_c.request_id: ids_c}, |
| 1119 | ) |
| 1120 | meta2 = sched.build_connector_meta(sched_out2) |
| 1121 | assert meta2.store_event >= 0 |
| 1122 | simulate_store_completion(sched, meta2.store_event) |
| 1123 | |
| 1124 | # req_a hashes should be evicted from CPU (they were LRU) |
| 1125 | for i, bhash in enumerate(req_a.block_hashes[:2]): |
| 1126 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1127 | cache_map = sched.cpu_block_pool.cached_block_hash_to_block |
| 1128 | cached = cache_map.get_one_block(bhash_with_group) |
| 1129 | assert cached is None, f"req_a block {i} should have been evicted" |
| 1130 | |
| 1131 | # req_b and req_c hashes should be present |
| 1132 | for i, bhash in enumerate(req_b.block_hashes[:2]): |
| 1133 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1134 | cache_map = sched.cpu_block_pool.cached_block_hash_to_block |
| 1135 | cached = cache_map.get_one_block(bhash_with_group) |
| 1136 | assert cached is not None, f"req_b block {i} should still be cached" |
| 1137 | |
| 1138 | for i, bhash in enumerate(req_c.block_hashes[:2]): |
| 1139 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1140 | cache_map = sched.cpu_block_pool.cached_block_hash_to_block |
| 1141 | cached = cache_map.get_one_block(bhash_with_group) |
| 1142 | assert cached is not None, f"req_c block {i} should still be cached" |
| 1143 | |
| 1144 | |
| 1145 | # --------------------------------------------------------------------------- |
| 1146 | # Test 4: Touched blocks survive eviction |
| 1147 | # --------------------------------------------------------------------------- |
| 1148 | def test_touched_blocks_survive_eviction() -> None: |
| 1149 | """Touching CPU blocks updates their LRU position, protecting them from eviction.""" |
| 1150 | # 5 total = 4 usable (null_block takes 1) |
| 1151 | fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False) |
| 1152 | sched = fix.scheduler |
| 1153 | |
| 1154 | # Fill CPU with 4 blocks (req_a: 2, req_b: 2) in LRU order |
| 1155 | req_a = make_request(num_blocks=2) |
| 1156 | req_b = make_request(num_blocks=2) |
| 1157 | |
| 1158 | kv_a = _alloc_and_register(fix, req_a, 2) |
| 1159 | kv_b = _alloc_and_register(fix, req_b, 2) |
| 1160 | sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0) |
| 1161 | sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0) |
| 1162 | |
| 1163 | ids_a = kv_a.get_block_ids() |
| 1164 | ids_b = kv_b.get_block_ids() |
| 1165 | sched_out = make_scheduler_output( |
| 1166 | { |
| 1167 | req_a.request_id: 2 * BLOCK_SIZE, |
| 1168 | req_b.request_id: 2 * BLOCK_SIZE, |
| 1169 | }, |
| 1170 | new_reqs={ |
| 1171 | req_a.request_id: ids_a, |
| 1172 | req_b.request_id: ids_b, |
| 1173 | }, |
| 1174 | ) |
| 1175 | meta = sched.build_connector_meta(sched_out) |
| 1176 | simulate_store_completion(sched, meta.store_event) |
| 1177 | |
| 1178 | # Touch req_a's CPU blocks to make them most-recently-used |
| 1179 | cpu_pool = sched.cpu_block_pool |
| 1180 | for bhash in req_a.block_hashes[:2]: |
| 1181 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1182 | cached_blk = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group) |
| 1183 | assert cached_blk is not None |
| 1184 | cpu_pool.touch([cached_blk]) |
| 1185 | # Undo touch to return ref_cnt to 0 |
| 1186 | # (so it's a free candidate but at MRU position) |
| 1187 | cpu_pool.free_blocks([cached_blk]) |
| 1188 | |
| 1189 | # Now store 2 more blocks; req_b (LRU front) should be evicted, not req_a |
| 1190 | req_c = make_request(num_blocks=2) |
| 1191 | kv_c = _alloc_and_register(fix, req_c, 2) |
| 1192 | sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0) |
| 1193 | |
| 1194 | ids_c = kv_c.get_block_ids() |
| 1195 | sched_out2 = make_scheduler_output( |
| 1196 | {req_c.request_id: 2 * BLOCK_SIZE}, |
| 1197 | new_reqs={req_c.request_id: ids_c}, |
| 1198 | ) |
| 1199 | meta2 = sched.build_connector_meta(sched_out2) |
| 1200 | simulate_store_completion(sched, meta2.store_event) |
| 1201 | |
| 1202 | # req_b should be evicted (LRU), req_a and req_c should survive |
| 1203 | for i, bhash in enumerate(req_b.block_hashes[:2]): |
| 1204 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1205 | cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group) |
| 1206 | assert cached is None, f"req_b block {i} should have been evicted (it was LRU)" |
| 1207 | |
| 1208 | for i, bhash in enumerate(req_a.block_hashes[:2]): |
| 1209 | bhash_with_group = make_block_hash_with_group_id(bhash, 0) |
| 1210 | cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group) |
| 1211 | assert cached is not None, f"req_a block {i} should survive (was touched/MRU)" |
| 1212 | |
| 1213 | |
| 1214 | # --------------------------------------------------------------------------- |
| 1215 | # Test 5: Preemption no CPU block leak |
| 1216 | # --------------------------------------------------------------------------- |
| 1217 | def test_preemption_no_cpu_block_leak() -> None: |
| 1218 | """request_finished during in-flight load defers cleanup; |
| 1219 | completes after load done.""" |
| 1220 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1221 | sched = fix.scheduler |
| 1222 | |
| 1223 | num_blocks = 2 |
| 1224 | |
| 1225 | # First: store blocks to CPU |
| 1226 | req = make_request(num_blocks=num_blocks) |
| 1227 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 1228 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1229 | block_ids = kv_blocks.get_block_ids() |
| 1230 | sched_out = make_scheduler_output( |
| 1231 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1232 | new_reqs={req.request_id: block_ids}, |
| 1233 | ) |
| 1234 | meta = sched.build_connector_meta(sched_out) |
| 1235 | simulate_store_completion(sched, meta.store_event) |
| 1236 | |
| 1237 | # Create new request with same tokens, check hit |
| 1238 | req2 = Request( |
| 1239 | request_id="req-preempt-load", |
| 1240 | prompt_token_ids=req.prompt_token_ids, |
| 1241 | sampling_params=req.sampling_params, |
| 1242 | pooling_params=None, |
| 1243 | mm_features=None, |
| 1244 | block_hasher=req._block_hasher, |
| 1245 | ) |
| 1246 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 1247 | assert hit_tokens is not None and hit_tokens > 0 |
| 1248 | |
| 1249 | gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) |
| 1250 | kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) |
| 1251 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 1252 | |
| 1253 | # Assign load_event via build_connector_meta |
| 1254 | block_ids2 = kv_blocks2.get_block_ids() |
| 1255 | sched_out2 = make_scheduler_output( |
| 1256 | {req2.request_id: 1}, |
| 1257 | new_reqs={req2.request_id: block_ids2}, |
| 1258 | ) |
| 1259 | meta2 = sched.build_connector_meta(sched_out2) |
| 1260 | assert meta2.load_event >= 0 |
| 1261 | |
| 1262 | # Request finishes BEFORE load completes -> deferred |
| 1263 | sched.request_finished(req2, block_ids=[]) |
| 1264 | assert req2.request_id in sched._reqs_to_load |
| 1265 | assert sched._reqs_to_load[req2.request_id].finished is True |
| 1266 | |
| 1267 | # Now simulate load completion -> cleanup fires |
| 1268 | simulate_load_completion(sched, {req2.request_id}) |
| 1269 | assert req2.request_id not in sched._reqs_to_load |
| 1270 | |
| 1271 | |
| 1272 | # --------------------------------------------------------------------------- |
| 1273 | # Test 6: Eager store preemption cleanup |
| 1274 | # --------------------------------------------------------------------------- |
| 1275 | def test_eager_store_preemption_cleanup() -> None: |
| 1276 | """In eager mode, finishing a request during in-flight store defers cleanup.""" |
| 1277 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1278 | sched = fix.scheduler |
| 1279 | |
| 1280 | num_blocks = 2 |
| 1281 | req = make_request(num_blocks=num_blocks) |
| 1282 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 1283 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1284 | |
| 1285 | block_ids = kv_blocks.get_block_ids() |
| 1286 | sched_out = make_scheduler_output( |
| 1287 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1288 | new_reqs={req.request_id: block_ids}, |
| 1289 | ) |
| 1290 | meta = sched.build_connector_meta(sched_out) |
| 1291 | store_event = meta.store_event |
| 1292 | assert store_event >= 0 |
| 1293 | |
| 1294 | # The request gets store_events populated |
| 1295 | assert req.request_id in sched._reqs_to_store |
| 1296 | store_state = sched._reqs_to_store[req.request_id] |
| 1297 | assert store_event in store_state.store_events |
| 1298 | |
| 1299 | # Finish request while store still in-flight -> deferred |
| 1300 | sched.request_finished(req, block_ids=[]) |
| 1301 | assert req.request_id in sched._reqs_to_store |
| 1302 | assert sched._reqs_to_store[req.request_id].finished is True |
| 1303 | |
| 1304 | # Simulate store completion -> deferred cleanup fires |
| 1305 | simulate_store_completion(sched, store_event) |
| 1306 | assert req.request_id not in sched._reqs_to_store |
| 1307 | |
| 1308 | |
| 1309 | # --------------------------------------------------------------------------- |
| 1310 | # Test 7: In-flight finish deferred cleanup (load variant) |
| 1311 | # --------------------------------------------------------------------------- |
| 1312 | def test_inflight_finish_deferred_cleanup() -> None: |
| 1313 | """Store, then start a load, request_finished defers, |
| 1314 | load completion fires cleanup.""" |
| 1315 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1316 | sched = fix.scheduler |
| 1317 | |
| 1318 | num_blocks = 2 |
| 1319 | |
| 1320 | # Store |
| 1321 | req = make_request(num_blocks=num_blocks) |
| 1322 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 1323 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1324 | block_ids = kv_blocks.get_block_ids() |
| 1325 | sched_out = make_scheduler_output( |
| 1326 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1327 | new_reqs={req.request_id: block_ids}, |
| 1328 | ) |
| 1329 | meta = sched.build_connector_meta(sched_out) |
| 1330 | simulate_store_completion(sched, meta.store_event) |
| 1331 | |
| 1332 | # Load |
| 1333 | req2 = Request( |
| 1334 | request_id="req-inflight-load", |
| 1335 | prompt_token_ids=req.prompt_token_ids, |
| 1336 | sampling_params=req.sampling_params, |
| 1337 | pooling_params=None, |
| 1338 | mm_features=None, |
| 1339 | block_hasher=req._block_hasher, |
| 1340 | ) |
| 1341 | hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 1342 | assert hit_tokens is not None and hit_tokens > 0 |
| 1343 | |
| 1344 | gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) |
| 1345 | kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) |
| 1346 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 1347 | |
| 1348 | block_ids2 = kv_blocks2.get_block_ids() |
| 1349 | sched_out2 = make_scheduler_output( |
| 1350 | {req2.request_id: 1}, |
| 1351 | new_reqs={req2.request_id: block_ids2}, |
| 1352 | ) |
| 1353 | meta2 = sched.build_connector_meta(sched_out2) |
| 1354 | assert meta2.load_event >= 0 |
| 1355 | |
| 1356 | # Finish before load completes |
| 1357 | sched.request_finished(req2, block_ids=[]) |
| 1358 | assert req2.request_id in sched._reqs_to_load |
| 1359 | |
| 1360 | # Simulate load completion -> request removed |
| 1361 | simulate_load_completion(sched, {req2.request_id}) |
| 1362 | assert req2.request_id not in sched._reqs_to_load |
| 1363 | |
| 1364 | |
| 1365 | # --------------------------------------------------------------------------- |
| 1366 | # Test 8: Null GPU blocks are skipped in store and load transfer pairs |
| 1367 | # --------------------------------------------------------------------------- |
| 1368 | def test_multi_group_null_blocks_skipped() -> None: |
| 1369 | """The null block id never appears in store or load pairs. |
| 1370 | |
| 1371 | A null block-table slot whose hash is still cached on the GPU is stored |
| 1372 | from the cached block instead; the null block itself is never copied. |
| 1373 | """ |
| 1374 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False) |
| 1375 | sched = fix.scheduler |
| 1376 | gpu_pool = fix.gpu_block_pool |
| 1377 | |
| 1378 | num_blocks = 2 |
| 1379 | req = make_request(num_blocks=num_blocks) |
| 1380 | |
| 1381 | # Allocate real blocks (with hashes) and use the null_block as a placeholder |
| 1382 | gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 1383 | null_block = gpu_pool.null_block |
| 1384 | |
| 1385 | # Mix: [real_block, null_block] — null_block has no hash, should be skipped |
| 1386 | mixed_blocks = [gpu_blocks[0], null_block] |
| 1387 | kv_blocks = KVCacheBlocks(blocks=(mixed_blocks,)) |
| 1388 | req.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 1389 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1390 | |
| 1391 | block_ids = kv_blocks.get_block_ids() |
| 1392 | sched_out = make_scheduler_output( |
| 1393 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1394 | new_reqs={req.request_id: block_ids}, |
| 1395 | ) |
| 1396 | meta = sched.build_connector_meta(sched_out) |
| 1397 | |
| 1398 | # Null block's ID should NOT appear in store_gpu_blocks |
| 1399 | null_block_id = null_block.block_id |
| 1400 | assert null_block_id not in meta.store_gpu_blocks, ( |
| 1401 | f"Null block id {null_block_id} should not appear in store transfer pairs" |
| 1402 | ) |
| 1403 | |
| 1404 | # The nulled slot's hash is still cached on the GPU, so that block is |
| 1405 | # recovered from the pool and stored alongside the real one. |
| 1406 | assert len(meta.store_gpu_blocks) == 2 |
| 1407 | assert gpu_blocks[0].block_id in meta.store_gpu_blocks |
| 1408 | assert gpu_blocks[1].block_id in meta.store_gpu_blocks |
| 1409 | |
| 1410 | # Complete the store |
| 1411 | assert meta.store_event >= 0 |
| 1412 | simulate_store_completion(sched, meta.store_event) |
| 1413 | |
| 1414 | # Create matching request and get load hit |
| 1415 | req2 = Request( |
| 1416 | request_id="req-null-load", |
| 1417 | prompt_token_ids=req.prompt_token_ids, |
| 1418 | sampling_params=req.sampling_params, |
| 1419 | pooling_params=None, |
| 1420 | mm_features=None, |
| 1421 | block_hasher=req._block_hasher, |
| 1422 | ) |
| 1423 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 1424 | assert hit_tokens == 2 * BLOCK_SIZE |
| 1425 | assert is_async is True |
| 1426 | |
| 1427 | # Allocate new GPU blocks for the load |
| 1428 | gpu_blocks2 = gpu_pool.get_new_blocks(2) |
| 1429 | kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) |
| 1430 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 1431 | |
| 1432 | sched_out2 = make_scheduler_output({req2.request_id: 1}) |
| 1433 | meta2 = sched.build_connector_meta(sched_out2) |
| 1434 | |
| 1435 | # Null block's ID should NOT appear in load_gpu_blocks |
| 1436 | assert null_block_id not in meta2.load_gpu_blocks, ( |
| 1437 | f"Null block id {null_block_id} should not appear in load transfer pairs" |
| 1438 | ) |
| 1439 | |
| 1440 | |
| 1441 | # --------------------------------------------------------------------------- |
| 1442 | # Test 8b: Non-prefix-cacheable scratch groups take no part in store or load |
| 1443 | # --------------------------------------------------------------------------- |
| 1444 | def test_scratch_group_excluded_from_store_and_load() -> None: |
| 1445 | """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded. |
| 1446 | |
| 1447 | Its single per-request block carries no hash, so the store must offload |
| 1448 | only the attention blocks and the load must pair only attention blocks. |
| 1449 | """ |
| 1450 | fix = make_scheduler( |
| 1451 | num_cpu_blocks=8, |
| 1452 | num_gpu_blocks=16, |
| 1453 | kv_cache_config=_make_scratch_kv_cache_config(16), |
| 1454 | ) |
| 1455 | sched = fix.scheduler |
| 1456 | gpu_pool = fix.gpu_block_pool |
| 1457 | num_blocks = 2 |
| 1458 | |
| 1459 | req = make_request(num_blocks=num_blocks) |
| 1460 | fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 1461 | scratch_block = gpu_pool.get_new_blocks(1) |
| 1462 | kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) |
| 1463 | req.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 1464 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1465 | sched_out = make_scheduler_output( |
| 1466 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1467 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 1468 | ) |
| 1469 | meta = sched.build_connector_meta(sched_out) |
| 1470 | assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) |
| 1471 | simulate_store_completion(sched, meta.store_event) |
| 1472 | |
| 1473 | req2 = Request( |
| 1474 | request_id="req-scratch-load", |
| 1475 | prompt_token_ids=req.prompt_token_ids, |
| 1476 | sampling_params=req.sampling_params, |
| 1477 | pooling_params=None, |
| 1478 | mm_features=None, |
| 1479 | block_hasher=req._block_hasher, |
| 1480 | ) |
| 1481 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 1482 | assert hit_tokens == num_blocks * BLOCK_SIZE |
| 1483 | assert is_async is True |
| 1484 | |
| 1485 | fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) |
| 1486 | scratch_block2 = gpu_pool.get_new_blocks(1) |
| 1487 | kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2)) |
| 1488 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 1489 | sched_out2 = make_scheduler_output( |
| 1490 | {req2.request_id: 1}, |
| 1491 | new_reqs={req2.request_id: kv_blocks2.get_block_ids()}, |
| 1492 | ) |
| 1493 | meta2 = sched.build_connector_meta(sched_out2) |
| 1494 | assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2) |
| 1495 | assert len(meta2.load_cpu_blocks) == num_blocks |
| 1496 | |
| 1497 | |
| 1498 | # --------------------------------------------------------------------------- |
| 1499 | # Test 9: Chunked prefill accumulates block_ids across steps |
| 1500 | # --------------------------------------------------------------------------- |
| 1501 | def test_chunked_prefill_reads_live_block_ids() -> None: |
| 1502 | """With chunked prefill, block IDs accumulate across scheduler steps. |
| 1503 | _prepare_eager_store_specs reads block IDs from scheduler_output via |
| 1504 | yield_req_data, so the store should reflect the updated (larger) block |
| 1505 | list, not a stale snapshot.""" |
| 1506 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1507 | sched = fix.scheduler |
| 1508 | |
| 1509 | num_blocks = 4 |
| 1510 | req = make_request(num_blocks=num_blocks) |
| 1511 | |
| 1512 | # First chunk: allocate 2 blocks |
| 1513 | kv_blocks_first = _alloc_and_register(fix, req, 2) |
| 1514 | sched.update_state_after_alloc(req, kv_blocks_first, num_external_tokens=0) |
| 1515 | |
| 1516 | assert req.request_id in sched._reqs_to_store |
| 1517 | # Should still be exactly 1 entry in _reqs_to_store |
| 1518 | assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1 |
| 1519 | |
| 1520 | # Build connector meta with 2 blocks — stores the first 2 |
| 1521 | ids_first = kv_blocks_first.get_block_ids() |
| 1522 | sched_out1 = make_scheduler_output( |
| 1523 | {req.request_id: 2 * BLOCK_SIZE}, |
| 1524 | new_reqs={req.request_id: ids_first}, |
| 1525 | ) |
| 1526 | meta1 = sched.build_connector_meta(sched_out1) |
| 1527 | assert meta1.store_event >= 0 |
| 1528 | assert len(meta1.store_gpu_blocks) == 2 |
| 1529 | simulate_store_completion(sched, meta1.store_event) |
| 1530 | |
| 1531 | # Second chunk: allocate 4 blocks total (2 new ones) |
| 1532 | kv_blocks_second = _alloc_and_register(fix, req, num_blocks) |
| 1533 | # update_state_after_alloc is idempotent for store registration |
| 1534 | sched.update_state_after_alloc(req, kv_blocks_second, num_external_tokens=0) |
| 1535 | |
| 1536 | # Still exactly 1 entry |
| 1537 | assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1 |
| 1538 | |
| 1539 | # The second chunk's NEW block IDs (positions 2,3) are passed as |
| 1540 | # cached_req_new_blocks. The full block_ids include both old and new, |
| 1541 | # but yield_req_data only appends the new_block_ids for cached reqs. |
| 1542 | ids_second_full = kv_blocks_second.get_block_ids() |
| 1543 | # New blocks are those beyond the first chunk |
| 1544 | new_block_ids = tuple(ids_second_full[g][2:] for g in range(len(ids_second_full))) |
| 1545 | sched_out2 = make_scheduler_output( |
| 1546 | {req.request_id: 2 * BLOCK_SIZE}, |
| 1547 | cached_req_new_blocks={req.request_id: new_block_ids}, |
| 1548 | ) |
| 1549 | meta2 = sched.build_connector_meta(sched_out2) |
| 1550 | assert meta2.store_event >= 0 |
| 1551 | # Only the 2 NEW blocks should be stored (first 2 already done) |
| 1552 | assert len(meta2.store_gpu_blocks) == 2 |
| 1553 | |
| 1554 | |
| 1555 | # --------------------------------------------------------------------------- |
| 1556 | # Test 10: Partial GPU prefix hit + CPU load + new compute blocks |
| 1557 | # --------------------------------------------------------------------------- |
| 1558 | def test_partial_gpu_prefix_plus_cpu_load() -> None: |
| 1559 | """When GPU has a prefix cache hit for the first N blocks, CPU has a |
| 1560 | hit for the next M blocks, and there are P new blocks needing fresh |
| 1561 | compute, the block layout is: |
| 1562 | |
| 1563 | | comp (N) | ext_comp (M) | new (P) | |
| 1564 | |
| 1565 | External blocks sit in the middle — not at the beginning or end. |
| 1566 | The load path must target hashes at positions [N, N+M). |
| 1567 | |
| 1568 | Request: 6 blocks (0..5). |
| 1569 | - Store all 6 to CPU. |
| 1570 | - New request: GPU prefix cache hits blocks 0,1 (hashed). |
| 1571 | CPU hits blocks 2,3. Blocks 4,5 are new (need compute). |
| 1572 | - update_state_after_alloc receives 6 GPU blocks: |
| 1573 | [0,1] hashed (comp), [2,3] unhashed (ext_comp), [4,5] unhashed (new). |
| 1574 | - Load must target hash positions 2,3. |
| 1575 | """ |
| 1576 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1577 | sched = fix.scheduler |
| 1578 | gpu_pool = fix.gpu_block_pool |
| 1579 | |
| 1580 | num_blocks = 6 |
| 1581 | req = make_request(num_blocks=num_blocks) |
| 1582 | |
| 1583 | # Store all 6 blocks to CPU via eager store. |
| 1584 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 1585 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1586 | block_ids = kv_blocks.get_block_ids() |
| 1587 | sched_out = make_scheduler_output( |
| 1588 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1589 | new_reqs={req.request_id: block_ids}, |
| 1590 | ) |
| 1591 | meta = sched.build_connector_meta(sched_out) |
| 1592 | assert meta.store_event >= 0 |
| 1593 | simulate_store_completion(sched, meta.store_event) |
| 1594 | |
| 1595 | # New request with same tokens — but only partial GPU prefix hit. |
| 1596 | req2 = Request( |
| 1597 | request_id="req-partial-gpu", |
| 1598 | prompt_token_ids=req.prompt_token_ids, |
| 1599 | sampling_params=req.sampling_params, |
| 1600 | pooling_params=None, |
| 1601 | mm_features=None, |
| 1602 | block_hasher=req._block_hasher, |
| 1603 | ) |
| 1604 | |
| 1605 | # GPU prefix cache hits the first 2 blocks. |
| 1606 | gpu_local_computed = 2 * BLOCK_SIZE |
| 1607 | hit_tokens, is_async = sched.get_num_new_matched_tokens( |
| 1608 | req2, num_computed_tokens=gpu_local_computed |
| 1609 | ) |
| 1610 | # CPU has all 6 blocks stored. make_request pads num_tokens by +1, so |
| 1611 | # the manager's num_tokens - 1 cap leaves all full blocks intact: |
| 1612 | # remaining hashable range = 6 - 2 = 4 blocks, all hit. |
| 1613 | num_cpu_hit_blocks = 4 |
| 1614 | assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, ( |
| 1615 | f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}" |
| 1616 | ) |
| 1617 | assert is_async is True |
| 1618 | |
| 1619 | # Simulate what the real scheduler does: only accept 2 of the 4 CPU hit |
| 1620 | # blocks as external (e.g. due to budget constraints), leaving 2 new |
| 1621 | # blocks for fresh compute. |
| 1622 | num_ext_blocks = 2 |
| 1623 | num_new_blocks = 2 |
| 1624 | external_tokens = num_ext_blocks * BLOCK_SIZE |
| 1625 | |
| 1626 | # Build block list matching real layout: | comp(2) | ext_comp(2) | new(2) | |
| 1627 | # comp: GPU prefix cache hit — blocks with hashes |
| 1628 | gpu_comp = _allocate_gpu_blocks(gpu_pool, req2, 2, group_id=0) |
| 1629 | # ext_comp + new: freshly allocated, no hashes |
| 1630 | gpu_ext_and_new = gpu_pool.get_new_blocks(num_ext_blocks + num_new_blocks) |
| 1631 | all_gpu_blocks = gpu_comp + gpu_ext_and_new |
| 1632 | kv_blocks2 = KVCacheBlocks(blocks=(all_gpu_blocks,)) |
| 1633 | |
| 1634 | # Critical call: with 2 hashed comp blocks and 2 external tokens worth |
| 1635 | # of blocks, the manager must derive skipped=2 and load hashes [2,3]. |
| 1636 | sched.update_state_after_alloc( |
| 1637 | req2, kv_blocks2, num_external_tokens=external_tokens |
| 1638 | ) |
| 1639 | |
| 1640 | block_ids2 = kv_blocks2.get_block_ids() |
| 1641 | sched_out2 = make_scheduler_output( |
| 1642 | {req2.request_id: num_new_blocks * BLOCK_SIZE}, |
| 1643 | new_reqs={req2.request_id: block_ids2}, |
| 1644 | ) |
| 1645 | meta2 = sched.build_connector_meta(sched_out2) |
| 1646 | assert meta2.load_event >= 0, "Expected a load event for partial GPU + CPU hit" |
| 1647 | assert len(meta2.load_gpu_blocks) == num_ext_blocks |
| 1648 | assert len(meta2.load_cpu_blocks) == num_ext_blocks |
| 1649 | |
| 1650 | # Verify the load targets the ext_comp GPU blocks (positions 2,3), |
| 1651 | # not the comp blocks (0,1) or new blocks (4,5). |
| 1652 | ext_block_ids = [b.block_id for b in gpu_ext_and_new[:num_ext_blocks]] |
| 1653 | for bid in meta2.load_gpu_blocks: |
| 1654 | assert bid in ext_block_ids, ( |
| 1655 | f"Load GPU block {bid} should be an ext_comp block, not a comp or new block" |
| 1656 | ) |
| 1657 | |
| 1658 | |
| 1659 | # --------------------------------------------------------------------------- |
| 1660 | # Test 11: TOCTOU between Phase A and Phase B (regression for #39702) |
| 1661 | # --------------------------------------------------------------------------- |
| 1662 | def test_toctou_cpu_hit_evicted_between_phases_no_crash() -> None: |
| 1663 | """Regression for vllm-project/vllm#39702. |
| 1664 | |
| 1665 | When ``get_num_new_matched_tokens`` (Phase A) reports a CPU cache hit |
| 1666 | of ``N`` tokens but ``update_state_after_alloc`` (Phase B) runs after |
| 1667 | other requests have caused LRU eviction of those exact blocks, the |
| 1668 | second ``find_longest_cache_hit`` call returns 0 while |
| 1669 | ``num_external_tokens`` is still ``N``, triggering |
| 1670 | ``AssertionError: Expected N hit tokens, got 0``. |
| 1671 | |
| 1672 | Setup: ``num_cpu_blocks=5`` (4 usable; null_block takes 1). |
| 1673 | 1. Store req_a's 2 blocks to CPU. CPU: [a0, a1, _, _]. |
| 1674 | 2. Phase A on req_b (same prompt as req_a) reports a 2-block hit. |
| 1675 | Without the fix this does NOT pin a0/a1 — they remain at LRU front. |
| 1676 | 3. Store req_c (2) + req_d (2) — 4 more blocks into 4 slots. With |
| 1677 | a0/a1 unpinned and at LRU front, they get evicted. CPU: [c0, c1, |
| 1678 | d0, d1]. |
| 1679 | 4. Phase B on req_b: re-searches the CPU coordinator, finds 0 |
| 1680 | cached hashes, asserts. |
| 1681 | |
| 1682 | After the fix, Phase A pins the hit blocks so step 3 evicts req_c's |
| 1683 | blocks instead, and Phase B reads the cached ``(cpu_hit_blocks, |
| 1684 | hit_length)`` tuple without re-searching. |
| 1685 | """ |
| 1686 | # 5 total = 4 usable (null_block takes 1) |
| 1687 | fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False) |
| 1688 | sched = fix.scheduler |
| 1689 | |
| 1690 | # --- Step 1: Store req_a's 2 blocks to CPU cache --- |
| 1691 | req_a = make_request(num_blocks=2) |
| 1692 | kv_a = _alloc_and_register(fix, req_a, 2) |
| 1693 | sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0) |
| 1694 | sched_out_a = make_scheduler_output( |
| 1695 | {req_a.request_id: 2 * BLOCK_SIZE}, |
| 1696 | new_reqs={req_a.request_id: kv_a.get_block_ids()}, |
| 1697 | ) |
| 1698 | meta_a = sched.build_connector_meta(sched_out_a) |
| 1699 | assert meta_a.store_event >= 0 |
| 1700 | simulate_store_completion(sched, meta_a.store_event) |
| 1701 | |
| 1702 | # --- Step 2: Phase A — req_b (same prompt as req_a) reports a CPU hit --- |
| 1703 | req_b = Request( |
| 1704 | request_id="req-b-toctou", |
| 1705 | prompt_token_ids=req_a.prompt_token_ids, |
| 1706 | sampling_params=req_a.sampling_params, |
| 1707 | pooling_params=None, |
| 1708 | mm_features=None, |
| 1709 | block_hasher=req_a._block_hasher, |
| 1710 | ) |
| 1711 | hit_tokens, is_async = sched.get_num_new_matched_tokens( |
| 1712 | req_b, num_computed_tokens=0 |
| 1713 | ) |
| 1714 | assert hit_tokens == 2 * BLOCK_SIZE, ( |
| 1715 | f"Phase A should report 2 blocks of CPU hit, got {hit_tokens}" |
| 1716 | ) |
| 1717 | assert is_async is True |
| 1718 | |
| 1719 | # --- Step 3: TOCTOU window — fill CPU cache so LRU evicts req_a's blocks |
| 1720 | # (in production this corresponds to other concurrent requests landing |
| 1721 | # between Phase A and Phase B for req_b). 4 usable slots, req_a occupies 2; |
| 1722 | # req_c (2) + req_d (2) require evicting 2 LRU blocks. --- |
| 1723 | req_c = make_request(num_blocks=2) |
| 1724 | req_d = make_request(num_blocks=2) |
| 1725 | kv_c = _alloc_and_register(fix, req_c, 2) |
| 1726 | kv_d = _alloc_and_register(fix, req_d, 2) |
| 1727 | sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0) |
| 1728 | sched.update_state_after_alloc(req_d, kv_d, num_external_tokens=0) |
| 1729 | sched_out_pressure = make_scheduler_output( |
| 1730 | { |
| 1731 | req_c.request_id: 2 * BLOCK_SIZE, |
| 1732 | req_d.request_id: 2 * BLOCK_SIZE, |
| 1733 | }, |
| 1734 | new_reqs={ |
| 1735 | req_c.request_id: kv_c.get_block_ids(), |
| 1736 | req_d.request_id: kv_d.get_block_ids(), |
| 1737 | }, |
| 1738 | ) |
| 1739 | meta_pressure = sched.build_connector_meta(sched_out_pressure) |
| 1740 | assert meta_pressure.store_event >= 0 |
| 1741 | simulate_store_completion(sched, meta_pressure.store_event) |
| 1742 | |
| 1743 | # --- Step 4: Phase B — must not crash --- |
| 1744 | # Before fix: AssertionError: Expected 32 hit tokens, got 0 |
| 1745 | # After fix: Phase A pinned the hits → Step 3 evicted req_c instead, |
| 1746 | # Phase B consumes the cached (cpu_hit_blocks, hit_length) tuple. |
| 1747 | gpu_blocks_b = fix.gpu_block_pool.get_new_blocks(2) |
| 1748 | kv_blocks_b = KVCacheBlocks(blocks=(gpu_blocks_b,)) |
| 1749 | sched.update_state_after_alloc(req_b, kv_blocks_b, num_external_tokens=hit_tokens) |
| 1750 | |
| 1751 | # --- Step 5: with the fix, the load is queued correctly --- |
| 1752 | sched_out_b = make_scheduler_output( |
| 1753 | {req_b.request_id: 1}, |
| 1754 | new_reqs={req_b.request_id: kv_blocks_b.get_block_ids()}, |
| 1755 | ) |
| 1756 | meta_b = sched.build_connector_meta(sched_out_b) |
| 1757 | assert meta_b.load_event >= 0, ( |
| 1758 | "Phase B should queue a load event using the pinned CPU hit blocks" |
| 1759 | ) |
| 1760 | assert len(meta_b.load_gpu_blocks) == 2 |
| 1761 | assert len(meta_b.load_cpu_blocks) == 2 |
| 1762 | |
| 1763 | |
| 1764 | # --------------------------------------------------------------------------- |
| 1765 | # Test 12: Reset with pending eager stores waits for completion |
| 1766 | # --------------------------------------------------------------------------- |
| 1767 | def test_reset_pending_eager_stores() -> None: |
| 1768 | """Eager mode: reset() abandons in-flight stores until they complete.""" |
| 1769 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1770 | sched = fix.scheduler |
| 1771 | gpu_pool = fix.gpu_block_pool |
| 1772 | |
| 1773 | num_blocks = 2 |
| 1774 | req = make_request(num_blocks=num_blocks) |
| 1775 | |
| 1776 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 1777 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1778 | block_ids = kv_blocks.get_block_ids() |
| 1779 | sched_out = make_scheduler_output( |
| 1780 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1781 | new_reqs={req.request_id: block_ids}, |
| 1782 | ) |
| 1783 | |
| 1784 | meta = sched.build_connector_meta(sched_out) |
| 1785 | assert meta.store_event >= 0 |
| 1786 | assert len(sched._store_event_to_blocks) > 0 |
| 1787 | |
| 1788 | # GPU blocks should have elevated ref_cnt from touch() |
| 1789 | for bid in meta.store_gpu_blocks: |
| 1790 | assert gpu_pool.blocks[bid].ref_cnt > 0 |
| 1791 | |
| 1792 | # Free the request's own block refs (simulates preemption) |
| 1793 | gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids[0]) |
| 1794 | |
| 1795 | # Reset should keep DMA refs pinned until the worker reports completion. |
| 1796 | assert sched.reset() is False |
| 1797 | assert len(sched._store_event_to_blocks) == 0 |
| 1798 | assert len(sched._abandoned_store_event_to_blocks) == 1 |
| 1799 | assert len(sched._reqs_to_store) == 0 |
| 1800 | assert len(sched._store_event_to_reqs) == 0 |
| 1801 | |
| 1802 | num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks() |
| 1803 | assert num_used > 1 |
| 1804 | |
| 1805 | simulate_store_completion(sched, meta.store_event) |
| 1806 | assert len(sched._abandoned_store_event_to_blocks) == 0 |
| 1807 | |
| 1808 | # All GPU blocks should now be free (ref_cnt == 0) except null block |
| 1809 | num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks() |
| 1810 | assert num_used == 1, f"Expected only null block in use, got {num_used}" |
| 1811 | |
| 1812 | # GPU prefix cache reset should now succeed |
| 1813 | assert gpu_pool.reset_prefix_cache() is True |
| 1814 | assert sched.reset() is True |
| 1815 | |
| 1816 | |
| 1817 | # --------------------------------------------------------------------------- |
| 1818 | # Test 13: Reset with pending lazy stores waits for completion |
| 1819 | # --------------------------------------------------------------------------- |
| 1820 | def test_reset_pending_lazy_stores() -> None: |
| 1821 | """Lazy mode: reset() abandons in-flight stores until they complete.""" |
| 1822 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True) |
| 1823 | sched = fix.scheduler |
| 1824 | gpu_pool = fix.gpu_block_pool |
| 1825 | |
| 1826 | num_blocks = 2 |
| 1827 | req = make_request(num_blocks=num_blocks) |
| 1828 | |
| 1829 | # Allocate, hash, and free — blocks move to free queue with hashes |
| 1830 | gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 1831 | gpu_pool.free_blocks(gpu_blocks) |
| 1832 | |
| 1833 | # Push hashed blocks to LRU head |
| 1834 | fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5) |
| 1835 | |
| 1836 | # Lazy scanner offloads old hashed blocks |
| 1837 | sched_out = make_scheduler_output({}) |
| 1838 | meta = sched.build_connector_meta(sched_out) |
| 1839 | assert meta.store_event >= 0 |
| 1840 | assert len(sched._store_event_to_blocks) > 0 |
| 1841 | |
| 1842 | gpu_pool.free_blocks(fillers) |
| 1843 | |
| 1844 | # Reset should keep DMA refs pinned until the worker reports completion. |
| 1845 | assert sched.reset() is False |
| 1846 | assert len(sched._store_event_to_blocks) == 0 |
| 1847 | assert len(sched._abandoned_store_event_to_blocks) == 1 |
| 1848 | assert sched._cursor is None |
| 1849 | |
| 1850 | simulate_store_completion(sched, meta.store_event) |
| 1851 | assert len(sched._abandoned_store_event_to_blocks) == 0 |
| 1852 | assert sched.reset() is True |
| 1853 | |
| 1854 | # No CPU cache hits after reset |
| 1855 | req2 = Request( |
| 1856 | request_id="req-after-lazy-reset", |
| 1857 | prompt_token_ids=req.prompt_token_ids, |
| 1858 | sampling_params=req.sampling_params, |
| 1859 | pooling_params=None, |
| 1860 | mm_features=None, |
| 1861 | block_hasher=req._block_hasher, |
| 1862 | ) |
| 1863 | hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 1864 | assert hit_tokens == 0, "CPU cache should be empty after reset" |
| 1865 | |
| 1866 | |
| 1867 | # --------------------------------------------------------------------------- |
| 1868 | # Test 14: Reset with pending loads waits for completion |
| 1869 | # --------------------------------------------------------------------------- |
| 1870 | def test_reset_pending_loads() -> None: |
| 1871 | """reset() abandons in-flight loads until they complete.""" |
| 1872 | fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False) |
| 1873 | sched = fix.scheduler |
| 1874 | gpu_pool = fix.gpu_block_pool |
| 1875 | |
| 1876 | num_blocks = 2 |
| 1877 | |
| 1878 | # First store blocks to CPU |
| 1879 | req = make_request(num_blocks=num_blocks) |
| 1880 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 1881 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 1882 | block_ids = kv_blocks.get_block_ids() |
| 1883 | sched_out = make_scheduler_output( |
| 1884 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 1885 | new_reqs={req.request_id: block_ids}, |
| 1886 | ) |
| 1887 | meta = sched.build_connector_meta(sched_out) |
| 1888 | simulate_store_completion(sched, meta.store_event) |
| 1889 | |
| 1890 | # Start a load — CPU cache hit |
| 1891 | req2 = Request( |
| 1892 | request_id="req-load-reset", |
| 1893 | prompt_token_ids=req.prompt_token_ids, |
| 1894 | sampling_params=req.sampling_params, |
| 1895 | pooling_params=None, |
| 1896 | mm_features=None, |
| 1897 | block_hasher=req._block_hasher, |
| 1898 | ) |
| 1899 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 1900 | assert hit_tokens is not None and hit_tokens > 0 |
| 1901 | |
| 1902 | gpu_blocks2 = gpu_pool.get_new_blocks(num_blocks) |
| 1903 | kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) |
| 1904 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 1905 | |
| 1906 | block_ids2 = kv_blocks2.get_block_ids() |
| 1907 | sched_out2 = make_scheduler_output( |
| 1908 | {req2.request_id: 1}, |
| 1909 | new_reqs={req2.request_id: block_ids2}, |
| 1910 | ) |
| 1911 | meta2 = sched.build_connector_meta(sched_out2) |
| 1912 | assert meta2.load_event >= 0 |
| 1913 | assert req2.request_id in sched._reqs_to_load |
| 1914 | |
| 1915 | # Free request block refs (simulates preemption) |
| 1916 | gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids[0]) |
| 1917 | gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids2[0]) |
| 1918 | |
| 1919 | # Reset should keep load touch refs until the worker reports completion. |
| 1920 | assert sched.reset() is False |
| 1921 | assert len(sched._reqs_to_load) == 0 |
| 1922 | assert len(sched._abandoned_reqs_to_load) == 1 |
| 1923 | assert len(sched._load_event_to_reqs) == 1 |
| 1924 | |
| 1925 | num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks() |
| 1926 | assert num_used > 1 |
| 1927 | |
| 1928 | simulate_load_completion(sched, {req2.request_id}) |
| 1929 | assert len(sched._abandoned_reqs_to_load) == 0 |
| 1930 | assert len(sched._load_event_to_reqs) == 0 |
| 1931 | assert sched.reset() is True |
| 1932 | |
| 1933 | # All GPU blocks free |
| 1934 | num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks() |
| 1935 | assert num_used == 1, f"Expected only null block in use, got {num_used}" |
| 1936 | |
| 1937 | |
| 1938 | def _make_cp_vllm_config( |
| 1939 | dcp_world_size: int = 1, |
| 1940 | pcp_world_size: int = 1, |
| 1941 | ) -> VllmConfig: |
| 1942 | """VllmConfig with context-parallel sizes set for scheduler-only tests.""" |
| 1943 | cfg = _make_vllm_config() |
| 1944 | |
| 1945 | cfg.parallel_config.decode_context_parallel_size = dcp_world_size |
| 1946 | cfg.parallel_config.prefill_context_parallel_size = pcp_world_size |
| 1947 | return cfg |
| 1948 | |
| 1949 | |
| 1950 | def _make_cp_scheduler( |
| 1951 | *, |
| 1952 | dcp_world_size: int = 1, |
| 1953 | pcp_world_size: int = 1, |
| 1954 | num_cpu_blocks: int = 8, |
| 1955 | num_gpu_blocks: int = 16, |
| 1956 | lazy: bool = False, |
| 1957 | ) -> SchedulerFixture: |
| 1958 | """Build a SimpleCPUOffloadScheduler with DCP-scaled block size.""" |
| 1959 | virtual_block_size = BLOCK_SIZE * dcp_world_size |
| 1960 | |
| 1961 | kv_cache_config = _make_kv_cache_config(num_gpu_blocks) |
| 1962 | vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size) |
| 1963 | cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks |
| 1964 | |
| 1965 | sched = SimpleCPUOffloadScheduler( |
| 1966 | vllm_config=vllm_config, |
| 1967 | kv_cache_config=kv_cache_config, |
| 1968 | cpu_capacity_bytes=cpu_capacity_bytes, |
| 1969 | scheduler_block_size=virtual_block_size, |
| 1970 | hash_block_size=virtual_block_size, |
| 1971 | lazy_offload=lazy, |
| 1972 | ) |
| 1973 | |
| 1974 | gpu_block_pool = BlockPool( |
| 1975 | num_gpu_blocks=num_gpu_blocks, |
| 1976 | enable_caching=True, |
| 1977 | hash_block_size=virtual_block_size, |
| 1978 | ) |
| 1979 | sched.bind_gpu_block_pool(gpu_block_pool) |
| 1980 | |
| 1981 | return SchedulerFixture( |
| 1982 | scheduler=sched, |
| 1983 | gpu_block_pool=gpu_block_pool, |
| 1984 | vllm_config=vllm_config, |
| 1985 | kv_cache_config=kv_cache_config, |
| 1986 | ) |
| 1987 | |
| 1988 | |
| 1989 | def _make_cp_request( |
| 1990 | num_blocks: int, |
| 1991 | virtual_block_size: int, |
| 1992 | request_id: str | None = None, |
| 1993 | ) -> Request: |
| 1994 | """Create a request whose block hashes are computed at the virtual |
| 1995 | (CP-scaled) block size, matching what the real scheduler does. |
| 1996 | """ |
| 1997 | global _req_counter |
| 1998 | _req_counter += 1 |
| 1999 | if request_id is None: |
| 2000 | request_id = f"req-cp-{_req_counter}" |
| 2001 | |
| 2002 | num_tokens = num_blocks * virtual_block_size + 1 |
| 2003 | start = _req_counter * 10000 |
| 2004 | prompt_token_ids = list(range(start, start + num_tokens)) |
| 2005 | sampling_params = SamplingParams(max_tokens=1) |
| 2006 | |
| 2007 | return Request( |
| 2008 | request_id=request_id, |
| 2009 | prompt_token_ids=prompt_token_ids, |
| 2010 | sampling_params=sampling_params, |
| 2011 | pooling_params=None, |
| 2012 | mm_features=None, |
| 2013 | block_hasher=get_request_block_hasher(virtual_block_size, sha256), |
| 2014 | ) |
| 2015 | |
| 2016 | |
| 2017 | def _allocate_cp_gpu_blocks( |
| 2018 | gpu_block_pool: BlockPool, |
| 2019 | request: Request, |
| 2020 | num_blocks: int, |
| 2021 | virtual_block_size: int, |
| 2022 | group_id: int = 0, |
| 2023 | ) -> list: |
| 2024 | """Allocate GPU blocks and cache them using the CP-scaled block size.""" |
| 2025 | blocks = gpu_block_pool.get_new_blocks(num_blocks) |
| 2026 | num_full = min(num_blocks, len(request.block_hashes)) |
| 2027 | if num_full > 0: |
| 2028 | gpu_block_pool.cache_full_blocks( |
| 2029 | request=request, |
| 2030 | blocks=blocks, |
| 2031 | num_cached_blocks=0, |
| 2032 | num_full_blocks=num_full, |
| 2033 | block_size=virtual_block_size, |
| 2034 | kv_cache_group_id=group_id, |
| 2035 | ) |
| 2036 | return blocks |
| 2037 | |
| 2038 | |
| 2039 | # --------------------------------------------------------------------------- |
| 2040 | # Test 15: CP block size scaling is correct |
| 2041 | # --------------------------------------------------------------------------- |
| 2042 | @pytest.mark.parametrize( |
| 2043 | "dcp_world_size, pcp_world_size", |
| 2044 | [ |
| 2045 | (2, 1), # DCP only |
| 2046 | (1, 2), # PCP only |
| 2047 | (2, 2), # DCP + PCP |
| 2048 | ], |
| 2049 | ) |
| 2050 | def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None: |
| 2051 | """Verify block size scaling follows DCP ownership.""" |
| 2052 | fix = _make_cp_scheduler( |
| 2053 | dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size |
| 2054 | ) |
| 2055 | sched = fix.scheduler |
| 2056 | |
| 2057 | assert sched.cp_world_size == dcp_world_size |
| 2058 | assert sched.block_size == BLOCK_SIZE * dcp_world_size |
| 2059 | |
| 2060 | |
| 2061 | # --------------------------------------------------------------------------- |
| 2062 | # Test 16: CP eager store-and-load roundtrip |
| 2063 | # --------------------------------------------------------------------------- |
| 2064 | @pytest.mark.parametrize( |
| 2065 | "dcp_world_size, pcp_world_size", |
| 2066 | [ |
| 2067 | (2, 1), |
| 2068 | (1, 2), |
| 2069 | ], |
| 2070 | ) |
| 2071 | def test_cp_eager_store_and_load_roundtrip( |
| 2072 | dcp_world_size: int, pcp_world_size: int |
| 2073 | ) -> None: |
| 2074 | """With CP enabled, store blocks to CPU and reload them for a new request |
| 2075 | with matching tokens. Verifies that hash matching and transfer-pair |
| 2076 | construction work with the virtual block size.""" |
| 2077 | fix = _make_cp_scheduler( |
| 2078 | dcp_world_size=dcp_world_size, |
| 2079 | pcp_world_size=pcp_world_size, |
| 2080 | num_cpu_blocks=8, |
| 2081 | num_gpu_blocks=16, |
| 2082 | lazy=False, |
| 2083 | ) |
| 2084 | sched = fix.scheduler |
| 2085 | vbs = BLOCK_SIZE * dcp_world_size |
| 2086 | |
| 2087 | num_blocks = 2 |
| 2088 | req = _make_cp_request(num_blocks, vbs) |
| 2089 | |
| 2090 | # Allocate GPU blocks and register hashes |
| 2091 | gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs) |
| 2092 | kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,)) |
| 2093 | req.num_computed_tokens = num_blocks * vbs |
| 2094 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2095 | |
| 2096 | block_ids = kv_blocks.get_block_ids() |
| 2097 | sched_out = make_scheduler_output( |
| 2098 | {req.request_id: num_blocks * vbs}, |
| 2099 | new_reqs={req.request_id: block_ids}, |
| 2100 | ) |
| 2101 | |
| 2102 | meta = sched.build_connector_meta(sched_out) |
| 2103 | assert meta.store_event >= 0, "Expected a store event" |
| 2104 | assert len(meta.store_gpu_blocks) == num_blocks |
| 2105 | assert len(meta.store_cpu_blocks) == num_blocks |
| 2106 | simulate_store_completion(sched, meta.store_event) |
| 2107 | |
| 2108 | # New request with same tokens — should get a full CPU cache hit. |
| 2109 | req2 = Request( |
| 2110 | request_id="req-cp-load", |
| 2111 | prompt_token_ids=req.prompt_token_ids, |
| 2112 | sampling_params=req.sampling_params, |
| 2113 | pooling_params=None, |
| 2114 | mm_features=None, |
| 2115 | block_hasher=req._block_hasher, |
| 2116 | ) |
| 2117 | |
| 2118 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 2119 | assert hit_tokens == num_blocks * vbs |
| 2120 | assert is_async is True |
| 2121 | |
| 2122 | # Allocate fresh GPU blocks for the load. |
| 2123 | gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks) |
| 2124 | kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) |
| 2125 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 2126 | |
| 2127 | sched_out2 = make_scheduler_output( |
| 2128 | {req2.request_id: 1}, |
| 2129 | new_reqs={req2.request_id: kv_blocks2.get_block_ids()}, |
| 2130 | ) |
| 2131 | meta2 = sched.build_connector_meta(sched_out2) |
| 2132 | assert meta2.load_event >= 0, "Expected a load event" |
| 2133 | assert len(meta2.load_gpu_blocks) == num_blocks |
| 2134 | assert len(meta2.load_cpu_blocks) == num_blocks |
| 2135 | |
| 2136 | |
| 2137 | # --------------------------------------------------------------------------- |
| 2138 | # Test 18: CP store and load use effective block size |
| 2139 | # --------------------------------------------------------------------------- |
| 2140 | @pytest.mark.parametrize( |
| 2141 | "dcp_world_size, pcp_world_size", |
| 2142 | [ |
| 2143 | (2, 1), |
| 2144 | (1, 2), |
| 2145 | (2, 2), |
| 2146 | ], |
| 2147 | ) |
| 2148 | def test_cp_effective_block_size_store_and_load( |
| 2149 | dcp_world_size: int, pcp_world_size: int |
| 2150 | ) -> None: |
| 2151 | """Verify store/load physical block counts scale with DCP, not PCP.""" |
| 2152 | fix = _make_cp_scheduler( |
| 2153 | dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size |
| 2154 | ) |
| 2155 | sched = fix.scheduler |
| 2156 | gpu_pool = fix.gpu_block_pool |
| 2157 | vbs = BLOCK_SIZE * dcp_world_size |
| 2158 | expected_blocks = 1 |
| 2159 | |
| 2160 | # Store one DCP-scaled logical block worth of tokens. |
| 2161 | req = _make_cp_request(num_blocks=2, virtual_block_size=vbs) |
| 2162 | gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs) |
| 2163 | kv = KVCacheBlocks(blocks=(gpu_blocks,)) |
| 2164 | req.num_computed_tokens = vbs |
| 2165 | sched.update_state_after_alloc(req, kv, num_external_tokens=0) |
| 2166 | m1 = sched.build_connector_meta( |
| 2167 | make_scheduler_output( |
| 2168 | {req.request_id: vbs}, |
| 2169 | new_reqs={req.request_id: kv.get_block_ids()}, |
| 2170 | ) |
| 2171 | ) |
| 2172 | assert len(m1.store_gpu_blocks) == expected_blocks |
| 2173 | assert len(m1.store_cpu_blocks) == expected_blocks |
| 2174 | simulate_store_completion(sched, m1.store_event) |
| 2175 | |
| 2176 | # Load one DCP-scaled logical block from a two-block CPU hit. |
| 2177 | req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs) |
| 2178 | kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),)) |
| 2179 | req2.num_computed_tokens = 2 * vbs |
| 2180 | sched.update_state_after_alloc(req2, kv2, num_external_tokens=0) |
| 2181 | m2 = sched.build_connector_meta( |
| 2182 | make_scheduler_output( |
| 2183 | {req2.request_id: 2 * vbs}, |
| 2184 | new_reqs={req2.request_id: kv2.get_block_ids()}, |
| 2185 | ) |
| 2186 | ) |
| 2187 | simulate_store_completion(sched, m2.store_event) |
| 2188 | |
| 2189 | req3 = Request( |
| 2190 | request_id="req-cp-partial-load", |
| 2191 | prompt_token_ids=req2.prompt_token_ids, |
| 2192 | sampling_params=req2.sampling_params, |
| 2193 | pooling_params=None, |
| 2194 | mm_features=None, |
| 2195 | block_hasher=req2._block_hasher, |
| 2196 | ) |
| 2197 | hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0) |
| 2198 | assert hit == 2 * vbs |
| 2199 | |
| 2200 | kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(expected_blocks),)) |
| 2201 | sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs) |
| 2202 | m3 = sched.build_connector_meta( |
| 2203 | make_scheduler_output( |
| 2204 | {req3.request_id: vbs}, |
| 2205 | new_reqs={req3.request_id: kv3.get_block_ids()}, |
| 2206 | ) |
| 2207 | ) |
| 2208 | assert m3.load_event >= 0 |
| 2209 | assert len(m3.load_gpu_blocks) == expected_blocks |
| 2210 | assert len(m3.load_cpu_blocks) == expected_blocks |
| 2211 | assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]] |
| 2212 | |
| 2213 | |
| 2214 | # --------------------------------------------------------------------------- |
| 2215 | # Test 17: CP lazy target blocks are scaled correctly |
| 2216 | # --------------------------------------------------------------------------- |
| 2217 | @pytest.mark.parametrize("cp_world_size", [1, 2, 4]) |
| 2218 | def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None: |
| 2219 | """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1 |
| 2220 | because each virtual block covers more tokens.""" |
| 2221 | kv_cache_config = _make_kv_cache_config(num_blocks=16) |
| 2222 | max_batched = 64 |
| 2223 | |
| 2224 | target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( |
| 2225 | kv_cache_config, max_batched, cp_world_size=1 |
| 2226 | ) |
| 2227 | target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( |
| 2228 | kv_cache_config, max_batched, cp_world_size=cp_world_size |
| 2229 | ) |
| 2230 | |
| 2231 | if cp_world_size == 1: |
| 2232 | assert target_cp == target_base |
| 2233 | else: |
| 2234 | assert target_cp < target_base, ( |
| 2235 | f"cp_world_size={cp_world_size}: target_cp={target_cp} should be " |
| 2236 | f"less than target_base={target_base}" |
| 2237 | ) |
| 2238 | |
| 2239 | |
| 2240 | def _make_hybrid_attention_mamba_scheduler( |
| 2241 | *, |
| 2242 | num_cpu_blocks: int = 8, |
| 2243 | num_gpu_blocks: int = 16, |
| 2244 | attention_block_size: int = BLOCK_SIZE, |
| 2245 | block_size: int = 4 * BLOCK_SIZE, |
| 2246 | scheduler_block_size: int | None = None, |
| 2247 | hash_block_size: int | None = None, |
| 2248 | dcp_world_size: int = 4, |
| 2249 | lazy: bool = False, |
| 2250 | mamba_cache_mode: MambaCacheMode = "align", |
| 2251 | enable_kv_cache_events: bool = False, |
| 2252 | ) -> SchedulerFixture: |
| 2253 | """Build a scheduler for one attention group plus one Mamba group.""" |
| 2254 | scheduler_block_size = scheduler_block_size or block_size |
| 2255 | hash_block_size = hash_block_size or block_size |
| 2256 | attention_spec = FullAttentionSpec( |
| 2257 | block_size=attention_block_size, |
| 2258 | num_kv_heads=NUM_KV_HEADS, |
| 2259 | head_size=HEAD_SIZE, |
| 2260 | dtype=DTYPE, |
| 2261 | ) |
| 2262 | mamba_spec = MambaSpec( |
| 2263 | block_size=block_size, |
| 2264 | shapes=((1, 1),), |
| 2265 | dtypes=(torch.float32,), |
| 2266 | mamba_cache_mode=mamba_cache_mode, |
| 2267 | ) |
| 2268 | groups = [ |
| 2269 | KVCacheGroupSpec(["attention"], attention_spec), |
| 2270 | KVCacheGroupSpec(["mamba"], mamba_spec), |
| 2271 | ] |
| 2272 | tensors = [ |
| 2273 | KVCacheTensor( |
| 2274 | size=spec.page_size_bytes * num_gpu_blocks, |
| 2275 | layers=group.layer_names, |
| 2276 | layer_stride=spec.page_size_bytes * num_gpu_blocks, |
| 2277 | block_stride=spec.page_size_bytes, |
| 2278 | ) |
| 2279 | for group, spec in zip(groups, (attention_spec, mamba_spec)) |
| 2280 | ] |
| 2281 | kv_cache_config = KVCacheConfig( |
| 2282 | num_blocks=num_gpu_blocks, |
| 2283 | kv_cache_tensors=tensors, |
| 2284 | kv_cache_groups=groups, |
| 2285 | ) |
| 2286 | vllm_config = _make_cp_vllm_config(dcp_world_size=dcp_world_size) |
| 2287 | vllm_config.cache_config.prefix_cache_retention_interval = 0 |
| 2288 | vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode |
| 2289 | if enable_kv_cache_events: |
| 2290 | vllm_config.kv_events_config = KVEventsConfig(enable_kv_cache_events=True) |
| 2291 | # _derive_cpu_config() scales the requested capacity against |
| 2292 | # kv_cache_tensors[0].size, so express it as a multiple of that tensor's |
| 2293 | # per-block size to get exactly num_cpu_blocks. |
| 2294 | cpu_capacity_bytes = tensors[0].size // num_gpu_blocks * num_cpu_blocks |
| 2295 | sched = SimpleCPUOffloadScheduler( |
| 2296 | vllm_config=vllm_config, |
| 2297 | kv_cache_config=kv_cache_config, |
| 2298 | cpu_capacity_bytes=cpu_capacity_bytes, |
| 2299 | scheduler_block_size=scheduler_block_size, |
| 2300 | hash_block_size=hash_block_size, |
| 2301 | lazy_offload=lazy, |
| 2302 | ) |
| 2303 | gpu_block_pool = BlockPool( |
| 2304 | num_gpu_blocks=num_gpu_blocks, |
| 2305 | enable_caching=True, |
| 2306 | hash_block_size=hash_block_size, |
| 2307 | ) |
| 2308 | sched.bind_gpu_block_pool(gpu_block_pool) |
| 2309 | return SchedulerFixture( |
| 2310 | scheduler=sched, |
| 2311 | gpu_block_pool=gpu_block_pool, |
| 2312 | vllm_config=vllm_config, |
| 2313 | kv_cache_config=kv_cache_config, |
| 2314 | ) |
| 2315 | |
| 2316 | |
| 2317 | def test_hybrid_store_uses_resolved_group_block_sizes() -> None: |
| 2318 | """Replicated groups must not be scaled by the DCP world size. |
| 2319 | |
| 2320 | The spec's ``dcp_sharded`` flag controls its effective block size. |
| 2321 | Reconstructing a group's geometry as |
| 2322 | ``spec.block_size * cp_world_size`` over-scales the replicated groups, so |
| 2323 | the eager store scan believes far fewer of their blocks are ready and |
| 2324 | silently offloads only a fraction of them. |
| 2325 | """ |
| 2326 | attention_block_size = BLOCK_SIZE |
| 2327 | mamba_block_size = 4 * BLOCK_SIZE |
| 2328 | dcp_world_size = 2 |
| 2329 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2330 | num_cpu_blocks=32, |
| 2331 | num_gpu_blocks=32, |
| 2332 | attention_block_size=attention_block_size, |
| 2333 | block_size=mamba_block_size, |
| 2334 | # Every prefix-cacheable group's resolved block size must be a multiple |
| 2335 | # of the hash block, so it cannot exceed the smallest of them. |
| 2336 | hash_block_size=BLOCK_SIZE, |
| 2337 | dcp_world_size=dcp_world_size, |
| 2338 | mamba_cache_mode="none", |
| 2339 | ) |
| 2340 | sched = fix.scheduler |
| 2341 | gpu_pool = fix.gpu_block_pool |
| 2342 | attention_size, mamba_size = sched.group_block_sizes |
| 2343 | assert attention_size == attention_block_size * dcp_world_size |
| 2344 | assert mamba_size == mamba_block_size |
| 2345 | |
| 2346 | confirmed = 2 * sched.block_size |
| 2347 | req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE) |
| 2348 | attention_blocks = _allocate_cp_gpu_blocks( |
| 2349 | gpu_pool, req, confirmed // attention_size, attention_size, group_id=0 |
| 2350 | ) |
| 2351 | mamba_blocks = _allocate_cp_gpu_blocks( |
| 2352 | gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1 |
| 2353 | ) |
| 2354 | kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) |
| 2355 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2356 | sched.build_connector_meta( |
| 2357 | make_scheduler_output( |
| 2358 | {req.request_id: confirmed}, |
| 2359 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 2360 | ) |
| 2361 | ) |
| 2362 | req.num_computed_tokens = confirmed |
| 2363 | meta = sched.build_connector_meta( |
| 2364 | make_scheduler_output( |
| 2365 | {req.request_id: 1}, |
| 2366 | cached_req_new_blocks={req.request_id: None}, |
| 2367 | ) |
| 2368 | ) |
| 2369 | |
| 2370 | # Both groups are positionally stable here, so every confirmed block of |
| 2371 | # both must be offloaded. Over-scaling the replicated group halves its |
| 2372 | # count. |
| 2373 | assert sched._reqs_to_store[req.request_id].num_stored_blocks == [ |
| 2374 | confirmed // attention_size, |
| 2375 | confirmed // mamba_size, |
| 2376 | ] |
| 2377 | assert set(meta.store_gpu_blocks) == { |
| 2378 | *(b.block_id for b in attention_blocks), |
| 2379 | *(b.block_id for b in mamba_blocks), |
| 2380 | } |
| 2381 | |
| 2382 | |
| 2383 | def test_dcp_mixed_cache_loads_fine_grained_external_hit() -> None: |
| 2384 | """Map a fine-grained external suffix after a nonzero local prefix.""" |
| 2385 | hash_block_size = BLOCK_SIZE |
| 2386 | mamba_block_size = 4 * BLOCK_SIZE |
| 2387 | local_tokens = 4 * BLOCK_SIZE |
| 2388 | # Span two Mamba blocks and end partially in the second one. This exercises |
| 2389 | # both cdiv(external_tokens, mamba_block_size) > 1 and the partial tail. |
| 2390 | external_tokens = 6 * BLOCK_SIZE |
| 2391 | total_cached_tokens = local_tokens + external_tokens |
| 2392 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2393 | num_cpu_blocks=32, |
| 2394 | num_gpu_blocks=32, |
| 2395 | block_size=mamba_block_size, |
| 2396 | hash_block_size=hash_block_size, |
| 2397 | dcp_world_size=2, |
| 2398 | ) |
| 2399 | sched = fix.scheduler |
| 2400 | gpu_pool = fix.gpu_block_pool |
| 2401 | |
| 2402 | assert external_tokens % mamba_block_size != 0 |
| 2403 | assert sched.cpu_coordinator.enable_partial_hash_hits |
| 2404 | |
| 2405 | producer = _make_cp_request(num_blocks=10, virtual_block_size=hash_block_size) |
| 2406 | attention_blocks = _allocate_cp_gpu_blocks( |
| 2407 | gpu_pool, |
| 2408 | producer, |
| 2409 | num_blocks=5, |
| 2410 | virtual_block_size=2 * BLOCK_SIZE, |
| 2411 | group_id=0, |
| 2412 | ) |
| 2413 | mamba_blocks = gpu_pool.get_new_blocks(3) |
| 2414 | gpu_pool.cache_full_blocks( |
| 2415 | request=producer, |
| 2416 | blocks=mamba_blocks, |
| 2417 | num_cached_blocks=0, |
| 2418 | num_full_blocks=2, |
| 2419 | block_size=mamba_block_size, |
| 2420 | kv_cache_group_id=1, |
| 2421 | ) |
| 2422 | gpu_pool.cache_partial_block( |
| 2423 | request=producer, |
| 2424 | block=mamba_blocks[2], |
| 2425 | num_tokens=total_cached_tokens, |
| 2426 | kv_cache_group_id=1, |
| 2427 | block_size=mamba_block_size, |
| 2428 | ) |
| 2429 | producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) |
| 2430 | sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0) |
| 2431 | |
| 2432 | # Positional scanning stores append-only FA blocks. Mamba align blocks use |
| 2433 | # exact handoffs for every retained boundary, including the partial tail. |
| 2434 | store_output = make_scheduler_output( |
| 2435 | {producer.request_id: total_cached_tokens}, |
| 2436 | new_reqs={producer.request_id: producer_blocks.get_block_ids()}, |
| 2437 | ) |
| 2438 | store_output.kv_connector_block_state = KVConnectorBlockState( |
| 2439 | req_ids=set(), |
| 2440 | resolve_block_ids=lambda _: (), |
| 2441 | boundary_state_offloads={ |
| 2442 | producer.request_id: [ |
| 2443 | (1, mamba_blocks[0].block_id, mamba_block_size), |
| 2444 | (1, mamba_blocks[1].block_id, 2 * mamba_block_size), |
| 2445 | (1, mamba_blocks[2].block_id, total_cached_tokens), |
| 2446 | ] |
| 2447 | }, |
| 2448 | ) |
| 2449 | store_meta = sched.build_connector_meta(store_output) |
| 2450 | # The exact handoffs land in the step they arrive; the append-only FA scan |
| 2451 | # picks its blocks up once the scheduler has committed the tokens. |
| 2452 | assert set(store_meta.store_gpu_blocks) == { |
| 2453 | block.block_id for block in mamba_blocks |
| 2454 | } |
| 2455 | simulate_store_completion(sched, store_meta.store_event) |
| 2456 | producer.num_computed_tokens = total_cached_tokens |
| 2457 | fa_meta = sched.build_connector_meta( |
| 2458 | make_scheduler_output( |
| 2459 | {producer.request_id: 1}, |
| 2460 | cached_req_new_blocks={producer.request_id: None}, |
| 2461 | ) |
| 2462 | ) |
| 2463 | assert set(fa_meta.store_gpu_blocks) == { |
| 2464 | block.block_id for block in attention_blocks |
| 2465 | } |
| 2466 | simulate_store_completion(sched, fa_meta.store_event) |
| 2467 | |
| 2468 | consumer = Request( |
| 2469 | request_id="req-dcp-mixed-fine-load", |
| 2470 | prompt_token_ids=producer.prompt_token_ids, |
| 2471 | sampling_params=producer.sampling_params, |
| 2472 | pooling_params=None, |
| 2473 | mm_features=None, |
| 2474 | block_hasher=producer._block_hasher, |
| 2475 | ) |
| 2476 | hit_tokens, is_async = sched.get_num_new_matched_tokens( |
| 2477 | consumer, num_computed_tokens=local_tokens |
| 2478 | ) |
| 2479 | assert hit_tokens == external_tokens |
| 2480 | assert is_async is True |
| 2481 | pending_blocks, pending_tokens, _ = sched._pending_cpu_hits[consumer.request_id] |
| 2482 | assert pending_tokens == external_tokens |
| 2483 | assert len(pending_blocks[1]) == 2 |
| 2484 | assert pending_blocks[1][0].is_null |
| 2485 | assert not pending_blocks[1][1].is_null |
| 2486 | |
| 2487 | local_attention = attention_blocks[:2] |
| 2488 | local_mamba = mamba_blocks[:1] |
| 2489 | load_blocks = KVCacheBlocks( |
| 2490 | blocks=( |
| 2491 | local_attention + gpu_pool.get_new_blocks(3), |
| 2492 | local_mamba + gpu_pool.get_new_blocks(2), |
| 2493 | ) |
| 2494 | ) |
| 2495 | sched.update_state_after_alloc( |
| 2496 | consumer, |
| 2497 | load_blocks, |
| 2498 | num_external_tokens=hit_tokens, |
| 2499 | ) |
| 2500 | load_meta = sched.build_connector_meta( |
| 2501 | make_scheduler_output( |
| 2502 | {consumer.request_id: 1}, |
| 2503 | new_reqs={consumer.request_id: load_blocks.get_block_ids()}, |
| 2504 | ) |
| 2505 | ) |
| 2506 | assert load_meta.load_gpu_blocks == [ |
| 2507 | *(block.block_id for block in load_blocks.blocks[0][2:5]), |
| 2508 | load_blocks.blocks[1][2].block_id, |
| 2509 | ] |
| 2510 | assert len(load_meta.load_cpu_blocks) == 4 |
| 2511 | |
| 2512 | |
| 2513 | def test_finished_eager_store_caches_fa_partial_tail_for_hybrid_hit() -> None: |
| 2514 | """A finished non-LCM prompt needs the FA partial tail in CPU too.""" |
| 2515 | hash_block_size = BLOCK_SIZE |
| 2516 | mamba_block_size = BLOCK_SIZE |
| 2517 | attention_block_size = 4 * BLOCK_SIZE |
| 2518 | dcp_world_size = 2 |
| 2519 | scheduler_block_size = attention_block_size * dcp_world_size |
| 2520 | total_cached_tokens = 5 * hash_block_size |
| 2521 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2522 | num_cpu_blocks=32, |
| 2523 | num_gpu_blocks=32, |
| 2524 | attention_block_size=attention_block_size, |
| 2525 | block_size=mamba_block_size, |
| 2526 | scheduler_block_size=scheduler_block_size, |
| 2527 | hash_block_size=hash_block_size, |
| 2528 | dcp_world_size=dcp_world_size, |
| 2529 | ) |
| 2530 | sched = fix.scheduler |
| 2531 | gpu_pool = fix.gpu_block_pool |
| 2532 | assert sched.fa_block_size == scheduler_block_size |
| 2533 | assert total_cached_tokens % sched.fa_block_size != 0 |
| 2534 | |
| 2535 | producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size) |
| 2536 | attention_blocks = gpu_pool.get_new_blocks(1) |
| 2537 | gpu_pool.cache_partial_block( |
| 2538 | request=producer, |
| 2539 | block=attention_blocks[0], |
| 2540 | num_tokens=total_cached_tokens, |
| 2541 | kv_cache_group_id=0, |
| 2542 | block_size=scheduler_block_size, |
| 2543 | ) |
| 2544 | mamba_blocks = gpu_pool.get_new_blocks(5) |
| 2545 | gpu_pool.cache_full_blocks( |
| 2546 | request=producer, |
| 2547 | blocks=mamba_blocks, |
| 2548 | num_cached_blocks=0, |
| 2549 | num_full_blocks=5, |
| 2550 | kv_cache_group_id=1, |
| 2551 | block_size=mamba_block_size, |
| 2552 | ) |
| 2553 | producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) |
| 2554 | sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0) |
| 2555 | |
| 2556 | store_output = make_scheduler_output( |
| 2557 | {producer.request_id: total_cached_tokens}, |
| 2558 | new_reqs={producer.request_id: producer_blocks.get_block_ids()}, |
| 2559 | ) |
| 2560 | store_output.kv_connector_block_state = KVConnectorBlockState( |
| 2561 | req_ids=set(), |
| 2562 | resolve_block_ids=lambda _: (), |
| 2563 | boundary_state_offloads={ |
| 2564 | producer.request_id: [ |
| 2565 | (1, mamba_blocks[4].block_id, total_cached_tokens), |
| 2566 | ] |
| 2567 | }, |
| 2568 | ) |
| 2569 | store_meta = sched.build_connector_meta(store_output) |
| 2570 | simulate_store_completion(sched, store_meta.store_event) |
| 2571 | |
| 2572 | producer.num_computed_tokens = total_cached_tokens |
| 2573 | sched.request_finished_all_groups(producer, producer_blocks.get_block_ids()) |
| 2574 | finish_meta = sched.build_connector_meta(make_scheduler_output({})) |
| 2575 | assert attention_blocks[0].block_id in finish_meta.store_gpu_blocks |
| 2576 | simulate_store_completion(sched, finish_meta.store_event) |
| 2577 | |
| 2578 | consumer = Request( |
| 2579 | request_id="req-finished-partial-tail-load", |
| 2580 | prompt_token_ids=producer.prompt_token_ids, |
| 2581 | sampling_params=producer.sampling_params, |
| 2582 | pooling_params=None, |
| 2583 | mm_features=None, |
| 2584 | block_hasher=producer._block_hasher, |
| 2585 | ) |
| 2586 | hit_tokens, is_async = sched.get_num_new_matched_tokens( |
| 2587 | consumer, num_computed_tokens=0 |
| 2588 | ) |
| 2589 | assert hit_tokens == total_cached_tokens |
| 2590 | assert is_async is True |
| 2591 | |
| 2592 | |
| 2593 | def test_finished_eager_store_does_not_duplicate_completed_partial_tail() -> None: |
| 2594 | """Decode can fill the boundary block; the scan then already covers it.""" |
| 2595 | hash_block_size = BLOCK_SIZE |
| 2596 | mamba_block_size = BLOCK_SIZE |
| 2597 | attention_block_size = 4 * BLOCK_SIZE |
| 2598 | dcp_world_size = 2 |
| 2599 | scheduler_block_size = attention_block_size * dcp_world_size |
| 2600 | prompt_tokens = 5 * hash_block_size |
| 2601 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2602 | num_cpu_blocks=32, |
| 2603 | num_gpu_blocks=32, |
| 2604 | attention_block_size=attention_block_size, |
| 2605 | block_size=mamba_block_size, |
| 2606 | scheduler_block_size=scheduler_block_size, |
| 2607 | hash_block_size=hash_block_size, |
| 2608 | dcp_world_size=dcp_world_size, |
| 2609 | ) |
| 2610 | sched = fix.scheduler |
| 2611 | gpu_pool = fix.gpu_block_pool |
| 2612 | assert prompt_tokens % sched.fa_block_size != 0 |
| 2613 | |
| 2614 | producer = _make_cp_request(num_blocks=5, virtual_block_size=hash_block_size) |
| 2615 | attention_blocks = gpu_pool.get_new_blocks(1) |
| 2616 | gpu_pool.cache_partial_block( |
| 2617 | request=producer, |
| 2618 | block=attention_blocks[0], |
| 2619 | num_tokens=prompt_tokens, |
| 2620 | kv_cache_group_id=0, |
| 2621 | block_size=scheduler_block_size, |
| 2622 | ) |
| 2623 | mamba_blocks = gpu_pool.get_new_blocks(5) |
| 2624 | gpu_pool.cache_full_blocks( |
| 2625 | request=producer, |
| 2626 | blocks=mamba_blocks, |
| 2627 | num_cached_blocks=0, |
| 2628 | num_full_blocks=5, |
| 2629 | kv_cache_group_id=1, |
| 2630 | block_size=mamba_block_size, |
| 2631 | ) |
| 2632 | producer_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) |
| 2633 | sched.update_state_after_alloc(producer, producer_blocks, num_external_tokens=0) |
| 2634 | |
| 2635 | # Decode carries the request past the physical attention block boundary, so |
| 2636 | # the positional scan now considers the same block the partial tail names. |
| 2637 | producer.num_computed_tokens = sched.fa_block_size |
| 2638 | sched.request_finished_all_groups(producer, producer_blocks.get_block_ids()) |
| 2639 | finish_meta = sched.build_connector_meta(make_scheduler_output({})) |
| 2640 | |
| 2641 | stored = finish_meta.store_gpu_blocks |
| 2642 | assert attention_blocks[0].block_id in stored |
| 2643 | assert len(stored) == len(set(stored)) |
| 2644 | assert len(finish_meta.store_cpu_blocks) == len(stored) |
| 2645 | |
| 2646 | |
| 2647 | def test_external_lookup_rejects_misaligned_local_prefix(caplog_vllm) -> None: |
| 2648 | scheduler_block_size = 4 * BLOCK_SIZE |
| 2649 | fix = _make_hybrid_attention_mamba_scheduler(block_size=scheduler_block_size) |
| 2650 | |
| 2651 | request = _make_cp_request(num_blocks=2, virtual_block_size=scheduler_block_size) |
| 2652 | # The ``vllm`` logger sets propagate=False, so plain ``caplog`` can observe |
| 2653 | # nothing; ``caplog_vllm`` re-enables propagation for the duration. |
| 2654 | with caplog_vllm.at_level(logging.WARNING): |
| 2655 | hit_tokens, is_async = fix.scheduler.get_num_new_matched_tokens( |
| 2656 | request, num_computed_tokens=BLOCK_SIZE |
| 2657 | ) |
| 2658 | |
| 2659 | assert (hit_tokens, is_async) == (0, False) |
| 2660 | assert "requires scheduler-block-aligned local tokens" in caplog_vllm.text |
| 2661 | |
| 2662 | |
| 2663 | def test_fine_grained_external_hits_require_eager_offload() -> None: |
| 2664 | hash_block_size = BLOCK_SIZE |
| 2665 | scheduler_block_size = 4 * BLOCK_SIZE |
| 2666 | |
| 2667 | eager = _make_hybrid_attention_mamba_scheduler( |
| 2668 | block_size=scheduler_block_size, |
| 2669 | hash_block_size=hash_block_size, |
| 2670 | ) |
| 2671 | lazy = _make_hybrid_attention_mamba_scheduler( |
| 2672 | block_size=scheduler_block_size, |
| 2673 | hash_block_size=hash_block_size, |
| 2674 | lazy=True, |
| 2675 | ) |
| 2676 | |
| 2677 | assert eager.scheduler.cpu_coordinator.enable_partial_hash_hits |
| 2678 | assert not lazy.scheduler.cpu_coordinator.enable_partial_hash_hits |
| 2679 | |
| 2680 | |
| 2681 | def test_eager_store_does_not_scan_mamba_blocks_positionally() -> None: |
| 2682 | """Mamba stores must use exact handoffs, not historical block positions.""" |
| 2683 | block_size = 4 * BLOCK_SIZE |
| 2684 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2685 | num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size |
| 2686 | ) |
| 2687 | sched = fix.scheduler |
| 2688 | gpu_pool = fix.gpu_block_pool |
| 2689 | attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers |
| 2690 | assert attention_manager.has_positionally_stable_blocks |
| 2691 | assert not mamba_manager.has_positionally_stable_blocks |
| 2692 | |
| 2693 | req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) |
| 2694 | attn_blocks = _allocate_cp_gpu_blocks( |
| 2695 | gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 |
| 2696 | ) |
| 2697 | gpu_blocks = gpu_pool.get_new_blocks(2) |
| 2698 | gpu_blocks[1].set_block_hash(make_block_hash_with_group_id(req.block_hashes[1], 1)) |
| 2699 | kv_blocks = KVCacheBlocks(blocks=(attn_blocks, gpu_blocks)) |
| 2700 | req.num_computed_tokens = 0 |
| 2701 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2702 | |
| 2703 | output = make_scheduler_output( |
| 2704 | {req.request_id: 2 * block_size}, |
| 2705 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 2706 | ) |
| 2707 | output.kv_connector_block_state = KVConnectorBlockState( |
| 2708 | req_ids=set(), |
| 2709 | resolve_block_ids=lambda _: (), |
| 2710 | boundary_state_offloads={ |
| 2711 | req.request_id: [(1, gpu_blocks[1].block_id, 2 * block_size)] |
| 2712 | }, |
| 2713 | ) |
| 2714 | meta1 = sched.build_connector_meta(output) |
| 2715 | assert meta1.store_event >= 0 |
| 2716 | # Only the exact handoff lands this step. build_connector_meta runs before |
| 2717 | # the scheduler commits this step's tokens, so the positional scan has |
| 2718 | # nothing confirmed yet. |
| 2719 | assert set(meta1.store_gpu_blocks) == {gpu_blocks[1].block_id} |
| 2720 | assert sched._reqs_to_store[req.request_id].num_stored_blocks == [0, 0] |
| 2721 | simulate_store_completion(sched, meta1.store_event) |
| 2722 | |
| 2723 | # Next step the attention blocks are confirmed and scanned positionally. A |
| 2724 | # later hash update still must not make the mamba align snapshot a valid |
| 2725 | # source; the explicit boundary-handoff tests cover the supported path. |
| 2726 | gpu_blocks[0].set_block_hash(make_block_hash_with_group_id(req.block_hashes[0], 1)) |
| 2727 | req.num_computed_tokens = 2 * block_size |
| 2728 | meta2 = sched.build_connector_meta( |
| 2729 | make_scheduler_output( |
| 2730 | {req.request_id: 1}, |
| 2731 | cached_req_new_blocks={req.request_id: None}, |
| 2732 | ) |
| 2733 | ) |
| 2734 | assert set(meta2.store_gpu_blocks) == {block.block_id for block in attn_blocks} |
| 2735 | assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0] |
| 2736 | simulate_store_completion(sched, meta2.store_event) |
| 2737 | |
| 2738 | # Nothing further: mamba align is never picked up positionally. |
| 2739 | meta3 = sched.build_connector_meta( |
| 2740 | make_scheduler_output( |
| 2741 | {req.request_id: 1}, |
| 2742 | cached_req_new_blocks={req.request_id: None}, |
| 2743 | ) |
| 2744 | ) |
| 2745 | assert meta3.store_event == -1 |
| 2746 | assert meta3.store_gpu_blocks == [] |
| 2747 | assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0] |
| 2748 | |
| 2749 | |
| 2750 | @pytest.mark.parametrize("gone_via", ["preempted", "finished", "unregistered"]) |
| 2751 | def test_boundary_handoff_dropped_for_departing_request(gone_via: str) -> None: |
| 2752 | """Handoffs for a request leaving this step must not be read. |
| 2753 | |
| 2754 | The scheduler drains boundary handoffs without filtering by request |
| 2755 | liveness, so a preempted or finished request can still offer block IDs. |
| 2756 | Those blocks are being released and may already back another request, so |
| 2757 | copying them would publish unrelated KV under a valid hash. |
| 2758 | """ |
| 2759 | block_size = 4 * BLOCK_SIZE |
| 2760 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2761 | num_cpu_blocks=16, num_gpu_blocks=24, block_size=block_size |
| 2762 | ) |
| 2763 | sched = fix.scheduler |
| 2764 | gpu_pool = fix.gpu_block_pool |
| 2765 | |
| 2766 | req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) |
| 2767 | attn_blocks = _allocate_cp_gpu_blocks( |
| 2768 | gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 |
| 2769 | ) |
| 2770 | mamba_blocks = gpu_pool.get_new_blocks(2) |
| 2771 | mamba_blocks[1].set_block_hash( |
| 2772 | make_block_hash_with_group_id(req.block_hashes[1], 1) |
| 2773 | ) |
| 2774 | kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks)) |
| 2775 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2776 | |
| 2777 | output = make_scheduler_output( |
| 2778 | {req.request_id: 2 * block_size}, |
| 2779 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 2780 | ) |
| 2781 | output.kv_connector_block_state = KVConnectorBlockState( |
| 2782 | req_ids=set(), |
| 2783 | resolve_block_ids=lambda _: (), |
| 2784 | boundary_state_offloads={ |
| 2785 | req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)] |
| 2786 | }, |
| 2787 | ) |
| 2788 | if gone_via == "preempted": |
| 2789 | output.preempted_req_ids = {req.request_id} |
| 2790 | elif gone_via == "finished": |
| 2791 | output.finished_req_ids = {req.request_id} |
| 2792 | else: |
| 2793 | sched._reqs_to_store.pop(req.request_id) |
| 2794 | |
| 2795 | meta = sched.build_connector_meta(output) |
| 2796 | |
| 2797 | # The exact-handoff block is never a store source for a departing request. |
| 2798 | assert mamba_blocks[1].block_id not in meta.store_gpu_blocks |
| 2799 | stats = sched.get_boundary_store_stats() |
| 2800 | assert stats.dropped_request_gone == 1 |
| 2801 | assert stats.stored == 0 |
| 2802 | |
| 2803 | |
| 2804 | def test_boundary_handoff_keeps_block_meta_index_parallel() -> None: |
| 2805 | """Boundary stores must extend block_meta alongside the block ids. |
| 2806 | |
| 2807 | ``_process_store_completion`` indexes ``block_meta`` by position, so a |
| 2808 | boundary handoff that appends only block ids shortens the list, trips its |
| 2809 | length assertion and shifts every following entry's metadata. |
| 2810 | """ |
| 2811 | block_size = 4 * BLOCK_SIZE |
| 2812 | fix = _make_hybrid_attention_mamba_scheduler( |
| 2813 | num_cpu_blocks=16, |
| 2814 | num_gpu_blocks=24, |
| 2815 | block_size=block_size, |
| 2816 | enable_kv_cache_events=True, |
| 2817 | ) |
| 2818 | sched = fix.scheduler |
| 2819 | gpu_pool = fix.gpu_block_pool |
| 2820 | assert sched.enable_kv_cache_events |
| 2821 | |
| 2822 | req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) |
| 2823 | attn_blocks = _allocate_cp_gpu_blocks( |
| 2824 | gpu_pool, req, 2, virtual_block_size=block_size, group_id=0 |
| 2825 | ) |
| 2826 | mamba_blocks = gpu_pool.get_new_blocks(2) |
| 2827 | mamba_blocks[1].set_block_hash( |
| 2828 | make_block_hash_with_group_id(req.block_hashes[1], 1) |
| 2829 | ) |
| 2830 | kv_blocks = KVCacheBlocks(blocks=(attn_blocks, mamba_blocks)) |
| 2831 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2832 | # Make the FA blocks positionally eligible in the same event as the |
| 2833 | # explicit Mamba boundary handoff. |
| 2834 | req.num_computed_tokens = 2 * block_size |
| 2835 | |
| 2836 | output = make_scheduler_output( |
| 2837 | {req.request_id: 2 * block_size}, |
| 2838 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 2839 | ) |
| 2840 | output.kv_connector_block_state = KVConnectorBlockState( |
| 2841 | req_ids=set(), |
| 2842 | resolve_block_ids=lambda _: (), |
| 2843 | boundary_state_offloads={ |
| 2844 | req.request_id: [(1, mamba_blocks[1].block_id, 2 * block_size)] |
| 2845 | }, |
| 2846 | ) |
| 2847 | |
| 2848 | gpu_ids, cpu_ids, req_ids, block_meta = sched.prepare_store_specs(output) |
| 2849 | assert mamba_blocks[1].block_id in gpu_ids |
| 2850 | assert set(gpu_ids) == { |
| 2851 | *(block.block_id for block in attn_blocks), |
| 2852 | mamba_blocks[1].block_id, |
| 2853 | } |
| 2854 | assert req_ids == [req.request_id] |
| 2855 | assert block_meta is not None |
| 2856 | assert len(block_meta) == len(gpu_ids) == len(cpu_ids) |
| 2857 | |
| 2858 | |
| 2859 | def test_finished_eager_store_recovers_nulled_window_tail() -> None: |
| 2860 | """A sliding-window group nulls pages that left the window before the |
| 2861 | connector sees the finished request's block table, while the GPU keeps |
| 2862 | them hashed in its free queue. The eager store must copy those pages, or |
| 2863 | the CPU lookup (which needs the whole window) misses every time.""" |
| 2864 | fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False) |
| 2865 | sched = fix.scheduler |
| 2866 | gpu_pool = fix.gpu_block_pool |
| 2867 | null_block = gpu_pool.null_block |
| 2868 | |
| 2869 | num_blocks = 4 # == sliding window of the test SWA group |
| 2870 | req = make_request(num_blocks=num_blocks) |
| 2871 | fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 2872 | swa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=1) |
| 2873 | # remove_skipped_blocks() ran: the two oldest window pages are freed (still |
| 2874 | # hashed) and their slots hold the null block. |
| 2875 | gpu_pool.free_blocks(swa_blocks[:2]) |
| 2876 | swa_table = [null_block, null_block, *swa_blocks[2:]] |
| 2877 | kv_blocks = KVCacheBlocks(blocks=(fa_blocks, swa_table)) |
| 2878 | req.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 2879 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2880 | |
| 2881 | block_ids = kv_blocks.get_block_ids() |
| 2882 | sched.request_finished_all_groups(req, block_ids) |
| 2883 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 2884 | |
| 2885 | assert null_block.block_id not in meta.store_gpu_blocks |
| 2886 | assert set(meta.store_gpu_blocks) == {b.block_id for b in (*fa_blocks, *swa_blocks)} |
| 2887 | assert meta.store_event >= 0 |
| 2888 | simulate_store_completion(sched, meta.store_event) |
| 2889 | |
| 2890 | req2 = Request( |
| 2891 | request_id="req-window-tail-load", |
| 2892 | prompt_token_ids=req.prompt_token_ids, |
| 2893 | sampling_params=req.sampling_params, |
| 2894 | pooling_params=None, |
| 2895 | mm_features=None, |
| 2896 | block_hasher=req._block_hasher, |
| 2897 | ) |
| 2898 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 2899 | assert hit_tokens == num_blocks * BLOCK_SIZE |
| 2900 | assert is_async is True |
| 2901 | |
| 2902 | |
| 2903 | QSA_RING_BLOCK_SIZE = 4 |
| 2904 | |
| 2905 | |
| 2906 | def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: |
| 2907 | register_all_kvcache_specs(vllm_config=None) |
| 2908 | full_spec = FullAttentionSpec( |
| 2909 | block_size=BLOCK_SIZE, |
| 2910 | num_kv_heads=NUM_KV_HEADS, |
| 2911 | head_size=HEAD_SIZE, |
| 2912 | dtype=DTYPE, |
| 2913 | ) |
| 2914 | ring_spec = CircularBufferSpec( |
| 2915 | block_size=QSA_RING_BLOCK_SIZE, |
| 2916 | num_kv_heads=NUM_KV_HEADS, |
| 2917 | head_size=HEAD_SIZE, |
| 2918 | head_size_v=0, |
| 2919 | dtype=DTYPE, |
| 2920 | ) |
| 2921 | groups = [ |
| 2922 | KVCacheGroupSpec(["layer_0"], full_spec), |
| 2923 | KVCacheGroupSpec(["qsa_ring"], ring_spec), |
| 2924 | ] |
| 2925 | tensors = [ |
| 2926 | KVCacheTensor( |
| 2927 | size=group.kv_cache_spec.page_size_bytes * num_blocks, |
| 2928 | layers=list(group.layer_names), |
| 2929 | layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks, |
| 2930 | block_stride=group.kv_cache_spec.page_size_bytes, |
| 2931 | ) |
| 2932 | for group in groups |
| 2933 | ] |
| 2934 | return KVCacheConfig( |
| 2935 | num_blocks=num_blocks, |
| 2936 | kv_cache_tensors=tensors, |
| 2937 | kv_cache_groups=groups, |
| 2938 | ) |
| 2939 | |
| 2940 | |
| 2941 | def _make_qsa_scheduler( |
| 2942 | lazy: bool = False, |
| 2943 | ) -> tuple[SimpleCPUOffloadScheduler, BlockPool]: |
| 2944 | kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16) |
| 2945 | sched = SimpleCPUOffloadScheduler( |
| 2946 | vllm_config=_make_vllm_config(), |
| 2947 | kv_cache_config=kv_cache_config, |
| 2948 | cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2, |
| 2949 | scheduler_block_size=BLOCK_SIZE, |
| 2950 | hash_block_size=BLOCK_SIZE, |
| 2951 | lazy_offload=lazy, |
| 2952 | ) |
| 2953 | gpu_pool = BlockPool( |
| 2954 | num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE |
| 2955 | ) |
| 2956 | sched.bind_gpu_block_pool(gpu_pool) |
| 2957 | return sched, gpu_pool |
| 2958 | |
| 2959 | |
| 2960 | def test_qsa_ring_group_is_never_stored_or_loaded() -> None: |
| 2961 | sched, gpu_pool = _make_qsa_scheduler() |
| 2962 | assert sched.prefix_cacheable_group_ids == (0,) |
| 2963 | |
| 2964 | num_blocks = 2 |
| 2965 | req = make_request(num_blocks=num_blocks) |
| 2966 | fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 2967 | ring_block = gpu_pool.get_new_blocks(1)[0] |
| 2968 | kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block])) |
| 2969 | req.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 2970 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 2971 | |
| 2972 | sched_out = make_scheduler_output( |
| 2973 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 2974 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 2975 | ) |
| 2976 | meta = sched.build_connector_meta(sched_out) |
| 2977 | assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks} |
| 2978 | assert ring_block.block_id not in meta.store_gpu_blocks |
| 2979 | assert meta.store_event >= 0 |
| 2980 | simulate_store_completion(sched, meta.store_event) |
| 2981 | |
| 2982 | req2 = Request( |
| 2983 | request_id="req-qsa-load", |
| 2984 | prompt_token_ids=req.prompt_token_ids, |
| 2985 | sampling_params=req.sampling_params, |
| 2986 | pooling_params=None, |
| 2987 | mm_features=None, |
| 2988 | block_hasher=req._block_hasher, |
| 2989 | ) |
| 2990 | hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) |
| 2991 | assert hit_tokens == num_blocks * BLOCK_SIZE |
| 2992 | assert is_async is True |
| 2993 | |
| 2994 | fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) |
| 2995 | ring_block2 = gpu_pool.get_new_blocks(1)[0] |
| 2996 | kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2])) |
| 2997 | sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) |
| 2998 | |
| 2999 | meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1})) |
| 3000 | assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2} |
| 3001 | assert ring_block2.block_id not in meta2.load_gpu_blocks |
| 3002 | assert len(meta2.load_cpu_blocks) == num_blocks |
| 3003 | |
| 3004 | |
| 3005 | def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: |
| 3006 | with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16) |
| 3007 | without_ring = KVCacheConfig( |
| 3008 | num_blocks=16, |
| 3009 | kv_cache_tensors=with_ring.kv_cache_tensors[:1], |
| 3010 | kv_cache_groups=with_ring.kv_cache_groups[:1], |
| 3011 | ) |
| 3012 | max_batched = 64 |
| 3013 | target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( |
| 3014 | with_ring, max_batched |
| 3015 | ) |
| 3016 | target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( |
| 3017 | without_ring, max_batched |
| 3018 | ) |
| 3019 | assert target_with == target_without |