Further Speed up FA3 Backend

[hebiao064]

We explored and discussed some ideas and we want to write it down for tracking, also welcome community developer to try out those unfinished

• (Good first issue) Skip len operation, get it directly from forward batch: FA3 speed up: skip len operation and get batch size directly from forward batch #5969 @lifuhuang
• GQA head packing: https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_attn_interface.py#L658 Change it to True and run benchmark.
• Split-KV. aka Flash Decoding: We already enabled it, it is indeed faster in lower batch and long context scenario. Benchmark will be attached.
• PDL: Dao-AILab/flash-attention@000090d
• (Won't do) Prepare Scheduler Metadata: Dao-AILab/flash-attention@fa60e7c (From Tri Dao's note, it can only speed up 2us, we can keep an eye on this, not recommending adopting this)
• For Llama Models, we observed that Spec Decoding with Top K > 1 is slightly slower than Flash Infer backend, we need comprehensive profiling and optimize it @MrAta
• Replace Pad operation by Copy: Optimize a pad operation to accelerate 25us #5945
• Remove is_fa3_supported from fa3 kernel: Remove unecessary is_fa3_supported check #6112
• Remove pad operation for all decode cases: optimize pad operations in fa3 to accelarate 100+us #6077

[sanchit-ahuja]

@hebiao064 Can I pick this up?

[hebiao064]

@sanchit-ahuja sure, please let me know which part you want to pick up

[sanchit-ahuja]

I can pick up the first and second to-do items for now and can look into PDL and the llama spec decoding issue later. Also, what benchmark do we need to run for the second item?

[hebiao064]

@sanchit-ahuja #5611 (comment)
Refer to this one, but you don't need to run it against Deepseek so you don't need to compile DeepGemm first.

[hebiao064]

Another potential issue for CUDA Graph: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/flashattention_backend.py#L1590

need to change to

metadata.cu_seqlens_k.copy_(
                torch.nn.functional.pad(
                    torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0)
                )
            )

to ensure CUDA Graph stability

[sanchit-ahuja]

Got it. Thanks!

[sanchit-ahuja]

@hebiao064 Can you let me know what do we need to do for PDL and how you want me to run profiling for Lllama models?

[sanchit-ahuja]

Also, I was looking into the pack_gqa issue and found out that this flag is part of the sgl-kernel within the repo. How can I expose this flag to the sglang framework for me to run the benchmark? There's also a test_flash_attention.py file within the the sgl-kernel's tests folder which disables pack_gqa test explicitly with a DISABLE_PACKGQA set to True. Do you want me to run that as well?

[sanchit-ahuja]

I did try running the test test_flash_attention.py but I think my GPU is not supported to run this. I have an A100 machine and I am getting this error: flash_attn at sgl-kernel is only supported on sm90 and above. Let me know how you want me to proceed ahead with this. Although, FA3 does support A100 now ref.
Thanks!

[hebiao064]

I did try running the test test_flash_attention.py but I think my GPU is not supported to run this. I have an A100 machine and I am getting this error: flash_attn at sgl-kernel is only supported on sm90 and above. Let me know how you want me to proceed ahead with this. Although, FA3 does support A100 now ref. Thanks!

#5536 (comment)

[sanchit-ahuja]

(Good first issue) Skip len operation, get it directly from forward batch: https://github.com/hebiao064/sglang/blob/main/python/sglang/srt/layers/attention/flashattention_backend.py#L341 , https://github.com/hebiao064/sglang/blob/main/python/sglang/srt/model_executor/forward_batch_info.py#L144

Another potential issue for CUDA Graph

These are done.

Do you want me to raise a PR for these fixes? In the meantime, I can wait for a fix for the Ampere GPUs. Will see if I can get access to an H100 machine somehow.