Collectives, topology, custom all-reduce

Tensor parallelism buys you a smaller weight-read per GPU and charges you two all-reduces per transformer block (§5.1 derives where they land). The bill looks trivial when you write it down as bytes and terrifying when you write it down as events.

Take Llama-3-8B, $L = 32$, $d = 4096$, bf16, TP=8 on one H100 SXM node, batch $t = 8$ decode tokens. Each all-reduce carries $t \cdot d \cdot 2 = 65{,}536$ bytes. A ring all-reduce over $N = 8$ ranks puts $2(N{-}1)/N \cdot S = 1.75 \times 64\,\mathrm{KiB} = 114{,}688$ bytes on the wire per rank, and NVLink 4 gives each H100 900 GB/s bidirectional — 450 GB/s per direction (NVIDIA H100). That is 0.26 µs of transfer. Sixty-four of them: 16 µs.

Now the other term. Every collective costs something fixed before a byte moves: a kernel launch, and a rendezvous in which all eight ranks agree they have arrived. §0.3 put bare launch overhead at 0.5–0.9 µs pipelined and 6.7 µs synchronised; the cross-device rendezvous is on top. Call the total $\alpha$ and, for now, take $\alpha = 3\ \mu\mathrm{s}$ — §5.4.3 shows that both engines' hard-coded thresholds imply a number in that band. Then:

Derived, not measured — the arithmetic is above and it scales linearly in $\alpha$. But the shape of the answer does not depend on the exact $\alpha$: 92% of the all-reduce cost in a decode step is fixed overhead, not bandwidth. Doubling NVLink bandwidth would save you 8 µs a step; halving $\alpha$ would save you 96.

This is the opposite of the training regime everyone's intuition comes from. A 100 MB gradient bucket all-reduced over the same eight GPUs puts 175 MB on the wire, 389 µs of transfer, against the same 3 µs of fixed cost — 0.8%. Same collective, same hardware, same NCCL. The message is fifteen hundred times bigger and the engineering problem inverts completely.

A collective has two costs that do not mix. One is proportional to bytes and is set by the slowest link the algorithm touches. The other is a constant $\alpha$ per invocation, set by launch machinery and by how many times the ranks have to synchronise with each other. Ring algorithms minimise the first at the expense of the second: they achieve the information-theoretic minimum wire traffic but need $2(N{-}1)$ lock-step rounds to do it. At $N = 8$ that is fourteen sequential handshakes to move 229 KB.

Direct peer-to-peer algorithms invert the trade. If every GPU can address every other GPU's memory — which NVSwitch plus CUDA IPC gives you inside a node — a rank can simply read all seven peer buffers and sum them locally. That moves $(N{-}1)S$ bytes instead of $1.75S$, four times more at $N=8$, but it needs exactly one synchronisation instead of fourteen. Below some message size the extra bytes are cheaper than the extra handshakes. That size is the crossover, and both engines hard-code it.

The reason this matters so much for inference and so little for training is that the decode all-reduce sits on the critical path between two dependent matmuls. The output projection produces a partial sum; the residual add needs the reduced sum; nothing else in that layer is ready to run. There is no other work to hide the collective behind.

Define the symbols once. $N$ is the number of ranks in the group. $S$ is the message size in bytes — for an all-reduce, the size of the tensor each rank contributes and receives. $B$ is the achievable unidirectional bandwidth of the link a rank uses, in bytes/second. $\alpha$ is the fixed per-invocation cost: kernel launch plus one rendezvous. $\alpha_{\mathrm{step}}$ is the incremental cost of one additional lock-step synchronisation inside the algorithm.

Ring all-reduce

A ring all-reduce splits the tensor into $N$ chunks of $S/N$ bytes and runs two phases. Phase one, reduce-scatter: $N{-}1$ steps, in each of which every rank sends one chunk to its successor and accumulates the chunk it receives. After $N{-}1$ steps each rank owns exactly one fully reduced chunk. Phase two, all-gather: $N{-}1$ more steps circulating those finished chunks. Every step every rank sends $S/N$ bytes, so:

