[Speculative Decoding] Add FA4-based Spec Support (#21080)

Co-authored-by: luhongyu.4869 <luhongyu.4869@bytedance.com>

python/sglang/jit_kernel/flash_attention_v4.py

@@ -83,6 +83,7 @@ def flash_attn_varlen_func(
         pack_gqa=pack_gqa,
         score_mod=score_mod,
         aux_tensors=aux_tensors,
+        return_lse=return_softmax_lse,
     )
 
     if return_softmax_lse:

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

@@ -2421,7 +2421,11 @@ def prepare_swa_spec_page_table_triton(
 class FlashAttentionMultiStepBackend:
 
     def __init__(
-        self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+        fa_impl_ver: int = 3,
     ):
         self.model_runner = model_runner
         self.topk = topk
@@ -2434,6 +2438,7 @@ class FlashAttentionMultiStepBackend:
                     speculative_step_id=i,
                     topk=self.topk,
                     speculative_num_steps=self.speculative_num_steps,
+                    fa_impl_ver=fa_impl_ver,
                 )
             )
 

python/sglang/srt/speculative/draft_utils.py

@@ -55,6 +55,7 @@ class DraftBackendFactory:
             "trtllm_mla": self._create_trtllm_mla_decode_backend,
             "nsa": self._create_nsa_decode_backend,
             "ascend": self._create_ascend_decode_backend,
+            "fa4": self._create_fa4_decode_backend,
         }
 
         return self._create_backend(
@@ -79,6 +80,7 @@ class DraftBackendFactory:
             "trtllm_mla": self._create_trtllm_mla_prefill_backend,
             "nsa": self._create_nsa_prefill_backend,
             "ascend": self._create_ascend_prefill_backend,
+            "fa4": self._create_fa4_prefill_backend,
         }
         backend_name = (
             "decode_attention_backend"
@@ -139,15 +141,24 @@ class DraftBackendFactory:
             self.draft_model_runner, self.topk, self.speculative_num_steps
         )
 
-    def _create_fa3_decode_backend(self):
+    def _create_fa_decode_backend(self, fa_impl_ver: int = 3):
         from sglang.srt.layers.attention.flashattention_backend import (
             FlashAttentionMultiStepBackend,
         )
 
         return FlashAttentionMultiStepBackend(
-            self.draft_model_runner, self.topk, self.speculative_num_steps
+            self.draft_model_runner,
+            self.topk,
+            self.speculative_num_steps,
+            fa_impl_ver=fa_impl_ver,
         )
 
+    def _create_fa3_decode_backend(self):
+        return self._create_fa_decode_backend(fa_impl_ver=3)
+
+    def _create_fa4_decode_backend(self):
+        return self._create_fa_decode_backend(fa_impl_ver=4)
+
     def _create_flashmla_decode_backend(self):
         from sglang.srt.layers.attention.flashmla_backend import (
             FlashMLAMultiStepDraftBackend,
@@ -213,12 +224,20 @@ class DraftBackendFactory:
 
         return AiterAttnBackend(self.draft_model_runner, skip_prefill=False)
 
-    def _create_fa3_prefill_backend(self):
+    def _create_fa_prefill_backend(self, fa_impl_ver: int = 3):
         from sglang.srt.layers.attention.flashattention_backend import (
             FlashAttentionBackend,
         )
 
-        return FlashAttentionBackend(self.draft_model_runner, skip_prefill=False)
+        return FlashAttentionBackend(
+            self.draft_model_runner, skip_prefill=False, fa_impl_ver=fa_impl_ver
+        )
+
+    def _create_fa3_prefill_backend(self):
+        return self._create_fa_prefill_backend(fa_impl_ver=3)
+
+    def _create_fa4_prefill_backend(self):
+        return self._create_fa_prefill_backend(fa_impl_ver=4)
 
     def _create_trtllm_mha_prefill_backend(self):
         from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend