@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 785b990116

[Distributed] Enable custom all-reduce under VLLM_BATCH_INVARIANT (#58623)

Signed-off-by: Flora Feng <4florafeng@gmail.com> Signed-off-by: sfeng33 <4florafeng@gmail.com>

Flora Feng committed Oct 1, 2026 at 22:25 UTC 785b990116ab2b7c77f5d946b2f52c2a8f9b9515

7 files changed +130 -31

docs/features/batch_invariance.md

+6 -2

@@ -146,11 +146,15 @@ When batch invariance is enabled, vLLM:

146

147 1. Uses deterministic kernel implementations for attention and other operations

148 2. Ensures consistent numerical behavior across different batch sizes

149 -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)

150 -4. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),

149 +3. Disables certain optimizations that may introduce non-determinism (such as sequence parallelism / async TP, whose reduce-scatter path is not batch-invariant)

150 +4. Under tensor parallelism, keeps custom all-reduce on with a fixed reduction order (the 1-stage kernel is pinned, and large inputs are reduced in fixed-size chunks), and disables FlashInfer, AITER and QuickReduce all-reduce

151 +5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),

152 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

153 (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).

154

155 +!!! warning

156 + Batch invariance under tensor parallelism is not yet supported for all-reduces whose size isn't a multiple of 16 bytes, for example a hidden size that isn't a multiple of 8 in fp16/bf16. Such tensors can switch between NCCL and custom all-reduce depending on batch size. All validated models meet this requirement.

157 +

158 !!! note

159 Enabling batch invariance may impact performance compared to the default non-deterministic mode. This trade-off is intentional to guarantee reproducibility.

160

tests/distributed/test_custom_all_reduce.py

+90

@@ -205,6 +205,42 @@ def test_custom_allreduce_filters_dtype(

205 assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) is expected

206

207

208 +@pytest.mark.parametrize("batch_invariant", [False, True])

209 +def test_custom_allreduce_size_gate_ignored_under_batch_invariance(

210 + batch_invariant: bool,

211 +) -> None:

212 + """Batch invariance must not switch backends based on the input size."""

213 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)

214 + communicator.disabled = False

215 + communicator.world_size = 2

216 + communicator.max_size = 1024

217 + communicator.batch_invariant = batch_invariant

218 + communicator._ptr = 0

219 +

220 + oversized = torch.empty(1024, dtype=torch.float16)

221 + assert communicator.should_custom_ar(oversized) is batch_invariant

222 +

223 +

224 +@pytest.mark.parametrize("batch_invariant", [False, True])

225 +def test_custom_reduce_scatter_disabled_under_batch_invariance(

226 + monkeypatch, batch_invariant: bool

227 +) -> None:

228 + """Reduce-scatter stays on one backend under batch invariance."""

229 + monkeypatch.setattr(car.current_platform, "is_cuda", lambda: True)

230 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)

231 + communicator.disabled = False

232 + communicator.world_size = 2

233 + communicator.fully_connected = True

234 + communicator.mnnvl_only = False

235 + communicator.mnnvl_multicast_ptr = 0

236 + communicator.max_reduce_scatter_size = 1024

237 + communicator.max_mnnvl_reduce_scatter_size = 1024

238 + communicator.batch_invariant = batch_invariant

239 +

240 + in_range = torch.empty(64, dtype=torch.float16)

241 + assert communicator.should_custom_reduce_scatter(in_range) is not batch_invariant

242 +

243 +

