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

python/sglang/srt/configs/model_config.py

@@ -160,6 +160,7 @@ class ModelConfig:
         self.attention_chunk_size = getattr(
             self.hf_text_config, "attention_chunk_size", None
         )
+        self.sliding_window_size = self._get_sliding_window_size()
         self.is_generation = is_generation_model(
             self.hf_config.architectures, is_embedding
         )
@@ -670,6 +671,12 @@ class ModelConfig:
         else:
             return "fp8"  # Default fallback
 
+    def _get_sliding_window_size(self) -> Optional[int]:
+        sliding_window_size = getattr(self.hf_text_config, "sliding_window_size", None)
+        if sliding_window_size is None:
+            sliding_window_size = getattr(self.hf_text_config, "sliding_window", None)
+        return sliding_window_size
+
     def _validate_quantize_and_serve_config(self):
         """Validate quantize_and_serve configuration."""
         if not self.quantize_and_serve:

python/sglang/srt/managers/scheduler.py

@@ -648,10 +648,8 @@ class Scheduler(
             else:
                 from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
 
-                params.is_local_attention = (
+                params.sliding_window_size = self.model_config.sliding_window_size
-                    "Llama4ForConditionalGeneration"
+                params.attention_chunk_size = self.model_config.attention_chunk_size
-                    in self.model_config.hf_config.architectures
-                )
 
                 self.tree_cache = SWAChunkCache(params)
         else:

python/sglang/srt/mem_cache/cache_init_params.py

@@ -26,4 +26,7 @@ class CacheInitParams:
     enable_kv_cache_events: bool = False
 
     enable_mamba_extra_buffer: bool = False
-    is_local_attention: bool = False
+
+    # For SWAChunkCache
+    sliding_window_size: Optional[int] = None
+    attention_chunk_size: Optional[int] = None

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.sliding_window_size is not None:
-        if self.is_local_attention:
+            # Sliding window attention (e.g. mimo-v2-flash, gpt-oss)
-            thresh -= attention_chunk_size
+            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:
+        if new_evicted_seqlen_local > req.evicted_seqlen_local:
-            new_evicted_seqlen_local = (
-                prelen // attention_chunk_size * attention_chunk_size
-            ) - (attention_chunk_size if not self.is_local_attention else 0)
             free_slots = self.req_to_token_pool.req_to_token[
                 req.req_pool_idx, req.evicted_seqlen_local : new_evicted_seqlen_local
             ]

python/sglang/srt/mem_cache/common.py

@@ -342,9 +342,7 @@ def alloc_for_extend(
     # free out-of-window swa tokens
     if isinstance(batch.tree_cache, SWAChunkCache):
         for req, pre_len in zip(batch.reqs, batch.prefix_lens):
-            batch.tree_cache.evict_swa(
+            batch.tree_cache.evict_swa(req, pre_len)
-                req, pre_len, batch.model_config.attention_chunk_size
-            )
 
     bs = len(batch.reqs)
     prefix_tensors = [r.prefix_indices for r in batch.reqs]
@@ -437,9 +435,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
     """
     if isinstance(batch.tree_cache, SWAChunkCache):
         for req in batch.reqs:
-            batch.tree_cache.evict_swa(
+            batch.tree_cache.evict_swa(req, req.seqlen - 1)
-                req, req.seqlen - 1, batch.model_config.attention_chunk_size
-            )
 
     bs = batch.seq_lens.shape[0]
 

python/sglang/srt/speculative/eagle_info_v2.py

@@ -82,9 +82,7 @@ class EagleDraftInputV2Mixin:
     def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
         if isinstance(batch.tree_cache, SWAChunkCache):
             for req in batch.reqs:
-                batch.tree_cache.evict_swa(
+                batch.tree_cache.evict_swa(req, req.seqlen - 1)
-                    req, req.seqlen - 1, batch.model_config.attention_chunk_size
-                )
 
         from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
 

python/sglang/srt/speculative/eagle_worker.py

@@ -368,9 +368,7 @@ class EAGLEWorker(TpModelWorker):
     def _draft_preprocess_decode(self, batch: ScheduleBatch):
         if isinstance(batch.tree_cache, SWAChunkCache):
             for req in batch.reqs:
-                batch.tree_cache.evict_swa(
+                batch.tree_cache.evict_swa(req, req.seqlen - 1)
-                    req, req.seqlen - 1, batch.model_config.attention_chunk_size
-                )
 
         # Parse args
         num_seqs = batch.batch_size()

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -348,9 +348,7 @@ class MultiLayerEagleWorker(TpModelWorker):
     def _draft_preprocess_decode(self, batch: ScheduleBatch):
         if isinstance(batch.tree_cache, SWAChunkCache):
             for req in batch.reqs:
-                batch.tree_cache.evict_swa(
+                batch.tree_cache.evict_swa(req, req.seqlen - 1)
-                    req, req.seqlen - 1, batch.model_config.attention_chunk_size
-                )
 
         # Parse args
         num_seqs = batch.batch_size()