[Roadmap] FlashAttention3 Support as SGLang Attention Backend

[hebiao064]

Functionality

• Basic FA3 support including MHA Models (Llama, QWen and etc), Cuda Graph, Sliding Window (Gemma): Support FA3 as Attention backend by using --attention-backend fa3 #4680 @hebiao064 @qingquansong
• Support Page Size > 1 Support Page Size > 1 for FA3 #4832 @hebiao064
• Support MLA for Deepseek-like models [Feature] Support FA3 backend for MLA #4831 @Fridge003
• Support Speculative Decoding [WIP] Add Eagle Speculative Decoding to FA3 Backend #4951 @qingquansong
• Figure out how to build FA3 into SGLang: [feat] add fa3 in sgl-kernel #4902 @yinfan98
• Add E2E Test like sglang/test/srt/test_triton_attention_backend.py: Add integration test for Flash Attention 3 #4760 @yubofredwang
• Support Multimodal @mickqian @yizhang2077
• Support FP8 @hebiao064 @yundai424
• MInference @yinfan98 @ZhangJianwei0311 @minminsun @zhyncs

Documentation:

• Add Docs about attention backend, maybe we need to add a comprehensive one with Support Matrix such as: #4865

Perf Optimization and Accuracy Problems

• Fix Cuda Graph Accuracy problem for Page Size > 1: Fix fa3 cuda graph page_size > 1 precision and page_size=1 speed #4855 @qingquansong
• Optimizing Decoding by remove `item() device sync: [FA3 Attn Backend] Remove Unnecessary Device Sync for FA3 #4745 @hebiao064
• Optimizing Prefill by remove item() device sync: [Fix] avoid stream sync and torch compile in prefill for fa3 backend #4932 @Fridge003

Success Criteria:

• The latency should be on par with vLLM FlashAttention3 and SGLang's FlashInfer implementation
• The accuracy should be on par with vLLM FlashAttention3 and SGLang's FlashInfer implementation

Other issues we surfaced but not scoped in this task:

• Flash Infer accuracy is bad for Gemma 2 Models
• VSCode Test Explorer is broken since some circular dependency: Fix circular imports in gptq.py and unblock test explorer #4736 @hebiao064

[zcnrex]

Will work on speculative decoding

[zhyncs]

ref #4686

[hebiao064]

I'll add accuracy and latency benchmark after each major feature introduction in this issue:

Accuracy:

After initial PR of #4680

Model	FA3 Accuracy	Flash Infer Accuracy
Meta-Llama-3.1-8B-Instruct	0.793	0.789
Qwen2.5-7B-Instruct	0.823	0.789
Gemma-2-9B	0.724 (Torch Native is 0.730)	0.132 (potential bug!)

After #4832 and #4855 (Page Size > 1):

Accuracy: with Cuda Graph:

Model	FA3 Accuracy	Flash Infer Accuracy
Meta-Llama-3.1-8B-Instruct	0.792/0.796	0.792/0.792
Qwen2.5-7B-Instruct	0.819/0.818	0.809/0.810

Note: 0.792/0.796 means 0.792 for page_size = 1, 0.796 for page_size = 128

From this we can conclude that page size shouldn't have impact on accuracy

[hebiao064]

Latency:

Benchmark Command I used:

python -m sglang.bench_one_batch --model /path/to/Meta-Llama-3.1-8B-Instruct  --batch-size 16 --input 1024 --output 512 --attention-backend fa3

After initial PR of #4680

Model	FA3 Latency	Flash Infer Latency
Meta-Llama-3.1-8B-Instruct with Cuda Graph	45328.51/1967.94	43511.35/1960.41
Meta-Llama-3.1-8B-Instruct w/o Cuda Graph	45648.71/1296.51	43664.38/1237.78

After couple of changes (introduced more logic, might have negative impact on latency) plus one optimization PR against Prefill: #4932 and one against Decode: #4745

Model	FA3 Latency	Flash Infer Latency
Meta-Llama-3.1-8B-Instruct with Cuda Graph	44392.57/2477.80	44736.24/2415.29
Meta-Llama-3.1-8B-Instruct w/o Cuda Graph	44637.53/1335.13	44796.29/1244.60