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 3, 19:00 UTC
[Bugfix][PD] Fix multi-node TP (TP>8) (#39907)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
May 13, 2026 at 07:20 UTC
71bcd02ef369866403c8a8bc03f7dc86e4f68f46
2 files changed
+29
-2
tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py
+13
-2
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from unittest.mock import MagicMock, patch
5
+
6
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa: E501
7
ExampleConnectorMetadata,
8
)
@@ -38,11 +40,20 @@ def test_kv_connector_mixin_clears_metadata():
40
vllm_config.kv_transfer_config.kv_role = "kv_both"
41
vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "unit"
42
41
- # Initialize the global connector instance
43
kv_cache_config = KVCacheConfig(
44
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
45
)
45
- ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
46
+ # Initialize the global connector instance.
47
+ # kv_transfer init now syncs engine_id across TP, so unit tests need
48
+ # a minimal mocked TP group.
49
+ mock_tp_group = MagicMock()
50
+ mock_tp_group.broadcast_object.side_effect = lambda value, src=0: value
51
+
52
+ with patch(
53
+ "vllm.distributed.parallel_state.get_tp_group",
54
+ return_value=mock_tp_group,
55
+ ):
56
+ ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
57
58
try:
59
# Minimal scheduler output with empty metadata; mixin should still
vllm/distributed/kv_transfer/kv_transfer_state.py
+16
@@ -48,6 +48,20 @@ def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> boo
48
return isinstance(connector, KVConnectorBase_V1)
49
50
51
+def _sync_engine_id_across_tp(vllm_config: "VllmConfig") -> None:
52
+ """Broadcast engine_id from TP rank 0 so all workers in a
53
+ multi-node TP group share the same value."""
54
+ from vllm.distributed.parallel_state import (
55
+ get_tp_group,
56
+ )
57
+
58
+ assert vllm_config.kv_transfer_config is not None
59
+ synced_id = get_tp_group().broadcast_object(
60
+ vllm_config.kv_transfer_config.engine_id, src=0
61
+ )
62
+ vllm_config.kv_transfer_config.engine_id = synced_id
63
+
64
+
65
def ensure_kv_transfer_initialized(
66
vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"
67
) -> None:
@@ -64,6 +78,8 @@ def ensure_kv_transfer_initialized(
78
vllm_config.kv_transfer_config.is_kv_transfer_instance
79
and _KV_CONNECTOR_AGENT is None
80
):
81
+ _sync_engine_id_across_tp(vllm_config)
82
+
83
_KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector(
84
config=vllm_config,
85
role=KVConnectorRole.WORKER,