@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d32575a2d2

[ROCm][P/D] Support MoRIIO heterogeneous TP fan-in (#46332)

Signed-off-by: Tan Pin Siang <tanpinsiang@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Jun Kang Chow <junkangchow@gmail.com> Co-authored-by: Chun Fang <chun.fang@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>

Tan Pin Siang committed Jun 23, 2026 at 18:33 UTC d32575a2d2d83362a69d8bcf4544589cdbc11a9c
6 files changed +544 -33
tests/v1/kv_connector/unit/test_moriio_connector.py
+8 -1
@@ -195,7 +195,14 @@ class FakeMoRIIOWrapper:
195 def _handle_completion_message(self, msg: str):
196 pass
197
198 - def send_notify(self, req_ids, remote_ip, remote_port, message_type=None):
198 + def send_notify(
199 + self,
200 + req_ids,
201 + remote_ip,
202 + remote_port,
203 + message_type=None,
204 + message_fields=None,
205 + ):
206 pass
207
208 def pop_finished_req_ids(self):
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+30 -10
@@ -36,6 +36,7 @@ msgpack = importlib.import_module("msgpack")
36
37 ROLE = moriio_common.ROLE
38 MoRIIOError = moriio_common.MoRIIOError
39 +MoRIIOTransferAck = moriio_common.MoRIIOTransferAck
40 RemoteAllocInfo = moriio_common.RemoteAllocInfo
41 WriteTask = moriio_common.WriteTask
42 set_role = moriio_common.set_role
@@ -440,9 +441,16 @@ def test_write_completion_notifies_once_after_all_sealed_writes_finish():
441 def _mark_transfer_terminal_locked(self, transfer_id):
442 self._terminal_transfer_ids[transfer_id] = None
443
443 - def send_notify(self, transfer_id, remote_ip, remote_port, message_type=None):
444 + def send_notify(
445 + self,
446 + transfer_id,
447 + remote_ip,
448 + remote_port,
449 + message_type=None,
450 + message_fields=None,
451 + ):
452 self.notifications.append(
445 - (transfer_id, remote_ip, remote_port, message_type)
453 + (transfer_id, remote_ip, remote_port, message_type, message_fields)
454 )
455
456 wrapper = FakeWrapper()
@@ -464,8 +472,8 @@ def test_write_completion_notifies_once_after_all_sealed_writes_finish():
472 writer._mark_write_done("xfer", request_info)
473 writer._finalize_if_complete("xfer", request_info)
474
467 - assert wrapper.notifications == [("xfer", "127.0.0.1", 7002, "write_done")]
468 - assert wrapper.done_req_ids == ["xfer"]
475 + assert wrapper.notifications == [("xfer", "127.0.0.1", 7002, "write_done", None)]
476 + assert wrapper.done_req_ids == [MoRIIOTransferAck("xfer")]
477 assert wrapper.done_remote_allocate_req_dict == {}
478 assert wrapper.wait_count == 1
479 assert wrapper.waited_statuses == [["status-a", "status-b"]]
@@ -509,7 +517,7 @@ def test_write_failure_marks_terminal_and_clears_scheduled_state():
517
518 writer._mark_request_done("xfer")
519
512 - assert wrapper.done_req_ids == ["xfer"]
520 + assert wrapper.done_req_ids == [MoRIIOTransferAck("xfer")]
521 assert wrapper.done_remote_allocate_req_dict == {}
522 assert wrapper._is_transfer_terminal_locked("xfer")
523 assert "xfer" not in writer._scheduled_writes
@@ -581,9 +589,21 @@ def test_late_remote_blocks_message_is_ignored_after_transfer_done():
589 pytest.param(
590 ROLE.PRODUCER,
591 msgpack.dumps({"type": "release", "transfer_id": "xfer"}),
584 - "release",
592 + MoRIIOTransferAck("xfer"),
593 id="release",
594 ),
595 + pytest.param(
596 + ROLE.PRODUCER,
597 + msgpack.dumps(
598 + {
599 + "type": "release",
600 + "transfer_id": "xfer",
601 + "consumer_tp_size": 8,
602 + }
603 + ),
604 + MoRIIOTransferAck("xfer", 8),
605 + id="release-consumer-tp-size",
606 + ),
607 pytest.param(None, b"xfer", "plain", id="plain-string"),
608 ],
609 )
@@ -603,11 +623,11 @@ def test_moriio_wrapper_routes_valid_messages(role, payload, expected):
623 assert request_info.decode_dp_rank == 3
624 elif expected == "write_done":
625 assert wrapper.done_write_cache_req_ids == ["xfer"]
606 - elif expected == "release":
607 - assert wrapper.done_req_ids == ["xfer"]
608 - assert wrapper._is_transfer_terminal_locked("xfer")
609 - else:
626 + elif expected == "plain":
627 assert completions == ["xfer"]
628 + else:
629 + assert wrapper.done_req_ids == [expected]
630 + assert wrapper._is_transfer_terminal_locked("xfer")
631
632
633 @pytest.mark.parametrize(
tests/v1/kv_connector/unit/test_moriio_tp_ack.py new
+310
@@ -0,0 +1,310 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import threading
5 +
6 +import pytest
7 +
8 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
9 + MoRIIOMode,
10 + MoRIIOTransferAck,
11 +)
12 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
13 + MoRIIOConnectorWorker,
14 + get_moriio_expected_ack_count,
15 + get_moriio_remote_tp_rank,
16 + resolve_moriio_transfer_ack,
17 + validate_moriio_heterogeneous_tp_kv_heads,
18 +)
19 +
20 +
21 +def test_remote_tp_rank_same_tp_maps_to_self():
22 + assert [get_moriio_remote_tp_rank(rank, 4, 4) for rank in range(4)] == [
23 + 0,
24 + 1,
25 + 2,
26 + 3,
27 + ]
28 +
29 +
30 +def test_remote_tp_rank_p4_d8_floor_maps_decode_to_prefill():
31 + assert [get_moriio_remote_tp_rank(rank, 8, 4) for rank in range(8)] == [
32 + 0,
33 + 0,
34 + 1,
35 + 1,
36 + 2,
37 + 2,
38 + 3,
39 + 3,
40 + ]
41 +
42 +
43 +def test_remote_tp_rank_p8_d4_maps_to_first_prefill_rank_per_pair():
44 + assert [get_moriio_remote_tp_rank(rank, 4, 8) for rank in range(4)] == [
45 + 0,
46 + 2,
47 + 4,
48 + 6,
49 + ]
50 +
51 +
52 +@pytest.mark.parametrize(
53 + ("local_tp_rank", "local_tp_size", "remote_tp_size"),
54 + [
55 + (0, 6, 4),
56 + (0, 4, 6),
57 + ],
58 +)
59 +def test_remote_tp_rank_invalid_non_multiple_tp_raises(
60 + local_tp_rank: int, local_tp_size: int, remote_tp_size: int
61 +):
62 + with pytest.raises(ValueError, match="multiple"):
63 + get_moriio_remote_tp_rank(local_tp_rank, local_tp_size, remote_tp_size)
64 +
65 +
66 +@pytest.mark.parametrize(
67 + ("local_tp_size", "remote_tp_size", "total_num_kv_heads"),
68 + [
69 + (4, 4, 8),
70 + (8, 4, 4),
71 + (4, 8, 4),
72 + ],
73 +)
74 +def test_heterogeneous_tp_head_guard_allows_supported_layouts(
75 + local_tp_size: int, remote_tp_size: int, total_num_kv_heads: int
76 +):
77 + validate_moriio_heterogeneous_tp_kv_heads(
78 + local_tp_size,
79 + remote_tp_size,
80 + total_num_kv_heads,
81 + is_mla=False,
82 + )
83 +
84 +
85 +def test_heterogeneous_tp_head_guard_allows_mla_layouts():
86 + validate_moriio_heterogeneous_tp_kv_heads(
87 + local_tp_size=2,
88 + remote_tp_size=4,
89 + total_num_kv_heads=4,
90 + is_mla=True,
91 + )
92 +
93 +
94 +@pytest.mark.parametrize(
95 + ("local_tp_size", "remote_tp_size", "total_num_kv_heads"),
96 + [
97 + (4, 2, 4),
98 + (2, 4, 4),
99 + ],
100 +)
101 +def test_heterogeneous_tp_head_guard_rejects_split_kv_heads(
102 + local_tp_size: int, remote_tp_size: int, total_num_kv_heads: int
103 +):
104 + with pytest.raises(NotImplementedError, match="replicated KV heads"):
105 + validate_moriio_heterogeneous_tp_kv_heads(
106 + local_tp_size,
107 + remote_tp_size,
108 + total_num_kv_heads,
109 + is_mla=False,
110 + )
111 +
112 +
113 +def test_expected_ack_count_for_homogeneous_or_smaller_consumer_tp_is_one():
114 + assert get_moriio_expected_ack_count(4, 4) == 1
115 + assert get_moriio_expected_ack_count(8, 4) == 1
116 +
117 +
118 +def test_expected_ack_count_for_decode_fan_in():
119 + assert get_moriio_expected_ack_count(4, 8) == 2
120 +
121 +
122 +def test_expected_ack_count_rejects_non_multiple_fan_in():
123 + with pytest.raises(ValueError, match="multiple"):
124 + get_moriio_expected_ack_count(4, 6)
125 +
126 +
127 +def test_plain_string_ack_is_backward_compatible_single_ack():
128 + notification_counts: dict[str, int] = {}
129 + completed_transfer_ids: set[str] = set()
130 +
131 + assert (
132 + resolve_moriio_transfer_ack(
133 + "tx-plain",
134 + producer_tp_size=4,
135 + live_transfer_ids={"tx-plain"},
136 + notification_counts=notification_counts,
137 + completed_transfer_ids=completed_transfer_ids,
138 + )
139 + == "tx-plain"
140 + )
141 + assert notification_counts == {}
142 + assert completed_transfer_ids == {"tx-plain"}
143 +
144 +
145 +def test_structured_release_ack_waits_for_all_expected_acks():
146 + ack = MoRIIOTransferAck("tx-fanin", consumer_tp_size=8)
147 + notification_counts: dict[str, int] = {}
148 + completed_transfer_ids: set[str] = set()
149 +
150 + assert (
151 + resolve_moriio_transfer_ack(
152 + ack,
153 + producer_tp_size=4,
154 + live_transfer_ids={"tx-fanin"},
155 + notification_counts=notification_counts,
156 + completed_transfer_ids=completed_transfer_ids,
157 + )
158 + is None
159 + )
160 + assert notification_counts == {"tx-fanin": 1}
161 + assert completed_transfer_ids == set()
162 +
163 + assert (
164 + resolve_moriio_transfer_ack(
165 + ack,
166 + producer_tp_size=4,
167 + live_transfer_ids={"tx-fanin"},
168 + notification_counts=notification_counts,
169 + completed_transfer_ids=completed_transfer_ids,
170 + )
171 + == "tx-fanin"
172 + )
173 + assert notification_counts == {}
174 + assert completed_transfer_ids == {"tx-fanin"}
175 +
176 +
177 +def test_duplicate_ack_after_completion_does_not_resolve_twice():
178 + ack = MoRIIOTransferAck("tx-dup", consumer_tp_size=8)
179 + notification_counts: dict[str, int] = {}
180 + completed_transfer_ids: set[str] = set()
181 +
182 + assert (
183 + resolve_moriio_transfer_ack(
184 + ack,
185 + producer_tp_size=4,
186 + live_transfer_ids={"tx-dup"},
187 + notification_counts=notification_counts,
188 + completed_transfer_ids=completed_transfer_ids,
189 + )
190 + is None
191 + )
192 + assert (
193 + resolve_moriio_transfer_ack(
194 + ack,
195 + producer_tp_size=4,
196 + live_transfer_ids={"tx-dup"},
197 + notification_counts=notification_counts,
198 + completed_transfer_ids=completed_transfer_ids,
199 + )
200 + == "tx-dup"
201 + )
202 + assert (
203 + resolve_moriio_transfer_ack(
204 + ack,
205 + producer_tp_size=4,
206 + live_transfer_ids={"tx-dup"},
207 + notification_counts=notification_counts,
208 + completed_transfer_ids=completed_transfer_ids,
209 + )
210 + is None
211 + )
212 + assert notification_counts == {}
213 + assert completed_transfer_ids == {"tx-dup"}
214 +
215 +
216 +def test_ack_for_non_live_transfer_is_ignored():
217 + notification_counts: dict[str, int] = {}
218 + completed_transfer_ids: set[str] = set()
219 +
220 + assert (
221 + resolve_moriio_transfer_ack(
222 + MoRIIOTransferAck("tx-stale", consumer_tp_size=8),
223 + producer_tp_size=4,
224 + live_transfer_ids={"tx-live"},
225 + notification_counts=notification_counts,
226 + completed_transfer_ids=completed_transfer_ids,
227 + )
228 + is None
229 + )
230 + assert notification_counts == {}
231 + assert completed_transfer_ids == set()
232 +
233 +
234 +def test_worker_get_finished_counts_structured_release_fan_in():
235 + class FakeWrapper:
236 + def __init__(self):
237 + self.batches = [
238 + [MoRIIOTransferAck("tx-fanin", consumer_tp_size=8)],
239 + [MoRIIOTransferAck("tx-fanin", consumer_tp_size=8)],
240 + ]
241 +
242 + def pop_finished_req_ids(self):
243 + return self.batches.pop(0)
244 +
245 + def shutdown(self):
246 + pass
247 +
248 + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
249 + worker.is_producer = True
250 + worker.mode = MoRIIOMode.READ
251 + worker.world_size = 4
252 + worker.moriio_wrapper = FakeWrapper()
253 + worker.transfer_id_to_request_id = {"tx-fanin": "req-fanin"}
254 + worker._consumer_notification_counts = {}
255 + worker._completed_consumer_notifications = set()
256 +
257 + assert worker.get_finished() == (set(), set())
258 + assert worker._consumer_notification_counts == {"tx-fanin": 1}
259 +
260 + assert worker.get_finished() == ({"req-fanin"}, set())
261 + assert worker._consumer_notification_counts == {}
262 + assert worker._completed_consumer_notifications == {"tx-fanin"}
263 +
264 +
265 +def test_read_completion_sends_structured_release_with_consumer_tp_size():
266 + class DoneStatus:
267 + def Succeeded(self):
268 + return True
269 +
270 + def Failed(self):
271 + return False
272 +
273 + class FakeWrapper:
274 + def __init__(self):
275 + self.lock = threading.Lock()
276 + self.sent = []
277 +
278 + def send_notify(
279 + self,
280 + transfer_id,
281 + host,
282 + port,
283 + message_type=None,
284 + message_fields=None,
285 + ):
286 + self.sent.append((transfer_id, host, port, message_type, message_fields))
287 +
288 + def shutdown(self):
289 + pass
290 +
291 + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
292 + worker.world_size = 8
293 + worker.moriio_wrapper = FakeWrapper()
294 + worker._recving_transfers = {"req": [DoneStatus()]}
295 + worker._recving_transfers_callback_addr = {
296 + "req": ("127.0.0.1", "7000", "tx-release")
297 + }
298 +
299 + assert worker._pop_done_transfers() == {"tx-release"}
300 + assert worker.moriio_wrapper.sent == [
301 + (
302 + "tx-release",
303 + "127.0.0.1",
304 + "7000",
305 + "release",
306 + {"consumer_tp_size": 8},
307 + )
308 + ]
309 + assert worker._recving_transfers == {}
310 + assert worker._recving_transfers_callback_addr == {}
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+6 -1
@@ -6,7 +6,7 @@ import threading
6 import time
7 from collections.abc import Iterator
8 from dataclasses import dataclass
9 -from typing import TYPE_CHECKING, Any
9 +from typing import TYPE_CHECKING, Any, NamedTuple
10
11 import msgspec
12 import regex as re
@@ -43,6 +43,11 @@ ReqId = str
43 TransferId = str
44
45
46 +class MoRIIOTransferAck(NamedTuple):
47 + transfer_id: TransferId
48 + consumer_tp_size: int = 1
49 +
50 +
51 @dataclass
52 class WriteTask:
53 request_id: ReqId
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+163 -10
@@ -6,6 +6,7 @@ import queue
6 import threading
7 import time
8 from collections import defaultdict
9 +from collections.abc import Collection
10 from concurrent.futures import Future, ThreadPoolExecutor
11 from typing import TYPE_CHECKING, Any
12
@@ -30,6 +31,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
31 MoRIIOConnectorMetadata,
32 MoRIIOConstants,
33 MoRIIOMode,
34 + MoRIIOTransferAck,
35 ReqId,
36 ReqMeta,
37 TransferId,
@@ -98,6 +100,92 @@ def is_moriio_available() -> bool:
100 return MoRIIO_enabled
101
102
103 +def get_moriio_remote_tp_rank(
104 + local_tp_rank: int, local_tp_size: int, remote_tp_size: int
105 +) -> int:
106 + if local_tp_size <= 0 or remote_tp_size <= 0:
107 + raise ValueError("TP sizes must be positive")
108 + if local_tp_rank < 0 or local_tp_rank >= local_tp_size:
109 + raise ValueError(
110 + f"local_tp_rank {local_tp_rank} must be in [0, {local_tp_size})"
111 + )
112 + if remote_tp_size == local_tp_size:
113 + return local_tp_rank
114 + if remote_tp_size > local_tp_size:
115 + if remote_tp_size % local_tp_size != 0:
116 + raise ValueError(
117 + f"remote tp_size {remote_tp_size} must be a multiple of local "
118 + f"tp_size {local_tp_size} for heterogeneous-TP P/D"
119 + )
120 + return local_tp_rank * (remote_tp_size // local_tp_size)
121 + if local_tp_size % remote_tp_size != 0:
122 + raise ValueError(
123 + f"local tp_size {local_tp_size} must be a multiple of remote "
124 + f"tp_size {remote_tp_size} for heterogeneous-TP P/D"
125 + )
126 + return local_tp_rank // (local_tp_size // remote_tp_size)
127 +
128 +
129 +def validate_moriio_heterogeneous_tp_kv_heads(
130 + local_tp_size: int,
131 + remote_tp_size: int,
132 + total_num_kv_heads: int,
133 + is_mla: bool,
134 +) -> None:
135 + if is_mla or local_tp_size == remote_tp_size:
136 + return
137 + if local_tp_size <= 0 or remote_tp_size <= 0 or total_num_kv_heads <= 0:
138 + raise ValueError("TP sizes and total_num_kv_heads must be positive")
139 + if min(local_tp_size, remote_tp_size) >= total_num_kv_heads:
140 + return
141 + raise NotImplementedError(
142 + "MoRIIO heterogeneous TP requires replicated KV heads on both "
143 + f"prefill and decode. Got total_num_kv_heads={total_num_kv_heads}, "
144 + f"local_tp_size={local_tp_size}, remote_tp_size={remote_tp_size}."
145 + )
146 +
147 +
148 +def get_moriio_expected_ack_count(producer_tp_size: int, consumer_tp_size: int) -> int:
149 + if producer_tp_size <= 0 or consumer_tp_size <= 0:
150 + raise ValueError("TP sizes must be positive")
151 + if consumer_tp_size <= producer_tp_size:
152 + return 1
153 + if consumer_tp_size % producer_tp_size != 0:
154 + raise ValueError(
155 + f"consumer tp_size {consumer_tp_size} must be a multiple of "
156 + f"producer tp_size {producer_tp_size} for heterogeneous-TP P/D"
157 + )
158 + return consumer_tp_size // producer_tp_size
159 +
160 +
161 +def resolve_moriio_transfer_ack(
162 + ack: MoRIIOTransferAck | TransferId,
163 + producer_tp_size: int,
164 + live_transfer_ids: Collection[TransferId],
165 + notification_counts: dict[TransferId, int],
166 + completed_transfer_ids: set[TransferId],
167 +) -> TransferId | None:
168 + if isinstance(ack, str):
169 + ack = MoRIIOTransferAck(ack)
170 + transfer_id = ack.transfer_id
171 + if transfer_id not in live_transfer_ids:
172 + return None
173 + if transfer_id in completed_transfer_ids:
174 + return None
175 +
176 + expected_acks = get_moriio_expected_ack_count(
177 + producer_tp_size, ack.consumer_tp_size
178 + )
179 + count = notification_counts.get(transfer_id, 0) + 1
180 + if count < expected_acks:
181 + notification_counts[transfer_id] = count
182 + return None
183 +
184 + notification_counts.pop(transfer_id, None)
185 + completed_transfer_ids.add(transfer_id)
186 + return transfer_id
187 +
188 +
189 class MoRIIOConnector(KVConnectorBase_V1):
190 def __init__(
191 self,
@@ -798,6 +886,11 @@ class MoRIIOConnectorWorker:
886 # Completions that arrived before transfer_id_to_request_id was populated.
887 # Retried each step until the mapping is established.
888 self._unmatched_write_completions: set[str] = set()
889 + # Producer-side READ-mode ACK fan-in. When decode TP is larger than
890 + # prefill TP, multiple decode ranks can read from one prefill rank and
891 + # notify the same transfer_id. Blocks are reusable only after all ACKs.
892 + self._consumer_notification_counts: dict[TransferId, int] = {}
893 + self._completed_consumer_notifications: set[TransferId] = set()
894
895 role = "producer" if self.is_producer else "consumer"
896 engine_suffix = (
@@ -1161,7 +1254,9 @@ class MoRIIOConnectorWorker:
1254 # a hack to keep us moving. We will switch when moving to etcd
1255 # or where we have a single ZMQ socket in the scheduler.
1256
1164 - port_offset = get_port_offset(remote_dp_rank, self.tp_rank)
1257 + port_offset = get_port_offset(
1258 + remote_dp_rank, self._remote_tp_rank(remote_tp_size)
1259 + )
1260 path = make_zmq_path("tcp", host, port + port_offset)
1261 logger.debug("handshake Querying metadata on path: %s", path)
1262
@@ -1226,6 +1321,9 @@ class MoRIIOConnectorWorker:
1321
1322 return {remote_agent_name}
1323
1324 + def _remote_tp_rank(self, remote_tp_size: int) -> int:
1325 + return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size)
1326 +
1327 def _background_moriio_handshake(
1328 self, req_id: ReqId, remote_engine_id: EngineId, meta: ReqMeta
1329 ):
@@ -1435,13 +1533,32 @@ class MoRIIOConnectorWorker:
1533 done_sending, done_recving = set(), set()
1534
1535 if self.is_producer:
1438 - # pop_finished_req_ids returns transfer_ids (the ZMQ payload sent
1439 - # by decode via send_notify); map back to req_ids for the scheduler.
1440 - finished_transfer_ids = self.moriio_wrapper.pop_finished_req_ids()
1536 + # pop_finished_req_ids returns release ACKs sent by decode. Keep
1537 + # duplicate ACKs because heterogeneous TP can fan multiple decode
1538 + # ranks into one prefill rank for the same transfer_id.
1539 + finished_acks = self.moriio_wrapper.pop_finished_req_ids()
1540 + resolved_transfer_ids: set[TransferId] = set()
1541 + for ack in finished_acks:
1542 + transfer_id = ack if isinstance(ack, str) else ack.transfer_id
1543 + if transfer_id not in self.transfer_id_to_request_id:
1544 + logger.warning(
1545 + "Could not find %s in transfer_id_to_request_id "
1546 + "lookup table. This could lead to a possible hang.",
1547 + transfer_id,
1548 + )
1549 + continue
1550 + resolved_transfer_id = resolve_moriio_transfer_ack(
1551 + ack,
1552 + producer_tp_size=self.world_size,
1553 + live_transfer_ids=self.transfer_id_to_request_id.keys(),
1554 + notification_counts=self._consumer_notification_counts,
1555 + completed_transfer_ids=(self._completed_consumer_notifications),
1556 + )
1557 + if resolved_transfer_id is not None:
1558 + resolved_transfer_ids.add(resolved_transfer_id)
1559 done_sending = {
1560 self.transfer_id_to_request_id[xfer_id]
1443 - for xfer_id in finished_transfer_ids
1444 - if xfer_id in self.transfer_id_to_request_id
1561 + for xfer_id in resolved_transfer_ids
1562 }
1563 else:
1564 if self.mode == MoRIIOMode.WRITE:
@@ -1486,7 +1603,13 @@ class MoRIIOConnectorWorker:
1603 if last.Succeeded():
1604 host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
1605 done_req_ids.add(xfer_id)
1489 - self.moriio_wrapper.send_notify(xfer_id, host, port)
1606 + self.moriio_wrapper.send_notify(
1607 + xfer_id,
1608 + host,
1609 + port,
1610 + message_type="release",
1611 + message_fields={"consumer_tp_size": self.world_size},
1612 + )
1613 to_remove.append(req_id)
1614 elif last.Failed():
1615 logger.error(
@@ -1499,7 +1622,13 @@ class MoRIIOConnectorWorker:
1622 )
1623 host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
1624 try:
1502 - self.moriio_wrapper.send_notify(xfer_id, host, port)
1625 + self.moriio_wrapper.send_notify(
1626 + xfer_id,
1627 + host,
1628 + port,
1629 + message_type="release",
1630 + message_fields={"consumer_tp_size": self.world_size},
1631 + )
1632 except Exception:
1633 logger.exception(
1634 "Failed to send error notification for request %s",
@@ -1585,6 +1714,15 @@ class MoRIIOConnectorWorker:
1714 """
1715 self.transfer_id_to_request_id = metadata.transfer_id_to_request_id
1716 if self.is_producer:
1717 + live_transfer_ids = set(self.transfer_id_to_request_id)
1718 + self._consumer_notification_counts = {
1719 + transfer_id: count
1720 + for transfer_id, count in self._consumer_notification_counts.items()
1721 + if transfer_id in live_transfer_ids
1722 + }
1723 + self._completed_consumer_notifications.intersection_update(
1724 + live_transfer_ids
1725 + )
1726 self.moriio_wrapper.async_wait_reqid()
1727 return
1728 if self.mode == MoRIIOMode.WRITE:
@@ -1663,6 +1801,7 @@ class MoRIIOConnectorWorker:
1801 remote_block_ids=meta.remote_block_ids,
1802 remote_host=meta.remote_host,
1803 remote_notify_port=meta.remote_notify_port,
1804 + remote_tp_size=meta.tp_size,
1805 )
1806
1807 def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
@@ -1756,6 +1895,7 @@ class MoRIIOConnectorWorker:
1895 local_block_ids: list[int],
1896 remote_block_ids: list[int],
1897 remote_moriio_meta: MoRIIOAgentMetadata,
1898 + remote_tp_size: int | None = None,
1899 ) -> tuple[list[int], list[int], list[int]]:
1900 """Compute transfer offsets for block data.
1901
@@ -1767,6 +1907,14 @@ class MoRIIOConnectorWorker:
1907 Returns:
1908 Tuple of (local_offsets, remote_offsets, transfer_sizes)
1909 """
1910 + validate_moriio_heterogeneous_tp_kv_heads(
1911 + local_tp_size=self.world_size,
1912 + remote_tp_size=(
1913 + remote_tp_size if remote_tp_size is not None else self.world_size
1914 + ),
1915 + total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
1916 + is_mla=self._is_mla_cache_layer(layer_name),
1917 + )
1918 return compute_block_transfer_offsets(
1919 layer_name=layer_name,
1920 kv_cache=self.kv_caches[layer_name],
@@ -1788,6 +1936,7 @@ class MoRIIOConnectorWorker:
1936 transfer_id: str,
1937 remote_host: str,
1938 remote_notify_port: int,
1939 + remote_tp_size: int,
1940 ) -> None:
1941 if self.mode == MoRIIOMode.WRITE:
1942 return
@@ -1800,7 +1949,11 @@ class MoRIIOConnectorWorker:
1949 layer_name
1950 )
1951 offs = self._compute_block_transfer_offsets(
1803 - layer_name, local_block_ids, remote_block_ids, remote_moriio_meta
1952 + layer_name,
1953 + local_block_ids,
1954 + remote_block_ids,
1955 + remote_moriio_meta,
1956 + remote_tp_size=remote_tp_size,
1957 )
1958 # TODO : apply multi-session batch-read when moriio support it
1959 transfer_status = self.moriio_wrapper.read_remote_data(
@@ -1810,6 +1963,6 @@ class MoRIIOConnectorWorker:
1963 self._recving_transfers[request_id].append(transfer_status)
1964 self._recving_transfers_callback_addr[request_id] = (
1965 remote_host,
1813 - str(remote_notify_port + self.tp_rank),
1966 + str(remote_notify_port + self._remote_tp_rank(remote_tp_size)),
1967 transfer_id,
1968 )
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+27 -11
@@ -27,6 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
27 MoRIIOAgentMetadata,
28 MoRIIOConstants,
29 MoRIIOError,
30 + MoRIIOTransferAck,
31 RemoteAllocInfo,
32 TransferError,
33 TransferId,
@@ -261,7 +262,7 @@ class MoRIIOWriter:
262 """Mark a request done so its blocks are freed, even on transfer failure."""
263 wrapper = self.worker.moriio_wrapper
264 with wrapper.lock:
264 - wrapper.done_req_ids.append(transfer_id)
265 + wrapper.done_req_ids.append(MoRIIOTransferAck(transfer_id))
266 wrapper.done_remote_allocate_req_dict.pop(transfer_id, None)
267 wrapper._mark_transfer_terminal_locked(transfer_id)
268 self._clear_transfer_state(transfer_id)
@@ -466,7 +467,9 @@ class MoRIIOWriter:
467 )
468 # mark request as done, then we can free the blocks
469 with self.worker.moriio_wrapper.lock:
469 - self.worker.moriio_wrapper.done_req_ids.append(transfer_id)
470 + self.worker.moriio_wrapper.done_req_ids.append(
471 + MoRIIOTransferAck(transfer_id)
472 + )
473 self.worker.moriio_wrapper.done_remote_allocate_req_dict.pop(
474 transfer_id, None
475 )
@@ -508,7 +511,7 @@ class MoRIIOWrapper:
511 self.remote_engine_ip: str | None = None
512 self.notify_port: int | None = None
513 self.lock = threading.Lock()
511 - self.done_req_ids: list[str] = []
514 + self.done_req_ids: list[MoRIIOTransferAck] = []
515 self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {}
516 self.done_write_cache_req_ids: list[str] = []
517 self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict()
@@ -784,15 +787,20 @@ class MoRIIOWrapper:
787 "Only prefill can get transfer release messages"
788 )
789 transfer_id = data["transfer_id"]
790 + consumer_tp_size = int(data.get("consumer_tp_size", 1))
791 + if consumer_tp_size <= 0:
792 + raise MoRIIOError(
793 + f"Invalid consumer_tp_size in release message: {consumer_tp_size}"
794 + )
795 with self.lock:
788 - self.done_req_ids.append(transfer_id)
796 + self.done_req_ids.append(MoRIIOTransferAck(transfer_id, consumer_tp_size))
797 self.done_remote_allocate_req_dict.pop(transfer_id, None)
798 self._mark_transfer_terminal_locked(transfer_id)
799
800 def _handle_completion_message(self, msg: str):
801 with self.lock:
802 if get_role() == ROLE.PRODUCER:
795 - self.done_req_ids.append(msg)
803 + self.done_req_ids.append(MoRIIOTransferAck(msg))
804 self.done_remote_allocate_req_dict.pop(msg, None)
805 self._mark_transfer_terminal_locked(msg)
806 else:
@@ -808,7 +816,12 @@ class MoRIIOWrapper:
816 self._terminal_transfer_ids.popitem(last=False)
817
818 def send_notify(
811 - self, req_ids, remote_ip, remote_port, message_type: str | None = None
819 + self,
820 + req_ids,
821 + remote_ip,
822 + remote_port,
823 + message_type: str | None = None,
824 + message_fields: dict[str, Any] | None = None,
825 ):
826 if not remote_ip or not remote_port:
827 logger.warning("Missing remote_ip or remote_port for notification")
@@ -836,18 +849,21 @@ class MoRIIOWrapper:
849 if message_type is None:
850 sock.send(req_id.encode("utf-8"))
851 else:
839 - sock.send(
840 - msgpack.dumps({"type": message_type, "transfer_id": req_id})
841 - )
852 + payload = {"type": message_type, "transfer_id": req_id}
853 + if message_fields:
854 + payload.update(message_fields)
855 + sock.send(msgpack.dumps(payload))
856 except Exception as e:
857 logger.error("Failed to send notification to %s: %s", path, e)
858 self.paths.pop(path, None)
859 raise
860
861 def pop_finished_req_ids(self):
848 - # producer invocation: get the set of completed requests at the decode
862 + # Producer invocation: return every completion message since the last
863 + # call. Do not dedupe: heterogeneous TP can produce multiple release
864 + # ACKs for the same transfer_id and the caller must count each one.
865 with self.lock:
850 - done_send = set(self.done_req_ids)
866 + done_send = list(self.done_req_ids)
867 self.done_req_ids = []
868 return done_send
869