| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Metadata dataclasses and helpers for the NIXL connector.""" |
| 4 | |
| 5 | from dataclasses import dataclass, field |
| 6 | from typing import Any |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds, EngineId |
| 10 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 11 | KVConnectorHandshakeMetadata, |
| 12 | KVConnectorMetadata, |
| 13 | ) |
| 14 | from vllm.logger import init_logger |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | TransferHandle = int |
| 19 | ReqId = str |
| 20 | |
| 21 | GET_META_MSG = b"get_meta_msg" |
| 22 | |
| 23 | # Push-mode (WRITE-based) registration notification. |
| 24 | # Sent worker-to-worker over NIXL: D worker -> P worker, encoded as |
| 25 | # PUSH_REG_NOTIF_PREFIX + msgpack(registration_data). |
| 26 | PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" |
| 27 | # |
| 28 | # NIXL Connector Version |
| 29 | # |
| 30 | # Increment this version whenever there is an incompatible change to: |
| 31 | # - NixlAgentMetadata schema |
| 32 | # - kv_transfer_params schema or semantics |
| 33 | # - NIXL transfer protocol or wire format |
| 34 | # - KV cache memory layout or block organization |
| 35 | # - Any other change that breaks P/D interoperability |
| 36 | # |
| 37 | # Version History: |
| 38 | # 1: Initial version with compatibility checking |
| 39 | # 2: Add remote_request_id to kv_transfer_params |
| 40 | # 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata |
| 41 | # 4: Add KV block lease renewal through heartbeats |
| 42 | # 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake |
| 43 | # clock-sync timestamp |
| 44 | # 6: Validate EAGLE/MTP speculative configuration compatibility |
| 45 | # 7: Include NIXL transfer mode (push vs pull) in the compatibility hash |
| 46 | # 8: Add dcp_size and pcp_size to NixlAgentMetadata |
| 47 | # 9: Add block_strides |
| 48 | # 10: Add dense virtual transfer pages for compressed MLA caches |
| 49 | # 11: Add per-region transfer geometry and memory types to NixlAgentMetadata |
| 50 | # 12: Add per-region member names for PP push |
| 51 | # 13: Add packed-member layouts and order-independent packed-push backend hashes |
| 52 | # |
| 53 | NIXL_CONNECTOR_VERSION: int = 13 |
| 54 | |
| 55 | |
| 56 | @dataclass |
| 57 | class NixlAgentMetadata: |
| 58 | engine_id: str |
| 59 | agent_metadata: bytes |
| 60 | kv_caches_base_addr: list[int] |
| 61 | device_id: int |
| 62 | num_blocks: int |
| 63 | block_lens: list[int] |
| 64 | block_strides: list[int] |
| 65 | kv_cache_layout: str |
| 66 | block_size: int |
| 67 | ssm_sizes: tuple[int, int] |
| 68 | attn_backend_name: str |
| 69 | physical_blocks_per_logical_kv_block: int |
| 70 | region_num_blocks: list[int] | None = None |
| 71 | region_group_ids: list[int] | None = None |
| 72 | region_names: list[str] | None = None |
| 73 | region_mem_types: list[str] | None = None |
| 74 | dcp_size: int = 1 |
| 75 | pcp_size: int = 1 |
| 76 | # Layer names sharing each advertised region, in region order. |
| 77 | region_members: list[list[str]] = field(default_factory=list) |
| 78 | # Packed member -> (byte offset in its region's block, bytes per page). |
| 79 | packed_member_layouts: dict[str, tuple[int, int]] = field(default_factory=dict) |
| 80 | |
| 81 | |
| 82 | @dataclass |
| 83 | class NixlHandshakePayload(KVConnectorHandshakeMetadata): |
| 84 | """Wrapper for NIXL handshake sent over the wire. |
| 85 | |
| 86 | Enables two-phase decoding for graceful compatibility checking: |
| 87 | 1. Decode NixlHandshakePayload to get compatibility_hash |
| 88 | 2. Compute local hash and compare |
| 89 | 3. Only if hashes match, decode agent_metadata_bytes |
| 90 | |
| 91 | This prevents decoder errors when NixlAgentMetadata schema is |
| 92 | incompatible, allowing graceful failure with clear error message. |
| 93 | """ |
| 94 | |
| 95 | compatibility_hash: str |
| 96 | agent_metadata_bytes: bytes # NixlAgentMetadata encoded |
| 97 | |
| 98 | |
| 99 | def _get_speculative_compatibility_factors( |
| 100 | vllm_config: VllmConfig, |
| 101 | ) -> dict[str, Any] | None: |
| 102 | """Return NIXL compatibility factors for hidden-state-based speculators.""" |
| 103 | speculative_config = vllm_config.speculative_config |
| 104 | if speculative_config is None or not speculative_config.use_eagle(): |
| 105 | return None |
| 106 | |
| 107 | draft_model_config = speculative_config.draft_model_config |
| 108 | assert draft_model_config is not None |
| 109 | auxiliary_layer_ids = getattr( |
| 110 | draft_model_config.hf_config, |
| 111 | "eagle_aux_hidden_state_layer_ids", |
| 112 | None, |
| 113 | ) |
| 114 | |
| 115 | # kv_cache_dtype is a user override that defaults to None, meaning "inherit |
| 116 | # the target's --kv-cache-dtype". Resolve it to the effective value so an |
| 117 | # explicit setting on one side and inheritance on the other (same effective |
| 118 | # dtype) don't spuriously mismatch. |
| 119 | kv_cache_dtype = ( |
| 120 | speculative_config.kv_cache_dtype or vllm_config.cache_config.cache_dtype |
| 121 | ) |
| 122 | |
| 123 | # Note: the draft attention_backend is intentionally not hashed. Its only |
| 124 | # transfer-relevant effect is the KV block layout/size, which is validated |
| 125 | # per region at runtime in _validate_remote_agent_handshake. The connector |
| 126 | # only sees the raw override here (usually None = auto-select), never the |
| 127 | # resolved backend, so hashing it would cause false mismatches without |
| 128 | # catching anything the runtime layout check misses. |
| 129 | return { |
| 130 | "method": speculative_config.method, |
| 131 | "model": draft_model_config.model, |
| 132 | "revision": draft_model_config.revision, |
| 133 | "code_revision": draft_model_config.code_revision, |
| 134 | "parallel_drafting": speculative_config.parallel_drafting, |
| 135 | "kv_cache_dtype": str(kv_cache_dtype), |
| 136 | "auxiliary_layer_ids": ( |
| 137 | tuple(auxiliary_layer_ids) if auxiliary_layer_ids is not None else None |
| 138 | ), |
| 139 | } |
| 140 | |
| 141 | |
| 142 | def compute_nixl_compatibility_hash( |
| 143 | vllm_config: VllmConfig, |
| 144 | attn_backend_name: str, |
| 145 | transfer_mode: str = "pull", |
| 146 | ) -> str: |
| 147 | """Compute compatibility hash for NIXL KV transfer. |
| 148 | |
| 149 | Hash only the factors that affect whether two NIXL instances can |
| 150 | successfully transfer KV cache data. |
| 151 | |
| 152 | Factors included: |
| 153 | - vLLM version and NIXL connector version |
| 154 | - Model architecture (name, dtype, KV heads, layers) |
| 155 | - KV cache format (dtype, sliding window) |
| 156 | - Attention backend |
| 157 | - EAGLE/MTP configuration that affects transferred state |
| 158 | - Transfer mode (push vs pull) |
| 159 | |
| 160 | The transfer mode is included because the push (WRITE) and pull (READ) |
| 161 | connectors use incompatible transfer protocols; a push connector and a |
| 162 | pull connector must never complete a handshake with each other. |
| 163 | |
| 164 | Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout |
| 165 | are validated at runtime in _validate_remote_agent_handshake and are not |
| 166 | included in this hash to support heterogeneous deployments. |
| 167 | |
| 168 | Note - the set of factors are likely to evolve significantly over |
| 169 | time to be more or less permissive. |
| 170 | |
| 171 | Returns: |
| 172 | SHA-256 hex digest |
| 173 | |
| 174 | """ |
| 175 | from vllm import __version__ as vllm_version |
| 176 | from vllm.config.utils import hash_factors |
| 177 | |
| 178 | model_config = vllm_config.model_config |
| 179 | cache_config = vllm_config.cache_config |
| 180 | is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager |
| 181 | |
| 182 | factors = { |
| 183 | # Version compatibility |
| 184 | "vllm_version": vllm_version, |
| 185 | "nixl_connector_version": NIXL_CONNECTOR_VERSION, |
| 186 | # Model architecture - affects KV cache shape |
| 187 | "model": model_config.model, |
| 188 | "dtype": str(model_config.dtype), |
| 189 | "num_kv_heads": model_config.get_total_num_kv_heads(), |
| 190 | "head_size": model_config.get_head_size(), |
| 191 | "num_hidden_layers": model_config.get_total_num_hidden_layers(), |
| 192 | # Attention backend and KV cache dtype affect memory layout |
| 193 | "attn_backend_name": attn_backend_name, |
| 194 | "cache_dtype": str(cache_config.cache_dtype), |
| 195 | "is_hma_enabled": is_hma_enabled, |
| 196 | "speculative_config": _get_speculative_compatibility_factors(vllm_config), |
| 197 | # push (WRITE) and pull (READ) connectors are protocol-incompatible |
| 198 | "transfer_mode": transfer_mode, |
| 199 | } |
| 200 | |
| 201 | compat_hash = hash_factors(factors) |
| 202 | logger.debug( |
| 203 | "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, " |
| 204 | "cache_dtype=%s, attn_backend=%s)", |
| 205 | compat_hash, |
| 206 | factors["model"], |
| 207 | factors["dtype"], |
| 208 | factors["num_kv_heads"], |
| 209 | factors["cache_dtype"], |
| 210 | attn_backend_name, |
| 211 | ) |
| 212 | return compat_hash |
| 213 | |
| 214 | |
| 215 | @dataclass |
| 216 | class HeartbeatInfo: |
| 217 | """Heartbeat data for a single remote engine, sent from D worker to P.""" |
| 218 | |
| 219 | req_ids: set[ReqId] |
| 220 | host: str |
| 221 | port: int |
| 222 | tp_size: int |
| 223 | dcp_size: int = 1 |
| 224 | pp_size: int = 1 |
| 225 | |
| 226 | |
| 227 | @dataclass |
| 228 | class RemoteMeta: |
| 229 | block_ids: BlockIds |
| 230 | host: str |
| 231 | port: int |
| 232 | engine_id: str |
| 233 | request_id: str |
| 234 | blocks_expiry_time: float | None = None |
| 235 | num_tokens: int | None = None |
| 236 | |
| 237 | |
| 238 | @dataclass |
| 239 | class ReqMeta: |
| 240 | local_block_ids: BlockIds |
| 241 | # To be used when logical block size does not match the kernel block size |
| 242 | local_physical_block_ids: BlockIds |
| 243 | tp_size: int |
| 244 | dcp_size: int = 1 |
| 245 | # Per-KV-cache-group logical blocks this rank already holds, i.e. its |
| 246 | # prefix-cache hit. Fixes where this rank's DCP slice starts relative to |
| 247 | # the remote's; kept per-group since hybrid models (e.g. SWA+FA) can have |
| 248 | # different cache-hit counts per group. |
| 249 | local_num_computed_blocks: tuple[int, ...] = () |
| 250 | remote: RemoteMeta | None = None |
| 251 | # Remote block size, discovered during NIXL handshake (push mode). |
| 252 | remote_block_size: int | None = None |
| 253 | # Remote producer pipeline-parallel size (push mode, D side). |
| 254 | pp_size: int = 1 |
| 255 | # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS |
| 256 | # and expects it in finished_recving; notify-only recvs must not be reported. |
| 257 | awaiting_kvs: bool = False |
| 258 | # Worker-only, per-region physical pages to zero after a successful pull. |
| 259 | # None selects group-based completion; empty lists mean no zeroing. |
| 260 | region_blocks_to_zero: BlockIds | None = None |
| 261 | |
| 262 | |
| 263 | class NixlConnectorMetadata(KVConnectorMetadata): |
| 264 | def __init__(self): |
| 265 | self.reqs_to_recv: dict[ReqId, ReqMeta] = {} |
| 266 | self.reqs_to_save: dict[ReqId, ReqMeta] = {} |
| 267 | self.reqs_to_send: dict[ReqId, float] = {} |
| 268 | # The scheduler process's time.perf_counter() when this metadata was |
| 269 | # built. reqs_to_send deadlines are stamped with the scheduler's |
| 270 | # clock, which is NOT comparable across processes (perf_counter is |
| 271 | # process/boot-local): workers must rebase the remaining TTL onto |
| 272 | # their own clock via this reference. 0.0 = unset (legacy metadata). |
| 273 | self.scheduler_clock: float = 0.0 |
| 274 | self.reqs_in_batch: set[ReqId] = set() |
| 275 | self.reqs_not_processed: set[ReqId] = set() |
| 276 | # Heartbeat data grouped by remote engine, sent by D worker to P. |
| 277 | self.heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {} |
| 278 | # Push mode (D side): registration data the D worker should send to |
| 279 | # P workers via NIXL notification on this step. |
| 280 | self.push_registrations: dict[ReqId, dict[str, Any]] = {} |
| 281 | # Push mode (P side): newly finished request blocks to be matched |
| 282 | # against pending D registrations on the P worker. |
| 283 | self.push_finished_blocks: dict[ReqId, BlockIds] = {} |
| 284 | |
| 285 | def _add_new_req( |
| 286 | self, |
| 287 | local_block_ids: BlockIds, |
| 288 | kv_transfer_params: dict[str, Any], |
| 289 | local_num_computed_blocks: tuple[int, ...] = (), |
| 290 | awaiting_kvs: bool = False, |
| 291 | ) -> ReqMeta: |
| 292 | return ReqMeta( |
| 293 | local_block_ids=local_block_ids, |
| 294 | local_physical_block_ids=local_block_ids, |
| 295 | # P workers don't need to receive these from proxy here. |
| 296 | tp_size=kv_transfer_params.get("tp_size", 1), |
| 297 | dcp_size=kv_transfer_params.get("dcp_size", 1), |
| 298 | remote_block_size=kv_transfer_params.get("remote_block_size"), |
| 299 | pp_size=kv_transfer_params.get("pp_size", 1), |
| 300 | local_num_computed_blocks=local_num_computed_blocks, |
| 301 | awaiting_kvs=awaiting_kvs, |
| 302 | ) |
| 303 | |
| 304 | def add_new_req_to_save( |
| 305 | self, |
| 306 | request_id: ReqId, |
| 307 | local_block_ids: BlockIds, |
| 308 | kv_transfer_params: dict[str, Any], |
| 309 | ): |
| 310 | self.reqs_to_save[request_id] = self._add_new_req( |
| 311 | local_block_ids, kv_transfer_params |
| 312 | ) |
| 313 | |
| 314 | def add_new_req_to_recv( |
| 315 | self, |
| 316 | request_id: ReqId, |
| 317 | local_block_ids: BlockIds, |
| 318 | kv_transfer_params: dict[str, Any], |
| 319 | local_num_computed_blocks: tuple[int, ...] = (), |
| 320 | awaiting_kvs: bool = False, |
| 321 | ): |
| 322 | req = self._add_new_req( |
| 323 | local_block_ids, |
| 324 | kv_transfer_params, |
| 325 | local_num_computed_blocks, |
| 326 | awaiting_kvs, |
| 327 | ) |
| 328 | req.remote = RemoteMeta( |
| 329 | block_ids=kv_transfer_params["remote_block_ids"], |
| 330 | engine_id=kv_transfer_params["remote_engine_id"], |
| 331 | request_id=kv_transfer_params["remote_request_id"], |
| 332 | host=kv_transfer_params["remote_host"], |
| 333 | port=kv_transfer_params["remote_port"], |
| 334 | blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"), |
| 335 | num_tokens=kv_transfer_params.get("remote_num_tokens"), |
| 336 | ) |
| 337 | self.reqs_to_recv[request_id] = req |