diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index f807deedb..78aaf5c65 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -189,6 +189,9 @@ class ModelConfig: and is_multimodal_chunked_prefill_supported(self.hf_config.architectures) ) self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures) + self.is_local_attention_model = is_local_attention_model( + self.hf_config.architectures + ) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) # Derive context length and model shapes @@ -1123,6 +1126,10 @@ def is_encoder_decoder_model(model_architectures: List[str]): return "MllamaForConditionalGeneration" in model_architectures +def is_local_attention_model(model_architectures: List[str]): + return "Llama4ForConditionalGeneration" in model_architectures + + def is_multimodal_chunked_prefill_supported(model_architectures: List[str]): """Check if chunked prefill is supported for a MultiModal model.""" unsupported = [ diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index bb2d6a4ad..53e3bfe9d 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -357,11 +357,12 @@ class FlashAttentionBackend(AttentionBackend): self.fa_impl_ver = fa_impl_ver # Local attention settings - self.attention_chunk_size = ( - model_runner.attention_chunk_size - if hasattr(model_runner, "attention_chunk_size") - else None - ) + self.has_local_attention = model_runner.model_config.is_local_attention_model + if self.has_local_attention: + assert ( + model_runner.attention_chunk_size is not None + ), "Attention chunk size is required for local attention" + self.attention_chunk_size = model_runner.attention_chunk_size # For each layer, the sliding_window_size can be different. This is only used for preparing SWA metadata. # We use `layer.sliding_window_size` to decide whether to use SWA for each layer. @@ -470,7 +471,7 @@ class FlashAttentionBackend(AttentionBackend): forward_batch.req_pool_indices, : metadata.max_seq_len_k ] # TODO: we need to test this part for llama 4 eagle case - self._init_local_attn_metadata(forward_batch, metadata, device) + self._maybe_init_local_attn_metadata(forward_batch, metadata, device) elif forward_batch.forward_mode.is_target_verify(): if self.topk <= 1: metadata.cache_seqlens_int32 = ( @@ -498,7 +499,7 @@ class FlashAttentionBackend(AttentionBackend): forward_batch.req_pool_indices, : metadata.max_seq_len_k ] - self._init_local_attn_metadata(forward_batch, metadata, device) + self._maybe_init_local_attn_metadata(forward_batch, metadata, device) else: metadata.cache_seqlens_int32 = forward_batch.seq_lens.to(torch.int32) metadata.max_seq_len_q = self.speculative_num_draft_tokens @@ -624,7 +625,7 @@ class FlashAttentionBackend(AttentionBackend): # Setup local attention if enabled if forward_batch.forward_mode == ForwardMode.EXTEND: - self._init_local_attn_metadata(forward_batch, metadata, device) + self._maybe_init_local_attn_metadata(forward_batch, metadata, device) # Encoder metadata for cross attention if forward_batch.encoder_lens is not None: @@ -778,7 +779,8 @@ class FlashAttentionBackend(AttentionBackend): # Check if we should use local attention use_local_attn = ( - self.attention_chunk_size is not None + self.has_local_attention + and self.attention_chunk_size is not None and metadata.local_attn_metadata is not None and (hasattr(layer, "use_irope") and layer.use_irope) ) @@ -1078,7 +1080,8 @@ class FlashAttentionBackend(AttentionBackend): metadata = self.forward_metadata local_attn_metadata = getattr(metadata, "local_attn_metadata", None) use_local_attn = ( - self.attention_chunk_size is not None + self.has_local_attention + and self.attention_chunk_size is not None and local_attn_metadata is not None and (hasattr(layer, "use_irope") and layer.use_irope) ) @@ -1350,7 +1353,7 @@ class FlashAttentionBackend(AttentionBackend): } # Only allocate local attention buffers if local attention is enabled # This prevents OOM errors when local attention is not being used - if self.attention_chunk_size is not None: + if self.has_local_attention: # Estimate maximum sizes for local attention metadata max_seq_len = self.max_context_len page_size = self.page_size or 1 @@ -1692,8 +1695,7 @@ class FlashAttentionBackend(AttentionBackend): ) self.decode_cuda_graph_metadata[bs] = metadata - if self.attention_chunk_size is not None: - self._update_local_attn_metadata_for_capture(metadata, batch_size) + self._maybe_update_local_attn_metadata_for_capture(metadata, batch_size) elif forward_mode.is_target_verify(): if self.topk <= 1: @@ -1941,7 +1943,7 @@ class FlashAttentionBackend(AttentionBackend): self.token_to_kv_pool if self.use_sliding_window_kv_pool else None, ) - self._update_local_attn_metadata_for_replay( + self._maybe_update_local_attn_metadata_for_replay( metadata, bs, ) @@ -2158,11 +2160,11 @@ class FlashAttentionBackend(AttentionBackend): """Get the fill value for sequence length in CUDA graph.""" return 1 - def _init_local_attn_metadata( + def _maybe_init_local_attn_metadata( self, forwardbatch: ForwardBatch, metadata: FlashAttentionMetadata, device ): """Centralized utility to initialize local_attn_metadata if chunked attention is enabled.""" - if self.attention_chunk_size is None: + if not self.has_local_attention: metadata.local_attn_metadata = None return @@ -2202,7 +2204,7 @@ class FlashAttentionBackend(AttentionBackend): ) metadata.local_attn_metadata = local_metadata - def _update_local_attn_metadata_for_capture( + def _maybe_update_local_attn_metadata_for_capture( self, metadata: FlashAttentionMetadata, bs: int ): """Update local attention metadata during CUDA graph capture phase. @@ -2211,6 +2213,9 @@ class FlashAttentionBackend(AttentionBackend): during the CUDA graph capture phase, optimizing memory usage by creating views of pre-allocated buffers with exactly the sizes needed. """ + if not self.has_local_attention: + return + seq_lens_capture = metadata.cache_seqlens_int32 max_seq_len = int(seq_lens_capture.max().item()) page_table_capture = metadata.page_table @@ -2258,13 +2263,13 @@ class FlashAttentionBackend(AttentionBackend): local_max_seq_len=max_seq_len, ) - def _update_local_attn_metadata_for_replay( + def _maybe_update_local_attn_metadata_for_replay( self, metadata: FlashAttentionMetadata, bs: int, ): """Update preallocated local attention metadata in-place before CUDA graph replay.""" - if self.attention_chunk_size is None: + if not self.has_local_attention: return # Access preallocated buffers