diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 1272f3a3b..fe0d65103 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ec4c4a279..2c04e955e 100644 --- a/python/sglang/srt/server_args.py +++ b/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", diff --git a/test/registered/distributed/test_disaggregation_decode_radix_cache.py b/test/registered/distributed/test_disaggregation_decode_radix_cache.py index 53356a6cb..3296f847f 100644 --- a/test/registered/distributed/test_disaggregation_decode_radix_cache.py +++ b/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() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index a62591158..16b9ac415 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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")