A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 23:45 UTC
[NIXL] Add compatibility checking to NIXL KV connector handshake (#29503)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Dec 5, 2025 at 14:52 UTC
949a6a19d26f2179c0e5fc5d94f7d033e4c6a695
3 files changed
+380
-26
tests/v1/kv_connector/unit/test_nixl_connector.py
+213
-11
@@ -9,8 +9,10 @@ import textwrap
9
import time
10
import uuid
11
from collections import defaultdict
12
-from unittest.mock import patch
12
+from typing import Any
13
+from unittest.mock import MagicMock, patch
14
15
+import msgspec
16
import pytest
17
import ray
18
import torch
@@ -18,6 +20,7 @@ import torch
20
from vllm import LLM
21
from vllm.config import KVTransferConfig
22
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
23
+from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
24
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
25
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
26
MultiKVConnectorStats,
@@ -29,7 +32,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
32
NixlConnectorMetadata,
33
NixlConnectorScheduler,
34
NixlConnectorWorker,
35
+ NixlHandshakePayload,
36
NixlKVConnectorStats,
37
+ compute_nixl_compatibility_hash,
38
)
39
from vllm.distributed.kv_transfer.kv_transfer_state import (
40
ensure_kv_transfer_shutdown,
@@ -317,13 +322,19 @@ def test_kv_transfer_handshake(dist_init):
322
}
323
prefill_connector.register_kv_caches(kv_caches)
324
320
- # Simulate EngineCore initialization that would
321
- # gather connector metadata from all workers, the scheduler connector
322
- # expects metadata to be in dict[int, KVConnectorHandshakeMetadata],
323
- # where the first key is the dp_rank, the second key is the tp_rank.
324
- metadata = {0: prefill_connector.get_handshake_metadata()}
325
+ # Simulate EngineCore initialization that would gather connector
326
+ # metadata from all workers
327
+ metadata = prefill_connector.get_handshake_metadata()
328
+
329
+ # metadata is a NixlHandshakePayload, decode it to get NixlAgentMetadata
330
+ decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
331
+ expected_agent_metadata = decoder.decode(metadata.agent_metadata_bytes)
332
+
333
+ # The scheduler connector expects metadata to be in
334
+ # dict[int, KVConnectorHandshakeMetadata], where the first key is
335
+ # the dp_rank, the second key is the tp_rank.
336
scheduler_connector = scheduler.get_kv_connector()
326
- scheduler_connector.set_xfer_handshake_metadata(metadata)
337
+ scheduler_connector.set_xfer_handshake_metadata({0: metadata})
338
339
# Simulate a request that finishes prefill, which returns
340
# corresponding NixlConnectorMetadata for decode instance.
@@ -362,9 +373,9 @@ def test_kv_transfer_handshake(dist_init):
373
)
374
375
received_metadata = mock_add_remote_agent.call_args.args
376
+ assert received_metadata[0] == expected_agent_metadata
377
assert received_metadata[1] == 0 # remote_tp_rank
378
assert received_metadata[2] == 1 # remote_tp_size
367
- assert metadata[0] == received_metadata[0]
379
380
# Need to shutdown the background thread to release NIXL side channel port
381
scheduler_connector.shutdown()
@@ -403,7 +414,6 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
414
device_id=0,
415
num_blocks=1,
416
block_lens=self.block_len_per_layer,
406
- attn_backend_name=self.backend_name,
417
# `self.kv_cache_layout` is only forced to HND when vllm engine
418
# is started. We mock HND here.
419
kv_cache_layout="HND",
@@ -651,7 +661,6 @@ class TestNixlHandshake:
661
device_id=0,
662
num_blocks=1,
663
block_lens=worker.block_len_per_layer,
654
- attn_backend_name=worker.backend_name,
664
kv_cache_layout=mismatched_layout,
665
block_size=worker.block_size,
666
)
@@ -706,7 +715,6 @@ class TestNixlHandshake:
715
num_blocks=1,
716
# prefill TP=1, decode TP=2, remote block_lens is double to local
717
block_lens=[i * 2 for i in worker.block_len_per_layer],
709
- attn_backend_name=worker.backend_name,
718
kv_cache_layout="HND",
719
block_size=worker.block_size,
720
)
@@ -1168,6 +1176,9 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1176
mock_wrapper_instance = mock_nixl_wrapper.return_value
1177
connector.connector_worker.nixl_wrapper = mock_wrapper_instance
1178
1179
+ # Appease NixlHandshakePayload encoding with some bytes
1180
+ mock_wrapper_instance.get_agent_metadata.return_value = b"fake_agent_metadata"
1181
+
1182
# Reassure the shutdown() check that the thread is terminated
1183
mock_thread.return_value.is_alive.return_value = False
1184
@@ -1534,3 +1545,194 @@ def test_transfer_setup_failure_returns_finished(dist_init):
1545
# ensure request appears in get_finished
1546
_, done_recving = connector.get_finished(finished_req_ids=set())
1547
assert request_id in done_recving
1548
+
1549
+
1550
+@pytest.mark.parametrize(
1551
+ "mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat",
1552
+ [
1553
+ ("vllm_version", {}, {"vllm_version": "0.6.1"}, True, True),
1554
+ ("nixl_connector_version", {}, {"connector_version": 37}, True, True),
1555
+ ("model_name", {"model": "facebook/opt-350m"}, {}, True, True),
1556
+ ("dtype", {"dtype": "bfloat16"}, {}, True, True),
1557
+ ("cache_dtype", {"cache_dtype": "fp8"}, {}, True, True),
1558
+ ("num_kv_heads", {"hf_overrides": {"num_key_value_heads": 8}}, {}, True, True),
1559
+ (
1560
+ "num_hidden_layers",
1561
+ {"hf_overrides": {"num_hidden_layers": 24}},
1562
+ {},
1563
+ True,
1564
+ True,
1565
+ ),
1566
+ ("hidden_size", {"hf_overrides": {"hidden_size": 1536}}, {}, True, True),
1567
+ ("block_size", {"block_size": 8}, {}, False, True),
1568
+ ("matching_config", {}, {}, False, True),
1569
+ ("escape_hatch", {"model": "facebook/opt-350m"}, {}, False, False),
1570
+ ],
1571
+)
1572
+@patch(
1573
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1574
+ FakeNixlWrapper,
1575
+)
1576
+def test_compatibility_hash_validation(
1577
+ dist_init,
1578
+ mismatch_type,
1579
+ config_overrides,
1580
+ version_override,
1581
+ should_fail,
1582
+ enforce_handshake_compat,
1583
+):
1584
+ """
1585
+ Test NIXL compatibility hash validation during handshake.
1586
+
1587
+ Parameters:
1588
+ mismatch_type: description of what is being tested
1589
+ config_overrides: dict of config to override for the remote instance
1590
+ version_override: version dict e.g. {"vllm_version": "0.6.1"}
1591
+ should_fail: whether the handshake should fail
1592
+ enforce_handshake_compat: whether to enforce compatibility checking
1593
+ """
1594
+ local_vllm_config = create_vllm_config(
1595
+ model="facebook/opt-125m",
1596
+ block_size=16,
1597
+ kv_connector_extra_config={
1598
+ "enforce_handshake_compat": enforce_handshake_compat
1599
+ },
1600
+ )
1601
+ decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
1602
+ decode_worker = decode_connector.connector_worker
1603
+
1604
+ remote_config_params: dict[str, Any] = {
1605
+ "model": "facebook/opt-125m",
1606
+ "block_size": 16,
1607
+ **config_overrides,
1608
+ }
1609
+ remote_vllm_config = create_vllm_config(**remote_config_params)
1610
+
1611
+ with contextlib.ExitStack() as stack:
1612
+ if "vllm_version" in version_override:
1613
+ stack.enter_context(
1614
+ patch("vllm.__version__", version_override["vllm_version"])
1615
+ )
1616
+ elif "connector_version" in version_override:
1617
+ stack.enter_context(
1618
+ patch.object(
1619
+ nixl_connector,
1620
+ "NIXL_CONNECTOR_VERSION",
1621
+ version_override["connector_version"],
1622
+ )
1623
+ )
1624
+ remote_hash = compute_nixl_compatibility_hash(
1625
+ remote_vllm_config, decode_worker.backend_name
1626
+ )
1627
+
1628
+ prefill_block_size = config_overrides.get("block_size", 16)
1629
+ prefill_metadata = NixlAgentMetadata(
1630
+ engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1631
+ agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1632
+ kv_caches_base_addr=[0],
1633
+ device_id=0,
1634
+ num_blocks=1,
1635
+ block_lens=[4096 * prefill_block_size], # slot_size * block_size
1636
+ kv_cache_layout="HND",
1637
+ block_size=prefill_block_size,
1638
+ )
1639
+ handshake_payload = NixlHandshakePayload(
1640
+ compatibility_hash=remote_hash,
1641
+ agent_metadata_bytes=msgspec.msgpack.encode(prefill_metadata),
1642
+ )
1643
+
1644
+ # Mock ZMQ socket to return our handshake payload
1645
+ mock_socket = MagicMock()
1646
+ mock_socket.recv.return_value = msgspec.msgpack.encode(handshake_payload)
1647
+
1648
+ # Mock add_remote_agent to avoid actual NIXL operations
1649
+ # Patch zmq_ctx to return our mock socket
1650
+ with (
1651
+ patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
1652
+ patch.object(nixl_connector, "zmq_ctx") as mock_zmq_ctx,
1653
+ ):
1654
+ mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
1655
+
1656
+ if should_fail:
1657
+ with pytest.raises(RuntimeError, match="compatibility hash mismatch"):
1658
+ decode_worker._nixl_handshake(
1659
+ host="localhost",
1660
+ port=1234,
1661
+ remote_tp_size=1,
1662
+ expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1663
+ )
1664
+ else:
1665
+ result = decode_worker._nixl_handshake(
1666
+ host="localhost",
1667
+ port=1234,
1668
+ remote_tp_size=1,
1669
+ expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1670
+ )
1671
+ # Verify handshake returned agent mapping
1672
+ assert isinstance(result, dict)
1673
+ assert len(result) == 1
1674
+
1675
+
1676
+@pytest.mark.parametrize(
1677
+ "error_scenario",
1678
+ [
1679
+ "handshake_decode_error",
1680
+ "handshake_validation_error",
1681
+ "metadata_decode_error",
1682
+ "metadata_validation_error",
1683
+ ],
1684
+)
1685
+@patch(
1686
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1687
+ FakeNixlWrapper,
1688
+)
1689
+def test_handshake_decode_errors(dist_init, error_scenario):
1690
+ """
1691
+ Test that msgspec decode errors are properly handled during handshake.
1692
+
1693
+ Tests both DecodeError and ValidationError for both decoders:
1694
+ - NixlHandshakePayload decoder
1695
+ - NixlAgentMetadata decoder
1696
+ """
1697
+ local_vllm_config = create_vllm_config(
1698
+ model="facebook/opt-125m",
1699
+ block_size=16,
1700
+ )
1701
+ decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
1702
+ decode_worker = decode_connector.connector_worker
1703
+
1704
+ if error_scenario == "handshake_decode_error":
1705
+ msg_bytes = b"this is not valid msgpack data"
1706
+ elif error_scenario == "handshake_validation_error":
1707
+ msg_bytes = msgspec.msgpack.encode({"wrong_field": "value"})
1708
+ elif error_scenario == "metadata_decode_error":
1709
+ valid_handshake = NixlHandshakePayload(
1710
+ compatibility_hash=decode_worker.compat_hash,
1711
+ agent_metadata_bytes=b"invalid msgpack for metadata",
1712
+ )
1713
+ msg_bytes = msgspec.msgpack.encode(valid_handshake)
1714
+
1715
+ elif error_scenario == "metadata_validation_error":
1716
+ valid_handshake = NixlHandshakePayload(
1717
+ compatibility_hash=decode_worker.compat_hash,
1718
+ agent_metadata_bytes=msgspec.msgpack.encode({"missing": "fields"}),
1719
+ )
1720
+ msg_bytes = msgspec.msgpack.encode(valid_handshake)
1721
+ else:
1722
+ raise AssertionError(f"{error_scenario} not a valid scenario")
1723
+
1724
+ mock_socket = MagicMock()
1725
+ mock_socket.recv.return_value = msg_bytes
1726
+ with (
1727
+ patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
1728
+ patch.object(nixl_connector, "zmq_ctx") as mock_zmq_ctx,
1729
+ ):
1730
+ mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
1731
+
1732
+ with pytest.raises(RuntimeError):
1733
+ decode_worker._nixl_handshake(
1734
+ host="localhost",
1735
+ port=1234,
1736
+ remote_tp_size=1,
1737
+ expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1738
+ )
tests/v1/kv_connector/unit/utils.py
+8
-2
@@ -90,13 +90,18 @@ def create_vllm_config(
90
max_model_len: int = 10000,
91
enable_chunked_prefill: bool = True,
92
enable_permute_local_kv: bool = False,
93
+ kv_connector_extra_config: dict[str, Any] | None = None,
94
+ dtype: str = "float16",
95
+ cache_dtype: str = "auto",
96
+ hf_overrides: dict[str, Any] | None = None,
97
) -> VllmConfig:
98
"""Initialize VllmConfig For Testing."""
99
model_config = ModelConfig(
100
model=model,
101
trust_remote_code=True,
98
- dtype="float16",
102
+ dtype=dtype,
103
seed=42,
104
+ hf_overrides=hf_overrides or {},
105
)
106
scheduler_config = SchedulerConfig(
107
max_num_seqs=max_num_seqs,
@@ -110,13 +115,14 @@ def create_vllm_config(
115
block_size=block_size,
116
gpu_memory_utilization=0.9,
117
swap_space=0,
113
- cache_dtype="auto",
118
+ cache_dtype=cache_dtype,
119
enable_prefix_caching=True,
120
)
121
kv_transfer_config = KVTransferConfig(
122
kv_connector="NixlConnector",
123
kv_role="kv_both",
124
enable_permute_local_kv=enable_permute_local_kv,
125
+ kv_connector_extra_config=kv_connector_extra_config or {},
126
)
127
return VllmConfig(
128
scheduler_config=scheduler_config,
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+159
-13
@@ -59,6 +59,21 @@ Transfer = tuple[int, float] # (xfer_handle, start_time)
59
EngineId = str
60
ReqId = str
61
62
+#
63
+# NIXL Connector Version
64
+#
65
+# Increment this version whenever there is an incompatible change to:
66
+# - NixlAgentMetadata schema
67
+# - kv_transfer_params schema or semantics
68
+# - NIXL transfer protocol or wire format
69
+# - KV cache memory layout or block organization
70
+# - Any other change that breaks P/D interoperability
71
+#
72
+# Version History:
73
+# 1: Initial version with compatibility checking
74
+#
75
+NIXL_CONNECTOR_VERSION: int = 1
76
+
77
GET_META_MSG = b"get_meta_msg"
78
79
logger = init_logger(__name__)
@@ -97,18 +112,95 @@ _NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices())
112
113
114
@dataclass
100
-class NixlAgentMetadata(KVConnectorHandshakeMetadata):
115
+class NixlAgentMetadata:
116
engine_id: str
117
agent_metadata: bytes
118
kv_caches_base_addr: list[int]
119
device_id: int
120
num_blocks: int
121
block_lens: list[int]
107
- attn_backend_name: str
122
kv_cache_layout: str
123
block_size: int
124
125
126
+@dataclass
127
+class NixlHandshakePayload(KVConnectorHandshakeMetadata):
128
+ """
129
+ Wrapper for NIXL handshake sent over the wire.
130
+
131
+ Enables two-phase decoding for graceful compatibility checking:
132
+ 1. Decode NixlHandshakePayload to get compatibility_hash
133
+ 2. Compute local hash and compare
134
+ 3. Only if hashes match, decode agent_metadata_bytes
135
+
136
+ This prevents decoder errors when NixlAgentMetadata schema is
137
+ incompatible, allowing graceful failure with clear error message.
138
+ """
139
+
140
+ compatibility_hash: str
141
+ agent_metadata_bytes: bytes # NixlAgentMetadata encoded
142
+
143
+
144
+def compute_nixl_compatibility_hash(
145
+ vllm_config: VllmConfig, attn_backend_name: str
146
+) -> str:
147
+ """
148
+ Compute compatibility hash for NIXL KV transfer.
149
+
150
+ Hash only the factors that affect whether two NIXL instances can
151
+ successfully transfer KV cache data.
152
+
153
+ Factors included:
154
+ - vLLM version and NIXL connector version
155
+ - Model architecture (name, dtype, KV heads, layers)
156
+ - KV cache format (dtype, sliding window)
157
+ - Attention backend
158
+
159
+ Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout
160
+ are validated at runtime in _validate_remote_agent_handshake and are not
161
+ included in this hash to support heterogeneous deployments.
162
+
163
+ Note - the set of factors are likely to evolve significantly over
164
+ time to be more or less permissive.
165
+
166
+ Returns:
167
+ SHA-256 hex digest
168
+ """
169
+ from vllm import __version__ as vllm_version
170
+ from vllm.config.utils import hash_factors
171
+
172
+ model_config = vllm_config.model_config
173
+ cache_config = vllm_config.cache_config
174
+
175
+ factors = {
176
+ # Version compatibility
177
+ "vllm_version": vllm_version,
178
+ "nixl_connector_version": NIXL_CONNECTOR_VERSION,
179
+ # Model architecture - affects KV cache shape
180
+ "model": model_config.model,
181
+ "dtype": str(model_config.dtype),
182
+ "num_kv_heads": model_config.get_total_num_kv_heads(),
183
+ "head_size": model_config.get_head_size(),
184
+ "num_hidden_layers": model_config.get_total_num_hidden_layers(),
185
+ # Attention backend and KV cache dtype affect memory layout
186
+ "attn_backend_name": attn_backend_name,
187
+ "cache_dtype": str(cache_config.cache_dtype),
188
+ }
189
+
190
+ compat_hash = hash_factors(factors)
191
+ logger.info(
192
+ "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, "
193
+ "cache_dtype=%s, attn_backend=%s)",
194
+ compat_hash,
195
+ factors["model"],
196
+ factors["dtype"],
197
+ factors["num_kv_heads"],
198
+ factors["cache_dtype"],
199
+ attn_backend_name,
200
+ )
201
+ return compat_hash
202
+
203
+
204
@dataclass
205
class ReqMeta:
206
local_block_ids: list[int]
@@ -396,14 +488,14 @@ class NixlConnectorScheduler:
488
encoded_data: dict[int, bytes] = {}
489
encoder = msgspec.msgpack.Encoder()
490
for tp_rank, rank_metadata in metadata.items():
399
- if not isinstance(rank_metadata, NixlAgentMetadata):
491
+ if not isinstance(rank_metadata, NixlHandshakePayload):
492
raise ValueError(
401
- "NixlConnectorScheduler expects NixlAgentMetadata for "
493
+ "NixlConnectorScheduler expects NixlHandshakePayload for "
494
"handshake metadata."
495
)
496
encoded_data[tp_rank] = encoder.encode(rank_metadata)
497
logger.debug(
406
- "Tp rank %d: encoded NixlAgentMetadata size: %s bytes",
498
+ "Tp rank %d: encoded NixlHandshakePayload size: %s bytes",
499
tp_rank,
500
str(len(encoded_data[tp_rank])),
501
)
@@ -794,7 +886,7 @@ class NixlConnectorWorker:
886
self._failed_recv_reqs: set[ReqId] = set()
887
888
# Handshake metadata of this worker for NIXL transfers.
797
- self.xfer_handshake_metadata: NixlAgentMetadata | None = None
889
+ self.xfer_handshake_metadata: NixlHandshakePayload | None = None
890
# Background thread for initializing new NIXL handshakes.
891
self._handshake_initiation_executor = ThreadPoolExecutor(
892
# NIXL is not guaranteed to be thread-safe, limit 1 worker.
@@ -829,6 +921,13 @@ class NixlConnectorWorker:
921
logger.debug("Detected attention backend %s", self.backend_name)
922
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
923
924
+ self.compat_hash = compute_nixl_compatibility_hash(
925
+ self.vllm_config, self.backend_name
926
+ )
927
+ self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
928
+ "enforce_handshake_compat", True
929
+ )
930
+
931
self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
932
self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
933
# With heterogeneous TP, P must wait for all assigned D TP workers to
@@ -877,14 +976,58 @@ class NixlConnectorWorker:
976
# Set receive timeout to 5 seconds to avoid hanging on dead server
977
sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds
978
sock.send(msg)
880
- metadata_bytes = sock.recv()
881
- decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
882
- metadata = decoder.decode(metadata_bytes)
979
+ handshake_bytes = sock.recv()
980
+
981
+ # Decode handshake payload to get compatibility hash
982
+ handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)
983
+ try:
984
+ handshake_payload = handshake_decoder.decode(handshake_bytes)
985
+ except (msgspec.DecodeError, msgspec.ValidationError) as e:
986
+ raise RuntimeError(
987
+ f"Failed to decode NixlHandshakePayload. This likely indicates "
988
+ f"an incompatibility between connector version. Error: {e}"
989
+ ) from e
990
+
991
got_metadata_time = time.perf_counter()
992
logger.debug(
993
"NIXL handshake: get metadata took: %s", got_metadata_time - start_time
994
)
995
996
+ # Check compatibility hash BEFORE decoding agent metadata
997
+ if (
998
+ self.enforce_compat_hash
999
+ and handshake_payload.compatibility_hash != self.compat_hash
1000
+ ):
1001
+ raise RuntimeError(
1002
+ f"NIXL compatibility hash mismatch. "
1003
+ f"Local: {self.compat_hash}, "
1004
+ f"Remote: {handshake_payload.compatibility_hash}. "
1005
+ f"Prefill and decode instances have incompatible configurations. "
1006
+ f"This may be due to: different vLLM versions, models, dtypes, "
1007
+ f"KV cache layouts, attention backends, etc. "
1008
+ f"Both instances must use identical configurations."
1009
+ f"Disable this check using "
1010
+ f'--kv-transfer-config \'{{"kv_connector_extra_config": '
1011
+ f'{{"enforce_handshake_compat": false}}}}\''
1012
+ )
1013
+
1014
+ logger.info(
1015
+ "NIXL compatibility check passed (hash: %s)",
1016
+ handshake_payload.compatibility_hash,
1017
+ )
1018
+
1019
+ # Decode agent metadata
1020
+ metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
1021
+ try:
1022
+ metadata = metadata_decoder.decode(
1023
+ handshake_payload.agent_metadata_bytes
1024
+ )
1025
+ except (msgspec.DecodeError, msgspec.ValidationError) as e:
1026
+ # This should not happen if hash matched
1027
+ raise RuntimeError(
1028
+ f"Failed to decode NixlAgentMetadata. Error: {e}"
1029
+ ) from e
1030
+
1031
# Ensure engine id matches.
1032
if metadata.engine_id != expected_engine_id:
1033
raise RuntimeError(
@@ -1175,19 +1318,24 @@ class NixlConnectorWorker:
1318
assert len(self.block_window_per_layer) == self.num_layers
1319
1320
# After KV Caches registered, listen for new connections.
1178
- self.xfer_handshake_metadata = NixlAgentMetadata(
1321
+ agent_metadata = NixlAgentMetadata(
1322
engine_id=self.engine_id,
1323
agent_metadata=self.nixl_wrapper.get_agent_metadata(),
1324
kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],
1325
device_id=self.device_id,
1326
num_blocks=self.num_blocks,
1327
block_lens=self.block_len_per_layer,
1185
- attn_backend_name=self.backend_name,
1328
kv_cache_layout=self.kv_cache_layout
1329
if not self.use_host_buffer
1330
else self.host_buffer_kv_cache_layout,
1331
block_size=self.block_size,
1332
)
1333
+ # Wrap metadata in payload with hash for defensive decoding
1334
+ encoder = msgspec.msgpack.Encoder()
1335
+ self.xfer_handshake_metadata = NixlHandshakePayload(
1336
+ compatibility_hash=self.compat_hash,
1337
+ agent_metadata_bytes=encoder.encode(agent_metadata),
1338
+ )
1339
1340
def register_local_xfer_handler(
1341
self,
@@ -1402,8 +1550,6 @@ class NixlConnectorWorker:
1550
remote_engine_id = nixl_agent_meta.engine_id
1551
1552
assert self._tp_size[remote_engine_id] == remote_tp_size
1405
- # TODO We may eventually want to skip enforcing the same attn backend.
1406
- assert nixl_agent_meta.attn_backend_name == self.backend_name
1553
1554
tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
1555
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(