The $2(N{-}1)/N$ coefficient is the reason ring is the default: it is provably the minimum number of bytes any rank must send, and it tends to $2S$ as $N$ grows rather than to $NS$. The $2(N{-}1)\alpha_{\mathrm{step}}$ term is the reason it is wrong for decode: it grows linearly in $N$ and does not shrink with $S$.

All-gather and reduce-scatter are each exactly half an all-reduce — which is why DP-attention's gather-then-scatter pair costs the same wire bytes as the TP all-reduce it replaces, and why the choice between them is about where the tokens live, not about bandwidth (§5.3).

Two-shot: the ring's byte count without the ring's rounds

The third row is the interesting one. Two-shot does a direct-P2P reduce-scatter followed by a direct-P2P all-gather. Rank $r$ reads the $r$-th shard from all $N$ peers and reduces it — $(N{-}1)S/N$ bytes pulled — then, after one barrier, every rank pulls the $N$ finished shards — another $(N{-}1)S/N$. Total $2(N{-}1)S/N$: identical to ring. The difference is that it needs two synchronisation rounds instead of fourteen, because a full-mesh fabric lets a rank talk to every peer simultaneously instead of only its ring successor.

One-shot wins while its extra bytes cost less than two-shot's extra barrier. Setting the two equal:

vLLM does not tune this at runtime; it hard-codes $S^{*}$ per world size. So run the formula backwards and ask what barrier cost those constants imply. At $N = 8$ the constant is 256 KiB, giving $\alpha_{\mathrm{step}} = 262144 \times 42 / (8 \times 450\times10^{9}) = 3.06\ \mu\mathrm{s}$. At $N = 4$ the constant is 512 KiB, giving $524288 \times 6/(4 \times 450\times10^{9}) = 1.75\ \mu\mathrm{s}$. Two independently chosen magic numbers, one formula, and the implied barrier cost lands in the 1.7–3.1 µs band and grows with $N$ — exactly what an $N$-way flag rendezvous should do, and exactly the microsecond scale §0.3 established for anything involving a launch or a device-wide sync. This is a consistency check on an ideal relative-cost model, not an independent measurement. The illustrative tables and plotted curves below use an effective fixed collective cost $\alpha_{\mathrm{eff}}=3\,\mu\mathrm{s}$, folding internal synchronization into that assumed constant; they do not substitute $3\,\mu\mathrm{s}$ for every round of the detailed ring equation.

Where a decode step actually sits on this curve

Read the last two rows together. At batch 256 the wire term finally dominates and you are in the a bandwidth-sensitive regime where algorithm choice still depends on topology and implementation. At batch 1 to 8 — representative small-batch latency-oriented examples — the fixed cost is 96× the wire cost for the small model. The custom all-reduce exists to attack $\alpha$, not $B$.

Every number above assumed one link speed. Real clusters have three, an order of magnitude apart, and which one a collective lands on is decided entirely by how you assign ranks to GPUs.

Derive the rule instead of asserting it. Put a TP=8 group across two nodes, four GPUs each. Lay the ring out 0–3 on node 0 and 4–7 on node 1: exactly two of the eight ring links cross the InfiniBand fabric. Each of the fourteen steps pushes $S/8$ down every link, so a crossing link carries $14 \times S/8 = 1.75S$ over the whole collective — 229 KB at $S = 128$ KiB. At 50 GB/s that is 4.59 µs against 0.51 µs all-NVLink, and because ring steps are lock-step the slow link sets the pace for everyone. Nine times longer for the same bytes, and the fixed cost per step is worse still, because an IB hop is not a device-side flag poll.

Multiply by 160 collectives per decode step: roughly 736 µs of pure cross-fabric transfer, added to a 5.2 ms Llama-3-70B step, before counting the additional $\alpha$. Now the alternative. PP=2 × TP=4 with each TP group inside one node: all 160 all-reduces stay on NVLink at $1.5S/B = 0.44\ \mu\mathrm{s}$ each, and the cross-fabric traffic collapses to one point-to-point send per micro-batch per stage boundary. The traced path sends both hidden state and residual, sharded across TP: $2S/4=64$ KiB per rank at batch 8, or 1.31 µs on an independent 50 GB/s link. The four ranks send 256 KiB in aggregate; a shared NIC needs an aggregate bandwidth budget. The ratio of cross-fabric events per decode step is 160 to 1. That is the entire argument, and it does not depend on $\alpha$ at all.

