| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | import torch |
| 6 | from typing_extensions import override |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.utils.math_utils import round_up |
| 10 | from vllm.v1.kv_offload.base import ( |
| 11 | CanonicalKVCaches, |
| 12 | OffloadingCounterMetadata, |
| 13 | OffloadingGaugeMetadata, |
| 14 | OffloadingHistogramMetadata, |
| 15 | OffloadingManager, |
| 16 | OffloadingMetricMetadata, |
| 17 | OffloadingSpec, |
| 18 | OffloadingWorker, |
| 19 | ) |
| 20 | from vllm.v1.kv_offload.config import OffloadingConfig |
| 21 | from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics |
| 22 | from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker |
| 23 | from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager |
| 24 | from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion |
| 25 | |
| 26 | |
| 27 | def _all_workers_barrier() -> None: |
| 28 | """Block until every worker rank has reached this point (gloo cpu group). |
| 29 | |
| 30 | A superset of the node-local mmap openers suffices: once the barrier |
| 31 | releases, every worker sharing the region file has mapped it.""" |
| 32 | from vllm.distributed.parallel_state import ( |
| 33 | get_inner_dp_world_group, |
| 34 | get_world_group, |
| 35 | ) |
| 36 | |
| 37 | try: |
| 38 | group = get_inner_dp_world_group() |
| 39 | except AssertionError: |
| 40 | group = get_world_group() |
| 41 | group.barrier() |
| 42 | |
| 43 | |
| 44 | class CPUOffloadingSpec(OffloadingSpec): |
| 45 | BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT |
| 46 | |
| 47 | @classmethod |
| 48 | def build_metric_definitions( |
| 49 | cls, extra_config: dict[str, Any] |
| 50 | ) -> dict[str, OffloadingMetricMetadata]: |
| 51 | definitions: dict[str, OffloadingMetricMetadata] = { |
| 52 | CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC: OffloadingGaugeMetadata( |
| 53 | documentation=( |
| 54 | "Fraction of CPU KV-cache space currently pinned by active " |
| 55 | "transfers (0.0 = idle, 1.0 = saturated). Sustained high " |
| 56 | "values indicate transfers (stores or promotions) may be " |
| 57 | "dropped due to insufficient capacity." |
| 58 | ), |
| 59 | ), |
| 60 | CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC: OffloadingGaugeMetadata( |
| 61 | documentation=( |
| 62 | "Fraction of CPU KV-cache space currently pinned by " |
| 63 | "in-flight stores that have not yet " |
| 64 | "completed (0.0 = idle, 1.0 = saturated)." |
| 65 | ), |
| 66 | ), |
| 67 | CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC: OffloadingGaugeMetadata( |
| 68 | documentation=( |
| 69 | "Fraction of CPU KV-cache space currently pinned by " |
| 70 | "in-flight loads that have not yet " |
| 71 | "completed (0.0 = idle, 1.0 = saturated)." |
| 72 | ), |
| 73 | ), |
| 74 | CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata( |
| 75 | documentation=( |
| 76 | "Histogram of the number of CPU chunks requested by each " |
| 77 | "KV offload prepare_store call." |
| 78 | ), |
| 79 | buckets=(1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144), |
| 80 | ), |
| 81 | } |
| 82 | store_threshold = int(extra_config.get("store_threshold", 0)) |
| 83 | if store_threshold >= 2: |
| 84 | definitions[CPUOffloadingMetrics.STORES_SKIPPED] = ( |
| 85 | OffloadingCounterMetadata( |
| 86 | documentation=( |
| 87 | "Number of KV offload stores skipped because the reuse " |
| 88 | "threshold was not reached." |
| 89 | ), |
| 90 | ) |
| 91 | ) |
| 92 | return definitions |
| 93 | |
| 94 | def __init__(self, config: OffloadingConfig): |
| 95 | super().__init__(config) |
| 96 | |
| 97 | cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use") |
| 98 | if not cpu_bytes_to_use: |
| 99 | raise Exception( |
| 100 | "cpu_bytes_to_use must be specified in kv_connector_extra_config" |
| 101 | ) |
| 102 | |
| 103 | world_size = config.parallel.world_size |
| 104 | self.num_chunks = 0 |
| 105 | self.kv_bytes_per_chunk = 0 |
| 106 | self.cpu_page_size_per_worker = 0 |
| 107 | self.replicated_layout = config.replicated_layout and self._uses_shared_region() |
| 108 | if config.worker_kv_bytes_per_block > 0 and world_size > 0: |
| 109 | num_copies = 1 if self.replicated_layout else world_size |
| 110 | kv_bytes_per_block = config.worker_kv_bytes_per_block * num_copies |
| 111 | kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk |
| 112 | |
| 113 | # calculate cpu_page_size_per_worker |
| 114 | self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies |
| 115 | |
| 116 | # calculate num_chunks |
| 117 | aligned_kv_bytes_per_chunk = round_up( |
| 118 | kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT |
| 119 | ) |
| 120 | self.num_chunks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk |
| 121 | |
| 122 | # Expose aligned_kv_bytes_per_chunk as |
| 123 | # kv_bytes_per_chunk. Note that this might contain |
| 124 | # some padding. i.e. each offloaded chunk is of the form, |
| 125 | # |--- W0-C0---|---- W1-C0---| ... |---- Wn-C0---| *** maybe-pad *** | |
| 126 | # or |--- C0 (single copy) ---| *** maybe-pad *** | |
| 127 | self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk |
| 128 | |
| 129 | # scheduler-side |
| 130 | self._manager: OffloadingManager | None = None |
| 131 | |
| 132 | # worker-side |
| 133 | self._worker: CPUOffloadingWorker | None = None |
| 134 | |
| 135 | self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru") |
| 136 | self.cache_policy_module_path: str | None = self.extra_config.get( |
| 137 | "cache_policy_module_path" |
| 138 | ) |
| 139 | |
| 140 | @override |
| 141 | def get_manager(self) -> OffloadingManager: |
| 142 | if not self._manager: |
| 143 | # store_threshold: how many times a chunk must be offered for |
| 144 | # storage before it is eligible for CPU offloading. Values < 2 |
| 145 | # disable filtering (a threshold of 1 equals no filter; 0 is the |
| 146 | # default). |
| 147 | store_threshold = int(self.extra_config.get("store_threshold", 0)) |
| 148 | |
| 149 | # Maximum entries in the internal tracker's LRU table. |
| 150 | max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000)) |
| 151 | |
| 152 | self._manager = CPUOffloadingManager( |
| 153 | num_chunks=self.num_chunks, |
| 154 | cache_policy=self.eviction_policy, |
| 155 | cache_policy_module_path=self.cache_policy_module_path, |
| 156 | enable_events=self.kv_events_config.enable_kv_cache_events, |
| 157 | store_threshold=store_threshold, |
| 158 | max_tracker_size=max_tracker_size, |
| 159 | ) |
| 160 | return self._manager |
| 161 | |
| 162 | def _uses_shared_region(self) -> bool: |
| 163 | """Whether the worker CPU buffer is the shared mmap region (vs a private |
| 164 | per-rank tensor); replicated-layout dedup is gated on this being True.""" |
| 165 | return current_platform.is_cuda_alike() and not current_platform.is_rocm() |
| 166 | |
| 167 | def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: |
| 168 | mmap_region: SharedOffloadRegion | None = None |
| 169 | # num_chunks == 0 would size the region to zero bytes, which cannot be |
| 170 | # mmap'd; fall back to the tensor path (empty tensors) as before. |
| 171 | if self._uses_shared_region() and self.num_chunks > 0: |
| 172 | # Replicated layout puts all ranks on slot 0 (single MLA copy); |
| 173 | # otherwise each rank takes its own slot by physical device index. |
| 174 | if self.replicated_layout: |
| 175 | rank = 0 |
| 176 | else: |
| 177 | world_size = self.config.parallel.world_size |
| 178 | rank = torch.accelerator.current_device_index() % world_size |
| 179 | mmap_region = SharedOffloadRegion( |
| 180 | engine_id=self.config.engine_id, |
| 181 | num_chunks=self.num_chunks, |
| 182 | rank=rank, |
| 183 | kv_bytes_per_chunk=self.kv_bytes_per_chunk, |
| 184 | cpu_page_size=self.cpu_page_size_per_worker, |
| 185 | barrier=_all_workers_barrier, |
| 186 | ) |
| 187 | try: |
| 188 | return CPUOffloadingWorker( |
| 189 | kv_caches=kv_caches, |
| 190 | blocks_per_chunk=self.blocks_per_chunk, |
| 191 | num_cpu_chunks=self.num_chunks, |
| 192 | mmap_region=mmap_region, |
| 193 | ) |
| 194 | except Exception: |
| 195 | if mmap_region is not None: |
| 196 | mmap_region.cleanup() |
| 197 | raise |
| 198 | |
| 199 | @override |
| 200 | def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: |
| 201 | if not self._worker: |
| 202 | if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): |
| 203 | raise Exception( |
| 204 | "CPU Offloading is currently only supported on CUDA-alike " |
| 205 | "and XPU GPUs" |
| 206 | ) |
| 207 | self._worker = self.create_worker(kv_caches) |
| 208 | |
| 209 | assert self._worker is not None |
| 210 | return self._worker |