Skip to content

Commit 173201f

Browse files
authored
fix: vLLM xfer bug (#4745)
Signed-off-by: alec-flowers <[email protected]>
1 parent 8601ccd commit 173201f

File tree

1 file changed

+16
-1
lines changed

1 file changed

+16
-1
lines changed

lib/bindings/kvbm/python/kvbm/vllm_integration/connector/pd_connector.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,10 @@
44
from typing import TYPE_CHECKING, Optional, Type
55

66
from kvbm.vllm_integration.connector.dynamo_connector import DynamoConnector
7-
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
7+
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
8+
KVConnectorHandshakeMetadata,
9+
KVConnectorRole,
10+
)
811
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
912
MultiConnector,
1013
MultiKVConnectorMetadata,
@@ -84,6 +87,18 @@ def __init__(
8487
# Worker-side methods
8588
# ==============================
8689

90+
def set_xfer_handshake_metadata(
91+
self, metadata: dict[int, KVConnectorHandshakeMetadata]
92+
) -> None:
93+
"""
94+
Propagate handshake metadata to child connectors.
95+
96+
This is required for NIXL connector to start its handshake listener
97+
which decode workers connect to for KV transfer coordination.
98+
"""
99+
for c in self._connectors:
100+
c.set_xfer_handshake_metadata(metadata)
101+
87102
def bind_connector_metadata(self, connector_metadata: PdConnectorMetadata) -> None:
88103
assert isinstance(connector_metadata, PdConnectorMetadata)
89104
if connector_metadata.extra_async_saves:

0 commit comments

Comments
 (0)