Retain SWA down to the last state checkpoint (#34729)

python/sglang/srt/managers/schedule_batch.py

@@ -3386,6 +3386,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req_to_token_pool=self.req_to_token_pool,
             token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
             is_chunk_cache=self.tree_cache.is_chunk_cache(),
+            retain_floor=self.tree_cache.swa_retain_floor(req),
         )
 
     def __str__(self):

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -377,6 +377,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
     def supports_swa(self) -> bool:
         return False
 
+    def swa_retain_floor(self, req) -> int | None:
+        # A match lands on a state checkpoint rather than on the tail, so a cache
+        # that pairs SWA with mamba/conv checkpoints has to keep the window behind
+        # the last checkpoint. Those caches override this. Everyone else has
+        # nothing deeper than the tail to protect.
+        return None
+
     def swa_reprefill_tail_tokens(self) -> int:
         # Only the unified_kv compress-only HiCache layout needs to hold back a
         # trailing sliding window for re-prefill; every other cache keeps SWA

python/sglang/srt/mem_cache/common.py

@@ -53,6 +53,7 @@ def free_swa_out_of_window_slots(
     req_to_token_pool: ReqToTokenPool,
     token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
     is_chunk_cache: bool = False,
+    retain_floor: int | None = None,
 ) -> None:
     if req.kv is None:
         return
@@ -76,6 +77,12 @@ def free_swa_out_of_window_slots(
         # boundary (page_floor(seq_len)) so the last leaf is never all-tombstone.
         # No extra page margin is needed.
         evict_threshold = pre_len - max(sliding_window_size, page_size)
+    if retain_floor is not None and not is_chunk_cache:
+        # The caller owns where the floor is (see BasePrefixCache.swa_retain_floor);
+        # this only promises not to free past it. Chunk cache has no tree, so a
+        # retained checkpoint could never be matched and holding it is pure cost.
+        evict_threshold = min(evict_threshold, retain_floor)
+
     new_swa_evicted_seqlen = max(
         req.kv.swa_evicted_seqlen,
         evict_threshold,

python/sglang/srt/mem_cache/unified_cache/components/swa_component.py

@@ -755,6 +755,7 @@ class SWAComponent(TreeComponent):
                 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,
+                retain_floor=self.cache.swa_retain_floor(req),
             )
         insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
 

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -2098,6 +2098,14 @@ class UnifiedRadixCache(BasePrefixCache):
         )
         return swa.sliding_window_size if unified_compress_only_hicache else 0
 
+    def swa_retain_floor(self, req) -> int | None:
+        if not self.is_mamba_enabled or self._sliding_window_size is None:
+            return None
+        checkpoint = req.mamba_last_track_seqlen
+        if checkpoint is None:
+            return None
+        return checkpoint - self._sliding_window_size
+
     def supports_swa(self) -> bool:
         return self.is_swa_enabled