feat(disagg/mooncake): async per-layer KV transfer primitive (lever A foundation)

Add _transfer_layers_async: submit each layer's transfer non-blocking via the
async API (batch_transfer_async_submit), pipelining layers in the RDMA engine,
then wait for all once (wait_batch_transfers). Gated by the new
SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER env (default off); replaces the monolithic
all-layers batch_transfer_sync on that path. This is the transfer mechanism for
per-layer overlap (lever A) and removes the per-layer blocking-sync tax measured
in B1a; the forward-overlap hook (G2) builds on it next. Uses the safe async API,
never the OnCuda busy-wait/_exit path.

Unit-tested (test_per_layer_transfer.py, 5 cases): one submit per non-empty layer,
single wait-for-all, empty-layer skip, submit-failure drain + return -1, and
wait-status propagation.

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

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -268,6 +268,9 @@ class MooncakeKVManager(CommonKVManager):
             self.enable_custom_mem_pool, self.custom_mem_pool_type = (
                 check_mooncake_custom_mem_pool_enabled()
             )
+            self.enable_per_layer_async_transfer = (
+                envs.SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER.get()
+            )
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self.start_decode_thread()
 
@@ -321,6 +324,31 @@ class MooncakeKVManager(CommonKVManager):
             mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
         )
 
+    def _transfer_layers_async(
+        self, mooncake_session_id, layers_params, set_transfer_blocks
+    ):
+        # Submit each layer's transfer non-blocking (pipelined in the RDMA engine),
+        # then wait for all to complete once. Removes the per-layer blocking-sync tax
+        # (B1a) and is the transfer mechanism for per-layer overlap (lever A). Uses the
+        # safe async API (G1: batch_transfer_async_submit + wait_batch_transfers),
+        # never the OnCuda busy-wait/_exit path.
+        batch_ids = []
+        for src_ptr, dst_ptr, item_len in layers_params:
+            transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
+            if not transfer_blocks:
+                continue
+            src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
+            batch_id = self.engine.batch_transfer_async_submit(
+                mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
+            )
+            if batch_id < 0:
+                # Submit failed: drain whatever was already submitted, then fail.
+                if batch_ids:
+                    self.engine.wait_batch_transfers(batch_ids)
+                return -1
+            batch_ids.append(batch_id)
+        return self.engine.wait_batch_transfers(batch_ids)
+
     def _send_kvcache_generic(
         self,
         mooncake_session_id: str,
@@ -418,7 +446,11 @@ class MooncakeKVManager(CommonKVManager):
             return self._transfer_data(mooncake_session_id, transfer_blocks)
 
         start_time = time.perf_counter() if transfer_stats_enabled else 0.0
-        if self.enable_custom_mem_pool:
+        if self.enable_per_layer_async_transfer:
+            status = self._transfer_layers_async(
+                mooncake_session_id, layers_params, set_transfer_blocks
+            )
+        elif self.enable_custom_mem_pool:
             futures = [
                 executor.submit(
                     process_layer,

python/sglang/srt/environ.py

@@ -211,6 +211,10 @@ class Envs:
     SGLANG_DEBUG_SORT_NVTX = EnvBool(False)
     SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False)
     SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
+    # Per-layer overlapped KV transfer (lever A): submit each layer's KV transfer
+    # asynchronously (pipelined batch_transfer_async) instead of one monolithic
+    # all-layers batch_transfer_sync. Default off until validated.
+    SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER = EnvBool(False)
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)

test/registered/unit/disaggregation/test_per_layer_transfer.py

@@ -0,0 +1,91 @@
+"""Unit tests for A1: the async per-layer transfer primitive
+(MooncakeKVManager._transfer_layers_async).
+
+It must submit each layer's transfer non-blocking (one batch_transfer_async_submit
+per non-empty layer), then wait for all once; skip empty layers; and on a submit
+failure, drain what was submitted and return -1. Engine is mocked (no RDMA/GPU).
+"""
+
+import unittest
+
+from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
+
+
+class _FakeEngine:
+    def __init__(self, submit_rets=None, wait_ret=0):
+        self.submits = []
+        self.waits = []
+        self._submit_rets = submit_rets
+        self._wait_ret = wait_ret
+        self._n = 0
+
+    def batch_transfer_async_submit(self, sid, src, dst, lens):
+        self.submits.append((sid, list(src), list(dst), list(lens)))
+        if self._submit_rets is not None:
+            r = self._submit_rets[self._n]
+            self._n += 1
+            return r
+        return 100 + len(self.submits)
+
+    def wait_batch_transfers(self, ids):
+        self.waits.append(list(ids))
+        return self._wait_ret
+
+
+def _mgr(engine):
+    m = MooncakeKVManager.__new__(MooncakeKVManager)  # skip heavy init
+    m.engine = engine
+    return m
+
+
+def _one_block_per_layer(src, dst, item_len):
+    return [(src, dst, item_len)]
+
+
+class TestPerLayerAsyncTransfer(unittest.TestCase):
+    def test_one_submit_per_layer_then_single_wait(self):
+        eng = _FakeEngine()
+        status = _mgr(eng)._transfer_layers_async(
+            "sess", [(1000, 2000, 64), (1064, 2064, 64), (1128, 2128, 64)],
+            _one_block_per_layer,
+        )
+        self.assertEqual(status, 0)
+        self.assertEqual(len(eng.submits), 3)          # one async submit per layer
+        self.assertEqual(len(eng.waits), 1)            # single wait-for-all
+        self.assertEqual(eng.waits[0], [101, 102, 103])
+        self.assertEqual(eng.submits[0], ("sess", [1000], [2000], [64]))
+
+    def test_skips_empty_layers(self):
+        eng = _FakeEngine()
+
+        def blocks(src, dst, il):
+            return [] if src == 1064 else [(src, dst, il)]
+
+        _mgr(eng)._transfer_layers_async(
+            "s", [(1000, 2000, 64), (1064, 2064, 64), (1128, 2128, 64)], blocks
+        )
+        self.assertEqual(len(eng.submits), 2)          # middle layer skipped
+        self.assertEqual(eng.waits[0], [101, 102])
+
+    def test_submit_failure_drains_and_returns_minus_one(self):
+        eng = _FakeEngine(submit_rets=[101, -1])       # 2nd submit fails
+        status = _mgr(eng)._transfer_layers_async(
+            "s", [(1, 2, 3), (4, 5, 6), (7, 8, 9)], _one_block_per_layer
+        )
+        self.assertEqual(status, -1)
+        self.assertEqual(len(eng.submits), 2)          # stopped at the failure
+        self.assertEqual(eng.waits, [[101]])           # drained the first submit
+
+    def test_wait_status_propagates(self):
+        eng = _FakeEngine(wait_ret=-7)
+        status = _mgr(eng)._transfer_layers_async("s", [(1, 2, 3)], _one_block_per_layer)
+        self.assertEqual(status, -7)
+
+    def test_no_layers_is_zero(self):
+        eng = _FakeEngine()
+        self.assertEqual(_mgr(eng)._transfer_layers_async("s", [], _one_block_per_layer), 0)
+        self.assertEqual(eng.waits, [[]])
+
+
+if __name__ == "__main__":
+    unittest.main()