[SimpleCPUOffloadConnector] Fix remaining global→block conversions under PCP/DCP (#46394)
Signed-off-by: Jeff Ma <jeffjma@umich.edu>
Jeff (Junze) Ma committed
Jun 24, 2026 at 23:05 UTC
36fd7e8b8623f8077a5439bec6092dfd6f9223d7
2 files changed
+94
-6
tests/v1/simple_kv_offload/test_scheduler.py
+81
@@ -1655,6 +1655,7 @@ def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None
1655
expected_cp = dcp_world_size * pcp_world_size
1656
assert sched.cp_world_size == expected_cp
1657
assert sched.block_size == BLOCK_SIZE * expected_cp
1658
+ assert sched.fa_block_size == BLOCK_SIZE * expected_cp
1659
1660
1661
# ---------------------------------------------------------------------------
@@ -1734,6 +1735,86 @@ def test_cp_eager_store_and_load_roundtrip(
1735
assert len(meta2.load_cpu_blocks) == num_blocks
1736
1737
1738
+# ---------------------------------------------------------------------------
1739
+# Test 18: CP store and load use effective block size
1740
+# ---------------------------------------------------------------------------
1741
+@pytest.mark.parametrize(
1742
+ "dcp_world_size, pcp_world_size",
1743
+ [
1744
+ (2, 1),
1745
+ (1, 2),
1746
+ (2, 2),
1747
+ ],
1748
+)
1749
+def test_cp_effective_block_size_store_and_load(
1750
+ dcp_world_size: int, pcp_world_size: int
1751
+) -> None:
1752
+ """Verify ready_blocks_g (store) and n_take_g (load) use the effective
1753
+ (physical * cp) block size, not the per-rank physical size."""
1754
+ fix = _make_cp_scheduler(
1755
+ dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
1756
+ )
1757
+ sched = fix.scheduler
1758
+ gpu_pool = fix.gpu_block_pool
1759
+ cp = dcp_world_size * pcp_world_size
1760
+ vbs = BLOCK_SIZE * cp
1761
+
1762
+ # Store: allocate 2 blocks, confirm only 1. Without the fix,
1763
+ # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks.
1764
+ req = _make_cp_request(num_blocks=2, virtual_block_size=vbs)
1765
+ gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs)
1766
+ kv = KVCacheBlocks(blocks=(gpu_blocks,))
1767
+ req.num_computed_tokens = vbs
1768
+ sched.update_state_after_alloc(req, kv, num_external_tokens=0)
1769
+ m1 = sched.build_connector_meta(
1770
+ make_scheduler_output(
1771
+ {req.request_id: vbs},
1772
+ new_reqs={req.request_id: kv.get_block_ids()},
1773
+ )
1774
+ )
1775
+ assert len(m1.store_gpu_blocks) == 1
1776
+ assert len(m1.store_cpu_blocks) == 1
1777
+ simulate_store_completion(sched, m1.store_event)
1778
+
1779
+ # Load: store 2 blocks from a second request, accept only 1 as external.
1780
+ # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both.
1781
+ req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs)
1782
+ kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),))
1783
+ req2.num_computed_tokens = 2 * vbs
1784
+ sched.update_state_after_alloc(req2, kv2, num_external_tokens=0)
1785
+ m2 = sched.build_connector_meta(
1786
+ make_scheduler_output(
1787
+ {req2.request_id: 2 * vbs},
1788
+ new_reqs={req2.request_id: kv2.get_block_ids()},
1789
+ )
1790
+ )
1791
+ simulate_store_completion(sched, m2.store_event)
1792
+
1793
+ req3 = Request(
1794
+ request_id="req-cp-partial-load",
1795
+ prompt_token_ids=req2.prompt_token_ids,
1796
+ sampling_params=req2.sampling_params,
1797
+ pooling_params=None,
1798
+ mm_features=None,
1799
+ block_hasher=req2._block_hasher,
1800
+ )
1801
+ hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0)
1802
+ assert hit == 2 * vbs
1803
+
1804
+ kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),))
1805
+ sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs)
1806
+ m3 = sched.build_connector_meta(
1807
+ make_scheduler_output(
1808
+ {req3.request_id: vbs},
1809
+ new_reqs={req3.request_id: kv3.get_block_ids()},
1810
+ )
1811
+ )
1812
+ assert m3.load_event >= 0
1813
+ assert len(m3.load_gpu_blocks) == 1
1814
+ assert len(m3.load_cpu_blocks) == 1
1815
+ assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]]
1816
+
1817
+
1818
# ---------------------------------------------------------------------------
1819
# Test 17: CP lazy target blocks are scaled correctly
1820
# ---------------------------------------------------------------------------
vllm/v1/simple_kv_offload/manager.py
+13
-6
@@ -103,9 +103,12 @@ class SimpleCPUOffloadScheduler:
103
assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups)
104
# FA group's own block_size; divides scheduler_block_size (the LCM)
105
# but is NOT assumed to equal it.
106
- self.fa_block_size: int = self.cpu_kv_cache_config.kv_cache_groups[
107
- self.fa_gidx
108
- ].kv_cache_spec.block_size
106
+ self.fa_block_size: int = (
107
+ self.cpu_kv_cache_config.kv_cache_groups[
108
+ self.fa_gidx
109
+ ].kv_cache_spec.block_size
110
+ * self.cp_world_size
111
+ )
112
assert self.block_size % self.fa_block_size == 0
113
114
logger.info(
@@ -348,10 +351,12 @@ class SimpleCPUOffloadScheduler:
351
# the rest will be released along with the temp pin below.
352
cpu_hit_blocks: list[list[KVCacheBlock]] = []
353
for g in range(num_groups):
351
- g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
354
+ g_block_size = (
355
+ kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
356
+ )
357
assert num_external_tokens % g_block_size == 0, (
358
f"num_external_tokens={num_external_tokens} not aligned to "
354
- f"group {g} block_size={g_block_size}"
359
+ f"group {g} effective block_size={g_block_size}"
360
)
361
n_take_g = num_external_tokens // g_block_size
362
cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])
@@ -599,7 +604,9 @@ class SimpleCPUOffloadScheduler:
604
already_stored_g = state.num_stored_blocks[g]
605
group_gpu_ids = block_ids_by_group[g]
606
602
- g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
607
+ g_block_size = (
608
+ kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
609
+ )
610
ready_blocks_g = aligned_tokens // g_block_size
611
scannable = group_gpu_ids[already_stored_g:ready_blocks_g]
612