From 9a91c73e092cb8b64d268ad4b89467f74b04b288 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 6 Jun 2026 23:20:19 +0000 Subject: [PATCH] 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) --- .../mooncake_transfer_engine.py | 48 ++++++++++ .../test_mooncake_async_wrapper.py | 94 +++++++++++++++++++ 2 files changed, 142 insertions(+) create mode 100644 test/registered/unit/disaggregation/test_mooncake_async_wrapper.py diff --git a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py index ba20176a6..5cf60e1bf 100644 --- a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py +++ b/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 diff --git a/test/registered/unit/disaggregation/test_mooncake_async_wrapper.py b/test/registered/unit/disaggregation/test_mooncake_async_wrapper.py new file mode 100644 index 000000000..975b7f4cf --- /dev/null +++ b/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()