[kliuae]

Overview

In MoRIIO WRITE mode, the prefill (P) instance waits for a request's KV copies before it sends write_done to the decode (D) instance, and the READ barrier waits for its reads the same way. Both go through MoRIIOWrapper.waiting_for_transfer_complete which makes one MoRI wait_all call with the whole transfer_timeout. On MoRI's xGMI backend, that call polls each status once and then sleeps until its deadline, so a copy that is still running at that poll holds the wait for the full timeout (30s by default). This also happens on MoRI's scatter/gather copy path under load, where a wait that still sees IN_PROGRESS at its deadline fails the transfer and leaves the decode request waiting. This PR calls wait_all in 10ms intervals, and checks the statuses between the intervals, up to the same transfer_timeout.

Claims

• On MI355X with MoRI v1.2.3.post1, a prefill process writing 118 layers (128 blocks of 32 KiB each) into a decode process on another GPU over xGMI with MoRI's default copy path, transfer_timeout 30s: on main all 6 of 6 waits took 29.999 s, even though every one of the 118 statuses had succeeded. With this PR, 100 of 100 waits ended after 0.210 s (median; min 0.183 s, max 0.210 s), all succeeded.
• When a status fails, the wait raises at once with fail_fast=True (WRITE), or once no status is in flight with fail_fast=False (READ barrier), as the polling fallback already does. Timeouts and error messages are unchanged.
• With P/D serving, for DeepSeek-V2-Lite-Chat served with prefill TP1 and decode TP1 on MI355X over xGMI, with MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 (MoRI's scatter/gather copy path) on prefill, 16K-token prompts: on main, two of three runs at concurrency 8 left requests that never completed (92 and 95 out of 96), each after a completion wait that failed with IN_PROGRESS when its 30s transfer_timeout expired. In the run that finished, concurrency 16 reached only 2.28 req/s with a client TTFT p90 of 30.2 s. With this PR, all 96 requests completed in each of three runs, with no failed wait, at 6.96–7.20 req/s.

Validation

Unit tests

• tests/v1/kv_connector/unit/test_moriio_transfer_completion.py: the native wait returns once the statuses succeed, follows fail_fast after a failure, and calls wait_all in 10ms slices.
• tests/v1/kv_connector/unit/test_moriio_tp_ack.py: the READ barrier's native-timeout test runs on a fake clock and checks that the wait lasts transfer_timeout in slices of at most 10ms.

Completion wait over xGMI

On MI355X with MoRI 1.2.3.post1 (no MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD, so MoRI uses its default copy path). A prefill process on GPU 0 writes 118 layers of 128 blocks (32 KiB each) into a decode process's buffer on GPU 1, one batch_write per layer, and waits for them through MoRIIOWrapper.waiting_for_transfer_complete with transfer_timeout=30, as the WRITE producer does before write_done.

	main	PR
number of waits timed	6	100
wait, median (s)	29.999	0.210
wait, min / max (s)	29.999 / 29.999	0.183 / 0.210
statuses succeeded at the end	118 of 118 in every wait	118 of 118 in every wait

reproducer (moriio_wait_repro.py)

"""Time vLLM's MoRIIO WRITE completion wait over xGMI.

A prefill process on --p-gpu writes 118 layers of 128 blocks (32 KiB each) into a
decode process's buffer on --d-gpu, one MoRI batch_write per layer, then waits
for them with MoRIIOWrapper.waiting_for_transfer_complete (transfer_timeout 30 s),
as the WRITE producer does before it sends write_done.

    python moriio_wait_repro.py --p-gpu 0 --d-gpu 1 --requests 3
"""

import argparse
import multiprocessing as mp
import os
import pickle
import tempfile
import time

LAYERS, BLOCKS, BLOCK_BYTES = 118, 128, 32768
LAYER_BYTES = 2 * BLOCKS * BLOCK_BYTES


def setup(role, port, run_dir):
    import torch
    from mori.io import (
        BackendType,
        EngineDesc,
        IOEngine,
        IOEngineConfig,
        MemoryDesc,
        XgmiBackendConfig,
    )

    buf = torch.zeros(LAYERS * LAYER_BYTES, dtype=torch.uint8, device="cuda")
    engine = IOEngine(f"{role}-repro", IOEngineConfig("127.0.0.1", port))
    engine.create_backend(BackendType.XGMI, XgmiBackendConfig())
    mem = engine.register_torch_tensor(buf)
    mine = os.path.join(run_dir, role)
    peer = os.path.join(run_dir, {"p": "d", "d": "p"}[role])
    with open(mine + ".tmp", "wb") as f:
        pickle.dump((engine.get_engine_desc().pack(), mem.pack()), f)
    os.rename(mine + ".tmp", mine)
    while not os.path.exists(peer):
        time.sleep(0.1)
    with open(peer, "rb") as f:
        peer_engine, peer_mem = pickle.load(f)
    engine.register_remote_engine(EngineDesc.unpack(peer_engine))
    return engine, buf, mem, MemoryDesc.unpack(peer_mem)


def decode(port, run_dir):
    _ = setup("d", port + 1, run_dir)
    while not os.path.exists(os.path.join(run_dir, "done")):
        time.sleep(0.5)


def prefill(port, run_dir, requests):
    from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
        MoRIIOWrapper,
    )

    engine, _buf, mem, peer_mem = setup("p", port, run_dir)
    session = engine.create_session(mem, peer_mem)
    wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=30.0)
    sizes = [BLOCK_BYTES] * BLOCKS
    for request in range(requests):
        statuses = []
        for layer in range(LAYERS):
            offsets = [layer * LAYER_BYTES + 2 * b * BLOCK_BYTES for b in range(BLOCKS)]
            uid = engine.allocate_transfer_uid()
            statuses.append(session.batch_write(offsets, offsets, sizes, uid))
        start = time.monotonic()
        try:
            wrapper.waiting_for_transfer_complete(statuses)
            outcome = "ok"
        except Exception as e:
            outcome = f"{type(e).__name__}: {e}"
        succeeded = sum(s.Succeeded() for s in statuses)
        print(
            f"request {request}: wait {time.monotonic() - start:.3f} s, "
            f"{succeeded}/{len(statuses)} statuses succeeded, {outcome}",
            flush=True,
        )
    open(os.path.join(run_dir, "done"), "w").close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--p-gpu", type=int, default=0)
    parser.add_argument("--d-gpu", type=int, default=1)
    parser.add_argument("--port", type=int, default=29600)
    parser.add_argument("--requests", type=int, default=3)
    args = parser.parse_args()
    ctx = mp.get_context("spawn")
    run_dir = tempfile.mkdtemp()
    os.environ["HIP_VISIBLE_DEVICES"] = str(args.d_gpu)
    d = ctx.Process(target=decode, args=(args.port, run_dir))
    d.start()
    os.environ["HIP_VISIBLE_DEVICES"] = str(args.p_gpu)
    p = ctx.Process(target=prefill, args=(args.port, run_dir, args.requests))
    p.start()
    p.join()
    d.join()

