[RadixTree][8/N Refactor]: unify lock interface (#20330)

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 = result.swa_uuid_for_lock
-            req.swa_uuid_for_lock = swa_uuid_for_lock
-        else:
-            self.tree_cache.inc_lock_ref(req.last_node)
 
     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)
+                swa_uuid_for_lock = result.swa_uuid_for_lock
-            else:
-                self.tree_cache.inc_lock_ref(last_node)
             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)
 

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

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):
+    def inc_lock_ref(self, node: Any) -> IncLockRefResult:
-        return 0
+        return IncLockRefResult(delta=0)
 
-    def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None):
+    def dec_lock_ref(
-        return 0
+        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.

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

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

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

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_

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

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
+            return DecLockRefResult()
-        if swa_uuid_for_lock is not None:
+        return self.inner.dec_lock_ref(node, params)
-            return self.inner.dec_lock_ref(node, swa_uuid_for_lock)
-        return self.inner.dec_lock_ref(node)
 
     # -- 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)
 

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)
+        self.dec_lock_ref(
-        swa_uuid_for_lock = self.inc_lock_ref(new_last_node)
+            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)

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.inc_lock_ref.return_value = IncLockRefResult()
-        tree_cache.dec_lock_ref.return_value = None
+        tree_cache.dec_lock_ref.return_value = DecLockRefResult()
         return tree_cache
 
     def create_token_allocator(