Fix evict swa for overlap scheduler and page size > 1 (#16507)

python/sglang/srt/managers/schedule_batch.py

@@ -540,9 +540,13 @@ class Req:
         # for corss-endoder model
         self.token_type_ids = token_type_ids
 
-        # The length of KV that have been removed in local attention chunked prefill
+        # The length of KV that have been removed in swa chunk cache
         self.evicted_seqlen_local = 0
 
+        # The index of the extend / decode batch
+        self.extend_batch_idx = 0
+        self.decode_batch_idx = 0
+
         # For multi-http worker
         self.http_worker_ipc = http_worker_ipc
 
@@ -1470,6 +1474,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req.req_pool_idx = req_pool_indices[i]
             assert seq_len - pre_len == req.extend_input_len
 
+            req.extend_batch_idx += 1
+
             # update req-level memory management fields
             req.kv_committed_len = seq_len
             req.kv_allocated_len = seq_len
@@ -1898,6 +1904,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
         # Update req-level memory management fields
         for req in self.reqs:
+            req.decode_batch_idx += 1
             req.kv_committed_len += 1
             req.kv_allocated_len += 1
 

python/sglang/srt/managers/scheduler.py

@@ -646,6 +646,7 @@ class Scheduler(
             enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
             pp_rank=self.pp_rank,
             pp_size=self.pp_size,
+            chunked_prefill_size=server_args.chunked_prefill_size,
         )
 
         if (

python/sglang/srt/mem_cache/cache_init_params.py

@@ -33,3 +33,5 @@ class CacheInitParams:
 
     pp_rank: int = 0
     pp_size: int = 1
+
+    chunked_prefill_size: Optional[int] = None

python/sglang/srt/mem_cache/chunk_cache.py

@@ -102,6 +102,9 @@ class SWAChunkCache(ChunkCache):
 
         self.sliding_window_size = params.sliding_window_size
         self.attention_chunk_size = params.attention_chunk_size
+        self.window_size = self.sliding_window_size or self.attention_chunk_size
+
+        self.chunked_prefill_size = params.chunked_prefill_size
 
     def evict_swa(
         self,
@@ -120,6 +123,11 @@ class SWAChunkCache(ChunkCache):
                 prelen // self.attention_chunk_size * self.attention_chunk_size,
             )
 
+        if self.page_size > 1:
+            new_evicted_seqlen_local = (
+                new_evicted_seqlen_local // self.page_size
+            ) * self.page_size
+
         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

python/sglang/srt/mem_cache/common.py

@@ -342,7 +342,16 @@ 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(req, pre_len)
+            if batch.enable_overlap:
+                # In chunked prefill case, when the second extend batch is scheduling, the first extend batch is still running, so we cannot evict swa tokens
+                if req.extend_batch_idx < 2:
+                    continue
+                else:
+                    batch.tree_cache.evict_swa(
+                        req, pre_len - batch.tree_cache.chunked_prefill_size
+                    )
+            else:
+                batch.tree_cache.evict_swa(req, pre_len)
 
     bs = len(batch.reqs)
     prefix_tensors = [r.prefix_indices for r in batch.reqs]
@@ -435,7 +444,11 @@ 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(req, req.seqlen - 1)
+            # We set evict_swa condition here with two reasons:
+            # 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running.
+            # 2. Evict swa every window_size tokens to reduce the overhead.
+            if req.decode_batch_idx % batch.tree_cache.window_size == 1:
+                batch.tree_cache.evict_swa(req, req.seqlen - 1)
 
     bs = batch.seq_lens.shape[0]