[SimpleCPUOffloadConnector] Fix remaining global→block conversions under PCP/DCP (#46394)

Signed-off-by: Jeff Ma <jeffjma@umich.edu>

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
 # ---------------------------------------------------------------------------

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]