The PCIe case is worse and both engines simply refuse to play. A PCIe Gen5 x16 link is roughly 128 GB/s aggregate bidirectional, about 64 GB/s each way, and PCIe GPUs are not a full mesh — traffic hairpins through a root complex. At $S = 128$ KiB the two-shot byte cost becomes 3.6 µs, 7× the 450 GB/s one-way NVLink estimate, and with contention it is worse. vLLM detects this and turns the custom kernel off entirely rather than shipping a slow path:

        fully_connected = False
        if same_node:
            physical_device_id = (
                current_platform.visible_device_id_to_physical_device_id(device.index)
            )
# ...
            fully_connected = current_platform.is_fully_connected(physical_device_ids)
        if same_node and world_size > 2 and not fully_connected:
            logger.warning(
                "Custom allreduce is disabled because it's not supported on"
                " more than two PCIe-only GPUs. To silence this warning, "
                "specify disable_custom_all_reduce=True explicitly."
            )
            return

is_fully_connected is not a heuristic — it asks NVML for the pairwise NVLink P2P status of every pair in the group and demands all of them be one hop:

    @classmethod
    @with_nvml_context
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
        """
        query if the set of gpus are fully connected by nvlink (1 hop)
        """
        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
                            handle,
                            peer_handle,
                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
                        )

vLLM's custom all-reduce, from the CUDA up

The core file is 356 lines and contains two kernels. The one-shot kernel is the whole idea in fifteen lines: barrier, read all $N$ peer pointers, sum, barrier.

template <typename T, int ngpus>
__global__ void __launch_bounds__(512, 1)
    cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg,
                               T* __restrict__ result, int rank, int size) {
  using P = typename packed_t<T>::P;
  using A = typename packed_t<T>::A;
  // note: we don't reorder the address so the accumulation order is the same
  // for all ranks, ensuring bitwise identical results
  auto dp = *_dp;
  barrier_at_start<ngpus>(sg, self_sg, rank);
  // do the actual reduction
  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
       idx += gridDim.x * blockDim.x) {
    ((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
  }
  barrier_at_end<ngpus, true>(sg, self_sg, rank);
}

dp.ptrs[i] is a device pointer into rank i's memory, valid because the ranks exchanged cudaIpcMemHandle_t at setup and each opened the others with cudaIpcOpenMemHandle(..., cudaIpcMemLazyEnablePeerAccess). NCCL is not involved; the load instruction goes straight over NVLink. The comment on line 13 is worth pausing on: because every rank sums the pointers in the same order, every rank produces bit-identical output. NCCL all-reduce also returns the same reduction result on every rank of a successful collective. The separate question is reproducibility across runs, message sizes, algorithms and batch compositions; §5.4.8 discusses that stronger requirement. See the NCCL collective contract.

The barrier is a spin on peer flags — no NCCL, no host involvement, one 32-bit store and one poll per participating GPU per block:

template <int ngpus>
DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
                              int rank) {
  uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
  if (threadIdx.x < ngpus) {
    auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
    auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
    // Write the expected counter value to peer and wait for correct value
    // from peer.
    st_flag_volatile(peer_counter_ptr, flag);
    while (ld_flag_volatile(self_counter_ptr) != flag);
  }
  __syncthreads();
  // use one thread to update flag
  if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
}

The two-shot kernel at csrc/custom_all_reduce.cuh:L30-L75 does the reduce-scatter into a per-rank scratch region (get_tmp_buf, the few megabytes sitting immediately after each Signal), barriers, then all-gathers. Its stage-2 comment names a subtlety that is easy to get wrong: "it's important to match the tid between the two stages, because visibility across devices is only guaranteed between threads that have the same tid."

Algorithm selection is a compile-time macro with the crossovers written as literals:

