| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Abstract interfaces and data types for the secondary tiering layer.""" |
| 4 | |
| 5 | import time |
| 6 | from abc import ABC, abstractmethod |
| 7 | from collections.abc import Collection, Iterable |
| 8 | from dataclasses import dataclass, field |
| 9 | from typing import TYPE_CHECKING, Any, ClassVar |
| 10 | |
| 11 | import numpy as np |
| 12 | |
| 13 | from vllm.v1.kv_offload.base import ( |
| 14 | Locality, |
| 15 | LookupResult, |
| 16 | Medium, |
| 17 | OffloadingEvent, |
| 18 | OffloadingMetricMetadata, |
| 19 | OffloadKey, |
| 20 | ReqContext, |
| 21 | RequestOffloadingContext, |
| 22 | ScheduleEndContext, |
| 23 | ) |
| 24 | from vllm.v1.kv_offload.tiering.backpressure import ( |
| 25 | BackpressureDetector, |
| 26 | ) |
| 27 | |
| 28 | if TYPE_CHECKING: |
| 29 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( |
| 30 | OffloadingConnectorStats, |
| 31 | ) |
| 32 | from vllm.v1.kv_offload.base import OffloadingSpec |
| 33 | |
| 34 | |
| 35 | # Type alias for job IDs used in async transfer tracking |
| 36 | JobId = int |
| 37 | |
| 38 | |
| 39 | class TieringOffloadingMetrics: |
| 40 | """Metric names for TieringOffloadingManager.""" |
| 41 | |
| 42 | LOOKUP_SYNC_DELAY = "vllm:kv_offload_tiering_lookup_sync_delay_seconds" |
| 43 | LOOKUP_ASYNC_DELAY = "vllm:kv_offload_tiering_lookup_async_delay_seconds" |
| 44 | READ_BYTES = "vllm:kv_offload_tiering_read_bytes" |
| 45 | READ_TIME = "vllm:kv_offload_tiering_read_time" |
| 46 | WRITE_BYTES = "vllm:kv_offload_tiering_write_bytes" |
| 47 | WRITE_TIME = "vllm:kv_offload_tiering_write_time" |
| 48 | PROMOTION_JOB_FAILURES = "vllm:kv_offload_tiering_promotion_job_failures" |
| 49 | CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures" |
| 50 | CHUNK_QUERIES = "vllm:kv_offload_tiering_chunk_queries" |
| 51 | CHUNK_HITS = "vllm:kv_offload_tiering_chunk_hits" |
| 52 | PRIMARY_WRITE_USAGE_PERC = "vllm:kv_offload_tiering_primary_write_usage_perc" |
| 53 | PRIMARY_READ_USAGE_PERC = "vllm:kv_offload_tiering_primary_read_usage_perc" |
| 54 | PROMOTION_ALLOCATION_FAILURES = ( |
| 55 | "vllm:kv_offload_tiering_promotion_allocation_failures" |
| 56 | ) |
| 57 | ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs" |
| 58 | ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs" |
| 59 | BACKPRESSURE_STORE_LATENCY_EMA = ( |
| 60 | "vllm:kv_offload_tiering_backpressure_store_latency_ema" |
| 61 | ) |
| 62 | BACKPRESSURE_STORES_DROPPED = "vllm:kv_offload_tiering_backpressure_stores_dropped" |
| 63 | BACKPRESSURE_BLOCKS_DROPPED = "vllm:kv_offload_tiering_backpressure_blocks_dropped" |
| 64 | |
| 65 | |
| 66 | @dataclass |
| 67 | class TransferJob: |
| 68 | """Metadata for an in-flight async transfer job.""" |
| 69 | |
| 70 | job_id: JobId |
| 71 | keys: Collection[OffloadKey] |
| 72 | chunk_ids: np.ndarray |
| 73 | is_promotion: bool |
| 74 | req_context: ReqContext |
| 75 | submit_time: float = field(default_factory=time.monotonic) |
| 76 | |
| 77 | |
| 78 | @dataclass |
| 79 | class JobResult: |
| 80 | """Result of an async transfer job.""" |
| 81 | |
| 82 | job_id: JobId |
| 83 | # True if all keys succeeded; False if all or some failed. |
| 84 | success: bool |
| 85 | # Only applicable to promotion jobs. On partial failure, identifies the |
| 86 | # keys that were successfully loaded. None means all keys share the fate |
| 87 | # indicated by `success`. Must be a subset of the job's original keys. |
| 88 | successful_keys: Collection[OffloadKey] | None = None |
| 89 | transfer_time: float | None = None |
| 90 | transfer_bytes: int | None = None |
| 91 | |
| 92 | |
| 93 | class ParentManager(ABC): |
| 94 | """Interface for secondary tiers to call back into the tiering manager. |
| 95 | |
| 96 | Passed to secondary tiers via serve_external_requests() each step. |
| 97 | The _SecondaryTierFacingParent wrapper implements this, automatically |
| 98 | excluding the calling tier from fan-out operations. |
| 99 | |
| 100 | Required call sequence for each remote request: |
| 101 | 1. on_new_request(req_context) — set up per-request state |
| 102 | 2. lookup(key, req_context) — check chunk availability |
| 103 | (repeat per chunk) |
| 104 | 3. create_store_job(keys, req_context) — pin chunks and get a |
| 105 | job handle |
| 106 | 4. on_request_finished(req_context) — clean up per-request state |
| 107 | |
| 108 | Steps 2-3 may be interleaved. Step 4 must be called even if no |
| 109 | chunks were found, to avoid leaking async lookup state (e.g. in |
| 110 | the fs tier's AsyncLookupManager). |
| 111 | """ |
| 112 | |
| 113 | @abstractmethod |
| 114 | def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: ... |
| 115 | |
| 116 | @abstractmethod |
| 117 | def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: ... |
| 118 | |
| 119 | @abstractmethod |
| 120 | def create_store_job( |
| 121 | self, |
| 122 | keys: Collection[OffloadKey], |
| 123 | req_context: ReqContext, |
| 124 | ) -> TransferJob: ... |
| 125 | |
| 126 | @abstractmethod |
| 127 | def on_request_finished(self, req_context: ReqContext) -> None: ... |
| 128 | |
| 129 | |
| 130 | class SecondaryTierManager(ABC): |
| 131 | """Abstract interface for managing a single non-primary offloading tier. |
| 132 | |
| 133 | Secondary tiers cannot directly access GPU memory. All data transfers |
| 134 | must go through the CPU (primary) tier: |
| 135 | - Store: GPU → CPU (primary) → secondary (cascade) |
| 136 | - Load: secondary → CPU (primary) → GPU (promotion) |
| 137 | |
| 138 | IMPORTANT: All methods run in the Scheduler process and must be |
| 139 | lightweight and non-blocking. submit_load() and submit_store() submit |
| 140 | async jobs; get_finished_jobs() polls for completion. |
| 141 | """ |
| 142 | |
| 143 | medium: ClassVar[Medium | None] = None |
| 144 | |
| 145 | def __init__( |
| 146 | self, |
| 147 | offloading_spec: "OffloadingSpec", |
| 148 | primary_kv_view: memoryview, |
| 149 | tier_type: str, |
| 150 | backpressure_detector: BackpressureDetector | None = None, |
| 151 | ) -> None: |
| 152 | """Args: |
| 153 | offloading_spec: Offloading configuration. |
| 154 | primary_kv_view: Memoryview of the primary tier's CPU KV cache. |
| 155 | tier_type: Tier type identifier, set by SecondaryTierFactory |
| 156 | from the registered tier type. |
| 157 | backpressure_detector: Optional `BackpressureDetector`. |
| 158 | |
| 159 | """ |
| 160 | self._offloading_spec = offloading_spec |
| 161 | self._primary_kv_view: memoryview = primary_kv_view |
| 162 | assert primary_kv_view.strides is not None, ( |
| 163 | "primary_kv_view.strides cannot be None" |
| 164 | ) |
| 165 | self._block_size_bytes: int = primary_kv_view.strides[0] |
| 166 | self.tier_type = tier_type |
| 167 | self.locality: Locality | None = None |
| 168 | self._bp_detector = backpressure_detector |
| 169 | |
| 170 | @property |
| 171 | def block_size_bytes(self) -> int: |
| 172 | return self._block_size_bytes |
| 173 | |
| 174 | @property |
| 175 | def bp_detector(self) -> BackpressureDetector | None: |
| 176 | return self._bp_detector |
| 177 | |
| 178 | @abstractmethod |
| 179 | def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: |
| 180 | """Check whether a chunk exists in this secondary tier. |
| 181 | |
| 182 | Args: |
| 183 | key: Offload key to look up. |
| 184 | req_context: per-request context (e.g. kv_transfer_params). |
| 185 | |
| 186 | Returns: |
| 187 | HIT if the chunk is present and ready, |
| 188 | MISS if not found, |
| 189 | or RETRY if the chunk is being transferred (retry later). |
| 190 | |
| 191 | """ |
| 192 | pass |
| 193 | |
| 194 | @abstractmethod |
| 195 | def submit_store(self, job_metadata: TransferJob) -> None: |
| 196 | """Submit an async job to store chunks from the primary tier to this |
| 197 | secondary tier. |
| 198 | |
| 199 | This method must be lightweight and non-blocking: allocate metadata |
| 200 | and submit the transfer, but do NOT perform the data copy on the |
| 201 | calling thread. |
| 202 | |
| 203 | Preconditions (guaranteed by the framework): |
| 204 | - ``job_metadata.chunk_ids`` are valid primary-tier slots, pinned |
| 205 | (ref-counted) for the duration of the transfer. |
| 206 | |
| 207 | The implementation is responsible for: |
| 208 | 1. Filtering out chunks already present in this tier |
| 209 | 2. Evicting chunks if capacity is needed |
| 210 | 3. Allocating space in this tier |
| 211 | 4. Submitting the async transfer (read from primary via chunk_ids) |
| 212 | |
| 213 | Report completion via ``get_finished_jobs()``. |
| 214 | |
| 215 | Args: |
| 216 | job_metadata: Job metadata including job_id, keys, and chunk_ids |
| 217 | identifying the primary-tier slots to read from. |
| 218 | |
| 219 | """ |
| 220 | pass |
| 221 | |
| 222 | @abstractmethod |
| 223 | def submit_load(self, job_metadata: TransferJob) -> None: |
| 224 | """Submit an async job to load chunks from this secondary tier to the |
| 225 | primary tier. |
| 226 | |
| 227 | This method must be lightweight and non-blocking: mark chunks as |
| 228 | in-flight and submit the transfer, but do NOT perform the data copy |
| 229 | on the calling thread. |
| 230 | |
| 231 | Preconditions (guaranteed by the framework): |
| 232 | - ``job_metadata.chunk_ids`` are allocated primary-tier slots |
| 233 | ready to receive data. |
| 234 | |
| 235 | The implementation must copy data from this tier into the |
| 236 | primary-tier slots identified by ``chunk_ids``. |
| 237 | |
| 238 | Report completion via ``get_finished_jobs()``. |
| 239 | |
| 240 | Args: |
| 241 | job_metadata: Job metadata including job_id, keys, and chunk_ids |
| 242 | identifying the primary-tier slots to write into. |
| 243 | |
| 244 | """ |
| 245 | pass |
| 246 | |
| 247 | @abstractmethod |
| 248 | def get_finished_jobs(self) -> Iterable[JobResult]: |
| 249 | """Return all jobs (loads and stores) that completed since the last call. |
| 250 | |
| 251 | The framework uses these results to release resources and finalize |
| 252 | transfers. |
| 253 | |
| 254 | Returns: |
| 255 | Iterable of JobResult objects for jobs finished since the |
| 256 | last call. |
| 257 | |
| 258 | """ |
| 259 | pass |
| 260 | |
| 261 | def has_pending_work(self) -> bool: |
| 262 | """Whether this tier needs the engine to keep stepping. |
| 263 | |
| 264 | While True, on_schedule_end() and get_finished_jobs() continue |
| 265 | to be called even when no requests are scheduled. |
| 266 | """ |
| 267 | return False |
| 268 | |
| 269 | def take_events(self) -> Iterable[OffloadingEvent]: |
| 270 | """Take KV events for storage state owned by this tier.""" |
| 271 | return () |
| 272 | |
| 273 | def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): |
| 274 | """Mark chunks as recently used for eviction policy. |
| 275 | |
| 276 | Args: |
| 277 | keys: Offload keys to mark as recently used. |
| 278 | req_context: Per-request context. |
| 279 | |
| 280 | """ |
| 281 | return |
| 282 | |
| 283 | @abstractmethod |
| 284 | def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: |
| 285 | """Called when a new request is first seen by the scheduler. |
| 286 | |
| 287 | Returns a RequestOffloadingContext expressing this tier's preference |
| 288 | for how chunks should be offloaded for this request. |
| 289 | |
| 290 | Args: |
| 291 | req_context: Per-request context. |
| 292 | |
| 293 | """ |
| 294 | pass |
| 295 | |
| 296 | def on_request_finished(self, req_context: ReqContext) -> None: |
| 297 | """Called when a request has finished. |
| 298 | |
| 299 | By the time this is called, all per-request calls for this request |
| 300 | (submit_store, submit_load, touch) have already been issued, and none |
| 301 | will follow. Note this does NOT imply the tier's transfers have |
| 302 | completed: jobs already submitted may still be in flight and will |
| 303 | report via get_finished_jobs(). This is the right place to release |
| 304 | per-request bookkeeping. |
| 305 | |
| 306 | Args: |
| 307 | req_context: per-request context. |
| 308 | |
| 309 | """ |
| 310 | return |
| 311 | |
| 312 | def serve_external_requests(self, parent: ParentManager) -> None: |
| 313 | """Process remotely-originated requests using the parent manager. |
| 314 | |
| 315 | Called once per scheduler step, BEFORE _flush_pending_promotions(). |
| 316 | The parent handle is valid only for the duration of this call. |
| 317 | Tiers that don't serve external requests leave this as a no-op. |
| 318 | """ |
| 319 | return |
| 320 | |
| 321 | def on_schedule_end(self, context: ScheduleEndContext) -> None: |
| 322 | """Called once at the end of each scheduler step. |
| 323 | |
| 324 | Args: |
| 325 | context: Per-step context from the scheduler. |
| 326 | |
| 327 | """ |
| 328 | return |
| 329 | |
| 330 | @abstractmethod |
| 331 | def drain_jobs(self) -> None: |
| 332 | """Block until every submitted load/store job has completed or failed. |
| 333 | |
| 334 | After this returns, no tier I/O is touching the primary memoryview, |
| 335 | and every submitted job's result is available from `get_finished_jobs()` |
| 336 | (yielded by a prior call or queued for the next one). Used by |
| 337 | `TieringOffloadingManager.reset_cache` to release primary slots |
| 338 | without racing with in-flight transfers. |
| 339 | |
| 340 | Implementations must not abort a mid-flight transfer: a partial copy |
| 341 | would corrupt either the primary memoryview or the secondary backing |
| 342 | store. Queued (not-yet-started) transfers may be cancelled, but their |
| 343 | failure result must still appear in `get_finished_jobs()`. |
| 344 | """ |
| 345 | pass |
| 346 | |
| 347 | def shutdown(self) -> None: |
| 348 | """Release resources held by this tier (threads, connections, etc.).""" |
| 349 | return |
| 350 | |
| 351 | @classmethod |
| 352 | def build_metric_definitions( |
| 353 | cls, extra_config: dict[str, Any] |
| 354 | ) -> dict[str, OffloadingMetricMetadata]: |
| 355 | """Return Prometheus metric definitions emitted by this tier.""" |
| 356 | return {} |
| 357 | |
| 358 | def get_stats(self) -> "OffloadingConnectorStats | None": |
| 359 | """Return and reset metric observations collected by this tier.""" |
| 360 | return None |