A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix NixlConnector handshake block_len validation for GQA-replicated KV heads (#45879)
Signed-off-by: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Co-authored-by: waynehacking8 <waynehacking8@gmail.com>
Bryan Shan committed
Jun 17, 2026 at 15:11 UTC
0d339cf13551bd4e6c8f4c1cfce1f47c806aa5e3
2 files changed
+141
-15
tests/v1/kv_connector/unit/test_nixl_connector.py
+125
-8
@@ -517,17 +517,17 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
517
assert expected_engine_id == self.REMOTE_ENGINE_ID
518
519
# Adjust remote block length metadata to satisfy heterogeneous TP
520
- # invariants enforced during handshake validation.
520
+ # invariants enforced during handshake validation. Use per-rank
521
+ # head ratio (not tp_ratio) to account for GQA replication capping.
522
remote_block_lens = list(self.block_len_per_layer)
523
tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
523
- if remote_tp_size > self.world_size:
524
- # P TP > D TP case, block_len of remote is smaller
524
+ total_kv = self.transfer_topo.total_num_kv_heads
525
+ local_heads = self.transfer_topo.local_physical_heads
526
+ remote_heads = max(1, total_kv // remote_tp_size)
527
+ if remote_tp_size != self.world_size:
528
remote_block_lens = [
526
- block_len // (-tp_ratio) for block_len in remote_block_lens
527
- ]
528
- elif remote_tp_size < self.world_size:
529
- remote_block_lens = [
530
- block_len * tp_ratio for block_len in remote_block_lens
529
+ block_len * remote_heads // local_heads
530
+ for block_len in remote_block_lens
531
]
532
533
# When remote tp_size > local tp_size, handshake with multiple
@@ -1144,6 +1144,123 @@ class TestNixlHandshake:
1144
with pytest.raises(AssertionError):
1145
worker2.add_remote_agent(bad_meta, remote_tp_size=1)
1146
1147
+ @patch(
1148
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1149
+ FakeNixlWrapper,
1150
+ )
1151
+ def test_handshake_validates_gqa_replicated_block_len(
1152
+ self, default_vllm_config, dist_init
1153
+ ):
1154
+ """Regression test for #45330.
1155
+
1156
+ When tp_size > total_num_kv_heads, GQA replication caps per-rank
1157
+ KV heads at 1, so block_len stops scaling with 1/tp. With 8 KV
1158
+ heads and D_TP=16 pulling from P_TP=8, both sides hold one head
1159
+ per rank and report the *same* block_len; the old validation
1160
+ expected local_block_len * tp_ratio and rejected the valid
1161
+ handshake.
1162
+ """
1163
+ vllm_config = create_vllm_config()
1164
+
1165
+ with patch(
1166
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1167
+ return_value=16,
1168
+ ):
1169
+ connector = NixlConnector(
1170
+ vllm_config,
1171
+ KVConnectorRole.WORKER,
1172
+ make_kv_cache_config(block_size=16),
1173
+ )
1174
+ connector.connector_worker = FakeNixlConnectorWorker(
1175
+ vllm_config, connector.engine_id, hand_shake_latency=0
1176
+ )
1177
+ worker = connector.connector_worker
1178
+
1179
+ worker.transfer_topo.total_num_kv_heads = 8
1180
+ worker.transfer_topo.local_physical_heads = 1
1181
+ worker.kv_cache_layout = "HND"
1182
+
1183
+ worker.slot_size_per_layer = [4096]
1184
+ worker.block_len_per_layer = [4096 * worker.block_size]
1185
+ worker.num_blocks = 1
1186
+ worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1187
+
1188
+ # Remote P with TP=8 also has 1 head/rank -> identical
1189
+ # block_len despite tp_ratio == 2.
1190
+ meta = NixlAgentMetadata(
1191
+ engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1192
+ agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1193
+ kv_caches_base_addr=[0],
1194
+ device_id=0,
1195
+ num_blocks=1,
1196
+ block_lens=list(worker.block_len_per_layer),
1197
+ kv_cache_layout="HND",
1198
+ block_size=worker.block_size,
1199
+ ssm_sizes=(0, 0),
1200
+ attn_backend_name=worker.backend_name,
1201
+ physical_blocks_per_logical_kv_block=1,
1202
+ )
1203
+
1204
+ # Must validate cleanly (used to raise AssertionError).
1205
+ worker.add_remote_agent(meta, remote_tp_size=8)
1206
+
1207
+ @patch(
1208
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1209
+ FakeNixlWrapper,
1210
+ )
1211
+ def test_handshake_rejects_wrong_block_len_without_gqa_replication(
1212
+ self, default_vllm_config, dist_init
1213
+ ):
1214
+ """Ensure the head-ratio validation still rejects genuinely wrong
1215
+ block_lens when GQA replication is NOT in effect (32 KV heads,
1216
+ D_TP=4, P_TP=2: head_ratio=4, both sides have >1 head/rank).
1217
+ """
1218
+ vllm_config = create_vllm_config()
1219
+
1220
+ with patch(
1221
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1222
+ return_value=4,
1223
+ ):
1224
+ connector = NixlConnector(
1225
+ vllm_config,
1226
+ KVConnectorRole.WORKER,
1227
+ make_kv_cache_config(block_size=16),
1228
+ )
1229
+ connector.connector_worker = FakeNixlConnectorWorker(
1230
+ vllm_config, connector.engine_id, hand_shake_latency=0
1231
+ )
1232
+ worker = connector.connector_worker
1233
+
1234
+ worker.transfer_topo.total_num_kv_heads = 32
1235
+ worker.transfer_topo.local_physical_heads = 8 # 32 // 4
1236
+ worker.kv_cache_layout = "HND"
1237
+
1238
+ slot_size = 4096
1239
+ worker.slot_size_per_layer = [slot_size]
1240
+ worker.block_len_per_layer = [slot_size * worker.block_size]
1241
+ worker.num_blocks = 1
1242
+ worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1243
+
1244
+ # Remote P_TP=2 has 16 heads/rank -> head_ratio = 16/8 = 2.
1245
+ # Correct remote block_len = local * 2. Send local * 1
1246
+ # (wrong) to verify rejection.
1247
+ bad_meta = NixlAgentMetadata(
1248
+ engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1249
+ agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1250
+ kv_caches_base_addr=[0],
1251
+ device_id=0,
1252
+ num_blocks=1,
1253
+ block_lens=list(worker.block_len_per_layer),
1254
+ kv_cache_layout="HND",
1255
+ block_size=worker.block_size,
1256
+ ssm_sizes=(0, 0),
1257
+ attn_backend_name=worker.backend_name,
1258
+ physical_blocks_per_logical_kv_block=1,
1259
+ )
1260
+
1261
+ with pytest.raises(AssertionError):
1262
+ worker.add_remote_agent(bad_meta, remote_tp_size=2)
1263
+
1264
1265
# NOTE: resource cleanup in mp backend is a bit finicky, so the order in which
1266
# we put here is important. First run ray, it will clean up the resources, then
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+16
-7
@@ -1599,7 +1599,9 @@ class NixlBaseConnectorWorker:
1599
# Per-region block_len validation enforcing the P/D invariant.
1600
# REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)
1601
# only allow the number of blocks to differ; SPLIT regions scale with
1602
- # tp_ratio. Mamba uses the ssm_sizes counterpart, so skip block_len here.
1602
+ # the per-rank KV head ratio rather than the raw tp_ratio, because GQA
1603
+ # replication caps per-rank heads at 1 when tp > total_kv_heads
1604
+ # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here.
1605
if not self._has_mamba:
1606
assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
1607
"Number of KV layers must match between prefill and decode"
@@ -1607,6 +1609,9 @@ class NixlBaseConnectorWorker:
1609
model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(
1610
remote_engine_id
1611
)
1612
+ total_kv_heads = self.transfer_topo.total_num_kv_heads
1613
+ local_heads = self.transfer_topo.local_physical_heads
1614
+ remote_heads = max(1, total_kv_heads // remote_tp_size)
1615
for i, local_len in enumerate(self.block_len_per_layer):
1616
replicated = model_replicated or self._is_region_replicated(i)
1617
remote_len = nixl_agent_meta.block_lens[i]
@@ -1617,9 +1622,13 @@ class NixlBaseConnectorWorker:
1622
f"remote={remote_len}, bsr={block_size_ratio})."
1623
)
1624
elif tp_ratio > 0:
1620
- assert remote_len == (local_len * tp_ratio) // block_size_ratio, (
1625
+ assert (
1626
+ remote_len
1627
+ == (local_len * remote_heads // local_heads) // block_size_ratio
1628
+ ), (
1629
f"SPLIT region {i}: remote P KV block_len {remote_len} "
1622
- f"must equal local {local_len} * tp_ratio {tp_ratio} "
1630
+ f"must equal local {local_len} * remote_heads "
1631
+ f"{remote_heads} // local_heads {local_heads} "
1632
f"// block_size_ratio {block_size_ratio}."
1633
)
1634
else:
@@ -1627,10 +1636,10 @@ class NixlBaseConnectorWorker:
1636
"Different local/remote block sizes are not supported "
1637
"when P TP > D TP."
1638
)
1630
- assert remote_len == local_len // (-tp_ratio), (
1631
- f"SPLIT region {i}: remote P KV block_len "
1632
- f"{remote_len} must equal local {local_len} "
1633
- f"// |tp_ratio| {-tp_ratio}."
1639
+ assert remote_len == local_len * remote_heads // local_heads, (
1640
+ f"SPLIT region {i}: remote P KV block_len {remote_len} "
1641
+ f"must equal local {local_len} * remote_heads "
1642
+ f"{remote_heads} // local_heads {local_heads}."
1643
)
1644
1645
# TP workers that handhshake with same remote have same #blocks.