fix(cp_per_layer_transfer): gate per-layer path on a single non-dummy decode info

The transfer worker iterates every non-dummy decode info for a room and calls per_layer_mgr.finish() once per info, but register_per_layer_transfer registers exactly one context per room/chunk (built for one info's dst_kv_indices). This is only sound when there is exactly one non-dummy info (required_dst_info_num == 1). With decode attn_tp < prefill attn_tp a single prefill rank holds >1 non-dummy infos; finishing once-per-info would over-pop chunk contexts and under-deliver KV to the other infos. Make the assumption explicit: register only when there is one non-dummy info, otherwise fall back to the monolithic post-forward transfer (which fans out to all infos correctly). Found by an independent first-principles audit. Adds TestRegisterGuardSingleInfo (2-info fallback, 1-info register, all-dummy fallback) exercising the real MooncakeKVManager.register_per_layer_transfer. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-07 09:51:05 +00:00
parent e9354c41bc
commit 243f1b964c
2 changed files with 96 additions and 20 deletions
@@ -395,5 +395,71 @@ class TestChunkedDstMapping(unittest.TestCase):
        self.fail("no rank owned a chunk-1 page; adjust fixture")


+class TestRegisterGuardSingleInfo(unittest.TestCase):
+    """register_per_layer_transfer must register a context only when there is
+    exactly one non-dummy decode info. The transfer worker calls finish() once
+    per non-dummy info for the room while the per-layer path registers a single
+    context, so >1 non-dummy info (decode attn_tp < prefill attn_tp) must fall
+    back to the monolithic transfer."""
+
+    def _fake_self(self, infos, recorder):
+        import types
+
+        mgr = types.SimpleNamespace(
+            register=lambda room, ctx, chunk_key=0: recorder.append(
+                (room, chunk_key)
+            ),
+        )
+        return types.SimpleNamespace(
+            per_layer_transfer_manager=mgr,
+            server_args=types.SimpleNamespace(enable_nsa_prefill_cp_shared_kv=True),
+            transfer_infos={7: infos},
+            kv_args=types.SimpleNamespace(page_size=64),
+            attn_cp_size=8,
+            attn_cp_rank=0,
+            build_per_layer_context=lambda sid, owned, dst: ("ctx", sid),
+        )
+
+    def _info(self, is_dummy):
+        import types
+
+        import numpy as np
+
+        return types.SimpleNamespace(
+            is_dummy=is_dummy,
+            dst_kv_indices=np.arange(64, dtype=np.int32),
+            mooncake_session_id="s",
+        )
+
+    def _call(self, infos, recorder):
+        try:
+            from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
+        except Exception as e:  # pragma: no cover - env without mooncake
+            self.skipTest(f"mooncake conn not importable: {e}")
+        s = self._fake_self(infos, recorder)
+        # logical pages incl. page 1 which cp_rank 0 owns ((1-1)%8==0)
+        return MooncakeKVManager.register_per_layer_transfer(
+            s, 7, [0, 1, 2, 3], chunk_key=0
+        )
+
+    def test_two_non_dummy_infos_falls_back(self):
+        rec = []
+        infos = {"a": self._info(False), "b": self._info(False)}
+        self.assertFalse(self._call(infos, rec))
+        self.assertEqual(rec, [])  # nothing registered
+
+    def test_single_non_dummy_info_registers(self):
+        rec = []
+        infos = {"a": self._info(True), "b": self._info(False)}  # 1 dummy + 1 real
+        self.assertTrue(self._call(infos, rec))
+        self.assertEqual(len(rec), 1)
+
+    def test_all_dummy_falls_back(self):
+        rec = []
+        infos = {"a": self._info(True), "b": self._info(True)}
+        self.assertFalse(self._call(infos, rec))
+        self.assertEqual(rec, [])
+
+
 if __name__ == "__main__":
    unittest.main()