#define REDUCE_CASE(ngpus)                              \
  case ngpus: {                                         \
    if (force_1stage) {                                 \
      KL(ngpus, cross_device_reduce_1stage);            \
    } else if (force_2stage) {                          \
      KL(ngpus, cross_device_reduce_2stage);            \
    } else {                                            \
      if (world_size_ == 2) {                           \
        KL(ngpus, cross_device_reduce_1stage);          \
      } else if (fully_connected_) {                    \
        if ((world_size_ <= 4 && bytes < 512 * 1024) || \
            (world_size_ <= 8 && bytes < 256 * 1024)) { \
          KL(ngpus, cross_device_reduce_1stage);        \
        } else {                                        \
          KL(ngpus, cross_device_reduce_2stage);        \
        }                                               \
      }                                                 \
    }                                                   \
    break;                                              \
  }

Those are the 512 KiB and 256 KiB constants §5.4.3 inverted. At $N = 2$ two-shot is never used, which the formula predicts: $S^{*} \to \infty$ as $N \to 2$, because $(N{-}1)(N{-}2) = 0$ — with two ranks, one-shot and two-shot move identical bytes and one-shot has fewer barriers.

Three constraints bound the whole design. The grid is capped at 36 blocks (csrc/custom_collective_common.cuh:L34-L39, kMaxBlocks = 36), with the author's note that "too many SMs will cause contention on NVLink bus". The world size must be one of 2, 4, 6, 8, 16. And the message must be small enough to fit the registered staging buffer, whose default is 8 MiB (max_size=8192 * 1024). All three land in the Python gate:

    def should_custom_ar(self, inp: torch.Tensor):
        if self.disabled or self.world_size > 8:
            return False
        inp_size = inp.numel() * inp.element_size()
        # custom allreduce requires input byte size to be multiples of 16
        if inp_size % 16 != 0:
            return False
        if not is_weak_contiguous(inp):
            return False
        # for 4 or more non NVLink-capable GPUs, custom allreduce provides
        # little performance improvement over NCCL.
        if self.world_size == 2 or self.fully_connected:
            return inp_size < self.max_size
        return False

And the dispatch chain in cuda_communicator.py tries five accelerated paths before falling back to NCCL, in a fixed priority order:

        # always try quick reduce first, then flashinfer, then the AITER or vLLM
        # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*)
        qr_comm = self.qr_comm
# ...
        ca_comm = self.ca_comm
        if (
            ca_comm is not None
            and not ca_comm.disabled
            and ca_comm.should_custom_ar(input_)
        ):
            out = ca_comm.custom_all_reduce(input_)
            assert out is not None
            return out
        symm_mem_comm = self.symm_mem_comm
        if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_):
            out = symm_mem_comm.all_reduce(input_)

QuickReduce, first in that list, is ROCm-only and does something the CUDA path does not: it quantises the message inside a two-shot all-reduce, with Q8/Q6/Q4 codecs (csrc/quickreduce/quick_reduce.h:L208-L230). Being a bandwidth optimisation it turns on at the opposite end of the size axis — the bf16 $N=8$ minimum size is 16 MB for the unquantised regime (vllm/distributed/device_communicators/quick_all_reduce.py:L49-L62), well past where decode lives.

SGLang: the same idea, tuned per architecture

SGLang's v1 communicator is a close relative of vLLM's, down to the 8 MiB _MAX_CAR_SIZE = 8192 * 1024 and the same 16-byte alignment gate (python/sglang/srt/distributed/device_communicators/custom_all_reduce.py:L260-L283; the only substantive difference is <= where vLLM has <). The interesting divergence is v2, which replaces the two hard-coded literals with a per-architecture, per-world-size table and adds a third algorithm.

    def _pick_algo(
        self, nbytes: int, can_use_graph: bool
    ) -> Tuple[Optional[AllReduceAlgo], _PullMode]:
        heuristic = self.config.graph if can_use_graph else self.config.eager
        default_mode = _PullMode.GRAPH if can_use_graph else _PullMode.EAGER
        use_multicast = self.config.num_mc_blocks is not None
        if nbytes <= heuristic.one_shot_push_threshold:
            return AllReduceAlgo.ONE_SHOT_PUSH, _PullMode.EAGER
        if nbytes <= heuristic.one_shot_pull_threshold:
            return AllReduceAlgo.ONE_SHOT_PULL, default_mode
        if use_multicast and heuristic.mc.contains(nbytes):
            return AllReduceAlgo.TWO_SHOT_PULL, _PullMode.MULTICAST
        if nbytes <= heuristic.two_shot_pull_threshold:
            return AllReduceAlgo.TWO_SHOT_PULL, default_mode
        return None, _PullMode.EAGER

