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>
2026-06-07 09:51:05 +00:00
parent 6bd5d5760c
commit aa6acc9485
2 changed files with 113 additions and 0 deletions
@@ -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()