@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 96229fa99e
Fork 0

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 1, 21:15 UTC

[KVConnector][1/N] PP-aware handshake aggregation and intermediate-PP output plumbing (#43720)

Signed-off-by: zixi-qi <zixi@inferact.ai>

qizixi committed Jun 4, 2026 at 22:04 UTC 96229fa99e9942be9d6d808c9e07b81d2c256f74
9 files changed +459 -26
tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py new
+232
@@ -0,0 +1,232 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from typing import Any
6 +
7 +import pytest
8 +
9 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
10 + KVConnectorBase_V1,
11 + KVConnectorHandshakeMetadata,
12 +)
13 +from vllm.v1.engine import core as engine_core_module
14 +
15 +pytestmark = pytest.mark.cpu_test
16 +
17 +
18 +class _Metadata(KVConnectorHandshakeMetadata):
19 + pass
20 +
21 +
22 +class _FakeExecutor:
23 + handshake_metadata_src: (
24 + list[dict[tuple[int, int], KVConnectorHandshakeMetadata] | None] | None
25 + )
26 + last_instance: "_FakeExecutor | None" = None
27 +
28 + def __init__(
29 + self,
30 + vllm_config: Any,
31 + ) -> None:
32 + del vllm_config
33 + self.handshake_metadata = self.handshake_metadata_src
34 + self.handshake_calls = 0
35 + _FakeExecutor.last_instance = self
36 +
37 + def get_kv_connector_handshake_metadata(
38 + self,
39 + ) -> list[dict[tuple[int, int], KVConnectorHandshakeMetadata] | None] | None:
40 + self.handshake_calls += 1
41 + return self.handshake_metadata
42 +
43 + def init_kv_output_aggregator(self, connector: KVConnectorBase_V1) -> None:
44 + pass
45 +
46 +
47 +def _run_engine_core_handshake(
48 + monkeypatch: pytest.MonkeyPatch,
49 + connector: KVConnectorBase_V1,
50 + *,
51 + handshake_metadata: (
52 + list[dict[tuple[int, int], KVConnectorHandshakeMetadata] | None] | None
53 + ),
54 +) -> _FakeExecutor:
55 + class _FakeScheduler:
56 + def __init__(self, **kwargs: Any) -> None:
57 + self.connector = connector
58 +
59 + def get_kv_connector(self) -> KVConnectorBase_V1:
60 + return connector
61 +
62 + _FakeExecutor.handshake_metadata_src = handshake_metadata
63 + _FakeExecutor.last_instance = None
64 +
65 + monkeypatch.setattr("vllm.plugins.load_general_plugins", lambda: None)
66 + monkeypatch.setattr(
67 + engine_core_module.EngineCore,
68 + "_initialize_kv_caches",
69 + lambda self, vllm_config: SimpleNamespace(kv_cache_groups=[object()]),
70 + )
71 + monkeypatch.setattr(
72 + engine_core_module,
73 + "StructuredOutputManager",
74 + lambda vllm_config: object(),
75 + )
76 + monkeypatch.setattr(
77 + engine_core_module,
78 + "resolve_kv_cache_block_sizes",
79 + lambda kv_cache_config, vllm_config: (16, 16),
80 + )
81 + monkeypatch.setattr(
82 + engine_core_module,
83 + "MULTIMODAL_REGISTRY",
84 + SimpleNamespace(engine_receiver_cache_from_config=lambda vllm_config: None),
85 + )
86 + monkeypatch.setattr(engine_core_module, "freeze_gc_heap", lambda: None)
87 + monkeypatch.setattr(
88 + engine_core_module, "maybe_attach_gc_debug_callback", lambda: None
89 + )
90 + monkeypatch.setattr(engine_core_module, "enable_envs_cache", lambda: None)
91 + monkeypatch.setattr(engine_core_module, "get_hash_fn_by_name", lambda name: None)
92 + monkeypatch.setattr(engine_core_module, "init_none_hash", lambda hash_fn: None)
93 + monkeypatch.setattr(
94 + engine_core_module, "get_request_block_hasher", lambda *args: None
95 + )
96 +
97 + vllm_config = SimpleNamespace(
98 + parallel_config=SimpleNamespace(data_parallel_rank_local=0),
99 + scheduler_config=SimpleNamespace(
100 + get_scheduler_cls=lambda: _FakeScheduler,
101 + enable_chunked_prefill=False,
102 + async_scheduling=False,
103 + ),
104 + speculative_config=None,
105 + ec_transfer_config=None,
106 + max_concurrent_batches=1,
107 + model_config=SimpleNamespace(runner_type="generate"),
108 + cache_config=SimpleNamespace(
109 + enable_prefix_caching=False,
110 + prefix_caching_hash_algo="builtin",
111 + ),
112 + )
113 +
114 + engine_core_module.EngineCore(vllm_config, _FakeExecutor, log_stats=False)
115 + assert _FakeExecutor.last_instance is not None
116 + return _FakeExecutor.last_instance
117 +
118 +
119 +class _LegacyConnector(KVConnectorBase_V1):
120 + def __init__(self) -> None:
121 + self.legacy_metadata: dict[int, KVConnectorHandshakeMetadata] | None = None
122 +
123 + def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:
124 + pass
125 +
126 + def wait_for_layer_load(self, layer_name: str) -> None:
127 + pass
128 +
129 + def save_kv_layer(
130 + self,
131 + layer_name: str,
132 + kv_layer: Any,
133 + attn_metadata: Any,
134 + **kwargs: Any,
135 + ) -> None:
136 + pass
137 +
138 + def wait_for_save(self) -> None:
139 + pass
140 +
141 + def get_num_new_matched_tokens(
142 + self, request: Any, num_computed_tokens: int
143 + ) -> tuple[int | None, bool]:
144 + return 0, False
145 +
146 + def update_state_after_alloc(
147 + self, request: Any, blocks: Any, num_external_tokens: int
148 + ) -> None:
149 + pass
150 +
151 + def build_connector_meta(self, scheduler_output: Any) -> Any:
152 + raise NotImplementedError
153 +
154 + def set_xfer_handshake_metadata(
155 + self, metadata: dict[int, KVConnectorHandshakeMetadata]
156 + ) -> None:
157 + self.legacy_metadata = metadata
158 +
159 +
160 +class _PPAwareConnector(_LegacyConnector):
161 + def __init__(self) -> None:
162 + super().__init__()
163 + self.pp_aware_metadata: (
164 + dict[tuple[int, int], KVConnectorHandshakeMetadata] | None
165 + ) = None
166 +
167 + def set_xfer_handshake_metadata_pp_aware(
168 + self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
169 + ) -> None:
170 + self.pp_aware_metadata = metadata
171 +
172 +
173 +def test_engine_unwraps_handshake_metadata_for_legacy_connector(
174 + monkeypatch: pytest.MonkeyPatch,
175 +) -> None:
176 + """Engine core always asks workers for `(pp_rank, tp_rank)`-keyed metadata,
177 + then unwraps to `{tp_rank: metadata}` for a connector that has not opted
178 + into PP-aware handshake (single-PP producer, all `pp_rank == 0`)."""
179 + metadata_0 = _Metadata()
180 + metadata_1 = _Metadata()
181 + connector = _LegacyConnector()
182 +
183 + executor = _run_engine_core_handshake(
184 + monkeypatch,
185 + connector,
186 + handshake_metadata=[
187 + {(0, 0): metadata_0},
188 + None,
189 + {(0, 1): metadata_1},
190 + ],
191 + )
192 +
193 + assert executor.handshake_calls == 1
194 + assert connector.legacy_metadata == {0: metadata_0, 1: metadata_1}
195 +
196 +
197 +def test_engine_rejects_pp_producer_for_legacy_connector(
198 + monkeypatch: pytest.MonkeyPatch,
199 +) -> None:
200 + """A connector that has not opted into PP-aware handshake must not silently
201 + drop metadata from `pp_rank > 0`; engine core init raises instead."""
202 + connector = _LegacyConnector()
203 +
204 + with pytest.raises(ValueError, match="does not support PP-disaggregated"):
205 + _run_engine_core_handshake(
206 + monkeypatch,
207 + connector,
208 + handshake_metadata=[{(0, 0): _Metadata()}, {(1, 0): _Metadata()}],
209 + )
210 +
211 +
212 +def test_engine_passes_handshake_metadata_through_for_pp_aware_connector(
213 + monkeypatch: pytest.MonkeyPatch,
214 +) -> None:
215 + """A PP-aware connector receives the full `(pp_rank, tp_rank)`-keyed dict
216 + unchanged."""
217 + metadata_0 = _Metadata()
218 + metadata_1 = _Metadata()
219 + connector = _PPAwareConnector()
220 +
221 + executor = _run_engine_core_handshake(
222 + monkeypatch,
223 + connector,
224 + handshake_metadata=[{(0, 0): metadata_0}, {(1, 0): metadata_1}],
225 + )
226 +
227 + assert executor.handshake_calls == 1
228 + assert connector.legacy_metadata is None
229 + assert connector.pp_aware_metadata == {
230 + (0, 0): metadata_0,
231 + (1, 0): metadata_1,
232 + }
tests/v1/kv_connector/unit/test_multi_connector.py
+5 -4
@@ -261,11 +261,12 @@ def test_multi_example_connector_consistency():
261 storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"])
262 storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"])
263 # First event is bind_gpu_block_pool from initialization, then
264 - # set_xfer_handshake_metadata, then on_new_request when the request is enqueued,
265 - # then get_num_new_matched_tokens and update_state_after_alloc from generate().
264 + # set_xfer_handshake_metadata_pp_aware, then on_new_request when the request is
265 + # enqueued, then get_num_new_matched_tokens and update_state_after_alloc from
266 + # generate().
267 assert storage1_scheduler_events[:6] == [
268 "bind_gpu_block_pool",
268 - "set_xfer_handshake_metadata",
269 + "set_xfer_handshake_metadata_pp_aware",
270 "on_new_request",
271 "get_num_new_matched_tokens 0",
272 "update_state_after_alloc num_blocks=[0] 0",
@@ -285,7 +286,7 @@ def test_multi_example_connector_consistency():
286 ]
287 assert storage2_scheduler_events[:6] == [
288 "bind_gpu_block_pool",
288 - "set_xfer_handshake_metadata",
289 + "set_xfer_handshake_metadata_pp_aware",
290 "on_new_request",
291 "get_num_new_matched_tokens 0",
292 "update_state_after_alloc num_blocks=[0] 0",
tests/v1/kv_connector/unit/test_transfer_topology_sharded.py new
+146
@@ -0,0 +1,146 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.distributed.kv_transfer.kv_connector.utils import (
7 + EngineTransferInfo,
8 + TransferTopology,
9 +)
10 +
11 +pytestmark = pytest.mark.cpu_test
12 +
13 +
14 +class _FakeAttentionBackend:
15 + @staticmethod
16 + def get_kv_cache_shape(
17 + num_blocks: int,
18 + block_size: int,
19 + num_kv_heads: int,
20 + head_size: int,
21 + ) -> tuple[int, int, int, int, int]:
22 + return (2, num_blocks, num_kv_heads, block_size, head_size)
23 +
24 +
25 +def _make_topology(
26 + *,
27 + tp_rank: int = 1,
28 + tp_size: int = 4,
29 + total_num_kv_heads: int = 8,
30 +) -> TransferTopology:
31 + return TransferTopology(
32 + tp_rank=tp_rank,
33 + tp_size=tp_size,
34 + block_size=16,
35 + engine_id="local-engine",
36 + is_mla=False,
37 + is_mamba=False,
38 + total_num_kv_heads=total_num_kv_heads,
39 + attn_backends=[_FakeAttentionBackend],
40 + )
41 +
42 +
43 +def test_legacy_register_remote_engine_uses_pp_rank_zero() -> None:
44 + topology = _make_topology()
45 + info = EngineTransferInfo(
46 + remote_tp_size=2,
47 + remote_block_len=1024,
48 + remote_block_size=16,
49 + remote_physical_blocks_per_logical=1,
50 + )
51 +
52 + registered = topology.register_remote_engine("remote-engine", info)
53 +
54 + assert registered == info
55 + assert registered.remote_pp_rank == 0
56 + assert topology.get_engine_info("remote-engine") == info
57 + assert topology._engines[("remote-engine", 0)] == info
58 + assert topology.target_remote_ranks("remote-engine") == [0]
59 +
60 +
61 +def test_register_remote_engine_stores_pp_ranks_separately() -> None:
62 + topology = _make_topology(tp_rank=0, tp_size=2)
63 +
64 + info_0 = EngineTransferInfo(
65 + remote_tp_size=2,
66 + remote_block_len=1024,
67 + remote_block_size=16,
68 + remote_physical_blocks_per_logical=1,
69 + remote_pp_rank=0,
70 + start_layer=0,
71 + end_layer=16,
72 + )
73 + info_1 = EngineTransferInfo(
74 + remote_tp_size=1,
75 + remote_block_len=512,
76 + remote_block_size=8,
77 + remote_physical_blocks_per_logical=2,
78 + remote_pp_rank=1,
79 + start_layer=16,
80 + end_layer=32,
81 + )
82 +
83 + registered_0 = topology.register_remote_engine("remote-engine", info_0)
84 + registered_1 = topology.register_remote_engine("remote-engine", info_1)
85 +
86 + assert registered_0 == info_0
87 + assert registered_1 == info_1
88 + assert topology.get_engine_info("remote-engine") == info_0
89 + assert topology.get_engine_info("remote-engine", 0) == info_0
90 + assert topology.get_engine_info("remote-engine", 1) == info_1
91 + assert set(topology._engines) == {
92 + ("remote-engine", 0),
93 + ("remote-engine", 1),
94 + }
95 +
96 +
97 +def test_helpers_use_requested_pp_rank() -> None:
98 + topology = _make_topology(tp_rank=1, tp_size=2, total_num_kv_heads=2)
99 + topology.register_remote_engine(
100 + "remote-engine",
101 + EngineTransferInfo(
102 + remote_tp_size=1,
103 + remote_block_len=1024,
104 + remote_block_size=16,
105 + remote_physical_blocks_per_logical=1,
106 + remote_pp_rank=0,
107 + start_layer=0,
108 + end_layer=8,
109 + ),
110 + )
111 + topology.register_remote_engine(
112 + "remote-engine",
113 + EngineTransferInfo(
114 + remote_tp_size=4,
115 + remote_block_len=1024,
116 + remote_block_size=16,
117 + remote_physical_blocks_per_logical=1,
118 + remote_pp_rank=1,
119 + start_layer=8,
120 + end_layer=16,
121 + ),
122 + )
123 +
124 + assert not topology.is_kv_replicated("remote-engine", 0)
125 + assert topology.is_kv_replicated("remote-engine", 1)
126 + assert topology.replicates_kv_cache("remote-engine", 1)
127 + assert topology.target_remote_ranks("remote-engine", 0) == [0]
128 + assert topology.target_remote_ranks("remote-engine", 1) == [2, 3]
129 + assert "remote_pp=1" in topology.describe("remote-engine", 1)
130 +
131 +
132 +def test_engine_info_fields_have_backward_compatible_defaults() -> None:
133 + topology = _make_topology()
134 + info = EngineTransferInfo(
135 + remote_tp_size=2,
136 + remote_block_len=1024,
137 + remote_block_size=16,
138 + remote_physical_blocks_per_logical=1,
139 + )
140 +
141 + registered = topology.register_remote_engine("remote-engine", info)
142 +
143 + assert topology.get_engine_info("remote-engine") == registered
144 + assert registered.remote_pp_rank == 0
145 + assert registered.start_layer == 0
146 + assert registered.end_layer == 0
vllm/distributed/kv_transfer/kv_connector/utils.py
+37 -15
@@ -368,7 +368,7 @@ def get_current_attn_backend(
368 class EngineTransferInfo:
369 """Common per-remote-engine transfer state, computed at handshake.
370
371 - Stored per ``engine_id`` inside ``TransferTopology._engines``.
371 + Stored per ``(engine_id, pp_rank)`` inside ``TransferTopology._engines``.
372 """
373
374 remote_tp_size: int
@@ -382,6 +382,15 @@ class EngineTransferInfo:
382 remote_physical_blocks_per_logical: int
383 """Physical blocks per logical block."""
384
385 + remote_pp_rank: int = 0
386 + """Remote producer PP rank for this engine."""
387 +
388 + start_layer: int = 0
389 + """Global index of the first layer owned by this PP rank."""
390 +
391 + end_layer: int = 0
392 + """Exclusive global index after the last layer owned by this PP rank."""
393 +
394
395 # ---- Transfer topology ----
396
@@ -403,7 +412,7 @@ class TransferTopology:
412 def __post_init__(self):
413 self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size)
414
406 - self._engines: dict[EngineId, EngineTransferInfo] = {}
415 + self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {}
416
417 # Figure out whether the first dimension of the cache is K/V
418 # or num_blocks.
@@ -461,13 +470,16 @@ class TransferTopology:
470 f"Cannot register local engine {self.engine_id} as remote. "
471 f"Local identity is set via __init__ params."
472 )
464 - if remote_engine_id in self._engines:
465 - return self._engines[remote_engine_id]
466 - self._engines[remote_engine_id] = info
473 + engine_key = (remote_engine_id, info.remote_pp_rank)
474 + if engine_key in self._engines:
475 + return self._engines[engine_key]
476 + self._engines[engine_key] = info
477 return info
478
469 - def get_engine_info(self, remote_engine_id: EngineId) -> EngineTransferInfo:
470 - return self._engines[remote_engine_id]
479 + def get_engine_info(
480 + self, remote_engine_id: EngineId, remote_pp_rank: int = 0
481 + ) -> EngineTransferInfo:
482 + return self._engines[(remote_engine_id, remote_pp_rank)]
483
484 # ============================================================
485 # Layout properties
@@ -528,15 +540,22 @@ class TransferTopology:
540 )
541 return self.block_size // remote_block_size
542
531 - def is_kv_replicated(self, remote_engine_id: EngineId) -> bool:
543 + def is_kv_replicated(
544 + self, remote_engine_id: EngineId, remote_pp_rank: int = 0
545 + ) -> bool:
546 """Whether the KV cache is replicated across TP workers due to the
547 number of TP workers being greater than the number of KV heads.
548 """
535 - return self._engines[remote_engine_id].remote_tp_size > self.total_num_kv_heads
549 + return (
550 + self._engines[(remote_engine_id, remote_pp_rank)].remote_tp_size
551 + > self.total_num_kv_heads
552 + )
553
537 - def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool:
554 + def replicates_kv_cache(
555 + self, remote_engine_id: EngineId, remote_pp_rank: int = 0
556 + ) -> bool:
557 # MLA is always replicated as the hidden dim can't be split.
539 - return self.is_mla or self.is_kv_replicated(remote_engine_id)
558 + return self.is_mla or self.is_kv_replicated(remote_engine_id, remote_pp_rank)
559
560 @property
561 def local_replicates_kv_cache(self) -> bool:
@@ -555,12 +574,14 @@ class TransferTopology:
574 abs_ratio = -tp_ratio
575 return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
576
558 - def target_remote_ranks(self, remote_engine_id: EngineId) -> list[int]:
577 + def target_remote_ranks(
578 + self, remote_engine_id: EngineId, remote_pp_rank: int = 0
579 + ) -> list[int]:
580 """Get the remote TP rank(s) that the current local TP rank will
581 read from. When remote tp_size > local tp_size, reads from
582 multiple remote ranks.
583 """
563 - info = self._engines[remote_engine_id]
584 + info = self._engines[(remote_engine_id, remote_pp_rank)]
585 tp_ratio = self.tp_ratio(info.remote_tp_size)
586 if tp_ratio > 0:
587 return [self.tp_rank // tp_ratio]
@@ -593,15 +614,16 @@ class TransferTopology:
614 # Regular case: backends like FA register K/V in separate regions
615 return cache if self.split_k_and_v else [cache]
616
596 - def describe(self, remote_engine_id: EngineId) -> str:
617 + def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
618 """One-line summary of transfer config for logging."""
598 - info = self._engines[remote_engine_id]
619 + info = self._engines[(remote_engine_id, remote_pp_rank)]
620 return (
621 f"TransferTopology("
622 f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, "
623 f"num_kv_heads={self.total_num_kv_heads if not self.is_mla else 1}, "
624 f"local_tp={self.tp_size}, "
625 f"remote_tp={info.remote_tp_size}, "
626 + f"remote_pp={remote_pp_rank}, "
627 f"local_rank={self.tp_rank}, "
628 f"remote_block_len={info.remote_block_len})"
629 )
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+17
@@ -644,6 +644,23 @@ class KVConnectorBase_V1(ABC):
644 """
645 return None
646
647 + def set_xfer_handshake_metadata_pp_aware(
648 + self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
649 + ) -> None:
650 + """
651 + Set handshake metadata keyed by (pp_rank, tp_rank).
652 + - Default implementation assumes pp_rank is always 0
653 + - PP-aware connectors override this to consume all PP producer shards.
654 + """
655 + if any(pp_rank != 0 for pp_rank, _ in metadata):
656 + raise ValueError(
657 + f"{type(self).__name__} received pp_rank > 0 handshake metadata "
658 + "but does not support PP-disaggregated KV transfer."
659 + )
660 + self.set_xfer_handshake_metadata(
661 + {tp_rank: meta for (_, tp_rank), meta in metadata.items()}
662 + )
663 +
664 @classmethod
665 def build_prom_metrics(
666 cls,
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+6
@@ -471,6 +471,12 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
471 for c in self._connectors:
472 c.set_xfer_handshake_metadata(metadata)
473
474 + def set_xfer_handshake_metadata_pp_aware(
475 + self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
476 + ) -> None:
477 + for c in self._connectors:
478 + c.set_xfer_handshake_metadata_pp_aware(metadata)
479 +
480 def _aggregate_request_finished(
481 self,
482 request: "Request",
vllm/v1/engine/core.py
+3 -3
@@ -177,13 +177,13 @@ class EngineCore:
177
178 if xfer_handshake_metadata:
179 # xfer_handshake_metadata is list of dicts from workers
180 - # Each dict already has structure {tp_rank: metadata}
180 + # Each dict already has structure {(pp_rank, tp_rank): metadata}
181 # Merge all worker dicts into a single dict
182 - content: dict[int, Any] = {}
182 + content: dict[tuple[int, int], Any] = {}
183 for worker_dict in xfer_handshake_metadata:
184 if worker_dict is not None:
185 content.update(worker_dict)
186 - kv_connector.set_xfer_handshake_metadata(content)
186 + kv_connector.set_xfer_handshake_metadata_pp_aware(content)
187
188 # Setup batch queue for pipeline parallelism.
189 # Batch queue for scheduled batches. This enables us to asynchronously
vllm/v1/executor/abstract.py
+1 -1
@@ -203,7 +203,7 @@ class Executor(ABC):
203
204 def get_kv_connector_handshake_metadata(
205 self,
206 - ) -> list[dict[int, KVConnectorHandshakeMetadata]]:
206 + ) -> list[dict[tuple[int, int], KVConnectorHandshakeMetadata]]:
207 return self.collective_rpc("get_kv_connector_handshake_metadata")
208
209 @overload
vllm/v1/worker/gpu_worker.py
+12 -3
@@ -34,6 +34,9 @@ from vllm.distributed.kv_transfer import (
34 get_kv_transfer_group,
35 has_kv_transfer_group,
36 )
37 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
38 + KVConnectorHandshakeMetadata,
39 +)
40 from vllm.distributed.parallel_state import (
41 Handle,
42 get_pp_group,
@@ -513,8 +516,13 @@ class Worker(WorkerBase):
516
517 return int(self.available_kv_cache_memory_bytes)
518
516 - def get_kv_connector_handshake_metadata(self) -> dict | None:
517 - """Get KV connector metadata from this worker if available."""
519 + def get_kv_connector_handshake_metadata(
520 + self,
521 + ) -> dict[tuple[int, int], KVConnectorHandshakeMetadata] | None:
522 + """Get KV connector metadata from this worker if available.
523 +
524 + Returned dict is keyed by `(pp_rank, tp_rank)`.
525 + """
526
527 if not has_kv_transfer_group():
528 return None
@@ -525,8 +533,9 @@ class Worker(WorkerBase):
533 if (metadata := connector.get_handshake_metadata()) is None:
534 return None
535
536 + pp_rank = get_pp_group().rank_in_group
537 tp_rank = get_tp_group().rank_in_group
529 - return {tp_rank: metadata}
538 + return {(pp_rank, tp_rank): metadata}
539
540 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
541 return self.model_runner.get_kv_cache_spec()