diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md index 461554575..98d07d31a 100644 --- a/docs/advanced_features/attention_backend.md +++ b/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)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | diff --git a/python/sglang/jit_kernel/flash_attention_v4.py b/python/sglang/jit_kernel/flash_attention_v4.py index dcd5f2334..0a79614ee 100644 --- a/python/sglang/jit_kernel/flash_attention_v4.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 2b3a7d3c1..ff170c390 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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, ) ) diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index 9c630da72..271dc00d0 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/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 diff --git a/test/registered/attention/test_flash_attention_4.py b/test/registered/attention/test_flash_attention_4.py index 3c9c4242b..ce61ec252 100644 --- a/test/registered/attention/test_flash_attention_4.py +++ b/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()