| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """P2P KV cache sharing protocol constants and documentation. |
| 4 | |
| 5 | Protocol Overview |
| 6 | ================= |
| 7 | |
| 8 | Two session types communicate over a bidirectional message channel: |
| 9 | - P2PClientSession (requests blocks from a server) |
| 10 | - P2PServerSession (serves blocks to a client) |
| 11 | |
| 12 | Connection Lifecycle |
| 13 | -------------------- |
| 14 | |
| 15 | 1. Client opens a connection and sends ConnectMsg with its identity, |
| 16 | RDMA metadata, memory layout, and config fingerprint. |
| 17 | 2. Server validates block_len and config_fingerprint, registers the |
| 18 | RDMA peer, and replies with ConnectAckMsg. |
| 19 | 3. Client receives ConnectAckMsg and transitions to ready state |
| 20 | (flushes any queued messages). |
| 21 | 4. Either side may send DisconnectMsg to gracefully close. |
| 22 | |
| 23 | Block Transfer Flow (happy path) |
| 24 | --------------------------------- |
| 25 | |
| 26 | 1. Client sends FetchMsg with a kv_request_id and lists of |
| 27 | block keys + remote indexes where it wants the data written. |
| 28 | A request may run several lookup→fetch rounds; symmetric-P2P |
| 29 | messages carry ROUND_SEQ so each round's supply, demand, and |
| 30 | completion stay isolated. The terminal empty FetchMsg is the |
| 31 | server-side "request finished" signal for the id: parked |
| 32 | LookupMsg batches are popped and ``cb.finish_request`` fires on |
| 33 | each. |
| 34 | 2. Server matches requested blocks against locally stored blocks: |
| 35 | - Blocks already available are transferred immediately via RDMA. |
| 36 | - Blocks not yet available are recorded as "demanded" and |
| 37 | transferred when the server later stores them. |
| 38 | 3. When all blocks for a kv_request_id are transferred, the server |
| 39 | sends TransferDoneMsg (success=True) to the client. |
| 40 | 4. Client reports the load job as complete. |
| 41 | |
| 42 | Abort Flow (timeout path) |
| 43 | -------------------------- |
| 44 | |
| 45 | 1. If the client times out waiting for TransferDoneMsg, it sends |
| 46 | AbortFetchMsg to cancel the request. |
| 47 | 2. Server cancels inflight transfers for that kv_request_id and |
| 48 | replies with AbortAckMsg. |
| 49 | 3. Client receives AbortAckMsg and reports the load job as failed. |
| 50 | 4. If AbortAckMsg itself times out, the client fails the job anyway. |
| 51 | |
| 52 | Message Format |
| 53 | -------------- |
| 54 | |
| 55 | All messages are dicts serialized with msgpack. Every message has a |
| 56 | TYPE_KEY key identifying its type. Additional fields depend on the |
| 57 | message type (see per-message class docstrings below). |
| 58 | |
| 59 | Security |
| 60 | -------- |
| 61 | |
| 62 | Sessions wrap all incoming message handling in try/except to guard |
| 63 | against malformed messages from adversarial or buggy peers. Invalid |
| 64 | messages are logged and dropped without crashing the session. |
| 65 | |
| 66 | The config_fingerprint field in ConnectMsg ensures peers have |
| 67 | compatible model configurations (model, dtype, block sizes). Mismatches |
| 68 | are rejected during the handshake. |
| 69 | """ |
| 70 | |
| 71 | TYPE_KEY = "type" |
| 72 | |
| 73 | |
| 74 | # --------------------------------------------------------------------------- |
| 75 | # Validation helpers |
| 76 | # --------------------------------------------------------------------------- |
| 77 | |
| 78 | |
| 79 | def _require(msg: dict, key: str, typ: type, *, name: str = "") -> None: |
| 80 | """Raise ValueError if msg[key] is missing or not isinstance(typ).""" |
| 81 | val = msg.get(key) |
| 82 | if not isinstance(val, typ): |
| 83 | label = name or key |
| 84 | raise ValueError(f"{label}: expected {typ.__name__}, got {type(val).__name__}") |
| 85 | |
| 86 | |
| 87 | def _require_pos_int(msg: dict, key: str, *, name: str = "") -> None: |
| 88 | """Raise ValueError if msg[key] is not a positive int.""" |
| 89 | val = msg.get(key) |
| 90 | if not isinstance(val, int) or val <= 0: |
| 91 | label = name or key |
| 92 | raise ValueError(f"{label}: expected positive int, got {val!r}") |
| 93 | |
| 94 | |
| 95 | def _require_non_neg_int(msg: dict, key: str, *, name: str = "") -> None: |
| 96 | """Raise ValueError if msg[key] is not a non-negative int.""" |
| 97 | val = msg.get(key) |
| 98 | if not isinstance(val, int) or val < 0: |
| 99 | label = name or key |
| 100 | raise ValueError(f"{label}: expected non-negative int, got {val!r}") |
| 101 | |
| 102 | |
| 103 | def _require_list(msg: dict, key: str, *, name: str = "") -> None: |
| 104 | """Raise ValueError if msg[key] is not a list.""" |
| 105 | val = msg.get(key) |
| 106 | if not isinstance(val, list): |
| 107 | label = name or key |
| 108 | raise ValueError(f"{label}: expected list, got {type(val).__name__}") |
| 109 | |
| 110 | |
| 111 | # --------------------------------------------------------------------------- |
| 112 | # Message classes |
| 113 | # --------------------------------------------------------------------------- |
| 114 | |
| 115 | |
| 116 | class ConnectMsg: |
| 117 | """Client → Server: initial handshake request. |
| 118 | |
| 119 | Fields: |
| 120 | PEER_ID: Local peer identity string. |
| 121 | AGENT_METADATA: RDMA agent metadata (opaque bytes). |
| 122 | BASE_ADDR: Base memory address of the KV block region. |
| 123 | NUM_BLOCKS: Number of blocks in the KV block region. |
| 124 | BLOCK_LEN: Size in bytes of each block (must match between peers). |
| 125 | CONFIG_FINGERPRINT: SHA-256 prefix of the model configuration. |
| 126 | Peers with different fingerprints are incompatible. |
| 127 | HASH_SEED: The peer's effective prefix-cache hash seed (PYTHONHASHSEED |
| 128 | if set, otherwise the built-in default). Block hashes chain from a |
| 129 | seed derived from it, so peers with different values compute |
| 130 | different hashes for identical content and must not exchange blocks. |
| 131 | """ |
| 132 | |
| 133 | TYPE = "connect" |
| 134 | PEER_ID = "peer_id" |
| 135 | AGENT_METADATA = "agent_metadata" |
| 136 | BASE_ADDR = "base_addr" |
| 137 | NUM_BLOCKS = "num_blocks" |
| 138 | BLOCK_LEN = "block_len" |
| 139 | CONFIG_FINGERPRINT = "config_fingerprint" |
| 140 | HASH_SEED = "hash_seed" |
| 141 | |
| 142 | @staticmethod |
| 143 | def validate(msg: dict) -> None: |
| 144 | """Raise ValueError if any field has an invalid type or value.""" |
| 145 | _require(msg, ConnectMsg.PEER_ID, str) |
| 146 | _require(msg, ConnectMsg.AGENT_METADATA, bytes) |
| 147 | _require_non_neg_int(msg, ConnectMsg.BASE_ADDR) |
| 148 | _require_pos_int(msg, ConnectMsg.NUM_BLOCKS) |
| 149 | _require_pos_int(msg, ConnectMsg.BLOCK_LEN) |
| 150 | _require(msg, ConnectMsg.HASH_SEED, str) |
| 151 | |
| 152 | |
| 153 | class ConnectAckMsg: |
| 154 | """Server → Client: handshake acknowledgement. |
| 155 | |
| 156 | Fields: |
| 157 | PEER_ID: Server's peer identity string. |
| 158 | """ |
| 159 | |
| 160 | TYPE = "connect_ack" |
| 161 | PEER_ID = "peer_id" |
| 162 | |
| 163 | @staticmethod |
| 164 | def validate(msg: dict) -> None: |
| 165 | """Raise ValueError if any field has an invalid type or value.""" |
| 166 | _require(msg, ConnectAckMsg.PEER_ID, str) |
| 167 | |
| 168 | |
| 169 | class DisconnectMsg: |
| 170 | """Either → Either: graceful connection close. |
| 171 | |
| 172 | No additional fields beyond TYPE_KEY. |
| 173 | """ |
| 174 | |
| 175 | TYPE = "disconnect" |
| 176 | |
| 177 | |
| 178 | class FetchMsg: |
| 179 | """Client → Server: request blocks for one lookup round. |
| 180 | |
| 181 | A non-empty fetch closes only its round. The terminal empty FetchMsg |
| 182 | (``KEYS`` and ``BLOCK_INDEXES`` both empty) is the "request |
| 183 | finished" signal: the server pops parked LookupMsg batches, calls |
| 184 | ``cb.finish_request`` on each, and drains any leftover supply. |
| 185 | |
| 186 | Fields: |
| 187 | KV_REQUEST_ID: Identifies this block transfer request. |
| 188 | KEYS: List of block keys (OffloadKey bytes). May be empty. |
| 189 | BLOCK_INDEXES: List of remote block indexes (same length as KEYS). |
| 190 | ROUND_SEQ: Lookup round this fetch closes. PD clients never probe |
| 191 | and stay on their single round 0. |
| 192 | """ |
| 193 | |
| 194 | TYPE = "fetch" |
| 195 | KV_REQUEST_ID = "kv_request_id" |
| 196 | KEYS = "keys" |
| 197 | BLOCK_INDEXES = "block_indexes" |
| 198 | ROUND_SEQ = "round_seq" |
| 199 | |
| 200 | @staticmethod |
| 201 | def validate(msg: dict) -> None: |
| 202 | """Raise ValueError if any field has an invalid type or value.""" |
| 203 | _require(msg, FetchMsg.KV_REQUEST_ID, str) |
| 204 | _require_non_neg_int(msg, FetchMsg.ROUND_SEQ) |
| 205 | _require_list(msg, FetchMsg.KEYS) |
| 206 | _require_list(msg, FetchMsg.BLOCK_INDEXES) |
| 207 | keys = msg[FetchMsg.KEYS] |
| 208 | indexes = msg[FetchMsg.BLOCK_INDEXES] |
| 209 | if len(keys) != len(indexes): |
| 210 | raise ValueError( |
| 211 | f"keys/block_indexes length mismatch: {len(keys)} vs {len(indexes)}" |
| 212 | ) |
| 213 | for idx in indexes: |
| 214 | if not isinstance(idx, int) or idx < 0: |
| 215 | raise ValueError(f"block_indexes: invalid index {idx!r}") |
| 216 | |
| 217 | |
| 218 | class LookupMsg: |
| 219 | """Client → Server: probe which block keys the peer holds. |
| 220 | |
| 221 | Sent on the consumer side under symmetric P2P (do_p2p_fetch=true) |
| 222 | after the consumer has aggregated per-block lookups across a |
| 223 | scheduler step. The producer replies with one or more LookupRespMsg |
| 224 | covering the requested keys. |
| 225 | |
| 226 | Fields: |
| 227 | KV_REQUEST_ID: Identifies this lookup transaction. |
| 228 | KEYS: List of block keys (OffloadKey bytes) to probe. |
| 229 | ROUND_SEQ: Lookup round these probes belong to; pinned supply is |
| 230 | parked under it for that round's fetch. |
| 231 | """ |
| 232 | |
| 233 | TYPE = "lookup" |
| 234 | KV_REQUEST_ID = "kv_request_id" |
| 235 | KEYS = "keys" |
| 236 | ROUND_SEQ = "round_seq" |
| 237 | |
| 238 | @staticmethod |
| 239 | def validate(msg: dict) -> None: |
| 240 | """Raise ValueError if any field has an invalid type or value.""" |
| 241 | _require(msg, LookupMsg.KV_REQUEST_ID, str) |
| 242 | _require_list(msg, LookupMsg.KEYS) |
| 243 | _require_non_neg_int(msg, LookupMsg.ROUND_SEQ) |
| 244 | |
| 245 | |
| 246 | class LookupRespMsg: |
| 247 | """Server → Client: per-key hit/miss answer for a prior LookupMsg. |
| 248 | |
| 249 | Carries two parallel arrays of equal length so each (key, |
| 250 | hit) pair is self-describing. The producer is free to split or |
| 251 | coalesce responses across multiple LookupMsgs for the same |
| 252 | KV_REQUEST_ID — the consumer matches each pair back to its |
| 253 | pending entry by (KV_REQUEST_ID, key). |
| 254 | |
| 255 | Fields: |
| 256 | KV_REQUEST_ID: The lookup transaction this responds to. |
| 257 | KEYS: List of block keys answered by this message. |
| 258 | HITS: Parallel list of bools — True if the producer holds the |
| 259 | corresponding block, False otherwise. |
| 260 | """ |
| 261 | |
| 262 | TYPE = "lookup_resp" |
| 263 | KV_REQUEST_ID = "kv_request_id" |
| 264 | KEYS = "keys" |
| 265 | HITS = "hits" |
| 266 | |
| 267 | @staticmethod |
| 268 | def validate(msg: dict) -> None: |
| 269 | """Raise ValueError if any field has an invalid type or value.""" |
| 270 | _require(msg, LookupRespMsg.KV_REQUEST_ID, str) |
| 271 | _require_list(msg, LookupRespMsg.KEYS) |
| 272 | _require_list(msg, LookupRespMsg.HITS) |
| 273 | keys = msg[LookupRespMsg.KEYS] |
| 274 | hits = msg[LookupRespMsg.HITS] |
| 275 | if len(keys) != len(hits): |
| 276 | raise ValueError(f"keys/hits length mismatch: {len(keys)} vs {len(hits)}") |
| 277 | for hit in hits: |
| 278 | if not isinstance(hit, bool): |
| 279 | raise ValueError(f"hits: invalid value {hit!r}") |
| 280 | |
| 281 | |
| 282 | class TransferDoneMsg: |
| 283 | """Server → Client: all blocks transferred for a request. |
| 284 | |
| 285 | Fields: |
| 286 | KV_REQUEST_ID: The request that completed. |
| 287 | SUCCESS: Whether the transfer completed successfully. |
| 288 | ROUND_SEQ: The fetch round that completed. Several loads can be |
| 289 | in flight per id (the scheduler submits loads incrementally), |
| 290 | so completions are matched by round. |
| 291 | """ |
| 292 | |
| 293 | TYPE = "transfer_done" |
| 294 | KV_REQUEST_ID = "kv_request_id" |
| 295 | SUCCESS = "success" |
| 296 | ROUND_SEQ = "round_seq" |
| 297 | |
| 298 | @staticmethod |
| 299 | def validate(msg: dict) -> None: |
| 300 | """Raise ValueError if any field has an invalid type or value.""" |
| 301 | _require(msg, TransferDoneMsg.KV_REQUEST_ID, str) |
| 302 | _require(msg, TransferDoneMsg.SUCCESS, bool) |
| 303 | _require_non_neg_int(msg, TransferDoneMsg.ROUND_SEQ) |
| 304 | |
| 305 | |
| 306 | class AbortFetchMsg: |
| 307 | """Client → Server: cancel a pending request. |
| 308 | |
| 309 | Fields: |
| 310 | KV_REQUEST_ID: The request to cancel. |
| 311 | ROUND_SEQ: The fetch round to cancel. |
| 312 | """ |
| 313 | |
| 314 | TYPE = "abort_fetch" |
| 315 | KV_REQUEST_ID = "kv_request_id" |
| 316 | ROUND_SEQ = "round_seq" |
| 317 | |
| 318 | @staticmethod |
| 319 | def validate(msg: dict) -> None: |
| 320 | """Raise ValueError if any field has an invalid type or value.""" |
| 321 | _require(msg, AbortFetchMsg.KV_REQUEST_ID, str) |
| 322 | _require_non_neg_int(msg, AbortFetchMsg.ROUND_SEQ) |
| 323 | |
| 324 | |
| 325 | class AbortAckMsg: |
| 326 | """Server → Client: acknowledge cancellation. |
| 327 | |
| 328 | Fields: |
| 329 | KV_REQUEST_ID: The request that was cancelled. |
| 330 | ROUND_SEQ: The round that was cancelled; echoes AbortFetchMsg. |
| 331 | """ |
| 332 | |
| 333 | TYPE = "abort_ack" |
| 334 | KV_REQUEST_ID = "kv_request_id" |
| 335 | ROUND_SEQ = "round_seq" |
| 336 | |
| 337 | @staticmethod |
| 338 | def validate(msg: dict) -> None: |
| 339 | """Raise ValueError if any field has an invalid type or value.""" |
| 340 | _require(msg, AbortAckMsg.KV_REQUEST_ID, str) |
| 341 | _require_non_neg_int(msg, AbortAckMsg.ROUND_SEQ) |