feat(disagg/mooncake): add non-blocking async transfer wrapper bindings

Add batch_transfer_async_submit() (non-blocking submit -> batch_id) and
wait_batch_transfers() (block-until-all, graceful) to MooncakeTransferEngine,
exposing the mooncake async API for the upcoming per-layer overlapped transfer.

Deliberately avoids batch_transfer_*_on_cuda: its CUDA host callback busy-waits
on the stream and calls _exit(1) on transfer failure (verified in the mooncake
source transfer_engine_py.cpp), which is unsafe for production (a decode-side
abort mid-transfer would crash the prefill). The async submit + wait pair lets
per-layer submits pipeline in the RDMA engine, then waits once on a background
thread with graceful failure.

Both bindings raise a clear upgrade error if the engine predates the API and
degrade to -1 on transient errors. Mock-engine unit tested (8 cases).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py

@@ -251,6 +251,54 @@ class MooncakeTransferEngine:
             )
         return ret
 
+    def batch_transfer_async_submit(
+        self,
+        session_id: str,
+        buffers: List[int],
+        peer_buffer_addresses: List[int],
+        lengths: List[int],
+    ) -> int:
+        """Submit a batch WRITE transfer WITHOUT blocking; returns a mooncake batch_id
+        handle (>= 0) or -1 on failure.
+
+        Unlike batch_transfer_sync (which submits and waits in one call), this returns
+        immediately after submitting, so multiple per-layer submits pipeline in the RDMA
+        engine before a single wait. Pair with wait_batch_transfers([batch_id, ...]).
+
+        Chosen over batch_transfer_*_on_cuda, whose host callback busy-waits on the CUDA
+        stream and calls _exit(1) on failure (verified in the mooncake source) — unsafe.
+        """
+        try:
+            return self.engine.batch_transfer_async_write(
+                session_id, buffers, peer_buffer_addresses, lengths
+            )
+        except Exception:
+            if not hasattr(self.engine, "batch_transfer_async_write"):
+                raise RuntimeError(
+                    "Mooncake's async transfer requires mooncake-transfer-engine "
+                    ">= 0.3.9. Please upgrade Mooncake."
+                )
+            logger.debug("Mooncake batch_transfer_async_write failed.")
+            return -1
+
+    def wait_batch_transfers(self, batch_ids: List[int]) -> int:
+        """Block until ALL submitted batch_ids complete; returns 0 on success, non-zero
+        on failure/timeout. Frees the batch_ids. Graceful (no process _exit). Call this
+        on a background thread, never the scheduler/compute thread.
+        """
+        if not batch_ids:
+            return 0
+        try:
+            return self.engine.get_batch_transfer_status(list(batch_ids))
+        except Exception:
+            if not hasattr(self.engine, "get_batch_transfer_status"):
+                raise RuntimeError(
+                    "Mooncake's async transfer status requires mooncake-transfer-engine "
+                    ">= 0.3.9. Please upgrade Mooncake."
+                )
+            logger.debug("Mooncake get_batch_transfer_status failed.")
+            return -1
+
     def get_session_id(self):
         return self.session_id
 

test/registered/unit/disaggregation/test_mooncake_async_wrapper.py

@@ -0,0 +1,94 @@
+"""Unit tests for the async/non-blocking mooncake wrapper bindings (G1).
+
+batch_transfer_async_submit() must do a non-blocking submit returning a batch_id;
+wait_batch_transfers() must block-until-all and return 0/non-zero. Both must raise a
+clear upgrade error if the engine lacks the method, and degrade to -1 on transient
+engine errors. No GPU / real RDMA needed — the engine is mocked.
+"""
+
+import unittest
+
+from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+    MooncakeTransferEngine,
+)
+
+
+class _GoodEngine:
+    def __init__(self, status_ret=0):
+        self.calls = []
+        self._status_ret = status_ret
+
+    def batch_transfer_async_write(self, sid, buffers, peers, lengths):
+        self.calls.append(("submit", sid, list(buffers), list(peers), list(lengths)))
+        return 7  # batch_id handle
+
+    def get_batch_transfer_status(self, ids):
+        self.calls.append(("wait", list(ids)))
+        return self._status_ret
+
+
+class _NoAsyncEngine:
+    pass  # lacks both methods
+
+
+class _RaisingEngine:
+    def batch_transfer_async_write(self, *a):
+        raise RuntimeError("transient rdma error")
+
+    def get_batch_transfer_status(self, *a):
+        raise RuntimeError("transient rdma error")
+
+
+def _wrapper(engine):
+    w = MooncakeTransferEngine.__new__(MooncakeTransferEngine)  # skip real init/connect
+    w.engine = engine
+    return w
+
+
+class TestAsyncSubmit(unittest.TestCase):
+    def test_submit_returns_batch_id_and_dispatches(self):
+        eng = _GoodEngine()
+        w = _wrapper(eng)
+        bid = w.batch_transfer_async_submit("sess", [1, 2], [10, 20], [64, 64])
+        self.assertEqual(bid, 7)
+        self.assertEqual(eng.calls[0], ("submit", "sess", [1, 2], [10, 20], [64, 64]))
+
+    def test_submit_missing_method_raises_upgrade(self):
+        w = _wrapper(_NoAsyncEngine())
+        with self.assertRaises(RuntimeError):
+            w.batch_transfer_async_submit("s", [1], [2], [3])
+
+    def test_submit_transient_error_returns_minus_one(self):
+        w = _wrapper(_RaisingEngine())
+        self.assertEqual(w.batch_transfer_async_submit("s", [1], [2], [3]), -1)
+
+
+class TestWait(unittest.TestCase):
+    def test_wait_dispatches_and_returns_status(self):
+        eng = _GoodEngine(status_ret=0)
+        w = _wrapper(eng)
+        self.assertEqual(w.wait_batch_transfers([7, 8]), 0)
+        self.assertEqual(eng.calls[0], ("wait", [7, 8]))
+
+    def test_wait_empty_is_zero_without_engine_call(self):
+        eng = _GoodEngine()
+        w = _wrapper(eng)
+        self.assertEqual(w.wait_batch_transfers([]), 0)
+        self.assertEqual(eng.calls, [])
+
+    def test_wait_nonzero_status_propagates(self):
+        w = _wrapper(_GoodEngine(status_ret=-1))
+        self.assertEqual(w.wait_batch_transfers([7]), -1)
+
+    def test_wait_missing_method_raises_upgrade(self):
+        w = _wrapper(_NoAsyncEngine())
+        with self.assertRaises(RuntimeError):
+            w.wait_batch_transfers([7])
+
+    def test_wait_transient_error_returns_minus_one(self):
+        w = _wrapper(_RaisingEngine())
+        self.assertEqual(w.wait_batch_transfers([7]), -1)
+
+
+if __name__ == "__main__":
+    unittest.main()