[PD]: Support incremental transfer for mooncake transfer engine (#24257)

Co-authored-by: Shangming Cai <csmthu@gmail.com>

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -81,6 +81,7 @@ class TransferInfo:
     dst_state_indices: List[int]
     required_dst_info_num: int
     is_dummy: bool
+    decode_prefix_len: Optional[int] = None
     # Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed)
     staging: Optional[StagingTransferInfo] = None
 
@@ -109,6 +110,9 @@ class TransferInfo:
             dst_state_indices=dst_state_indices,
             required_dst_info_num=int(msg[7].decode("ascii")),
             is_dummy=is_dummy,
+            decode_prefix_len=(
+                int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None
+            ),
         )
 
 
@@ -1217,7 +1221,9 @@ class MooncakeKVManager(CommonKVManager):
                         target_rank_registration_info: KVArgsRegisterInfo = (
                             self.decode_kv_args_table[req.mooncake_session_id]
                         )
-                        if self.is_mla_backend or (
+                        if len(kv_chunk.prefill_kv_indices) == 0:
+                            ret = 0
+                        elif self.is_mla_backend or (
                             self.attn_tp_size
                             == target_rank_registration_info.dst_attn_tp_size
                         ):
@@ -1340,6 +1346,7 @@ class MooncakeKVManager(CommonKVManager):
                 ):
                     if kv_chunk.room in self.transfer_infos:
                         self.transfer_infos.pop(kv_chunk.room)
+                    self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
 
             except Exception as e:
                 # NOTE(shangming): Remove this when we make sure the transfer thread is bug-free
@@ -1423,6 +1430,14 @@ class MooncakeKVManager(CommonKVManager):
                     )
                     # NOTE: after bootstrapping we can mark the req as waiting for input
                     if len(self.transfer_infos[room]) == required_dst_info_num:
+                        self.req_to_decode_prefix_len[room] = next(
+                            (
+                                info.decode_prefix_len
+                                for info in self.transfer_infos[room].values()
+                                if info.decode_prefix_len is not None
+                            ),
+                            0,
+                        )
                         self.update_status(room, KVPoll.WaitingForInput)
 
         threading.Thread(target=bootstrap_thread).start()
@@ -1650,6 +1665,12 @@ class MooncakeKVSender(CommonKVSender):
         self.conclude_state = None
         self.init_time = time.time()
 
+    def pop_decode_prefix_len(self) -> int:
+        return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
+
+    def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
+        return num_pages > 0 or last_chunk
+
     def send(
         self,
         kv_indices: npt.NDArray[np.int32],
@@ -1868,6 +1889,7 @@ class MooncakeKVReceiver(CommonKVReceiver):
                             else b""
                         ),
                         str(self.required_dst_info_num).encode("ascii"),
+                        str(decode_prefix_len or 0).encode("ascii"),
                     ]
                 )
         self.init_time = time.time()

python/sglang/srt/server_args.py

@@ -3744,10 +3744,12 @@ class ServerArgs:
                         "--disaggregation-decode-enable-radix-cache is incompatible "
                         "with --enable-hisparse"
                     )
