| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """KV cache event emission tests for SimpleCPUOffloadScheduler. |
| 4 | |
| 5 | Verifies that BlockStored/BlockRemoved events carry the correct storage |
| 6 | medium (MEDIUM_CPU / MEDIUM_STORAGE) and per-group metadata when |
| 7 | ``enable_kv_cache_events`` is set, and that no events are emitted when it |
| 8 | is unset. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | from tests.v1.simple_kv_offload.test_scheduler import ( |
| 14 | _BYTES_PER_BLOCK, |
| 15 | BLOCK_SIZE, |
| 16 | DTYPE, |
| 17 | HEAD_SIZE, |
| 18 | NUM_KV_HEADS, |
| 19 | SchedulerFixture, |
| 20 | _alloc_and_register, |
| 21 | _allocate_cp_gpu_blocks, |
| 22 | _allocate_gpu_blocks, |
| 23 | _make_cp_request, |
| 24 | _make_kv_cache_config, |
| 25 | _make_scratch_kv_cache_config, |
| 26 | _make_vllm_config, |
| 27 | make_request, |
| 28 | make_scheduler_output, |
| 29 | simulate_store_completion, |
| 30 | ) |
| 31 | from vllm import SamplingParams |
| 32 | from vllm.config.kv_events import KVEventsConfig |
| 33 | from vllm.distributed.kv_events import ( |
| 34 | MEDIUM_CPU, |
| 35 | MEDIUM_STORAGE, |
| 36 | BlockRemoved, |
| 37 | BlockStored, |
| 38 | ) |
| 39 | from vllm.lora.request import LoRARequest |
| 40 | from vllm.utils.hashing import sha256 |
| 41 | from vllm.v1.core.block_pool import BlockPool |
| 42 | from vllm.v1.core.kv_cache_manager import KVCacheBlocks |
| 43 | from vllm.v1.core.kv_cache_utils import ( |
| 44 | get_block_hash, |
| 45 | get_request_block_hasher, |
| 46 | make_block_hash_with_group_id, |
| 47 | maybe_convert_block_hash, |
| 48 | ) |
| 49 | from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs |
| 50 | from vllm.v1.kv_cache_interface import ( |
| 51 | FullAttentionSpec, |
| 52 | KVCacheConfig, |
| 53 | KVCacheGroupSpec, |
| 54 | KVCacheTensor, |
| 55 | MambaSpec, |
| 56 | SlidingWindowSpec, |
| 57 | ) |
| 58 | from vllm.v1.request import Request |
| 59 | from vllm.v1.simple_kv_offload.manager import ( |
| 60 | SimpleCPUOffloadScheduler, |
| 61 | StoreRequestState, |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | def _make_mixed_kv_cache_config( |
| 66 | num_blocks: int, |
| 67 | sliding_window: int | None = None, |
| 68 | with_mamba: bool = False, |
| 69 | mamba_block_size: int = BLOCK_SIZE, |
| 70 | ) -> KVCacheConfig: |
| 71 | """Build a KVCacheConfig with a FullAttention group plus optional |
| 72 | SlidingWindow and Mamba groups for per-group metadata tests.""" |
| 73 | register_all_kvcache_specs(vllm_config=None) |
| 74 | groups = [] |
| 75 | tensors = [] |
| 76 | |
| 77 | # Group 0 is always FullAttention (the scheduler requires one FA group). |
| 78 | fa_layers = ["layer_fa"] |
| 79 | fa_spec = FullAttentionSpec( |
| 80 | block_size=BLOCK_SIZE, |
| 81 | num_kv_heads=NUM_KV_HEADS, |
| 82 | head_size=HEAD_SIZE, |
| 83 | dtype=DTYPE, |
| 84 | ) |
| 85 | groups.append(KVCacheGroupSpec(fa_layers, fa_spec)) |
| 86 | fa_bytes = _BYTES_PER_BLOCK * num_blocks |
| 87 | tensors.append( |
| 88 | KVCacheTensor( |
| 89 | size=fa_bytes, |
| 90 | layers=fa_layers, |
| 91 | layer_stride=fa_bytes, |
| 92 | block_stride=_BYTES_PER_BLOCK, |
| 93 | ) |
| 94 | ) |
| 95 | |
| 96 | if sliding_window is not None: |
| 97 | sw_layers = ["layer_sw"] |
| 98 | sw_spec = SlidingWindowSpec( |
| 99 | block_size=BLOCK_SIZE, |
| 100 | num_kv_heads=NUM_KV_HEADS, |
| 101 | head_size=HEAD_SIZE, |
| 102 | dtype=DTYPE, |
| 103 | sliding_window=sliding_window, |
| 104 | ) |
| 105 | groups.append(KVCacheGroupSpec(sw_layers, sw_spec)) |
| 106 | sw_bytes = _BYTES_PER_BLOCK * num_blocks |
| 107 | tensors.append( |
| 108 | KVCacheTensor( |
| 109 | size=sw_bytes, |
| 110 | layers=sw_layers, |
| 111 | layer_stride=sw_bytes, |
| 112 | block_stride=_BYTES_PER_BLOCK, |
| 113 | ) |
| 114 | ) |
| 115 | |
| 116 | if with_mamba: |
| 117 | m_layers = ["layer_m"] |
| 118 | m_spec = MambaSpec( |
| 119 | block_size=mamba_block_size, |
| 120 | shapes=((NUM_KV_HEADS, HEAD_SIZE, HEAD_SIZE),), |
| 121 | dtypes=(DTYPE,), |
| 122 | mamba_cache_mode="align", |
| 123 | ) |
| 124 | groups.append(KVCacheGroupSpec(m_layers, m_spec)) |
| 125 | m_bytes = _BYTES_PER_BLOCK * num_blocks |
| 126 | tensors.append( |
| 127 | KVCacheTensor( |
| 128 | size=m_bytes, |
| 129 | layers=m_layers, |
| 130 | layer_stride=m_bytes, |
| 131 | block_stride=_BYTES_PER_BLOCK, |
| 132 | ) |
| 133 | ) |
| 134 | |
| 135 | return KVCacheConfig( |
| 136 | num_blocks=num_blocks, |
| 137 | kv_cache_tensors=tensors, |
| 138 | kv_cache_groups=groups, |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | def make_events_scheduler( |
| 143 | num_cpu_blocks: int = 8, |
| 144 | num_gpu_blocks: int = 16, |
| 145 | num_groups: int = 1, |
| 146 | lazy: bool = False, |
| 147 | disk_capacity_bytes: int = 0, |
| 148 | dcp_world_size: int = 1, |
| 149 | kv_cache_config: KVCacheConfig | None = None, |
| 150 | enable_events: bool = True, |
| 151 | ) -> SchedulerFixture: |
| 152 | """Build a SimpleCPUOffloadScheduler with kv cache events enabled.""" |
| 153 | if kv_cache_config is None: |
| 154 | kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) |
| 155 | vllm_config = _make_vllm_config() |
| 156 | if enable_events: |
| 157 | vllm_config.kv_events_config = KVEventsConfig(enable_kv_cache_events=True) |
| 158 | vllm_config.parallel_config.decode_context_parallel_size = dcp_world_size |
| 159 | |
| 160 | virtual_block_size = BLOCK_SIZE * dcp_world_size |
| 161 | cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups |
| 162 | |
| 163 | sched = SimpleCPUOffloadScheduler( |
| 164 | vllm_config=vllm_config, |
| 165 | kv_cache_config=kv_cache_config, |
| 166 | cpu_capacity_bytes=cpu_capacity_bytes, |
| 167 | scheduler_block_size=virtual_block_size, |
| 168 | hash_block_size=virtual_block_size, |
| 169 | lazy_offload=lazy, |
| 170 | disk_capacity_bytes=disk_capacity_bytes, |
| 171 | ) |
| 172 | |
| 173 | gpu_block_pool = BlockPool( |
| 174 | num_gpu_blocks=num_gpu_blocks, |
| 175 | enable_caching=True, |
| 176 | hash_block_size=virtual_block_size, |
| 177 | ) |
| 178 | sched.bind_gpu_block_pool(gpu_block_pool) |
| 179 | return SchedulerFixture( |
| 180 | scheduler=sched, |
| 181 | gpu_block_pool=gpu_block_pool, |
| 182 | vllm_config=vllm_config, |
| 183 | kv_cache_config=kv_cache_config, |
| 184 | num_groups=num_groups, |
| 185 | ) |
| 186 | |
| 187 | |
| 188 | def _do_eager_store( |
| 189 | fix: SchedulerFixture, |
| 190 | num_blocks: int = 2, |
| 191 | lora_request: LoRARequest | None = None, |
| 192 | ) -> tuple[object, Request]: |
| 193 | """Run an eager store of ``num_blocks`` and complete it. |
| 194 | |
| 195 | Returns (meta, req) for further assertions. |
| 196 | """ |
| 197 | sched = fix.scheduler |
| 198 | if lora_request is not None: |
| 199 | # lora_request must be set in the constructor so block_hashes |
| 200 | # (computed in __init__) include the lora extra_keys. |
| 201 | req = Request( |
| 202 | request_id="req-lora", |
| 203 | prompt_token_ids=list(range(num_blocks * BLOCK_SIZE + 1)), |
| 204 | sampling_params=SamplingParams(max_tokens=1), |
| 205 | pooling_params=None, |
| 206 | mm_features=None, |
| 207 | lora_request=lora_request, |
| 208 | block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256), |
| 209 | ) |
| 210 | else: |
| 211 | req = make_request(num_blocks=num_blocks) |
| 212 | kv_blocks = _alloc_and_register(fix, req, num_blocks) |
| 213 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 214 | block_ids = kv_blocks.get_block_ids() |
| 215 | sched_out = make_scheduler_output( |
| 216 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 217 | new_reqs={req.request_id: block_ids}, |
| 218 | ) |
| 219 | meta = sched.build_connector_meta(sched_out) |
| 220 | assert meta.store_event >= 0 |
| 221 | simulate_store_completion(sched, meta.store_event) |
| 222 | return meta, req |
| 223 | |
| 224 | |
| 225 | # --------------------------------------------------------------------------- |
| 226 | # events disabled -> no events |
| 227 | # --------------------------------------------------------------------------- |
| 228 | def test_events_disabled_no_events() -> None: |
| 229 | """When enable_kv_cache_events is False, take_events returns nothing.""" |
| 230 | fix = make_events_scheduler(enable_events=False) |
| 231 | sched = fix.scheduler |
| 232 | assert sched.enable_kv_cache_events is False |
| 233 | _do_eager_store(fix, num_blocks=2) |
| 234 | events = list(sched.take_events()) |
| 235 | assert events == [], f"expected no events, got {events}" |
| 236 | |
| 237 | |
| 238 | # --------------------------------------------------------------------------- |
| 239 | # eager store emits BlockStored with per-group metadata |
| 240 | # --------------------------------------------------------------------------- |
| 241 | def test_block_stored_per_group_metadata_full_attention() -> None: |
| 242 | """BlockStored carries medium, group_idx, kind, sliding_window, block_size, |
| 243 | and the correct block hash.""" |
| 244 | fix = make_events_scheduler() |
| 245 | sched = fix.scheduler |
| 246 | assert sched.kv_event_medium == MEDIUM_CPU |
| 247 | _, req = _do_eager_store(fix, num_blocks=2) |
| 248 | events = list(sched.take_events()) |
| 249 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 250 | assert len(stored) == 2 |
| 251 | ev = stored[0] |
| 252 | assert ev.medium == MEDIUM_CPU |
| 253 | assert ev.group_idx == 0 |
| 254 | assert ev.kv_cache_spec_kind == "full_attention" |
| 255 | assert ev.kv_cache_spec_sliding_window is None |
| 256 | assert ev.block_size == BLOCK_SIZE |
| 257 | assert ev.locality == "LOCAL" |
| 258 | expected_hash_0 = maybe_convert_block_hash( |
| 259 | get_block_hash(make_block_hash_with_group_id(req.block_hashes[0], 0)) |
| 260 | ) |
| 261 | assert ev.block_hashes == [expected_hash_0] |
| 262 | assert ev.parent_block_hash is None |
| 263 | assert req.prompt_token_ids is not None |
| 264 | assert ev.token_ids == req.prompt_token_ids[0:BLOCK_SIZE] |
| 265 | |
| 266 | # Second block's parent_block_hash equals the first block's hash. |
| 267 | ev1 = stored[1] |
| 268 | expected_hash_1 = maybe_convert_block_hash( |
| 269 | get_block_hash(make_block_hash_with_group_id(req.block_hashes[1], 0)) |
| 270 | ) |
| 271 | assert ev1.block_hashes == [expected_hash_1] |
| 272 | assert ev1.parent_block_hash == expected_hash_0 |
| 273 | assert ev1.token_ids == req.prompt_token_ids[BLOCK_SIZE : 2 * BLOCK_SIZE] |
| 274 | |
| 275 | |
| 276 | def test_eager_store_lora_metadata() -> None: |
| 277 | """BlockStored carries lora_id and lora_name from req.lora_request.""" |
| 278 | fix = make_events_scheduler() |
| 279 | sched = fix.scheduler |
| 280 | |
| 281 | _do_eager_store( |
| 282 | fix, |
| 283 | num_blocks=1, |
| 284 | lora_request=LoRARequest( |
| 285 | lora_name="test-lora", |
| 286 | lora_int_id=42, |
| 287 | lora_path="/fake/lora/path", |
| 288 | ), |
| 289 | ) |
| 290 | |
| 291 | events = list(sched.take_events()) |
| 292 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 293 | assert len(stored) == 1 |
| 294 | ev = stored[0] |
| 295 | assert ev.lora_id == 42, f"expected lora_id=42, got {ev.lora_id}" |
| 296 | assert ev.lora_name == "test-lora", ( |
| 297 | f"expected lora_name='test-lora', got {ev.lora_name!r}" |
| 298 | ) |
| 299 | assert ev.extra_keys is not None, "expected extra_keys for lora request" |
| 300 | assert any(keys is not None and "test-lora" in keys for keys in ev.extra_keys), ( |
| 301 | f"expected 'test-lora' in extra_keys, got {ev.extra_keys}" |
| 302 | ) |
| 303 | assert len(ev.token_ids) == BLOCK_SIZE |
| 304 | |
| 305 | |
| 306 | # --------------------------------------------------------------------------- |
| 307 | # disk mode -> MEDIUM_STORAGE |
| 308 | # --------------------------------------------------------------------------- |
| 309 | def test_disk_mode_medium_storage_on_block_stored() -> None: |
| 310 | """disk_capacity_bytes>0 sets kv_event_medium=STORAGE on BlockStored.""" |
| 311 | fix = make_events_scheduler( |
| 312 | disk_capacity_bytes=_BYTES_PER_BLOCK * 8, |
| 313 | enable_events=True, |
| 314 | ) |
| 315 | sched = fix.scheduler |
| 316 | assert sched.kv_event_medium == MEDIUM_STORAGE |
| 317 | _do_eager_store(fix, num_blocks=1) |
| 318 | events = list(sched.take_events()) |
| 319 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 320 | assert len(stored) == 1 |
| 321 | assert stored[0].medium == MEDIUM_STORAGE |
| 322 | assert stored[0].locality == "LOCAL" |
| 323 | |
| 324 | |
| 325 | def test_finished_eager_store_emits_all_storage_events() -> None: |
| 326 | """Finished stores preserve and announce every registered hash.""" |
| 327 | fix = make_events_scheduler( |
| 328 | disk_capacity_bytes=_BYTES_PER_BLOCK * 8, |
| 329 | enable_events=True, |
| 330 | ) |
| 331 | sched = fix.scheduler |
| 332 | req = make_request(num_blocks=2) |
| 333 | kv_blocks = _alloc_and_register(fix, req, num_blocks=2) |
| 334 | gpu_blocks = kv_blocks.blocks[0] |
| 335 | primary_hashes = [] |
| 336 | for gpu_block in gpu_blocks: |
| 337 | assert gpu_block.block_hash is not None |
| 338 | primary_hashes.append(gpu_block.block_hash) |
| 339 | |
| 340 | fine_grained_req = Request( |
| 341 | request_id="req-finished-fine-grained-hash", |
| 342 | prompt_token_ids=req.prompt_token_ids, |
| 343 | sampling_params=req.sampling_params, |
| 344 | pooling_params=None, |
| 345 | mm_features=None, |
| 346 | block_hasher=get_request_block_hasher(BLOCK_SIZE // 2, sha256), |
| 347 | ) |
| 348 | secondary_hash = make_block_hash_with_group_id(fine_grained_req.block_hashes[0], 0) |
| 349 | fix.gpu_block_pool._insert_block_hash( |
| 350 | secondary_hash, gpu_blocks[0], num_tokens=BLOCK_SIZE // 2 |
| 351 | ) |
| 352 | |
| 353 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 354 | sched.request_finished_all_groups(req, kv_blocks.get_block_ids()) |
| 355 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 356 | assert meta.store_event >= 0 |
| 357 | simulate_store_completion(sched, meta.store_event) |
| 358 | |
| 359 | primary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 360 | primary_hashes[0] |
| 361 | ) |
| 362 | secondary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 363 | secondary_hash |
| 364 | ) |
| 365 | assert primary_cpu_block is not None |
| 366 | assert secondary_cpu_block is primary_cpu_block |
| 367 | |
| 368 | stored = [e for e in sched.take_events() if isinstance(e, BlockStored)] |
| 369 | primary_event_hashes = [ |
| 370 | maybe_convert_block_hash(get_block_hash(block_hash)) |
| 371 | for block_hash in primary_hashes |
| 372 | ] |
| 373 | secondary_event_hash = maybe_convert_block_hash(get_block_hash(secondary_hash)) |
| 374 | by_hash = {event.block_hashes[0]: event for event in stored} |
| 375 | assert set(by_hash) == {*primary_event_hashes, secondary_event_hash} |
| 376 | assert req.prompt_token_ids is not None |
| 377 | assert ( |
| 378 | by_hash[primary_event_hashes[0]].token_ids == req.prompt_token_ids[:BLOCK_SIZE] |
| 379 | ) |
| 380 | assert ( |
| 381 | by_hash[primary_event_hashes[1]].token_ids |
| 382 | == req.prompt_token_ids[BLOCK_SIZE : 2 * BLOCK_SIZE] |
| 383 | ) |
| 384 | assert by_hash[primary_event_hashes[1]].parent_block_hash == primary_event_hashes[0] |
| 385 | secondary_event = by_hash[secondary_event_hash] |
| 386 | assert secondary_event.medium == MEDIUM_STORAGE |
| 387 | assert secondary_event.block_size == 0 |
| 388 | assert secondary_event.token_ids == [] |
| 389 | assert secondary_event.parent_block_hash is None |
| 390 | |
| 391 | sched.cpu_block_pool.evict_blocks({primary_cpu_block.block_id}) |
| 392 | removed = [ |
| 393 | event for event in sched.take_events() if isinstance(event, BlockRemoved) |
| 394 | ] |
| 395 | assert {event.block_hashes[0] for event in removed} == { |
| 396 | primary_event_hashes[0], |
| 397 | secondary_event_hash, |
| 398 | } |
| 399 | assert all(event.medium == MEDIUM_STORAGE for event in removed) |
| 400 | |
| 401 | |
| 402 | def test_secondary_hash_block_stored_metadata() -> None: |
| 403 | """Secondary hashes carry their own token range and block size.""" |
| 404 | fix = make_events_scheduler( |
| 405 | disk_capacity_bytes=_BYTES_PER_BLOCK * 8, |
| 406 | enable_events=True, |
| 407 | ) |
| 408 | sched = fix.scheduler |
| 409 | sched.hash_block_size = BLOCK_SIZE // 2 |
| 410 | req = Request( |
| 411 | request_id="req-secondary-hash", |
| 412 | prompt_token_ids=list(range(BLOCK_SIZE + 1)), |
| 413 | sampling_params=SamplingParams(max_tokens=1), |
| 414 | pooling_params=None, |
| 415 | mm_features=None, |
| 416 | block_hasher=get_request_block_hasher(BLOCK_SIZE // 2, sha256), |
| 417 | ) |
| 418 | req.num_computed_tokens = BLOCK_SIZE |
| 419 | gpu_block = fix.gpu_block_pool.get_new_blocks(1)[0] |
| 420 | primary_hash = make_block_hash_with_group_id(req.block_hashes[1], 0) |
| 421 | secondary_hash = make_block_hash_with_group_id(req.block_hashes[0], 0) |
| 422 | fix.gpu_block_pool._insert_block_hash( |
| 423 | primary_hash, gpu_block, num_tokens=BLOCK_SIZE |
| 424 | ) |
| 425 | fix.gpu_block_pool._insert_block_hash( |
| 426 | secondary_hash, gpu_block, num_tokens=BLOCK_SIZE // 2 |
| 427 | ) |
| 428 | state = StoreRequestState( |
| 429 | request=req, |
| 430 | block_ids=([gpu_block.block_id],), |
| 431 | num_stored_blocks=[0], |
| 432 | ) |
| 433 | |
| 434 | gpu_block_ids, advanced_per_group, block_meta = sched._select_eager_blocks_to_store( |
| 435 | state, |
| 436 | ([gpu_block.block_id],), |
| 437 | ) |
| 438 | assert gpu_block_ids == [gpu_block.block_id] |
| 439 | assert advanced_per_group == [1] |
| 440 | assert block_meta is not None |
| 441 | assert set(block_meta[0]) == {primary_hash, secondary_hash} |
| 442 | |
| 443 | cpu_block = sched.cpu_block_pool.get_new_blocks(1)[0] |
| 444 | sched._process_store_completion( |
| 445 | gpu_block_ids, |
| 446 | [cpu_block.block_id], |
| 447 | block_meta, |
| 448 | ) |
| 449 | |
| 450 | primary_event_hash = maybe_convert_block_hash(get_block_hash(primary_hash)) |
| 451 | secondary_event_hash = maybe_convert_block_hash(get_block_hash(secondary_hash)) |
| 452 | stored = [event for event in sched.take_events() if isinstance(event, BlockStored)] |
| 453 | by_hash = {event.block_hashes[0]: event for event in stored} |
| 454 | assert set(by_hash) == {primary_event_hash, secondary_event_hash} |
| 455 | assert by_hash[primary_event_hash].block_size == BLOCK_SIZE |
| 456 | assert req.prompt_token_ids is not None |
| 457 | assert by_hash[primary_event_hash].token_ids == req.prompt_token_ids[:BLOCK_SIZE] |
| 458 | assert by_hash[secondary_event_hash].block_size == BLOCK_SIZE // 2 |
| 459 | assert ( |
| 460 | by_hash[secondary_event_hash].token_ids |
| 461 | == req.prompt_token_ids[: BLOCK_SIZE // 2] |
| 462 | ) |
| 463 | |
| 464 | primary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 465 | primary_hash |
| 466 | ) |
| 467 | secondary_cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block( |
| 468 | secondary_hash |
| 469 | ) |
| 470 | assert primary_cpu_block is not None |
| 471 | assert secondary_cpu_block is primary_cpu_block |
| 472 | |
| 473 | |
| 474 | # --------------------------------------------------------------------------- |
| 475 | # BlockRemoved relabeled to correct medium |
| 476 | # --------------------------------------------------------------------------- |
| 477 | def test_block_removed_relabeled_medium_cpu() -> None: |
| 478 | """Evicted BlockRemoved (emitted as GPU by BlockPool) is relabeled CPU.""" |
| 479 | # 5 total = 4 usable (null_block takes 1). Fill 4, then store 2 more to |
| 480 | # evict 2 LRU blocks, producing BlockRemoved events. |
| 481 | fix = make_events_scheduler(num_cpu_blocks=5, num_gpu_blocks=16) |
| 482 | sched = fix.scheduler |
| 483 | |
| 484 | # Fill CPU with 4 blocks (2 + 2). |
| 485 | _do_eager_store(fix, num_blocks=2) |
| 486 | # Drain store events so they don't clutter the eviction check. |
| 487 | list(sched.take_events()) |
| 488 | _do_eager_store(fix, num_blocks=2) |
| 489 | list(sched.take_events()) |
| 490 | |
| 491 | # Now store 2 more -> evicts 2 LRU blocks -> BlockRemoved emitted by pool |
| 492 | # with MEDIUM_GPU, then relabeled by scheduler.take_events(). |
| 493 | _do_eager_store(fix, num_blocks=2) |
| 494 | events = list(sched.take_events()) |
| 495 | |
| 496 | removed = [e for e in events if isinstance(e, BlockRemoved)] |
| 497 | assert len(removed) >= 2, ( |
| 498 | f"expected >=2 BlockRemoved from eviction, got {len(removed)}" |
| 499 | ) |
| 500 | for ev in removed: |
| 501 | assert ev.medium == MEDIUM_CPU, ( |
| 502 | f"BlockRemoved should be relabeled CPU, got {ev.medium}" |
| 503 | ) |
| 504 | assert ev.locality == "LOCAL", ( |
| 505 | f"BlockRemoved should be LOCAL, got {ev.locality}" |
| 506 | ) |
| 507 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 508 | assert len(stored) == 2 |
| 509 | |
| 510 | |
| 511 | # --------------------------------------------------------------------------- |
| 512 | # lazy store emits BlockStored |
| 513 | # --------------------------------------------------------------------------- |
| 514 | def test_lazy_store_emits_block_stored() -> None: |
| 515 | """Lazy-mode store completion also emits BlockStored with MEDIUM_CPU.""" |
| 516 | fix = make_events_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True) |
| 517 | sched = fix.scheduler |
| 518 | gpu_pool = fix.gpu_block_pool |
| 519 | |
| 520 | # Allocate, hash, free -> hashed blocks in free queue. |
| 521 | req = make_request(num_blocks=2) |
| 522 | gpu_blocks = [] |
| 523 | for blk in gpu_pool.get_new_blocks(2): |
| 524 | gpu_blocks.append(blk) |
| 525 | gpu_pool.cache_full_blocks( |
| 526 | request=req, |
| 527 | blocks=gpu_blocks, |
| 528 | num_cached_blocks=0, |
| 529 | num_full_blocks=2, |
| 530 | block_size=BLOCK_SIZE, |
| 531 | kv_cache_group_id=0, |
| 532 | ) |
| 533 | gpu_pool.free_blocks(gpu_blocks) |
| 534 | |
| 535 | # Push hashed blocks to LRU head (8 total - 1 null = 7 usable; 2 freed |
| 536 | # hashed + 5 other free -> consume 5 fillers). |
| 537 | fillers = gpu_pool.get_new_blocks(5) |
| 538 | |
| 539 | meta = sched.build_connector_meta(make_scheduler_output({})) |
| 540 | assert meta.store_event >= 0 |
| 541 | simulate_store_completion(sched, meta.store_event) |
| 542 | gpu_pool.free_blocks(fillers) |
| 543 | |
| 544 | events = list(sched.take_events()) |
| 545 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 546 | assert len(stored) == 2, f"expected 2 BlockStored, got {len(stored)}" |
| 547 | for ev in stored: |
| 548 | assert ev.medium == MEDIUM_CPU |
| 549 | assert ev.group_idx == 0 |
| 550 | assert ev.kv_cache_spec_kind == "full_attention" |
| 551 | assert ev.locality == "LOCAL" |
| 552 | assert ev.token_ids == [] |
| 553 | assert ev.parent_block_hash is None |
| 554 | |
| 555 | |
| 556 | # --------------------------------------------------------------------------- |
| 557 | # cp_world_size scales block_size for non-Mamba |
| 558 | # --------------------------------------------------------------------------- |
| 559 | def test_cp_world_size_scales_block_size() -> None: |
| 560 | """block_size on BlockStored = spec.block_size * cp_world_size (FA).""" |
| 561 | fix = make_events_scheduler(dcp_world_size=2) |
| 562 | sched = fix.scheduler |
| 563 | gpu_pool = fix.gpu_block_pool |
| 564 | vbs = BLOCK_SIZE * 2 |
| 565 | assert sched.cp_world_size == 2 |
| 566 | |
| 567 | req = _make_cp_request(num_blocks=1, virtual_block_size=vbs) |
| 568 | gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 1, vbs) |
| 569 | kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,)) |
| 570 | req.num_computed_tokens = vbs |
| 571 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 572 | sched_out = make_scheduler_output( |
| 573 | {req.request_id: vbs}, |
| 574 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 575 | ) |
| 576 | meta = sched.build_connector_meta(sched_out) |
| 577 | assert meta.store_event >= 0 |
| 578 | simulate_store_completion(sched, meta.store_event) |
| 579 | |
| 580 | events = list(sched.take_events()) |
| 581 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 582 | assert len(stored) == 1 |
| 583 | # spec.block_size=BLOCK_SIZE, cp_world_size=2 -> BLOCK_SIZE * 2 |
| 584 | assert stored[0].block_size == BLOCK_SIZE * 2 |
| 585 | assert stored[0].locality == "LOCAL" |
| 586 | |
| 587 | |
| 588 | # --------------------------------------------------------------------------- |
| 589 | # Mamba block_size NOT scaled by cp_world_size |
| 590 | # --------------------------------------------------------------------------- |
| 591 | def test_mamba_block_size_not_scaled() -> None: |
| 592 | """Mamba group: block_size = spec.block_size (no cp scaling). |
| 593 | |
| 594 | With cp_world_size=2, scheduler_block_size=hash_block_size=32. The |
| 595 | coordinator constraints force Mamba spec.block_size=32 (must divide |
| 596 | scheduler_block_size and be a multiple of hash_block_size). The event |
| 597 | block_size for Mamba must be 32 (spec.block_size * 1), NOT 64 |
| 598 | (spec.block_size * cp_world_size) -- that is the bug this guards. |
| 599 | """ |
| 600 | kv_cfg = _make_mixed_kv_cache_config( |
| 601 | num_blocks=16, with_mamba=True, mamba_block_size=32 |
| 602 | ) |
| 603 | fix = make_events_scheduler( |
| 604 | num_cpu_blocks=8, |
| 605 | num_gpu_blocks=16, |
| 606 | kv_cache_config=kv_cfg, |
| 607 | dcp_world_size=2, |
| 608 | num_groups=2, |
| 609 | ) |
| 610 | sched = fix.scheduler |
| 611 | |
| 612 | cpu_pool = sched.cpu_block_pool |
| 613 | gpu_pool = fix.gpu_block_pool |
| 614 | |
| 615 | cpu_blk = cpu_pool.get_new_blocks(1)[0] |
| 616 | gpu_blk = gpu_pool.get_new_blocks(1)[0] |
| 617 | |
| 618 | raw_hash = make_request(num_blocks=1).block_hashes[0] |
| 619 | mamba_hash = make_block_hash_with_group_id(raw_hash, group_id=1) |
| 620 | gpu_blk._block_hash = mamba_hash # type: ignore[attr-defined] |
| 621 | |
| 622 | sched._process_store_completion( |
| 623 | gpu_block_ids=[gpu_blk.block_id], |
| 624 | cpu_block_ids=[cpu_blk.block_id], |
| 625 | ) |
| 626 | |
| 627 | events = list(sched.take_events()) |
| 628 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 629 | assert len(stored) == 1 |
| 630 | ev = stored[0] |
| 631 | assert ev.group_idx == 1 |
| 632 | assert ev.kv_cache_spec_kind == "mamba" |
| 633 | assert ev.kv_cache_spec_sliding_window is None |
| 634 | assert ev.locality == "LOCAL" |
| 635 | assert ev.block_size == 32, ( |
| 636 | f"Mamba block_size should be unscaled (32), got {ev.block_size}" |
| 637 | ) |
| 638 | |
| 639 | |
| 640 | # --------------------------------------------------------------------------- |
| 641 | # multi-group FA+SW stores carry correct group_idx and kind |
| 642 | # --------------------------------------------------------------------------- |
| 643 | def test_multi_group_fa_block_stored_group_idx() -> None: |
| 644 | """Two distinct groups (FA + SW): each BlockStored carries its group_idx |
| 645 | and kind. HybridKVCacheCoordinator requires >=2 distinct spec groups, so |
| 646 | we use FA (group 0) + SlidingWindow (group 1).""" |
| 647 | kv_cfg = _make_mixed_kv_cache_config(num_blocks=16, sliding_window=32) |
| 648 | fix = make_events_scheduler( |
| 649 | num_cpu_blocks=8, |
| 650 | num_gpu_blocks=16, |
| 651 | kv_cache_config=kv_cfg, |
| 652 | num_groups=2, |
| 653 | ) |
| 654 | sched = fix.scheduler |
| 655 | |
| 656 | cpu_pool = sched.cpu_block_pool |
| 657 | gpu_pool = fix.gpu_block_pool |
| 658 | raw_hash = make_request(num_blocks=1).block_hashes[0] |
| 659 | |
| 660 | # Store one block for each group via direct _process_store_completion. |
| 661 | for gidx in (0, 1): |
| 662 | cpu_blk = cpu_pool.get_new_blocks(1)[0] |
| 663 | gpu_blk = gpu_pool.get_new_blocks(1)[0] |
| 664 | gpu_blk._block_hash = make_block_hash_with_group_id( # type: ignore |
| 665 | raw_hash, group_id=gidx |
| 666 | ) |
| 667 | sched._process_store_completion( |
| 668 | gpu_block_ids=[gpu_blk.block_id], |
| 669 | cpu_block_ids=[cpu_blk.block_id], |
| 670 | ) |
| 671 | |
| 672 | events = list(sched.take_events()) |
| 673 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 674 | assert len(stored) == 2 |
| 675 | by_group = {ev.group_idx: ev for ev in stored} |
| 676 | assert set(by_group) == {0, 1} |
| 677 | assert by_group[0].kv_cache_spec_kind == "full_attention" |
| 678 | assert by_group[0].kv_cache_spec_sliding_window is None |
| 679 | assert by_group[1].kv_cache_spec_kind == "sliding_window" |
| 680 | assert by_group[1].kv_cache_spec_sliding_window == 32 |
| 681 | for ev in stored: |
| 682 | assert ev.locality == "LOCAL" |
| 683 | |
| 684 | |
| 685 | def test_disk_mode_block_removed_relabeled_storage() -> None: |
| 686 | """Disk mode: evicted BlockRemoved relabeled to MEDIUM_STORAGE.""" |
| 687 | fix = make_events_scheduler( |
| 688 | num_cpu_blocks=5, |
| 689 | num_gpu_blocks=16, |
| 690 | disk_capacity_bytes=_BYTES_PER_BLOCK * 5, |
| 691 | ) |
| 692 | sched = fix.scheduler |
| 693 | assert sched.kv_event_medium == MEDIUM_STORAGE |
| 694 | |
| 695 | # Fill disk with 4 blocks (2 + 2). |
| 696 | _do_eager_store(fix, num_blocks=2) |
| 697 | list(sched.take_events()) |
| 698 | _do_eager_store(fix, num_blocks=2) |
| 699 | list(sched.take_events()) |
| 700 | |
| 701 | # Store 2 more -> evicts 2 LRU blocks -> BlockRemoved relabeled STORAGE. |
| 702 | _do_eager_store(fix, num_blocks=2) |
| 703 | events = list(sched.take_events()) |
| 704 | |
| 705 | removed = [e for e in events if isinstance(e, BlockRemoved)] |
| 706 | assert len(removed) >= 2, ( |
| 707 | f"expected >=2 BlockRemoved from eviction, got {len(removed)}" |
| 708 | ) |
| 709 | for ev in removed: |
| 710 | assert ev.medium == MEDIUM_STORAGE, ( |
| 711 | f"BlockRemoved should be relabeled STORAGE, got {ev.medium}" |
| 712 | ) |
| 713 | assert ev.locality == "LOCAL", ( |
| 714 | f"BlockRemoved should be LOCAL, got {ev.locality}" |
| 715 | ) |
| 716 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 717 | assert len(stored) == 2 |
| 718 | for ev in stored: |
| 719 | assert ev.medium == MEDIUM_STORAGE |
| 720 | |
| 721 | |
| 722 | def test_mamba_align_skips_positional_event_metadata() -> None: |
| 723 | """Mamba align emits events only from explicit boundary handoffs.""" |
| 724 | vbs = BLOCK_SIZE * 2 |
| 725 | kv_cfg = _make_mixed_kv_cache_config( |
| 726 | num_blocks=16, with_mamba=True, mamba_block_size=32 |
| 727 | ) |
| 728 | fix = make_events_scheduler( |
| 729 | num_cpu_blocks=8, |
| 730 | num_gpu_blocks=16, |
| 731 | kv_cache_config=kv_cfg, |
| 732 | dcp_world_size=2, |
| 733 | num_groups=2, |
| 734 | ) |
| 735 | sched = fix.scheduler |
| 736 | gpu_pool = fix.gpu_block_pool |
| 737 | |
| 738 | req = _make_cp_request(num_blocks=4, virtual_block_size=vbs) |
| 739 | fa_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=0) |
| 740 | mamba_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs, group_id=1) |
| 741 | kv_blocks = KVCacheBlocks(blocks=(fa_blocks, mamba_blocks)) |
| 742 | req.num_computed_tokens = 4 * vbs |
| 743 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 744 | |
| 745 | block_ids = kv_blocks.get_block_ids() |
| 746 | sched_out = make_scheduler_output( |
| 747 | {req.request_id: 4 * vbs}, |
| 748 | new_reqs={req.request_id: block_ids}, |
| 749 | ) |
| 750 | meta = sched.build_connector_meta(sched_out) |
| 751 | assert meta.store_event >= 0 |
| 752 | simulate_store_completion(sched, meta.store_event) |
| 753 | |
| 754 | events = list(sched.take_events()) |
| 755 | stored = [e for e in events if isinstance(e, BlockStored)] |
| 756 | assert len(stored) == 2 |
| 757 | assert {ev.group_idx for ev in stored} == {0} |
| 758 | assert req.prompt_token_ids is not None |
| 759 | for idx, ev in enumerate(stored): |
| 760 | assert ev.block_size == vbs |
| 761 | assert ev.token_ids == req.prompt_token_ids[idx * vbs : (idx + 1) * vbs] |
| 762 | expected_hash = maybe_convert_block_hash( |
| 763 | get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0)) |
| 764 | ) |
| 765 | assert ev.block_hashes == [expected_hash] |
| 766 | |
| 767 | |
| 768 | def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None: |
| 769 | """Event metadata is resolved only for prefix-cacheable groups. |
| 770 | |
| 771 | A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not |
| 772 | divide the hash block size, so viewing the request hashes at its block |
| 773 | size is impossible; the store must skip it and still emit BlockStored |
| 774 | for the attention blocks. |
| 775 | """ |
| 776 | fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16)) |
| 777 | sched = fix.scheduler |
| 778 | gpu_pool = fix.gpu_block_pool |
| 779 | num_blocks = 2 |
| 780 | |
| 781 | req = make_request(num_blocks=num_blocks) |
| 782 | fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) |
| 783 | scratch_block = gpu_pool.get_new_blocks(1) |
| 784 | kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) |
| 785 | req.num_computed_tokens = num_blocks * BLOCK_SIZE |
| 786 | sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) |
| 787 | sched_out = make_scheduler_output( |
| 788 | {req.request_id: num_blocks * BLOCK_SIZE}, |
| 789 | new_reqs={req.request_id: kv_blocks.get_block_ids()}, |
| 790 | ) |
| 791 | meta = sched.build_connector_meta(sched_out) |
| 792 | assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) |
| 793 | simulate_store_completion(sched, meta.store_event) |
| 794 | |
| 795 | events = list(sched.take_events()) |
| 796 | assert len(events) == num_blocks |
| 797 | assert all(isinstance(e, BlockStored) for e in events) |
| 798 | assert {e.group_idx for e in events} == {0} |
| 799 | assert all(len(e.token_ids) == BLOCK_SIZE for e in events) |