[Feature][torch.compile] Add pass to rearrange AllGather for FP8 models in sequence parallel for better Async TP fusion

[jasonlizhengjian]

Purpose

Currently, the Async TP AllGather GEMM pass only fuses patterns of AllGather - GEMM , however in FP8 this is broken up by a type conversion and currently looks like AllGather (BF16) - Convert (BF16 to FP8) - GEMM (FP8) such as in this profile (the convert shows up as a blue triton kernel):

This PR introduces a new pass to convert/rearrange AllGather (BF16) - Convert (BF16 to FP8) to Convert (BF16 to FP8) - AllGather (FP8) which speeds up the AllGather operation and allows for the Async TP fusion to happen.

Profile of after the FP8 AllGather Pass + Async TP

Enable this pass with "pass_config": {"enable_async_tp": true, "enable_sequence_parallelism": true, "enable_fp8_allgather_opt": true}

this pass must be used with sequence parallelism
Only static quantization scales are supported (can be matched by the pattern)

This PR depends on #26038

Test Plan

Added test of this pass to tests/compile/test_async_tp.py

Benchmarked to show e2e speedups

Test Result

Benchmarks on H100 with TP = 4

vllm bench latency --model nvidia/Llama-3.3-70B-Instruct-FP8 --output-len 1 --input-len 8192 --batch-size 1 --tensor-parallel-size 4 --load-format dummy --num_iters_warmup 5 --num_iters 10 -O "$COMPILE_CONFIG" --no-enable-prefix-caching

# Baseline (without Async TP)
COMPILE_CONFIG='{"level":3, "compile_sizes": [8192]}'
Avg latency: 0.405582538433373 seconds
10% percentile latency: 0.40465721269138155 seconds
25% percentile latency: 0.40539184131193906 seconds
50% percentile latency: 0.40564539190381765 seconds
75% percentile latency: 0.40587058768142015 seconds
90% percentile latency: 0.40601281519047916 seconds
99% percentile latency: 0.4066646806197241 seconds

# Async TP without the new pass
COMPILE_CONFIG='{"level":3, "compile_sizes": [8192], "pass_config": {"enable_async_tp": true, "enable_sequence_parallelism": true, "enable_fp8_allgather_opt": false}}'
Avg latency: 0.38048242605291305 seconds
10% percentile latency: 0.37953354166820646 seconds
25% percentile latency: 0.3799488676013425 seconds
50% percentile latency: 0.3806821608450264 seconds
75% percentile latency: 0.3810685377102345 seconds
90% percentile latency: 0.3813340303953737 seconds
99% percentile latency: 0.3813573355646804 seconds

# Async TP with this new pass (This PR)
COMPILE_CONFIG='{"level":3, "compile_sizes": [8192], "pass_config": {"enable_async_tp": true, "enable_sequence_parallelism": true, "enable_fp8_allgather_opt": true}}'
Avg latency: 0.30300158439204095 seconds
10% percentile latency: 0.3005924086552113 seconds
25% percentile latency: 0.3008677661418915 seconds
50% percentile latency: 0.3024391559883952 seconds
75% percentile latency: 0.3050330716650933 seconds
90% percentile latency: 0.3055407572071999 seconds
99% percentile latency: 0.3067375366995111 seconds

1.27x speedup over Async TP without the pass introduced by this PR

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jasonlizhengjian.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces a new compiler pass to optimize AllGather operations for FP8 models by reordering them with quantization, which is a valuable performance optimization. The implementation, including the new pattern matching pass and supporting changes to configuration and communication libraries, appears solid. The addition of a correctness test with logprob comparison is also a good practice.

However, there is a critical issue in the PR description's benchmark results. The provided numbers indicate a 1.27x slowdown when the optimization is enabled, which contradicts the claim of a '1.27x speedup'. Please verify these results and update the description, as this is crucial for evaluating the feature's impact. It's possible the results for the two configurations were swapped in the description.

I have also left a comment on a potentially overly strict assertion in the new test logic in tests/utils.py that could lead to flaky tests.

[gemini-code-assist]

The assertion for top-k token overlap seems overly strict and contradicts the comment on lines 712-714. The comment states 'Require at least 95% overlap (at most 1 token differs in top-5)', but for a top-5 list, allowing one token to differ corresponds to an 80% overlap (4/5). The current threshold of 0.95 requires a perfect match (5/5 common tokens), which might lead to flaky tests due to minor numerical precision differences inherent with FP8. I suggest lowering the threshold to 0.8 to align with the intention of allowing one differing token.

Suggested change

	assert overlap >= 0.95, (
	assert overlap >= 0.8, (

[jasonlizhengjian]

However, there is a critical issue in the PR description's benchmark results. The provided numbers indicate a 1.27x slowdown when the optimization is enabled, which contradicts the claim of a '1.27x speedup'. Please verify these results and update the description, as this is crucial for evaluating the feature's impact. It's possible the results for the two configurations were swapped in the description.

fixed the description to unswap these

[jasonlizhengjian]

@cascade812 please take a look if you have time

[jasonlizhengjian]

cc @andoorve

[andoorve]

@ProExpertProg can you comment if this aligns with what we were discussing re: #25179

Cc: @tlrmchlsmth @mgoin

[ProExpertProg]

@andoorve I don't think we need an additional pass; I think we should be able to fix the current SequenceParallelismPass to handle this case. Also it will be easier to do after my PR #24604, which should hopefully get merged sometime this week.

[andoorve]

Sounds good, we've been meaning to take a look at that PR this week. We'll see what we can do and update here.