-                if self.disaggregation_transfer_backend != "nixl":
+                if self.disaggregation_transfer_backend not in ("nixl", "mooncake"):
                     raise ValueError(
                         "--disaggregation-decode-enable-radix-cache currently "
-                        "requires --disaggregation-transfer-backend nixl"
+                        "requires --disaggregation-transfer-backend in "
+                        "('nixl', 'mooncake'), but got "
+                        f"{self.disaggregation_transfer_backend!r}"
                     )
                 if self.speculative_algorithm is not None:
                     raise ValueError(
@@ -6453,7 +6455,7 @@ class ServerArgs:
         parser.add_argument(
             "--disaggregation-decode-enable-radix-cache",
             action="store_true",
-            help="Enable radix cache on decode server (PD mode). Caches KV prefixes to avoid redundant transfers. Requires --disaggregation-transfer-backend nixl and is incompatible with --enable-hisparse.",
+            help="Enable radix cache on decode server (PD mode). Caches KV prefixes to avoid redundant transfers. Requires --disaggregation-transfer-backend nixl or mooncake and is incompatible with --enable-hisparse.",
         )
         parser.add_argument(
             "--disaggregation-decode-enable-offload-kvcache",

test/registered/distributed/test_disaggregation_decode_radix_cache.py

@@ -16,7 +16,7 @@ from sglang.test.test_utils import (
     try_cached_model,
 )
 
-register_cuda_ci(est_time=120, suite="stage-c-test-8-gpu-h20")
+register_cuda_ci(est_time=300, suite="stage-c-test-8-gpu-h20")
 
 
 def _has_nixl():
@@ -27,18 +27,26 @@ def _has_nixl():
     return True
 
 
-@unittest.skipUnless(
-    is_in_ci() or _has_nixl(),
-    "NIXL is required for decode radix cache disaggregation coverage.",
-)
-class TestDisaggregationDecodeRadixCache(PDDisaggregationServerBase):
+def _has_mooncake():
+    try:
+        import mooncake.engine  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+class DisaggregationDecodeRadixCacheTestMixin:
     extra_decode_args = ["--disaggregation-decode-enable-radix-cache"]
+    transfer_backend_name = None
 
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
         cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
-        cls.transfer_backend = ["--disaggregation-transfer-backend", "nixl"]
+        cls.transfer_backend = [
+            "--disaggregation-transfer-backend",
+            cls.transfer_backend_name,
+        ]
         cls.launch_all()
 
     def _assert_process_healthy(self, name, process, url):
@@ -99,11 +107,9 @@ class TestDisaggregationDecodeRadixCache(PDDisaggregationServerBase):
         metrics_second = run_eval(args)
         print(f"Second run metrics: {metrics_second}")
 
-        # Both runs should have reasonable accuracy
         self.assertGreater(metrics_first["score"], 0.80)
         self.assertGreater(metrics_second["score"], 0.80)
 
-        # Second run accuracy should not drop more than 3% compared to first run
         accuracy_drop = metrics_first["score"] - metrics_second["score"]
         self.assertLessEqual(
             accuracy_drop,
@@ -114,5 +120,25 @@ class TestDisaggregationDecodeRadixCache(PDDisaggregationServerBase):
         )
 
 
+@unittest.skipUnless(
+    is_in_ci() or _has_nixl(),
+    "NIXL is required for decode radix cache disaggregation coverage.",
+)
+class TestDisaggregationDecodeRadixCacheNixl(
+    DisaggregationDecodeRadixCacheTestMixin, PDDisaggregationServerBase
+):
+    transfer_backend_name = "nixl"
+
+
+@unittest.skipUnless(
+    is_in_ci() or _has_mooncake(),
+    "Mooncake is required for decode radix cache disaggregation coverage.",
+)
+class TestDisaggregationDecodeRadixCacheMooncake(
+    DisaggregationDecodeRadixCacheTestMixin, PDDisaggregationServerBase
+):
+    transfer_backend_name = "mooncake"
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/server_args/test_server_args.py

@@ -63,6 +63,28 @@ class TestLoadBalanceMethod(unittest.TestCase):
             str(context.exception),
         )
 
+    def test_pd_decode_radix_cache_allows_mooncake(self):
+        server_args = ServerArgs(
+            model_path="dummy",
+            disaggregation_mode="decode",
+            disaggregation_decode_enable_radix_cache=True,
+            disaggregation_transfer_backend="mooncake",
+        )
+
+        self.assertFalse(server_args.disable_radix_cache)
+
+    def test_pd_decode_radix_cache_rejects_unknown_backend(self):
+        with self.assertRaises(ValueError) as context:
+            ServerArgs(
+                model_path="dummy",
+                disaggregation_mode="decode",
+                disaggregation_decode_enable_radix_cache=True,
+                disaggregation_transfer_backend="fake",
+            )
+
+        self.assertIn("('nixl', 'mooncake')", str(context.exception))
+        self.assertIn("'fake'", str(context.exception))
+
 
 class TestPortArgs(unittest.TestCase):
     @patch("sglang.srt.server_args.get_free_port")