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

A high-throughput and memory-efficient inference and serving engine for LLMs

[KV Connector][Mooncake] Pipeline-parallel support for PD-disaggregated serving with Mooncake connector (#44528)

Signed-off-by: hanhan.hank <hanhan.hank@bytedance.com> Signed-off-by: Hank Han <hanhan7630@outlook.com>

Hank Han committed Jun 16, 2026 at 19:35 UTC d53f4593cec92f9114ed6574dec24b0d69deb65e
4 files changed +680 -34
tests/v1/kv_connector/unit/test_mooncake_connector.py
+503 -5
@@ -4,22 +4,29 @@
4 import asyncio
5 import contextlib
6 import time
7 +from types import SimpleNamespace
8 from unittest.mock import AsyncMock, MagicMock, patch
9
10 import pytest
11 import torch
12 import zmq.asyncio
13
14 +from vllm import envs
15 from vllm.config import set_current_vllm_config
16 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
17 KVConnectorRole,
18 MooncakeConnector,
19 MooncakeConnectorMetadata,
20 + MooncakeConnectorWorker,
21 MooncakeXferMetadata,
22 MooncakeXferResponse,
23 MooncakeXferResponseStatus,
24 PullReqMeta,
25 SendBlockMeta,
26 + TransferRegion,
27 + _align_transfer_regions,
28 + get_mooncake_bootstrap_addr,
29 + should_launch_bootstrap_server,
30 )
31 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
32 MooncakeBootstrapServer,
@@ -57,6 +64,272 @@ class FakeMooncakeWrapper:
64 return 0
65
66
67 +def test_align_transfer_regions_uses_layer_name_occurrences():
68 + """Repeated layer names should align by occurrence order."""
69 +
70 + local_regions = [
71 + TransferRegion(
72 + layer_name="model.layers.1.self_attn",
73 + layer_index=1,
74 + base_addr=0x1000,
75 + block_len=256,
76 + kv_block_len=128,
77 + ),
78 + TransferRegion(
79 + layer_name="model.layers.1.self_attn",
80 + layer_index=1,
81 + base_addr=0x1100,
82 + block_len=256,
83 + kv_block_len=128,
84 + ),
85 + ]
86 + remote_regions = [
87 + TransferRegion(
88 + layer_name="model.layers.0.self_attn",
89 + layer_index=0,
90 + base_addr=0xA000,
91 + block_len=256,
92 + kv_block_len=128,
93 + ),
94 + TransferRegion(
95 + layer_name="model.layers.1.self_attn",
96 + layer_index=1,
97 + base_addr=0xB000,
98 + block_len=256,
99 + kv_block_len=128,
100 + ),
101 + TransferRegion(
102 + layer_name="model.layers.1.self_attn",
103 + layer_index=1,
104 + base_addr=0xB100,
105 + block_len=256,
106 + kv_block_len=128,
107 + ),
108 + ]
109 +
110 + aligned_local, aligned_remote, err = _align_transfer_regions(
111 + local_regions, remote_regions
112 + )
113 +
114 + assert err is None
115 + assert [r.base_addr for r in aligned_local] == [0x1000, 0x1100]
116 + assert [r.base_addr for r in aligned_remote] == [0xB000, 0xB100]
117 +
118 +
119 +@pytest.mark.asyncio
120 +async def test_build_transfer_params_separates_prefill_pp_layers():
121 + """Each producer PP stage should send only its registered layer shard."""
122 +
123 + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
124 + worker.async_zmq_ctx = MagicMock()
125 + worker.is_kv_consumer = True
126 + worker.is_kv_producer = True
127 + worker.tp_rank = 0
128 + worker.tp_size = 1
129 + worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False)
130 +
131 + block_len = 256
132 + remote_regions = [
133 + TransferRegion(
134 + layer_name=f"model.layers.{layer_index}.self_attn",
135 + layer_index=layer_index,
136 + base_addr=base_addr,
137 + block_len=block_len,
138 + kv_block_len=block_len,
139 + )
140 + for layer_index, base_addr in [
141 + (0, 0xA000),
142 + (1, 0xB000),
143 + (2, 0xC000),
144 + (3, 0xD000),
145 + ]
146 + ]
147 + producer_pp_regions = {
148 + 0: [
149 + TransferRegion(
150 + layer_name="model.layers.0.self_attn",
151 + layer_index=0,
152 + base_addr=0x1000,
153 + block_len=block_len,
154 + kv_block_len=block_len,
155 + ),
156 + TransferRegion(
157 + layer_name="model.layers.1.self_attn",
158 + layer_index=1,
159 + base_addr=0x2000,
160 + block_len=block_len,
161 + kv_block_len=block_len,
162 + ),
163 + ],
164 + 1: [
165 + TransferRegion(
166 + layer_name="model.layers.2.self_attn",
167 + layer_index=2,
168 + base_addr=0x3000,
169 + block_len=block_len,
170 + kv_block_len=block_len,
171 + ),
172 + TransferRegion(
173 + layer_name="model.layers.3.self_attn",
174 + layer_index=3,
175 + base_addr=0x4000,
176 + block_len=block_len,
177 + kv_block_len=block_len,
178 + ),
179 + ],
180 + }
181 + expected_by_pp_rank = {
182 + 0: {
183 + "layers": [0, 1],
184 + "src_ptrs": [0x1000 + 10 * block_len, 0x2000 + 10 * block_len],
185 + "dst_ptrs": [0xA000 + 20 * block_len, 0xB000 + 20 * block_len],
186 + },
187 + 1: {
188 + "layers": [2, 3],
189 + "src_ptrs": [0x3000 + 10 * block_len, 0x4000 + 10 * block_len],
190 + "dst_ptrs": [0xC000 + 20 * block_len, 0xD000 + 20 * block_len],
191 + },
192 + }
193 +
194 + transfer_id = "xfer-pp-split"
195 + send_meta = SendBlockMeta(
196 + p_req_id="p-req-pp",
197 + transfer_id=transfer_id,
198 + local_block_ids=[[10, 11]],
199 + ready=asyncio.Event(),
200 + )
201 + xfer_meta = MooncakeXferMetadata(
202 + remote_hostname="consumer-host",
203 + remote_port=54321,
204 + remote_tp_size=1,
205 + remote_tp_rank=0,
206 + req_blocks={"d-req-pp": (transfer_id, [[20, 21]])},
207 + kv_caches_base_addr=[region.base_addr for region in remote_regions],
208 + block_lens=[region.block_len for region in remote_regions],
209 + registered_layer_names=[region.layer_name for region in remote_regions],
210 + registered_layer_indices=[region.layer_index for region in remote_regions],
211 + )
212 +
213 + for pp_rank, local_regions in producer_pp_regions.items():
214 + aligned_local, aligned_remote, err = _align_transfer_regions(
215 + local_regions, remote_regions
216 + )
217 +
218 + assert err is None
219 + assert [r.layer_index for r in aligned_local] == (
220 + expected_by_pp_rank[pp_rank]["layers"]
221 + )
222 + assert [r.layer_index for r in aligned_remote] == (
223 + expected_by_pp_rank[pp_rank]["layers"]
224 + )
225 +
226 + (
227 + src_ptrs,
228 + dst_ptrs,
229 + lengths,
230 + err_reqs,
231 + err_msg,
232 + ) = await worker._build_transfer_params(
233 + ready_reqs=[("d-req-pp", send_meta)],
234 + agent_meta=xfer_meta,
235 + local_regions=aligned_local,
236 + remote_regions=aligned_remote,
237 + )
238 +
239 + assert err_reqs == []
240 + assert err_msg is None
241 + assert src_ptrs == expected_by_pp_rank[pp_rank]["src_ptrs"]
242 + assert dst_ptrs == expected_by_pp_rank[pp_rank]["dst_ptrs"]
243 + assert lengths == [2 * block_len, 2 * block_len]
244 +
245 +
246 +@pytest.mark.asyncio
247 +async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
248 + monkeypatch,
249 +):
250 + """Producer sends its PP layer shard to the matching consumer layer address."""
251 +
252 + monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
253 + vllm_config = create_vllm_config(
254 + kv_connector="MooncakeConnector", kv_role="kv_producer"
255 + )
256 +
257 + with set_current_vllm_config(vllm_config), patch_worker_dependencies():
258 + prefill_connector = MooncakeConnector(
259 + vllm_config,
260 + KVConnectorRole.WORKER,
261 + _make_test_kv_cache_config(),
262 + )
263 + prefill_worker = prefill_connector.connector_worker
264 +
265 + block_len = 4096
266 + kv_half = block_len // 2
267 + prefill_worker.kv_caches_base_addr = [0x1000]
268 + prefill_worker.block_len_per_layer = [block_len]
269 + prefill_worker.registered_layer_names = ["model.layers.1.self_attn"]
270 + prefill_worker.registered_layer_indices = [1]
271 +
272 + class InlineSenderLoop:
273 + async def run_in_executor(self, executor, func, *args):
274 + return func(*args)
275 +
276 + origin_sender_loop = prefill_worker.sender_loop
277 + prefill_worker.sender_loop = InlineSenderLoop()
278 +
279 + transfer_id = "xfer-layer-align"
280 + send_meta = SendBlockMeta(
281 + p_req_id="p-req-layer-align",
282 + transfer_id=transfer_id,
283 + local_block_ids=[[10]],
284 + ready=asyncio.Event(),
285 + )
286 + prefill_worker.reqs_need_send[transfer_id] = send_meta
287 + send_meta.ready.set()
288 +
289 + xfer_meta = MooncakeXferMetadata(
290 + remote_hostname="consumer-host",
291 + remote_port=54321,
292 + remote_tp_size=1,
293 + remote_tp_rank=0,
294 + req_blocks={"d-req-layer-align": (transfer_id, [[20]])},
295 + kv_caches_base_addr=[0xA000, 0xB000],
296 + block_lens=[block_len, block_len],
297 + registered_layer_names=[
298 + "model.layers.0.self_attn",
299 + "model.layers.1.self_attn",
300 + ],
301 + registered_layer_indices=[0, 1],
302 + )
303 + mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
304 + mock_socket.send_multipart = AsyncMock()
305 + identity = b"consumer-layer-align"
306 +
307 + with patch.object(
308 + prefill_worker, "_send_blocks", return_value=0
309 + ) as mock_send_blocks:
310 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
311 +
312 + src_ptrs, dst_ptrs, lengths = mock_send_blocks.call_args[0][1:]
313 + assert src_ptrs == [
314 + 0x1000 + 10 * block_len,
315 + 0x1000 + 10 * block_len + kv_half,
316 + ]
317 + assert dst_ptrs == [
318 + 0xB000 + 20 * block_len,
319 + 0xB000 + 20 * block_len + kv_half,
320 + ]
321 + assert lengths == [kv_half, kv_half]
322 +
323 + sent_identity, sent_payload = mock_socket.send_multipart.call_args[0][0]
324 + assert sent_identity == identity
325 + response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
326 + assert response.status == MooncakeXferResponseStatus.FINISH
327 + assert response.ok_reqs == ["d-req-layer-align"]
328 +
329 + prefill_worker.sender_loop = origin_sender_loop
330 + prefill_worker.shutdown()
331 +
332 +
333 def test_basic_interface():
334 """Unit test for basic MooncakeConnector interface functionality."""
335
@@ -174,7 +447,7 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
447 assert response.status_code == 200
448 assert response.json() == {}
449
177 - # Register a worker
450 + # Register multiple PP workers from the same producer engine.
451 payload1 = {
452 "engine_id": "eng-1",
453 "dp_rank": 0,
@@ -187,7 +460,19 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
460 assert response.status_code == 200
461 assert response.json() == {"status": "ok"}
462
190 - # Query after registration
463 + payload2 = {
464 + "engine_id": "eng-1",
465 + "dp_rank": 0,
466 + "tp_rank": 0,
467 + "pp_rank": 1,
468 + "addr": "tcp://2.2.2.2:2222",
469 + }
470 + async with httpx.AsyncClient() as client:
471 + response = await client.post(f"{base_url}/register", json=payload2)
472 + assert response.status_code == 200
473 + assert response.json() == {"status": "ok"}
474 +
475 + # Query after registration should preserve the PP dimension.
476 async with httpx.AsyncClient() as client:
477 response = await client.get(f"{base_url}/query")
478 assert response.status_code == 200
@@ -195,6 +480,7 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
480 assert "0" in data
481 assert data["0"]["engine_id"] == "eng-1"
482 assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111"
483 + assert data["0"]["worker_addr"]["0"]["1"] == "tcp://2.2.2.2:2222"
484
485 # Test failure: re-registering the same worker
486 async with httpx.AsyncClient() as client:
@@ -216,6 +502,104 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
502 assert "Engine ID mismatch" in response.text
503
504
505 +def _make_bootstrap_vllm_config(
506 + *,
507 + local_engines_only: bool = False,
508 + data_parallel_rank_local: int = 0,
509 + data_parallel_index: int = 0,
510 + nnodes_within_dp: int = 1,
511 +) -> SimpleNamespace:
512 + return SimpleNamespace(
513 + parallel_config=SimpleNamespace(
514 + local_engines_only=local_engines_only,
515 + data_parallel_rank_local=data_parallel_rank_local,
516 + data_parallel_index=data_parallel_index,
517 + nnodes_within_dp=nnodes_within_dp,
518 + master_addr="model-parallel-master",
519 + data_parallel_master_ip="data-parallel-master",
520 + )
521 + )
522 +
523 +
524 +@pytest.mark.parametrize(
525 + (
526 + "tp_rank",
527 + "pp_rank",
528 + "local_engines_only",
529 + "data_parallel_rank_local",
530 + "data_parallel_index",
531 + "expected",
532 + ),
533 + [
534 + (1, 0, False, 0, 0, False),
535 + (0, 1, False, 0, 0, False),
536 + (0, 0, True, 0, 1, True),
537 + (0, 0, True, 1, 0, False),
538 + (0, 0, False, 0, 0, True),
539 + (0, 0, False, 0, 1, False),
540 + ],
541 + ids=[
542 + "nonzero_tp_rank",
543 + "nonzero_pp_rank",
544 + "local_engine_rank_zero",
545 + "local_engine_nonzero_rank",
546 + "internal_lb_first_dp_engine",
547 + "internal_lb_nonzero_dp_engine",
548 + ],
549 +)
550 +def test_should_launch_bootstrap_server_selects_single_owner(
551 + tp_rank: int,
552 + pp_rank: int,
553 + local_engines_only: bool,
554 + data_parallel_rank_local: int,
555 + data_parallel_index: int,
556 + expected: bool,
557 +):
558 + vllm_config = _make_bootstrap_vllm_config(
559 + local_engines_only=local_engines_only,
560 + data_parallel_rank_local=data_parallel_rank_local,
561 + data_parallel_index=data_parallel_index,
562 + )
563 + with (
564 + patch(
565 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
566 + "mooncake_connector.get_tensor_model_parallel_rank",
567 + return_value=tp_rank,
568 + ),
569 + patch(
570 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
571 + "mooncake_connector.get_pp_group"
572 + ) as mock_pp_group,
573 + ):
574 + mock_pp_group.return_value.rank_in_group = pp_rank
575 + assert should_launch_bootstrap_server(vllm_config) is expected
576 +
577 +
578 +@pytest.mark.parametrize(
579 + ("local_engines_only", "nnodes_within_dp", "expected_host"),
580 + [
581 + (True, 2, "127.0.0.1"),
582 + (False, 2, "model-parallel-master"),
583 + (False, 1, "data-parallel-master"),
584 + ],
585 + ids=["local_engine", "multi_node_tp_or_pp", "single_node_internal_lb"],
586 +)
587 +def test_get_mooncake_bootstrap_addr_selects_expected_host(
588 + local_engines_only: bool,
589 + nnodes_within_dp: int,
590 + expected_host: str,
591 +):
592 + vllm_config = _make_bootstrap_vllm_config(
593 + local_engines_only=local_engines_only,
594 + nnodes_within_dp=nnodes_within_dp,
595 + )
596 +
597 + assert get_mooncake_bootstrap_addr(vllm_config) == (
598 + expected_host,
599 + envs.VLLM_MOONCAKE_BOOTSTRAP_PORT,
600 + )
601 +
602 +
603 def test_scheduler_request_finished():
604 """
605 Tests the scheduler-side logic when a request finishes.
@@ -273,7 +657,6 @@ def patch_worker_dependencies():
657 patch(
658 "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_pp_group"
659 ) as mock_pp,
276 - patch("vllm.distributed.parallel_state.is_local_first_rank", return_value=True),
660 patch(
661 "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.should_launch_bootstrap_server",
662 return_value=False,
@@ -307,6 +690,93 @@ def patch_worker_dependencies():
690 }
691
692
693 +@pytest.mark.asyncio
694 +@pytest.mark.parametrize(
695 + ("local_pp_size", "local_pp_rank", "expected_addrs"),
696 + [
697 + (1, 0, ["tcp://producer-pp0:1234", "tcp://producer-pp1:1234"]),
698 + (2, 1, ["tcp://producer-pp1:1234"]),
699 + ],
700 + ids=["heterogeneous_pp_pulls_all_remote_pp", "matching_pp_pulls_same_rank"],
701 +)
702 +async def test_receive_kv_selects_remote_pp_workers(
703 + local_pp_size: int,
704 + local_pp_rank: int,
705 + expected_addrs: list[str],
706 +):
707 + """Decode workers should not hard-code producer pp_rank 0."""
708 +
709 + vllm_config = create_vllm_config(
710 + kv_connector="MooncakeConnector", kv_role="kv_consumer"
711 + )
712 +
713 + with set_current_vllm_config(vllm_config), patch_worker_dependencies():
714 + decode_connector = MooncakeConnector(
715 + vllm_config,
716 + KVConnectorRole.WORKER,
717 + _make_test_kv_cache_config(),
718 + )
719 + decode_worker = decode_connector.connector_worker
720 + decode_worker.pp_size = local_pp_size
721 + decode_worker.pp_rank = local_pp_rank
722 + decode_worker._remote_agents = {
723 + "p-engine": {
724 + 0: {
725 + 0: "tcp://producer-pp0:1234",
726 + 1: "tcp://producer-pp1:1234",
727 + }
728 + }
729 + }
730 + decode_worker._tp_size["p-engine"] = 1
731 +
732 + pull_metas = {
733 + "d-req-1": PullReqMeta(
734 + d_req_id="d-req-1",
735 + transfer_id="xfer-req-1",
736 + local_block_ids=[[100, 101]],
737 + remote_engine_id="p-engine",
738 + remote_bootstrap_addr="http://bootstrap:33333",
739 + )
740 + }
741 + seen_addrs: list[str] = []
742 +
743 + async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]):
744 + seen_addrs.append(worker_addr)
745 + for meta in metas.values():
746 + meta.pull_tasks_count -= 1
747 +
748 + with patch.object(
749 + decode_worker,
750 + "receive_kv_from_single_worker",
751 + side_effect=fake_receive,
752 + ):
753 + decode_worker.receive_kv("p-engine", pull_metas)
754 + await asyncio.sleep(0)
755 +
756 + assert seen_addrs == expected_addrs
757 + assert pull_metas["d-req-1"].pull_tasks_count == 0
758 + decode_worker.shutdown()
759 +
760 +
761 +def test_resolve_need_send_accounts_for_remote_tp_fanout():
762 + """Producer-side completion waits for every paired consumer TP pull."""
763 +
764 + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
765 + worker.async_zmq_ctx = MagicMock()
766 + worker.is_kv_consumer = True
767 + worker.is_kv_producer = True
768 + send_meta = SendBlockMeta(
769 + p_req_id="p-req-1",
770 + transfer_id="xfer-req-1",
771 + local_block_ids=[[1]],
772 + ready=asyncio.Event(),
773 + )
774 +
775 + worker.resolve_need_send(send_meta, remote_tp_ranks=[0, 1])
776 +
777 + assert send_meta.need_send == 2
778 +
779 +
780 @pytest.mark.asyncio
781 @patch(
782 "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
@@ -335,6 +805,8 @@ async def test_kv_producer(monkeypatch):
805 prefill_worker.kv_caches_base_addr = [0x1000]
806 block_len = 4096
807 prefill_worker.block_len_per_layer = [block_len]
808 + prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
809 + prefill_worker.registered_layer_indices = [0]
810
811 # Override loop to use current test loop
812 origin_sender_loop = prefill_worker.sender_loop
@@ -360,6 +832,8 @@ async def test_kv_producer(monkeypatch):
832 req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
833 kv_caches_base_addr=[0x2000],
834 block_lens=[block_len],
835 + registered_layer_names=["model.layers.0.self_attn"],
836 + registered_layer_indices=[0],
837 )
838
839 mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
@@ -506,6 +980,9 @@ async def test_kv_consumuer(monkeypatch):
980 )
981 decode_worker = decode_connector.connector_worker
982 decode_worker.kv_caches_base_addr = [0x1000]
983 + decode_worker.block_len_per_layer = [4096]
984 + decode_worker.registered_layer_names = ["model.layers.0.self_attn"]
985 + decode_worker.registered_layer_indices = [0]
986 decode_worker.rpc_port = 54321
987
988 # A request to pull data arrives.
@@ -547,6 +1024,10 @@ async def test_kv_consumuer(monkeypatch):
1024 assert sent_meta.remote_hostname == "127.0.0.1"
1025 assert sent_meta.remote_port == 54321
1026 assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [[100, 101]])
1027 + assert sent_meta.kv_caches_base_addr == [0x1000]
1028 + assert sent_meta.block_lens == [4096]
1029 + assert sent_meta.registered_layer_names == ["model.layers.0.self_attn"]
1030 + assert sent_meta.registered_layer_indices == [0]
1031
1032 # Verify internal state is updated correctly.
1033 assert "d-req-1" in decode_worker.finished_recving_reqs
@@ -626,7 +1107,10 @@ def test_register_kv_caches():
1107 )
1108 tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
1109 tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
629 - kv_caches = {"layer0": tensor1, "layer1": tensor2}
1110 + kv_caches = {
1111 + "model.layers.0.self_attn": tensor1,
1112 + "model.layers.1.self_attn": tensor2,
1113 + }
1114
1115 with patch.object(
1116 worker.engine, "batch_register_memory", return_value=0
@@ -643,6 +1127,8 @@ def test_register_kv_caches():
1127 assert len(worker.block_len_per_layer) == len(registered_ptrs)
1128 for bl in worker.block_len_per_layer:
1129 assert bl == tensor1.nbytes // tensor1.shape[0]
1130 + assert worker.registered_layer_names == list(kv_caches)
1131 + assert worker.registered_layer_indices == [0, 1]
1132
1133
1134 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
@@ -677,7 +1163,10 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
1163 mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)
1164 # Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size.
1165 eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16)
680 - kv_caches = {"mla_layer": mla_cache, "eagle_layer": eagle_cache}
1166 + kv_caches = {
1167 + "model.layers.0.mla_attn": mla_cache,
1168 + "model.layers.1.eagle_attn": eagle_cache,
1169 + }
1170
1171 with patch.object(
1172 worker.engine, "batch_register_memory", return_value=0
@@ -692,6 +1181,11 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
1181 mla_cache.nbytes // mla_cache.shape[0],
1182 eagle_cache.nbytes // eagle_cache.shape[0],
1183 ]
1184 + assert worker.registered_layer_names == [
1185 + "model.layers.0.mla_attn",
1186 + "model.layers.1.eagle_attn",
1187 + ]
1188 + assert worker.registered_layer_indices == [0, 1]
1189
1190
1191 @pytest.mark.asyncio
@@ -742,6 +1236,8 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1236
1237 prefill_worker.kv_caches_base_addr = [0x1000]
1238 prefill_worker.block_len_per_layer = [local_block_len]
1239 + prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
1240 + prefill_worker.registered_layer_indices = [0]
1241
1242 origin_sender_loop = prefill_worker.sender_loop
1243 prefill_worker.sender_loop = asyncio.get_event_loop()
@@ -787,6 +1283,8 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
1283 },
1284 kv_caches_base_addr=[0x2000],
1285 block_lens=[remote_block_len],
1286 + registered_layer_names=["model.layers.0.self_attn"],
1287 + registered_layer_indices=[0],
1288 )
1289
1290 mock_send_blocks.reset_mock()
tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
+20 -4
@@ -261,12 +261,20 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch):
261
262 local_regions = [
263 TransferRegion(
264 - base_addr=0x1000, block_len=block_len, kv_block_len=block_len
264 + layer_name="model.layers.0.self_attn",
265 + layer_index=0,
266 + base_addr=0x1000,
267 + block_len=block_len,
268 + kv_block_len=block_len,
269 ),
270 ]
271 remote_regions = [
272 TransferRegion(
269 - base_addr=0x2000, block_len=block_len, kv_block_len=block_len
273 + layer_name="model.layers.0.self_attn",
274 + layer_index=0,
275 + base_addr=0x2000,
276 + block_len=block_len,
277 + kv_block_len=block_len,
278 ),
279 ]
280
@@ -344,12 +352,20 @@ async def test_build_transfer_params_group_count_mismatch(monkeypatch):
352
353 local_regions = [
354 TransferRegion(
347 - base_addr=0x1000, block_len=block_len, kv_block_len=block_len
355 + layer_name="model.layers.0.self_attn",
356 + layer_index=0,
357 + base_addr=0x1000,
358 + block_len=block_len,
359 + kv_block_len=block_len,
360 ),
361 ]
362 remote_regions = [
363 TransferRegion(
352 - base_addr=0x2000, block_len=block_len, kv_block_len=block_len
364 + layer_name="model.layers.0.self_attn",
365 + layer_index=0,
366 + base_addr=0x2000,
367 + block_len=block_len,
368 + kv_block_len=block_len,
369 ),
370 ]
371
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+149 -24
@@ -43,10 +43,10 @@ from vllm.distributed.parallel_state import (
43 get_pp_group,
44 get_tensor_model_parallel_rank,
45 get_tensor_model_parallel_world_size,
46 - is_local_first_rank,
46 )
47 from vllm.forward_context import ForwardContext
48 from vllm.logger import init_logger
49 +from vllm.model_executor.models.utils import extract_layer_index
50 from vllm.platforms import current_platform
51 from vllm.utils.math_utils import cdiv
52 from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
@@ -80,6 +80,8 @@ TransferId = str # KV transfer coordination ID (shared by P/D)
80
81 @dataclass(frozen=True)
82 class TransferRegion:
83 + layer_name: str
84 + layer_index: int
85 base_addr: int
86 block_len: int
87 kv_block_len: int
@@ -109,18 +111,28 @@ def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:
111 def _expand_transfer_regions(
112 base_addrs: list[int],
113 block_lens: list[int],
114 + layer_names: list[str],
115 + layer_indices: list[int],
116 is_kv_layout_blocks_first: bool,
117 ) -> list[TransferRegion]:
118 """Expand registered KV tensors into the regions transferred by Mooncake."""
115 - assert len(base_addrs) == len(block_lens), (
116 - "Mooncake transfer regions require matching numbers of base addresses "
117 - f"and block lengths, got {len(base_addrs)} and {len(block_lens)}."
119 + assert (
120 + len(base_addrs) == len(block_lens) == len(layer_names) == len(layer_indices)
121 + ), (
122 + "Mooncake transfer regions require matching metadata lengths, got "
123 + f"base_addrs={len(base_addrs)}, block_lens={len(block_lens)}, "
124 + f"layer_names={len(layer_names)}, "
125 + f"layer_indices={len(layer_indices)}."
126 )
127 regions: list[TransferRegion] = []
120 - for base_addr, block_len in zip(base_addrs, block_lens):
128 + for base_addr, block_len, layer_name, layer_index in zip(
129 + base_addrs, block_lens, layer_names, layer_indices
130 + ):
131 kv_block_len = block_len // 2 if is_kv_layout_blocks_first else block_len
132 regions.append(
133 TransferRegion(
134 + layer_name=layer_name,
135 + layer_index=layer_index,
136 base_addr=base_addr,
137 block_len=block_len,
138 kv_block_len=kv_block_len,
@@ -129,6 +141,8 @@ def _expand_transfer_regions(
141 if is_kv_layout_blocks_first:
142 regions.append(
143 TransferRegion(
144 + layer_name=layer_name,
145 + layer_index=layer_index,
146 base_addr=base_addr + kv_block_len,
147 block_len=block_len,
148 kv_block_len=kv_block_len,
@@ -244,6 +258,62 @@ def _validate_asymmetric_region_lengths(
258 return None
259
260
261 +def _align_transfer_regions(
262 + local_regions: list[TransferRegion],
263 + remote_regions: list[TransferRegion],
264 +) -> tuple[list[TransferRegion], list[TransferRegion], str | None]:
265 + """Align KV transfer regions by registered layer-name occurrence.
266 +
267 + PP shards own different layer subsets. Positional matching is therefore
268 + wrong once producer and consumer have different PP layouts. Multiple
269 + registered transfer buffers for the same layer are represented by repeated
270 + layer names and matched by occurrence order.
271 + """
272 +
273 + def keyed_regions(
274 + regions: list[TransferRegion],
275 + ) -> list[tuple[tuple[str, int], TransferRegion]]:
276 + counts: dict[str, int] = defaultdict(int)
277 + keyed: list[tuple[tuple[str, int], TransferRegion]] = []
278 + for region in regions:
279 + occurrence = counts[region.layer_name]
280 + counts[region.layer_name] += 1
281 + keyed.append(((region.layer_name, occurrence), region))
282 + return keyed
283 +
284 + local_keyed = keyed_regions(local_regions)
285 + remote_keyed = keyed_regions(remote_regions)
286 + remote_by_key = dict(remote_keyed)
287 + aligned_local: list[TransferRegion] = []
288 + aligned_remote: list[TransferRegion] = []
289 + for key, local_region in local_keyed:
290 + remote_region = remote_by_key.get(key)
291 + if remote_region is None:
292 + return (
293 + [],
294 + [],
295 + (
296 + "Mooncake producer registered layer has no matching "
297 + f"consumer occurrence: {key[0]} occurrence {key[1]}."
298 + ),
299 + )
300 + if local_region.layer_index != remote_region.layer_index:
301 + return (
302 + [],
303 + [],
304 + (
305 + "Mooncake registered layer index mismatch for "
306 + f"{local_region.layer_name}: producer="
307 + f"{local_region.layer_index}, consumer="
308 + f"{remote_region.layer_index}."
309 + ),
310 + )
311 + aligned_local.append(local_region)
312 + aligned_remote.append(remote_region)
313 +
314 + return aligned_local, aligned_remote, None
315 +
316 +
317 def _get_tensor_dense_flag(tensor: torch.Tensor) -> bool | None:
318 is_dense = getattr(tensor, "is_non_overlapping_and_dense", None)
319 if callable(is_dense):
@@ -262,6 +332,8 @@ class MooncakeXferMetadata(
332 req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]]
333 kv_caches_base_addr: list[int]
334 block_lens: list[int]
335 + registered_layer_names: list[str] = msgspec.field(default_factory=list)
336 + registered_layer_indices: list[int] = msgspec.field(default_factory=list)
337
338
339 class MooncakeXferResponseStatus(IntEnum):
@@ -782,17 +854,15 @@ class MooncakeConnectorWorker:
854 self.tp_size = get_tensor_model_parallel_world_size()
855 self.num_blocks = 0
856 self.block_len_per_layer: list[int] = []
857 + self.registered_layer_names: list[str] = []
858 + self.registered_layer_indices: list[int] = []
859 self.seen_base_addresses: list[int] = []
860
861 assert (parallel_config := vllm_config.parallel_config)
862 dp_rank = parallel_config.data_parallel_index
863 dp_local_rank = parallel_config.data_parallel_rank_local
864 self.dp_rank = dp_local_rank if parallel_config.local_engines_only else dp_rank
791 - pp_size = vllm_config.parallel_config.pipeline_parallel_size
792 - if pp_size > 1:
793 - raise ValueError(
794 - "Mooncake Transfer Engine does not support pipeline parallelism yet."
795 - )
865 + self.pp_size = vllm_config.parallel_config.pipeline_parallel_size
866 self.pp_rank = get_pp_group().rank_in_group
867
868 self.kv_caches_base_addr: list[int] = []
@@ -1020,11 +1090,27 @@ class MooncakeConnectorWorker:
1090 await sock.send_multipart((identity, self._encoder.encode(response)))
1091 return
1092 local_regions = self._get_transfer_regions(
1023 - self.kv_caches_base_addr, self.block_len_per_layer
1093 + self.kv_caches_base_addr,
1094 + self.block_len_per_layer,
1095 + self.registered_layer_names,
1096 + self.registered_layer_indices,
1097 )
1098 remote_regions = self._get_transfer_regions(
1026 - meta.kv_caches_base_addr, meta.block_lens
1099 + meta.kv_caches_base_addr,
1100 + meta.block_lens,
1101 + meta.registered_layer_names,
1102 + meta.registered_layer_indices,
1103 + )
1104 + local_regions, remote_regions, align_err = _align_transfer_regions(
1105 + local_regions, remote_regions
1106 )
1107 + if align_err is not None:
1108 + response = MooncakeXferResponse(
1109 + status=MooncakeXferResponseStatus.ERROR,
1110 + err_msg=align_err,
1111 + )
1112 + await sock.send_multipart((identity, self._encoder.encode(response)))
1113 + return
1114 validation_err = _validate_asymmetric_region_lengths(
1115 local_regions=local_regions,
1116 remote_regions=remote_regions,
@@ -1171,11 +1257,15 @@ class MooncakeConnectorWorker:
1257 )
1258 await sock.send_multipart((identity, self._encoder.encode(response)))
1259
1174 - def resolve_need_send(self, send_meta: SendBlockMeta, remote_tp_ranks: list[int]):
1260 + def resolve_need_send(
1261 + self,
1262 + send_meta: SendBlockMeta,
1263 + remote_tp_ranks: list[int],
1264 + ):
1265 # Prepare for heterogeneous TP (one P pairs to multiple D)
1266 send_meta.need_send = len(remote_tp_ranks)
1267 logger.debug(
1178 - "Mooncake request %s will be served by %d consumer TP workers: %s",
1268 + "Mooncake request %s will be served by %d consumer TP workers: TP ranks=%s",
1269 send_meta.transfer_id,
1270 send_meta.need_send,
1271 remote_tp_ranks,
@@ -1394,10 +1484,13 @@ class MooncakeConnectorWorker:
1484 kv_data_lens = []
1485 seen_base_addresses = []
1486 self.block_len_per_layer = []
1487 + self.registered_layer_names = []
1488 + self.registered_layer_indices = []
1489
1490 split_k_and_v = self.transfer_topo.split_k_and_v
1491 tensor_size_bytes = None
1492 for layer_name, cache_or_caches in kv_caches.items():
1493 + layer_index = extract_layer_index(layer_name)
1494 cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1495 logger.debug(
1496 "registering layer %s with %d cache tensor(s)",
@@ -1428,6 +1521,8 @@ class MooncakeConnectorWorker:
1521 block_len = cache.stride(0) * cache.element_size()
1522
1523 self.block_len_per_layer.append(block_len)
1524 + self.registered_layer_names.append(layer_name)
1525 + self.registered_layer_indices.append(layer_index)
1526 kv_data_ptrs.append(base_addr)
1527 kv_data_lens.append(self.num_blocks * block_len)
1528
@@ -1547,6 +1642,8 @@ class MooncakeConnectorWorker:
1642 },
1643 kv_caches_base_addr=self.kv_caches_base_addr,
1644 block_lens=self.block_len_per_layer,
1645 + registered_layer_names=self.registered_layer_names,
1646 + registered_layer_indices=self.registered_layer_indices,
1647 )
1648
1649 encoded_data = self._encoder.encode(metadata)
@@ -1648,16 +1745,28 @@ class MooncakeConnectorWorker:
1745 remote_tp_ranks = self.transfer_topo.handshake_target_ranks(
1746 self._tp_size[remote_engine_id]
1747 )
1651 - count = len(remote_tp_ranks)
1748 + worker_addrs: list[str] = []
1749 + selected_remote_pp: dict[int, list[int]] = {}
1750 + for remote_tp_rank in remote_tp_ranks:
1751 + pp_to_addr = self._remote_agents[remote_engine_id][remote_tp_rank]
1752 + if self.pp_size == len(pp_to_addr) and self.pp_rank in pp_to_addr:
1753 + pp_ranks = [self.pp_rank]
1754 + else:
1755 + pp_ranks = sorted(pp_to_addr)
1756 + selected_remote_pp[remote_tp_rank] = pp_ranks
1757 + worker_addrs.extend(pp_to_addr[pp_rank] for pp_rank in pp_ranks)
1758 +
1759 + count = len(worker_addrs)
1760 logger.debug(
1653 - "Receiving Mooncake KV for engine %s from producer TP ranks %s",
1761 + "Receiving Mooncake KV for engine %s from producer TP ranks %s "
1762 + "and PP ranks %s",
1763 remote_engine_id,
1764 remote_tp_ranks,
1765 + selected_remote_pp,
1766 )
1767 for pull_meta in pull_metas.values():
1768 pull_meta.pull_tasks_count = count
1659 - for remote_tp_rank in remote_tp_ranks:
1660 - worker_addr = self._remote_agents[remote_engine_id][remote_tp_rank][0]
1769 + for worker_addr in worker_addrs:
1770 asyncio.create_task(
1771 self.receive_kv_from_single_worker(worker_addr, pull_metas)
1772 )
@@ -1740,11 +1849,17 @@ class MooncakeConnectorWorker:
1849 return self.transfer_topo.local_replicates_kv_cache
1850
1851 def _get_transfer_regions(
1743 - self, base_addrs: list[int], block_lens: list[int]
1852 + self,
1853 + base_addrs: list[int],
1854 + block_lens: list[int],
1855 + layer_names: list[str],
1856 + layer_indices: list[int],
1857 ) -> list[TransferRegion]:
1858 return _expand_transfer_regions(
1859 base_addrs=base_addrs,
1860 block_lens=block_lens,
1861 + layer_names=layer_names,
1862 + layer_indices=layer_indices,
1863 is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks,
1864 )
1865
@@ -1816,14 +1931,20 @@ def _async_loop(loop: asyncio.AbstractEventLoop):
1931
1932 def should_launch_bootstrap_server(vllm_config: VllmConfig) -> bool:
1933 assert (parallel_config := vllm_config.parallel_config)
1934 + # Only the TP=0, PP=0 worker of the designated engine should launch it.
1935 + if get_tensor_model_parallel_rank() != 0:
1936 + return False
1937 + if get_pp_group().rank_in_group != 0:
1938 + return False
1939 +
1940 # In hybrid or external LB mode,
1941 # each instance should have its own bootstrap server.
1821 - #
1942 + if parallel_config.local_engines_only:
1943 + return parallel_config.data_parallel_rank_local == 0
1944 +
1945 # In internal LB mode,
1823 - # only the real global first rank need to launch the bootstrap server.
1824 - return is_local_first_rank() and (
1825 - parallel_config.local_engines_only or parallel_config.data_parallel_index == 0
1826 - )
1946 + # only the first data-parallel engine should launch the bootstrap server.
1947 + return parallel_config.data_parallel_index == 0
1948
1949
1950 def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]:
@@ -1836,6 +1957,10 @@ def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]:
1957 if parallel_config.local_engines_only:
1958 # In hybrid or external LB mode, connect to local server.
1959 host = "127.0.0.1"
1960 + elif parallel_config.nnodes_within_dp > 1:
1961 + # Internal LB multi-node TP/PP uses the model-parallel master as the
1962 + # single bootstrap endpoint for all ranks in the engine.
1963 + host = parallel_config.master_addr
1964 else:
1965 host = parallel_config.data_parallel_master_ip
1966 port = envs.VLLM_MOONCAKE_BOOTSTRAP_PORT
vllm/distributed/kv_transfer/kv_transfer_state.py
+8 -1
@@ -50,8 +50,13 @@ def is_v1_kv_transfer_group(connector: KVConnectorBaseType | None = None) -> boo
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."""
53 + multi-node TP group share the same value.
54 +
55 + When PP is enabled, also broadcast across PP ranks so all workers in
56 + the same model-parallel engine share the same value.
57 + """
58 from vllm.distributed.parallel_state import (
59 + get_pp_group,
60 get_tp_group,
61 )
62
@@ -59,6 +64,8 @@ def _sync_engine_id_across_tp(vllm_config: "VllmConfig") -> None:
64 synced_id = get_tp_group().broadcast_object(
65 vllm_config.kv_transfer_config.engine_id, src=0
66 )
67 + if vllm_config.parallel_config.pipeline_parallel_size > 1:
68 + synced_id = get_pp_group().broadcast_object(synced_id, src=0)
69 vllm_config.kv_transfer_config.engine_id = synced_id
70
71