[Bugfix][KV-transfer] MoRIIO: per-layer READ-completion barrier in wait_for_layer_load (#48534)
Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io> Signed-off-by: Redwan Khan <redwan.khan@mangoboost.io> Signed-off-by: harishk-mangoboost <harish.kambhampaty@mangoboost.io> Signed-off-by: limeward <32970461+edwinlim0919@users.noreply.github.com> Signed-off-by: QinPR <1905873179@qq.com> Co-authored-by: Jaeyoun Kim <jaeyoun.kim@mangoboost.io> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Redwan Khan <redwan.khan@mangoboost.io> Co-authored-by: harishk-mangoboost <harish.kambhampaty@mangoboost.io> Co-authored-by: rkhan055 <redwankhan055@gmail.com> Co-authored-by: Peiran Qin <66068739+QinPR@users.noreply.github.com> Co-authored-by: QinPR <1905873179@qq.com>
limeward committed
Aug 7, 2026 at 07:23 UTC
47228db84ce59321e6e464f20fabe0ef86e6aef5
2 files changed
+86
-13
tests/v1/kv_connector/unit/test_moriio_tp_ack.py
+14
-1
@@ -10,6 +10,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
10
MoRIIOTransferAck,
11
)
12
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
13
+ MoRIIOConnector,
14
MoRIIOConnectorWorker,
15
get_moriio_expected_ack_count,
16
get_moriio_remote_tp_rank,
@@ -292,7 +293,7 @@ def test_read_completion_sends_structured_release_with_consumer_tp_size():
293
worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
294
worker.world_size = 8
295
worker.moriio_wrapper = FakeWrapper()
295
- worker._recving_transfers = {"req": [DoneStatus()]}
296
+ worker._recving_transfers = {"req": {"layer0": DoneStatus()}}
297
worker._recving_transfers_callback_addr = {
298
"req": ("127.0.0.1", "7000", "tx-release")
299
}
@@ -311,3 +312,15 @@ def test_read_completion_sends_structured_release_with_consumer_tp_size():
312
]
313
assert worker._recving_transfers == {}
314
assert worker._recving_transfers_callback_addr == {}
315
+
316
+
317
+def test_requested_cudagraph_mode_is_never_overridden():
318
+ # The configured cudagraph mode is always honored: the barrier fires when
319
+ # the operator sets cudagraph_mode=PIECEWISE, and READ mode with full
320
+ # graphs only warns instead of silently forcing PIECEWISE.
321
+ assert (
322
+ MoRIIOConnector.requires_piecewise_for_cudagraph({"read_mode": True}) is False
323
+ )
324
+ assert (
325
+ MoRIIOConnector.requires_piecewise_for_cudagraph({"read_mode": False}) is False
326
+ )
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+72
-12
@@ -17,7 +17,7 @@ import numpy as np
17
import torch
18
import zmq
19
20
-from vllm.config import VllmConfig
20
+from vllm.config import CUDAGraphMode, VllmConfig
21
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
22
KVConnectorBase_V1,
23
KVConnectorMetadata,
@@ -65,7 +65,7 @@ from vllm.distributed.parallel_state import (
65
get_tp_group,
66
get_world_group,
67
)
68
-from vllm.forward_context import ForwardContext
68
+from vllm.forward_context import ForwardContext, get_forward_context
69
from vllm.logger import init_logger
70
from vllm.utils.network_utils import (
71
make_zmq_path,
@@ -210,6 +210,19 @@ class MoRIIOConnector(KVConnectorBase_V1):
210
+ str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
211
)
212
self.mode = get_moriio_mode(self.kv_transfer_config)
213
+ if (
214
+ self.mode == MoRIIOMode.READ
215
+ and self.kv_transfer_config.is_kv_consumer
216
+ and vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
217
+ ):
218
+ # warn only; kv-read barrier requires PIECEWISE cudagraph mode
219
+ logger.warning_once(
220
+ "MoRIIO READ mode is running with %s CUDA graphs: per-layer "
221
+ "KV-read barrier can't fire inside full graph; accuracy may "
222
+ "degrade at high concurrency. Set cudagraph_mode=PIECEWISE "
223
+ "in --compilation-config.",
224
+ vllm_config.compilation_config.cudagraph_mode.name,
225
+ )
226
if role == KVConnectorRole.SCHEDULER:
227
self.connector_scheduler: MoRIIOConnectorScheduler | None = (
228
MoRIIOConnectorScheduler(vllm_config, self.engine_id)
@@ -299,7 +312,8 @@ class MoRIIOConnector(KVConnectorBase_V1):
312
self.connector_worker.start_load_kv(self._connector_metadata)
313
314
def wait_for_layer_load(self, layer_name: str) -> None:
302
- pass
315
+ assert self.connector_worker is not None
316
+ self.connector_worker.wait_for_layer_load(layer_name)
317
318
def save_kv_layer(
319
self,
@@ -1225,8 +1239,8 @@ class MoRIIOConnectorWorker:
1239
# Map of engine_id -> num_blocks. All ranks in the same deployment will
1240
# have the same number of blocks.
1241
self.dst_num_blocks: dict[EngineId, int] = {}
1228
- # In progress transfers.
1229
- self._recving_transfers: defaultdict[ReqId, list] = defaultdict(list)
1242
+ # In-progress READ transfers: req_id -> {layer_name: status}.
1243
+ self._recving_transfers: defaultdict[ReqId, dict] = defaultdict(dict)
1244
# Values are (remote_host, remote_notify_port, transfer_id).
1245
self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {}
1246
# Monotonic-clock start times for each in-flight recv transfer.
@@ -1897,14 +1911,60 @@ class MoRIIOConnectorWorker:
1911
1912
return done_sending, done_recving
1913
1914
+ def wait_for_layer_load(self, layer_name: str) -> None:
1915
+ """Block until all in-flight READs of this layer have landed.
1916
+
1917
+ A host-side blocking wait must not run during full-graph capture.
1918
+ """
1919
+ if self.is_producer or self.mode != MoRIIOMode.READ:
1920
+ return
1921
+
1922
+ if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL:
1923
+ return
1924
+
1925
+ deadline = time.monotonic() + self.moriio_config.transfer_timeout
1926
+ while True:
1927
+ with self.moriio_wrapper.lock:
1928
+ pending = [
1929
+ status_by_layer[layer_name]
1930
+ for status_by_layer in self._recving_transfers.values()
1931
+ if layer_name in status_by_layer
1932
+ ]
1933
+
1934
+ if not pending:
1935
+ return
1936
+
1937
+ still_running = False
1938
+ for status in pending:
1939
+ # A failed read is dropped in _pop_done_transfers.
1940
+ if status.Succeeded() or status.Failed():
1941
+ continue
1942
+ still_running = True
1943
+
1944
+ if not still_running:
1945
+ return
1946
+
1947
+ if time.monotonic() > deadline:
1948
+ logger.warning(
1949
+ "MoRIIO READ barrier timed out for layer %s; proceeding "
1950
+ "(request dropped via get_finished).",
1951
+ layer_name,
1952
+ )
1953
+ return
1954
+
1955
+ time.sleep(0.001)
1956
+
1957
def _pop_done_transfers(self) -> set[str]:
1958
done_req_ids: set[str] = set()
1959
_xfer_timeout = int(os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S", "120"))
1960
with self.moriio_wrapper.lock:
1961
to_remove = []
1905
- for req_id, status_list in self._recving_transfers.items():
1906
- last = status_list[-1]
1907
- if last.Succeeded():
1962
+ for req_id, status_by_layer in self._recving_transfers.items():
1963
+ statuses = list(status_by_layer.values())
1964
+ failed_status = next(
1965
+ (status for status in statuses if status.Failed()), None
1966
+ )
1967
+ if statuses and all(status.Succeeded() for status in statuses):
1968
host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
1969
done_req_ids.add(xfer_id)
1970
self.moriio_wrapper.send_notify(
@@ -1915,14 +1975,14 @@ class MoRIIOConnectorWorker:
1975
message_fields={"consumer_tp_size": self.world_size},
1976
)
1977
to_remove.append(req_id)
1918
- elif last.Failed():
1978
+ elif failed_status is not None:
1979
logger.error(
1980
"RDMA transfer failed for request %s: %s (code=%s). "
1981
"Notifying prefill to free blocks; request will be "
1982
"aborted by timeout.",
1983
req_id,
1924
- last.Message(),
1925
- last.Code(),
1984
+ failed_status.Message(),
1985
+ failed_status.Code(),
1986
)
1987
host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
1988
try:
@@ -2564,7 +2624,7 @@ class MoRIIOConnectorWorker:
2624
time.sleep(_backoff)
2625
_backoff = min(_backoff * 2, 0.05)
2626
with self.moriio_wrapper.lock:
2567
- self._recving_transfers[request_id].append(transfer_status)
2627
+ self._recving_transfers[request_id][layer_name] = transfer_status
2628
self._recving_transfers_start.setdefault(request_id, time.monotonic())
2629
self._recving_transfers_callback_addr[request_id] = (
2630
remote_host,