diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 8cc95abc5..f807deedb 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bf8db9bcf..120d5696e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 = ( - "Llama4ForConditionalGeneration" - in self.model_config.hf_config.architectures - ) + params.sliding_window_size = self.model_config.sliding_window_size + params.attention_chunk_size = self.model_config.attention_chunk_size self.tree_cache = SWAChunkCache(params) else: diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index a833bec98..d410dda0e 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 756dd3024..09d2c951c 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/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 ] diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index f8bbac4fe..b06e3a401 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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( - req, pre_len, batch.model_config.attention_chunk_size - ) + batch.tree_cache.evict_swa(req, pre_len) 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( - req, req.seqlen - 1, batch.model_config.attention_chunk_size - ) + batch.tree_cache.evict_swa(req, req.seqlen - 1) bs = batch.seq_lens.shape[0] diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index caacbfff5..a14cddb05 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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( - req, req.seqlen - 1, batch.model_config.attention_chunk_size - ) + batch.tree_cache.evict_swa(req, req.seqlen - 1) from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 7514d7338..c508b552d 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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( - req, req.seqlen - 1, batch.model_config.attention_chunk_size - ) + batch.tree_cache.evict_swa(req, req.seqlen - 1) # Parse args num_seqs = batch.batch_size() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 9dc3c5fd7..6aa79e7a8 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/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( - req, req.seqlen - 1, batch.model_config.attention_chunk_size - ) + batch.tree_cache.evict_swa(req, req.seqlen - 1) # Parse args num_seqs = batch.batch_size()