From f648eed23dc48fcc8ba64be0c4182f8e775b5bfa Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sun, 20 Sep 2026 18:58:55 +0800 Subject: [PATCH] [Bugfix][KV Offload] Skip non-prefix-cacheable groups in SimpleCPUOffload (GLM-5.3-Flash kpool tail and QSA) (#56810) Signed-off-by: Jared Wen Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: Zhewen Li --- tests/v1/simple_kv_offload/test_kv_events.py | 36 +++ tests/v1/simple_kv_offload/test_scheduler.py | 223 ++++++++++++++++++- vllm/v1/simple_kv_offload/manager.py | 10 +- 3 files changed, 267 insertions(+), 2 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_kv_events.py b/tests/v1/simple_kv_offload/test_kv_events.py index f21dae7e0b..68121f17e0 100644 --- a/tests/v1/simple_kv_offload/test_kv_events.py +++ b/tests/v1/simple_kv_offload/test_kv_events.py @@ -19,8 +19,10 @@ from tests.v1.simple_kv_offload.test_scheduler import ( SchedulerFixture, _alloc_and_register, _allocate_cp_gpu_blocks, + _allocate_gpu_blocks, _make_cp_request, _make_kv_cache_config, + _make_scratch_kv_cache_config, _make_vllm_config, make_request, make_scheduler_output, @@ -757,3 +759,37 @@ def test_mamba_align_skips_positional_event_metadata() -> None: get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0)) ) assert ev.block_hashes == [expected_hash] + + +def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None: + """Event metadata is resolved only for prefix-cacheable groups. + + A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not + divide the hash block size, so viewing the request hashes at its block + size is impossible; the store must skip it and still emit BlockStored + for the attention blocks. + """ + fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16)) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + num_blocks = 2 + + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + scratch_block = gpu_pool.get_new_blocks(1) + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) + simulate_store_completion(sched, meta.store_event) + + events = list(sched.take_events()) + assert len(events) == num_blocks + assert all(isinstance(e, BlockStored) for e in events) + assert {e.group_idx for e in events} == {0} + assert all(len(e.token_ids) == BLOCK_SIZE for e in events) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index ffc8b20478..ed07718de6 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -41,7 +41,9 @@ from vllm.v1.core.single_type_kv_cache_manager import ( register_all_kvcache_specs, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, + KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -125,6 +127,45 @@ def _make_kv_cache_config( ) +def _make_scratch_kv_cache_config( + num_blocks: int, scratch_block_size: int = 4 +) -> KVCacheConfig: + """FullAttention group plus a non-prefix-cacheable scratch group. + + Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request + scratch block of ``index_kpool`` tokens: it is excluded from prefix + caching, so its block size neither divides nor is divided by the + hash block size in general. + """ + register_all_kvcache_specs(vllm_config=None) + fa_config = _make_kv_cache_config(num_blocks, num_groups=1) + scratch_layers = ["layer_scratch"] + scratch_spec = KpoolTailSpec( + block_size=scratch_block_size, + num_kv_heads=2, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + sliding_window=scratch_block_size, + ) + assert not scratch_spec.prefix_cacheable + scratch_bytes = scratch_spec.page_size_bytes * num_blocks + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=fa_config.kv_cache_tensors + + [ + KVCacheTensor( + size=scratch_bytes, + layers=scratch_layers, + layer_stride=scratch_bytes, + block_stride=scratch_spec.page_size_bytes, + ) + ], + kv_cache_groups=fa_config.kv_cache_groups + + [KVCacheGroupSpec(scratch_layers, scratch_spec)], + ) + + def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig: """Minimal VllmConfig for scheduler tests (no GPU).""" model_config = ModelConfig( @@ -174,9 +215,13 @@ def make_scheduler( num_gpu_blocks: int = 16, num_groups: int = 1, lazy: bool = False, + kv_cache_config: KVCacheConfig | None = None, ) -> SchedulerFixture: """Build a SimpleCPUOffloadScheduler with small block pools.""" - kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) + if kv_cache_config is None: + kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups) + else: + num_groups = len(kv_cache_config.kv_cache_groups) vllm_config = _make_vllm_config() cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups @@ -1392,6 +1437,63 @@ def test_multi_group_null_blocks_skipped() -> None: ) +# --------------------------------------------------------------------------- +# Test 8b: Non-prefix-cacheable scratch groups take no part in store or load +# --------------------------------------------------------------------------- +def test_scratch_group_excluded_from_store_and_load() -> None: + """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded. + + Its single per-request block carries no hash, so the store must offload + only the attention blocks and the load must pair only attention blocks. + """ + fix = make_scheduler( + num_cpu_blocks=8, + num_gpu_blocks=16, + kv_cache_config=_make_scratch_kv_cache_config(16), + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + num_blocks = 2 + + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + scratch_block = gpu_pool.get_new_blocks(1) + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks) + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-scratch-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + scratch_block2 = gpu_pool.get_new_blocks(1) + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2)) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + sched_out2 = make_scheduler_output( + {req2.request_id: 1}, + new_reqs={req2.request_id: kv_blocks2.get_block_ids()}, + ) + meta2 = sched.build_connector_meta(sched_out2) + assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2) + assert len(meta2.load_cpu_blocks) == num_blocks + + # --------------------------------------------------------------------------- # Test 9: Chunked prefill accumulates block_ids across steps # --------------------------------------------------------------------------- @@ -2844,3 +2946,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None: hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) assert hit_tokens == num_blocks * BLOCK_SIZE assert is_async is True + + +QSA_RING_BLOCK_SIZE = 4 + + +def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig: + register_all_kvcache_specs(vllm_config=None) + full_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring_spec = CircularBufferSpec( + block_size=QSA_RING_BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + head_size_v=0, + dtype=DTYPE, + ) + groups = [ + KVCacheGroupSpec(["layer_0"], full_spec), + KVCacheGroupSpec(["qsa_ring"], ring_spec), + ] + tensors = [ + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * num_blocks, + layers=list(group.layer_names), + layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks, + block_stride=group.kv_cache_spec.page_size_bytes, + ) + for group in groups + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + + +def _make_qsa_scheduler( + lazy: bool = False, +) -> tuple[SimpleCPUOffloadScheduler, BlockPool]: + kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + sched = SimpleCPUOffloadScheduler( + vllm_config=_make_vllm_config(), + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + lazy_offload=lazy, + ) + gpu_pool = BlockPool( + num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE + ) + sched.bind_gpu_block_pool(gpu_pool) + return sched, gpu_pool + + +def test_qsa_ring_group_is_never_stored_or_loaded() -> None: + sched, gpu_pool = _make_qsa_scheduler() + assert sched.prefix_cacheable_group_ids == (0,) + + num_blocks = 2 + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + ring_block = gpu_pool.get_new_blocks(1)[0] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block])) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + sched_out = make_scheduler_output( + {req.request_id: num_blocks * BLOCK_SIZE}, + new_reqs={req.request_id: kv_blocks.get_block_ids()}, + ) + meta = sched.build_connector_meta(sched_out) + assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks} + assert ring_block.block_id not in meta.store_gpu_blocks + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-qsa-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True + + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks) + ring_block2 = gpu_pool.get_new_blocks(1)[0] + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2])) + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) + + meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1})) + assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2} + assert ring_block2.block_id not in meta2.load_gpu_blocks + assert len(meta2.load_cpu_blocks) == num_blocks + + +def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: + with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16) + without_ring = KVCacheConfig( + num_blocks=16, + kv_cache_tensors=with_ring.kv_cache_tensors[:1], + kv_cache_groups=with_ring.kv_cache_groups[:1], + ) + max_batched = 64 + target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + with_ring, max_batched + ) + target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks( + without_ring, max_batched + ) + assert target_with == target_without diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 85dde88f2d..96bfb35a67 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -165,6 +165,9 @@ class SimpleCPUOffloadScheduler: kv_cache_config, offload_capacity ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks + self.prefix_cacheable_group_ids = ( + self.cpu_kv_cache_config.prefix_cacheable_group_ids + ) self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU # Find the full attention kv group for prefix cache matching. self.fa_gidx = -1 @@ -297,7 +300,7 @@ class SimpleCPUOffloadScheduler: """GPU blocks to keep available (free/offloaded) per step in lazy mode.""" WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks target = 0 - for g in kv_cache_config.kv_cache_groups: + for g in kv_cache_config.prefix_cacheable_groups: spec = g.kv_cache_spec # Only full attention is sharded across DCP ranks; replicated specs # (mamba, sliding window, chunked-local) keep their own block size. @@ -442,6 +445,9 @@ class SimpleCPUOffloadScheduler: # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): + if g not in self.prefix_cacheable_group_ids: + cpu_hit_blocks.append([]) + continue g_block_size = self.group_block_sizes[g] n_take_g = cdiv(num_external_tokens, g_block_size) cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -828,6 +834,8 @@ class SimpleCPUOffloadScheduler: num_free = self.cpu_block_pool.get_num_free_blocks() for g, group_gpu_ids in enumerate(block_ids_by_group): + if g not in self.prefix_cacheable_group_ids: + continue if len(gpu_block_ids) >= num_free: break group_manager = self.cpu_coordinator.single_type_managers[g]