python3 moriio_wait_repro.py --p-gpu 0 --d-gpu 1 --requests 6     # main
python3 moriio_wait_repro.py --p-gpu 0 --d-gpu 1 --requests 100   # this PR

P/D serving over xGMI

DeepSeek-V2-Lite-Chat, prefill and decode both served with TP1, MoRIIO WRITE over xGMI, using MoRIIO toy proxy, MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 on the prefill instance, benchmarked with ISL/OSL=16k/16.

	main (runs 1 / 2 / 3)	PR (runs 1 / 2 / 3)
concurrency 8: requests completed (out of 96)	92 / 96 / 95	96 / 96 / 96
completion waits that failed with IN_PROGRESS	1 / 0 / 1	0 / 0 / 0
concurrency 8: client TTFT mean (ms)	– / 1,089.8 / –	1,074.3 / 1,085.7 / 1,075.3
concurrency 8: client TTFT p99 (ms)	– / 1,490.3 / –	1,584.5 / 1,607.7 / 1,588.6
concurrency 16: requests completed (of 192)	– / 192 / –	192 / 192 / 192
concurrency 16: request throughput (req/s)	– / 2.28 / –	7.12 / 7.17 / 7.20
concurrency 16: client TTFT median (ms)	– / 2,191.7 / –	2,174.4 / 2,164.0 / 2,159.1
concurrency 16: client TTFT p90 (ms)	– / 30,151.4 / –	2,317.5 / 2,212.6 / 2,187.5
concurrency 16: client TTFT p99 (ms)	– / 31,936.9 / –	3,571.5 / 3,560.5 / 3,569.3

On main the unfinished requests are the ones whose transfer failed with the prefill write worker raising TransferError ("MoRIIO batch wait failed (StatusCode.IN_PROGRESS): RDMA transfer timed out after 30s") while the decode instance being told nothing, so the decode request kept waiting for its KV. In main's second run at concurrency 16, the waits returned at their 30s deadline instead of failing. As the write worker handles one transfer at a time, every transfer queued behind such wait was delayed too.

serving and benchmark commands

export VLLM_ROCM_USE_AITER=1 VLLM_USE_BREAKABLE_CUDAGRAPH=0
ARGS="--tensor-parallel-size 1 --block-size 64 --kv-cache-dtype auto --attention-backend ROCM_AITER_MLA \
  --gpu-memory-utilization 0.90 --enable-chunked-prefill --max-num-batched-tokens 32768 --no-enable-prefix-caching \
  --max-num-seqs 64 --compilation-config {\"cudagraph_mode\":\"FULL_DECODE_ONLY\"} --max-model-len 40960"

python3 examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py --port 11350 --service-discovery-port 11353

