diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 38c27d28a..08e50d655 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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)) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index f0df9e989..0efc27d44 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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, - host_hit_length: int, + params: InitLoadBackParams, ) -> Tuple[torch.Tensor, Any]: """ Preparing KV cache loading from host to device. diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 7ce972a85..13398b430 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/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, - host_hit_length: int, - mem_quota: Optional[int] = None, + params: InitLoadBackParams, ): + 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: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index d1e9cf1ae..e13d43c37 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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, - host_hit_length: int, - mem_quota: Optional[int] = None, + params: InitLoadBackParams, ): - _ = 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: diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 54664387c..10fa2fae0 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/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): - return self.inner.init_load_back(last_host_node, host_hit_length) + def init_load_back(self, params: InitLoadBackParams): + return self.inner.init_load_back(params) def ready_to_load_host_cache(self): return self.inner.ready_to_load_host_cache()