| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import functools |
| 4 | import time |
| 5 | from collections import deque |
| 6 | from collections.abc import Sequence |
| 7 | from dataclasses import dataclass |
| 8 | from typing import NamedTuple |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | |
| 13 | from vllm import _custom_ops as ops |
| 14 | from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.triton_utils import HAS_TRITON, triton |
| 18 | from vllm.utils.math_utils import cdiv |
| 19 | from vllm.utils.torch_utils import PIN_MEMORY |
| 20 | from vllm.v1.kv_offload.base import ( |
| 21 | BlockIDsLoadStoreSpec, |
| 22 | CanonicalKVCacheRef, |
| 23 | CanonicalKVCaches, |
| 24 | CanonicalPageMapping, |
| 25 | GPULoadStoreSpec, |
| 26 | LoadStoreSpec, |
| 27 | OffloadingWorker, |
| 28 | TransferResult, |
| 29 | ) |
| 30 | from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion |
| 31 | from vllm.v1.kv_offload.cpu.swap_blocks_triton import ( |
| 32 | THRESHOLD_BYTES, |
| 33 | swap_blocks_batch, |
| 34 | ) |
| 35 | |
| 36 | logger = init_logger(__name__) |
| 37 | |
| 38 | |
| 39 | def _select_swap_blocks_fn( |
| 40 | layer_refs_per_group: list[list[CanonicalKVCacheRef]], |
| 41 | gpu_to_cpu: bool, |
| 42 | host_memory_is_pinned: bool = True, |
| 43 | ): |
| 44 | """Resolve the swap_blocks function for a handler at init time.""" |
| 45 | # GPU->CPU is bandwidth-bound; the dedicated copy engine beats Triton. |
| 46 | # The Triton kernel dereferences CPU pointers on the GPU, which is only |
| 47 | # valid for pinned host memory. |
| 48 | if gpu_to_cpu or not host_memory_is_pinned: |
| 49 | return ops.swap_blocks_batch |
| 50 | # Fall back to the C++ DMA path on platforms where Triton isn't usable |
| 51 | # (e.g. ROCm host mappings) or where GPU kernels cannot directly |
| 52 | # dereference CPU pointers (XPU lacks CUDA's unified virtual address space, |
| 53 | # so the Triton kernel's tl.load(cpu_ptr) is invalid on XPU). |
| 54 | if not HAS_TRITON or current_platform.is_xpu() or current_platform.is_rocm(): |
| 55 | return ops.swap_blocks_batch |
| 56 | page_sizes = [r.page_size_bytes for g in layer_refs_per_group for r in g] |
| 57 | # Triton wins only on small, 8-byte-aligned payloads. |
| 58 | if ( |
| 59 | not page_sizes |
| 60 | or max(page_sizes) >= THRESHOLD_BYTES |
| 61 | or any(s % 8 for s in page_sizes) |
| 62 | ): |
| 63 | return ops.swap_blocks_batch |
| 64 | chunk = min(triton.next_power_of_2(max(page_sizes)), 8192) |
| 65 | return functools.partial(swap_blocks_batch, bytes_per_chunk=chunk) |
| 66 | |
| 67 | |
| 68 | @dataclass |
| 69 | class Transfer: |
| 70 | job_id: int |
| 71 | stream: torch.cuda.Stream |
| 72 | start_event: torch.Event |
| 73 | end_event: torch.Event |
| 74 | num_bytes: int |
| 75 | batch_src: torch.Tensor |
| 76 | batch_dst: torch.Tensor |
| 77 | batch_sizes: torch.Tensor |
| 78 | |
| 79 | |
| 80 | def compute_sub_block_ptrs( |
| 81 | block_ids: np.ndarray, |
| 82 | blocks_per_chunk: int, |
| 83 | output: np.ndarray, |
| 84 | tensor: torch.Tensor, |
| 85 | skip_count: int = 0, |
| 86 | ): |
| 87 | """Compute byte pointers for sub-blocks of the given block IDs. |
| 88 | |
| 89 | Each block in block_ids contains blocks_per_chunk sub-blocks. |
| 90 | The pointer for sub-block j of block b is: |
| 91 | base_ptr + b * row_stride + j * block_page_size |
| 92 | |
| 93 | where block_page_size = tensor.shape[1] // blocks_per_chunk (gpu page size). |
| 94 | |
| 95 | This handles tensors where row_stride != blocks_per_chunk * block_page_size |
| 96 | (e.g. non-contiguous CPU tensors). |
| 97 | |
| 98 | Args: |
| 99 | block_ids: array of block IDs at the tensor's native granularity. |
| 100 | blocks_per_chunk: number of sub-blocks per block. |
| 101 | output: pre-allocated pointer array to write pointers into. |
| 102 | tensor: the source or destination tensor. |
| 103 | skip_count: sub-blocks to skip in the first block. |
| 104 | |
| 105 | """ |
| 106 | assert skip_count < blocks_per_chunk |
| 107 | |
| 108 | num_sub_blocks = len(output) |
| 109 | base_ptr = tensor.data_ptr() |
| 110 | row_stride = tensor.stride(0) |
| 111 | |
| 112 | if blocks_per_chunk == 1: |
| 113 | # Fast path: 1:1 mapping, no sub-block expansion needed. |
| 114 | output[:] = base_ptr + block_ids.astype(np.uint64)[:num_sub_blocks] * row_stride |
| 115 | return |
| 116 | |
| 117 | # Vectorized expansion for blocks_per_chunk > 1. |
| 118 | assert tensor.shape[1] % blocks_per_chunk == 0 |
| 119 | block_page_size = tensor.shape[1] // blocks_per_chunk |
| 120 | sub_offsets = np.arange(blocks_per_chunk, dtype=np.uint64) * block_page_size |
| 121 | # (num_blocks, 1) + (1, blocks_per_chunk) -> (num_blocks, blocks_per_chunk) |
| 122 | all_ptrs = ( |
| 123 | base_ptr + block_ids.astype(np.uint64)[:, np.newaxis] * row_stride |
| 124 | ) + sub_offsets[np.newaxis, :] |
| 125 | # Flatten and apply skip_count / truncation |
| 126 | flat = all_ptrs.ravel() |
| 127 | output[:] = flat[skip_count : skip_count + num_sub_blocks] |
| 128 | |
| 129 | |
| 130 | class CopyPlan(NamedTuple): |
| 131 | """Precomputed fragment-copy template for one data ref under the canonical |
| 132 | CPU layout, unrolled from the ref's mapped runs. Offsets are relative to |
| 133 | the per-block base pointers on each side.""" |
| 134 | |
| 135 | frag_offsets_src: np.ndarray |
| 136 | frag_offsets_dst: np.ndarray |
| 137 | frag_sizes: np.ndarray |
| 138 | total_bytes: int |
| 139 | |
| 140 | @property |
| 141 | def num_frags(self) -> int: |
| 142 | return len(self.frag_sizes) |
| 143 | |
| 144 | |
| 145 | def _build_copy_plan(ref: CanonicalKVCacheRef, gpu_to_cpu: bool) -> CopyPlan: |
| 146 | """Unroll one data ref's mapped runs into a per-fragment CopyPlan.""" |
| 147 | mapping = ref.mapping |
| 148 | assert mapping is not None |
| 149 | local: list[int] = [] |
| 150 | canonical: list[int] = [] |
| 151 | sizes: list[int] = [] |
| 152 | for run in mapping.runs: |
| 153 | for i in range(run.num_fragments): |
| 154 | local.append(run.local_offset + i * run.local_stride) |
| 155 | canonical.append(run.canonical_offset + i * run.canonical_stride) |
| 156 | sizes.append(run.fragment_size) |
| 157 | src, dst = (local, canonical) if gpu_to_cpu else (canonical, local) |
| 158 | return CopyPlan( |
| 159 | frag_offsets_src=np.asarray(src, dtype=np.uint64), |
| 160 | frag_offsets_dst=np.asarray(dst, dtype=np.uint64), |
| 161 | frag_sizes=np.asarray(sizes, dtype=np.int64), |
| 162 | total_bytes=sum(sizes), |
| 163 | ) |
| 164 | |
| 165 | |
| 166 | def _canonical_page_ids( |
| 167 | block_ids: np.ndarray, blocks_per_chunk: int, count: int, skip_count: int |
| 168 | ) -> np.ndarray: |
| 169 | """Global canonical page ids matching compute_sub_block_ptrs' enumeration. |
| 170 | These identify canonical pages consistently across ranks, so they key |
| 171 | CanonicalPageMapping.is_writer rotation.""" |
| 172 | if blocks_per_chunk == 1: |
| 173 | return block_ids[:count] |
| 174 | flat = ( |
| 175 | block_ids[:, np.newaxis] * blocks_per_chunk + np.arange(blocks_per_chunk) |
| 176 | ).ravel() |
| 177 | return flat[skip_count : skip_count + count] |
| 178 | |
| 179 | |
| 180 | def _canonical_block_sizes( |
| 181 | layer_refs_per_group: list[list[CanonicalKVCacheRef]], num_tensors: int |
| 182 | ) -> list[int]: |
| 183 | """Canonical CPU bytes per GPU block for each tensor, taken from the refs' |
| 184 | mappings. Requires every ref to carry a mapping.""" |
| 185 | canonical_bytes_per_block = [0] * num_tensors |
| 186 | for layer_refs in layer_refs_per_group: |
| 187 | for ref in layer_refs: |
| 188 | assert ref.mapping is not None |
| 189 | canonical_bytes_per_block[ref.tensor_idx] = max( |
| 190 | canonical_bytes_per_block[ref.tensor_idx], |
| 191 | ref.mapping.canonical_page_size_bytes, |
| 192 | ) |
| 193 | assert all(size > 0 for size in canonical_bytes_per_block) |
| 194 | return canonical_bytes_per_block |
| 195 | |
| 196 | |
| 197 | # Bound registration size to avoid driver limits on large host allocations. |
| 198 | MAX_HOST_REGISTER_CHUNK_BYTES = 64 * 1024**3 |
| 199 | |
| 200 | |
| 201 | def pin_mmap_region(region: SharedOffloadRegion) -> None: |
| 202 | """Register row-aligned chunks, rolling back on failure.""" |
| 203 | if not current_platform.is_cuda_alike(): |
| 204 | logger.info( |
| 205 | "Skipping mmap host registration on %s; cudaHostRegister is only " |
| 206 | "available on CUDA/ROCm.", |
| 207 | current_platform.device_name, |
| 208 | ) |
| 209 | return |
| 210 | |
| 211 | rank = region.rank |
| 212 | try: |
| 213 | cudart = CudaRTLibrary() |
| 214 | except (AssertionError, AttributeError, OSError): |
| 215 | logger.warning( |
| 216 | "Could not load the CUDA runtime for host registration on rank=%d; " |
| 217 | "the offload region stays pageable", |
| 218 | rank, |
| 219 | exc_info=True, |
| 220 | ) |
| 221 | return |
| 222 | |
| 223 | base_ptr = region._base.data_ptr() |
| 224 | total_size = region.total_size_bytes |
| 225 | # Chunks end on block-row boundaries, which are page aligned, so neither the |
| 226 | # driver's page rounding nor any single block transfer straddles two |
| 227 | # registrations. |
| 228 | rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1) |
| 229 | chunk_size = rows_per_chunk * region._row_stride |
| 230 | |
| 231 | # Register, drain and roll back through the same runtime handle, so a |
| 232 | # failed chunk leaves neither a pending error nor a partly pinned region. |
| 233 | addresses: list[int] = [] |
| 234 | for offset in range(0, total_size, chunk_size): |
| 235 | address = base_ptr + offset |
| 236 | size = min(chunk_size, total_size - offset) |
| 237 | result = cudart.cudaHostRegister(address, size) |
| 238 | if result == 0: |
| 239 | addresses.append(address) |
| 240 | continue |
| 241 | cudart.cudaGetLastError() |
| 242 | logger.warning( |
| 243 | "cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); " |
| 244 | "the offload region stays pageable", |
| 245 | rank, |
| 246 | offset / 1e9, |
| 247 | total_size / 1e9, |
| 248 | result, |
| 249 | ) |
| 250 | for registered in reversed(addresses): |
| 251 | unregister_result = cudart.cudaHostUnregister(registered) |
| 252 | if unregister_result != 0: |
| 253 | cudart.cudaGetLastError() |
| 254 | logger.warning( |
| 255 | "cudaHostUnregister failed for rank=%d at %#x (code=%d); " |
| 256 | "that chunk stays registered until the process exits", |
| 257 | rank, |
| 258 | registered, |
| 259 | unregister_result, |
| 260 | ) |
| 261 | return |
| 262 | |
| 263 | region.pinned_addresses.extend(addresses) |
| 264 | region.is_pinned = True |
| 265 | logger.debug( |
| 266 | "cudaHostRegister rank=%d %.2f GB in %d chunk(s)", |
| 267 | rank, |
| 268 | total_size / 1e9, |
| 269 | len(addresses), |
| 270 | ) |
| 271 | |
| 272 | |
| 273 | def _new_descriptor_buffers( |
| 274 | num_copy_ops: int, |
| 275 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 276 | pin = PIN_MEMORY |
| 277 | # CUDA cache_kernels.cu requires int64; XPU DMA engine requires uint64. |
| 278 | ptr_dtype = torch.uint64 if current_platform.is_xpu() else torch.int64 |
| 279 | return ( |
| 280 | torch.empty(num_copy_ops, dtype=ptr_dtype, pin_memory=pin), |
| 281 | torch.empty(num_copy_ops, dtype=ptr_dtype, pin_memory=pin), |
| 282 | torch.empty(num_copy_ops, dtype=ptr_dtype, pin_memory=pin), |
| 283 | ) |
| 284 | |
| 285 | |
| 286 | class SingleDirectionOffloadingHandler: |
| 287 | """Handles transfers for a single direction, either CPU->GPU or GPU->CPU. |
| 288 | Transfers are guaranteed to be executed in order of their submission. |
| 289 | Each transfer uses a unique CUDA stream, and its stream will start |
| 290 | executing only after the streams of previous transfers have finished. |
| 291 | """ |
| 292 | |
| 293 | def __init__( |
| 294 | self, |
| 295 | gpu_tensors: list[torch.Tensor], |
| 296 | cpu_tensors: list[torch.Tensor], |
| 297 | blocks_per_chunk: int, |
| 298 | layer_refs_per_group: list[list[CanonicalKVCacheRef]], |
| 299 | gpu_to_cpu: bool, |
| 300 | canonical_layout: bool = False, |
| 301 | host_memory_is_pinned: bool = True, |
| 302 | ): |
| 303 | """Initialize a SingleDirectionOffloadingHandler. |
| 304 | |
| 305 | Args: |
| 306 | gpu_tensors: list of GPU KV cache tensors. |
| 307 | Each of shape (num_gpu_blocks, gpu_page_size_bytes) with dtype int8. |
| 308 | cpu_tensors: list of CPU KV cache tensors. |
| 309 | Each of shape (num_cpu_chunks, cpu_page_size_bytes) with dtype int8. |
| 310 | Order should match gpu_tensors. |
| 311 | blocks_per_chunk: number of blocks transferred per chunk. |
| 312 | layer_refs_per_group: list of CanonicalKVCacheRef per group. |
| 313 | gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU. |
| 314 | canonical_layout: if True, CPU pages use the canonical layout |
| 315 | described by the refs' mappings. |
| 316 | host_memory_is_pinned: whether the CPU tensors are pinned, so GPU |
| 317 | kernels may dereference them directly. |
| 318 | |
| 319 | """ |
| 320 | assert len(gpu_tensors) == len(cpu_tensors) |
| 321 | assert len(gpu_tensors) > 0 |
| 322 | |
| 323 | canonical_bytes_per_block = ( |
| 324 | _canonical_block_sizes(layer_refs_per_group, len(gpu_tensors)) |
| 325 | if canonical_layout |
| 326 | else None |
| 327 | ) |
| 328 | |
| 329 | # assert input tensors are as expected |
| 330 | for t_idx, (gpu_tensor, cpu_tensor) in enumerate(zip(gpu_tensors, cpu_tensors)): |
| 331 | assert gpu_tensor.dtype == torch.int8 |
| 332 | assert gpu_tensor.ndim == 2 |
| 333 | assert gpu_tensor.is_cuda or gpu_tensor.is_xpu |
| 334 | assert cpu_tensor.dtype == torch.int8 |
| 335 | assert cpu_tensor.ndim == 2 |
| 336 | assert cpu_tensor.device.type == "cpu" |
| 337 | _, gpu_page_size = gpu_tensor.shape |
| 338 | _, cpu_page_size = cpu_tensor.shape |
| 339 | if canonical_bytes_per_block is not None: |
| 340 | assert ( |
| 341 | cpu_page_size == canonical_bytes_per_block[t_idx] * blocks_per_chunk |
| 342 | ) |
| 343 | else: |
| 344 | assert cpu_page_size == gpu_page_size * blocks_per_chunk |
| 345 | |
| 346 | self.src_tensors: list[torch.Tensor] = ( |
| 347 | gpu_tensors if gpu_to_cpu else cpu_tensors |
| 348 | ) |
| 349 | self.dst_tensors: list[torch.Tensor] = ( |
| 350 | cpu_tensors if gpu_to_cpu else gpu_tensors |
| 351 | ) |
| 352 | self.gpu_to_cpu: bool = gpu_to_cpu |
| 353 | self.layer_refs_per_group = layer_refs_per_group |
| 354 | self._swap_blocks_batch = _select_swap_blocks_fn( |
| 355 | layer_refs_per_group, gpu_to_cpu, host_memory_is_pinned |
| 356 | ) |
| 357 | |
| 358 | # GPU blocks may be smaller |
| 359 | # cpu_page_size = gpu_page_size * blocks_per_chunk. |
| 360 | self.src_blocks_per_chunk = 1 if self.gpu_to_cpu else blocks_per_chunk |
| 361 | self.dst_blocks_per_chunk = blocks_per_chunk if self.gpu_to_cpu else 1 |
| 362 | |
| 363 | # Per (group, ref) static copy plans for the canonical layout |
| 364 | self._canonical_copy_plans: list[list[CopyPlan]] | None = ( |
| 365 | [ |
| 366 | [_build_copy_plan(ref, gpu_to_cpu) for ref in layer_refs] |
| 367 | for layer_refs in layer_refs_per_group |
| 368 | ] |
| 369 | if canonical_layout |
| 370 | else None |
| 371 | ) |
| 372 | self._fill_group_ops = ( |
| 373 | self._fill_canonical_ops if canonical_layout else self._fill_direct_ops |
| 374 | ) |
| 375 | # Reusable per-block base-pointer scratch for the canonical fill, |
| 376 | # sized to the largest possible group (grown on demand) |
| 377 | num_scratch_blocks = gpu_tensors[0].shape[0] if canonical_layout else 0 |
| 378 | self._scratch_bases_src = np.empty(num_scratch_blocks, dtype=np.uint64) |
| 379 | self._scratch_bases_dst = np.empty(num_scratch_blocks, dtype=np.uint64) |
| 380 | |
| 381 | # job_id -> event |
| 382 | self._transfer_events: dict[int, torch.Event] = {} |
| 383 | # queue of transfers (job_id, stream, event) |
| 384 | self._transfers: deque[Transfer] = deque() |
| 385 | # list of CUDA streams available for re-use |
| 386 | self._stream_pool: list[torch.cuda.Stream] = [] |
| 387 | # list of CUDA events available for re-use |
| 388 | self._event_pool: list[torch.Event] = [] |
| 389 | # list of pinned descriptor buffer sets available for re-use |
| 390 | self._buffer_pool: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] |
| 391 | |
| 392 | def _estimate_max_copy_ops(self, group_sizes: Sequence[int]) -> int: |
| 393 | """Upper bound on the number of copy descriptors for a transfer. |
| 394 | |
| 395 | Exact for the direct layout. The canonical path may fill fewer: |
| 396 | writer rotation later drops the blocks this rank does not write.""" |
| 397 | num_copy_ops = 0 |
| 398 | for g_idx, (group_size, layer_refs) in enumerate( |
| 399 | zip(group_sizes, self.layer_refs_per_group) |
| 400 | ): |
| 401 | if self._canonical_copy_plans is None: |
| 402 | num_copy_ops += group_size * len(layer_refs) |
| 403 | else: |
| 404 | num_copy_ops += group_size * sum( |
| 405 | plan.num_frags for plan in self._canonical_copy_plans[g_idx] |
| 406 | ) |
| 407 | return num_copy_ops |
| 408 | |
| 409 | def _fill_direct_ops( |
| 410 | self, |
| 411 | g_idx: int, |
| 412 | group_src: np.ndarray, |
| 413 | group_dst: np.ndarray, |
| 414 | group_size: int, |
| 415 | src_skip_count: int, |
| 416 | dst_skip_count: int, |
| 417 | all_src: np.ndarray, |
| 418 | all_dst: np.ndarray, |
| 419 | all_sizes: np.ndarray, |
| 420 | op_idx: int, |
| 421 | ) -> tuple[int, int]: |
| 422 | """Fill one group's copy descriptors for the direct (worker-private) |
| 423 | layout: one whole-page copy per (block, ref). |
| 424 | |
| 425 | Returns (op_idx past the filled descriptors, bytes added).""" |
| 426 | num_bytes = 0 |
| 427 | for data_ref in self.layer_refs_per_group[g_idx]: |
| 428 | t_idx = data_ref.tensor_idx |
| 429 | end_idx = op_idx + group_size |
| 430 | |
| 431 | compute_sub_block_ptrs( |
| 432 | group_src, |
| 433 | self.src_blocks_per_chunk, |
| 434 | all_src[op_idx:end_idx], |
| 435 | self.src_tensors[t_idx], |
| 436 | skip_count=src_skip_count, |
| 437 | ) |
| 438 | compute_sub_block_ptrs( |
| 439 | group_dst, |
| 440 | self.dst_blocks_per_chunk, |
| 441 | all_dst[op_idx:end_idx], |
| 442 | self.dst_tensors[t_idx], |
| 443 | skip_count=dst_skip_count, |
| 444 | ) |
| 445 | |
| 446 | all_sizes[op_idx:end_idx] = data_ref.page_size_bytes |
| 447 | num_bytes += group_size * data_ref.page_size_bytes |
| 448 | op_idx = end_idx |
| 449 | return op_idx, num_bytes |
| 450 | |
| 451 | def _fill_canonical_ops( |
| 452 | self, |
| 453 | g_idx: int, |
| 454 | group_src: np.ndarray, |
| 455 | group_dst: np.ndarray, |
| 456 | group_size: int, |
| 457 | src_skip_count: int, |
| 458 | dst_skip_count: int, |
| 459 | all_src: np.ndarray, |
| 460 | all_dst: np.ndarray, |
| 461 | all_sizes: np.ndarray, |
| 462 | op_idx: int, |
| 463 | ) -> tuple[int, int]: |
| 464 | """Fill one group's copy descriptors for the canonical layout: |
| 465 | scatter each block through the ref's precomputed CopyPlan, keeping |
| 466 | only the blocks this rank writes. |
| 467 | |
| 468 | Returns (op_idx past the filled descriptors, bytes added).""" |
| 469 | assert self._canonical_copy_plans is not None |
| 470 | # Zero-copy reinterpretation for pointer arithmetic: uint64 and the |
| 471 | # buffers' int64 are bit-equivalent for addresses |
| 472 | all_src_u64 = all_src.view(np.uint64) |
| 473 | all_dst_u64 = all_dst.view(np.uint64) |
| 474 | if group_size > len(self._scratch_bases_src): |
| 475 | self._scratch_bases_src = np.empty(group_size, dtype=np.uint64) |
| 476 | self._scratch_bases_dst = np.empty(group_size, dtype=np.uint64) |
| 477 | |
| 478 | num_bytes = 0 |
| 479 | for plan, data_ref in zip( |
| 480 | self._canonical_copy_plans[g_idx], self.layer_refs_per_group[g_idx] |
| 481 | ): |
| 482 | if plan.num_frags == 0: |
| 483 | continue |
| 484 | t_idx = data_ref.tensor_idx |
| 485 | |
| 486 | # 1. Base byte pointer of every block on each side |
| 487 | block_bases_src = self._scratch_bases_src[:group_size] |
| 488 | block_bases_dst = self._scratch_bases_dst[:group_size] |
| 489 | compute_sub_block_ptrs( |
| 490 | group_src, |
| 491 | self.src_blocks_per_chunk, |
| 492 | block_bases_src, |
| 493 | self.src_tensors[t_idx], |
| 494 | skip_count=src_skip_count, |
| 495 | ) |
| 496 | compute_sub_block_ptrs( |
| 497 | group_dst, |
| 498 | self.dst_blocks_per_chunk, |
| 499 | block_bases_dst, |
| 500 | self.dst_tensors[t_idx], |
| 501 | skip_count=dst_skip_count, |
| 502 | ) |
| 503 | |
| 504 | # 2. On store, keep only the blocks this rank is elected to write |
| 505 | mapping = data_ref.mapping |
| 506 | assert mapping is not None |
| 507 | if self.gpu_to_cpu and mapping.num_writers > 1: |
| 508 | block_bases_src, block_bases_dst = self._filter_writer_blocks( |
| 509 | block_bases_src, |
| 510 | block_bases_dst, |
| 511 | mapping, |
| 512 | group_dst, |
| 513 | group_size, |
| 514 | dst_skip_count, |
| 515 | ) |
| 516 | num_active_blocks = len(block_bases_src) |
| 517 | |
| 518 | # 3. Expand (block base + fragment offset) into one descriptor |
| 519 | # per (block, fragment), writing straight into the descriptor |
| 520 | # buffers: reshaping a contiguous 1D slice is a view, so the |
| 521 | # broadcasts below allocate nothing |
| 522 | end_idx = op_idx + num_active_blocks * plan.num_frags |
| 523 | np.add( |
| 524 | block_bases_src[:, None], |
| 525 | plan.frag_offsets_src[None, :], |
| 526 | out=all_src_u64[op_idx:end_idx].reshape( |
| 527 | num_active_blocks, plan.num_frags |
| 528 | ), |
| 529 | ) |
| 530 | np.add( |
| 531 | block_bases_dst[:, None], |
| 532 | plan.frag_offsets_dst[None, :], |
| 533 | out=all_dst_u64[op_idx:end_idx].reshape( |
| 534 | num_active_blocks, plan.num_frags |
| 535 | ), |
| 536 | ) |
| 537 | all_sizes[op_idx:end_idx].reshape(num_active_blocks, plan.num_frags)[:] = ( |
| 538 | plan.frag_sizes |
| 539 | ) |
| 540 | num_bytes += num_active_blocks * plan.total_bytes |
| 541 | op_idx = end_idx |
| 542 | return op_idx, num_bytes |
| 543 | |
| 544 | def _filter_writer_blocks( |
| 545 | self, |
| 546 | block_bases_src: np.ndarray, |
| 547 | block_bases_dst: np.ndarray, |
| 548 | mapping: CanonicalPageMapping, |
| 549 | group_dst: np.ndarray, |
| 550 | group_size: int, |
| 551 | dst_skip_count: int, |
| 552 | ) -> tuple[np.ndarray, np.ndarray]: |
| 553 | """Keep only the blocks this rank writes: replicated ranks take turns |
| 554 | writing shared canonical pages, keyed by the rank-consistent CPU-side |
| 555 | canonical page id.""" |
| 556 | cpu_page_ids = _canonical_page_ids( |
| 557 | group_dst, |
| 558 | self.dst_blocks_per_chunk, |
| 559 | group_size, |
| 560 | dst_skip_count, |
| 561 | ) |
| 562 | writer_mask = cpu_page_ids % mapping.num_writers == mapping.writer_index |
| 563 | return block_bases_src[writer_mask], block_bases_dst[writer_mask] |
| 564 | |
| 565 | def transfer_async( |
| 566 | self, job_id: int, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec |
| 567 | ) -> bool: |
| 568 | assert isinstance(src_spec, BlockIDsLoadStoreSpec) |
| 569 | assert isinstance(dst_spec, BlockIDsLoadStoreSpec) |
| 570 | |
| 571 | src_blocks = src_spec.block_ids |
| 572 | dst_blocks = dst_spec.block_ids |
| 573 | assert src_blocks.ndim == 1 |
| 574 | assert dst_blocks.ndim == 1 |
| 575 | |
| 576 | num_src_blocks = len(src_blocks) |
| 577 | num_dst_blocks = len(dst_blocks) |
| 578 | |
| 579 | # There are 2 types of transfers: |
| 580 | # 1. GPU -> CPU |
| 581 | # 2. CPU -> GPU |
| 582 | # |
| 583 | # transfers are also to CPU chunks, EXCEPT MAYBE for the first and last chunk. |
| 584 | # i.e. the first and last CPU chunks in src_blocks can match against |
| 585 | # a smaller (byte-wise) set of GPU blocks in dst_blocks. |
| 586 | # In such cases, we may need to skip some gpu-sized sub-blocks, |
| 587 | # and start reading/writing from the middle of the first CPU chunk. |
| 588 | # If we have multiple KV cache groups (when using HMA with hybrid models), |
| 589 | # we may have a partial first/last CPU chunk per each group. |
| 590 | # The group_sizes parameter encodes the size of each group of blocks |
| 591 | # in the GPU dst_blocks. |
| 592 | # If group_sizes is None, we assume all blocks belong to a single group. |
| 593 | # The logical_offset parameter maps each group of blocks to its logical |
| 594 | # offset inside the request, counting in GPU blocks. |
| 595 | # This allows us to find the correct starting position |
| 596 | # in the matching first CPU chunk. |
| 597 | |
| 598 | # extract group_sizes from the GPU spec |
| 599 | gpu_spec = src_spec if self.gpu_to_cpu else dst_spec |
| 600 | assert isinstance(gpu_spec, GPULoadStoreSpec) |
| 601 | group_sizes = gpu_spec.group_sizes |
| 602 | assert len(group_sizes) == len(self.layer_refs_per_group) |
| 603 | |
| 604 | # extract block indices from the GPU spec |
| 605 | block_indices = gpu_spec.block_indices |
| 606 | assert len(block_indices) == len(self.layer_refs_per_group) |
| 607 | |
| 608 | num_copy_ops = self._estimate_max_copy_ops(group_sizes) |
| 609 | |
| 610 | # reuse a pooled buffer set, growing it if this transfer needs more room |
| 611 | batch_src, batch_dst, batch_sizes = ( |
| 612 | self._buffer_pool.pop() |
| 613 | if self._buffer_pool |
| 614 | else _new_descriptor_buffers(num_copy_ops) |
| 615 | ) |
| 616 | if batch_src.numel() < num_copy_ops: |
| 617 | batch_src, batch_dst, batch_sizes = _new_descriptor_buffers(num_copy_ops) |
| 618 | |
| 619 | src = batch_src[:num_copy_ops] |
| 620 | dst = batch_dst[:num_copy_ops] |
| 621 | sizes = batch_sizes[:num_copy_ops] |
| 622 | all_src = src.numpy() |
| 623 | all_dst = dst.numpy() |
| 624 | all_sizes = sizes.numpy() |
| 625 | |
| 626 | src_offset = 0 |
| 627 | dst_offset = 0 |
| 628 | op_idx = 0 |
| 629 | # count total number of bytes copied |
| 630 | num_transfer_bytes = 0 |
| 631 | for g_idx, (group_size, block_idx) in enumerate( |
| 632 | zip(group_sizes, block_indices) |
| 633 | ): |
| 634 | if group_size == 0: |
| 635 | continue |
| 636 | |
| 637 | src_logical_blocks_to_skip = block_idx % self.src_blocks_per_chunk |
| 638 | dst_logical_blocks_to_skip = block_idx % self.dst_blocks_per_chunk |
| 639 | src_logical_blocks_count = group_size + src_logical_blocks_to_skip |
| 640 | dst_logical_blocks_count = group_size + dst_logical_blocks_to_skip |
| 641 | |
| 642 | dst_blocks_count = cdiv(dst_logical_blocks_count, self.dst_blocks_per_chunk) |
| 643 | dst_end_offset = dst_offset + dst_blocks_count |
| 644 | assert dst_end_offset <= num_dst_blocks |
| 645 | |
| 646 | src_blocks_count = cdiv(src_logical_blocks_count, self.src_blocks_per_chunk) |
| 647 | src_end_offset = src_offset + src_blocks_count |
| 648 | assert src_end_offset <= num_src_blocks |
| 649 | |
| 650 | op_idx, group_bytes = self._fill_group_ops( |
| 651 | g_idx, |
| 652 | group_src=src_blocks[src_offset:src_end_offset], |
| 653 | group_dst=dst_blocks[dst_offset:dst_end_offset], |
| 654 | group_size=group_size, |
| 655 | src_skip_count=src_logical_blocks_to_skip, |
| 656 | dst_skip_count=dst_logical_blocks_to_skip, |
| 657 | all_src=all_src, |
| 658 | all_dst=all_dst, |
| 659 | all_sizes=all_sizes, |
| 660 | op_idx=op_idx, |
| 661 | ) |
| 662 | num_transfer_bytes += group_bytes |
| 663 | |
| 664 | src_offset = src_end_offset |
| 665 | dst_offset = dst_end_offset |
| 666 | |
| 667 | assert src_offset == num_src_blocks |
| 668 | assert dst_offset == num_dst_blocks |
| 669 | # Writer rotation may skip non-writer blocks, leaving op_idx below |
| 670 | # the sized upper bound |
| 671 | assert op_idx <= num_copy_ops |
| 672 | src = src[:op_idx] |
| 673 | dst = dst[:op_idx] |
| 674 | sizes = sizes[:op_idx] |
| 675 | |
| 676 | stream = ( |
| 677 | self._stream_pool.pop() if self._stream_pool else current_platform.Stream() |
| 678 | ) |
| 679 | start_event = ( |
| 680 | self._event_pool.pop() |
| 681 | if self._event_pool |
| 682 | else torch.Event(enable_timing=True) |
| 683 | ) |
| 684 | end_event = ( |
| 685 | self._event_pool.pop() |
| 686 | if self._event_pool |
| 687 | else torch.Event(enable_timing=True) |
| 688 | ) |
| 689 | |
| 690 | # Stores must wait for the model to finish writing the KV they read. |
| 691 | # Loads must wait for pending writes (including zeroing) to their |
| 692 | # destination blocks: the scheduler's _skip_zero_block_ids only edits |
| 693 | # the step being scheduled and cannot retract zeroing shipped in an |
| 694 | # earlier step for a since-reallocated block; with async scheduling |
| 695 | # nothing else orders that zeroing against this copy. |
| 696 | stream.wait_stream(current_platform.current_stream()) |
| 697 | if self._transfers: |
| 698 | last_transfer: Transfer = self._transfers[-1] |
| 699 | last_event = last_transfer.end_event |
| 700 | # assure job will start only after the previous one completes |
| 701 | stream.wait_event(last_event) |
| 702 | # CPU->GPU reads from host memory, which is never written |
| 703 | # by a concurrent GPU stream, so CU_MEMCPY_SRC_ACCESS_ORDER_ANY is |
| 704 | # safe and lets the driver pipeline source reads. GPU->CPU reads |
| 705 | # from the live GPU KV cache, which the compute stream keeps |
| 706 | # writing; we must keep STREAM ordering so source reads are gated |
| 707 | # by the transfer stream's wait_stream(compute) barrier. |
| 708 | is_src_access_order_any = not self.gpu_to_cpu |
| 709 | with current_platform.stream(stream): |
| 710 | start_event.record(stream) |
| 711 | if op_idx > 0: |
| 712 | self._swap_blocks_batch( |
| 713 | src, |
| 714 | dst, |
| 715 | sizes, |
| 716 | is_src_access_order_any=is_src_access_order_any, |
| 717 | ) |
| 718 | end_event.record(stream) |
| 719 | |
| 720 | self._transfer_events[job_id] = end_event |
| 721 | self._transfers.append( |
| 722 | Transfer( |
| 723 | job_id=job_id, |
| 724 | stream=stream, |
| 725 | start_event=start_event, |
| 726 | end_event=end_event, |
| 727 | num_bytes=num_transfer_bytes, |
| 728 | batch_src=batch_src, |
| 729 | batch_dst=batch_dst, |
| 730 | batch_sizes=batch_sizes, |
| 731 | ) |
| 732 | ) |
| 733 | |
| 734 | # success |
| 735 | return True |
| 736 | |
| 737 | def get_finished(self) -> list[TransferResult]: |
| 738 | results: list[TransferResult] = [] |
| 739 | while self._transfers and self._transfers[0].end_event.query(): |
| 740 | transfer = self._transfers.popleft() |
| 741 | transfer_time = ( |
| 742 | transfer.start_event.elapsed_time(transfer.end_event) * 1e-3 |
| 743 | ) # elapsed_time is in milliseconds |
| 744 | result = TransferResult( |
| 745 | job_id=transfer.job_id, |
| 746 | success=True, |
| 747 | transfer_size=transfer.num_bytes, |
| 748 | transfer_time=transfer_time, |
| 749 | ) |
| 750 | |
| 751 | results.append(result) |
| 752 | self._stream_pool.append(transfer.stream) |
| 753 | self._event_pool.append(transfer.end_event) |
| 754 | self._event_pool.append(transfer.start_event) |
| 755 | self._buffer_pool.append( |
| 756 | (transfer.batch_src, transfer.batch_dst, transfer.batch_sizes) |
| 757 | ) |
| 758 | del self._transfer_events[transfer.job_id] |
| 759 | return results |
| 760 | |
| 761 | def wait(self, job_ids: set[int]): |
| 762 | for job_id in job_ids: |
| 763 | event = self._transfer_events.get(job_id) |
| 764 | if event is not None: |
| 765 | event.synchronize() |
| 766 | |
| 767 | def shutdown(self) -> None: |
| 768 | """Drain this direction and release its transfer-side resources.""" |
| 769 | sync_error: Exception | None = None |
| 770 | while self._transfers: |
| 771 | transfer = self._transfers[0] |
| 772 | try: |
| 773 | transfer.end_event.synchronize() |
| 774 | except Exception as e: |
| 775 | logger.exception( |
| 776 | "Failed to synchronize transfer end event; " |
| 777 | "skipping %d remaining transfers", |
| 778 | len(self._transfers) - 1, |
| 779 | ) |
| 780 | self._transfers.clear() |
| 781 | sync_error = e |
| 782 | break |
| 783 | self._transfers.popleft() |
| 784 | |
| 785 | self._transfer_events.clear() |
| 786 | self._stream_pool.clear() |
| 787 | self._event_pool.clear() |
| 788 | self._buffer_pool.clear() |
| 789 | self.src_tensors.clear() |
| 790 | self.dst_tensors.clear() |
| 791 | if sync_error is not None: |
| 792 | raise sync_error |
| 793 | |
| 794 | |
| 795 | class CPUOffloadingWorker(OffloadingWorker): |
| 796 | """OffloadingWorker for CPU offloading. |
| 797 | |
| 798 | Composes two SingleDirectionOffloadingHandler instances (one for each |
| 799 | direction) and exposes them through the explicit submit_store / |
| 800 | submit_load API. |
| 801 | """ |
| 802 | |
| 803 | def __init__( |
| 804 | self, |
| 805 | kv_caches: CanonicalKVCaches, |
| 806 | blocks_per_chunk: int, |
| 807 | num_cpu_chunks: int, |
| 808 | mmap_region: SharedOffloadRegion | None = None, |
| 809 | canonical_layout: bool = False, |
| 810 | ): |
| 811 | assert not canonical_layout or mmap_region is not None |
| 812 | # The caller owns mmap_region until this constructor returns. After a |
| 813 | # successful construction, the worker is the sole owner and releases |
| 814 | # it after both transfer directions have stopped. |
| 815 | self._mmap_region = mmap_region |
| 816 | pin_memory = PIN_MEMORY |
| 817 | logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors)) |
| 818 | if mmap_region is not None and pin_memory: |
| 819 | pin_mmap_region(mmap_region) |
| 820 | host_memory_is_pinned = pin_memory and ( |
| 821 | mmap_region is None or mmap_region.is_pinned |
| 822 | ) |
| 823 | |
| 824 | canonical_bytes_per_block = ( |
| 825 | _canonical_block_sizes(kv_caches.group_data_refs, len(kv_caches.tensors)) |
| 826 | if canonical_layout |
| 827 | else None |
| 828 | ) |
| 829 | |
| 830 | gpu_tensors: list[torch.Tensor] = [] |
| 831 | cpu_tensors: list[torch.Tensor] = [] |
| 832 | for t_idx, kv_cache_tensor in enumerate(kv_caches.tensors): |
| 833 | gpu_page_size_bytes = kv_cache_tensor.page_size_bytes |
| 834 | gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view( |
| 835 | (-1, gpu_page_size_bytes) |
| 836 | ) |
| 837 | cpu_page_size_bytes = gpu_page_size_bytes * blocks_per_chunk |
| 838 | |
| 839 | if canonical_bytes_per_block is not None: |
| 840 | assert mmap_region is not None |
| 841 | cpu_tensor = mmap_region.create_next_canonical_view( |
| 842 | canonical_bytes_per_block[t_idx] * blocks_per_chunk |
| 843 | ) |
| 844 | elif mmap_region is not None: |
| 845 | cpu_tensor = mmap_region.create_next_worker_view(cpu_page_size_bytes) |
| 846 | else: |
| 847 | t0 = time.monotonic() |
| 848 | cpu_tensor = torch.zeros( |
| 849 | (num_cpu_chunks, cpu_page_size_bytes), |
| 850 | dtype=torch.int8, |
| 851 | device="cpu", |
| 852 | pin_memory=pin_memory, |
| 853 | ) |
| 854 | logger.debug( |
| 855 | "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s", |
| 856 | num_cpu_chunks, |
| 857 | cpu_page_size_bytes, |
| 858 | num_cpu_chunks * cpu_page_size_bytes / 1e9, |
| 859 | time.monotonic() - t0, |
| 860 | ) |
| 861 | |
| 862 | gpu_tensors.append(gpu_tensor) |
| 863 | cpu_tensors.append(cpu_tensor) |
| 864 | |
| 865 | self._store_handler = SingleDirectionOffloadingHandler( |
| 866 | gpu_tensors=gpu_tensors, |
| 867 | cpu_tensors=cpu_tensors, |
| 868 | blocks_per_chunk=blocks_per_chunk, |
| 869 | layer_refs_per_group=kv_caches.group_data_refs, |
| 870 | gpu_to_cpu=True, |
| 871 | canonical_layout=canonical_layout, |
| 872 | host_memory_is_pinned=host_memory_is_pinned, |
| 873 | ) |
| 874 | |
| 875 | self._load_handler = SingleDirectionOffloadingHandler( |
| 876 | gpu_tensors=gpu_tensors, |
| 877 | cpu_tensors=cpu_tensors, |
| 878 | blocks_per_chunk=blocks_per_chunk, |
| 879 | layer_refs_per_group=kv_caches.group_data_refs, |
| 880 | gpu_to_cpu=False, |
| 881 | canonical_layout=canonical_layout, |
| 882 | host_memory_is_pinned=host_memory_is_pinned, |
| 883 | ) |
| 884 | |
| 885 | def submit_store( |
| 886 | self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec |
| 887 | ) -> bool: |
| 888 | """Async GPU -> CPU.""" |
| 889 | return self._store_handler.transfer_async(job_id, src_spec, dst_spec) |
| 890 | |
| 891 | def submit_load( |
| 892 | self, job_id: int, src_spec: LoadStoreSpec, dst_spec: GPULoadStoreSpec |
| 893 | ) -> bool: |
| 894 | """Async CPU -> GPU.""" |
| 895 | return self._load_handler.transfer_async(job_id, src_spec, dst_spec) |
| 896 | |
| 897 | def get_finished(self) -> list[TransferResult]: |
| 898 | return self._store_handler.get_finished() + self._load_handler.get_finished() |
| 899 | |
| 900 | def wait(self, job_ids: set[int]) -> None: |
| 901 | self._store_handler.wait(job_ids) |
| 902 | self._load_handler.wait(job_ids) |
| 903 | |
| 904 | def shutdown(self) -> None: |
| 905 | handler_failed = False |
| 906 | try: |
| 907 | self._store_handler.shutdown() |
| 908 | except Exception: |
| 909 | logger.exception("Failed to shut down store offloading handler") |
| 910 | handler_failed = True |
| 911 | |
| 912 | try: |
| 913 | self._load_handler.shutdown() |
| 914 | except Exception: |
| 915 | logger.exception("Failed to shut down load offloading handler") |
| 916 | handler_failed = True |
| 917 | |
| 918 | if self._mmap_region is not None: |
| 919 | if handler_failed: |
| 920 | try: |
| 921 | torch.accelerator.synchronize() |
| 922 | except Exception: |
| 923 | logger.warning( |
| 924 | "Device sync before mmap cleanup failed; " |
| 925 | "proceeding with cleanup anyway", |
| 926 | exc_info=True, |
| 927 | ) |
| 928 | self._mmap_region.cleanup() |
| 929 | self._mmap_region = None |