[sgl] proactively release out-of-window SWA slots after chunked prefill (#27402)

Co-authored-by: ispobock <ispobaoke@gmail.com>

python/sglang/srt/environ.py

@@ -782,6 +782,9 @@ class Envs:
     SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW = EnvBool(False)
     SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN = EnvBool(False)
 
+    # Unified radix cache
+    SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS = EnvBool(False)
+
     # DeepGemm Mega MoE
     SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE = EnvBool(False)
     SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK = EnvInt(1024)

python/sglang/srt/managers/schedule_batch.py

@@ -83,6 +83,7 @@ from sglang.srt.mem_cache.common import (
     alloc_for_decode,
     alloc_for_extend,
     evict_from_tree_cache,
+    free_swa_out_of_window_slots,
     get_alloc_reserve_per_decode,
     release_kv_cache,
 )
@@ -2835,41 +2836,15 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     def _evict_swa(self, req: Req, pre_len: int):
         assert self.tree_cache.supports_swa(), "prefix cache must support swa"
-        sliding_window_size = self.tree_cache.sliding_window_size
+        free_swa_out_of_window_slots(
-
+            req,
-        # For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
+            pre_len,
-        assert (
+            sliding_window_size=self.tree_cache.sliding_window_size,
-            req.cache_protected_len % self.tree_cache.page_size == 0
+            page_size=self.tree_cache.page_size,
-        ), "cache_protected_len must be page aligned"
+            req_to_token_pool=self.req_to_token_pool,
-        req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, req.cache_protected_len)
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
-
-        # Subtract an extra page_size so the eviction frontier never reaches the
-        # radix tree insert boundary (page_floor(seq_len)). This keeps at least one
-        # page of non-evicted SWA KV for the tree to store as a non-tombstone node,
-        # preserving cache reuse in multi-turn scenarios. Without this, leaf nodes
-        # may become tombstoned, causing SWA memory leak.
-        # See also: _insert_helper case 3 in swa_radix_cache.py (defensive counterpart).
-        if envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get():
-            evict_threshold = pre_len - sliding_window_size
-        else:
-            evict_threshold = pre_len - sliding_window_size - self.tree_cache.page_size
-        new_swa_evicted_seqlen = max(
-            req.swa_evicted_seqlen,
-            evict_threshold,
         )
 
-        if self.tree_cache.page_size > 1:
-            new_swa_evicted_seqlen = (
-                new_swa_evicted_seqlen // self.tree_cache.page_size
-            ) * self.tree_cache.page_size
-
-        if new_swa_evicted_seqlen > req.swa_evicted_seqlen:
-            free_slots = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, req.swa_evicted_seqlen : new_swa_evicted_seqlen
-            ]
-            self.token_to_kv_pool_allocator.free_swa(free_slots)
-            req.swa_evicted_seqlen = new_swa_evicted_seqlen
-
     def __str__(self):
         return (
             f"ScheduleBatch(forward_mode={self.forward_mode.name if self.forward_mode else 'None'}, "

python/sglang/srt/mem_cache/common.py

@@ -28,6 +28,7 @@ _is_hip = is_hip()
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 
 # Needs 2 + 1 slots for mamba request with prefix cache. 2 for ping pong cache, 1 for running mamba state.
 MAMBA_STATE_PER_REQ_PREFIX_CACHE = 3
@@ -54,6 +55,49 @@ def page_align_floor(length: int, page_size: int) -> int:
     return (length // page_size) * page_size
 
 
+def free_swa_out_of_window_slots(
+    req: Req,
+    pre_len: int,
+    *,
+    sliding_window_size: int,
+    page_size: int,
+    req_to_token_pool: ReqToTokenPool,
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+) -> None:
+    from sglang.srt.environ import envs
+
+    # For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
+    assert (
+        req.cache_protected_len % page_size == 0
+    ), "cache_protected_len must be page aligned"
+    req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, req.cache_protected_len)
+
+    # Subtract an extra page_size so the eviction frontier never reaches the
+    # radix tree insert boundary (page_floor(seq_len)). This keeps at least one
+    # page of non-evicted SWA KV for the tree to store as a non-tombstone node,
+    # preserving cache reuse in multi-turn scenarios. Without this, leaf nodes
+    # may become tombstoned, causing SWA memory leak.
+    # See also: _insert_helper case 3 in swa_radix_cache.py (defensive counterpart).
+    if envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get():
+        evict_threshold = pre_len - sliding_window_size
+    else:
+        evict_threshold = pre_len - sliding_window_size - page_size
+    new_swa_evicted_seqlen = max(
+        req.swa_evicted_seqlen,
+        evict_threshold,
+    )
+
+    if page_size > 1:
+        new_swa_evicted_seqlen = (new_swa_evicted_seqlen // page_size) * page_size
+
+    if new_swa_evicted_seqlen > req.swa_evicted_seqlen:
+        free_slots = req_to_token_pool.req_to_token[
+            req.req_pool_idx, req.swa_evicted_seqlen : new_swa_evicted_seqlen
+        ]
+        token_to_kv_pool_allocator.free_swa(free_slots)
+        req.swa_evicted_seqlen = new_swa_evicted_seqlen
+
+
 def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
     if getattr(req, "skip_radix_cache_insert", False):
         return

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

@@ -13,6 +13,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
     MatchResult,
 )
+from sglang.srt.mem_cache.common import free_swa_out_of_window_slots
 from sglang.srt.mem_cache.hicache_storage import (
     PoolHitPolicy,
     PoolName,
@@ -524,6 +525,20 @@ class SWAComponent(TreeComponent):
             insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
         return None
 
+    def free_out_of_window_slots(
+        self, req: Req, pre_len: int, insert_params: InsertParams
+    ) -> None:
+        if self.sliding_window_size is not None:
+            free_swa_out_of_window_slots(
+                req,
+                pre_len,
+                sliding_window_size=self.sliding_window_size,
+                page_size=self.cache.page_size,
+                req_to_token_pool=self.cache.req_to_token_pool,
+                token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
+            )
+        insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen
+
     # ---- HiCache Hooks ----
 
     def build_hicache_transfers(

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

@@ -380,6 +380,11 @@ class TreeComponent(ABC):
         paths it is still provided so components can free their resources."""
         pass
 
+    def free_out_of_window_slots(
+        self, req: Req, pre_len: int, insert_params: InsertParams
+    ) -> None:
+        pass
+
     # ---- HiCache Hooks ----
 
     def build_hicache_transfers(

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Iterator, Optional, TypeVar
 import torch
 
 from sglang.srt.disaggregation.kv_events import StorageMedium
+from sglang.srt.environ import envs
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
     DecLockRefParams,
@@ -756,6 +757,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             if cl is not None:
                 effective_cache_len = min(effective_cache_len, cl)
 
+        if envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get():
+            for comp in self._components_tuple:
+                comp.free_out_of_window_slots(
+                    req, effective_cache_len - 1, insert_params
+                )
+
         if effective_cache_len <= 0:
             req.prefix_indices = kv_indices_orig.to(dtype=torch.int64, copy=True)
             for comp in self._components_tuple:

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -1522,6 +1522,98 @@ class UnifiedRadixCacheSuite:
         self.assertLess(b_pos, a_pos, "A was below B in pre, must stay below")
         tree.sanity_check()
 
+    def test_swa_eager_eviction_on_unfinished_req(self):
+        if not self.cfg.has_swa or self.cfg.has_mamba:
+            self.skipTest(
+                "requires SWA without Mamba (mamba alters effective_cache_len)"
+            )
+        if self.cfg.page_size != 1 or self.cfg.sliding_window_size != 4:
+            self.skipTest("requires page_size=1, sliding_window_size=4")
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+
+        pre_len = 20
+        req = self._make_req(req_to_token_pool)
+        tokens = self._make_seq(1, pre_len)
+        req.origin_input_ids = tokens
+        req.output_ids = []
+        req.full_untruncated_fill_ids = array("q", tokens)
+        req.fill_len = len(req.full_untruncated_fill_ids)
+        kv_indices = self._alloc(allocator, pre_len)
+        req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
+        req.kv_committed_len = pre_len
+        req.last_node = tree.root_node
+        req.cache_protected_len = 0
+        req.swa_uuid_for_lock = None
+        req.extra_key = None
+        req.swa_evicted_seqlen = 0
+
+        swa_avail_before = allocator.swa_attn_allocator.available_size()
+
+        with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
+            tree.cache_unfinished_req(req)
+
+        cushion = self.cfg.sliding_window_size + self.cfg.page_size
+        expected_evicted = (pre_len - 1) - cushion
+        self.assertEqual(
+            req.swa_evicted_seqlen,
+            expected_evicted,
+            f"swa_evicted_seqlen should advance to (pre_len-1) - cushion = "
+            f"{expected_evicted}, got {req.swa_evicted_seqlen}",
+        )
+
+        swa_avail_after = allocator.swa_attn_allocator.available_size()
+        self.assertGreaterEqual(
+            swa_avail_after - swa_avail_before,
+            expected_evicted,
+            f"SWA pool should have freed at least {expected_evicted} slots; "
+            f"before={swa_avail_before}, after={swa_avail_after}",
+        )
+
+        tree.dec_lock_ref(
+            req.last_node,
+            DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)),
+        )
+        tree.sanity_check()
+
+    def test_swa_eager_eviction_noop_when_within_window(self):
+        if not self.cfg.has_swa or self.cfg.has_mamba:
+            self.skipTest("requires SWA without Mamba")
+        if self.cfg.page_size != 1 or self.cfg.sliding_window_size != 4:
+            self.skipTest("requires page_size=1, sliding_window_size=4")
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+
+        cushion = self.cfg.sliding_window_size + self.cfg.page_size  # = 5
+        pre_len = cushion  # exactly at the boundary, nothing slid out
+        req = self._make_req(req_to_token_pool)
+        tokens = self._make_seq(1, pre_len)
+        req.origin_input_ids = tokens
+        req.output_ids = []
+        req.full_untruncated_fill_ids = array("q", tokens)
+        req.fill_len = len(req.full_untruncated_fill_ids)
+        kv_indices = self._alloc(allocator, pre_len)
+        req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
+        req.kv_committed_len = pre_len
+        req.last_node = tree.root_node
+        req.cache_protected_len = 0
+        req.swa_uuid_for_lock = None
+        req.extra_key = None
+        req.swa_evicted_seqlen = 0
+
+        with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
+            tree.cache_unfinished_req(req)
+
+        self.assertEqual(
+            req.swa_evicted_seqlen,
+            0,
+            "Nothing should be evicted when prefill fits inside the cushion",
+        )
+
+        tree.dec_lock_ref(
+            req.last_node,
+            DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)),
+        )
+        tree.sanity_check()
+
     def test_swa_sanity_check_passes_after_deep_match(self):
         if not self._swa_pinning_cfg_supported():
             self.skipTest("requires SWA-only config with node size >= cushion")