| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """TieringOffloadingSpec: Spec for multi-tier KV cache offloading. |
| 4 | |
| 5 | This spec creates a TieringOffloadingManager with a CPU primary tier |
| 6 | and configurable secondary tiers (e.g., Storage, Network). |
| 7 | |
| 8 | Configuration via kv_connector_extra_config: |
| 9 | - cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier |
| 10 | - block_size: (optional) Tokens per offloaded chunk (default: GPU block size) |
| 11 | - eviction_policy: (optional) Primary tier eviction policy: built-in "lru"/ |
| 12 | "arc", or the name of a policy registered via CachePolicyFactory, or an |
| 13 | out-of-tree CachePolicy class name paired with cache_policy_module_path |
| 14 | (default: "lru") |
| 15 | - cache_policy_module_path: (optional) Python import path to load |
| 16 | eviction_policy from when it names an out-of-tree CachePolicy not |
| 17 | registered via CachePolicyFactory |
| 18 | - secondary_tiers: (optional) List of secondary tier configurations |
| 19 | Each secondary tier config is a dict with: |
| 20 | - type: (required) Type of secondary tier (e.g., "example", "fs", |
| 21 | "p2p", "obj"), or the class name of an out-of-tree |
| 22 | SecondaryTierManager paired with module_path. |
| 23 | - module_path: (optional) Python import path to load 'type' from |
| 24 | when it names an out-of-tree SecondaryTierManager not registered |
| 25 | via SecondaryTierFactory.register_tier() |
| 26 | - Additional tier-specific parameters are passed directly to the tier |
| 27 | constructor. See each tier's documentation for supported parameters. |
| 28 | |
| 29 | Example configuration: |
| 30 | { |
| 31 | "cpu_bytes_to_use": 10737418240, # 10 GB |
| 32 | "block_size": 16, |
| 33 | "eviction_policy": "lru", |
| 34 | "secondary_tiers": [ |
| 35 | { |
| 36 | "type": "example", |
| 37 | "custom_param": 67 |
| 38 | } |
| 39 | ] |
| 40 | } |
| 41 | |
| 42 | Example out-of-tree tier configuration: |
| 43 | { |
| 44 | "cpu_bytes_to_use": 10737418240, |
| 45 | "secondary_tiers": [ |
| 46 | { |
| 47 | "type": "MyCustomTier", |
| 48 | "module_path": "my_package.my_module", |
| 49 | "custom_param": "value" |
| 50 | } |
| 51 | ] |
| 52 | } |
| 53 | """ |
| 54 | |
| 55 | from typing import Any |
| 56 | |
| 57 | import torch |
| 58 | from typing_extensions import override |
| 59 | |
| 60 | from vllm.logger import init_logger |
| 61 | from vllm.v1.kv_offload.base import ( |
| 62 | CanonicalKVCaches, |
| 63 | OffloadingCounterMetadata, |
| 64 | OffloadingGaugeMetadata, |
| 65 | OffloadingHistogramMetadata, |
| 66 | OffloadingManager, |
| 67 | OffloadingMetricMetadata, |
| 68 | ) |
| 69 | from vllm.v1.kv_offload.config import OffloadingConfig |
| 70 | from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker |
| 71 | from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion |
| 72 | from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec |
| 73 | from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics |
| 74 | from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory |
| 75 | from vllm.v1.kv_offload.tiering.manager import ( |
| 76 | CPUPrimaryTierOffloadingManager, |
| 77 | TieringOffloadingManager, |
| 78 | ) |
| 79 | |
| 80 | logger = init_logger(__name__) |
| 81 | |
| 82 | |
| 83 | class TieringOffloadingSpec(CPUOffloadingSpec): |
| 84 | """Spec for multi-tier KV cache offloading. |
| 85 | |
| 86 | Creates a TieringOffloadingManager with: |
| 87 | - Primary tier: CPU (LRU or ARC eviction policy) |
| 88 | - Secondary tiers: Configurable via extra_config |
| 89 | |
| 90 | The CPU primary tier has direct GPU access and serves as the gateway for |
| 91 | all GPU↔offload operations. Secondary tiers cannot directly access GPU |
| 92 | memory and must transfer data through the primary tier. |
| 93 | """ |
| 94 | |
| 95 | BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT |
| 96 | |
| 97 | @classmethod |
| 98 | @override |
| 99 | def build_metric_definitions( |
| 100 | cls, extra_config: dict[str, Any] |
| 101 | ) -> dict[str, OffloadingMetricMetadata]: |
| 102 | metrics = super().build_metric_definitions(extra_config) |
| 103 | metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY] = ( |
| 104 | OffloadingHistogramMetadata( |
| 105 | documentation=( |
| 106 | "Histogram of blocking time spent in a per-chunk tier lookup " |
| 107 | "that resolved as a hit or miss, labeled by tier, in seconds." |
| 108 | ), |
| 109 | labelnames=("tier",), |
| 110 | buckets=( |
| 111 | 0.00001, |
| 112 | 0.00005, |
| 113 | 0.0001, |
| 114 | 0.0005, |
| 115 | 0.001, |
| 116 | 0.005, |
| 117 | 0.01, |
| 118 | 0.05, |
| 119 | 0.1, |
| 120 | 0.5, |
| 121 | 1, |
| 122 | ), |
| 123 | ) |
| 124 | ) |
| 125 | metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY] = ( |
| 126 | OffloadingHistogramMetadata( |
| 127 | documentation=( |
| 128 | "Histogram of wall-clock time from a per-chunk tier lookup " |
| 129 | "first returning retry until that same tier lookup resolves " |
| 130 | "as a hit or miss, labeled by tier, in seconds." |
| 131 | ), |
| 132 | labelnames=("tier",), |
| 133 | buckets=( |
| 134 | 0.0001, |
| 135 | 0.0005, |
| 136 | 0.001, |
| 137 | 0.005, |
| 138 | 0.01, |
| 139 | 0.05, |
| 140 | 0.1, |
| 141 | 0.5, |
| 142 | 1, |
| 143 | 5, |
| 144 | 10, |
| 145 | ), |
| 146 | ) |
| 147 | ) |
| 148 | metrics[TieringOffloadingMetrics.READ_BYTES] = OffloadingCounterMetadata( |
| 149 | documentation=( |
| 150 | "Total bytes read from secondary tiers into the primary tier, " |
| 151 | "labeled by tier." |
| 152 | ), |
| 153 | labelnames=("tier",), |
| 154 | ) |
| 155 | metrics[TieringOffloadingMetrics.READ_TIME] = OffloadingCounterMetadata( |
| 156 | documentation=( |
| 157 | "Total time spent reading from secondary tiers into the primary " |
| 158 | "tier, in seconds, labeled by tier." |
| 159 | ), |
| 160 | labelnames=("tier",), |
| 161 | ) |
| 162 | metrics[TieringOffloadingMetrics.WRITE_BYTES] = OffloadingCounterMetadata( |
| 163 | documentation=( |
| 164 | "Total bytes written from the primary tier to secondary tiers, " |
| 165 | "labeled by tier." |
| 166 | ), |
| 167 | labelnames=("tier",), |
| 168 | ) |
| 169 | metrics[TieringOffloadingMetrics.WRITE_TIME] = OffloadingCounterMetadata( |
| 170 | documentation=( |
| 171 | "Total time spent writing from the primary tier to secondary " |
| 172 | "tiers, in seconds, labeled by tier." |
| 173 | ), |
| 174 | labelnames=("tier",), |
| 175 | ) |
| 176 | metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES] = ( |
| 177 | OffloadingCounterMetadata( |
| 178 | documentation=( |
| 179 | "Number of failed secondary-tier promotion jobs, labeled by tier." |
| 180 | ), |
| 181 | labelnames=("tier",), |
| 182 | ) |
| 183 | ) |
| 184 | metrics[TieringOffloadingMetrics.CASCADE_JOB_FAILURES] = ( |
| 185 | OffloadingCounterMetadata( |
| 186 | documentation=( |
| 187 | "Number of failed secondary-tier cascade jobs, labeled by tier." |
| 188 | ), |
| 189 | labelnames=("tier",), |
| 190 | ) |
| 191 | ) |
| 192 | metrics[TieringOffloadingMetrics.CHUNK_QUERIES] = OffloadingCounterMetadata( |
| 193 | documentation=( |
| 194 | "Number of chunk lookup queries sent to a tier, labeled by tier." |
| 195 | ), |
| 196 | labelnames=("tier",), |
| 197 | ) |
| 198 | metrics[TieringOffloadingMetrics.CHUNK_HITS] = OffloadingCounterMetadata( |
| 199 | documentation="Number of chunk lookup hits in a tier, labeled by tier.", |
| 200 | labelnames=("tier",), |
| 201 | ) |
| 202 | metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES] = ( |
| 203 | OffloadingCounterMetadata( |
| 204 | documentation=( |
| 205 | "Number of promotion attempts that failed because the " |
| 206 | "primary tier could not allocate space." |
| 207 | ), |
| 208 | ) |
| 209 | ) |
| 210 | metrics[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC] = ( |
| 211 | OffloadingGaugeMetadata( |
| 212 | documentation=( |
| 213 | "Current fraction of primary-tier space used by writes from " |
| 214 | "secondary tiers, labeled by tier." |
| 215 | ), |
| 216 | labelnames=("tier",), |
| 217 | ) |
| 218 | ) |
| 219 | metrics[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC] = ( |
| 220 | OffloadingGaugeMetadata( |
| 221 | documentation=( |
| 222 | "Current fraction of primary-tier space used by reads to " |
| 223 | "secondary tiers, labeled by tier." |
| 224 | ), |
| 225 | labelnames=("tier",), |
| 226 | ) |
| 227 | ) |
| 228 | metrics[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS] = ( |
| 229 | OffloadingGaugeMetadata( |
| 230 | documentation=( |
| 231 | "Number of active secondary-tier promotion jobs, labeled by tier." |
| 232 | ), |
| 233 | labelnames=("tier",), |
| 234 | ) |
| 235 | ) |
| 236 | metrics[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS] = OffloadingGaugeMetadata( |
| 237 | documentation=( |
| 238 | "Number of active secondary-tier cascade jobs, labeled by tier." |
| 239 | ), |
| 240 | labelnames=("tier",), |
| 241 | ) |
| 242 | secondary_tier_configs = extra_config.get("secondary_tiers", []) |
| 243 | if not isinstance(secondary_tier_configs, list): |
| 244 | raise ValueError("secondary_tiers must be a list of tier configurations") |
| 245 | |
| 246 | for tier_config in secondary_tier_configs: |
| 247 | assert isinstance(tier_config, dict) |
| 248 | tier_cls = SecondaryTierFactory.get_tier_class(tier_config) |
| 249 | metrics.update(tier_cls.build_metric_definitions(tier_config)) |
| 250 | |
| 251 | metrics[TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA] = ( |
| 252 | OffloadingGaugeMetadata( |
| 253 | documentation=( |
| 254 | "Exponential moving average of store latency " |
| 255 | "for back-pressure detection, in s/MiB." |
| 256 | ), |
| 257 | labelnames=("tier",), |
| 258 | ) |
| 259 | ) |
| 260 | metrics[TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED] = ( |
| 261 | OffloadingCounterMetadata( |
| 262 | documentation=( |
| 263 | "Number of store operations dropped due to " |
| 264 | "back-pressure on a secondary tier." |
| 265 | ), |
| 266 | labelnames=("tier",), |
| 267 | ) |
| 268 | ) |
| 269 | metrics[TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED] = ( |
| 270 | OffloadingCounterMetadata( |
| 271 | documentation=( |
| 272 | "Number of blocks dropped due to back-pressure on a secondary tier." |
| 273 | ), |
| 274 | labelnames=("tier",), |
| 275 | ) |
| 276 | ) |
| 277 | |
| 278 | return metrics |
| 279 | |
| 280 | def __init__(self, config: OffloadingConfig): |
| 281 | super().__init__(config) |
| 282 | # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it |
| 283 | self._manager: OffloadingManager | None = None |
| 284 | |
| 285 | # Parse secondary tier configurations |
| 286 | self.secondary_tier_configs = self.extra_config.get("secondary_tiers", []) |
| 287 | if not isinstance(self.secondary_tier_configs, list): |
| 288 | raise ValueError("secondary_tiers must be a list of tier configurations") |
| 289 | |
| 290 | # Backpressure config is merged field-by-field in priority order |
| 291 | # (highest first): |
| 292 | # 1. Per-tier ``backpressure`` dict in the tier config |
| 293 | # 2. Top-level ``backpressure`` in kv_connector_extra_config |
| 294 | # Merging per field (rather than per whole dict) means a partial |
| 295 | # tier override still inherits missing fields from the top-level |
| 296 | # default, so the resolved dict reaching the factory is complete. |
| 297 | # Within each tier's resolved dict, tier-type-aware water marks |
| 298 | # are filled in last so a bare ``"backpressure": {}`` picks up |
| 299 | # sensible thresholds for the storage medium. |
| 300 | bp_defaults = self.extra_config.get("backpressure") |
| 301 | |
| 302 | for tier_cfg in self.secondary_tier_configs: |
| 303 | tier_override = tier_cfg.get("backpressure") |
| 304 | # Overlay from lowest to highest precedence so higher-precedence |
| 305 | # fields win while lower-precedence ones fill in the gaps. |
| 306 | merged: dict[str, Any] = {} |
| 307 | for source in (bp_defaults, tier_override): |
| 308 | if source: |
| 309 | merged.update(source) |
| 310 | # Only set a resolved dict when at least one source contributed |
| 311 | # (or an explicit ``backpressure`` key was present, e.g. ``{}``), |
| 312 | # so tiers without any backpressure config stay unconfigured. |
| 313 | if merged or "backpressure" in tier_cfg: |
| 314 | tier_cfg["backpressure"] = merged |
| 315 | |
| 316 | # Scheduler-side mmap (rank=None); kept for cleanup |
| 317 | self._scheduler_mmap: SharedOffloadRegion | None = None |
| 318 | |
| 319 | # Set by create_worker when canonical_layout is enabled: True when |
| 320 | # every layer's canonical bytes are parallelism-agnostic (portable), |
| 321 | # False when some layers use the opaque fallback (exact-topology only) |
| 322 | self.all_layers_portable: bool | None = None |
| 323 | |
| 324 | # engine_id is unique per DP replica (suffixed with _dp{rank} in both |
| 325 | # the Ray and multiprocessing paths), so it names a per-replica offload |
| 326 | # region. |
| 327 | self._engine_id = config.engine_id |
| 328 | |
| 329 | @override |
| 330 | def get_manager(self) -> OffloadingManager: |
| 331 | """Get the TieringOffloadingManager. |
| 332 | |
| 333 | Creates a TieringOffloadingManager with: |
| 334 | - Primary tier: CPU (LRU or ARC) |
| 335 | - Secondary tiers: As configured in extra_config |
| 336 | |
| 337 | Returns: |
| 338 | TieringOffloadingManager instance |
| 339 | |
| 340 | """ |
| 341 | if not self._manager: |
| 342 | if int(self.extra_config.get("store_threshold", 0)) >= 2: |
| 343 | raise ValueError( |
| 344 | "store_threshold is not supported for TieringOffloadingSpec" |
| 345 | ) |
| 346 | |
| 347 | scheduler_mmap: SharedOffloadRegion | None = None |
| 348 | primary_tier: CPUPrimaryTierOffloadingManager | None = None |
| 349 | secondary_tiers = [] |
| 350 | try: |
| 351 | # Create scheduler-side SharedOffloadRegion (rank=None) so the |
| 352 | # primary tier can eagerly create a memoryview over _base. |
| 353 | scheduler_mmap = SharedOffloadRegion( |
| 354 | engine_id=self._engine_id, |
| 355 | num_chunks=self.num_chunks, |
| 356 | rank=None, |
| 357 | kv_bytes_per_chunk=self.kv_bytes_per_chunk, |
| 358 | cpu_page_size=self.cpu_page_size_per_worker, |
| 359 | ) |
| 360 | self._scheduler_mmap = scheduler_mmap |
| 361 | |
| 362 | # Create primary tier (CPU-based) |
| 363 | primary_tier = CPUPrimaryTierOffloadingManager( |
| 364 | num_chunks=self.num_chunks, |
| 365 | cache_policy=self.eviction_policy, |
| 366 | cache_policy_module_path=self.cache_policy_module_path, |
| 367 | enable_events=self.kv_events_config.enable_kv_cache_events, |
| 368 | mmap_region=scheduler_mmap, |
| 369 | ) |
| 370 | |
| 371 | # Create secondary tiers |
| 372 | primary_kv_view = primary_tier.get_kv_memoryview() |
| 373 | for i, tier_config in enumerate(self.secondary_tier_configs): |
| 374 | tier = SecondaryTierFactory.create_secondary_tier( |
| 375 | tier_config, primary_kv_view, self |
| 376 | ) |
| 377 | secondary_tiers.append(tier) |
| 378 | logger.info( |
| 379 | "Created secondary tier #%d (%s)", |
| 380 | i, |
| 381 | tier.tier_type, |
| 382 | ) |
| 383 | |
| 384 | # Create TieringOffloadingManager. GPU↔CPU transfers use the inherited |
| 385 | # get_worker(). Secondary tier transfers are handled by the |
| 386 | # secondary tier managers and need no additional workers here. |
| 387 | tiering_manager = TieringOffloadingManager( |
| 388 | primary_tier=primary_tier, |
| 389 | secondary_tiers=secondary_tiers, |
| 390 | ) |
| 391 | self._manager = tiering_manager |
| 392 | except Exception: |
| 393 | for tier in reversed(secondary_tiers): |
| 394 | try: |
| 395 | tier.shutdown() |
| 396 | except Exception: |
| 397 | logger.exception( |
| 398 | "Failed to shut down secondary tier during " |
| 399 | "initialization cleanup" |
| 400 | ) |
| 401 | if primary_tier is not None: |
| 402 | try: |
| 403 | primary_tier.shutdown() |
| 404 | except Exception: |
| 405 | logger.exception( |
| 406 | "Failed to shut down primary tier during " |
| 407 | "initialization cleanup" |
| 408 | ) |
| 409 | elif scheduler_mmap is not None: |
| 410 | try: |
| 411 | scheduler_mmap.cleanup() |
| 412 | except Exception: |
| 413 | logger.exception( |
| 414 | "Failed to clean up scheduler mmap during " |
| 415 | "initialization cleanup" |
| 416 | ) |
| 417 | self._scheduler_mmap = None |
| 418 | raise |
| 419 | |
| 420 | logger.info( |
| 421 | "Created TieringOffloadingManager with primary tier " |
| 422 | "(%s, %s chunks) and %s secondary tier(s)", |
| 423 | self.eviction_policy, |
| 424 | self.num_chunks, |
| 425 | len(secondary_tiers), |
| 426 | ) |
| 427 | |
| 428 | return self._manager |
| 429 | |
| 430 | @override |
| 431 | def _uses_shared_region(self) -> bool: |
| 432 | # Tiering always allocates on the shared region (every platform), so the |
| 433 | # replicated-layout gate must not be narrowed by the CPU spec's |
| 434 | # CUDA-alike check. |
| 435 | return True |
| 436 | |
| 437 | @override |
| 438 | def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: |
| 439 | world_size = self.config.parallel.world_size |
| 440 | if self.replicated_layout: |
| 441 | rank = 0 |
| 442 | else: |
| 443 | # Fold the global physical device index into the replica-local |
| 444 | # [0, world_size) slot range. |
| 445 | rank = torch.accelerator.current_device_index() % world_size |
| 446 | worker_mmap = SharedOffloadRegion( |
| 447 | engine_id=self._engine_id, |
| 448 | num_chunks=self.num_chunks, |
| 449 | rank=rank, |
| 450 | kv_bytes_per_chunk=self.kv_bytes_per_chunk, |
| 451 | cpu_page_size=self.cpu_page_size_per_worker, |
| 452 | ) |
| 453 | try: |
| 454 | if self.config.canonical_layout: |
| 455 | self._validate_canonical_refs(kv_caches) |
| 456 | return CPUOffloadingWorker( |
| 457 | kv_caches=kv_caches, |
| 458 | blocks_per_chunk=self.blocks_per_chunk, |
| 459 | num_cpu_chunks=self.num_chunks, |
| 460 | mmap_region=worker_mmap, |
| 461 | canonical_layout=self.config.canonical_layout, |
| 462 | ) |
| 463 | except Exception: |
| 464 | worker_mmap.cleanup() |
| 465 | raise |
| 466 | |
| 467 | def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None: |
| 468 | """Require a mapping on every ref and record layer portability. |
| 469 | |
| 470 | Fails loudly rather than persist direct-layout bytes under a |
| 471 | canonical format identity.""" |
| 472 | all_refs = [ |
| 473 | ref for group_refs in kv_caches.group_data_refs for ref in group_refs |
| 474 | ] |
| 475 | if any(ref.mapping is None for ref in all_refs): |
| 476 | raise RuntimeError( |
| 477 | "canonical_layout was requested but the KV cache layout " |
| 478 | "could not be certified for canonical offload (offload " |
| 479 | "workers must be exactly the TP group, and packed / " |
| 480 | "cross-layer KV layouts are not supported). Remove " |
| 481 | "canonical_layout from kv_connector_extra_config." |
| 482 | ) |
| 483 | self.all_layers_portable = all( |
| 484 | ref.mapping is not None and ref.mapping.parallelism_agnostic |
| 485 | for ref in all_refs |
| 486 | ) |
| 487 | if self.config.parallel.is_parallelism_agnostic and not ( |
| 488 | self.all_layers_portable |
| 489 | ): |
| 490 | # The scheduler-side storage namespace was already collapsed on |
| 491 | # the static portability claim; opaque per-topology bytes must |
| 492 | # not land in it. |
| 493 | raise RuntimeError( |
| 494 | "canonical_layout could not certify every layer as " |
| 495 | "parallelism-agnostic, but the storage namespace is shared " |
| 496 | "across topologies. Remove canonical_layout from " |
| 497 | "kv_connector_extra_config or disable parallel-agnostic " |
| 498 | "secondary tiers." |
| 499 | ) |
| 500 | logger.info( |
| 501 | "Canonical KV layout enabled (all_layers_portable=%s)", |
| 502 | self.all_layers_portable, |
| 503 | ) |