From 3c1f9eafa544b6522dc8e302516a68c35b7b8b83 Mon Sep 17 00:00:00 2001 From: Bi Xue Date: Thu, 11 Jun 2026 23:17:18 -0700 Subject: [PATCH] [sgl] proactively release out-of-window SWA slots after chunked prefill (#27402) Co-authored-by: ispobock --- python/sglang/srt/environ.py | 3 + python/sglang/srt/managers/schedule_batch.py | 41 ++------- python/sglang/srt/mem_cache/common.py | 44 +++++++++ .../unified_cache_components/swa_component.py | 15 +++ .../tree_component.py | 5 + .../srt/mem_cache/unified_radix_cache.py | 7 ++ .../test_unified_radix_cache_unittest.py | 92 +++++++++++++++++++ 7 files changed, 174 insertions(+), 33 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a4c8dff0b..86cc29506 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 38807cd8c..23b41ce06 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 - - # 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 % self.tree_cache.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 - self.tree_cache.page_size - new_swa_evicted_seqlen = max( - req.swa_evicted_seqlen, - evict_threshold, + free_swa_out_of_window_slots( + req, + pre_len, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.tree_cache.page_size, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, ) - 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'}, " diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index bf9e0561a..358d14e36 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 99c8451be..7856cdf20 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py index 6bf5f8ef3..f83e37f2f 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index baa239257..e96ff48f2 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index da06f6294..e2553e67d 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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")