diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 536fafe2b..d7d47b3ec 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -349,6 +349,50 @@ class MooncakeKVManager(CommonKVManager): batch_ids.append(batch_id) return self.engine.wait_batch_transfers(batch_ids) + def build_per_layer_context( + self, + mooncake_session_id: str, + prefill_kv_indices: npt.NDArray[np.int32], + dst_kv_indices: npt.NDArray[np.int32], + ): + """Build a PerLayerTransferContext (lever A) for one request's main-KV pages, + from the SAME CP-filtered (prefill_kv_indices, dst_kv_indices) the post-forward + transfer would use — so the bytes moved are identical to the monolithic path. + The caller supplies the indices (reusing the send() CP filter, so this does NOT + re-derive the CP owner mapping — the #1 correctness risk). Mirrors the MLA + branch of _send_kvcache_generic exactly. Returns None if not applicable + (MHA / no decode registration yet / empty owned set).""" + if not self.is_mla_backend: + return None # first lever-A impl targets MLA (the production GLM/NSA path) + reg = self.decode_kv_args_table.get(mooncake_session_id) + if reg is None or len(prefill_kv_indices) == 0: + return None + from sglang.srt.disaggregation.cp_per_layer_transfer import ( + PerLayerTransferContext, + build_layer_blocks, + ) + + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + src_kv_ptrs, dst_kv_ptrs, n_layers = self.get_mla_kv_ptrs_with_pp( + self.kv_args.kv_data_ptrs, reg.dst_kv_ptrs + ) + item_lens = self.kv_args.kv_item_lens + + def get_blocks(layer_id): + if layer_id >= n_layers: + return None + return build_layer_blocks( + src_kv_ptrs[layer_id], + dst_kv_ptrs[layer_id], + item_lens[layer_id], + prefill_kv_blocks, + dst_kv_blocks, + ) + + return PerLayerTransferContext(self.engine, mooncake_session_id, get_blocks) + def _send_kvcache_generic( self, mooncake_session_id: str, diff --git a/test/registered/unit/disaggregation/test_per_layer_transfer.py b/test/registered/unit/disaggregation/test_per_layer_transfer.py index 5a54fd76b..b3662a52c 100644 --- a/test/registered/unit/disaggregation/test_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_per_layer_transfer.py @@ -87,5 +87,74 @@ class TestPerLayerAsyncTransfer(unittest.TestCase): self.assertEqual(eng.waits, [[]]) +class _Reg: + pass + + +class _KVArgs: + pass + + +def _mgr_for_ctx(src_ptrs, dst_ptrs, item_lens, is_mla=True, registered=True): + m = MooncakeKVManager.__new__(MooncakeKVManager) + m.is_mla_backend = is_mla + m.engine = object() + kv = _KVArgs() + kv.kv_data_ptrs = src_ptrs + kv.kv_item_lens = item_lens + kv.prefill_start_layer = 0 + m.kv_args = kv + if registered: + reg = _Reg() + reg.dst_kv_ptrs = dst_ptrs + m.decode_kv_args_table = {"sess": reg} + else: + m.decode_kv_args_table = {} + return m + + +class TestBuildPerLayerContext(unittest.TestCase): + def test_get_blocks_matches_per_layer_math(self): + import numpy as np + + m = _mgr_for_ctx([1000, 9000], [5000, 8000], [64, 64]) + ctx = m.build_per_layer_context( + "sess", + np.array([0, 1, 2], dtype=np.int32), # one contiguous prefill run + np.array([10, 11, 12], dtype=np.int32), # -> one contiguous dst run + ) + self.assertIsNotNone(ctx) + # layer 0: src=1000+0*64, dst=5000+10*64, len=64*3 (3-page run) + self.assertEqual(ctx.get_blocks(0), ([1000], [5000 + 10 * 64], [64 * 3])) + self.assertEqual(ctx.get_blocks(1), ([9000], [8000 + 10 * 64], [64 * 3])) + self.assertIsNone(ctx.get_blocks(2)) # only 2 layers + + def test_none_for_mha(self): + import numpy as np + + m = _mgr_for_ctx([1], [2], [3], is_mla=False) + self.assertIsNone( + m.build_per_layer_context("sess", np.array([0], np.int32), np.array([0], np.int32)) + ) + + def test_none_when_decode_not_registered(self): + import numpy as np + + m = _mgr_for_ctx([1], [2], [3], registered=False) + self.assertIsNone( + m.build_per_layer_context("sess", np.array([0], np.int32), np.array([0], np.int32)) + ) + + def test_none_for_empty_indices(self): + import numpy as np + + m = _mgr_for_ctx([1], [2], [3]) + self.assertIsNone( + m.build_per_layer_context( + "sess", np.array([], np.int32), np.array([], np.int32) + ) + ) + + if __name__ == "__main__": unittest.main()