[sgl] wire SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW on Unified Cache. (#28161)

python/sglang/srt/mem_cache/unified_cache_components/swa_component.py

@@ -514,6 +514,49 @@ class SWAComponent(TreeComponent):
                 dec_swa = False
             cur = cur.parent
 
+    def release_window_lock(
+        self,
+        node: UnifiedTreeNode,
+        swa_uuid_for_lock: Optional[int] = None,
+    ) -> None:
+        """Early-release the SWA lock along [node, swa_uuid_for_lock] while
+        leaving Full and Mamba locks intact.
+
+        Called when a request's decode position has advanced past the sliding
+        window — the SWA portion of the tree lock is no longer needed but the
+        Full lock must stay so the request's prefix is protected.
+
+        Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is
+        invoked at most once per (node, swa_uuid_for_lock) pair.
+        """
+        ct = self.component_type
+        root = self.cache.root_node
+
+        cur = node
+        while cur is not root:
+            cd = cur.component_data[ct]
+            # Acquire skips tombstoned nodes; release must skip them too. Same
+            # for nodes with lock_ref == 0 — acquire never credited them.
+            if cd.value is None or cd.lock_ref == 0:
+                if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock:
+                    break
+                cur = cur.parent
+                continue
+
+            cd.lock_ref -= 1
+            if cd.lock_ref == 0:
+                key_len = len(cur.key)
+                self.cache.component_protected_size_[ct] -= key_len
+                self.cache.component_evictable_size_[ct] += key_len
+                if self.cache._is_device_leaf(cur):
+                    self.cache._evict_component_and_detach_lru(
+                        cur, self, target=EvictLayer.DEVICE
+                    )
+
+            if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock:
+                break
+            cur = cur.parent
+
     def prepare_for_caching_req(
         self,
         req: Req,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -638,7 +638,10 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         return result
 
     def dec_lock_ref(
-        self, node: Any, params: Optional[DecLockRefParams] = None
+        self,
+        node: Any,
+        params: Optional[DecLockRefParams] = None,
+        skip_swa: bool = False,
     ) -> DecLockRefResult:
         result = self.session.try_dec_lock_ref(node, params)
         if result is not None:
@@ -646,12 +649,36 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         if self.disable:
             return DecLockRefResult()
         for component in self._components_tuple:
+            if skip_swa and component.component_type == ComponentType.SWA:
+                continue
             component.release_component_lock(node=node, params=params)
 
         self._update_evictable_leaf_sets(node)
         # TODO: delta is not aggregated from components; no caller uses it yet.
         return DecLockRefResult()
 
+    def dec_swa_lock_only(
+        self,
+        node: UnifiedTreeNode,
+        swa_uuid_for_lock: Optional[int] = None,
+    ) -> None:
+        """Early-release the SWA portion of a request's tree lock, plus any
+        strictly-lower-priority locks (e.g. Mamba) co-located on `node`.
+        """
+        if self.disable:
+            return
+        swa_component = self.components.get(ComponentType.SWA)
+        if swa_component is None:
+            return
+        swa_component.release_window_lock(node, swa_uuid_for_lock)
+
+        # Drop strictly-lower-priority locks (e.g. Mamba) co-located on `node`.
+        swa_priority = swa_component.eviction_priority(is_leaf=False)
+        dec_params = DecLockRefParams(swa_uuid_for_lock=swa_uuid_for_lock)
+        for comp in self._components_tuple:
+            if comp.eviction_priority(is_leaf=False) < swa_priority:
+                comp.release_component_lock(node, dec_params)
+
     def inc_host_lock_ref(self, node: Any) -> IncLockRefResult:
         if self.disable:
             return IncLockRefResult()
@@ -741,6 +768,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         self.dec_lock_ref(
             req.last_node,
             DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)),
+            skip_swa=getattr(req, "swa_prefix_lock_released", False),
         )
 
         # cleanup
@@ -1248,6 +1276,18 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             if comp.eviction_priority(is_leaf) <= trigger_priority:
                 if comp is not trigger and comp.node_has_component_data(node, target):
                     cd = node.component_data[comp.component_type]
+                    # A comp whose TRUE internal priority outranks the trigger
+                    # is only in this loop because leaf-collapse flattened
+                    # priorities; a lock on it is a legit pin and must be
+                    # spared. A lock on a strictly-lower-priority tier is a
+                    # real strand — fall through to the assert below.
+                    if comp.eviction_priority(
+                        is_leaf=False
+                    ) >= trigger.eviction_priority(is_leaf=False):
+                        if EvictLayer.DEVICE in target and cd.lock_ref != 0:
+                            continue
+                        if EvictLayer.HOST in target and cd.host_lock_ref != 0:
+                            continue
                     if EvictLayer.DEVICE in target:
                         assert cd.lock_ref == 0
                     if EvictLayer.HOST in target: