@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7e47fb72b5

[ROCm][P/D] Fix MoRIIO WRITE mode for mixed KV layouts (#46290)

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>

Tan Pin Siang committed Jun 23, 2026 at 12:12 UTC 7e47fb72b568bd18062ca14e71271dd32edb372c
6 files changed +1104 -195
tests/v1/kv_connector/unit/test_moriio_connector.py
+75 -2
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import importlib.util
4 +import socket
5 import uuid
6 from unittest.mock import MagicMock, patch
7
@@ -9,7 +10,6 @@ import pytest
10 import torch
11 import zmq
12
12 -from tests.conftest import _find_free_port
13 from vllm.config import (
14 CacheConfig,
15 DeviceConfig,
@@ -23,12 +23,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
23 MoRIIOAgentMetadata,
24 MoRIIOConnectorMetadata,
25 MoRIIOConstants,
26 + MoRIIOMode,
27 resolve_host_ip,
28 zmq_ctx,
29 )
30 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
31 KVConnectorRole,
32 MoRIIOConnector,
33 + MoRIIOConnectorScheduler,
34 MoRIIOConnectorWorker,
35 )
36 from vllm.platforms import current_platform
@@ -46,6 +48,12 @@ from vllm.v1.kv_cache_interface import (
48 from .utils import create_request, create_scheduler
49
50
51 +def _find_free_port() -> int:
52 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
53 + sock.bind(("", 0))
54 + return sock.getsockname()[1]
55 +
56 +
57 def _make_test_kv_cache_config() -> KVCacheConfig:
58 layer_names = ["layer0", "layer1", "layer2"]
59 return KVCacheConfig(
@@ -121,6 +129,16 @@ def _setup_kv_transfer_request(
129 return request
130
131
132 +def _write_consumer_scheduler_for_finished_request(tp_size: int = 2):
133 + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler)
134 + scheduler.is_producer = False
135 + scheduler.mode = MoRIIOMode.WRITE
136 + scheduler.tp_size = tp_size
137 + scheduler._reqs_need_recv = {}
138 + scheduler.unmap_request_id = MagicMock()
139 + return scheduler
140 +
141 +
142 class FakeMoRIIOWrapper:
143 # A fake MoRIIOWrapper for testing purposes
144 def __init__(self, *args, **kwargs):
@@ -177,7 +195,7 @@ class FakeMoRIIOWrapper:
195 def _handle_completion_message(self, msg: str):
196 pass
197
180 - def send_notify(self, req_ids, remote_ip, remote_port):
198 + def send_notify(self, req_ids, remote_ip, remote_port, message_type=None):
199 pass
200
201 def pop_finished_req_ids(self):
@@ -434,6 +452,61 @@ def test_read_mode_loads_remote_block_ids():
452 assert block_id == block.block_id, f"{block_id} != {block.block_id}"
453
454
455 +@pytest.mark.parametrize(
456 + ("transfer_id", "extra_params", "expected_notifications"),
457 + [
458 + pytest.param(
459 + "xfer-7",
460 + {"remote_host": "127.0.0.1", "remote_notify_port": 7000},
461 + [
462 + ("xfer-7", "127.0.0.1", 7000),
463 + ("xfer-7", "127.0.0.1", 7001),
464 + ],
465 + id="address-available",
466 + ),
467 + pytest.param("xfer-8", {}, [], id="address-unavailable-plain-id"),
468 + ],
469 +)
470 +def test_write_mode_finished_before_alloc_releases_prefill_blocks(
471 + transfer_id, extra_params, expected_notifications
472 +):
473 + scheduler = _write_consumer_scheduler_for_finished_request(tp_size=2)
474 + notifications = []
475 + scheduler._send_transfer_release = lambda transfer_id, host, port: (
476 + notifications.append((transfer_id, host, port))
477 + )
478 + request = create_request(request_id=7, do_remote_prefill=True)
479 + request.request_id = "plain-decode-id"
480 + request.kv_transfer_params = {
481 + "do_remote_prefill": True,
482 + "do_remote_decode": False,
483 + "transfer_id": transfer_id,
484 + } | extra_params
485 +
486 + delay_free, new_params = scheduler.request_finished(request, block_ids=[])
487 +
488 + assert not delay_free
489 + assert new_params is None
490 + assert request.kv_transfer_params["do_remote_prefill"] is False
491 + assert scheduler._reqs_need_recv == {}
492 + assert notifications == expected_notifications
493 +
494 +
495 +def test_send_transfer_release_sends_structured_release_message():
496 + scheduler = _write_consumer_scheduler_for_finished_request()
497 + path = make_zmq_path("tcp", "127.0.0.1", 7000)
498 + sock = MagicMock()
499 + scheduler.paths = {path: sock}
500 +
501 + scheduler._send_transfer_release("xfer-7", "127.0.0.1", 7000)
502 +
503 + payload = sock.send.call_args.args[0]
504 + assert msgspec.msgpack.decode(payload) == {
505 + "type": "release",
506 + "transfer_id": "xfer-7",
507 + }
508 +
509 +
510 @pytest.mark.skipif(
511 not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
512 )
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+536 -41
@@ -2,7 +2,11 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import importlib.util
5 +import threading
6 +from collections import OrderedDict, defaultdict
7 +from queue import Queue
8 from types import SimpleNamespace
9 +from typing import Any
10
11 import pytest
12 import torch
@@ -19,16 +23,33 @@ if not (current_platform.is_rocm() and mori_available):
23 allow_module_level=True,
24 )
25
26 +moriio_common = importlib.import_module(
27 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
28 +)
29 +moriio_engine = importlib.import_module(
30 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine"
31 +)
32 moriio_layout = importlib.import_module(
33 "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
34 )
35 +msgpack = importlib.import_module("msgpack")
36 +
37 +ROLE = moriio_common.ROLE
38 +MoRIIOError = moriio_common.MoRIIOError
39 +RemoteAllocInfo = moriio_common.RemoteAllocInfo
40 +WriteTask = moriio_common.WriteTask
41 +set_role = moriio_common.set_role
42 +MoRIIOWrapper = moriio_engine.MoRIIOWrapper
43 +MoRIIOWriter = moriio_engine.MoRIIOWriter
44
45
27 -def _full_spec(block_size: int = 4) -> FullAttentionSpec:
46 +def _full_spec(
47 + block_size: int = 4, num_kv_heads: int = 2, head_size: int = 3
48 +) -> FullAttentionSpec:
49 return FullAttentionSpec(
50 block_size=block_size,
30 - num_kv_heads=2,
31 - head_size=3,
51 + num_kv_heads=num_kv_heads,
52 + head_size=head_size,
53 dtype=torch.bfloat16,
54 )
55
@@ -59,55 +80,216 @@ def _remote_meta(num_blocks: int = 16) -> SimpleNamespace:
80 return SimpleNamespace(num_blocks=num_blocks)
81
82
62 -def test_separated_kv_layout_uses_kv_axis_zero_and_block_axis_one():
63 - cache = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
64 - worker = _worker({"layer": cache}, {"layer": _full_spec()})
65 -
66 - geometry = moriio_layout.get_layer_transfer_geometry(
67 - "layer", cache, worker.layer_to_spec, remote_num_blocks=16
83 +def _writer_with_fake_worker(fake_worker: Any) -> Any:
84 + writer = MoRIIOWriter.__new__(MoRIIOWriter)
85 + writer._worker_ref = lambda: fake_worker
86 + writer._write_task_q = Queue()
87 + writer._write_state_lock = threading.Lock()
88 + writer._scheduled_writes = defaultdict(int)
89 + writer._scheduled_layers = defaultdict(set)
90 + writer._sealed_writes = {}
91 + writer.ensure_worker_started = lambda: None
92 + return writer
93 +
94 +
95 +def _wrapper_for_messages() -> Any:
96 + wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper)
97 + wrapper.lock = threading.Lock()
98 + wrapper.done_remote_allocate_req_dict = {}
99 + wrapper.done_req_ids = []
100 + wrapper.done_write_cache_req_ids = []
101 + wrapper._terminal_transfer_ids = OrderedDict()
102 + return wrapper
103 +
104 +
105 +def _write_task(layer_name: str, transfer_id: str = "xfer") -> Any:
106 + return WriteTask(
107 + request_id="req",
108 + transfer_id=transfer_id,
109 + dst_engine_id="remote-engine",
110 + local_block_ids=[1, 3],
111 + remote_block_ids_hint=None,
112 + layer_name=layer_name,
113 + event=None,
114 + remote_notify_port=7000,
115 + remote_ip="127.0.0.1",
116 )
69 - assert geometry.block_stride == 24
70 - assert geometry.local_kv_stride == 192
71 - assert geometry.remote_kv_stride == 384
72 - assert geometry.split_kv_regions
117
74 - assert moriio_layout.compute_block_transfer_offsets(
75 - "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
76 - ) == ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48])
118
78 -
79 -def test_interleaved_kv_layout_uses_block_axis_zero_and_kv_axis_one():
80 - cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
81 - worker = _worker({"layer": cache}, {"layer": _full_spec()})
119 +@pytest.mark.parametrize(
120 + ("shape", "spec", "remote_num_blocks", "expected_geometry", "expected_offsets"),
121 + [
122 + pytest.param(
123 + (2, 8, 4, 2, 3),
124 + _full_spec(),
125 + 16,
126 + {
127 + "block_stride": 24,
128 + "local_kv_stride": 192,
129 + "remote_kv_stride": 384,
130 + "split_kv_regions": True,
131 + },
132 + ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48]),
133 + id="separated",
134 + ),
135 + pytest.param(
136 + (8, 2, 4, 2, 3),
137 + _full_spec(),
138 + 16,
139 + {
140 + "block_stride": 48,
141 + "local_kv_stride": 24,
142 + "remote_kv_stride": 24,
143 + "split_kv_regions": False,
144 + },
145 + ([96, 288], [384, 480], [96, 96]),
146 + id="interleaved",
147 + ),
148 + pytest.param(
149 + (2, 8, 2, 4, 3),
150 + _full_spec(),
151 + 16,
152 + {
153 + "block_size": 4,
154 + "block_stride": 24,
155 + "local_kv_stride": 192,
156 + "remote_kv_stride": 384,
157 + "split_kv_regions": True,
158 + },
159 + ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48]),
160 + id="shuffled-separated",
161 + ),
162 + pytest.param(
163 + (8, 2, 2, 4, 3),
164 + _full_spec(),
165 + 16,
166 + {
167 + "block_size": 4,
168 + "block_stride": 48,
169 + "local_kv_stride": 24,
170 + "remote_kv_stride": 24,
171 + "split_kv_regions": False,
172 + },
173 + ([96, 288], [384, 480], [96, 96]),
174 + id="shuffled-interleaved",
175 + ),
176 + pytest.param(
177 + (2, 16, 2, 2, 3),
178 + _full_spec(),
179 + 16,
180 + {
181 + "num_blocks": 8,
182 + "block_size": 4,
183 + "block_stride": 24,
184 + "local_kv_stride": 192,
185 + "remote_kv_stride": 384,
186 + "split_kv_regions": True,
187 + },
188 + ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48]),
189 + id="separated-kernel-blocks",
190 + ),
191 + pytest.param(
192 + (16, 2, 2, 2, 3),
193 + _full_spec(),
194 + 16,
195 + {
196 + "num_blocks": 8,
197 + "block_size": 4,
198 + "block_stride": 48,
199 + "local_kv_stride": None,
200 + "remote_kv_stride": None,
201 + "transfers_per_block": 1,
202 + },
203 + ([96, 288], [384, 480], [96, 96]),
204 + id="interleaved-kernel-blocks",
205 + ),
206 + pytest.param(
207 + (2, 32, 8, 2, 3),
208 + _full_spec(block_size=16, num_kv_heads=8),
209 + 8,
210 + {
211 + "num_blocks": 4,
212 + "block_size": 16,
213 + "block_len": 768,
214 + "block_stride": 384,
215 + "local_kv_stride": 1536,
216 + "remote_kv_stride": 3072,
217 + "split_kv_regions": True,
218 + },
219 + (
220 + [768, 2304, 3840, 5376],
221 + [3072, 3840, 9216, 9984],
222 + [768, 768, 768, 768],
223 + ),
224 + id="separated-kernel-axis-from-spec",
225 + ),
226 + pytest.param(
227 + (32, 2, 8, 2, 3),
228 + _full_spec(block_size=16, num_kv_heads=8),
229 + 8,
230 + {
231 + "num_blocks": 4,
232 + "block_size": 16,
233 + "block_len": 1536,
234 + "block_stride": 768,
235 + "local_kv_stride": None,
236 + "remote_kv_stride": None,
237 + "transfers_per_block": 1,
238 + },
239 + ([1536, 4608], [6144, 7680], [1536, 1536]),
240 + id="interleaved-kernel-axis-from-spec",
241 + ),
242 + pytest.param(
243 + (8, 4, 3),
244 + _mla_spec(),
245 + 16,
246 + {
247 + "block_stride": 12,
248 + "local_kv_stride": None,
249 + "remote_kv_stride": None,
250 + "transfers_per_block": 1,
251 + },
252 + ([24, 72], [96, 120], [24, 24]),
253 + id="mla-key-only",
254 + ),
255 + ],
256 +)
257 +def test_supported_layouts_compute_expected_geometry_and_offsets(
258 + shape, spec, remote_num_blocks, expected_geometry, expected_offsets
259 +):
260 + cache = torch.empty(shape, dtype=torch.bfloat16)
261 + worker = _worker({"layer": cache}, {"layer": spec})
262
263 geometry = moriio_layout.get_layer_transfer_geometry(
84 - "layer", cache, worker.layer_to_spec, remote_num_blocks=16
264 + "layer", cache, worker.layer_to_spec, remote_num_blocks=remote_num_blocks
265 )
86 - assert geometry.block_stride == 48
87 - assert geometry.local_kv_stride == 24
88 - assert geometry.remote_kv_stride == 24
89 - assert not geometry.split_kv_regions
90 -
91 - assert moriio_layout.compute_block_transfer_offsets(
92 - "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
93 - ) == ([96, 288], [384, 480], [96, 96])
266 + for field, expected in expected_geometry.items():
267 + assert getattr(geometry, field) == expected
268
269 + assert (
270 + moriio_layout.compute_block_transfer_offsets(
271 + "layer",
272 + cache,
273 + worker.layer_to_spec,
274 + [1, 3],
275 + [4, 5],
276 + remote_num_blocks,
277 + )
278 + == expected_offsets
279 + )
280
96 -def test_mla_key_only_layout_transfers_one_slab_per_block():
97 - cache = torch.empty((8, 4, 3), dtype=torch.bfloat16)
98 - worker = _worker({"layer": cache}, {"layer": _mla_spec()})
281
100 - geometry = moriio_layout.get_layer_transfer_geometry(
101 - "layer", cache, worker.layer_to_spec, remote_num_blocks=16
282 +def test_kernel_block_layout_without_spec_dimensions_rejects_ambiguous_axes():
283 + cache = torch.empty((2, 32, 8, 2, 3), dtype=torch.bfloat16)
284 + worker = _worker(
285 + {"layer": cache},
286 + {"layer": SimpleNamespace(block_size=16)},
287 )
103 - assert geometry.block_stride == 12
104 - assert geometry.local_kv_stride is None
105 - assert geometry.remote_kv_stride is None
106 - assert geometry.transfers_per_block == 1
288
108 - assert moriio_layout.compute_block_transfer_offsets(
109 - "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
110 - ) == ([24, 72], [96, 120], [24, 24])
289 + with pytest.raises(ValueError, match="Ambiguous MoRIIO kernel-block"):
290 + moriio_layout.get_layer_transfer_geometry(
291 + "layer", cache, worker.layer_to_spec, remote_num_blocks=8
292 + )
293
294
295 def test_mixed_layers_compute_distinct_offsets_per_layer():
@@ -155,6 +337,319 @@ def test_mixed_layers_compute_distinct_offsets_per_layer():
337 assert interleaved != indexer
338
339
340 +def test_write_transfer_plan_caches_offsets_per_geometry():
341 + kv_caches = {
342 + "dense0": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
343 + "dense1": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
344 + "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
345 + }
346 + calls: list[str] = []
347 +
348 + class FakeWorker:
349 + kv_caches: dict[str, torch.Tensor]
350 + layer_name_to_local_kv_cache_metadata: dict[str, list[Any]]
351 +
352 + def _compute_block_transfer_offsets(
353 + self, layer_name, local_block_ids, remote_block_ids, remote_moriio_meta
354 + ):
355 + calls.append(layer_name)
356 + call_id = len(calls)
357 + return ([call_id], [call_id + 10], [call_id + 20])
358 +
359 + fake_worker = FakeWorker()
360 + fake_worker.kv_caches = kv_caches
361 + fake_worker.layer_name_to_local_kv_cache_metadata = {name: [] for name in kv_caches}
362 + writer = MoRIIOWriter.__new__(MoRIIOWriter)
363 + writer._worker_ref = lambda: fake_worker
364 + request_info = RemoteAllocInfo(block_ids=[4, 5])
365 + remote_meta = _remote_meta()
366 +
367 + dense0_plan = writer._prepare_transfer_plan(
368 + SimpleNamespace(
369 + layer_name="dense0",
370 + local_block_ids=[1, 3],
371 + request_id="req",
372 + transfer_id="xfer",
373 + ),
374 + request_info,
375 + remote_meta,
376 + )
377 + dense1_plan = writer._prepare_transfer_plan(
378 + SimpleNamespace(
379 + layer_name="dense1",
380 + local_block_ids=[1, 3],
381 + request_id="req",
382 + transfer_id="xfer",
383 + ),
384 + request_info,
385 + remote_meta,
386 + )
387 + indexer_plan = writer._prepare_transfer_plan(
388 + SimpleNamespace(
389 + layer_name="indexer",
390 + local_block_ids=[1, 3],
391 + request_id="req",
392 + transfer_id="xfer",
393 + ),
394 + request_info,
395 + remote_meta,
396 + )
397 +
398 + assert calls == ["dense0", "indexer"]
399 + assert dense0_plan.transfer_local_offsets == [1]
400 + assert dense1_plan.transfer_local_offsets == [1]
401 + assert indexer_plan.transfer_local_offsets == [2]
402 + assert len(request_info.transfer_offsets) == 2
403 +
404 +
405 +def test_write_scheduler_deduplicates_layers_and_seals_expected_count():
406 + request_info = RemoteAllocInfo(block_ids=[4, 5])
407 + wrapper = _wrapper_for_messages()
408 + wrapper.done_remote_allocate_req_dict["xfer"] = request_info
409 + writer = _writer_with_fake_worker(SimpleNamespace(moriio_wrapper=wrapper))
410 +
411 + assert writer.schedule_write(_write_task("dense0"))
412 + assert not writer.schedule_write(_write_task("dense0"))
413 + assert writer.schedule_write(_write_task("indexer"))
414 +
415 + assert writer._write_task_q.qsize() == 2
416 + writer.seal_pending_transfers()
417 +
418 + assert request_info.writes_expected == 2
419 + assert writer._sealed_writes["xfer"] == 2
420 +
421 +
422 +def test_write_completion_notifies_once_after_all_sealed_writes_finish():
423 + class FakeWrapper:
424 + def __init__(self):
425 + self.done_remote_allocate_req_dict = {}
426 + self.done_req_ids = []
427 + self.lock = threading.Lock()
428 + self.notifications = []
429 + self.wait_count = 0
430 + self.waited_statuses = []
431 + self._terminal_transfer_ids = OrderedDict()
432 +
433 + def waiting_for_transfer_complete(self, transfer_statuses=None):
434 + self.wait_count += 1
435 + self.waited_statuses.append(list(transfer_statuses or []))
436 +
437 + def _is_transfer_terminal_locked(self, transfer_id):
438 + return transfer_id in self._terminal_transfer_ids
439 +
440 + def _mark_transfer_terminal_locked(self, transfer_id):
441 + self._terminal_transfer_ids[transfer_id] = None
442 +
443 + def send_notify(self, transfer_id, remote_ip, remote_port, message_type=None):
444 + self.notifications.append(
445 + (transfer_id, remote_ip, remote_port, message_type)
446 + )
447 +
448 + wrapper = FakeWrapper()
449 + request_info = RemoteAllocInfo(block_ids=[4, 5], writes_expected=2)
450 + request_info.transfer_statuses.extend(["status-a", "status-b"])
451 + request_info.completion_request_id = "req"
452 + request_info.completion_remote_notify_port = 7000
453 + request_info.completion_remote_ip = "127.0.0.1"
454 + wrapper.done_remote_allocate_req_dict["xfer"] = request_info
455 + writer = _writer_with_fake_worker(
456 + SimpleNamespace(moriio_wrapper=wrapper, tp_rank=2)
457 + )
458 + writer._scheduled_writes["xfer"] = 2
459 + writer._scheduled_layers["xfer"] = {"dense0", "indexer"}
460 + writer._sealed_writes["xfer"] = 2
461 +
462 + writer._mark_write_done("xfer", request_info)
463 + assert wrapper.notifications == []
464 + writer._mark_write_done("xfer", request_info)
465 + writer._finalize_if_complete("xfer", request_info)
466 +
467 + assert wrapper.notifications == [("xfer", "127.0.0.1", 7002, "write_done")]
468 + assert wrapper.done_req_ids == ["xfer"]
469 + assert wrapper.done_remote_allocate_req_dict == {}
470 + assert wrapper.wait_count == 1
471 + assert wrapper.waited_statuses == [["status-a", "status-b"]]
472 + assert request_info.transfer_statuses == []
473 + assert wrapper._is_transfer_terminal_locked("xfer")
474 +
475 +
476 +def test_moriio_wrapper_waits_scoped_statuses_without_global_drain():
477 + class FakeStatus:
478 + def __init__(self):
479 + self.checked = 0
480 +
481 + def Succeeded(self):
482 + self.checked += 1
483 + return True
484 +
485 + def Failed(self):
486 + return False
487 +
488 + wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper)
489 + wrapper.lock = threading.Lock()
490 + wrapper._transfer_timeout = 1
491 + global_status = FakeStatus()
492 + scoped_status = FakeStatus()
493 + wrapper.transfer_status = [global_status]
494 +
495 + wrapper.waiting_for_transfer_complete([scoped_status])
496 +
497 + assert scoped_status.checked == 1
498 + assert global_status.checked == 0
499 + assert wrapper.transfer_status == [global_status]
500 +
501 +
502 +def test_write_failure_marks_terminal_and_clears_scheduled_state():
503 + wrapper = _wrapper_for_messages()
504 + wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5])
505 + writer = _writer_with_fake_worker(SimpleNamespace(moriio_wrapper=wrapper))
506 + writer._scheduled_writes["xfer"] = 2
507 + writer._scheduled_layers["xfer"] = {"dense0", "indexer"}
508 + writer._sealed_writes["xfer"] = 2
509 +
510 + writer._mark_request_done("xfer")
511 +
512 + assert wrapper.done_req_ids == ["xfer"]
513 + assert wrapper.done_remote_allocate_req_dict == {}
514 + assert wrapper._is_transfer_terminal_locked("xfer")
515 + assert "xfer" not in writer._scheduled_writes
516 + assert "xfer" not in writer._scheduled_layers
517 + assert "xfer" not in writer._sealed_writes
518 +
519 +
520 +def test_schedule_write_rejects_terminal_transfer_without_recreating_state():
521 + wrapper = _wrapper_for_messages()
522 + wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5])
523 + writer = _writer_with_fake_worker(SimpleNamespace(moriio_wrapper=wrapper))
524 + writer._scheduled_writes["xfer"] = 1
525 + writer._scheduled_layers["xfer"] = {"dense0"}
526 + writer._sealed_writes["xfer"] = 1
527 +
528 + writer._mark_request_done("xfer")
529 +
530 + assert not writer.schedule_write(_write_task("indexer"))
531 + assert writer._write_task_q.empty()
532 + assert "xfer" not in writer._scheduled_writes
533 + assert "xfer" not in writer._scheduled_layers
534 + assert "xfer" not in writer._sealed_writes
535 +
536 +
537 +def test_late_remote_blocks_message_is_ignored_after_transfer_done():
538 + set_role(ROLE.PRODUCER)
539 + wrapper = _wrapper_for_messages()
540 + with wrapper.lock:
541 + wrapper._mark_transfer_terminal_locked("xfer")
542 +
543 + wrapper._handle_message(
544 + msgpack.dumps(
545 + {
546 + "type": "remote_blocks",
547 + "req_id": "req",
548 + "transfer_id": "xfer",
549 + "block_notify_list": [4, 5],
550 + "decode_rank": 3,
551 + }
552 + )
553 + )
554 +
555 + assert "xfer" not in wrapper.done_remote_allocate_req_dict
556 +
557 +
558 +@pytest.mark.parametrize(
559 + ("role", "payload", "expected"),
560 + [
561 + pytest.param(
562 + ROLE.PRODUCER,
563 + msgpack.dumps(
564 + {
565 + "type": "remote_blocks",
566 + "req_id": "req",
567 + "transfer_id": "xfer",
568 + "block_notify_list": [4, 5],
569 + "decode_rank": 3,
570 + }
571 + ),
572 + "remote_blocks",
573 + id="remote-blocks",
574 + ),
575 + pytest.param(
576 + ROLE.CONSUMER,
577 + msgpack.dumps({"type": "write_done", "transfer_id": "xfer"}),
578 + "write_done",
579 + id="write-done",
580 + ),
581 + pytest.param(
582 + ROLE.PRODUCER,
583 + msgpack.dumps({"type": "release", "transfer_id": "xfer"}),
584 + "release",
585 + id="release",
586 + ),
587 + pytest.param(None, b"xfer", "plain", id="plain-string"),
588 + ],
589 +)
590 +def test_moriio_wrapper_routes_valid_messages(role, payload, expected):
591 + wrapper = _wrapper_for_messages()
592 + completions: list[str] = []
593 + if role is not None:
594 + set_role(role)
595 + if expected == "plain":
596 + wrapper._handle_completion_message = completions.append
597 +
598 + wrapper._handle_message(payload)
599 +
600 + if expected == "remote_blocks":
601 + request_info = wrapper.done_remote_allocate_req_dict["xfer"]
602 + assert request_info.block_ids == [4, 5]
603 + assert request_info.decode_dp_rank == 3
604 + elif expected == "write_done":
605 + 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:
610 + assert completions == ["xfer"]
611 +
612 +
613 +@pytest.mark.parametrize(
614 + ("role", "payload", "match"),
615 + [
616 + pytest.param(
617 + None,
618 + msgpack.dumps({"type": "unknown", "transfer_id": "xfer"}),
619 + "Unhandled structured message type",
620 + id="unknown-structured-type",
621 + ),
622 + pytest.param(
623 + ROLE.PRODUCER,
624 + msgpack.dumps(
625 + {
626 + "type": "remote_blocks",
627 + "req_id": "req",
628 + "transfer_id": "xfer",
629 + "block_notify_list": [],
630 + }
631 + ),
632 + "block_notify_list cannot be empty",
633 + id="empty-remote-blocks",
634 + ),
635 + pytest.param(
636 + None,
637 + b"",
638 + "Unhandled message format",
639 + id="empty-completion",
640 + ),
641 + ],
642 +)
643 +def test_moriio_wrapper_rejects_invalid_messages(role, payload, match):
644 + wrapper = _wrapper_for_messages()
645 + if role is not None:
646 + set_role(role)
647 + wrapper._handle_completion_message = lambda msg: None
648 +
649 + with pytest.raises(MoRIIOError, match=match):
650 + wrapper._handle_message(payload)
651 +
652 +
653 def test_block_id_length_mismatch_raises_value_error():
654 cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
655 worker = _worker({"layer": cache}, {"layer": _full_spec()})
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+28 -10
@@ -78,8 +78,17 @@ class RemoteAllocInfo:
78
79 block_ids: list[int]
80 writes_done: int = 0
81 + writes_expected: int | None = None
82 decode_dp_rank: int = 0
82 - transfer_offset: tuple[list[int], list[int], list[int]] | None = None
83 + completion_request_id: str | None = None
84 + completion_remote_notify_port: int | None = None
85 + completion_remote_ip: str | None = None
86 + completion_notified: bool = False
87 + transfer_statuses: list[Any] = field(default_factory=list)
88 + transfer_offsets: dict[
89 + tuple[tuple[int, ...], tuple[int, ...], torch.dtype],
90 + tuple[list[int], list[int], list[int]],
91 + ] = field(default_factory=dict)
92
93
94 class ROLE(Enum):
@@ -434,12 +443,21 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
443 ):
444 transfer_id = kv_transfer_params["transfer_id"]
445
437 - # Parse host/ports from the request_id. The router embeds both zmq_addresses
438 - # in the request_id
439 - peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
440 - remote_host, remote_handshake_port, remote_notify_port = (
441 - parse_moriio_zmq_address(peer_zmq)
442 - )
446 + remote_host = kv_transfer_params.get("remote_host")
447 + remote_handshake_port = kv_transfer_params.get("remote_handshake_port")
448 + remote_notify_port = kv_transfer_params.get("remote_notify_port")
449 + if (
450 + remote_host is None
451 + or remote_handshake_port is None
452 + or remote_notify_port is None
453 + ):
454 + # Parse host/ports from the request_id. The router embeds both
455 + # zmq_addresses in PD request IDs, but WRITE decode requests may carry
456 + # a plain request ID and get the remote address via kv_transfer_params.
457 + peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
458 + remote_host, remote_handshake_port, remote_notify_port = (
459 + parse_moriio_zmq_address(peer_zmq)
460 + )
461
462 _req = ReqMeta(
463 transfer_id=transfer_id,
@@ -447,9 +465,9 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
465 remote_block_ids=kv_transfer_params["remote_block_ids"],
466 remote_engine_id=kv_transfer_params["remote_engine_id"],
467 remote_host=remote_host,
450 - remote_port=remote_handshake_port,
451 - remote_handshake_port=remote_handshake_port,
452 - remote_notify_port=remote_notify_port,
468 + remote_port=int(remote_handshake_port),
469 + remote_handshake_port=int(remote_handshake_port),
470 + remote_notify_port=int(remote_notify_port),
471 tp_size=kv_transfer_params.get("tp_size", 1),
472 remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
473 )
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+109 -64
@@ -234,7 +234,13 @@ class MoRIIOConnector(KVConnectorBase_V1):
234 return None
235
236 def wait_for_save(self):
237 - pass
237 + if self.mode != MoRIIOMode.WRITE or get_role() != ROLE.PRODUCER:
238 + return
239 + assert self.connector_worker is not None
240 + assert isinstance(self._connector_metadata, MoRIIOConnectorMetadata), (
241 + "Connector metadata not initialized yet"
242 + )
243 + self.connector_worker.wait_for_save(self._connector_metadata)
244
245 def shutdown(self):
246 if self.connector_worker is not None:
@@ -390,6 +396,49 @@ class MoRIIOConnectorScheduler:
396 serialized_data = msgpack.dumps(data)
397 self.paths[path].send(serialized_data)
398
399 + def _send_transfer_release(self, transfer_id: TransferId, host: str, port: int):
400 + path = make_zmq_path("tcp", host, port)
401 + if path not in self.paths:
402 + ctx = zmq.Context.instance()
403 + sock = make_zmq_socket(
404 + ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False
405 + )
406 + self.paths[path] = sock
407 +
408 + self.paths[path].send(
409 + msgpack.dumps({"type": "release", "transfer_id": transfer_id})
410 + )
411 +
412 + def _release_write_prefill_blocks(self, request_id: ReqId, params: dict[str, Any]):
413 + transfer_id = params.get("transfer_id")
414 + if transfer_id is None:
415 + logger.warning(
416 + "Cannot release WRITE prefill blocks for request %s: "
417 + "missing transfer_id",
418 + request_id,
419 + )
420 + return
421 +
422 + remote_dp_rank = params.get("remote_dp_rank", 0)
423 + remote_host = params.get("remote_host")
424 + remote_notify_port = params.get("remote_notify_port")
425 + if remote_host is None or remote_notify_port is None:
426 + try:
427 + peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False)
428 + remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)
429 + except ValueError:
430 + logger.warning(
431 + "Cannot release WRITE prefill blocks for request %s: "
432 + "missing remote notify address",
433 + request_id,
434 + )
435 + return
436 +
437 + remote_notify_port = int(remote_notify_port)
438 + for tp_index in range(self.tp_size):
439 + target_port = remote_notify_port + get_port_offset(remote_dp_rank, tp_index)
440 + self._send_transfer_release(transfer_id, remote_host, target_port)
441 +
442 def update_state_after_alloc(
443 self,
444 request: "Request",
@@ -443,11 +492,20 @@ class MoRIIOConnectorScheduler:
492 )
493
494 remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0)
446 -
447 - peer_zmq = get_peer_zmq_from_request_id(
448 - request.request_id, is_producer=False
495 + remote_host = request.kv_transfer_params.get("remote_host")
496 + remote_notify_port = request.kv_transfer_params.get(
497 + "remote_notify_port"
498 )
450 - remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)
499 + if remote_host is None or remote_notify_port is None:
500 + peer_zmq = get_peer_zmq_from_request_id(
501 + request.request_id, is_producer=False
502 + )
503 + remote_host, _, remote_notify_port = parse_moriio_zmq_address(
504 + peer_zmq
505 + )
506 + remote_notify_port = int(remote_notify_port)
507 +
508 + block_ids = blocks.get_block_ids()[0]
509
510 for tp_index in range(self.tp_size):
511 target_port = remote_notify_port + get_port_offset(
@@ -457,7 +515,7 @@ class MoRIIOConnectorScheduler:
515 self.send_notify_block(
516 req_id=request.request_id,
517 transfer_id=request.kv_transfer_params["transfer_id"],
460 - block_notify_list=blocks.get_block_ids()[0],
518 + block_notify_list=block_ids,
519 host=remote_host,
520 port=target_port,
521 )
@@ -473,60 +531,31 @@ class MoRIIOConnectorScheduler:
531 meta = MoRIIOConnectorMetadata()
532 meta.transfer_id_to_request_id = self.transfer_id_to_request_id
533
476 - if self.mode == MoRIIOMode.WRITE:
477 - # when async_load_kv finished,
478 - # new reqs will be added to scheduler_output.scheduled_new_reqs
479 -
480 - if get_role() == ROLE.CONSUMER:
481 - for new_req in scheduler_output.scheduled_new_reqs:
482 - red_id = new_req.req_id
483 - local_block_ids = list(new_req.block_ids)[0]
484 - assert new_req.sampling_params is not None, (
485 - f"sampling_params is None for req {new_req.req_id}"
486 - )
487 - assert hasattr(new_req.sampling_params, "extra_args"), (
488 - f"sampling_params missing extra_args for req {new_req.req_id}"
489 - )
490 - kv_transfer_params = (
491 - new_req.sampling_params.extra_args.get("kv_transfer_params", {})
492 - if new_req.sampling_params.extra_args
493 - else {}
494 - )
495 - meta.add_new_req(
496 - red_id,
497 - local_block_ids,
498 - kv_transfer_params,
499 - )
500 - if get_role() == ROLE.PRODUCER:
501 - # This is the logic for checking against chunked prefill.
502 - # When the last chunk is identified,
503 - # It places the request metadata into the saving queue.
504 -
505 - for i, req_id in enumerate(
506 - scheduler_output.scheduled_cached_reqs.req_ids
507 - ):
508 - new_block_ids = (
509 - scheduler_output.scheduled_cached_reqs.new_block_ids[i]
510 - )
511 -
512 - if new_block_ids is not None:
513 - block_ids = new_block_ids[0]
514 - # TODO : hybrid attn, etc
515 - req, existing_blocks = self._reqs_need_pending_save[req_id]
516 - updated_blocks = list(existing_blocks) + (block_ids)
517 - self._reqs_need_pending_save[req_id] = (req, updated_blocks)
518 - if (
519 - len(self._reqs_need_pending_save[req_id][1])
520 - * self.block_size
521 - >= req.num_prompt_tokens
522 - ):
523 - meta.add_new_req(
524 - request_id=req_id,
525 - local_block_ids=self._reqs_need_pending_save[req_id][1],
526 - kv_transfer_params=req.kv_transfer_params or {},
527 - write_mode=True,
528 - )
529 - del self._reqs_need_pending_save[req_id]
534 + if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.PRODUCER:
535 + # This is the logic for checking against chunked prefill.
536 + # When the last chunk is identified,
537 + # It places the request metadata into the saving queue.
538 +
539 + for i, req_id in enumerate(scheduler_output.scheduled_cached_reqs.req_ids):
540 + new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i]
541 +
542 + if new_block_ids is not None:
543 + block_ids = new_block_ids[0]
544 + # TODO : hybrid attn, etc
545 + req, existing_blocks = self._reqs_need_pending_save[req_id]
546 + updated_blocks = list(existing_blocks) + (block_ids)
547 + self._reqs_need_pending_save[req_id] = (req, updated_blocks)
548 + if (
549 + len(self._reqs_need_pending_save[req_id][1]) * self.block_size
550 + >= req.num_prompt_tokens
551 + ):
552 + meta.add_new_req(
553 + request_id=req_id,
554 + local_block_ids=self._reqs_need_pending_save[req_id][1],
555 + kv_transfer_params=req.kv_transfer_params or {},
556 + write_mode=True,
557 + )
558 + del self._reqs_need_pending_save[req_id]
559
560 # Loop through scheduled reqs and convert to ReqMeta.
561 for req_id, (req, block_ids) in self._reqs_need_recv.items():
@@ -601,9 +630,15 @@ class MoRIIOConnectorScheduler:
630 # update_state_after_alloc must not have been called (the request
631 # must have been aborted before it was scheduled).
632 # To avoid stranding the prefill blocks in the prefill instance,
604 - # we must add empty block_ids to _reqs_need_recv so that our
605 - # worker side will notify and free blocks in the prefill instance.
606 - self._reqs_need_recv[request.request_id] = (request, [])
633 + # READ mode adds empty block_ids to _reqs_need_recv so the worker
634 + # side notifies the prefill instance. WRITE mode should notify the
635 + # producer directly: there is no decode allocation for the producer
636 + # to write into, and a plain request_id may not contain router-
637 + # embedded MoRIIO ZMQ addresses.
638 + if self.mode == MoRIIOMode.WRITE:
639 + self._release_write_prefill_blocks(request.request_id, params)
640 + else:
641 + self._reqs_need_recv[request.request_id] = (request, [])
642 params["do_remote_prefill"] = False
643 return False, None
644
@@ -635,6 +670,10 @@ class MoRIIOConnectorScheduler:
670 do_remote_decode=False,
671 remote_block_ids=computed_block_ids,
672 remote_engine_id=self.engine_id,
673 + remote_host=self.host_ip,
674 + remote_handshake_port=self.handshake_port,
675 + remote_notify_port=self.side_notify_port,
676 + remote_dp_size=self.vllm_config.parallel_config.data_parallel_size,
677 tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
678 transfer_id=params["transfer_id"],
679 )
@@ -1480,7 +1519,7 @@ class MoRIIOConnectorWorker:
1519 metadata: MoRIIOConnectorMetadata,
1520 layer_name: str,
1521 kv_layer: torch.Tensor,
1483 - attn_metadata: "AttentionMetadata",
1522 + attn_metadata: "AttentionMetadata | None",
1523 **kwargs,
1524 ):
1525 if not self.is_producer:
@@ -1604,6 +1643,12 @@ class MoRIIOConnectorWorker:
1643
1644 self._reqs_to_send.update(metadata.reqs_to_send)
1645
1646 + def wait_for_save(self, metadata: MoRIIOConnectorMetadata):
1647 + if self.mode == MoRIIOMode.WRITE and self.is_producer:
1648 + for layer_name, kv_layer in self.kv_caches.items():
1649 + self.save_kv_layer(metadata, layer_name, kv_layer, None)
1650 + self._writer.seal_pending_transfers()
1651 +
1652 def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
1653 logger.debug(
1654 "Remote agent %s available, calling _read_blocks for req %s",
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+247 -73
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import threading
4 import time
5 +from collections import OrderedDict, defaultdict
6 +from queue import Empty, Queue
7 from typing import TYPE_CHECKING, Any
8 from weakref import ref as weakref_ref
9
@@ -18,8 +20,6 @@ from vllm.utils.network_utils import (
20 if TYPE_CHECKING:
21 from mori.io import BackendType
22
21 -from queue import Empty, Queue
22 -
23 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
24 ROLE,
25 HandshakeError,
@@ -61,10 +61,25 @@ except ImportError:
61 """Write task execution logic for MoRIIO connector."""
62
63
64 +_MAX_TERMINAL_TRANSFER_IDS = 4096
65 +
66 +
67 +WriteGeometryKey = tuple[tuple[int, ...], tuple[int, ...], torch.dtype]
68 +
69 +
70 +def _get_write_geometry_key(kv_cache: torch.Tensor) -> WriteGeometryKey:
71 + return (tuple(kv_cache.shape), tuple(kv_cache.stride()), kv_cache.dtype)
72 +
73 +
74 class MoRIIOWriter:
75 """Handles write operations for KV cache transfers.
66 - Implements distributed KV cache transfer using the MoRIIO library
67 - for RDMA-based communication between prefill and decode instances."""
76 +
77 + WRITE mode state machine:
78 + D sends destination block allocation, P schedules one write per layer
79 + after the layer CUDA event, P seals the scheduled write count after
80 + forward, then P notifies D and releases P blocks after all scheduled
81 + writes complete.
82 + """
83
84 def __init__(self, worker: "MoRIIOConnectorWorker"):
85 """Initialize the writer.
@@ -76,7 +91,11 @@ class MoRIIOWriter:
91 self._write_task_q: Queue[WriteTask] = Queue()
92 self._write_worker_started = False
93 self._write_worker_lock = threading.Lock()
94 + self._write_state_lock = threading.Lock()
95 self._deferred_tasks: list[WriteTask] = []
96 + self._scheduled_writes: dict[TransferId, int] = defaultdict(int)
97 + self._scheduled_layers: dict[TransferId, set[str]] = defaultdict(set)
98 + self._sealed_writes: dict[TransferId, int] = {}
99 self._defer_timeout = worker.moriio_config.defer_timeout
100
101 @property
@@ -106,14 +125,55 @@ class MoRIIOWriter:
125 thread.start()
126 logger.info("Started MoRIIO write worker thread")
127
109 - def schedule_write(self, task: WriteTask) -> None:
128 + def schedule_write(self, task: WriteTask) -> bool:
129 """Schedule a write task.
130
131 Args:
132 task: The write task to schedule
133 """
134 self.ensure_worker_started()
135 + if self._is_transfer_terminal(task.transfer_id):
136 + return False
137 +
138 + with self._write_state_lock:
139 + if self._is_transfer_terminal(task.transfer_id):
140 + return False
141 + if task.layer_name in self._scheduled_layers[task.transfer_id]:
142 + return False
143 + self._scheduled_layers[task.transfer_id].add(task.layer_name)
144 + self._scheduled_writes[task.transfer_id] += 1
145 self._write_task_q.put(task)
146 + return True
147 +
148 + def is_scheduled(self, transfer_id: TransferId, layer_name: str) -> bool:
149 + with self._write_state_lock:
150 + return layer_name in self._scheduled_layers.get(transfer_id, set())
151 +
152 + def seal_pending_transfers(self) -> None:
153 + """Seal expected WRITE counts after the model forward has run.
154 +
155 + `save_kv_layer` is only invoked for attention layers whose backend uses
156 + the standard KV connector hook. Hybrid models can register more KV
157 + cache tensors than the number of hooks that fire in a forward, so WRITE
158 + completion must be based on the tasks actually queued for the transfer.
159 + """
160 + pending: list[tuple[TransferId, RemoteAllocInfo]] = []
161 + with self._write_state_lock:
162 + for transfer_id, write_count in self._scheduled_writes.items():
163 + if transfer_id in self._sealed_writes:
164 + continue
165 + self._sealed_writes[transfer_id] = write_count
166 + request_info = (
167 + self.worker.moriio_wrapper.done_remote_allocate_req_dict.get(
168 + transfer_id
169 + )
170 + )
171 + if request_info is not None:
172 + request_info.writes_expected = write_count
173 + pending.append((transfer_id, request_info))
174 +
175 + for transfer_id, request_info in pending:
176 + self._finalize_if_complete(transfer_id, request_info)
177
178 def _write_worker_loop(self) -> None:
179 """Main loop for the write worker thread."""
@@ -128,6 +188,9 @@ class MoRIIOWriter:
188 except Empty:
189 continue
190
191 + if self._is_transfer_terminal(task.transfer_id):
192 + continue
193 +
194 # Check if remote blocks are ready
195 if not self._is_remote_ready(task):
196 # task.retry_count += 1
@@ -158,6 +221,8 @@ class MoRIIOWriter:
221 still_deferred: list[WriteTask] = []
222
223 for task in self._deferred_tasks:
224 + if self._is_transfer_terminal(task.transfer_id):
225 + continue
226 if now - task.enqueue_time > defer_timeout:
227 logger.error(
228 "Deferred write task for request %s expired after %.1fs "
@@ -181,12 +246,25 @@ class MoRIIOWriter:
246
247 self._deferred_tasks = still_deferred
248
249 + def _clear_transfer_state(self, transfer_id: TransferId) -> None:
250 + with self._write_state_lock:
251 + self._scheduled_writes.pop(transfer_id, None)
252 + self._scheduled_layers.pop(transfer_id, None)
253 + self._sealed_writes.pop(transfer_id, None)
254 +
255 + def _is_transfer_terminal(self, transfer_id: TransferId) -> bool:
256 + wrapper = self.worker.moriio_wrapper
257 + with wrapper.lock:
258 + return wrapper._is_transfer_terminal_locked(transfer_id)
259 +
260 def _mark_request_done(self, transfer_id: str) -> None:
261 """Mark a request done so its blocks are freed, even on transfer failure."""
262 wrapper = self.worker.moriio_wrapper
263 with wrapper.lock:
264 wrapper.done_req_ids.append(transfer_id)
189 - wrapper.done_remote_allocate_req_dict.pop(transfer_id, None)
265 + wrapper.done_remote_allocate_req_dict.pop(transfer_id, None)
266 + wrapper._mark_transfer_terminal_locked(transfer_id)
267 + self._clear_transfer_state(transfer_id)
268
269 def _is_remote_ready(self, task: WriteTask) -> bool:
270 """Check if remote blocks are allocated for this task.
@@ -229,6 +307,12 @@ class MoRIIOWriter:
307 """
308 # Get remote allocation info
309 request_info = self._get_remote_alloc_info(task.transfer_id)
310 + with self._write_state_lock:
311 + request_info.completion_request_id = task.request_id
312 + request_info.completion_remote_notify_port = task.remote_notify_port
313 + request_info.completion_remote_ip = task.remote_ip
314 + if task.transfer_id in self._sealed_writes:
315 + request_info.writes_expected = self._sealed_writes[task.transfer_id]
316
317 if request_info.block_ids is None:
318 logger.debug(
@@ -259,10 +343,12 @@ class MoRIIOWriter:
343 plan = self._prepare_transfer_plan(task, request_info, remote_moriio_meta)
344
345 # Execute transfer
262 - self._do_layer_write(plan, sessions)
346 + transfer_statuses = self._do_layer_write(plan, sessions)
347 + with self._write_state_lock:
348 + request_info.transfer_statuses.extend(transfer_statuses)
349
350 # Finalize if all layers complete
265 - self._finalize_if_complete(task, request_info)
351 + self._mark_write_done(task.transfer_id, request_info)
352
353 def _prepare_transfer_plan(
354 self,
@@ -279,21 +365,23 @@ class MoRIIOWriter:
365 Returns:
366 The transfer plan
367 """
282 - # Compute offsets if not cached
283 - if request_info.transfer_offset is None:
368 + layer_cache = self.worker.kv_caches[task.layer_name]
369 + geometry_key = _get_write_geometry_key(layer_cache)
370 + offsets = request_info.transfer_offsets.get(geometry_key)
371 + if offsets is None:
372 offsets = self.worker._compute_block_transfer_offsets(
373 task.layer_name,
374 task.local_block_ids,
375 request_info.block_ids,
376 remote_moriio_meta,
377 )
290 - request_info.transfer_offset = offsets
378 + request_info.transfer_offsets[geometry_key] = offsets
379
380 # Get session index
381 layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys())
382 sess_idx = layer_names.index(task.layer_name)
383
296 - local_off, remote_off, sizes = request_info.transfer_offset
384 + local_off, remote_off, sizes = offsets
385
386 return LayerTransferPlan(
387 request_id=task.request_id,
@@ -306,7 +394,7 @@ class MoRIIOWriter:
394 use_batch=True,
395 )
396
309 - def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> None:
397 + def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]:
398 """Perform the actual layer write.
399
400 Args:
@@ -314,59 +402,82 @@ class MoRIIOWriter:
402 sessions: List of transfer sessions
403 """
404 if plan.use_batch:
317 - self.worker.moriio_wrapper.write_remote_data(
318 - plan.transfer_sizes,
319 - plan.transfer_local_offsets,
320 - plan.transfer_remote_offsets,
321 - sessions[plan.sess_idx],
322 - )
323 - else:
324 - for i in range(len(plan.transfer_local_offsets)):
405 + return [
406 + self.worker.moriio_wrapper.write_remote_data(
407 + plan.transfer_sizes,
408 + plan.transfer_local_offsets,
409 + plan.transfer_remote_offsets,
410 + sessions[plan.sess_idx],
411 + )
412 + ]
413 +
414 + transfer_statuses: list[Any] = []
415 + for i in range(len(plan.transfer_local_offsets)):
416 + transfer_statuses.append(
417 self.worker.moriio_wrapper.write_remote_data_single(
418 plan.transfer_sizes[i],
419 plan.transfer_local_offsets[i],
420 plan.transfer_remote_offsets[i],
421 plan.sess_idx,
422 )
423 + )
424 + return transfer_statuses
425
332 - def _finalize_if_complete(
333 - self, task: WriteTask, request_info: RemoteAllocInfo
426 + def _mark_write_done(
427 + self, transfer_id: TransferId, request_info: RemoteAllocInfo
428 ) -> None:
335 - """Finalize transfer if all layers are complete.
429 + """Record one completed WRITE task and finalize if sealed."""
430 + with self._write_state_lock:
431 + request_info.writes_done += 1
432 + self._finalize_if_complete(transfer_id, request_info)
433
337 - Args:
338 - task: The write task
339 - request_info: Remote allocation information
340 - """
341 - request_info.writes_done += 1
434 + def _finalize_if_complete(
435 + self, transfer_id: TransferId, request_info: RemoteAllocInfo
436 + ) -> None:
437 + """Finalize transfer if all scheduled writes are complete."""
438 + with self._write_state_lock:
439 + expected = request_info.writes_expected
440 + if expected is None or request_info.writes_done < expected:
441 + return
442 + if request_info.completion_notified:
443 + return
444 + request_id = request_info.completion_request_id
445 + remote_notify_port = request_info.completion_remote_notify_port
446 + remote_ip = request_info.completion_remote_ip
447 + if request_id is None or remote_notify_port is None or remote_ip is None:
448 + return
449 + transfer_statuses = list(request_info.transfer_statuses)
450 + request_info.transfer_statuses.clear()
451 + request_info.completion_notified = True
452
343 - if request_info.writes_done >= self.worker.num_layers:
344 - # Wait for transfer to complete
345 - self.worker.moriio_wrapper.waiting_for_transfer_complete()
453 + # Wait for this request's transfers to complete.
454 + self.worker.moriio_wrapper.waiting_for_transfer_complete(transfer_statuses)
455
347 - remote_port = task.remote_notify_port + get_port_offset(
348 - request_info.decode_dp_rank, self.worker.tp_rank
349 - )
350 - # Consider using RDMA immediate data in decode side
351 - # to eliminate the need for this notification.
352 - # Consider including the first gen token from prefill in the notification
456 + remote_port = remote_notify_port + get_port_offset(
457 + request_info.decode_dp_rank, self.worker.tp_rank
458 + )
459 + # Consider using RDMA immediate data in decode side
460 + # to eliminate the need for this notification.
461 + # Consider including the first gen token from prefill in the notification
462
354 - # Send completion notification
355 - self.worker.moriio_wrapper.send_notify(
356 - task.transfer_id, task.remote_ip, remote_port
357 - )
358 - # mark request as done, then we can free the blocks
359 - with self.worker.moriio_wrapper.lock:
360 - self.worker.moriio_wrapper.done_req_ids.append(task.transfer_id)
361 - del self.worker.moriio_wrapper.done_remote_allocate_req_dict[
362 - task.transfer_id
363 - ]
364 - logger.debug(
365 - "Completed transfer for (request, transfer) %s, %s, notified port %d",
366 - task.request_id,
367 - task.transfer_id,
368 - remote_port,
463 + # Send completion notification
464 + self.worker.moriio_wrapper.send_notify(
465 + transfer_id, remote_ip, remote_port, message_type="write_done"
466 + )
467 + # mark request as done, then we can free the blocks
468 + with self.worker.moriio_wrapper.lock:
469 + self.worker.moriio_wrapper.done_req_ids.append(transfer_id)
470 + self.worker.moriio_wrapper.done_remote_allocate_req_dict.pop(
471 + transfer_id, None
472 )
473 + self.worker.moriio_wrapper._mark_transfer_terminal_locked(transfer_id)
474 + self._clear_transfer_state(transfer_id)
475 + logger.debug(
476 + "Completed transfer for (request, transfer) %s, %s, notified port %d",
477 + request_id,
478 + transfer_id,
479 + remote_port,
480 + )
481
482
483 class MoRIIOWrapper:
@@ -400,6 +511,7 @@ class MoRIIOWrapper:
511 self.done_req_ids: list[str] = []
512 self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {}
513 self.done_write_cache_req_ids: list[str] = []
514 + self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict()
515 self._transfer_timeout = transfer_timeout
516 self.notify_thread: threading.Thread | None = None
517 self.sessions: list[IOEngine.Session] = []
@@ -506,8 +618,7 @@ class MoRIIOWrapper:
618 transfer_status = session.batch_write(
619 local_offset, remote_offset, transfer_size_byte, write_uid
620 )
509 - with self.lock:
510 - self.transfer_status.append(transfer_status)
621 + return transfer_status
622
623 def write_remote_data_single(
624 self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0
@@ -520,16 +631,18 @@ class MoRIIOWrapper:
631 transfer_size_byte,
632 self.moriio_engine.allocate_transfer_uid(),
633 )
523 - with self.lock:
524 - self.transfer_status.append(transfer_status)
634 + return transfer_status
635
526 - def waiting_for_transfer_complete(self):
527 - if not self.transfer_status:
528 - return
636 + def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None):
637 + if transfer_statuses is None:
638 + with self.lock:
639 + transfers_to_wait = self.transfer_status[:]
640 + self.transfer_status.clear()
641 + else:
642 + transfers_to_wait = list(transfer_statuses)
643
530 - with self.lock:
531 - transfers_to_wait = self.transfer_status[:]
532 - self.transfer_status.clear()
644 + if not transfers_to_wait:
645 + return
646
647 timeout = self._transfer_timeout
648 deadline = time.monotonic() + timeout
@@ -598,49 +711,105 @@ class MoRIIOWrapper:
711 # [read] mode: receives block release messages from decode side
712 # Decode Role:
713 # [write] mode: receives KV cache write completion notifications
714 + msg_str = repr(msg)
715 handled = False
716 try:
717 data = msgpack.loads(msg)
604 - if isinstance(data, dict) and "req_id" in data:
718 + if isinstance(data, dict):
719 self._handle_structured_message(data)
606 -
720 return
608 - except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException):
721 + except (
722 + msgpack.exceptions.ExtraData,
723 + msgpack.exceptions.UnpackException,
724 + ValueError,
725 + ):
726 logger.debug("Failed to decode msgpack message, will try as string")
727 pass
728
729 try:
730 msg_str = msg.decode("UTF-8")
614 - if msg_str.startswith(MoRIIOConstants.TRANSFER_PREFIX):
731 + if msg_str:
732 self._handle_completion_message(msg_str)
733 handled = True
734 except UnicodeDecodeError:
618 - logger.warning("Received non-UTF8 message: %s", msg_str)
735 + logger.warning("Received non-UTF8 message: %r", msg)
736 if not handled:
737 raise MoRIIOError(f"Unhandled message format: {msg_str}")
738
739 def _handle_structured_message(self, data: dict):
740 + message_type = data.get("type")
741 + if message_type is None and "req_id" in data:
742 + message_type = "remote_blocks"
743 +
744 + if message_type == "remote_blocks":
745 + self._handle_remote_blocks_message(data)
746 + elif message_type == "write_done":
747 + self._handle_write_done_message(data)
748 + elif message_type == "release":
749 + self._handle_release_message(data)
750 + else:
751 + raise MoRIIOError(f"Unhandled structured message type: {message_type}")
752 +
753 + def _handle_remote_blocks_message(self, data: dict):
754 assert get_role() == ROLE.PRODUCER, "Only prefill can get block messages"
755 transfer_id = data["transfer_id"]
756 block_notify_list = data.get("block_notify_list", [])
757 decode_dp_rank = data.get("decode_rank", 0)
627 - assert len(block_notify_list) > 0, (
628 - "block_notify_list cannot be empty in remote allocate message"
629 - )
758 + if not block_notify_list:
759 + raise MoRIIOError(
760 + "block_notify_list cannot be empty in remote allocate message"
761 + )
762
763 with self.lock:
764 + if self._is_transfer_terminal_locked(transfer_id):
765 + logger.debug(
766 + "Ignoring remote allocation for terminal transfer %s",
767 + transfer_id,
768 + )
769 + return
770 self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(
771 block_ids=block_notify_list, decode_dp_rank=decode_dp_rank
772 )
773
774 + def _handle_write_done_message(self, data: dict):
775 + assert get_role() != ROLE.PRODUCER, (
776 + "Only decode can get WRITE completion messages"
777 + )
778 + transfer_id = data["transfer_id"]
779 + with self.lock:
780 + self.done_write_cache_req_ids.append(transfer_id)
781 +
782 + def _handle_release_message(self, data: dict):
783 + assert get_role() == ROLE.PRODUCER, (
784 + "Only prefill can get transfer release messages"
785 + )
786 + transfer_id = data["transfer_id"]
787 + with self.lock:
788 + self.done_req_ids.append(transfer_id)
789 + self.done_remote_allocate_req_dict.pop(transfer_id, None)
790 + self._mark_transfer_terminal_locked(transfer_id)
791 +
792 def _handle_completion_message(self, msg: str):
793 with self.lock:
794 if get_role() == ROLE.PRODUCER:
795 self.done_req_ids.append(msg)
796 + self.done_remote_allocate_req_dict.pop(msg, None)
797 + self._mark_transfer_terminal_locked(msg)
798 else:
799 self.done_write_cache_req_ids.append(msg)
800
643 - def send_notify(self, req_ids, remote_ip, remote_port):
801 + def _is_transfer_terminal_locked(self, transfer_id: TransferId) -> bool:
802 + return transfer_id in self._terminal_transfer_ids
803 +
804 + def _mark_transfer_terminal_locked(self, transfer_id: TransferId) -> None:
805 + self._terminal_transfer_ids[transfer_id] = None
806 + self._terminal_transfer_ids.move_to_end(transfer_id)
807 + while len(self._terminal_transfer_ids) > _MAX_TERMINAL_TRANSFER_IDS:
808 + self._terminal_transfer_ids.popitem(last=False)
809 +
810 + def send_notify(
811 + self, req_ids, remote_ip, remote_port, message_type: str | None = None
812 + ):
813 if not remote_ip or not remote_port:
814 logger.warning("Missing remote_ip or remote_port for notification")
815 return
@@ -664,7 +833,12 @@ class MoRIIOWrapper:
833 "Invalid req_id type: %s, expected str", type(req_id)
834 )
835 continue
667 - sock.send(req_id.encode("utf-8"))
836 + if message_type is None:
837 + sock.send(req_id.encode("utf-8"))
838 + else:
839 + sock.send(
840 + msgpack.dumps({"type": message_type, "transfer_id": req_id})
841 + )
842 except Exception as e:
843 logger.error("Failed to send notification to %s: %s", path, e)
844 self.paths.pop(path, None)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
+109 -5
@@ -56,6 +56,42 @@ def is_mla_cache_layer(
56 return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
57
58
59 +def _spec_dim_matches(value: int, expected: int | None) -> bool:
60 + return expected is None or value == expected
61 +
62 +
63 +def _kernel_layout_matches(
64 + spec: KVCacheSpec, kernel_block_size: int, num_kv_heads: int, head_dim: int
65 +) -> bool:
66 + if kernel_block_size <= 0 or spec.block_size % kernel_block_size != 0:
67 + return False
68 + return _spec_dim_matches(
69 + num_kv_heads, getattr(spec, "num_kv_heads", None)
70 + ) and _spec_dim_matches(head_dim, getattr(spec, "head_size", None))
71 +
72 +
73 +def _select_kernel_block_layout(
74 + layer_name: str, shape: torch.Size, spec: KVCacheSpec
75 +) -> tuple[int, int, int]:
76 + axis2_matches = _kernel_layout_matches(spec, shape[2], shape[3], shape[4])
77 + axis3_matches = _kernel_layout_matches(spec, shape[3], shape[2], shape[4])
78 +
79 + if axis2_matches and axis3_matches and shape[2] != shape[3]:
80 + raise ValueError(
81 + f"Ambiguous MoRIIO kernel-block K/V cache shape for layer "
82 + f"{layer_name}: {tuple(shape)}"
83 + )
84 + if axis2_matches:
85 + return shape[2], shape[3], shape[4]
86 + if axis3_matches:
87 + return shape[3], shape[2], shape[4]
88 +
89 + raise ValueError(
90 + f"Unsupported MoRIIO K/V cache shape for layer {layer_name}: "
91 + f"{tuple(shape)} does not contain block size {spec.block_size}"
92 + )
93 +
94 +
95 def get_layer_transfer_geometry(
96 layer_name: str,
97 kv_cache: torch.Tensor,
@@ -65,6 +101,7 @@ def get_layer_transfer_geometry(
101 shape = kv_cache.shape
102 stride = kv_cache.stride()
103 element_size = kv_cache.element_size()
104 + spec = layer_to_spec[layer_name]
105 is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
106
107 if is_mla_cache and len(shape) == 3:
@@ -85,25 +122,92 @@ def get_layer_transfer_geometry(
122 )
123
124 if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
88 - _, num_blocks, block_size, num_kv_heads, head_dim = shape
125 + _, num_blocks = shape[:2]
126 + kernel_blocks_per_block = 1
127 + if shape[2] == spec.block_size:
128 + block_size, num_kv_heads, head_dim = shape[2:]
129 + elif shape[3] == spec.block_size:
130 + num_kv_heads, block_size, head_dim = shape[2:]
131 + else:
132 + kernel_num_blocks = num_blocks
133 + kernel_block_size, num_kv_heads, head_dim = _select_kernel_block_layout(
134 + layer_name, shape, spec
135 + )
136 + kernel_blocks_per_block = spec.block_size // kernel_block_size
137 + if kernel_num_blocks % kernel_blocks_per_block != 0:
138 + raise ValueError(
139 + f"Unsupported MoRIIO K/V cache shape for layer {layer_name}: "
140 + f"{tuple(shape)} has {kernel_num_blocks} kernel blocks, "
141 + f"not divisible by {kernel_blocks_per_block}"
142 + )
143 + num_blocks = kernel_num_blocks // kernel_blocks_per_block
144 + block_size = spec.block_size
145 slot_size_bytes = num_kv_heads * head_dim * element_size
146 block_len = block_size * slot_size_bytes
91 - remote_kv_stride = stride[1] * (remote_num_blocks or num_blocks)
147 return LayerTransferGeometry(
148 num_blocks=num_blocks,
149 block_size=block_size,
150 block_len=block_len,
151 slot_size_bytes=slot_size_bytes,
97 - block_stride=stride[1],
152 + block_stride=stride[1] * kernel_blocks_per_block,
153 local_kv_stride=stride[0],
99 - remote_kv_stride=remote_kv_stride,
154 + remote_kv_stride=(
155 + stride[1] * kernel_blocks_per_block * (remote_num_blocks or num_blocks)
156 + ),
157 transfers_per_block=2,
158 regions_per_block=1,
159 split_kv_regions=True,
160 )
161
162 if not is_mla_cache and len(shape) == 5 and shape[1] == 2:
106 - num_blocks, _, block_size, num_kv_heads, head_dim = shape
163 + num_blocks = shape[0]
164 + if shape[2] == spec.block_size:
165 + block_size, num_kv_heads, head_dim = shape[2:]
166 + slot_size_bytes = num_kv_heads * head_dim * element_size
167 + block_len = block_size * slot_size_bytes
168 + return LayerTransferGeometry(
169 + num_blocks=num_blocks,
170 + block_size=block_size,
171 + block_len=block_len,
172 + slot_size_bytes=slot_size_bytes,
173 + block_stride=stride[0],
174 + local_kv_stride=stride[1],
175 + remote_kv_stride=stride[1],
176 + transfers_per_block=2,
177 + regions_per_block=2,
178 + split_kv_regions=False,
179 + )
180 + elif shape[3] == spec.block_size:
181 + num_kv_heads, block_size, head_dim = shape[2:]
182 + else:
183 + kernel_num_blocks = num_blocks
184 + kernel_block_size, _, _ = _select_kernel_block_layout(
185 + layer_name, shape, spec
186 + )
187 + kernel_blocks_per_block = spec.block_size // kernel_block_size
188 + if kernel_num_blocks % kernel_blocks_per_block != 0:
189 + raise ValueError(
190 + f"Unsupported MoRIIO K/V cache shape for layer {layer_name}: "
191 + f"{tuple(shape)} has {kernel_num_blocks} kernel blocks, "
192 + f"not divisible by {kernel_blocks_per_block}"
193 + )
194 + num_blocks = kernel_num_blocks // kernel_blocks_per_block
195 + block_size = spec.block_size
196 + block_stride = stride[0] * kernel_blocks_per_block
197 + block_len = block_stride * element_size
198 + slot_size_bytes = block_len // block_size
199 + return LayerTransferGeometry(
200 + num_blocks=num_blocks,
201 + block_size=block_size,
202 + block_len=block_len,
203 + slot_size_bytes=slot_size_bytes,
204 + block_stride=block_stride,
205 + local_kv_stride=None,
206 + remote_kv_stride=None,
207 + transfers_per_block=1,
208 + regions_per_block=1,
209 + split_kv_regions=False,
210 + )
211 slot_size_bytes = num_kv_heads * head_dim * element_size
212 block_len = block_size * slot_size_bytes
213 return LayerTransferGeometry(