Push versus pull is the new axis. vLLM's one-shot is a pull: each rank reads seven peer buffers. A push has each rank write its contribution into every peer's workspace and then reduce locally — same byte count, but the write is fire-and-forget while a read stalls on the round trip. Push wins at the smallest sizes; pull wins once the workspace copy starts to matter. And the H100 (SM90) table at TP=8 puts both crossovers at exactly 128 KiB:

        8: config(
            8,
            graph=(128.0 * KB, 128.0 * KB, 32.00 * MB, Range(512 * KB, 128 * MB)),
            eager=(128.0 * KB, 128.0 * KB, 128.0 * MB, Range(0, 128 * MB)),
        ),

Read that as (one-shot-push threshold, one-shot-pull threshold, two-shot-pull threshold, multicast range). Two things stand out. First, the crossover shrinks monotonically with world size in that table — 16 MB at TP=2, 384 KB at TP=4, 128 KB at TP=8 — which is the $S^{*} \propto N/[(N{-}1)(N{-}2)]$ shape the formula predicts. Second, the graph and eager columns differ, but a crossover alone does not identify why. A launch cost common to both paths cancels when their times are equated. CUDA graph replay can change launch and implementation costs; matched timings are needed to separate them.

The fourth entry is a hardware-gated multicast implementation path. Do not generalize that path's Blackwell gate to NVLS as a whole: NCCL documents NVLS from version 2.17 on Hopper with third-generation NVSwitch. NVLink SHARP lets the switch itself perform the reduction, so a rank issues one multicast store and the fabric does the $N$-way sum. That eliminates the $(N{-}1)$ factor from the byte count entirely, but it needs NVLS-capable hardware and is gated by num_mc_blocks, which is None below TP=4.

SGLang carries one genuinely measured crossover in a source comment. It is Blackwell, not Hopper, so it does not transfer — but it is the only real number about these algorithms available from either checkout:

# push wins below ~0.5 MB on B200x8/GB300 (bs<=32 @ H=7168); NVLS 2shot wins
# above (measured: push 8.5us vs pull 11.0us at 448KB, 15.5 vs 11.0 at 896KB)
_PUSH_MAX_BYTES = 512 * 1024

Note the absolute scale: 8.5 µs for a 448 KB all-reduce on eight B200s. Bytes at an assumed 900 GB/s one-way link give $7S/\beta\approx3.6$ µs for the eight-rank one-shot traffic model, not $S/\beta$. The remaining time cannot be attributed uniquely to a barrier or launch without profiling.

SGLang also ships an mscclpp path (--enable-mscclpp, "for small messages"), gated by a tuned config lookup rather than a threshold, and explicitly barred from any piecewise-CUDA-graph phase because switching dispatch mid-compile triggers recompilation (python/sglang/srt/distributed/device_communicators/pymscclpp.py:L328-L353).

Llama-3-70B, TP=8, H100 SXM node, batch 8 decode, inside o_proj's row-parallel forward.

Change one thing — batch 8 to batch 32 — and $S$ becomes 524,288 bytes. Step F still passes (under 8 MiB), but step H now takes the else branch and launches cross_device_reduce_2stage: 917 KB pulled instead of 3.67 MB, at the price of one extra barrier. Change it to batch 512 and $S = 8{,}388{,}608 = $ max_size exactly; inp_size < self.max_size is now false, should_custom_ar returns False, and the message goes to pynccl. Three different code paths across one batch-size sweep, with no log line telling you which one you got.

The obvious optimisation is to put communication on its own stream and let compute proceed. For a decode step, the honest answer is that there is very little to overlap, and both engines' code reflects that.

