diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py index 3be1be18e534..b9eabeec9cb0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py @@ -30,11 +30,45 @@ logger = init_logger(__name__) +# Length of the suffix appended by InputProcessor.assign_request_id(): +# f"{external_req_id}-{random_uuid():.8}" +# The suffix is "-" + 8 hex chars = 9 characters total. +_INTERNAL_ID_SUFFIX_LEN = 9 + + +def _strip_internal_id_suffix(request_id: str) -> str: + """Recover the original proxy-generated request ID by stripping the + random suffix that InputProcessor.assign_request_id() appended. + + assign_request_id() produces: + f"{external_req_id}-{random_uuid():.8}" + + So the internal ID is always: + "-" + <8 hex chars>. + We strip the last 9 characters to get the original ID back. + + This is critical for P2P disaggregated serving: the Prefill and Decode + instances each call assign_request_id() independently, producing + different random suffixes. Using the stripped (original) ID as the + NCCL send/recv key ensures both sides agree on the same key. + """ + if ( + len(request_id) >= _INTERNAL_ID_SUFFIX_LEN + and request_id[-_INTERNAL_ID_SUFFIX_LEN] == "-" + and all(c in "0123456789abcdef" for c in request_id[-8:]) + ): + return request_id[:-_INTERNAL_ID_SUFFIX_LEN] + # Fallback: if the format doesn't match, return as-is. + return request_id + @dataclass class ReqMeta: - # Request Id + # Internal request id (used for scheduler-side tracking) request_id: str + # Transfer id: the original proxy-generated request id, identical on + # both Prefill and Decode instances. Used as the key for P2P NCCL + # send/recv matching and for parsing remote addresses. + transfer_id: str # Request block ids block_ids: torch.Tensor # Request num tokens @@ -47,6 +81,7 @@ def make_meta( block_ids_tensor = torch.tensor(block_ids) return ReqMeta( request_id=request_id, + transfer_id=_strip_internal_id_suffix(request_id), block_ids=block_ids_tensor, num_tokens=len(token_ids), ) @@ -201,8 +236,8 @@ def inject_kv_into_layer( # Load the KV for each request each layer for request in metadata.requests: - request_id = request.request_id - ip, port = self.parse_request_id(request_id, False) + transfer_id = request.transfer_id + ip, port = self.parse_request_id(transfer_id, False) remote_address = ip + ":" + str(port + self._rank) for layer_name in forward_context.no_compile_layers: layer = forward_context.no_compile_layers[layer_name] @@ -217,16 +252,14 @@ def inject_kv_into_layer( layer = kv_cache[forward_context.virtual_engine] kv_cache = self.p2p_nccl_engine.recv_tensor( - request.request_id + "#" + layer_name, remote_address + transfer_id + "#" + layer_name, remote_address ) if kv_cache is None: - logger.warning("🚧kv_cache is None, %s", request.request_id) + logger.warning("🚧kv_cache is None, %s", transfer_id) continue - inject_kv_into_layer( - layer, kv_cache, request.block_ids, request.request_id - ) + inject_kv_into_layer(layer, kv_cache, request.block_ids, transfer_id) def wait_for_layer_load(self, layer_name: str) -> None: """Blocking until the KV for a specific layer is loaded into vLLM's @@ -297,13 +330,13 @@ def extract_kv_from_layer( connector_metadata = self._get_connector_metadata() assert isinstance(connector_metadata, P2pNcclConnectorMetadata) for request in connector_metadata.requests: - request_id = request.request_id - ip, port = self.parse_request_id(request_id, True) + transfer_id = request.transfer_id + ip, port = self.parse_request_id(transfer_id, True) remote_address = ip + ":" + str(port + self._rank) kv_cache = extract_kv_from_layer(kv_layer, request.block_ids) self.p2p_nccl_engine.send_tensor( - request_id + "#" + layer_name, kv_cache, remote_address + transfer_id + "#" + layer_name, kv_cache, remote_address ) def wait_for_save(self):