@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 971dac2caa

[Bugfix][KV-transfer] MoRIIO: retry RDMA send-queue-full backpressure instead of failing the read (#47495)

Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io> Signed-off-by: Edwin Lim <edwinlim0919@gmail.com> Signed-off-by: harishk-mangoboost <harish.kambhampaty@mangoboost.io> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: harishk-mangoboost <harish.kambhampaty@mangoboost.io>

limeward committed Jul 16, 2026 at 13:02 UTC 971dac2caa827cb34dc682791ac6e15f6dbdab84

4 files changed +91 -14

tests/v1/kv_connector/unit/test_moriio_kv_layout.py

+11 -2

@@ -670,16 +670,25 @@ def test_moriio_wrapper_rejects_invalid_messages(role, payload, match):

670 wrapper._handle_message(payload)

671

672

673 -def test_block_id_length_mismatch_raises_value_error():

673 +def test_local_block_ids_longer_than_remote_raises_value_error():

674 cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)

675 worker = _worker({"layer": cache}, {"layer": _full_spec()})

676

677 - with pytest.raises(ValueError, match="must have the same length"):

677 + with pytest.raises(ValueError, match="longer than remote_block_ids"):

678 moriio_layout.compute_block_transfer_offsets(

679 "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks

680 )

681

682

683 +def test_empty_local_block_ids_is_free_only_noop():

684 + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)

685 + worker = _worker({"layer": cache}, {"layer": _full_spec()})

686 +

687 + assert moriio_layout.compute_block_transfer_offsets(

688 + "layer", cache, worker.layer_to_spec, [], [4, 5], _remote_meta().num_blocks

689 + ) == ([], [], [])

690 +

691 +

692 def test_registration_regions_do_not_split_interleaved_or_mla_cache():

693 separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)

694 interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)

tests/v1/kv_connector/unit/test_moriio_tp_ack.py

@@ -253,6 +253,7 @@ def test_worker_get_finished_counts_structured_release_fan_in():

253 worker.transfer_id_to_request_id = {"tx-fanin": "req-fanin"}

254 worker._consumer_notification_counts = {}

255 worker._completed_consumer_notifications = set()

256 + worker._pending_unmapped_acks = []

257

258 assert worker.get_finished() == (set(), set())

259 assert worker._consumer_notification_counts == {"tx-fanin": 1}

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

+70 -9

@@ -1030,6 +1030,16 @@ class MoRIIOConnectorWorker:

1030 use_mla=self.use_mla,

1031 )

1032 self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}

1033 + # READ-mode producer: a decode release-ACK can arrive BEFORE

1034 + # start_load_kv populates transfer_id_to_request_id (the notify races

1035 + # ahead of the scheduler->worker sync). Buffer such ACKs and retry them

1036 + # next get_finished tick instead of dropping them -- dropping loses the

1037 + # completion, so the request is never marked done_sending, its KV blocks

1038 + # leak, and the prefill KV cache wedges at high concurrency. Buffered

1039 + # BEFORE resolve_moriio_transfer_ack, so each ACK is counted exactly once

1040 + # (on the tick its mapping exists) -- the heterogeneous-TP ack-counting

1041 + # is preserved.

1042 + self._pending_unmapped_acks: list = []

1043

1044 # TODO: consider the integration of flashinfer or other backends.

1045 self.backend_name = backend.get_name()

@@ -1542,16 +1552,22 @@ class MoRIIOConnectorWorker:

1552 # pop_finished_req_ids returns release ACKs sent by decode. Keep

1553 # duplicate ACKs because heterogeneous TP can fan multiple decode

1554 # ranks into one prefill rank for the same transfer_id.

1545 - finished_acks = self.moriio_wrapper.pop_finished_req_ids()

1555 + # Combine freshly-arrived ACKs with any buffered from prior ticks

1556 + # whose transfer_id wasn't mapped yet (notify raced ahead of

1557 + # start_load_kv); retry the lookup every tick. Buffered before

1558 + # resolve_moriio_transfer_ack so each ACK is counted exactly once.

1559 + finished_acks = self._pending_unmapped_acks + list(

1560 + self.moriio_wrapper.pop_finished_req_ids()

1561 + )

1562 + self._pending_unmapped_acks = []

1563 resolved_transfer_ids: set[TransferId] = set()

1564 for ack in finished_acks:

1565 transfer_id = ack if isinstance(ack, str) else ack.transfer_id

1566 if transfer_id not in self.transfer_id_to_request_id:

1550 - logger.warning(

1551 - "Could not find %s in transfer_id_to_request_id "

1552 - "lookup table. This could lead to a possible hang.",

1553 - transfer_id,

1554 - )

1567 + # Mapping not populated yet -- buffer and retry next tick,

1568 + # do NOT drop (dropping leaks producer KV at high conc and

1569 + # wedges the prefill).

