[Bugfix] Disable sequence parallelism / async TP under batch invariance and add a TP regression test (#56377)
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lio Einaudi committed
Sep 27, 2026 at 10:19 UTC
0c87a197b856a25670cc711456533e5ef9f94d9f
2 files changed
+21
-1
docs/features/batch_invariance.md
+1
-1
@@ -144,7 +144,7 @@ When batch invariance is enabled, vLLM:
144
145
1. Uses deterministic kernel implementations for attention and other operations
146
2. Ensures consistent numerical behavior across different batch sizes
147
-3. Disables certain optimizations that may introduce non-determinism (such as custom all-reduce operations in tensor parallel mode)
147
+3. Disables certain optimizations that may introduce non-determinism (such as custom all-reduce operations in tensor parallel mode, and sequence parallelism / async TP, whose reduce-scatter path is not batch-invariant)
148
4. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),
149
runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out
150
(not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).
vllm/config/vllm.py
+20
@@ -1718,6 +1718,26 @@ class VllmConfig:
1718
)
1719
self.compilation_config.mode = CompilationMode.NONE
1720
1721
+ # TODO: This is a stopgap that turns sequence parallelism / async TP
1722
+ # off under batch invariance. Make the sequence-parallel reduce-scatter
1723
+ # path batch-invariant instead so they can stay enabled (#56370).
1724
+ # Sequence parallelism / async TP rewrite all_reduce + rms_norm into a
1725
+ # reduce-scatter whose reduction order depends on the batch, so they
1726
+ # are not batch-invariant. Decide this before the breakable-CUDA-graph
1727
+ # default below, which declines to auto-enable when they are on.
1728
+ pass_config = self.compilation_config.pass_config
1729
+ if envs.VLLM_BATCH_INVARIANT and (
1730
+ pass_config.enable_sp or pass_config.fuse_gemm_comms
1731
+ ):
1732
+ logger.warning_once(
1733
+ "Disabling sequence parallelism and async TP "
1734
+ "(pass_config.enable_sp / fuse_gemm_comms) when "
1735
+ "VLLM_BATCH_INVARIANT is enabled: the reduce-scatter path "
1736
+ "is not batch-invariant (see vllm-project/vllm#56370)."
1737
+ )
1738
+ pass_config.enable_sp = False
1739
+ pass_config.fuse_gemm_comms = False
1740
+
1741
breakable_cudagraph_enabled = self._maybe_enable_breakable_cudagraph()
1742
1743
if not breakable_cudagraph_enabled and (