| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State |
| 4 | communication in vLLM v1 |
| 5 | |
| 6 | The class provides the following primitives: |
| 7 | Scheduler-side: runs in the scheduler, binds metadata, which |
| 8 | is used by the worker-side to load/save KV cache. |
| 9 | get_num_new_matched_tokens() - get number of new tokens |
| 10 | that exist in the remote KV cache. Might be called multiple |
| 11 | times for a given request and should be side-effect free. |
| 12 | update_state_after_alloc() - update KVConnector state after |
| 13 | temporary buffer alloc by the CacheManager. |
| 14 | update_connector_output() - update KVConnector state after |
| 15 | output is received from worker-side connectors. |
| 16 | request_finished() - called once when a request is finished, |
| 17 | with the computed kv cache blocks for the request. |
| 18 | Returns whether KV cache should be freed now or if the |
| 19 | connector now assumes responsibility for freeing the |
| 20 | the blocks asynchronously. Also optionally returns KV |
| 21 | transfer params. |
| 22 | take_events() - returns new KV events that were collected |
| 23 | by the connector since the last call. |
| 24 | |
| 25 | Worker-side: runs in each worker, loads/saves KV cache to/from |
| 26 | the Connector based on the metadata. |
| 27 | handle_preemptions() - called for handling preempted requests |
| 28 | or request evicted blocks before they are overwritten |
| 29 | |
| 30 | start_load_kv() - starts loading all KVs (maybe async) |
| 31 | wait_for_layer_load() - blocks until layer i load is done |
| 32 | |
| 33 | save_kv_layer() - starts saving KV for layer i (maybe async) |
| 34 | wait_for_save() - submits saves and performs required synchronization |
| 35 | |
| 36 | get_transfer_results() - returns async send/receive completions and |
| 37 | receive failures in one snapshot. |
| 38 | build_connector_worker_meta() - builds metadata to be sent |
| 39 | back to the scheduler-side connector |
| 40 | """ |
| 41 | |
| 42 | import enum |
| 43 | from abc import ABC, abstractmethod |
| 44 | from collections.abc import Callable, Iterable |
| 45 | from contextlib import AbstractContextManager |
| 46 | from dataclasses import dataclass, field |
| 47 | from typing import TYPE_CHECKING, Any, Literal |
| 48 | |
| 49 | import torch |
| 50 | |
| 51 | from vllm.logger import init_logger |
| 52 | from vllm.v1.attention.backend import AttentionMetadata |
| 53 | from vllm.v1.core.sched.output import SchedulerOutput |
| 54 | from vllm.v1.outputs import KVConnectorOutput |
| 55 | |
| 56 | if TYPE_CHECKING: |
| 57 | from vllm.config import VllmConfig |
| 58 | from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents |
| 59 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( |
| 60 | KVConnectorPromMetrics, |
| 61 | KVConnectorStats, |
| 62 | PromMetric, |
| 63 | PromMetricT, |
| 64 | ) |
| 65 | from vllm.forward_context import ForwardContext |
| 66 | from vllm.v1.core.block_pool import BlockPool |
| 67 | from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager |
| 68 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 69 | from vllm.v1.request import Request |
| 70 | |
| 71 | # s_tensor_list, d_tensor_list, s_indices, d_indices, direction |
| 72 | CopyBlocksOp = Callable[ |
| 73 | [ |
| 74 | dict[str, torch.Tensor], |
| 75 | dict[str, torch.Tensor], |
| 76 | list[int], |
| 77 | list[int], |
| 78 | Literal["h2d", "d2h"], |
| 79 | ], |
| 80 | None, |
| 81 | ] |
| 82 | |
| 83 | logger = init_logger(__name__) |
| 84 | |
| 85 | |
| 86 | @dataclass |
| 87 | class KVConnectorTransferResults: |
| 88 | """Asynchronous transfer completions from one worker snapshot. |
| 89 | |
| 90 | Failed receives also appear in ``finished_recving`` so the scheduler can |
| 91 | release the request from its transfer wait state. |
| 92 | """ |
| 93 | |
| 94 | finished_sending: set[str] = field(default_factory=set) |
| 95 | finished_recving: set[str] = field(default_factory=set) |
| 96 | failed_recving: set[str] = field(default_factory=set) |
| 97 | |
| 98 | |
| 99 | class SupportsHMA(ABC): |
| 100 | """The class that indicates the corresponding connector supports hybrid memory |
| 101 | allocator (HMA). |
| 102 | This is required to use the connector together with hybrid memory allocator. |
| 103 | """ |
| 104 | |
| 105 | @abstractmethod |
| 106 | def request_finished_all_groups( |
| 107 | self, |
| 108 | request: "Request", |
| 109 | block_ids: tuple[list[int], ...], |
| 110 | ) -> tuple[bool, dict[str, Any] | None]: |
| 111 | """Called exactly once when a request has finished for all kv cache groups, |
| 112 | before its blocks are freed for each group. |
| 113 | |
| 114 | NOTE(Kuntai): This function is only supported by connectors that support HMA. |
| 115 | |
| 116 | The connector may assumes responsibility for freeing the blocks |
| 117 | asynchronously by returning True. |
| 118 | |
| 119 | Returns: |
| 120 | True if the request is being saved/sent asynchronously and blocks |
| 121 | should not be freed until the request_id is returned from |
| 122 | get_finished(). |
| 123 | Optional KVTransferParams to be included in the request outputs |
| 124 | returned by the engine. |
| 125 | |
| 126 | """ |
| 127 | raise NotImplementedError |
| 128 | |
| 129 | |
| 130 | def supports_hma(connector: Any) -> bool: |
| 131 | if isinstance(connector, type): |
| 132 | return issubclass(connector, SupportsHMA) |
| 133 | else: |
| 134 | return isinstance(connector, SupportsHMA) |
| 135 | |
| 136 | |
| 137 | class KVConnectorRole(enum.Enum): |
| 138 | # Connector running in the scheduler process |
| 139 | SCHEDULER = 0 |
| 140 | |
| 141 | # Connector running in the worker process |
| 142 | WORKER = 1 |
| 143 | |
| 144 | |
| 145 | class KVConnectorHandshakeMetadata(ABC): # noqa: B024 |
| 146 | """Metadata used for out of band connector handshake between |
| 147 | P/D workers. This needs to serializable. |
| 148 | """ |
| 149 | |
| 150 | pass |
| 151 | |
| 152 | |
| 153 | class KVConnectorMetadata(ABC): # noqa: B024 |
| 154 | """Abstract Metadata used to communicate |
| 155 | Scheduler KVConnector -> Worker KVConnector. |
| 156 | """ |
| 157 | |
| 158 | pass |
| 159 | |
| 160 | |
| 161 | class KVConnectorWorkerMetadata(ABC): |
| 162 | """Abstract Metadata used to communicate back |
| 163 | Worker KVConnector -> Scheduler KVConnector. |
| 164 | |
| 165 | Each worker can output its own metadata. |
| 166 | For a single engine step, all metadata objects returned by workers |
| 167 | will be aggregated using the `aggregate` method below, before |
| 168 | being passed to the Scheduler KVConnector. |
| 169 | """ |
| 170 | |
| 171 | @abstractmethod |
| 172 | def aggregate( |
| 173 | self, other: "KVConnectorWorkerMetadata" |
| 174 | ) -> "KVConnectorWorkerMetadata": |
| 175 | """Aggregate metadata with another `KVConnectorWorkerMetadata` object.""" |
| 176 | pass |
| 177 | |
| 178 | |
| 179 | class KVConnectorBase_V1(ABC): |
| 180 | """Base class for KV connectors.""" |
| 181 | |
| 182 | @property |
| 183 | def supports_divergent_local_hybrid_hits(self) -> bool: |
| 184 | """Whether external hits can complete divergent local hybrid hits. |
| 185 | |
| 186 | A capable connector restores lagging recurrent state when the local |
| 187 | full-attention group reaches a deeper boundary. Defaults to False. |
| 188 | """ |
| 189 | return False |
| 190 | |
| 191 | @property |
| 192 | def requires_kv_delivery(self) -> bool: |
| 193 | """Whether this connector hands off KV that must be reliably delivered. |
| 194 | |
| 195 | If True, a request preempted while its hand-off is still pending is |
| 196 | recomputed rather than allowed to finish and hand off blocks that the |
| 197 | preemption already freed. Defaults to the producer role, since only a |
| 198 | producer hands KV off when a request completes. Best-effort caches |
| 199 | return False, as a dropped save is just a future cache miss. |
| 200 | """ |
| 201 | return self._kv_transfer_config.is_kv_producer |
| 202 | |
| 203 | def __init__( |
| 204 | self, |
| 205 | vllm_config: "VllmConfig", |
| 206 | role: KVConnectorRole, |
| 207 | kv_cache_config: "KVCacheConfig", |
| 208 | ): |
| 209 | logger.warning( |
| 210 | "Initializing KVConnectorBase_V1. This API is experimental and " |
| 211 | "subject to change in the future as we iterate the design." |
| 212 | ) |
| 213 | self._connector_metadata: KVConnectorMetadata | None = None |
| 214 | self._vllm_config = vllm_config |
| 215 | if vllm_config.kv_transfer_config is not None: |
| 216 | self._kv_transfer_config = vllm_config.kv_transfer_config |
| 217 | else: |
| 218 | raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1") |
| 219 | self._kv_cache_config = kv_cache_config |
| 220 | self._kv_cache_manager: KVCacheManager | None = None |
| 221 | self._role = role |
| 222 | |
| 223 | @property |
| 224 | def role(self) -> KVConnectorRole: |
| 225 | return self._role |
| 226 | |
| 227 | # ============================== |
| 228 | # Worker-side methods |
| 229 | # ============================== |
| 230 | |
| 231 | def get_mem_pool_context(self) -> AbstractContextManager | None: |
| 232 | """Return a custom KV cache allocation context, if configured. |
| 233 | |
| 234 | Returning None uses the engine's default memory pool. |
| 235 | """ |
| 236 | return None |
| 237 | |
| 238 | def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None: |
| 239 | """Set the connector metadata from the scheduler. |
| 240 | |
| 241 | This function should be called by the model runner every time |
| 242 | before the model execution. The metadata will be used for runtime |
| 243 | KV cache loading and saving. |
| 244 | |
| 245 | Args: |
| 246 | connector_metadata (dict): the connector metadata. |
| 247 | |
| 248 | """ |
| 249 | self._connector_metadata = connector_metadata |
| 250 | |
| 251 | def clear_connector_metadata(self) -> None: |
| 252 | """Clear the connector metadata. |
| 253 | |
| 254 | This function should be called by the model runner every time |
| 255 | after the model execution. |
| 256 | """ |
| 257 | self._connector_metadata = None |
| 258 | |
| 259 | def _get_connector_metadata(self) -> KVConnectorMetadata: |
| 260 | """Get the connector metadata. |
| 261 | |
| 262 | This function should only be called inside the connector. |
| 263 | |
| 264 | Returns: |
| 265 | ConnectorMetadata: the connector metadata. |
| 266 | |
| 267 | """ |
| 268 | # Should only be called while set to valid metadata. |
| 269 | assert self._connector_metadata is not None |
| 270 | return self._connector_metadata |
| 271 | |
| 272 | def has_connector_metadata(self) -> bool: |
| 273 | """Check whether the connector metadata is currently set. |
| 274 | |
| 275 | Returns: |
| 276 | bool: True if connector metadata exists, False otherwise. |
| 277 | |
| 278 | """ |
| 279 | return self._connector_metadata is not None |
| 280 | |
| 281 | def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): |
| 282 | """Initialize with the KV caches. Useful for pre-registering the |
| 283 | KV Caches in the KVConnector (e.g. for NIXL). |
| 284 | |
| 285 | Args: |
| 286 | kv_caches: dictionary of layer names, kv cache |
| 287 | |
| 288 | """ |
| 289 | return |
| 290 | |
| 291 | def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): |
| 292 | """Set the xPU-specific ops for copying KV between host and device. |
| 293 | Needed when host buffer is used for kv transfer (e.g., in NixlConnector) |
| 294 | """ |
| 295 | return |
| 296 | |
| 297 | # TODO(NickLucche): group model-runner lifecycle hooks in the interface. |
| 298 | def finish_forward(self) -> None: |
| 299 | """Notify the connector that the model no longer reads this step's KV.""" |
| 300 | return |
| 301 | |
| 302 | def reset_capture_state(self) -> None: |
| 303 | """Reset worker state mutated while capturing CUDA graphs.""" |
| 304 | return |
| 305 | |
| 306 | def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata): |
| 307 | """Handle preempted requests or evicted blocks BEFORE they are overwritten. |
| 308 | Needed for connectors which use async saves (e.g., OffloadingConnector) |
| 309 | """ |
| 310 | return |
| 311 | |
| 312 | @abstractmethod |
| 313 | def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: |
| 314 | """Start loading the KV cache from the connector to vLLM's paged |
| 315 | KV buffer. Loads required by the current forward start before it; |
| 316 | independent asynchronous loads may start after it is submitted. |
| 317 | |
| 318 | Args: |
| 319 | forward_context (ForwardContext): the forward context. |
| 320 | **kwargs: additional arguments for the load operation |
| 321 | |
| 322 | Note: |
| 323 | The number of elements in kv_caches and layer_names should be |
| 324 | the same. |
| 325 | |
| 326 | """ |
| 327 | pass |
| 328 | |
| 329 | @abstractmethod |
| 330 | def wait_for_layer_load(self, layer_name: str) -> None: |
| 331 | """Block until the KV for a specific layer is loaded into vLLM's |
| 332 | paged buffer. This is called from within attention layer to ensure |
| 333 | async copying from start_load_kv is complete. |
| 334 | |
| 335 | This interface will be useful for layer-by-layer pipelining. |
| 336 | |
| 337 | Args: |
| 338 | layer_name: the name of that layer |
| 339 | |
| 340 | """ |
| 341 | pass |
| 342 | |
| 343 | @abstractmethod |
| 344 | def save_kv_layer( |
| 345 | self, |
| 346 | layer_name: str, |
| 347 | kv_layer: torch.Tensor, |
| 348 | attn_metadata: "AttentionMetadata", |
| 349 | **kwargs: Any, |
| 350 | ) -> None: |
| 351 | """Start saving a layer of KV cache from vLLM's paged buffer |
| 352 | to the connector. This is called from within attention layer to |
| 353 | enable async copying during execution. |
| 354 | |
| 355 | Args: |
| 356 | layer_name (str): the name of the layer. |
| 357 | kv_layer (torch.Tensor): the paged KV buffer of the current |
| 358 | layer in vLLM. |
| 359 | attn_metadata (AttentionMetadata): the attention metadata. |
| 360 | **kwargs: additional arguments for the save operation. |
| 361 | |
| 362 | """ |
| 363 | pass |
| 364 | |
| 365 | @abstractmethod |
| 366 | def wait_for_save(self): |
| 367 | """Submit saves and perform required synchronization. |
| 368 | |
| 369 | Called once per step, including empty steps, after any target and draft |
| 370 | forwards. Metadata remains bound during this call. Do nothing if there |
| 371 | is no save work. |
| 372 | |
| 373 | Async saves must retain their source blocks until completion is reported |
| 374 | to the scheduler. Otherwise, wait for saves to finish before returning. |
| 375 | """ |
| 376 | pass |
| 377 | |
| 378 | def get_finished( |
| 379 | self, finished_req_ids: set[str] |
| 380 | ) -> tuple[set[str] | None, set[str] | None]: |
| 381 | """Notifies worker-side connector ids of requests that have |
| 382 | finished generating tokens on the worker. |
| 383 | The scheduler process (via the Executors) will use this output |
| 384 | to track which workers are done. |
| 385 | |
| 386 | Returns: |
| 387 | ids of requests that have finished asynchronous transfer |
| 388 | (requests that previously returned True from request_finished()), |
| 389 | tuple of (sending/saving ids, recving/loading ids). |
| 390 | The finished saves/sends req ids must belong to a set provided in a |
| 391 | call to this method (this call or a prior one). |
| 392 | |
| 393 | """ |
| 394 | return None, None |
| 395 | |
| 396 | def get_transfer_results( |
| 397 | self, finished_req_ids: set[str] |
| 398 | ) -> KVConnectorTransferResults: |
| 399 | """Return completed sends, receives, and receive failures together.""" |
| 400 | finished_sending, finished_recving = self.get_finished(finished_req_ids) |
| 401 | return KVConnectorTransferResults( |
| 402 | finished_sending=set(finished_sending or ()), |
| 403 | finished_recving=set(finished_recving or ()), |
| 404 | ) |
| 405 | |
| 406 | def get_block_ids_with_load_errors(self) -> set[int]: |
| 407 | """Get the set of block IDs that failed to load. |
| 408 | |
| 409 | Returns: |
| 410 | Set of block IDs that encountered load errors. |
| 411 | Empty set if no load errors occurred. |
| 412 | |
| 413 | Notes: |
| 414 | - Applies to both sync- and async-loading requests. |
| 415 | - Async loading: failed blocks may be reported in any forward pass |
| 416 | up to and including the pass where the request ID is returned by |
| 417 | `get_transfer_results()`. Even if failures occur, the request |
| 418 | must still be reported as finished receiving, and the failed |
| 419 | block IDs must appear here no later than that same pass. |
| 420 | - Sync loading: failed blocks should be reported in the forward |
| 421 | pass in which they are detected. |
| 422 | |
| 423 | """ |
| 424 | return set() |
| 425 | |
| 426 | def shutdown(self): |
| 427 | """Shutdown the connector. This is called when the worker process |
| 428 | is shutting down to ensure that all the async operations are |
| 429 | completed and the connector is cleaned up properly. |
| 430 | """ |
| 431 | return None |
| 432 | |
| 433 | def get_kv_connector_stats(self) -> "KVConnectorStats | None": |
| 434 | """Get the KV connector stats collected during the last interval.""" |
| 435 | return None |
| 436 | |
| 437 | def get_kv_connector_kv_cache_events(self) -> "KVConnectorKVEvents | None": |
| 438 | """Get the KV connector kv cache events collected during the last interval. |
| 439 | This function should be called by the model runner every time after the |
| 440 | model execution and before cleanup. |
| 441 | """ |
| 442 | return None |
| 443 | |
| 444 | def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: |
| 445 | """Get the KVConnector handshake metadata for this connector. |
| 446 | This metadata is used for out-of-band connector handshake |
| 447 | between P/D workers. |
| 448 | |
| 449 | Returns: |
| 450 | KVConnectorHandshakeMetadata: the handshake metadata. |
| 451 | None if no handshake metadata is available. |
| 452 | |
| 453 | """ |
| 454 | return None |
| 455 | |
| 456 | def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: |
| 457 | """Build the KVConnector worker metadata for this engine step. |
| 458 | |
| 459 | Returns: |
| 460 | KVConnectorWorkerMetadata: the worker metadata. |
| 461 | None if no worker metadata is available. |
| 462 | |
| 463 | """ |
| 464 | return None |
| 465 | |
| 466 | # ============================== |
| 467 | # Scheduler-side methods |
| 468 | # ============================== |
| 469 | |
| 470 | def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None: |
| 471 | """Bind the scheduler's cache manager after it has been constructed.""" |
| 472 | self._kv_cache_manager = kv_cache_manager |
| 473 | self.bind_gpu_block_pool(kv_cache_manager.block_pool) |
| 474 | |
| 475 | def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: |
| 476 | """Bind the GPU block pool to the connector for per-GPU block status tracking. |
| 477 | For example, inc/dec ref counts, or iterate over the prefix cache blocks. |
| 478 | |
| 479 | Args: |
| 480 | gpu_block_pool: the GPU block pool. |
| 481 | |
| 482 | """ |
| 483 | return |
| 484 | |
| 485 | @abstractmethod |
| 486 | def get_num_new_matched_tokens( |
| 487 | self, |
| 488 | request: "Request", |
| 489 | num_computed_tokens: int, |
| 490 | ) -> tuple[int | None, bool]: |
| 491 | """Get number of new tokens that can be loaded from the |
| 492 | external KV cache beyond the num_computed_tokens. |
| 493 | |
| 494 | Args: |
| 495 | request (Request): the request object. |
| 496 | num_computed_tokens (int): the number of locally |
| 497 | computed tokens for this request |
| 498 | |
| 499 | Returns: |
| 500 | A tuple with the following elements: |
| 501 | - An optional number of tokens that can be loaded from the |
| 502 | external KV cache beyond what is already computed. |
| 503 | If None, it means that the connector needs more time to |
| 504 | determine the number of matched tokens, and the scheduler |
| 505 | should query for this request again later. |
| 506 | - `True` if external KV cache tokens will be loaded |
| 507 | asynchronously (between scheduler steps). Must be |
| 508 | 'False' if the first element is 0. |
| 509 | |
| 510 | Notes: |
| 511 | The connector should only consider the largest prefix of prompt- |
| 512 | tokens for which KV cache is actually available at the time of the |
| 513 | call. If the cache cannot be loaded for some tokens (e.g., due to |
| 514 | connectivity issues or eviction), those tokens must not be taken |
| 515 | into account. |
| 516 | |
| 517 | """ |
| 518 | pass |
| 519 | |
| 520 | @abstractmethod |
| 521 | def update_state_after_alloc( |
| 522 | self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int |
| 523 | ): |
| 524 | """Update KVConnector state after block allocation. |
| 525 | |
| 526 | If get_num_new_matched_tokens previously returned True for a |
| 527 | request, this function may be called twice for that same request - |
| 528 | first when blocks are allocated for the connector tokens to be |
| 529 | asynchronously loaded into, and second when any additional blocks |
| 530 | are allocated, after the load/transfer is complete. |
| 531 | |
| 532 | Decide whether to load based on ``num_external_tokens``, not on |
| 533 | whether ``blocks`` is empty: ``blocks`` may be non-empty even when |
| 534 | ``num_external_tokens == 0`` (e.g. a non-chosen sub-connector of |
| 535 | MultiConnector still receives the request's real blocks). |
| 536 | |
| 537 | Args: |
| 538 | request (Request): the request object. |
| 539 | blocks (KVCacheBlocks): the blocks allocated for the request. |
| 540 | num_external_tokens (int): the number of tokens to load from the |
| 541 | external KV cache. 0 means nothing should be loaded. |
| 542 | |
| 543 | """ |
| 544 | pass |
| 545 | |
| 546 | @abstractmethod |
| 547 | def build_connector_meta( |
| 548 | self, scheduler_output: SchedulerOutput |
| 549 | ) -> KVConnectorMetadata: |
| 550 | """Build the connector metadata for this step. |
| 551 | |
| 552 | This function should NOT modify fields in the scheduler_output. |
| 553 | Also, calling this function will reset the state of the connector. |
| 554 | |
| 555 | Args: |
| 556 | scheduler_output (SchedulerOutput): the scheduler output object. |
| 557 | |
| 558 | """ |
| 559 | pass |
| 560 | |
| 561 | def on_new_request(self, request: "Request") -> None: |
| 562 | """Called by the scheduler when a new request is added. |
| 563 | |
| 564 | Connectors can override this to inspect the request and perform |
| 565 | bookkeeping. The default implementation is a no-op. |
| 566 | """ |
| 567 | return |
| 568 | |
| 569 | def update_connector_output(self, connector_output: KVConnectorOutput): |
| 570 | """Update KVConnector state from worker-side connectors output. |
| 571 | |
| 572 | Args: |
| 573 | connector_output (KVConnectorOutput): the worker-side |
| 574 | connectors output. |
| 575 | |
| 576 | """ |
| 577 | return |
| 578 | |
| 579 | def request_finished( |
| 580 | self, |
| 581 | request: "Request", |
| 582 | block_ids: list[int], |
| 583 | ) -> tuple[bool, dict[str, Any] | None]: |
| 584 | """Called exactly once when a request has finished, before its blocks are |
| 585 | freed. |
| 586 | |
| 587 | The connector may assumes responsibility for freeing the blocks |
| 588 | asynchronously by returning True. |
| 589 | |
| 590 | Returns: |
| 591 | True if the request is being saved/sent asynchronously and blocks |
| 592 | should not be freed until the request_id is returned from |
| 593 | get_finished(). |
| 594 | Optional KVTransferParams to be included in the request outputs |
| 595 | returned by the engine. |
| 596 | |
| 597 | """ |
| 598 | return False, None |
| 599 | |
| 600 | def register_finished_partial_tail( |
| 601 | self, |
| 602 | request: "Request", |
| 603 | block_ids: tuple[list[int], ...], |
| 604 | partial_tail_offloads: list[tuple[int, int, int]], |
| 605 | ) -> bool: |
| 606 | """Register finish-time partial-tail sources before block cleanup. |
| 607 | |
| 608 | Returns True when the connector accepts responsibility for the sources |
| 609 | and the request's blocks must remain alive until ``get_finished()``. |
| 610 | """ |
| 611 | return False |
| 612 | |
| 613 | def take_events(self) -> Iterable["KVCacheEvent"]: |
| 614 | """Take the KV cache events from the connector. |
| 615 | |
| 616 | Yields: |
| 617 | New KV cache events since the last call. |
| 618 | |
| 619 | """ |
| 620 | return () |
| 621 | |
| 622 | def has_pending_block_frees(self) -> bool: |
| 623 | """Whether pending transfers can release blocks instead of preemption.""" |
| 624 | return False |
| 625 | |
| 626 | def has_pending_push_work(self) -> bool: |
| 627 | """Return True if the connector has push-mode work that requires |
| 628 | the engine main loop to keep stepping (e.g. a P-side request whose |
| 629 | KV blocks are waiting to be WRITTEN to a D node). |
| 630 | |
| 631 | Connectors that don't implement push-based KV transfer should |
| 632 | leave this as False. |
| 633 | """ |
| 634 | # TODO: replace with a more general connector hook for keeping the |
| 635 | # scheduler alive (e.g. extend has_unfinished_requests). |
| 636 | return False |
| 637 | |
| 638 | @classmethod |
| 639 | def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: |
| 640 | """Get the required KV cache layout for this connector. |
| 641 | |
| 642 | Args: |
| 643 | vllm_config (VllmConfig): the vllm config. |
| 644 | |
| 645 | Returns: |
| 646 | str: the required KV cache layout. e.g. HND, or NHD. |
| 647 | None if the connector does not require a specific layout. |
| 648 | |
| 649 | """ |
| 650 | if cls is KVConnectorBase_V1: |
| 651 | raise TypeError( |
| 652 | "get_required_kvcache_layout should not be called " |
| 653 | "on the abstract base class" |
| 654 | ) |
| 655 | return None |
| 656 | |
| 657 | @classmethod |
| 658 | def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool: |
| 659 | """Check if this connector requires PIECEWISE CUDA graph mode. |
| 660 | |
| 661 | Connectors that use asynchronous layer-by-layer operations |
| 662 | (wait_for_layer_load/save_kv_layer) should override this method |
| 663 | to return True when those operations are enabled. These operations |
| 664 | cannot be captured in CUDA graphs and will be skipped during replay, |
| 665 | causing data races. PIECEWISE mode allows Python code to execute |
| 666 | between graph pieces, ensuring proper synchronization. |
| 667 | |
| 668 | Args: |
| 669 | extra_config: The kv_connector_extra_config dict from |
| 670 | KVTransferConfig. |
| 671 | |
| 672 | Returns: |
| 673 | True if this connector requires PIECEWISE CUDA graph mode, |
| 674 | False otherwise. |
| 675 | |
| 676 | """ |
| 677 | return False |
| 678 | |
| 679 | def get_finished_count(self) -> int | None: |
| 680 | """Get the count of requests expected to complete send/receive operations |
| 681 | via this connector. This method is used to initialize the |
| 682 | KVOutputAggregator, overwriting the default world_size. |
| 683 | |
| 684 | Returns: |
| 685 | int: expected sending or receiving completion count. |
| 686 | |
| 687 | """ |
| 688 | return None |
| 689 | |
| 690 | @classmethod |
| 691 | def build_kv_connector_stats( |
| 692 | cls, data: dict[str, Any] | None = None |
| 693 | ) -> "KVConnectorStats | None": |
| 694 | """KVConnectorStats resolution method. This method allows dynamically |
| 695 | registered connectors to return their own KVConnectorStats object, |
| 696 | which can implement custom aggregation logic on the data dict. |
| 697 | """ |
| 698 | return None |
| 699 | |
| 700 | def set_xfer_handshake_metadata( |
| 701 | self, metadata: dict[int, KVConnectorHandshakeMetadata] |
| 702 | ) -> None: |
| 703 | """Set the KV connector handshake metadata for this connector. |
| 704 | |
| 705 | Args: |
| 706 | metadata (KVConnectorHandshakeMetadata): the handshake metadata to set. |
| 707 | |
| 708 | """ |
| 709 | return None |
| 710 | |
| 711 | def set_xfer_handshake_metadata_pp_aware( |
| 712 | self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata] |
| 713 | ) -> None: |
| 714 | """Set handshake metadata keyed by (pp_rank, tp_rank). |
| 715 | - Default implementation assumes pp_rank is always 0 |
| 716 | - PP-aware connectors override this to consume all PP producer shards. |
| 717 | """ |
| 718 | if any(pp_rank != 0 for pp_rank, _ in metadata): |
| 719 | raise ValueError( |
| 720 | f"{type(self).__name__} received pp_rank > 0 handshake metadata " |
| 721 | "but does not support PP-disaggregated KV transfer." |
| 722 | ) |
| 723 | self.set_xfer_handshake_metadata( |
| 724 | {tp_rank: meta for (_, tp_rank), meta in metadata.items()} |
| 725 | ) |
| 726 | |
| 727 | @classmethod |
| 728 | def build_prom_metrics( |
| 729 | cls, |
| 730 | vllm_config: "VllmConfig", |
| 731 | metric_types: dict[type["PromMetric"], type["PromMetricT"]], |
| 732 | labelnames: list[str], |
| 733 | per_engine_labelvalues: dict[int, list[object]], |
| 734 | ) -> "KVConnectorPromMetrics | None": |
| 735 | """Create a KVConnectorPromMetrics subclass which should register |
| 736 | per-connector Prometheus metrics and implement observe() to |
| 737 | expose connector transfer stats via Prometheus. |
| 738 | """ |
| 739 | return None |
| 740 | |
| 741 | def reset_cache(self) -> bool | None: |
| 742 | """Reset the connector's internal cache. |
| 743 | |
| 744 | Returns: |
| 745 | bool: True if the cache was successfully reset, False otherwise. |
| 746 | |
| 747 | """ |
| 748 | logger.debug( |
| 749 | "Connector cache reset requested, but %s does not implement reset_cache().", |
| 750 | type(self).__name__, |
| 751 | ) |
| 752 | |
| 753 | return None |