fix(swa): port tombstone-aware eviction + dec_swa_lock_only from upsream #23882

python/sglang/srt/environ.py

@@ -261,6 +261,7 @@ class Envs:
     SGLANG_REQ_WAITING_TIMEOUT = EnvFloat(-1)  # in seconds
     SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False)
     SGLANG_REQ_RUNNING_TIMEOUT = EnvFloat(-1)  # in seconds
+    SGLANG_SWA_EVICTION_INTERVAL_MULTIPLIER = EnvFloat(1.0)
     SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120)
 
     # Test: pd-disaggregation
@@ -519,6 +520,7 @@ class Envs:
     SGLANG_OPT_ALLOW_SHARED_EXPERT_DUAL_STREAM = EnvBool(False)  # verified in journal 2026-04-21-017
     SGLANG_OPT_CACHE_SWA_TRANSLATION = EnvBool(False)
     SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False)
+    SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW = EnvBool(False)
     SGLANG_OPT_MXFP4_FUSE_RSF_SHARED_ADD = EnvBool(False)
     SGLANG_OPT_MXFP4_STATIC_SCALE_ONES = EnvBool(False)
     SGLANG_OPT_MXFP4_SKIP_DISPATCHER_MAPPING = EnvBool(False)

python/sglang/srt/managers/schedule_batch.py

@@ -643,6 +643,7 @@ def __init__(
         self.storage_hit_length = 0
         # The node to lock until for swa radix tree lock ref
         self.swa_uuid_for_lock: Optional[int] = None
+        self.swa_prefix_lock_released: bool = False
         # The prefix length that is inserted into the tree cache
         self.cache_protected_len: int = 0
 
@@ -1091,6 +1092,7 @@ def reset_for_retract(self):
         self.indexer_topk = None
         self.last_node = None
         self.swa_uuid_for_lock = None
+        self.swa_prefix_lock_released = False
         self.extend_input_len = 0
         self.is_retracted = True
         self.retracted_stain = True
@@ -2269,24 +2271,39 @@ def maybe_evict_swa(self):
             sliding_window_size = self.tree_cache.sliding_window_size
             server_args = get_global_server_args()
 
-            if (
-                self.forward_mode.is_decode()
-                and server_args.enable_piecewise_cuda_graph
-                and not self.tree_cache.is_chunk_cache()
-            ):
-                return
+            release_leaf_lock = (
+                envs.SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW.get()
+                and hasattr(self.tree_cache, "dec_swa_lock_only")
+            )
 
+            page_size = self.tree_cache.page_size
+            eviction_interval = max(
+                page_size,
+                int(
+                    sliding_window_size
+                    * envs.SGLANG_SWA_EVICTION_INTERVAL_MULTIPLIER.get()
+                ),
+            )
+            eviction_interval = (eviction_interval // page_size) * page_size
             for idx, req in enumerate(self.reqs):
                 if self.forward_mode.is_decode():
-                    # We set evict_swa condition here with two reasons:
-                    # 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running.
-                    # 2. Evict swa every window_size tokens to reduce the overhead.
-                    if req.decode_batch_idx % sliding_window_size == 1:
+                    if req.decode_batch_idx % eviction_interval == 1:
                         self._evict_swa(req, req.seqlen - 1)
+
+                    if (
+                        release_leaf_lock
+                        and not req.swa_prefix_lock_released
+                        and req.swa_uuid_for_lock is not None
+                        and req.last_node is not None
+                        and req.decode_batch_idx >= sliding_window_size
+                    ):
+                        self.tree_cache.dec_swa_lock_only(
+                            req.last_node, req.swa_uuid_for_lock
+                        )
+                        req.swa_prefix_lock_released = True
                 elif self.forward_mode.is_extend() and self.tree_cache.is_chunk_cache():
                     pre_len = self.prefix_lens[idx]
                     if self.enable_overlap:
-                        # In chunked prefill case, when the second extend batch is scheduling, the first extend batch is still running, so we cannot evict swa tokens
                         if req.extend_batch_idx < 2:
                             continue
                         else:
@@ -2309,6 +2326,7 @@ def _evict_swa(self, req: Req, pre_len: int):
         ), "cache_protected_len must be page aligned"
         req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, req.cache_protected_len)
 
