minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: luhongyu.4869 <luhongyu.4869@bytedance.com>
This commit is contained in:
narutolhy
2026-04-04 02:09:45 -07:00
committed by GitHub
co-authored by luhongyu.4869
parent 34d5765e2f
commit 24763256b9
5 changed files with 94 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/attention_backend.md
+1 -1
View File
@@ -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
+1
View File
@@ -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
+6 -1
View File
@@ -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
+23 -4
View File
@@ -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
+63 -5
View File
@@ -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()