244 @pytest.mark.parametrize(

245 ("major", "local_multicast", "expected"),

246 [

@@ -622,6 +658,60 @@ def eager_allreduce(

658 torch.testing.assert_close(out, inp * (tp_size**num_communication))

659

660

661 +@ray.remote(num_gpus=1, max_calls=1)

662 +def chunked_allreduce(

663 + monkeypatch: pytest.MonkeyPatch,

664 + tp_size,

665 + pp_size,

666 + rank,

667 + distributed_init_port,

668 +):

669 + """Inputs above max_size are reduced in chunks and match the fixed-order

670 + reference bitwise, both eagerly and inside a CUDA graph."""

671 + with monkeypatch.context() as m:

672 + m.delenv("CUDA_VISIBLE_DEVICES", raising=False)

673 + m.delenv("HIP_VISIBLE_DEVICES", raising=False)

674 + m.setenv("VLLM_CUSTOM_ALLREDUCE_ALGO", "1stage")

675 + device = torch.device(f"cuda:{rank}")

676 + torch.accelerator.set_device_index(device)

677 + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)

678 + ensure_model_parallel_initialized(tp_size, pp_size)

679 + group = get_tp_group().device_group

680 + fa = get_tp_group().device_communicator.ca_comm

681 +

682 + # Two full chunks plus a partial one.

683 + dtype = torch.bfloat16

684 + chunk_numel = fa.max_size // dtype.itemsize

685 + inp = torch.randn(2 * chunk_numel + 4096, device=device).to(dtype)

686 + gathered = [torch.empty_like(inp) for _ in range(tp_size)]

687 + dist.all_gather(gathered, inp, group=group)

688 + ref = gathered[0].float()

689 + for peer in gathered[1:]:

690 + ref = ref + peer.float()

691 + ref = ref.to(torch.bfloat16)

692 +

693 + out = fa.all_reduce(inp, registered=False)

694 + assert torch.equal(out, ref)

695 +

696 + # Weak-contiguous but not C-contiguous: a transposed matrix.

697 + inp_t, ref_t = inp.view(-1, 4096).t(), ref.view(-1, 4096).t()

698 + assert torch.equal(fa.all_reduce(inp_t, registered=False), ref_t)

699 +

700 + with fa.capture():

701 + graph = torch.cuda.CUDAGraph()

702 + with torch.cuda.graph(graph):

703 + out = fa.all_reduce(inp, registered=True)

704 + graph.replay()

705 + assert torch.equal(out, ref)

706 +

707 +

708 +@pytest.mark.parametrize("tp_size", [2, 4])

709 +def test_custom_allreduce_chunked(monkeypatch: pytest.MonkeyPatch, tp_size):

710 + if tp_size > torch.accelerator.device_count():

711 + pytest.skip("Not enough GPUs to run the test.")

712 + multi_process_parallel(monkeypatch, tp_size, 1, chunked_allreduce)

713 +

714 +

715 @pytest.mark.parametrize("tp_size", [2])

716 @pytest.mark.parametrize("pipeline_parallel_size", [1, 2])

717 @pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])

tests/v1/determinism/test_batch_invariance.py

-20

@@ -176,17 +176,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(

176 random.seed(seed)

177 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))

178

179 - # For batch invariance, disable custom all-reduce to ensure deterministic

180 - # all-reduce operations (custom all-reduce may not be deterministic)

181 - import vllm.envs as envs

182 -

183 - disable_custom_ar = envs.VLLM_BATCH_INVARIANT

184 -

185 - if disable_custom_ar:

186 - print(f"\n{'=' * 80}")

187 - print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})")

188 - print(f"{'=' * 80}\n")

189 -

179 llm = LLM(

180 model=TEST_MODEL,

181 tensor_parallel_size=tp_size,

@@ -744,15 +733,6 @@ def test_decode_logprobs_match_prefill_logprobs(

733 random.seed(seed)

734 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))

735

747 - import vllm.envs as envs

748 -

749 - disable_custom_ar = envs.VLLM_BATCH_INVARIANT

750 -

751 - if disable_custom_ar:

752 - print(f"\n{'=' * 80}")

753 - print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})")

754 - print(f"{'=' * 80}\n")

755 -

