diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index e3575af63d..98159e38a7 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -411,6 +411,35 @@ Example (OpenAI-compatible completions request): } ``` +## Explicit pool sizing for `SimpleCPUOffloadConnector` + +`SimpleCPUOffloadConnector` accepts an optional `num_cpu_blocks` in +`kv_connector_extra_config`. The scheduler and every worker use this exact +number of offload blocks instead of automatic sizing from +`cpu_bytes_to_use_per_rank` or `cpu_bytes_to_use / world_size`. +Use the same positive integer on all ranks, chosen +to fit the smallest per-worker capacity. This supports pipeline stages with +different KV cache sizes without adding a collective during initialization. + +```json +{ + "kv_connector": "SimpleCPUOffloadConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "cpu_bytes_to_use_per_rank": 1073741824, + "num_cpu_blocks": 64, + "lazy_offload": false + } +} +``` + +The example count must be adjusted for the model, KV layout, and parallelism. +The configured byte budget remains an upper bound: startup fails if the +requested count exceeds a worker's capacity. For the `disk` backend the count +selects disk slots and is checked against `disk_capacity_bytes`. Omitting +`num_cpu_blocks` preserves automatic per-worker sizing. This option applies +to `SimpleCPUOffloadConnector`, separately from `OffloadingConnector` above. + ## Further Reading - [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput). diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 527f7bee05..31aa75d31f 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -245,6 +245,55 @@ def test_connector_reports_configured_cache_source( assert connector.get_external_cache_hit_sources(None, 0) == {} # type: ignore[arg-type] +@pytest.mark.parametrize("num_layers", [1, 2]) +@pytest.mark.parametrize("backend", ["cpu", "disk"]) +def test_connector_explicit_cpu_blocks_align_uneven_stages( + tmp_path, num_layers: int, backend: str +): + """The same configured pool size must fit lighter and heavier PP stages.""" + config = _make_kv_cache_config(num_blocks=16) + names = [f"layer_{i}" for i in range(num_layers)] + config.kv_cache_groups[0].layer_names = names + config.kv_cache_tensors[0].layers = names + config.kv_cache_tensors[0].size *= num_layers + vllm_config = _make_vllm_config() + extra = { + "cpu_bytes_to_use_per_rank": _BYTES_PER_BLOCK * 8, + "num_cpu_blocks": 4, + "kv_offload_backend": backend, + } + if backend == "disk": + extra.update( + disk_path=str(tmp_path / "offload.bin"), + disk_capacity_bytes=_BYTES_PER_BLOCK * 8, + ) + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config = extra + connector = SimpleCPUOffloadConnector( + vllm_config, KVConnectorRole.SCHEDULER, config + ) + assert connector.scheduler_manager is not None + assert connector.scheduler_manager.num_cpu_blocks == 4 + assert len(connector.scheduler_manager.cpu_block_pool.blocks) == 4 + + +@pytest.mark.parametrize( + "count, capacity_bytes", + [(count, _BYTES_PER_BLOCK * 8) for count in [0, -1, True, 1.5, "4", 9]] + + [(1, _BYTES_PER_BLOCK // 2)], +) +def test_connector_rejects_invalid_explicit_cpu_blocks(count, capacity_bytes): + config = _make_kv_cache_config(num_blocks=16) + vllm_config = _make_vllm_config() + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config = { + "cpu_bytes_to_use_per_rank": capacity_bytes, + "num_cpu_blocks": count, + } + with pytest.raises(ValueError, match="num_cpu_blocks"): + SimpleCPUOffloadConnector(vllm_config, KVConnectorRole.SCHEDULER, config) + + def make_scheduler( num_cpu_blocks: int = 8, num_gpu_blocks: int = 16, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 4f9a7c451f..c58cea6788 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -371,6 +371,69 @@ def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout): } +@pytest.mark.parametrize("num_layers", [1, 2]) +def test_explicit_cpu_blocks_bound_live_worker_allocation(monkeypatch, num_layers): + """Different live KV strides must honor the same explicit CPU pool size.""" + num_blocks = 4 + names = [f"layer.{i}" for i in range(num_layers)] + raw = torch.zeros( + num_layers * num_blocks * BLOCK_BYTES, dtype=torch.int8, device="cuda" + ) + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=raw.nbytes, + layers=names, + layer_stride=num_blocks * BLOCK_BYTES, + block_stride=BLOCK_BYTES, + ) + ], + kv_cache_groups=[], + ) + caches = dict(zip(names, raw.view(num_layers, num_blocks, BLOCK_BYTES))) + worker = SimpleCPUOffloadWorker(None, config, 8 * BLOCK_BYTES, num_cpu_blocks=4) + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False) + try: + worker.register_kv_caches(caches) + assert worker.num_cpu_blocks == 4 + assert worker.cpu_kv_caches is not None + assert all(t.shape[0] == 4 for t in worker.cpu_kv_caches.values()) + assert sum(t.nbytes for t in worker.cpu_kv_caches.values()) <= 8 * BLOCK_BYTES + finally: + if worker._backend is not None: + worker._backend.shutdown() + + +@pytest.mark.parametrize( + "capacity_bytes, count, local_capacity", + [(8 * BLOCK_BYTES, 5, 4), (BLOCK_BYTES, 1, 0)], +) +def test_explicit_cpu_blocks_reject_heavier_worker_before_allocation( + capacity_bytes, count, local_capacity +): + """A count that fits the scheduler's stage must fail on a smaller worker.""" + names = ["layer.0", "layer.1"] + raw = torch.zeros(2 * 4 * BLOCK_BYTES, dtype=torch.int8, device="cuda") + config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=raw.nbytes, + layers=names, + layer_stride=4 * BLOCK_BYTES, + block_stride=BLOCK_BYTES, + ) + ], + kv_cache_groups=[], + ) + worker = SimpleCPUOffloadWorker(None, config, capacity_bytes, num_cpu_blocks=count) + caches = dict(zip(names, raw.view(2, 4, BLOCK_BYTES))) + with pytest.raises(ValueError, match=f"local offload capacity {local_capacity}"): + worker.register_kv_caches(caches) + assert worker.cpu_kv_caches is None + + def test_register_kv_cache_storage_with_trailing_padding(monkeypatch): num_blocks = 4 block_bytes = 32 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 0bd8ebe3cd..49afa017f6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -101,6 +101,16 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): ) cpu_capacity_per_rank = explicit + # When set, num_cpu_blocks overrides automatic pool sizing from byte budgets + # (including cpu_bytes_to_use and cpu_bytes_to_use_per_rank), allowing ranks + # with different KV layouts (e.g. PP stages) to use the same pool size. + # Byte budgets still bound capacity. + num_cpu_blocks = extra_config.get("num_cpu_blocks") + if num_cpu_blocks is not None and ( + type(num_cpu_blocks) is not int or num_cpu_blocks <= 0 + ): + raise ValueError("num_cpu_blocks must be a positive integer") + lazy_offload = bool(extra_config.get("lazy_offload", False)) kv_offload_backend = str(extra_config.get("kv_offload_backend", "cpu")) @@ -160,6 +170,12 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): kv_offload_backend, disk_path or "none", ) + if num_cpu_blocks is not None: + logger.info( + "Using explicit num_cpu_blocks=%d for offload pool sizing " + "instead of automatic sizing from byte budgets.", + num_cpu_blocks, + ) if role == KVConnectorRole.SCHEDULER: from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes @@ -177,6 +193,7 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): lazy_offload=lazy_offload, disk_capacity_bytes=disk_capacity_bytes if disk_mode else 0, use_page_cache=use_page_cache if disk_mode else False, + num_cpu_blocks=num_cpu_blocks, ) elif role == KVConnectorRole.WORKER: self.worker_handler = SimpleCPUOffloadWorker( @@ -188,6 +205,7 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): disk_capacity_bytes=disk_capacity_bytes, disk_buffer_slots=disk_buffer_slots, use_page_cache=use_page_cache, + num_cpu_blocks=num_cpu_blocks, ) # --- Worker-side methods --- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index dfc24e420a..8cf5f4f985 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -151,6 +151,7 @@ class SimpleCPUOffloadScheduler: lazy_offload: bool = False, disk_capacity_bytes: int = 0, use_page_cache: bool = False, + num_cpu_blocks: int | None = None, ): self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config @@ -170,7 +171,7 @@ class SimpleCPUOffloadScheduler: # Derive a CPU KVCacheConfig from the GPU config and build a coordinator assert kv_cache_config is not None self.cpu_kv_cache_config = self._derive_cpu_config( - kv_cache_config, offload_capacity + kv_cache_config, offload_capacity, num_cpu_blocks ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks self.prefix_cacheable_group_ids = ( @@ -278,7 +279,9 @@ class SimpleCPUOffloadScheduler: @staticmethod def _derive_cpu_config( - gpu_config: "KVCacheConfig", cpu_capacity_bytes: int + gpu_config: "KVCacheConfig", + cpu_capacity_bytes: int, + num_cpu_blocks: int | None = None, ) -> "KVCacheConfig": """Derive a CPU KVCacheConfig from the GPU config. Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio.""" @@ -291,7 +294,16 @@ class SimpleCPUOffloadScheduler: # so its size is the total GPU KV cache size. gpu_total_bytes = gpu_config.kv_cache_tensors[0].size num_gpu_blocks = gpu_config.num_blocks - num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes) + local_capacity = num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes + if num_cpu_blocks is None: + num_cpu_blocks = max(1, local_capacity) + elif ( + type(num_cpu_blocks) is not int or not 0 < num_cpu_blocks <= local_capacity + ): + raise ValueError( + f"num_cpu_blocks={num_cpu_blocks!r} must be between 1 and " + f"local offload capacity {local_capacity}" + ) # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally. cpu_tensors = [ KVCacheTensor( diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index c120951eed..2ad318b987 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -36,10 +36,12 @@ class SimpleCPUOffloadWorker: disk_capacity_bytes: int = 0, disk_buffer_slots: int = 2, use_page_cache: bool = False, + num_cpu_blocks: int | None = None, ): self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config self.cpu_capacity_bytes = cpu_capacity_bytes + self._num_cpu_blocks_override = num_cpu_blocks self.disk_path = disk_path self.disk_capacity_bytes = disk_capacity_bytes self.disk_buffer_slots = disk_buffer_slots @@ -159,7 +161,10 @@ class SimpleCPUOffloadWorker: ] total_bytes_per_block = sum(per_tensor_bpb) - self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block) + self.num_cpu_blocks = self._resolve_num_cpu_blocks( + self.disk_capacity_bytes if self.disk_mode else self.cpu_capacity_bytes, + total_bytes_per_block, + ) # Use lowest priority so KV cache I/O yields to compute streams. low_pri, _ = torch.cuda.Stream.priority_range() @@ -173,13 +178,27 @@ class SimpleCPUOffloadWorker: else: self._init_cpu_mode(unique_gpu_caches, total_bytes_per_block, self.device) + def _resolve_num_cpu_blocks(self, capacity_bytes: int, bytes_per_block: int) -> int: + local_capacity = capacity_bytes // bytes_per_block + requested = self._num_cpu_blocks_override + if requested is None: + return max(1, local_capacity) + if type(requested) is not int or not 0 < requested <= local_capacity: + raise ValueError( + f"num_cpu_blocks={requested!r} must be between 1 and " + f"local offload capacity {local_capacity}" + ) + return requested + def _init_disk_mode( self, unique_gpu_caches: dict[str, torch.Tensor], total_bytes_per_block: int, device: torch.device, ) -> None: - num_disk_slots = max(1, self.disk_capacity_bytes // total_bytes_per_block) + num_disk_slots = self._resolve_num_cpu_blocks( + self.disk_capacity_bytes, total_bytes_per_block + ) self.num_cpu_blocks = num_disk_slots logger.info(