feat(disagg/mooncake): build_per_layer_context assembly (lever A, A3-step2)

Add MooncakeKVManager.build_per_layer_context: assembles a PerLayerTransferContext
from the SAME CP-filtered (prefill_kv_indices, dst_kv_indices) the post-forward
transfer uses — so the bytes moved are byte-identical to the monolithic path, and
the CP owner mapping is NOT re-derived (eliminating the #1 correctness risk). It
mirrors the MLA branch of _send_kvcache_generic exactly (get_mla_kv_ptrs_with_pp +
group_concurrent_contiguous + build_layer_blocks, verified set_transfer_blocks-
identical). Returns None for MHA / unregistered decode / empty owned set.

Unit-tested (4 cases): per-layer address correctness + the None guards. Remaining
A3-step3: call this in the send/scheduler flow (register before forward, finish
after, skip the main-KV monolithic send), then output-equality + TTFT verification.

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

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,

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()