| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Worker-side handler for SimpleCPUOffloadConnector.""" |
| 4 | |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.utils.torch_utils import PIN_MEMORY |
| 12 | from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend |
| 13 | from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor |
| 14 | from vllm.v1.simple_kv_offload.disk_backend import DiskBackend |
| 15 | from vllm.v1.simple_kv_offload.metadata import ( |
| 16 | SimpleCPUOffloadMetadata, |
| 17 | SimpleCPUOffloadWorkerMetadata, |
| 18 | ) |
| 19 | |
| 20 | if TYPE_CHECKING: |
| 21 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | class SimpleCPUOffloadWorker: |
| 27 | """Worker-side handler for CPU offloading transfers.""" |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | vllm_config: VllmConfig, |
| 32 | kv_cache_config: "KVCacheConfig | None", |
| 33 | cpu_capacity_bytes: int, |
| 34 | kv_offload_backend: str = "cpu", |
| 35 | disk_path: str | None = None, |
| 36 | disk_capacity_bytes: int = 0, |
| 37 | disk_buffer_slots: int = 2, |
| 38 | use_page_cache: bool = False, |
| 39 | ): |
| 40 | self.vllm_config = vllm_config |
| 41 | self.kv_cache_config = kv_cache_config |
| 42 | self.cpu_capacity_bytes = cpu_capacity_bytes |
| 43 | self.disk_path = disk_path |
| 44 | self.disk_capacity_bytes = disk_capacity_bytes |
| 45 | self.disk_buffer_slots = disk_buffer_slots |
| 46 | self.use_page_cache = use_page_cache |
| 47 | self.disk_mode = kv_offload_backend == "disk" |
| 48 | |
| 49 | self.gpu_kv_caches: dict[str, torch.Tensor] | None = None |
| 50 | self.cpu_kv_caches: dict[str, torch.Tensor] | None = None |
| 51 | self.device: torch.device | None = None |
| 52 | self.num_cpu_blocks: int = 0 |
| 53 | |
| 54 | # CUDA streams for the async transfers |
| 55 | self.load_stream: torch.cuda.Stream | None = None |
| 56 | self.store_stream: torch.cuda.Stream | None = None |
| 57 | |
| 58 | self._backend: DmaCopyBackend | DiskBackend | None = None |
| 59 | |
| 60 | # Ordered (event_idx, Event). Events pre-allocated on main thread. |
| 61 | self._load_events: list[tuple[int, torch.Event]] = [] |
| 62 | self._store_events: list[tuple[int, torch.Event]] = [] |
| 63 | # High-water marks: highest event_idx completed per stream. |
| 64 | # When the event list is empty, the hwm covers all prior events. |
| 65 | self._load_hwm: int = -1 |
| 66 | self._store_hwm: int = -1 |
| 67 | |
| 68 | # Metadata for the current step |
| 69 | self._connector_metadata: SimpleCPUOffloadMetadata | None = None |
| 70 | |
| 71 | # Compute-done event recorded before each store; reused across steps |
| 72 | # (get_finished runs once per step, copy queue is FIFO). |
| 73 | self._store_compute_done: torch.Event | None = None |
| 74 | |
| 75 | # Pending event index sets, populated in bind_connector_metadata |
| 76 | self._pending_load_event_indices: set[int] = set() |
| 77 | self._pending_store_event_indices: set[int] = set() |
| 78 | # Completed store events to report via build_connector_worker_meta |
| 79 | self._completed_store_events: dict[int, int] = {} |
| 80 | |
| 81 | def register_kv_caches( |
| 82 | self, |
| 83 | kv_caches: dict[str, torch.Tensor], |
| 84 | ) -> None: |
| 85 | """Register GPU KV caches and allocate pinned CPU tensors. |
| 86 | The worker will infer the underlying raw storage from the kv_caches. |
| 87 | |
| 88 | Args: |
| 89 | kv_caches: Per-layer GPU KV caches. Values are either a single |
| 90 | tensor (attention layers) or a list of tensors (Mamba layers |
| 91 | in hybrid models). All values are included for offloading |
| 92 | by resolving to their underlying raw storage. |
| 93 | |
| 94 | """ |
| 95 | if not kv_caches: |
| 96 | logger.warning("No KV caches to offload.") |
| 97 | return |
| 98 | |
| 99 | self.device = next(iter(kv_caches.values())).device |
| 100 | |
| 101 | assert self.kv_cache_config is not None |
| 102 | num_blocks = self.kv_cache_config.num_blocks |
| 103 | assert self.kv_cache_config.kv_cache_tensors |
| 104 | |
| 105 | # The DMA backend copies blocks as base + block_id * stride(0), so every |
| 106 | # region is a [num_blocks, block_bytes] view. Block bytes come from each |
| 107 | # tensor's placement -- layers in one allocation can differ in page size. |
| 108 | unique_gpu_caches: dict[str, torch.Tensor] = {} |
| 109 | seen: set[tuple[torch.device, int]] = set() |
| 110 | for kv_cache_tensor in self.kv_cache_config.kv_cache_tensors: |
| 111 | name = next((n for n in kv_cache_tensor.layers if n in kv_caches), None) |
| 112 | if name is None: |
| 113 | continue |
| 114 | num_layers = len(kv_cache_tensor.layers) |
| 115 | block_bytes = kv_cache_tensor.block_stride |
| 116 | blocks_span = num_blocks * block_bytes |
| 117 | layer_outermost = kv_cache_tensor.layer_stride >= blocks_span |
| 118 | if layer_outermost: |
| 119 | # One region per layer, or per KV head group under LHBNC. |
| 120 | start = kv_cache_tensor.offset |
| 121 | span = num_layers * kv_cache_tensor.layer_stride |
| 122 | layer_names = kv_cache_tensor.layers |
| 123 | else: |
| 124 | start = 0 |
| 125 | span = blocks_span |
| 126 | layer_names = kv_cache_tensor.layers[:1] |
| 127 | |
| 128 | cache = kv_caches[name] |
| 129 | raw = torch.empty(0, dtype=torch.int8, device=cache.device).set_( |
| 130 | cache.untyped_storage() |
| 131 | ) |
| 132 | assert raw.numel() >= start + span, ( |
| 133 | f"KV cache {name!r} storage has {raw.numel()} bytes, smaller than " |
| 134 | f"the {span} bytes it places at offset {start}" |
| 135 | ) |
| 136 | regions = raw[start : start + span].view(-1, num_blocks, block_bytes) |
| 137 | groups_per_layer, remainder = divmod(len(regions), len(layer_names)) |
| 138 | assert remainder == 0, ( |
| 139 | f"KV cache {name!r} places {len(layer_names)} layers in " |
| 140 | f"{len(regions)} block-strided regions" |
| 141 | ) |
| 142 | for idx, region in enumerate(regions): |
| 143 | key = (region.device, region.data_ptr()) |
| 144 | if key in seen: |
| 145 | continue |
| 146 | seen.add(key) |
| 147 | layer_name = layer_names[idx // groups_per_layer] |
| 148 | key_name = ( |
| 149 | layer_name |
| 150 | if groups_per_layer == 1 |
| 151 | else f"{layer_name}.{idx % groups_per_layer}" |
| 152 | ) |
| 153 | unique_gpu_caches[key_name] = region |
| 154 | |
| 155 | # Compute per-tensor bytes_per_block. Tensors may have different |
| 156 | # page_size_bytes (e.g., UniformTypeKVCacheSpecs with varying head_size). |
| 157 | per_tensor_bpb = [ |
| 158 | t.stride(0) * t.element_size() for t in unique_gpu_caches.values() |
| 159 | ] |
| 160 | total_bytes_per_block = sum(per_tensor_bpb) |
| 161 | |
| 162 | self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block) |
| 163 | |
| 164 | # Use lowest priority so KV cache I/O yields to compute streams. |
| 165 | low_pri, _ = torch.cuda.Stream.priority_range() |
| 166 | self.load_stream = torch.cuda.Stream(priority=low_pri) |
| 167 | self.store_stream = torch.cuda.Stream(priority=low_pri) |
| 168 | |
| 169 | self.gpu_kv_caches = unique_gpu_caches |
| 170 | |
| 171 | if self.disk_mode: |
| 172 | self._init_disk_mode(unique_gpu_caches, total_bytes_per_block, self.device) |
| 173 | else: |
| 174 | self._init_cpu_mode(unique_gpu_caches, total_bytes_per_block, self.device) |
| 175 | |
| 176 | def _init_disk_mode( |
| 177 | self, |
| 178 | unique_gpu_caches: dict[str, torch.Tensor], |
| 179 | total_bytes_per_block: int, |
| 180 | device: torch.device, |
| 181 | ) -> None: |
| 182 | num_disk_slots = max(1, self.disk_capacity_bytes // total_bytes_per_block) |
| 183 | self.num_cpu_blocks = num_disk_slots |
| 184 | |
| 185 | logger.info( |
| 186 | "SimpleCPUOffloadWorker [DISK]: %d tensors, %d disk slots (%.2f GB)", |
| 187 | len(unique_gpu_caches), |
| 188 | num_disk_slots, |
| 189 | (num_disk_slots * total_bytes_per_block) / (1024**3), |
| 190 | ) |
| 191 | |
| 192 | assert self.disk_path is not None |
| 193 | rank_path = f"{self.disk_path}.rank_{device.index or 0}" |
| 194 | self._backend = DiskBackend() |
| 195 | self._backend.init( |
| 196 | unique_gpu_caches, |
| 197 | device, |
| 198 | self.load_stream, |
| 199 | self.store_stream, |
| 200 | rank_path, |
| 201 | num_disk_slots, |
| 202 | total_bytes_per_block, |
| 203 | self.disk_buffer_slots, |
| 204 | self.use_page_cache, |
| 205 | ) |
| 206 | |
| 207 | def _init_cpu_mode( |
| 208 | self, |
| 209 | unique_gpu_caches: dict[str, torch.Tensor], |
| 210 | total_bytes_per_block: int, |
| 211 | device: torch.device, |
| 212 | ) -> None: |
| 213 | logger.info( |
| 214 | "SimpleCPUOffloadWorker [CPU]: %d tensors, %d CPU blocks (%.2f GB)", |
| 215 | len(unique_gpu_caches), |
| 216 | self.num_cpu_blocks, |
| 217 | (self.num_cpu_blocks * total_bytes_per_block) / (1024**3), |
| 218 | ) |
| 219 | |
| 220 | pin_memory = PIN_MEMORY |
| 221 | if not pin_memory: |
| 222 | logger.warning( |
| 223 | "Pinned memory not available. CPU offload performance may be degraded." |
| 224 | ) |
| 225 | |
| 226 | self.cpu_kv_caches = {} |
| 227 | for name, gpu_tensor in unique_gpu_caches.items(): |
| 228 | cpu_shape = (self.num_cpu_blocks,) + gpu_tensor.shape[1:] |
| 229 | # Allocate non-pinned first, then pin via cudaHostRegister to |
| 230 | # bypass PyTorch's CUDACachingHostAllocator which rounds up to |
| 231 | # the next power of 2 (e.g. 100 GB -> 128 GB). |
| 232 | tensor = torch.zeros(cpu_shape, dtype=gpu_tensor.dtype, device="cpu") |
| 233 | if pin_memory: |
| 234 | pin_tensor(tensor) |
| 235 | self.cpu_kv_caches[name] = tensor |
| 236 | |
| 237 | self._backend = DmaCopyBackend() |
| 238 | self._backend.init( |
| 239 | unique_gpu_caches, |
| 240 | self.cpu_kv_caches, |
| 241 | device, |
| 242 | self.load_stream, |
| 243 | self.store_stream, |
| 244 | ) |
| 245 | |
| 246 | def bind_connector_metadata(self, metadata: SimpleCPUOffloadMetadata) -> None: |
| 247 | self._connector_metadata = metadata |
| 248 | if metadata.load_event >= 0: |
| 249 | self._pending_load_event_indices.add(metadata.load_event) |
| 250 | if metadata.store_event >= 0: |
| 251 | self._pending_store_event_indices.add(metadata.store_event) |
| 252 | |
| 253 | def clear_connector_metadata(self) -> None: |
| 254 | self._connector_metadata = None |
| 255 | |
| 256 | def start_load_kv(self) -> None: |
| 257 | # Invoked after the forward launch on steps without sync loads |
| 258 | # (SchedulerOutput.has_sync_kv_loads), so the CPU-side block copy |
| 259 | # op overhead (~5ms) stays hidden behind GPU compute. Stores are |
| 260 | # issued in wait_for_save(). |
| 261 | metadata = self._connector_metadata |
| 262 | if metadata is not None and metadata.load_cpu_blocks: |
| 263 | backend = self._backend |
| 264 | assert backend is not None |
| 265 | backend.launch_copy( |
| 266 | metadata.load_cpu_blocks, |
| 267 | metadata.load_gpu_blocks, |
| 268 | is_store=False, |
| 269 | event_idx=metadata.load_event, |
| 270 | events_list=self._load_events, |
| 271 | ) |
| 272 | |
| 273 | def wait_for_save(self) -> None: |
| 274 | """Submit async stores. |
| 275 | |
| 276 | Stores (GPU->CPU) read the live KV cache, which the compute stream may |
| 277 | still be writing under v1 overlapped execution, so they are ordered |
| 278 | after a compute-done event recorded on the current stream. See |
| 279 | #45704 for the bug and #39306 for the srcAccessOrder rationale. |
| 280 | """ |
| 281 | metadata = self._connector_metadata |
| 282 | if metadata is not None and metadata.store_gpu_blocks: |
| 283 | backend = self._backend |
| 284 | assert backend is not None |
| 285 | if self._store_compute_done is None: |
| 286 | self._store_compute_done = torch.Event() |
| 287 | self._store_compute_done.record(torch.cuda.current_stream()) |
| 288 | backend.launch_copy( |
| 289 | metadata.store_gpu_blocks, |
| 290 | metadata.store_cpu_blocks, |
| 291 | is_store=True, |
| 292 | event_idx=metadata.store_event, |
| 293 | events_list=self._store_events, |
| 294 | wait_event=self._store_compute_done, |
| 295 | ) |
| 296 | |
| 297 | def get_finished( |
| 298 | self, finished_req_ids: set[str] |
| 299 | ) -> tuple[set[str] | None, set[str] | None]: |
| 300 | """Report completed transfer events to the scheduler. |
| 301 | |
| 302 | Loads are submitted in start_load_kv() and stores in wait_for_save(). |
| 303 | |
| 304 | Returns: |
| 305 | tuple of (finished_sending, finished_recving). |
| 306 | - finished_sending: always None (stores use worker metadata). |
| 307 | - finished_recving: req_ids whose loads have completed. |
| 308 | |
| 309 | """ |
| 310 | metadata = self._connector_metadata |
| 311 | finished_recving: set[str] = set() |
| 312 | |
| 313 | if self._pending_load_event_indices: |
| 314 | load_wm = self._poll_stream_events(is_store=False) |
| 315 | for j in [j for j in self._pending_load_event_indices if j <= load_wm]: |
| 316 | self._pending_load_event_indices.discard(j) |
| 317 | req_ids = ( |
| 318 | metadata.load_event_to_reqs.get(j) if metadata is not None else None |
| 319 | ) |
| 320 | if req_ids: |
| 321 | finished_recving.update(req_ids) |
| 322 | |
| 323 | if self._pending_store_event_indices: |
| 324 | store_wm = self._poll_stream_events(is_store=True) |
| 325 | for j in [j for j in self._pending_store_event_indices if j <= store_wm]: |
| 326 | self._pending_store_event_indices.discard(j) |
| 327 | self._completed_store_events[j] = 1 |
| 328 | |
| 329 | return None, finished_recving or None |
| 330 | |
| 331 | def build_connector_worker_meta(self) -> SimpleCPUOffloadWorkerMetadata | None: |
| 332 | """Return completed store events since the last call.""" |
| 333 | if not self._completed_store_events: |
| 334 | return None |
| 335 | meta = SimpleCPUOffloadWorkerMetadata( |
| 336 | completed_store_events=self._completed_store_events, |
| 337 | ) |
| 338 | self._completed_store_events = {} |
| 339 | return meta |
| 340 | |
| 341 | def handle_preemptions( |
| 342 | self, kv_connector_metadata: SimpleCPUOffloadMetadata |
| 343 | ) -> None: |
| 344 | """Sync all in-flight transfers before preempted blocks are reused.""" |
| 345 | if not kv_connector_metadata.need_flush: |
| 346 | return |
| 347 | self._flush_and_sync_all() |
| 348 | |
| 349 | def _flush_and_sync_all(self) -> None: |
| 350 | """Synchronize all in-flight transfer events.""" |
| 351 | for event_idx, event in self._load_events: |
| 352 | event.synchronize() |
| 353 | self._load_hwm = event_idx |
| 354 | self._load_events.clear() |
| 355 | |
| 356 | for event_idx, event in self._store_events: |
| 357 | event.synchronize() |
| 358 | self._store_hwm = event_idx |
| 359 | self._store_events.clear() |
| 360 | |
| 361 | def _poll_stream_events(self, is_store: bool) -> int: |
| 362 | """Non-blocking poll for completed events and return the high-water mark.""" |
| 363 | events = self._store_events if is_store else self._load_events |
| 364 | hwm = self._store_hwm if is_store else self._load_hwm |
| 365 | while events: |
| 366 | event_idx, event = events[0] |
| 367 | if not event.query(): |
| 368 | break |
| 369 | hwm = event_idx |
| 370 | events.pop(0) |
| 371 | if is_store: |
| 372 | self._store_hwm = hwm |
| 373 | else: |
| 374 | self._load_hwm = hwm |
| 375 | return hwm |