minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Skip local attn init metadata for mimo swa model (#16349)

Browse Source
This commit is contained in:
Ke Bao
2026-01-04 22:38:36 +08:00
committed by GitHub
parent bf32cd8397
commit b328cd20bb
2 changed files with 31 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+7
View File
@@ -189,6 +189,9 @@ class ModelConfig:
and is_multimodal_chunked_prefill_supported(self.hf_config.architectures) and is_multimodal_chunked_prefill_supported(self.hf_config.architectures)
) )
self.is_encoder_decoder = is_encoder_decoder_model(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) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
# Derive context length and model shapes # Derive context length and model shapes
@@ -1123,6 +1126,10 @@ def is_encoder_decoder_model(model_architectures: List[str]):
return "MllamaForConditionalGeneration" in model_architectures 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]): def is_multimodal_chunked_prefill_supported(model_architectures: List[str]):
"""Check if chunked prefill is supported for a MultiModal model.""" """Check if chunked prefill is supported for a MultiModal model."""
unsupported = [ unsupported = [
python/sglang/srt/layers/attention/flashattention_backend.py
+24 -19
View File
@@ -357,11 +357,12 @@ class FlashAttentionBackend(AttentionBackend):
self.fa_impl_ver = fa_impl_ver self.fa_impl_ver = fa_impl_ver
# Local attention settings # Local attention settings
self.attention_chunk_size = ( self.has_local_attention = model_runner.model_config.is_local_attention_model
model_runner.attention_chunk_size if self.has_local_attention:
if hasattr(model_runner, "attention_chunk_size") assert (
else None 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. # 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. # 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 forward_batch.req_pool_indices, : metadata.max_seq_len_k
] ]
# TODO: we need to test this part for llama 4 eagle case # 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(): elif forward_batch.forward_mode.is_target_verify():
if self.topk <= 1: if self.topk <= 1:
metadata.cache_seqlens_int32 = ( metadata.cache_seqlens_int32 = (
@@ -498,7 +499,7 @@ class FlashAttentionBackend(AttentionBackend):
forward_batch.req_pool_indices, : metadata.max_seq_len_k 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: else:
metadata.cache_seqlens_int32 = forward_batch.seq_lens.to(torch.int32) metadata.cache_seqlens_int32 = forward_batch.seq_lens.to(torch.int32)
metadata.max_seq_len_q = self.speculative_num_draft_tokens metadata.max_seq_len_q = self.speculative_num_draft_tokens
@@ -624,7 +625,7 @@ class FlashAttentionBackend(AttentionBackend):
# Setup local attention if enabled # Setup local attention if enabled
if forward_batch.forward_mode == ForwardMode.EXTEND: 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 # Encoder metadata for cross attention
if forward_batch.encoder_lens is not None: if forward_batch.encoder_lens is not None:
@@ -778,7 +779,8 @@ class FlashAttentionBackend(AttentionBackend):
# Check if we should use local attention # Check if we should use local attention
use_local_attn = ( 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 metadata.local_attn_metadata is not None
and (hasattr(layer, "use_irope") and layer.use_irope) and (hasattr(layer, "use_irope") and layer.use_irope)
) )
@@ -1078,7 +1080,8 @@ class FlashAttentionBackend(AttentionBackend):
metadata = self.forward_metadata metadata = self.forward_metadata
local_attn_metadata = getattr(metadata, "local_attn_metadata", None) local_attn_metadata = getattr(metadata, "local_attn_metadata", None)
use_local_attn = ( 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 local_attn_metadata is not None
and (hasattr(layer, "use_irope") and layer.use_irope) 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 # Only allocate local attention buffers if local attention is enabled
# This prevents OOM errors when local attention is not being used # 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 # Estimate maximum sizes for local attention metadata
max_seq_len = self.max_context_len max_seq_len = self.max_context_len
page_size = self.page_size or 1 page_size = self.page_size or 1
@@ -1692,8 +1695,7 @@ class FlashAttentionBackend(AttentionBackend):
) )
self.decode_cuda_graph_metadata[bs] = metadata self.decode_cuda_graph_metadata[bs] = metadata
if self.attention_chunk_size is not None: self._maybe_update_local_attn_metadata_for_capture(metadata, batch_size)
self._update_local_attn_metadata_for_capture(metadata, batch_size)
elif forward_mode.is_target_verify(): elif forward_mode.is_target_verify():
if self.topk <= 1: 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.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, metadata,
bs, bs,
) )
@@ -2158,11 +2160,11 @@ class FlashAttentionBackend(AttentionBackend):
"""Get the fill value for sequence length in CUDA graph.""" """Get the fill value for sequence length in CUDA graph."""
return 1 return 1
def _init_local_attn_metadata( def _maybe_init_local_attn_metadata(
self, forwardbatch: ForwardBatch, metadata: FlashAttentionMetadata, device self, forwardbatch: ForwardBatch, metadata: FlashAttentionMetadata, device
): ):
"""Centralized utility to initialize local_attn_metadata if chunked attention is enabled.""" """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 metadata.local_attn_metadata = None
return return
@@ -2202,7 +2204,7 @@ class FlashAttentionBackend(AttentionBackend):
) )
metadata.local_attn_metadata = local_metadata 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 self, metadata: FlashAttentionMetadata, bs: int
): ):
"""Update local attention metadata during CUDA graph capture phase. """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 during the CUDA graph capture phase, optimizing memory usage by creating views of
pre-allocated buffers with exactly the sizes needed. pre-allocated buffers with exactly the sizes needed.
""" """
if not self.has_local_attention:
return
seq_lens_capture = metadata.cache_seqlens_int32 seq_lens_capture = metadata.cache_seqlens_int32
max_seq_len = int(seq_lens_capture.max().item()) max_seq_len = int(seq_lens_capture.max().item())
page_table_capture = metadata.page_table page_table_capture = metadata.page_table
@@ -2258,13 +2263,13 @@ class FlashAttentionBackend(AttentionBackend):
local_max_seq_len=max_seq_len, local_max_seq_len=max_seq_len,
) )
def _update_local_attn_metadata_for_replay( def _maybe_update_local_attn_metadata_for_replay(
self, self,
metadata: FlashAttentionMetadata, metadata: FlashAttentionMetadata,
bs: int, bs: int,
): ):
"""Update preallocated local attention metadata in-place before CUDA graph replay.""" """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 return
# Access preallocated buffers # Access preallocated buffers