@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c88d3d4775

[SimpleCPUOffloadConnector] PCP + DCP support (#39831)

Signed-off-by: Jonathan Chen <chenleejonathan@gmail.com>

Jonathan Chen committed Jun 20, 2026 at 18:01 UTC c88d3d4775c41793543e97f1609d3161a5689905
2 files changed +243 -7
tests/v1/simple_kv_offload/test_scheduler.py
+230
@@ -6,6 +6,7 @@ from __future__ import annotations
6
7 from dataclasses import dataclass
8
9 +import pytest
10 import torch
11
12 from vllm import SamplingParams
@@ -1528,3 +1529,232 @@ def test_reset_pending_loads() -> None:
1529 # All GPU blocks free
1530 num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
1531 assert num_used == 1, f"Expected only null block in use, got {num_used}"
1532 +
1533 +
1534 +def _make_cp_vllm_config(
1535 + dcp_world_size: int = 1,
1536 + pcp_world_size: int = 1,
1537 +) -> VllmConfig:
1538 + """VllmConfig with context-parallel sizes set for scheduler-only tests."""
1539 + cfg = _make_vllm_config()
1540 +
1541 + cfg.parallel_config.decode_context_parallel_size = dcp_world_size
1542 + cfg.parallel_config.prefill_context_parallel_size = pcp_world_size
1543 + return cfg
1544 +
1545 +
1546 +def _make_cp_scheduler(
1547 + *,
1548 + dcp_world_size: int = 1,
1549 + pcp_world_size: int = 1,
1550 + num_cpu_blocks: int = 8,
1551 + num_gpu_blocks: int = 16,
1552 + lazy: bool = False,
1553 +) -> SchedulerFixture:
1554 + """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""
1555 + cp_world_size = dcp_world_size * pcp_world_size
1556 + virtual_block_size = BLOCK_SIZE * cp_world_size
1557 +
1558 + kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
1559 + vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)
1560 + cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks
1561 +
1562 + sched = SimpleCPUOffloadScheduler(
1563 + vllm_config=vllm_config,
1564 + kv_cache_config=kv_cache_config,
1565 + cpu_capacity_bytes=cpu_capacity_bytes,
1566 + scheduler_block_size=virtual_block_size,
1567 + hash_block_size=virtual_block_size,
1568 + lazy_offload=lazy,
1569 + )
1570 +
1571 + gpu_block_pool = BlockPool(
1572 + num_gpu_blocks=num_gpu_blocks,
1573 + enable_caching=True,
1574 + hash_block_size=virtual_block_size,
1575 + )
1576 + sched.bind_gpu_block_pool(gpu_block_pool)
1577 +
1578 + return SchedulerFixture(
1579 + scheduler=sched,
1580 + gpu_block_pool=gpu_block_pool,
1581 + vllm_config=vllm_config,
1582 + kv_cache_config=kv_cache_config,
1583 + )
1584 +
1585 +
1586 +def _make_cp_request(
1587 + num_blocks: int,
1588 + virtual_block_size: int,
1589 + request_id: str | None = None,
1590 +) -> Request:
1591 + """Create a request whose block hashes are computed at the virtual
1592 + (CP-scaled) block size, matching what the real scheduler does.
1593 + """
1594 + global _req_counter
1595 + _req_counter += 1
1596 + if request_id is None:
1597 + request_id = f"req-cp-{_req_counter}"
1598 +
1599 + num_tokens = num_blocks * virtual_block_size + 1
1600 + start = _req_counter * 10000
1601 + prompt_token_ids = list(range(start, start + num_tokens))
1602 + sampling_params = SamplingParams(max_tokens=1)
1603 +
1604 + return Request(
1605 + request_id=request_id,
1606 + prompt_token_ids=prompt_token_ids,
1607 + sampling_params=sampling_params,
1608 + pooling_params=None,
1609 + mm_features=None,
1610 + block_hasher=get_request_block_hasher(virtual_block_size, sha256),
1611 + )
1612 +
1613 +
1614 +def _allocate_cp_gpu_blocks(
1615 + gpu_block_pool: BlockPool,
1616 + request: Request,
1617 + num_blocks: int,
1618 + virtual_block_size: int,
1619 + group_id: int = 0,
1620 +) -> list:
1621 + """Allocate GPU blocks and cache them using the CP-scaled block size."""
1622 + blocks = gpu_block_pool.get_new_blocks(num_blocks)
1623 + num_full = min(num_blocks, len(request.block_hashes))
1624 + if num_full > 0:
1625 + gpu_block_pool.cache_full_blocks(
1626 + request=request,
1627 + blocks=blocks,
1628 + num_cached_blocks=0,
1629 + num_full_blocks=num_full,
1630 + block_size=virtual_block_size,
1631 + kv_cache_group_id=group_id,
1632 + )
1633 + return blocks
1634 +
1635 +
1636 +# ---------------------------------------------------------------------------
1637 +# Test 15: CP block size scaling is correct
1638 +# ---------------------------------------------------------------------------
1639 +@pytest.mark.parametrize(
1640 + "dcp_world_size, pcp_world_size",
1641 + [
1642 + (2, 1), # DCP only
1643 + (1, 2), # PCP only
1644 + (2, 2), # DCP + PCP
1645 + ],
1646 +)
1647 +def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:
1648 + """Verify that the scheduler's block_size and cp_world_size are correctly
1649 + scaled when context parallelism is enabled."""
1650 + fix = _make_cp_scheduler(
1651 + dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
1652 + )
1653 + sched = fix.scheduler
1654 +
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 +
1659 +
1660 +# ---------------------------------------------------------------------------
1661 +# Test 16: CP eager store-and-load roundtrip
1662 +# ---------------------------------------------------------------------------
1663 +@pytest.mark.parametrize(
1664 + "dcp_world_size, pcp_world_size",
1665 + [
1666 + (2, 1),
1667 + (1, 2),
1668 + ],
1669 +)
1670 +def test_cp_eager_store_and_load_roundtrip(
1671 + dcp_world_size: int, pcp_world_size: int
1672 +) -> None:
1673 + """With CP enabled, store blocks to CPU and reload them for a new request
1674 + with matching tokens. Verifies that hash matching and transfer-pair
1675 + construction work with the virtual block size."""
1676 + fix = _make_cp_scheduler(
1677 + dcp_world_size=dcp_world_size,
1678 + pcp_world_size=pcp_world_size,
1679 + num_cpu_blocks=8,
1680 + num_gpu_blocks=16,
1681 + lazy=False,
1682 + )
1683 + sched = fix.scheduler
1684 + cp = dcp_world_size * pcp_world_size
1685 + vbs = BLOCK_SIZE * cp
1686 +
1687 + num_blocks = 2
1688 + req = _make_cp_request(num_blocks, vbs)
1689 +
1690 + # Allocate GPU blocks and register hashes
1691 + gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs)
1692 + kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
1693 + req.num_computed_tokens = num_blocks * vbs
1694 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
1695 +
1696 + block_ids = kv_blocks.get_block_ids()
1697 + sched_out = make_scheduler_output(
1698 + {req.request_id: num_blocks * vbs},
1699 + new_reqs={req.request_id: block_ids},
1700 + )
1701 +
1702 + meta = sched.build_connector_meta(sched_out)
1703 + assert meta.store_event >= 0, "Expected a store event"
1704 + assert len(meta.store_gpu_blocks) == num_blocks
1705 + assert len(meta.store_cpu_blocks) == num_blocks
1706 + simulate_store_completion(sched, meta.store_event)
1707 +
1708 + # New request with same tokens — should get a full CPU cache hit.
1709 + req2 = Request(
1710 + request_id="req-cp-load",
1711 + prompt_token_ids=req.prompt_token_ids,
1712 + sampling_params=req.sampling_params,
1713 + pooling_params=None,
1714 + mm_features=None,
1715 + block_hasher=req._block_hasher,
1716 + )
1717 +
1718 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
1719 + assert hit_tokens == num_blocks * vbs
1720 + assert is_async is True
1721 +
1722 + # Allocate fresh GPU blocks for the load.
1723 + gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
1724 + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
1725 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
1726 +
1727 + sched_out2 = make_scheduler_output(
1728 + {req2.request_id: 1},
1729 + new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
1730 + )
1731 + meta2 = sched.build_connector_meta(sched_out2)
1732 + assert meta2.load_event >= 0, "Expected a load event"
1733 + assert len(meta2.load_gpu_blocks) == num_blocks
1734 + assert len(meta2.load_cpu_blocks) == num_blocks
1735 +
1736 +
1737 +# ---------------------------------------------------------------------------
1738 +# Test 17: CP lazy target blocks are scaled correctly
1739 +# ---------------------------------------------------------------------------
1740 +@pytest.mark.parametrize("cp_world_size", [1, 2, 4])
1741 +def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:
1742 + """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1
1743 + because each virtual block covers more tokens."""
1744 + kv_cache_config = _make_kv_cache_config(num_blocks=16)
1745 + max_batched = 64
1746 +
1747 + target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
1748 + kv_cache_config, max_batched, cp_world_size=1
1749 + )
1750 + target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
1751 + kv_cache_config, max_batched, cp_world_size=cp_world_size
1752 + )
1753 +
1754 + if cp_world_size == 1:
1755 + assert target_cp == target_base
1756 + else:
1757 + assert target_cp < target_base, (
1758 + f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "
1759 + f"less than target_base={target_base}"
1760 + )
vllm/v1/simple_kv_offload/manager.py
+13 -7
@@ -82,6 +82,9 @@ class SimpleCPUOffloadScheduler:
82 vllm_config.kv_events_config is not None
83 and vllm_config.kv_events_config.enable_kv_cache_events
84 )
85 + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
86 + pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
87 + self.cp_world_size = dcp_world_size * pcp_world_size
88 self.block_size = scheduler_block_size
89 self.hash_block_size = hash_block_size
90 assert self.block_size % self.hash_block_size == 0
@@ -113,9 +116,6 @@ class SimpleCPUOffloadScheduler:
116 )
117
118 # TODO (yifan): maybe need to enable kv_cache_events and metrics_collector here.
116 - dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
117 - pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
118 - assert dcp_world_size == 1 and pcp_world_size == 1
119 self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
120 kv_cache_config=self.cpu_kv_cache_config,
121 max_model_len=vllm_config.model_config.max_model_len,
@@ -155,6 +155,7 @@ class SimpleCPUOffloadScheduler:
155 self._target_free = self._estimate_lazy_target_blocks(
156 kv_cache_config,
157 vllm_config.scheduler_config.max_num_batched_tokens,
158 + self.cp_world_size,
159 )
160 else:
161 self._target_free = 0
@@ -215,19 +216,22 @@ class SimpleCPUOffloadScheduler:
216
217 @staticmethod
218 def _estimate_lazy_target_blocks(
218 - kv_cache_config: "KVCacheConfig", max_num_batched_tokens: int
219 + kv_cache_config: "KVCacheConfig",
220 + max_num_batched_tokens: int,
221 + cp_world_size: int = 1,
222 ) -> int:
223 """GPU blocks to keep available (free/offloaded) per step in lazy mode."""
224 WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks
225 target = 0
226 for g in kv_cache_config.kv_cache_groups:
227 spec = g.kv_cache_spec
228 + block_size = spec.block_size * cp_world_size
229 if isinstance(spec, MambaSpec):
230 target += 2
231 elif isinstance(spec, SlidingWindowSpec):
228 - target += cdiv(spec.sliding_window, spec.block_size) + 1
232 + target += cdiv(spec.sliding_window, block_size) + 1
233 else:
230 - target += cdiv(max_num_batched_tokens, spec.block_size)
234 + target += cdiv(max_num_batched_tokens, block_size)
235 return int(target * (1 + WATERMARK_RATIO))
236
237 def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
@@ -363,7 +367,9 @@ class SimpleCPUOffloadScheduler:
367 continue
368
369 # Number of blocks in the computed range for this group.
366 - g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
370 + g_block_size = (
371 + kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
372 + )
373 n_computed_g = cdiv(total_computed_tokens, g_block_size)
374
375 # Back-trace: ext blocks sit at the tail of the computed range.