1570 + self._pending_unmapped_acks.append(ack)

1571 continue

1572 resolved_transfer_id = resolve_moriio_transfer_ack(

1573 ack,

@@ -1933,6 +1949,23 @@ class MoRIIOConnectorWorker:

1949 ),

1950 )

1951

1952 + @staticmethod

1953 + def _is_sq_full_status(status) -> bool:

1954 + """True if a MoRIIO transfer status is a transient RDMA send-queue-full

1955 + rejection (retryable backpressure), not a terminal failure.

1956 +

1957 + read_remote_data posts the RDMA READ synchronously (the mori executor

1958 + joins its worker before returning and marks the status on the calling

1959 + thread), so a send-queue-full rejection is a Failed() status the moment

1960 + the call returns. mori surfaces it as a generic ERR_RDMA_OP carrying

1961 + "SQ full" in the message (no distinct code), so we match the message.

1962 + Only meaningful once status.Failed() is True.

1963 + """

1964 + try:

1965 + return bool(status.Failed()) and "SQ full" in (status.Message() or "")

1966 + except Exception:

1967 + return False

1968 +

1969 def _read_blocks(

1970 self,

1971 local_block_ids: list[int],

@@ -1950,6 +1983,8 @@ class MoRIIOConnectorWorker:

1983 dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)

1984 sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)

1985

1986 + # SQ-full backpressure deadline, shared across this request's layers.

1987 + _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout

1988 for layer_name in self.layer_name_to_local_kv_cache_metadata:

1989 sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(

1990 layer_name

@@ -1962,9 +1997,35 @@ class MoRIIOConnectorWorker:

1997 remote_tp_size=remote_tp_size,

1998 )

1999 # TODO : apply multi-session batch-read when moriio support it

1965 - transfer_status = self.moriio_wrapper.read_remote_data(

1966 - offs[2], offs[0], offs[1], sessions[sess_idx]

1967 - )

2000 + #

2001 + # SQ-full backpressure: read_remote_data posts the RDMA READ

2002 + # SYNCHRONOUSLY, so a send-queue-full rejection (per-QP HW cap) comes

2003 + # back as a Failed() status right here. A SEPARATE CQ-poll thread

2004 + # drains completions and frees SQ depth, so back off and RE-POST

2005 + # rather than let a transient rejection abort the request. No

2006 + # self-deadlock (the drain is off-thread); the reserve is

2007 + # all-or-nothing (nothing posted on a rejected attempt). Bounded by

2008 + # transfer_timeout; on sustained overload store the failed status and

2009 + # let get_finished handle it non-fatally (notify prefill + drop).

2010 + _backoff = 0.001

2011 + while True:

2012 + transfer_status = self.moriio_wrapper.read_remote_data(

2013 + offs[2], offs[0], offs[1], sessions[sess_idx]

2014 + )

2015 + if not self._is_sq_full_status(transfer_status):

2016 + break

2017 + if time.monotonic() > _sq_deadline:

2018 + logger.warning(

2019 + "MoRIIO READ send queue stayed full past "

2020 + "transfer_timeout for req %s layer %s; storing failed "

2021 + "status (get_finished notifies prefill and drops the "

2022 + "request). Raise qp_per_transfer if frequent.",

2023 + request_id,

2024 + layer_name,

2025 + )

2026 + break

2027 + time.sleep(_backoff)

2028 + _backoff = min(_backoff * 2, 0.05)

2029 with self.moriio_wrapper.lock:

2030 self._recving_transfers[request_id].append(transfer_status)

2031 self._recving_transfers_callback_addr[request_id] = (

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py

+9 -3

@@ -312,10 +312,16 @@ def compute_block_transfer_offsets(

312 [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]]

313 ] = merge_contiguous_offsets,

314 ) -> tuple[list[int], list[int], list[int]]:

315 - if len(local_block_ids) != len(remote_block_ids):

315 + # A shorter (or empty) local list is the READ-mode "drop the transfer, just

316 + # free the prefill blocks" case (full-prefix-hit / aborted-before-scheduled):

317 + # decode pulls fewer blocks than the prefill holds. The zip loop below pairs

318 + # local[i]<->remote[i] and sizes by len(local), so a short local transfers

319 + # only what decode allocated and an empty local is a no-op. A longer local

320 + # list is a genuine bug and still fails loudly.

321 + if len(local_block_ids) > len(remote_block_ids):

322 raise ValueError(

317 - "local_block_ids and remote_block_ids must have the same length: "

318 - f"{len(local_block_ids)} != {len(remote_block_ids)}"

323 + "local_block_ids longer than remote_block_ids: "

324 + f"{len(local_block_ids)} > {len(remote_block_ids)}"

325 )

326 geometry = get_layer_transfer_geometry(

327 layer_name, kv_cache, layer_to_spec, remote_num_blocks