[Bugfix][KV Offload] Allow explicit SimpleCPUOffload block counts (#61674)

Signed-off-by: Sherif Waly <sherif.waly@mistral.ai>
Co-authored-by: vllm-agent <claw@inferact.ai>

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

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(

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(