add entrance and solve conflict

Author: Fridge003

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

@@ -59,6 +59,7 @@ def __init__(
         self.device = model_runner.device
         self.decode_cuda_graph_metadata = {}
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.page_size = model_runner.page_size
         self.use_mla = (
             model_runner.model_config.attention_arch == AttentionArch.MLA
         ) and (not global_server_args_dict["disable_mla"])
@@ -83,6 +84,17 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
         metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
             forward_batch.req_pool_indices, : metadata.max_seq_len_k
         ]
+
+        # Precompute strided indices
+        # [0, page_size, 2 * page_size, ...]
+        if self.page_size > 1:
+            self.strided_indices = torch.arange(
+                0, metadata.page_table.shape[1], self.page_size, device=self.device
+            )
+            metadata.page_table = (
+                metadata.page_table[:, self.strided_indices] // self.page_size
+            )
+
         if forward_batch.forward_mode == ForwardMode.DECODE:
             # Precompute cumulative sequence lengths
             metadata.cu_seqlens_q = torch.arange(
@@ -141,16 +153,24 @@ def forward_extend(
             else (-1, -1)
         )
 
+        page_table = metadata.page_table
+
         # # Use Flash Attention for prefill
         if not self.use_mla:
             # Do multi-head attention
             kv_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
             key_cache, value_cache = kv_cache[0], kv_cache[1]
+            key_cache = key_cache.view(
+                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+            )
+            value_cache = value_cache.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+            )
             o = flash_attn_with_kvcache(
                 q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                k_cache=key_cache.unsqueeze(1),
-                v_cache=value_cache.unsqueeze(1),
-                page_table=metadata.page_table,
+                k_cache=key_cache,
+                v_cache=value_cache,
+                page_table=page_table,
                 cache_seqlens=metadata.cache_seqlens_int32,
                 cu_seqlens_q=metadata.cu_seqlens_q,
                 cu_seqlens_k_new=metadata.cu_seqlens_k,
@@ -176,7 +196,7 @@ def forward_extend(
                 k_cache=k_rope.unsqueeze(1),
                 v_cache=c_kv.unsqueeze(1),
                 qv=q_nope,
-                page_table=metadata.page_table,
+                page_table=page_table,
                 cache_seqlens=metadata.cache_seqlens_int32,
                 cu_seqlens_q=metadata.cu_seqlens_q,
                 cu_seqlens_k_new=metadata.cu_seqlens_k,
@@ -231,22 +251,30 @@ def forward_decode(
             else (-1, -1)
         )
 
+        page_table = metadata.page_table
+
         if not self.use_mla:
             # Do multi-head attention
 
             # Get KV cache
             kv_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
             key_cache, value_cache = kv_cache[0], kv_cache[1]
+            key_cache = key_cache.view(
+                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+            )
+            value_cache = value_cache.view(
+                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+            )
 
             # Pre-reshape query tensor
             q_reshaped = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
 
             # Run attention with precomputed values
             o = flash_attn_with_kvcache(
                 q=q_reshaped,
-                k_cache=key_cache.unsqueeze(1),
-                v_cache=value_cache.unsqueeze(1),
-                page_table=metadata.page_table,
+                k_cache=key_cache,
+                v_cache=value_cache,
+                page_table=page_table,
                 cache_seqlens=metadata.cache_seqlens_int32,
                 cu_seqlens_q=metadata.cu_seqlens_q,
                 cu_seqlens_k_new=metadata.cu_seqlens_k,
@@ -273,7 +301,7 @@ def forward_decode(
                 k_cache=k_rope.unsqueeze(1),
                 v_cache=c_kv.unsqueeze(1),
                 qv=q_nope,
-                page_table=metadata.page_table,
+                page_table=page_table,
                 cache_seqlens=metadata.cache_seqlens_int32,
                 cu_seqlens_q=metadata.cu_seqlens_q,
                 cu_seqlens_k_new=metadata.cu_seqlens_k,
@@ -300,7 +328,13 @@ def init_cuda_graph_state(self, max_bs: int):
         self.decode_cuda_graph_metadata = {
             # Page table for token mapping (batch_size, max_context_len)
             "page_table": torch.zeros(
-                max_bs, self.max_context_len, dtype=torch.int32, device=self.device
+                max_bs,
+                (self.max_context_len + self.page_size - 1) // self.page_size,
+                dtype=torch.int32,
+                device=self.device,
+            ),
+            "strided_indices": torch.arange(
+                0, self.max_context_len, self.page_size, device=self.device
             ),
         }
 

python/sglang/srt/model_executor/model_runner.py

@@ -230,6 +230,10 @@ def model_specific_adjustment(self):
                 elif server_args.enable_flashmla:
                     logger.info("MLA optimization is turned on. Use flashmla decode.")
                     server_args.attention_backend = "flashmla"
+                elif server_args.attention_backend == "fa3":
+                    logger.info(
+                        f"MLA optimization is turned on. Use flash attention 3 backend."
+                    )
                 else:
                     logger.info("MLA optimization is turned on. Use triton backend.")
                     server_args.attention_backend = "triton"
@@ -879,7 +883,7 @@ def init_attention_backend(self):
                 "Please use `--attention-backend flashinfer`."
             )
             logger.warning(
-                "FlashAttention v3 Backend is in Beta. Multimodal, Page > 1, FP8, MLA and Speculative Decoding are not supported."
+                "FlashAttention v3 Backend is in Beta. Multimodal, FP8, and Speculative Decoding are not supported."
             )
             from sglang.srt.layers.attention.flashattention_backend import (
                 FlashAttentionBackend,

python/sglang/srt/models/deepseek_v2.py

@@ -655,6 +655,7 @@ def __init__(
         self.flashinfer_mla_disable_ragged = global_server_args_dict[
             "flashinfer_mla_disable_ragged"
         ]
+        self.attention_backend = global_server_args_dict["attention_backend"]
         self.rocm_fused_decode_mla = os.getenv("SGLANG_ROCM_FUSED_DECODE_MLA") == "1"
 
     def no_absorb(self, forward_batch: ForwardBatch) -> bool:
@@ -667,6 +668,9 @@ def no_absorb(self, forward_batch: ForwardBatch) -> bool:
                 and not forward_batch.forward_mode.is_draft_extend()
                 and sum(forward_batch.extend_prefix_lens_cpu) == 0
             )
+        elif self.attention_backend == "fa3":
+            # Flash Attention: Keep absorbing for all extend/decode
+            return False
         else:
             # Triton: Use normal computation for prefill and use weight absorption for extend/decode
             return (