diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 2c0aa805a..536fafe2b 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d1053d5b0..79b024215 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/test/registered/unit/disaggregation/test_per_layer_transfer.py b/test/registered/unit/disaggregation/test_per_layer_transfer.py new file mode 100644 index 000000000..5a54fd76b --- /dev/null +++ b/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()