diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4d7deb474..d423e61d7 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5b3f3b301..92d286897 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 ( diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index 50430a5eb..2ed8803b0 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 3eafb7592..0d5fecf46 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 38a4eee8e..12b3937e3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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]