diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index 7b46d68f7..912c99c80 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -31,11 +31,15 @@ def handle_pd_disaggregation(server_args: "ServerArgs") -> None: "--disaggregation-decode-enable-radix-cache is incompatible " "with --enable-hisparse" ) - if server_args.disaggregation_transfer_backend not in ("nixl", "mooncake"): + if server_args.disaggregation_transfer_backend not in ( + "nixl", + "mooncake", + "mori", + ): raise ValueError( "--disaggregation-decode-enable-radix-cache currently " "requires --disaggregation-transfer-backend in " - "('nixl', 'mooncake'), but got " + "('nixl', 'mooncake', 'mori'), but got " f"{server_args.disaggregation_transfer_backend!r}" ) if server_args.speculative_algorithm is not None: diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index ad33b9b8b..8aada7c25 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -119,6 +119,11 @@ class TransferInfo: dst_state_indices: List[npt.NDArray[np.int32]] required_dst_info_num: int is_dummy: bool + # Number of tokens decode already holds in its radix cache; prefill should + # only send pages beyond this prefix. None means the receiver did not + # populate this field (older receiver or radix-cache feature off) -> treat + # as 0 (no prefix hit, full send) for backward compatibility. + decode_prefix_len: Optional[int] = None @classmethod def from_zmq(cls, payload: List[bytes]) -> TransferInfo: @@ -145,7 +150,19 @@ class TransferInfo: required_dst_info_num = ( int(payload[7].decode("ascii")) if len(payload) > 7 else 1 ) - is_dummy = dst_kv_indices.size == 0 and dst_aux_index < 0 + + if len(payload) > 8 and payload[8]: + decode_prefix_len: Optional[int] = int(payload[8].decode("ascii")) + else: + decode_prefix_len = None + + # A transfer is "dummy" only when the receiver does not need any + # kv/aux/state delivered. When decode_prefix_len > 0 and the delta is + # exactly zero (full prefix hit), dst_kv_indices is empty but aux is + # still needed -> not dummy. + is_dummy = ( + dst_kv_indices.size == 0 and dst_aux_index < 0 and not decode_prefix_len + ) return cls( room=room, endpoint=endpoint, @@ -156,6 +173,7 @@ class TransferInfo: dst_state_indices=dst_state_indices, required_dst_info_num=required_dst_info_num, is_dummy=is_dummy, + decode_prefix_len=decode_prefix_len, ) @@ -488,11 +506,37 @@ class MoriKVManager(CommonKVManager): infos[transfer_info.engine_key] = transfer_info if len(infos) >= transfer_info.required_dst_info_num: - logger.debug( - "Bootstrap room %s got enough transfer info (%s)", - transfer_info.room, - len(infos), + # All decode peers reported their dst metadata; pick a + # non-None decode_prefix_len if any peer set it (they + # should all agree, but be defensive). 0 means "no + # prefix hit", which is the same as "feature off". + chosen_prefix_len = next( + ( + info.decode_prefix_len + for info in infos.values() + if info.decode_prefix_len is not None + ), + 0, ) + self.req_to_decode_prefix_len[transfer_info.room] = ( + chosen_prefix_len + ) + if chosen_prefix_len > 0: + # Surface incremental KV transfer at INFO so it's + # visible without bumping the global log level. + logger.info( + "MoriKV incremental: room=%s prefix_len=%s peers=%s", + transfer_info.room, + chosen_prefix_len, + len(infos), + ) + else: + logger.debug( + "Bootstrap room %s got enough transfer info (%s), " + "decode_prefix_len=0", + transfer_info.room, + len(infos), + ) self.update_status(transfer_info.room, KVPoll.WaitingForInput) except Exception: logger.exception("Failed to parse transfer info message") @@ -1628,6 +1672,12 @@ class MoriKVReceiver(CommonKVReceiver): aux_bytes = str(aux_index).encode("ascii") if aux_index is not None else b"" normalized_state = _normalize_state_indices_per_component(state_indices) + decode_prefix_bytes = ( + str(int(decode_prefix_len)).encode("ascii") + if decode_prefix_len is not None and decode_prefix_len > 0 + else b"" + ) + for bootstrap_info in self.bootstrap_infos: sock, lock = self._connect_to_bootstrap_server(bootstrap_info) is_dummy = bootstrap_info.get("is_dummy", False) @@ -1647,6 +1697,7 @@ class MoriKVReceiver(CommonKVReceiver): aux_bytes if not is_dummy else b"", state_bytes, str(self.required_dst_info_num).encode("ascii"), + decode_prefix_bytes, ] ) self.init_time = time.time() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4bd7ff922..a15a62cb1 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7359,7 +7359,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 or mooncake 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, mooncake or mori and is incompatible with --enable-hisparse.", ) parser.add_argument( "--disaggregation-decode-enable-offload-kvcache", diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 348025739..c55d85ef4 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -118,7 +118,7 @@ class TestLoadBalanceMethod(unittest.TestCase): disaggregation_transfer_backend="fake", ) - self.assertIn("('nixl', 'mooncake')", str(context.exception)) + self.assertIn("('nixl', 'mooncake', 'mori')", str(context.exception)) self.assertIn("'fake'", str(context.exception))