[Bugfix][ROCm][MoRIIO] Wait for MoRIIO WRITE mode completions in short chunks

tests/v1/kv_connector/unit/test_moriio_tp_ack.py

@@ -655,10 +655,14 @@ def test_step_barrier_timeout_fails_closed_while_reads_in_flight(has_mamba, fake
     moriio_engine.StatusCode is None, reason="mori batched wait unavailable"
 )
 @pytest.mark.parametrize("has_mamba", [False, True])
-def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(has_mamba):
-    engine = SimpleNamespace(
-        wait_all=Mock(return_value=moriio_engine.StatusCode.IN_PROGRESS)
-    )
+def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(
+    has_mamba, fake_clock
+):
+    def wait_all(statuses, timeout_ms):
+        fake_clock.sleep(timeout_ms / 1000)
+        return moriio_engine.StatusCode.IN_PROGRESS
+
+    engine = SimpleNamespace(wait_all=Mock(side_effect=wait_all))
     worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
     worker._has_mamba = has_mamba
     worker.moriio_wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05)
@@ -667,7 +671,8 @@ def test_step_barrier_native_timeout_fails_closed_while_reads_in_flight(has_mamb
     with pytest.raises(TransferError, match="timed out"):
         worker._await_reads_issued_this_step()
 
-    engine.wait_all.assert_called_once()
+    assert fake_clock.now >= 0.05
+    assert all(c.kwargs["timeout_ms"] <= 10 for c in engine.wait_all.call_args_list)
 
 
 def test_step_barrier_noop_when_no_reads_posted():

tests/v1/kv_connector/unit/test_moriio_transfer_completion.py

@@ -131,7 +131,7 @@ def test_wait_blocks_until_terminal():
         timer.cancel()
 
 
-@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 50), (0.0001, 1)])
+@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 10), (0.004, 4), (0.0001, 1)])
 def test_native_wait_receives_all_statuses_and_a_positive_deadline(
     timeout, expected_ms
 ):
@@ -186,5 +186,44 @@ def test_native_wait_drains_only_the_unscoped_status_snapshot():
     wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05)
     wrapper.transfer_status.extend(statuses)
     wrapper.waiting_for_transfer_complete()
-    engine.wait_all.assert_called_once_with(statuses, timeout_ms=50)
+    engine.wait_all.assert_called_once_with(statuses, timeout_ms=10)
     assert wrapper.transfer_status == []
+
+
+def _wait_all_that_completes(status, on_call, result="IN_PROGRESS"):
+    def wait_all(statuses, timeout_ms):
+        if wait_all_mock.call_count >= on_call:
+            status.state = "success"
+        return getattr(moriio_engine.StatusCode, result)
+
+    wait_all_mock = Mock(side_effect=wait_all)
+    return wait_all_mock
+
+
+@pytest.mark.parametrize("on_call", [1, 3])
+def test_native_wait_returns_once_the_statuses_succeed(on_call):
+    """MoRI's xgmi wait_all can return IN_PROGRESS at its deadline although
+    every status has succeeded."""
+    late = _pending()
+    wait_all = _wait_all_that_completes(late, on_call)
+    wrapper = MoRIIOWrapper(
+        moriio_engine=SimpleNamespace(wait_all=wait_all), transfer_timeout=30.0
+    )
+    wrapper.waiting_for_transfer_complete([_ok(), late])
+    assert wait_all.call_count == on_call
+    assert all(c.kwargs["timeout_ms"] == 10 for c in wait_all.call_args_list)
+
+
+@pytest.mark.parametrize("fail_fast", [True, False])
+def test_native_wait_after_a_failure_waits_for_in_flight_statuses_unless_fail_fast(
+    fail_fast,
+):
+    late = _pending()
+    wait_all = _wait_all_that_completes(late, on_call=3, result="ERR_RDMA_OP")
+    wrapper = MoRIIOWrapper(
+        moriio_engine=SimpleNamespace(wait_all=wait_all), transfer_timeout=30.0
+    )
+    with pytest.raises(TransferError, match="SQ full"):
+        wrapper.waiting_for_transfer_complete([_bad(), late], fail_fast=fail_fast)
+    assert wait_all.call_count == (1 if fail_fast else 3)
+    assert late.Succeeded() is not fail_fast

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

@@ -74,6 +74,7 @@
 
 
 _MAX_TERMINAL_TRANSFER_IDS = 4096
+_NATIVE_WAIT_SLICE_S = 0.01
 
 
 WriteGeometryKey = tuple[tuple[int, ...], tuple[int, ...], torch.dtype]
@@ -911,25 +912,42 @@ def waiting_for_transfer_complete(
             self._poll_transfers_until_done(transfers_to_wait, fail_fast)
             return
 
-        # wait_all waits on every status until the deadline; a failure does not
-        # end it early, so fail_fast only applies to the polling fallback.
-        # Zero means a nonblocking progress call in Mori, not an immediate
-        # deadline. Round up so even sub-millisecond waits remain blocking.
-        timeout_ms = ceil(self._transfer_timeout * 1000)
-        result = wait_all(transfers_to_wait, timeout_ms=timeout_ms)
-        if result != StatusCode.SUCCESS:
-            details = [
-                f"{status.Message()} (code={status.Code()})"
-                for status in transfers_to_wait
-                if status.Failed()
-            ]
-            if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS):
-                details.append(
-                    f"RDMA transfer timed out after {self._transfer_timeout:g}s"
-                )
-            raise TransferError(
-                f"MoRIIO batch wait failed ({result}): " + "; ".join(details)
-            )
+        # On MoRI's xgmi backend a bounded wait_all polls each status once and
+        # then sleeps until its deadline.
+        remaining = self._transfer_timeout
+        deadline = time.monotonic() + remaining
+        while True:
+            # Zero means a nonblocking progress call in Mori, not an immediate
+            # deadline. Round up so even sub-millisecond waits remain blocking.
+            slice_ms = ceil(min(remaining, _NATIVE_WAIT_SLICE_S) * 1000)
+            result = wait_all(transfers_to_wait, timeout_ms=slice_ms)
+            if result == StatusCode.SUCCESS:
+                return
+            state = self.poll_transfer_batch(transfers_to_wait)
+            if state is TransferBatchState.DONE:
+                return
+            if state is TransferBatchState.FAILED:
+                if fail_fast or all(
+                    status.Succeeded() or status.Failed()
+                    for status in transfers_to_wait
+                ):
+                    break
+            elif result not in (StatusCode.INIT, StatusCode.IN_PROGRESS):
+                break
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
+
+        details = [
+            f"{status.Message()} (code={status.Code()})"
+            for status in transfers_to_wait
+            if status.Failed()
+        ]
+        if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS):
+            details.append(f"RDMA transfer timed out after {self._transfer_timeout:g}s")
+        raise TransferError(
+            f"MoRIIO batch wait failed ({result}): " + "; ".join(details)
+        )
 
     def _poll_transfers_until_done(
         self, transfers_to_wait: list[Any], fail_fast: bool = True