HIP_VISIBLE_DEVICES=0 VLLM_PORT=11358 MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 \
vllm serve deepseek-ai/DeepSeek-V2-Lite-Chat --port 11351 $ARGS \
  --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_connector_extra_config":{"backend":"xgmi","proxy_ip":"127.0.0.1","proxy_ping_port":"11353","http_port":"11351","handshake_port":"11354","notify_port":"11355"}}'

HIP_VISIBLE_DEVICES=1 VLLM_PORT=11362 \
vllm serve deepseek-ai/DeepSeek-V2-Lite-Chat --port 11352 $ARGS \
  --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"backend":"xgmi","proxy_ip":"127.0.0.1","proxy_ping_port":"11353","http_port":"11352","handshake_port":"11356","notify_port":"11357"}}'

for cn in 8:8 8:96 16:192; do   # warm-up, the concurrency-8 run, the concurrency-16 run
  vllm bench serve --backend vllm --model deepseek-ai/DeepSeek-V2-Lite-Chat --base-url http://127.0.0.1:11350 \
    --endpoint /v1/completions --dataset-name random --random-input-len 16384 --random-output-len 16 \
    --max-concurrency ${cn%:*} --num-prompts ${cn#*:} --ignore-eos --seed 0 \
    --percentile-metrics ttft,tpot,e2el --metric-percentiles 50,90,99
done
grep -c 'MoRIIO batch wait failed (StatusCode.IN_PROGRESS)' prefill.log   # failed completion waits

Details

On MoRIIO's WRITE path of the prefill instance under PD disaggregated serving, for one request:

1. D allocates blocks for the request and sends their ids to every P rank.
2. In P's forward, each layer's write is scheduled with the layer's CUDA event.
3. P's write worker waits for each event and posts the layer's copies (batch_write on a MoRI session).
4. After the last layer, _finalize_if_complete waits for every copy of the request (waiting_for_transfer_complete) and sends write_done.
5. D marks the request's KV as received and starts decoding it.

Step 4 makes one MoRI call with the whole timeout:

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py's MoRIIOWrapper.waiting_for_transfer_complete
timeout_ms = ceil(self._transfer_timeout * 1000)
result = wait_all(transfers_to_wait, timeout_ms=timeout_ms)

On MoRI's xGMI backend, a status completes only when it is polled, and a bounded wait_all polls each status once and then sleeps until its deadline. If the copies are still running at that one poll, the wait costs the whole transfer_timeout; when the statuses are still IN_PROGRESS at the deadline, the transfer fails with TransferError, and on main no failure notice reaches decode. With the default copy path the last layers are nearly always still running when step 4 starts, so the wait lasts the whole transfer_timeout and the request's write_done reaches D 30s late. As shown in the the reproducer above, however, every status had succeeded long before the call returned. On the scatter/gather path the copies are fast, but when prefill has already queued its next step, they complete only after that step, so under load the one poll misses them there as well.

This PR addresses this by calling wait_all with at most 10ms and checking the statuses with poll_transfer_batch between calls, until transfer_timeout:

• every status succeeded: return;
• a status failed: raise at once with fail_fast=True, or once nothing is in flight with fail_fast=False (on main the native path ignored fail_fast and always waited for the deadline);
• an error result without a failed status, or the deadline: raise, with the same messages as before.

AI agent was used in this PR.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[whx-sjtu]

No blocking issues. Reproduced the same stall on MiniMax-M3.1 1P1D over XGMI, which this PR covers.

The sliced wait adds at most 10 ms per request there, and our slowest wait (2.1 s) is far from the 30 s bound. Data below.

Evidence

• Setup: MiniMax-M3.1-Flash MXFP4, P TP4 + D TP4 on one 8x MI355X node, MoRIIOConnector WRITE mode, backend: xgmi.
• Unfixed: the prefill write worker sat in wait_all, and decode saw about one transfer finish every 30-60 s, matching the 30 s default deadline.
• Local workaround timeout_ms=-1 for XGMI only (blocks in Mori's wait callback): a 3600 s agentic run at concurrency 48 completed with 1/5967 request errors. It drops the transfer_timeout bound, so the sliced wait here is the better fix.
• Per-request XGMI wait at concurrency 48 (2847 requests): p50 6 ms, p90 0.33 s, p99 1.2 s, max 2.1 s.

[ChuanLi1101]

Looks correct — no required changes from me.

Two optional, non-blocking notes:

1. The 10ms slice is fine for xGMI; if another MoRI backend behaves differently later, we can make it configurable.
2. The root cause is still in MoRI (wait_all polls once then sleeps until the deadline). The sliced wait is a reasonable vLLM workaround; worth a MoRI-side issue so it does not sleep the full timeout after a single poll. Not a blocker for this PR.

[shen-shanshan]

/ci run --allow-stale

[github-actions]

✅ Triggered Buildkite CI #99493 for commit beeed848ffd0.

⚠️ This PR is 88 commits behind upstream main. Running CI at your own risk because --allow-stale was requested; outdated CI configuration may cause failures. Before merging, merge or rebase onto the latest main, then rerun /ci run on the latest PR commit.

[whx-sjtu]

lgtm