diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 7ebd5cd7b..38c27d28a 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor +from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams from sglang.srt.utils import get_bool_env_var _ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG") @@ -560,11 +561,9 @@ class PrefillAdder: self._update_prefill_budget(prefix_len, trunc_len, 0) def _req_inc_lock_ref(self, req: Req): + result = self.tree_cache.inc_lock_ref(req.last_node) if self.is_hybrid_swa: - swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node) - req.swa_uuid_for_lock = swa_uuid_for_lock - else: - self.tree_cache.inc_lock_ref(req.last_node) + req.swa_uuid_for_lock = result.swa_uuid_for_lock def add_dllm_staging_req(self, req: Req): assert self.dllm_config is not None @@ -624,14 +623,15 @@ class PrefillAdder: @contextmanager def _lock_node(self, last_node: TreeNode): try: + result = self.tree_cache.inc_lock_ref(last_node) if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache(): - swa_uuid_for_lock = self.tree_cache.inc_lock_ref(last_node) - else: - self.tree_cache.inc_lock_ref(last_node) + swa_uuid_for_lock = result.swa_uuid_for_lock yield None finally: if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache(): - self.tree_cache.dec_lock_ref(last_node, swa_uuid_for_lock) + self.tree_cache.dec_lock_ref( + last_node, DecLockRefParams(swa_uuid_for_lock=swa_uuid_for_lock) + ) else: self.tree_cache.dec_lock_ref(last_node) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index e5c251a83..4342ba82f 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -88,6 +88,28 @@ class EvictResult: mamba_num_evicted: int = 0 +@dataclasses.dataclass +class IncLockRefResult: + """Result of an inc_lock_ref operation.""" + + delta: Optional[int] = None + swa_uuid_for_lock: Optional[int] = None + + +@dataclasses.dataclass +class DecLockRefParams: + """Parameters for dec_lock_ref operation.""" + + swa_uuid_for_lock: Optional[int] = None + + +@dataclasses.dataclass +class DecLockRefResult: + """Result of an dec_lock_ref operation.""" + + delta: Optional[int] = None + + class MatchResult(NamedTuple): """Result of a prefix match operation. @@ -158,11 +180,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait): pass @abstractmethod - def inc_lock_ref(self, node: Any): + def inc_lock_ref(self, node: Any) -> IncLockRefResult: pass @abstractmethod - def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None): + def dec_lock_ref( + self, node: Any, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: pass def evictable_size(self): diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index f8e053285..5d58bcde5 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -9,8 +9,11 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, InsertParams, InsertResult, MatchPrefixParams, @@ -80,11 +83,13 @@ class ChunkCache(BasePrefixCache): def evict(self, params: EvictParams) -> EvictResult: return EvictResult() - def inc_lock_ref(self, node: Any): - return 0 + def inc_lock_ref(self, node: Any) -> IncLockRefResult: + return IncLockRefResult(delta=0) - def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None): - return 0 + def dec_lock_ref( + self, node: Any, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: + return DecLockRefResult(delta=0) def protected_size(self): # NOTE: no protected size in chunk cache. Chunk cache's eviction is the same with request's lifecycle. 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 e03d09569..2f738508c 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -14,8 +14,11 @@ import torch from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, MatchPrefixParams, MatchResult, ) @@ -198,7 +201,8 @@ class HiMambaRadixCache(MambaRadixCache): else: ancestor_node = node - delta = self.inc_lock_ref(ancestor_node) + result = self.inc_lock_ref(ancestor_node) + delta = result.delta full_host_indices = torch.cat([n.host_value for n in nodes_to_load]) if (len(full_host_indices) < self.load_back_threshold) or ( @@ -988,9 +992,9 @@ class HiMambaRadixCache(MambaRadixCache): return super().sanity_check() - def inc_lock_ref(self, node: TreeNode) -> Optional[int]: + def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: if self.disable: - return 0 + return IncLockRefResult(delta=0) delta = 0 if node.mamba_value is not None: @@ -1014,11 +1018,13 @@ class HiMambaRadixCache(MambaRadixCache): self.evictable_full_device_leaves.discard(node) node.full_lock_ref += 1 node = node.parent - return delta + return IncLockRefResult(delta=delta) - def dec_lock_ref(self, node: TreeNode): + def dec_lock_ref( + self, node: TreeNode, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: if self.disable: - return 0 + return DecLockRefResult(delta=0) delta = 0 @@ -1044,7 +1050,7 @@ class HiMambaRadixCache(MambaRadixCache): if node.full_lock_ref == 0: self._update_full_device_leaf_status(node) node = node.parent - return delta + return DecLockRefResult(delta=delta) # ---- L3 Support ---- diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 041318bf2..6f9190399 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -15,8 +15,11 @@ import torch from sglang.srt.environ import envs from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, InsertParams, InsertResult, MatchPrefixParams, @@ -817,9 +820,9 @@ class HiRadixCache(RadixCache): """ return RadixKey(token_ids=list(token_ids)) - def inc_lock_ref(self, node: TreeNode): + def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: if self.disable: - return 0 + return IncLockRefResult(delta=0) delta = 0 while node != self.root_node: @@ -831,11 +834,13 @@ class HiRadixCache(RadixCache): self._update_leaf_status(node) self._update_host_leaf_status(node) node = node.parent - return delta + return IncLockRefResult(delta=delta) - def dec_lock_ref(self, node: TreeNode): + def dec_lock_ref( + self, node: TreeNode, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: if self.disable: - return 0 + return DecLockRefResult(delta=0) delta = 0 while node != self.root_node: @@ -851,7 +856,7 @@ class HiRadixCache(RadixCache): node is self.root_node ), f"This request holds the node from another tree" node = node.parent - return delta + return DecLockRefResult(delta=delta) def _update_host_leaf_status(self, node: TreeNode): if not node.evicted or node.lock_ref > 0: @@ -1011,7 +1016,8 @@ class HiRadixCache(RadixCache): ancester_node = node # protect the ancestor nodes from eviction - delta = self.inc_lock_ref(ancester_node) + result = self.inc_lock_ref(ancester_node) + delta = result.delta # load it all or not at all host_indices = torch.cat([n.host_value for n in nodes_to_load]) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index ac27d93c1..147738e55 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -35,8 +35,11 @@ from sglang.srt.mem_cache.allocator import ( ) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, InsertParams, InsertResult, MatchPrefixParams, @@ -806,14 +809,14 @@ class MambaRadixCache(BasePrefixCache): return full_num_evicted - def inc_lock_ref(self, node: TreeNode) -> Optional[int]: + def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: """ Increment the lock reference count for the node. It locks the full_lock_ref for nodes between the [last node, root), exclusive. It locks the mamba_lock_ref for current node if its mamba_value exists. """ if self.disable: - return None + return IncLockRefResult() # protect mamba value in current node if it exists if node.mamba_value is not None: @@ -832,16 +835,18 @@ class MambaRadixCache(BasePrefixCache): self.full_protected_size_ += len(node.value) node.full_lock_ref += 1 node = node.parent - return None + return IncLockRefResult() - def dec_lock_ref(self, node: TreeNode): + def dec_lock_ref( + self, node: TreeNode, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: """ Decrement the lock reference count for the node. It unlocks the full_lock_ref for nodes between the [last node, root), exclusive. It unlocks the mamba_lock_ref for current node if its mamba_value exists. """ if self.disable: - return None + return DecLockRefResult() if node.mamba_value is not None: assert ( @@ -862,6 +867,8 @@ class MambaRadixCache(BasePrefixCache): node.full_lock_ref -= 1 node = node.parent + return DecLockRefResult() + def sanity_check(self): if self.disable: return diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 8e13ff4d1..501fc4223 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -42,8 +42,11 @@ from sglang.srt.disaggregation.kv_events import ( ) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, InsertParams, InsertResult, MatchPrefixParams, @@ -609,9 +612,9 @@ class RadixCache(BasePrefixCache): self.update_eviction_metrics(num_evicted, start_time) return EvictResult(num_tokens_evicted=num_evicted) - def inc_lock_ref(self, node: TreeNode): + def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: if self.disable: - return 0 + return IncLockRefResult(delta=0) delta = 0 while node != self.root_node: @@ -622,11 +625,13 @@ class RadixCache(BasePrefixCache): node.lock_ref += 1 self._update_leaf_status(node) node = node.parent - return delta + return IncLockRefResult(delta=delta) - def dec_lock_ref(self, node: TreeNode): + def dec_lock_ref( + self, node: TreeNode, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: if self.disable: - return 0 + return DecLockRefResult(delta=0) delta = 0 while node != self.root_node: @@ -641,7 +646,7 @@ class RadixCache(BasePrefixCache): node is self.root_node ), f"This request holds the node from another tree" node = node.parent - return delta + return DecLockRefResult(delta=delta) def evictable_size(self): return self.evictable_size_ diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 6ed293d1b..66f9fad96 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -2,14 +2,17 @@ from __future__ import annotations import logging import time -from typing import TYPE_CHECKING, List, Set +from typing import TYPE_CHECKING, List, Optional, Set import torch from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, MatchPrefixParams, MatchResult, ) @@ -123,21 +126,25 @@ class RadixCacheCpp(BasePrefixCache): raise NotImplementedError("Host cache is not supported yet") - def dec_lock_ref(self, node: TreeNodeCpp): + def dec_lock_ref( + self, node: TreeNodeCpp, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: """ Decrement the reference count of a node to root of the radix tree. Args: node (TreeNodeCpp): The handle of the node to decrement the reference count for. """ self.tree.lock_ref(node, False) # do not increment + return DecLockRefResult() - def inc_lock_ref(self, node: TreeNodeCpp): + def inc_lock_ref(self, node: TreeNodeCpp) -> IncLockRefResult: """ Increment the reference count of from a node to root of the radix tree. Args: node (TreeNodeCpp): The handle of the node to increment the reference count for. """ self.tree.lock_ref(node, True) + return IncLockRefResult() def evict(self, params: EvictParams) -> EvictResult: start_time = time.perf_counter() diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 873acf28d..3405918bb 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -7,8 +7,11 @@ import torch from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, MatchPrefixParams, MatchResult, ) @@ -223,17 +226,17 @@ class SessionAwareCache(BasePrefixCache): def evict(self, params: EvictParams) -> EvictResult: return self.inner.evict(params) - def inc_lock_ref(self, node: Any): + def inc_lock_ref(self, node: Any) -> IncLockRefResult: if isinstance(node, _VirtualNode): - return None + return IncLockRefResult() return self.inner.inc_lock_ref(node) - def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None): + def dec_lock_ref( + self, node: Any, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: if isinstance(node, _VirtualNode): - return - if swa_uuid_for_lock is not None: - return self.inner.dec_lock_ref(node, swa_uuid_for_lock) - return self.inner.dec_lock_ref(node) + return DecLockRefResult() + return self.inner.dec_lock_ref(node, params) # -- Session lifecycle -- @@ -245,7 +248,10 @@ class SessionAwareCache(BasePrefixCache): if slot.last_node is not None: if slot.swa_uuid_for_lock is not None: - self.inner.dec_lock_ref(slot.last_node, slot.swa_uuid_for_lock) + self.inner.dec_lock_ref( + slot.last_node, + DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock), + ) else: self.inner.dec_lock_ref(slot.last_node) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 50b24a313..34df1617d 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -30,8 +30,11 @@ from numpy import float64 from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, + DecLockRefParams, + DecLockRefResult, EvictParams, EvictResult, + IncLockRefResult, InsertParams, InsertResult, MatchPrefixParams, @@ -485,7 +488,9 @@ class SWARadixCache(BasePrefixCache): self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:]) # Remove req slot release the cache lock - self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock) + self.dec_lock_ref( + req.last_node, DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock) + ) def cache_unfinished_req(self, req: Req, chunked=False) -> None: """Cache request when it is unfinished.""" @@ -536,8 +541,11 @@ class SWARadixCache(BasePrefixCache): req.cache_protected_len = len(new_indices) - self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock) - swa_uuid_for_lock = self.inc_lock_ref(new_last_node) + self.dec_lock_ref( + req.last_node, DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock) + ) + result = self.inc_lock_ref(new_last_node) + swa_uuid_for_lock = result.swa_uuid_for_lock # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later if len(new_indices) < len(kv_indices): @@ -647,7 +655,7 @@ class SWARadixCache(BasePrefixCache): num_tokens_evicted=full_num_evicted, swa_num_tokens_evicted=swa_num_evicted ) - def inc_lock_ref(self, node: TreeNode) -> Optional[int]: + def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: """ Increment the lock reference count for the node. Returns the swa_uuid_for_lock, which needs to be passed to dec_lock_ref. @@ -655,7 +663,7 @@ class SWARadixCache(BasePrefixCache): It locks the swa_lock_ref for nodes between the [last node, swa_uuid_for_lock], inclusive. """ if self.disable: - return None + return IncLockRefResult() swa_lock_size = 0 swa_uuid_for_lock = None @@ -686,17 +694,21 @@ class SWARadixCache(BasePrefixCache): node.swa_uuid = gen_swa_uuid() swa_uuid_for_lock = node.swa_uuid node = node.parent - return swa_uuid_for_lock + return IncLockRefResult(swa_uuid_for_lock=swa_uuid_for_lock) - def dec_lock_ref(self, node: TreeNode, swa_uuid_for_lock: Optional[int] = None): + def dec_lock_ref( + self, node: TreeNode, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: """ Decrement the lock reference count for the node. It unlocks the full_lock_ref for nodes between the [last node, root), exclusive. It unlocks the swa_lock_ref for nodes between the [last node, swa_uuid_for_lock], inclusive. If swa_uuid_for_lock is None, it unlocks to the root, exclusive. """ + swa_uuid_for_lock = params.swa_uuid_for_lock if params is not None else None + if self.disable: - return + return DecLockRefResult() dec_lock_swa = True while node != self.root_node: @@ -725,6 +737,8 @@ class SWARadixCache(BasePrefixCache): node = node.parent + return DecLockRefResult() + def sanity_check(self): self.full_lru_list.sanity_check(self) self.swa_lru_list.sanity_check(self) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index f2fe6d193..ea1ed8bc1 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -4,6 +4,10 @@ from unittest.mock import MagicMock, patch from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefResult, + IncLockRefResult, +) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -34,8 +38,8 @@ class TestPrefillAdder(CustomTestCase): tree_cache.swa_evictable_size.return_value = swa_evictable_size tree_cache.evictable_size.return_value = evictable_size tree_cache.disable = False - tree_cache.inc_lock_ref.return_value = None - tree_cache.dec_lock_ref.return_value = None + tree_cache.inc_lock_ref.return_value = IncLockRefResult() + tree_cache.dec_lock_ref.return_value = DecLockRefResult() return tree_cache def create_token_allocator(