Clean up swa handling in fa3 backend (#15877)

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -11,6 +11,7 @@ import triton.language as tl
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.radix_attention import AttentionType
+from sglang.srt.mem_cache.memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput
@@ -335,12 +336,13 @@ class FlashAttentionBackend(AttentionBackend):
         self.page_size = model_runner.page_size
         self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
         self.skip_prefill = skip_prefill
-        self.is_hybrid_swa = model_runner.is_hybrid_swa
+
-        if self.is_hybrid_swa:
+        self.use_sliding_window_kv_pool = isinstance(
-            self.full_to_swa_index_mapping = (
+            model_runner.token_to_kv_pool, SWAKVPool
-                model_runner.token_to_kv_pool.full_to_swa_index_mapping
         )
+        if self.use_sliding_window_kv_pool:
             self.token_to_kv_pool = model_runner.token_to_kv_pool
+
         self.topk = model_runner.server_args.speculative_eagle_topk or 0
         self.speculative_num_steps = speculative_num_steps
         self.speculative_num_draft_tokens = (
@@ -732,10 +734,10 @@ class FlashAttentionBackend(AttentionBackend):
         # Calculate window size (can be moved to metadata if layer properties don't change)
         # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1
         # here is two side inclusive
-        is_hybrid_swa = (
+        is_swa_layer = (
             layer.sliding_window_size is not None and layer.sliding_window_size > -1
         )
-        window_size = (layer.sliding_window_size, 0) if is_hybrid_swa else (-1, -1)
+        window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1)
         k_descale, v_descale = None, None
         # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
         # has corresponding quantization method so that layer.k_scale is not None,
@@ -771,7 +773,7 @@ class FlashAttentionBackend(AttentionBackend):
         use_cascade_attn = (
             forward_batch.forward_mode.is_target_verify()
             and self.topk > 1
-            and not is_hybrid_swa
+            and not is_swa_layer
         )
 
         # For fa3 interface version compatibility, we put new fields into conditional keyword args
@@ -788,7 +790,7 @@ class FlashAttentionBackend(AttentionBackend):
             cu_seqlens_q = local_metadata.local_query_start_loc
             cache_seqlens = local_metadata.local_seqused_k
             max_seqlen_q = local_metadata.local_max_query_len
-        elif is_hybrid_swa and metadata.swa_spec_metadata is not None:
+        elif is_swa_layer and metadata.swa_spec_metadata is not None:
             swa_spec_metadata = metadata.swa_spec_metadata
             page_table = swa_spec_metadata.page_table
             cu_seqlens_q = swa_spec_metadata.cu_seqlens_q
@@ -797,15 +799,10 @@ class FlashAttentionBackend(AttentionBackend):
             cu_seqlens_k = swa_spec_metadata.cu_seqlens_k
         else:
             page_table = metadata.page_table
-            if self.is_hybrid_swa:
+            if is_swa_layer and self.use_sliding_window_kv_pool:
-                _, is_swa = forward_batch.token_to_kv_pool.layers_mapping[
-                    layer.layer_id
-                ]
-                if is_swa:
                 page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                     page_table
                 )
-                    window_size = (self.attention_chunk_size, 0)
             cu_seqlens_q = metadata.cu_seqlens_q
             cache_seqlens = metadata.cache_seqlens_int32
             max_seqlen_q = metadata.max_seq_len_q
@@ -1074,11 +1071,11 @@ class FlashAttentionBackend(AttentionBackend):
         # Calculate window size (can be moved to metadata if layer properties don't change)
         # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1
         # here is two side inclusive
-        window_size = (
+        is_swa_layer = (
-            (layer.sliding_window_size, 0)
+            layer.sliding_window_size is not None and layer.sliding_window_size > -1
-            if layer.sliding_window_size is not None and layer.sliding_window_size > -1
-            else (-1, -1)
         )
+        window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1)
+
         causal = True
         if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
             causal = False
@@ -1157,17 +1154,10 @@ class FlashAttentionBackend(AttentionBackend):
                 )
             else:
                 page_table = metadata.page_table
-                if self.is_hybrid_swa:
+                if is_swa_layer and self.use_sliding_window_kv_pool:
-                    _, is_swa = forward_batch.token_to_kv_pool.layers_mapping[
+                    page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
-                        layer.layer_id
-                    ]
-                    if is_swa:
-                        page_table = (
-                            self.token_to_kv_pool.translate_loc_from_full_to_swa(
                         page_table
                     )
-                        )
-                        window_size = (self.attention_chunk_size, 0)
                 cache_seqlens = metadata.cache_seqlens_int32
                 cu_seqlens_k = metadata.cu_seqlens_k
                 max_seqlen_q = metadata.max_seq_len_q
@@ -2139,9 +2129,9 @@ class FlashAttentionBackend(AttentionBackend):
 
         cu_seqlens_q = metadata.cu_seqlens_q
         cache_seqlens_int32 = metadata.cache_seqlens_int32
-        if self.is_hybrid_swa:
+        if self.use_sliding_window_kv_pool:
-            page_table = self.full_to_swa_index_mapping[metadata.page_table].to(
+            page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
-                torch.int32
+                metadata.page_table
             )
         else:
             page_table = metadata.page_table
@@ -2263,10 +2253,10 @@ class FlashAttentionBackend(AttentionBackend):
         # Without this slicing, the pre-allocated page_table may contain zeros or invalid indices
         # beyond the actual sequence length, leading to incorrect attention calculations
         max_seq_len = int(seqlens.max().item())
-        if self.is_hybrid_swa:
+        if self.use_sliding_window_kv_pool:
-            sliced_page_table = self.full_to_swa_index_mapping[
+            sliced_page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
                 metadata.page_table[:bs, :max_seq_len]
-            ].to(torch.int32)
+            )
         else:
             sliced_page_table = metadata.page_table[:bs, :max_seq_len]
 

python/sglang/srt/model_executor/model_runner.py

@@ -440,16 +440,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 self.model, self.remote_instance_transfer_engine
             )
 
-        # Check if the model is using hybrid SWA
-        if (
-            not self.server_args.disable_hybrid_swa_memory
-            and self.sliding_window_size is not None
-            and self.sliding_window_size > 0
-        ):
-            architectures = self.model_config.hf_config.architectures
-            if architectures and not any("Llama4" in arch for arch in architectures):
-                self.is_hybrid_swa = self.model_config.is_hybrid_swa = True
-
         # For MTP models like DeepSeek-V3 or GLM-4.5, the MTP layer(s) are used separately as draft
         # models for speculative decoding. In those cases, `num_nextn_predict_layers` is used to
         # determine the number of layers.
@@ -941,6 +931,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self.sliding_window_size = None
         if hasattr(self.model, "get_attention_sliding_window_size"):
             self.sliding_window_size = self.model.get_attention_sliding_window_size()
+        elif (
+            self.model_config.is_hybrid_swa
+            and self.model_config.sliding_window_size is not None
+        ):
+            # sliding window field in model config may have different meaning for different kinds of models (e.g., dllm), here we only consider the sliding window in SWA model
+            self.sliding_window_size = self.model_config.sliding_window_size
         elif self.model_config.attention_chunk_size is not None:
             self.sliding_window_size = self.model_config.attention_chunk_size
             logger.info(