[RadixTree][9/N Refactor]: Support unified init_load_back params (#20590)

python/sglang/srt/managers/schedule_policy.py

@@ -38,6 +38,7 @@ from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
+    InitLoadBackParams,
     InsertParams,
     MatchPrefixParams,
 )
@@ -769,7 +770,10 @@ class PrefillAdder:
 
             if req.host_hit_length > 0:
                 new_indices, req.last_node = self.tree_cache.init_load_back(
-                    req.last_host_node, req.host_hit_length
+                    InitLoadBackParams(
+                        last_host_node=req.last_host_node,
+                        host_hit_length=req.host_hit_length,
+                    )
                 )
                 req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
                 req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -110,6 +110,16 @@ class DecLockRefResult:
     delta: Optional[int] = None
 
 
+@dataclasses.dataclass
+class InitLoadBackParams:
+    """Unified parameters for init_load_back across different cache types"""
+
+    last_host_node: Any
+    host_hit_length: int
+    mem_quota: Optional[int] = None
+    req: Optional[Req] = None
+
+
 class MatchResult(NamedTuple):
     """Result of a prefix match operation.
 
@@ -215,8 +225,7 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
 
     def init_load_back(
         self,
-        last_host_node: Any,
+        params: InitLoadBackParams,
-        host_hit_length: int,
     ) -> Tuple[torch.Tensor, Any]:
         """
         Preparing KV cache loading from host to device.

python/sglang/srt/mem_cache/hi_mamba_radix_cache.py

@@ -19,6 +19,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     EvictParams,
     EvictResult,
     IncLockRefResult,
+    InitLoadBackParams,
     MatchPrefixParams,
     MatchResult,
 )
@@ -256,10 +257,10 @@ class HiMambaRadixCache(MambaRadixCache):
 
     def init_load_back(
         self,
-        last_node: TreeNode,
+        params: InitLoadBackParams,
-        host_hit_length: int,
-        mem_quota: Optional[int] = None,
     ):
+        last_node = params.last_host_node
+        mem_quota = params.mem_quota
         if last_node.evicted:
             loading_values = self.load_back(last_node, mem_quota)
             if loading_values is not None:

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -20,6 +20,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     EvictParams,
     EvictResult,
     IncLockRefResult,
+    InitLoadBackParams,
     InsertParams,
     InsertResult,
     MatchPrefixParams,
@@ -1066,11 +1067,10 @@ class HiRadixCache(RadixCache):
 
     def init_load_back(
         self,
-        last_node: TreeNode,
+        params: InitLoadBackParams,
-        host_hit_length: int,
-        mem_quota: Optional[int] = None,
     ):
-        _ = host_hit_length  # unused, but kept for compatibility
+        last_node = params.last_host_node
+        mem_quota = params.mem_quota
         if last_node.evicted:
             loading_values = self.load_back(last_node, mem_quota)
             if loading_values is not None:

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -12,6 +12,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     EvictParams,
     EvictResult,
     IncLockRefResult,
+    InitLoadBackParams,
     MatchPrefixParams,
     MatchResult,
 )
@@ -319,8 +320,8 @@ class SessionAwareCache(BasePrefixCache):
     def pretty_print(self):
         return self.inner.pretty_print()
 
-    def init_load_back(self, last_host_node, host_hit_length):
+    def init_load_back(self, params: InitLoadBackParams):
-        return self.inner.init_load_back(last_host_node, host_hit_length)
+        return self.inner.init_load_back(params)
 
     def ready_to_load_host_cache(self):
         return self.inner.ready_to_load_host_cache()