736 llm = LLM(

737 model=TEST_MODEL,

738 tensor_parallel_size=tp_size,

vllm/config/parallel.py

-4

@@ -1087,10 +1087,6 @@ class ParallelConfig:

1087 # Lazy import to avoid circular import

1088 from vllm.v1.executor import Executor

1089

1090 - # Enable batch invariance settings if requested

1091 - if envs.VLLM_BATCH_INVARIANT:

1092 - self.disable_custom_all_reduce = True

1093 -

1090 if (

1091 self.distributed_executor_backend is not None

1092 and not isinstance(self.distributed_executor_backend, str)

vllm/distributed/device_communicators/cuda_communicator.py

+11 -3

@@ -68,8 +68,11 @@ class CudaCommunicator(DeviceCommunicatorBase):

68 envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC

69 and not envs.VLLM_BATCH_INVARIANT

70 )

71 - use_aiter_allreduce = use_custom_allreduce and bool(

72 - rocm_aiter_ops.is_custom_all_reduce_enabled()

71 + # Neither AITER nor QuickReduce all-reduce has a fixed reduction order.

72 + use_aiter_allreduce = (

73 + use_custom_allreduce

74 + and not envs.VLLM_BATCH_INVARIANT

75 + and bool(rocm_aiter_ops.is_custom_all_reduce_enabled())

76 )

78 self.use_custom_allreduce = use_custom_allreduce

@@ -166,7 +169,12 @@ class CudaCommunicator(DeviceCommunicatorBase):

169 else:

170 self.use_aiter_ag_rs = True

171

169 - if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm():

172 + if (

173 + use_custom_allreduce

174 + and self.world_size > 1

175 + and current_platform.is_rocm()

176 + and not envs.VLLM_BATCH_INVARIANT

177 + ):

178 # Initialize a custom quick all-reduce implementation for AMD.

179 # Quick reduce is designed as a complement to custom allreduce

180 # (vLLM's or AITER's), so it is initialized for either backend.

vllm/distributed/device_communicators/custom_all_reduce.py

+21 -2

@@ -119,6 +119,7 @@ from vllm.distributed.utils import is_weak_contiguous # noqa: E402

119

120 class CustomAllreduce:

121 _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8, 16]

122 + batch_invariant = False

123 _DEFAULT_ALL_GATHER_MAX_SIZE = 2 * 1024 * 1024

124 _DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES = {

125 2: 8 * 1024 * 1024,

@@ -179,6 +180,7 @@ class CustomAllreduce:

180 self.mnnvl_multimem_rs_local_ptr = 0

181 self.mnnvl_multimem_rs_multicast_ptr = 0

182 self.mnnvl_only = False

183 + self.batch_invariant = envs.VLLM_BATCH_INVARIANT

184

185 if not custom_ar:

186 # disable because of missing custom allreduce library

@@ -504,7 +506,7 @@ class CustomAllreduce:

506 # for 4 or more non NVLink-capable GPUs, custom allreduce provides

507 # little performance improvement over NCCL.

508 if self.world_size == 2 or self.fully_connected:

507 - return inp_size < self.max_size

509 + return self.batch_invariant or inp_size < self.max_size

510 return False

511

512 def all_reduce(

@@ -518,13 +520,27 @@ class CustomAllreduce:

520 """

521 if out is None:

522 out = torch.empty_like(inp)

523 + chunk_numel = self.max_size // inp.element_size()

524 + if inp.numel() <= chunk_numel:

525 + self._all_reduce_chunk(inp, out, registered)

526 + return out

527 + # Storage-order views: the gate only admits one contiguous block.

528 + flat_inp = inp.as_strided((inp.numel(),), (1,), inp.storage_offset())

529 + flat_out = out.as_strided((out.numel(),), (1,), out.storage_offset())

530 + for start in range(0, flat_inp.numel(), chunk_numel):

531 + end = start + chunk_numel

532 + self._all_reduce_chunk(flat_inp[start:end], flat_out[start:end], registered)

533 + return out

534 +

535 + def _all_reduce_chunk(

536 + self, inp: torch.Tensor, out: torch.Tensor, registered: bool

537 + ) -> None:

538 if registered:

539 ops.all_reduce(self._ptr, inp, out, 0, 0)

540 else:

541 ops.all_reduce(

542 self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size

543 )

527 - return out

544

545 def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None:

546 """The main allreduce API that provides support for cuda graph."""

@@ -603,6 +619,9 @@ class CustomAllreduce:

619 ) -> _ReduceScatterBackend | None:

620 if self.disabled or not current_platform.is_cuda():

621 return None

622 + # The size gates below would switch backends per batch.

623 + if self.batch_invariant:

624 + return None

625 if self.world_size == 16 and not self.mnnvl_only:

626 return None

627 inp_size = inp.nbytes

vllm/model_executor/determinism/batch_invariant.py

@@ -1139,6 +1139,8 @@ def enable_batch_invariant_mode():

1139

1140 def override_envs_for_invariance():

1141 os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"

1142 + # Only the 1-stage kernel has a size- and rank-independent reduction order.

1143 + os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = "1stage"

1144

1145 os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

1146