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

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

docs/advanced_features/attention_backend.md

@@ -19,7 +19,7 @@ The support matrix is split into two parts: MHA (standard attention) and MLA (mu
 |---------------------------------|-----------------------------|------------------|-----------------|-----------------|-----------------|--------------------|----------------|
 | **FlashInfer**                  | ✅                          | ✅               | ❌              | ✅              | ✅              | ✅                 | ❌             |
 | **FA3 (FlashAttention 3)**      | ✅                          | ✅               | ❌              | ✅              | ✅              | ✅                 | ✅             |
-| **FA4 (FlashAttention 4)**      | 128                         | ❌               | ✅              | ❌              | ❌              | ❌                 | ✅             |
+| **FA4 (FlashAttention 4)**      | 128                         | ❌               | ✅              | ✅              | ✅              | ❌                 | ✅             |
 | **Triton**                      | ❌                          | ✅               | ✅              | ✅              | ✅              | ✅                 | ✅             |
 | **Torch Native (SDPA)**         | ❌                          | ✅               | ✅              | ❌              | ❌              | ❌                 | ✅             |
 | **FlexAttention (PyTorch)**     | ❌                          | ❌               | ✅              | ❌              | ❌              | ❌                 | ❌             |

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

test/registered/attention/test_flash_attention_4.py

@@ -1,6 +1,7 @@
 import unittest
 from types import SimpleNamespace
-from urllib.parse import urlparse
+
+import requests
 
 from sglang.srt.utils import get_device_sm, kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -23,10 +24,8 @@ class TestFlashAttention4(unittest.TestCase):
         cls.base_url = DEFAULT_URL_FOR_TEST
         other_args = [
             "--trust-remote-code",
-            "--prefill-attention-backend",
+            "--attention-backend",
             "fa4",
-            "--decode-attention-backend",
-            "flashinfer",
         ]
         cls.process = popen_launch_server(
             cls.model,
@@ -40,7 +39,6 @@ class TestFlashAttention4(unittest.TestCase):
         kill_process_tree(cls.process.pid)
 
     def test_gsm8k(self):
-        parsed_url = urlparse(self.base_url)
         args = SimpleNamespace(
             base_url=self.base_url,
             model=self.model,
@@ -56,5 +54,65 @@ class TestFlashAttention4(unittest.TestCase):
         self.assertGreater(metrics["score"], 0.89)
 
 
+@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
+class TestFlashAttention4SpeculativeDecodeTopk(unittest.TestCase):
+    """Test FlashAttention4 with EAGLE3 speculative decoding (topk > 1).
+
+    Verifies that FA4 + EAGLE3 topk > 1 produces correct outputs and
+    achieves meaningful speculative acceptance length.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--attention-backend",
+            "fa4",
+            "--speculative-algorithm",
+            "EAGLE3",
+            "--speculative-draft-model-path",
+            "lmsys/SGLang-EAGLE3-Qwen3-30B-A3B-Instruct-2507-SpecForge-Nex",
+            "--speculative-num-steps",
+            "5",
+            "--speculative-eagle-topk",
+            "4",
+            "--speculative-num-draft-tokens",
+            "8",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1319,
+            num_threads=200,
+        )
+        metrics = run_eval(args)
+        print(metrics)
+        self.assertGreater(metrics["score"], 0.89)
+
+        server_info = requests.get(self.base_url + "/server_info").json()
+        avg_spec_accept_length = server_info["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreater(avg_spec_accept_length, 1.5)
+
+
 if __name__ == "__main__":
     unittest.main()