[PD][AMD]: incremental KV transfer with decode radix cache (#26288)

Signed-off-by: inkcherry <mingzhi.liu@amd.com>
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>

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:

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()

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",

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))