Separate swa and local attention chunk cache eviction (#15820)

python/sglang/srt/mem_cache/chunk_cache.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 """Cache for chunked prefill, used when RadixCache is disabled."""
 
+import logging
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch
@@ -14,6 +15,9 @@ if TYPE_CHECKING:
     from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 
 
+logger = logging.getLogger(__name__)
+
+
 class ChunkCache(BasePrefixCache):
     def __init__(self, params: CacheInitParams):
         self.req_to_token_pool = params.req_to_token_pool
@@ -82,22 +86,41 @@ class SWAChunkCache(ChunkCache):
     def __init__(self, params: CacheInitParams):
         assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
         super().__init__(params)
-        self.is_local_attention = params.is_local_attention
+
+        assert (
+            params.sliding_window_size is not None
+            or params.attention_chunk_size is not None
+        ), "Sliding window size or attention chunk size must be set for SWAChunkCache"
+
+        if (
+            params.sliding_window_size is not None
+            and params.attention_chunk_size is not None
+        ):
+            logger.warning(
+                "Sliding window size and attention chunk size are both set, use sliding window size for chunk cache eviction."
+            )
+
+        self.sliding_window_size = params.sliding_window_size
+        self.attention_chunk_size = params.attention_chunk_size
 
     def evict_swa(
         self,
         req: Req,
         prelen: int,
-        attention_chunk_size: int,
     ):
-        thresh = req.evicted_seqlen_local + attention_chunk_size * 2
-        if self.is_local_attention:
-            thresh -= attention_chunk_size
+        if self.sliding_window_size is not None:
+            # Sliding window attention (e.g. mimo-v2-flash, gpt-oss)
+            new_evicted_seqlen_local = max(
+                req.evicted_seqlen_local, prelen - self.sliding_window_size
+            )
+        elif self.attention_chunk_size is not None:
+            # Local attention (e.g. llama4)
+            new_evicted_seqlen_local = max(
+                req.evicted_seqlen_local,
+                prelen // self.attention_chunk_size * self.attention_chunk_size,
+            )
 
-        if prelen >= thresh:
-            new_evicted_seqlen_local = (
-                prelen // attention_chunk_size * attention_chunk_size
-            ) - (attention_chunk_size if not self.is_local_attention else 0)
+        if new_evicted_seqlen_local > req.evicted_seqlen_local:
             free_slots = self.req_to_token_pool.req_to_token[
                 req.req_pool_idx, req.evicted_seqlen_local : new_evicted_seqlen_local
             ]