+        # -page_size: keep one extra page inside window to avoid evicting a partially-filled page
         new_swa_evicted_seqlen = max(
             req.swa_evicted_seqlen,
             pre_len - sliding_window_size - self.tree_cache.page_size,

python/sglang/srt/managers/schedule_policy.py

@@ -37,6 +37,7 @@
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
+    DecLockRefParams,
     InsertParams,
     MatchPrefixParams,
 )
@@ -573,11 +574,9 @@ def _add_dllm_req(self, req: Req, prefix_len: int):
         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
@@ -642,14 +641,14 @@ def add_chunked_req(self, req: Req):
     @contextmanager
     def _lock_node(self, last_node: TreeNode):
         try:
-            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)
+            result = 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)
 

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -70,11 +70,27 @@ class InsertResult:
     mamba_exist: bool = False
 
 
+@dataclasses.dataclass
+class IncLockRefResult:
+    delta: Optional[int] = None
+    swa_uuid_for_lock: Optional[int] = None
+
+
+@dataclasses.dataclass
+class DecLockRefParams:
+    swa_uuid_for_lock: Optional[int] = None
+
+
+@dataclasses.dataclass
+class DecLockRefResult:
+    delta: Optional[int] = None
+
+
 @dataclasses.dataclass
 class EvictParams:
     """Unified parameters for evict across different cache types"""
 
-    num_tokens: int
+    num_tokens: int = 0
     swa_num_tokens: int = 0
     mamba_num: int = 0
 
@@ -155,11 +171,13 @@ def evict(self, params: EvictParams) -> EvictResult:
         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, **kwargs
+    ) -> "DecLockRefResult":
         pass
 
     def evictable_size(self):

python/sglang/srt/mem_cache/chunk_cache.py

@@ -9,8 +9,11 @@
 
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
+    DecLockRefParams,
+    DecLockRefResult,
     EvictParams,
     EvictResult,
+    IncLockRefResult,
     InsertParams,
     InsertResult,
     MatchPrefixParams,
@@ -79,11 +82,13 @@ def cache_unfinished_req(self, req: Req, chunked=False):
     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, **kwargs
+    ) -> 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/hiradix_cache.py

@@ -14,8 +14,11 @@
 
 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,
@@ -720,9 +723,9 @@ def loading_check(self):
     def evictable_size(self):
         return self.evictable_size_
 
-    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:
@@ -734,11 +737,13 @@ def inc_lock_ref(self, node: TreeNode):
             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, **kwargs
+    ) -> DecLockRefResult:
         if self.disable:
-            return 0
+            return DecLockRefResult(delta=0)
 
         delta = 0
         while node != self.root_node:
@@ -754,7 +759,7 @@ def dec_lock_ref(self, node: TreeNode):
                     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:
@@ -887,7 +892,8 @@ def load_back(
             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.base_prefix_cache import (
     BasePrefixCache,
+    DecLockRefParams,
+    DecLockRefResult,
     EvictParams,
     EvictResult,
+    IncLockRefResult,
     InsertParams,
     InsertResult,
     MatchPrefixParams,
@@ -785,14 +788,14 @@ def evict_full(self, full_num_tokens: int) -> int:
 
         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:
@@ -811,16 +814,18 @@ def inc_lock_ref(self, node: TreeNode) -> Optional[int]:
                 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, **kwargs
+    ) -> 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 (
@@ -840,6 +845,7 @@ def dec_lock_ref(self, node: TreeNode):
                 self.full_protected_size_ -= len(node.value)
             node.full_lock_ref -= 1
             node = node.parent
+        return DecLockRefResult()
 
     def sanity_check(self):
         if self.disable:

python/sglang/srt/mem_cache/radix_cache.py

@@ -42,8 +42,11 @@
 )
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
+    DecLockRefParams,
+    DecLockRefResult,
     EvictParams,
     EvictResult,
+    IncLockRefResult,
     InsertParams,
     InsertResult,
     MatchPrefixParams,
@@ -591,9 +594,9 @@ def evict(self, params: EvictParams) -> EvictResult:
         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:
@@ -604,11 +607,13 @@ def inc_lock_ref(self, node: TreeNode):
             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, **kwargs
+    ) -> DecLockRefResult:
         if self.disable:
-            return 0
+            return DecLockRefResult(delta=0)
 
         delta = 0
         while node != self.root_node:
@@ -623,7 +628,7 @@ def dec_lock_ref(self, node: TreeNode):
                     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_