vLLM's pynccl calls take a stream argument that defaults to current_stream() — the compute stream:

        if stream is None:
            stream = current_stream()
        self.nccl.ncclAllReduce(
            buffer_type(in_tensor.data_ptr()),
            buffer_type(out_tensor.data_ptr()),
            in_tensor.numel(),
            ncclDataTypeEnum.from_torch(in_tensor.dtype),
            ncclRedOpTypeEnum.from_torch(op),
            self.comm,
            cudaStream_t(stream.cuda_stream),
        )

The one place a separate stream appears is graph capture, and it is a capture-mechanics stream, not an overlap stream — the code explicitly joins it to the current stream first (vllm/distributed/parallel_state.py:L653-L660). Inside the replayed graph, the all-reduce is a node between two matmul nodes with a data dependency on both sides. Figure 1 is not a simplification; it is the dependency graph.

So the real work goes into fusion rather than overlap: if you cannot hide the collective, make it cost one launch instead of two, or fold it into an adjacent kernel. SGLang pushes hardest here.

Fusion 1: all-reduce + residual + RMSNorm. The FlashInfer path submits the pattern kARResidualRMSNorm, so the reduction, the residual add and the norm are one kernel with one rendezvous instead of three launches:

    kwargs = dict(
        input=input_tensor,
        workspace=workspace_manager.workspace,
        pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
        launch_with_pdl=True,
        residual_out=residual_out,
        norm_out=norm_out,
        residual_in=residual,
        rms_gamma=weight,
        rms_eps=eps,
        use_oneshot=use_oneshot,
        fp32_acc=fp32_acc,
    )

The layer decides whether to take it in should_fuse_mlp_allreduce_with_next_layer (python/sglang/srt/layers/communicator.py:L809-L866), which is gated on batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE — 2048, at python/sglang/srt/layers/communicator.py:L162-L179 — and refuses under hybrid EP+TP, because the fused kernel reduces over one group and the post-experts reduction spans two disjoint ones, so fusing would "silently return under-reduced activations".

Fusion 2: GEMM + all-reduce in one kernel. This is the more radical one, and it is worth understanding because it removes the barrier from the critical path rather than amortising it. A row-parallel o_proj is a matmul whose output is immediately all-reduced. Those are two kernels with a full dependency between them: the matmul must finish, then the collective launches, barriers, moves bytes, barriers again. But the matmul finishes tile by tile. If the GEMM's epilogue pushes each finished output tile straight into the peers' comm region, the communication of early tiles overlaps the computation of later ones, and the only remaining serialisation is one flag boundary at the end:

