[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 PR1, PR2, PR3 PR4 @qingquansong @hebiao064 @zcnrex
• 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 unit test on page_size > 1 and mla and integration test for Flash Attention 3 #4760 @yubofredwang
• Support Multimodal [FA3 Feature] Support multi modal Llama-3.2-11B-Vision-Instruct #5103 @zcnrex @mickqian @yizhang2077
• Support FP8 Fix loading KV quantization scale; Enable modelopt kv cache #4686 @yundai424

Documentation and Benchmark:

• Add Docs about attention backend, maybe we need to add a comprehensive one with Support Matrix such as: #4865
• Benchmark Attention Backends with different features and models #5172 @zhyncs @hebiao064 Shivam (In Review, will be tracked offline)

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
• Optimizing Draft Decode and Target Verify CUDA Graph Latency: Refactor and Optimize FA3 Code #5090 @hebiao064

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 on H100:

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

** after some retrospective, I found above benchmark was done on H200, so we can ignore this section...**

After Spec Decode changes (introduced more logic, might have negative impact on latency)

Model	FA3 Latency	Flash Infer Latency
Meta-Llama-3.1-8B-Instruct with Cuda Graph	43566.63/1938.83	43838.53/1960.30
Meta-Llama-3.1-8B-Instruct w/o Cuda Graph	43468.50/1315.45	43782.75/1238.26