[Transformers backend] Add sequence parallelism and async TP

[hmellor]

Overview

Add opt-in sequence parallelism (SP) and async TP to the Transformers modeling backend, implemented eagerly with module hooks rather than as a torch.compile pass. This restores both features for Transformers backend models after #60517 removes the compile-based SequenceParallelismPass/AsyncTPPass, and unlike those passes it works on Model Runner V2 and with any token count.

Claims

• --enable-sequence-parallel shards the residual stream across TP ranks between the first and last decoder layer of each pipeline stage. Each child of a decoder layer that reduces across TP (attention, MLP, MoE, MLA, ...) all-gathers its input and reduce-scatters its output instead of all-reducing it, so norms and residual adds only process 1 / tp_size of the tokens.
• --enable-async-tp (implies SP) fuses the collectives with the adjacent unquantized GEMMs via torch.ops.symm_mem.fused_all_gather_matmul / fused_matmul_reduce_scatter. Quantized linears fall back to plain collectives.
• Works with dense and MoE models (internal and external routing, and use_sequence_parallel_moe, where the experts take the sharded tokens directly), MLA, quantized models, PP, Eagle3, torch.compile and CUDA graphs.
• Both flags only affect the Transformers backend; a warning is logged if they are set for a model that does not use it.
• Safety checks at init: a region that applies a non-commuting op (residual add, norm, activation, ...) to a partial sum raises, and a decoder layer that may mix tokens outside its regions logs a warning.
• fx_utils.trace now traces a fuser's per-instance forward (needed to find the fused QKV/gate-up linears for async TP).
• Not yet done: performance benchmarks, and FP8 async TP fusion (the removed pass supported FP8 static scales).

Validation

B200, TP=2, torch.compile + FULL_AND_PIECEWISE CUDA graphs unless stated.

Check	Result
test_distributed (Llama-3.2-1B vs HF): TP / SP / async TP	3/3 pass; async TP fuses 64/64 collectives
test_quantization with TP=2 + SP (AWQ, GPTQ × eager, compiled)	4/4 pass
OLMoE (MoE, internal routing): TP vs SP, greedy	3/4 prompts identical, 1 diverges at a near tie
OLMoE, TP=2 × DP=2 + EP (use_sequence_parallel_moe): EP vs EP + SP	top-2 logprobs within 0.1 for all compared tokens
Granite-3.1-1B-A400M (MoE, external routing): TP vs SP	4/5 identical
ERNIE-4.5-21B-A3B (external routing + shared experts): TP vs SP	4/5 identical, 1 diverges at a near tie (chosen logprob -0.885 vs -0.858)
Tiny random DeepSeek-V2 (MLA): TP vs SP	loads and runs; divergences only at near-uniform ties
VLLM_GPU_SYNC_CHECK=error (Llama SP, Llama async TP, OLMoE SP)	no synchronization detected; generations coherent
tests/models/transformers/fusers (CPU)	118 passed, 1 skipped

Plain TP is deterministic run to run on Granite, but SP varies by about one bf16 logit step at one near-tie token (probabilities 0.5/0.5 vs 0.62/0.38 across two runs). The cause is not yet identified.

pytest tests/models/transformers/test_backend.py -k test_distributed -v
pytest tests/models/transformers/fusers -q
vllm serve <model> --model-impl transformers -tp 2 --enable-sequence-parallel  # or --enable-async-tp

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--60653.org.readthedocs.build/en/60653/