"""K3 fused o_proj GEMM + all-reduce for decode (bf16, TP row-parallel).

# ... one line naming the external .cuh elided; see the Unverified note below ...
rank computes the local ``x_r [M, K] @ W_r [7168, K]^T`` partial AND the
cross-rank sum -- the epilogue pushes finished tiles straight into a
peer-mapped P2P comm region, one flag boundary, then a tile-local reduce
writes the fully reduced ``out [M, 7168]`` on every rank. Replaces the
o_proj GEMM + NCCL all-reduce pair with one launch (see GEMM_AR_README.md).

Contracts:

* bf16 only; ``out = sum_r bf16(x_r @ W_r^T)`` (partials round to bf16
  pre-sum -- same numerics as the unfused bf16 GEMM + ring AR).
* M in [1, 512]; internally rounded up to a tuned cell {8, 16, 32, 64,
  128, 256, 512}. ``out`` is allocated with ``cell`` rows and sliced.
# ...

The constraints are exactly what the technique implies. Decode shapes only — $M \le 512$ tokens, because the comm region must be preallocated and the flag ring sized per dispatch cell. Peer-mapped P2P, so full NVLink. SM100+ and $2 \le \mathrm{TP} \le 8$ (python/sglang/srt/layers/k3_gemm_ar.py:L36-L44). And the wrapper is careful about when it initialises: the comm region and JIT compile happen at model build, well before graph capture, because "capture must only see the ready-to-launch path" (python/sglang/srt/layers/k3_gemm_ar.py:L50-L55). It swaps o_proj.forward for a wrapper that falls back to the ordinary GEMM-plus-all-reduce whenever mod.fits(x) is false — prefill, non-2D input, wrong dtype.

vLLM's counterpart is flashinfer_all_reduce.py plus the symmetric-memory path, and it also maintains fused all-gather and reduce-scatter kernels in csrc/custom_all_gather_reduce_scatter.cuh for the DP-attention pattern — cross_device_all_gather is the same barrier/P2P-copy/barrier skeleton with the reduce removed (csrc/custom_all_gather_reduce_scatter.cuh:L13-L31), plus Lamport-flag variants for multi-node NVLink where a barrier is too expensive and the receiver polls for non-sentinel payload instead.

Diagnose rank agreement and transport separately

The symptom is a process that stops producing tokens, pins one SM at 100%, and never times out at the socket layer. SGLang's watchdog will eventually say so:

        pyspy_dump_schedulers()
        logger.error(
            f"{self.debug_name} watchdog timeout "
            f"({self.watchdog_timeout=}, {self.soft=})"
        )
        print(file=sys.stderr, flush=True)
        print(file=sys.stdout, flush=True)

        if not self.soft:
            # Wait for some time so that the parent process can print the error.
            time.sleep(5)
            self.parent_process.send_signal(signal.SIGQUIT)

First test a structural cause before blaming the fabric. A collective is a rendezvous: it completes when every rank has entered it with matching arguments. If one rank enters with a different tensor size, or takes a branch that skips the collective entirely, the others wait forever. The fabric is doing exactly what it was told. SGLang's own debugging notes name the four causes as size mismatch, branch divergence, cascading state drift, and one rank crashing (.claude/skills/debug-distributed-hang/SKILL.md), and record a real diagnosis where a CUDA coredump showed ncclDevKernel_AllGather_RING_LL as the stuck kernel and a py-spy stack pointed at LogitsProcessor.forward → tensor_model_parallel_all_gather — an all-gather size mismatch between TP ranks, not a bad cable.

This is why every fallback predicate in this chapter has to be rank-invariant, and why the code says so out loud:

    Every check here is rank-invariant by construction, and must stay that way:
    a rank that quietly falls back to NCCL while its peers enter the kernel
    mismatches and hangs. The unavailable flag and workspace initialization are
    cross-rank synced at init time (``_sync_allreduce_unavailable_across_tp``);
    the rest are pure functions of the group identity and of tensor metadata,
    which is identical on every rank of the group.

Read should_custom_ar again with that lens. Every clause — byte size, alignment, contiguity, world size, fully_connected — is either identical on all ranks by construction or synchronised at init. Not one of them reads a per-rank runtime value. That is not incidental; it is the invariant that keeps the fallback from deadlocking.

Diagnosis order

Non-determinism, and the price of removing it

Agreement across ranks within one successful all-reduce is distinct from reproducibility across different calls. Both custom and NCCL implementations must produce rank-consistent outputs; changing reduction order, channel count, message size or algorithm can change floating-point rounding across batch sizes or runs. If you need run-to-run determinism, SGLang throws the entire fast path away:

                else:
                    # CUDA: use NCCL tree algorithm
                    os.environ["NCCL_ALGO"] = "allreduce:tree"
                    self.disable_custom_all_reduce = True
                    # should_torch_symm_mem_allreduce() takes the
                    # symmetric-memory path only below a byte threshold, so
                    # which reduce runs would follow the token count.
                    self.enable_torch_symm_mem = False
                    # Each channel carries a differently shaped tree and the
                    # channel count is picked from the message size, so a
                    # token's reduction order would follow the token count.
                    nchannels = str(envs.SGLANG_DETERMINISTIC_NCCL_NCHANNELS.get())
                    os.environ["NCCL_MIN_NCHANNELS"] = nchannels
                    os.environ["NCCL_MAX_NCHANNELS"] = nchannels

That comment is the clearest statement of the problem anywhere in either repo: a token's reduction order would follow the token count. Batch a request differently and you get a different sum. Pinning the channel count and forcing the tree algorithm fixes it, and costs you the custom kernel, the symmetric-memory kernel, and NCCL's own size-adaptive tuning — every $\alpha$ optimisation in this chapter, in exchange for reproducibility. On AMD, SGLang keeps the one-stage kernel instead, precisely because "each GPU reads all data from all GPUs, reduces locally in fixed order".

The silent P2P failure

A driver can report can_device_access_peer == True while P2P does not actually work. vLLM found this the hard way and now runs a real cross-process IPC write before trusting it (vllm/distributed/device_communicators/all_reduce_utils.py:L236-L263, which cites issue #2728). If the check subprocess fails you get:

                raise RuntimeError(
                    f"Error happened when batch testing "
                    f"peer-to-peer access from {batch_src} to {batch_tgt}:\n"
                    f"{returned.stderr.decode()}"
                ) from e

It is cached, so a stale cache after a driver change is a real failure mode. VLLM_SKIP_P2P_CHECK=1 bypasses the test — and, on a broken driver, hands you silently wrong numerics instead of a warning.

1. Prove the fabric works before blaming the engine. vLLM ships the script; run it at your TP size. It exercises PyTorch NCCL, GLOO, pynccl, and pynccl inside a CUDA graph — four distinct failure surfaces:

NCCL_DEBUG=TRACE torchrun --nproc-per-node=<number-of-GPUs> test.py

2. Measure what the custom kernel is worth on your box. Both engines expose one flag, so you can A/B the whole mechanism:

vllm bench latency --model meta-llama/Meta-Llama-3-8B --tensor-parallel-size 8 \
    --input-len 128 --output-len 128 --batch-size 8
vllm bench latency --model meta-llama/Meta-Llama-3-8B --tensor-parallel-size 8 \
    --input-len 128 --output-len 128 --batch-size 8 --disable-custom-all-reduce

python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B --tp 8
python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B --tp 8 --disable-custom-all-reduce

The table in §5.4.3 predicts the gap is largest at small batch and small model, and shrinks toward zero as $t$ grows past a few hundred. If your measured gap does not have that shape, your $\alpha$ is not what this chapter assumed — and you now have a real measurement, which is worth more than the model.

3. Force each algorithm and find your own crossover. vLLM reads an environment variable inside the kernel dispatcher, so you can pin the algorithm without recompiling (csrc/custom_all_reduce.cuh:L270-L285):

VLLM_CUSTOM_ALLREDUCE_ALGO=oneshot vllm bench latency ...   # also: 1stage, 2stage, twoshot

Sweep --batch-size from 1 to 512 under each setting. The batch size where the two curves cross, times $d \times 2$ bytes, is your machine's $S^{*}$. Compare it with vLLM's hard-coded 256 KiB and with SGLang's 128 KiB for SM90 TP=8.

• vLLM PR #2192, "Implement custom all reduce kernels" — the commit that created csrc/custom_all_reduce.cuh. Its follow-ups #2642 (no repeated IPC open) and #8558 (unsafe synchronisation) are the best record of how the barrier reached its current form.
• SGLang PR #19880 — "Support JIT custom all reduce (rewrite as v2)", where the push/pull/multicast split and the per-architecture threshold table landed. #32541 is where the fused GEMM+all-reduce arrived.
• vLLM issue #2728 — why can_device_access_peer cannot be trusted, and the origin of the can_actually_p2p subprocess test. Read it before you set VLLM_SKIP_P2P_CHECK.
• The NCCL user guide, and in particular its Environment Variables page — the authoritative list. NCCL_ALGO, NCCL_MIN_NCHANNELS and NCCL_DEBUG_SUBSYS are the three that matter for the material in this chapter.
• Massively Scale Your Deep Learning Training with NCCL 2.4 — NVIDIA's own account of why the double binary tree beats ring at small message sizes. The argument is the latency-vs-bandwidth split of this chapter, made for a different workload.
• NVIDIA H100 product page and the NVLink and NVSwitch page — the source of the 900 GB/s bidirectional per-GPU figure used throughout.
• python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py — the most useful single artefact in either repo for this topic: a per-architecture, per-world-size crossover table someone actually measured, with the SM90 and SM100 numbers side by side.
• .claude/skills/debug-distributed-hang/SKILL.md in the SGLang tree — a genuinely good short methodology for collective hangs, including the CUDA user-triggered coredump recipe and the per-rank logging pattern for binary-searching a divergence point.