From 36fd7e8b8623f8077a5439bec6092dfd6f9223d7 Mon Sep 17 00:00:00 2001 From: "Jeff (Junze) Ma" <93145857+majunze2001@users.noreply.github.com> Date: Wed, 24 Jun 2026 23:05:24 -0700 Subject: [PATCH] =?UTF-8?q?[SimpleCPUOffloadConnector]=20Fix=20remaining?= =?UTF-8?q?=20global=E2=86=92block=20conversions=20under=20PCP/DCP=20(#463?= =?UTF-8?q?94)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Jeff Ma --- tests/v1/simple_kv_offload/test_scheduler.py | 81 ++++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 19 +++-- 2 files changed, 94 insertions(+), 6 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 1ec986eada..09586f5e6b 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1655,6 +1655,7 @@ def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None expected_cp = dcp_world_size * pcp_world_size assert sched.cp_world_size == expected_cp assert sched.block_size == BLOCK_SIZE * expected_cp + assert sched.fa_block_size == BLOCK_SIZE * expected_cp # --------------------------------------------------------------------------- @@ -1734,6 +1735,86 @@ def test_cp_eager_store_and_load_roundtrip( assert len(meta2.load_cpu_blocks) == num_blocks +# --------------------------------------------------------------------------- +# Test 18: CP store and load use effective block size +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "dcp_world_size, pcp_world_size", + [ + (2, 1), + (1, 2), + (2, 2), + ], +) +def test_cp_effective_block_size_store_and_load( + dcp_world_size: int, pcp_world_size: int +) -> None: + """Verify ready_blocks_g (store) and n_take_g (load) use the effective + (physical * cp) block size, not the per-rank physical size.""" + fix = _make_cp_scheduler( + dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + cp = dcp_world_size * pcp_world_size + vbs = BLOCK_SIZE * cp + + # Store: allocate 2 blocks, confirm only 1. Without the fix, + # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks. + req = _make_cp_request(num_blocks=2, virtual_block_size=vbs) + gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs) + kv = KVCacheBlocks(blocks=(gpu_blocks,)) + req.num_computed_tokens = vbs + sched.update_state_after_alloc(req, kv, num_external_tokens=0) + m1 = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: vbs}, + new_reqs={req.request_id: kv.get_block_ids()}, + ) + ) + assert len(m1.store_gpu_blocks) == 1 + assert len(m1.store_cpu_blocks) == 1 + simulate_store_completion(sched, m1.store_event) + + # Load: store 2 blocks from a second request, accept only 1 as external. + # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both. + req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs) + kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),)) + req2.num_computed_tokens = 2 * vbs + sched.update_state_after_alloc(req2, kv2, num_external_tokens=0) + m2 = sched.build_connector_meta( + make_scheduler_output( + {req2.request_id: 2 * vbs}, + new_reqs={req2.request_id: kv2.get_block_ids()}, + ) + ) + simulate_store_completion(sched, m2.store_event) + + req3 = Request( + request_id="req-cp-partial-load", + prompt_token_ids=req2.prompt_token_ids, + sampling_params=req2.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req2._block_hasher, + ) + hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0) + assert hit == 2 * vbs + + kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),)) + sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs) + m3 = sched.build_connector_meta( + make_scheduler_output( + {req3.request_id: vbs}, + new_reqs={req3.request_id: kv3.get_block_ids()}, + ) + ) + assert m3.load_event >= 0 + assert len(m3.load_gpu_blocks) == 1 + assert len(m3.load_cpu_blocks) == 1 + assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]] + + # --------------------------------------------------------------------------- # Test 17: CP lazy target blocks are scaled correctly # --------------------------------------------------------------------------- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 07978a9dd6..dfaa2234eb 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -103,9 +103,12 @@ class SimpleCPUOffloadScheduler: assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups) # FA group's own block_size; divides scheduler_block_size (the LCM) # but is NOT assumed to equal it. - self.fa_block_size: int = self.cpu_kv_cache_config.kv_cache_groups[ - self.fa_gidx - ].kv_cache_spec.block_size + self.fa_block_size: int = ( + self.cpu_kv_cache_config.kv_cache_groups[ + self.fa_gidx + ].kv_cache_spec.block_size + * self.cp_world_size + ) assert self.block_size % self.fa_block_size == 0 logger.info( @@ -348,10 +351,12 @@ 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): - g_block_size = kv_cache_groups[g].kv_cache_spec.block_size + g_block_size = ( + kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size + ) assert num_external_tokens % g_block_size == 0, ( f"num_external_tokens={num_external_tokens} not aligned to " - f"group {g} block_size={g_block_size}" + f"group {g} effective block_size={g_block_size}" ) n_take_g = num_external_tokens // g_block_size cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -599,7 +604,9 @@ class SimpleCPUOffloadScheduler: already_stored_g = state.num_stored_blocks[g] group_gpu_ids = block_ids_by_group[g] - g_block_size = kv_cache_groups[g].kv_cache_spec.block_size + g_block_size = ( + kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size + ) ready_blocks_g = aligned_tokens // g_block_size scannable = group_gpu_ids[already_stored_g:ready_blocks_g]