[AsyncTP] Set return_A=False in AllGatherGEMMPattern to unlock multim…

[tianrengao]

…em path

The AllGatherGEMMPattern replacement calls fused_all_gather_matmul but never uses the gathered tensor (ag_output). However, it relies on the default return_A=True, which prevents the multimem execution path in _should_use_multimem_all_gather_matmul from being selected.

Setting return_A=False signals that the gathered tensor is not needed, enabling the NVLink multicast-based multimem path for small M values. Microbenchmarks on 8xH100 show 1.5x-8.5x per-op speedup for M<=2048.

Note: PyTorch's native Inductor FX pass (micro_pipeline_tp.py) already performs this optimization by checking if ag_output has no downstream users, but vLLM's pattern matcher does not do this liveness analysis.

Purpose

Test Plan

Test Result

---

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, @tianrengao.

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 performance optimization by setting return_A=False in the fused_all_gather_matmul call within the AllGatherGEMMPattern. As explained in the description, this change is safe because the gathered tensor is not used downstream, and it unlocks a faster NVLink multicast-based execution path for small matrix sizes, resulting in significant speedups. The change is well-contained, correct, and the added code comment clearly explains the rationale. I've confirmed that this optimization is specific to fused_all_gather_matmul and does not apply to fused_all_gather_scaled_matmul used in other patterns, as the latter does not have the return_A parameter. The implementation is solid and I have no further suggestions.

[mergify]

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

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

[WoosukKwon]

Closing as superseded by #60517, which removed the torch.compile-based sequence parallelism pass and AsyncTP (along with enable_sp,
fuse_gemm_comms, sp_min_token_num, and their tests). The code this PR modifies no longer exists. Thanks for the contribution!