[Distributed] Route small TP all-gathers through custom_all_gather

[pst2154]

Overview

Opt-in VLLM_CUSTOM_ALL_GATHER_SMALL=1 (default off): CudaCommunicator.all_gather() sends TP all-gathers of at most VLLM_CUSTOM_ALL_GATHER_SMALL_MAX_BYTES per rank (default 64 KiB) to the existing CustomAllreduce.custom_all_gather, falling back to NCCL when it declines. Reworked after review: no new kernel or CustomAllreduce method, only routing. Related: #57856.

Claims

• Small, latency-bound gathers (e.g. the speculative drafter's per-step gathers) use the custom all-gather instead of an NCCL kernel. Today custom_all_gather is only reached from sp_all_gather; tensor_model_parallel_all_gather always uses NCCL.
• Microbenchmark, TP4 on 4x B200, CUDA graph, 64 B to 96 KiB per rank: about 2.7x to 3.5x lower latency per gather than NCCL (see Validation).
• Same values and layout as the NCCL path for any dim, except that on multicast-capable GPUs a -0.0 comes back as +0.0 (see Details).
• Flag off: behaviour unchanged.

Validation

On ac9126e58 (nightly image with this branch's vllm/ files overlaid):

• pytest -v tests/distributed/test_custom_all_gather_small.py (4x B200, Ray): 8 passed. Routing and layout for dim 0 / 1 / -1, threshold and decline fallbacks (unit); TP2 / TP4 comparison against the NCCL path on drafter-like shapes (bf16 [1..48, 1024], [1, 32768], fp32 [1, 16], [8, 16], [8, 2], and one size above the threshold), with NaN / inf / subnormal / signed-zero inputs, eagerly and over 20 CUDA graph replays. Values are compared bitwise with signed zeros canonicalized.
• Microbenchmark (TP4, 4x B200 with NVLink multicast, 50 back-to-back gathers captured in a CUDA graph, median per gather, relative to NCCL = 1.00):

per-rank input	NCCL	custom_all_gather (Lamport on this node)
fp32 [1, 16] (64 B)	1.00	0.35
fp32 [8, 16] (512 B)	1.00	0.37
bf16 [1, 1024] (2 KiB)	1.00	0.36
bf16 [8, 1024] (16 KiB)	1.00	0.31
bf16 [1, 32768] (64 KiB)	1.00	0.28
bf16 [48, 1024] (96 KiB)	1.00	0.29

No end-to-end measurement of this change alone on this commit.

Details

• Routing: after the existing NVLS symmetric-memory branch and before pynccl: 0 < nbytes <= cap and custom_all_gather() returns a tensor; otherwise NCCL. The dim-0 result gets the same movedim / reshape as the NCCL path (shared helper). The cap is read once at init.
• Kernel choice is custom_all_gather's: the MNNVL Lamport kernel when multicast is available, else the one-shot CUDA-IPC kernel. Both use buffers registered at init, so they are CUDA-graph safe. Inputs that custom_all_gather does not take (not a 16-byte multiple such as a single [1, 2] fp32 row, non-float dtypes, above its size limit) fall back to NCCL.
• Signed zero: the Lamport kernel reserves the 32-bit word 0x80000000 as its "not written yet" marker and stores such words as 0, so -0.0 arrives as +0.0 (and in bf16/fp16 a word holding -0.0 with a zero neighbour). All other bit patterns are preserved. For the drafter's gathers (hidden states, logits, argmax candidates) this does not change any result.
• Open questions: env var plus cap vs a default-on size-based policy; the 64 KiB default covers the drafter's gathers, larger sizes were not measured.
• Relation to [Perf] Use custom all-gather for full-vocabulary logits #57856 (open): it routes the full-vocabulary logits gather in LogitsProcessor to custom_all_gather(). This PR routes small gathers at the communicator level; the call sites are complementary.

AI assistance: Claude Code was used to write this change and its tests, run the validation and draft this description. The human submitter reviews every changed line before this leaves draft.

---

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

[pst2154]

Fair question, and you're right. I'd avoided custom_all_gather because on multicast-capable GPUs it uses the MNNVL Lamport kernel, which reserves 0x80000000 as its empty marker, so -0.0 comes back as +0.0. I wanted a pure byte copy, but values are otherwise identical and the drafter's gathers don't care. It's also much faster: on 4x B200 under CUDA graph, 3.6-4.2 µs vs 10.4-14.5 µs for NCCL and 12.4-13.6 µs for the one-shot IPC kernel this PR added, at 64 B-96 KiB per rank.

So I've reduced the PR locally to routing only: CudaCommunicator.all_gather() sends inputs up to the size limit to the existing custom_all_gather and falls back to NCCL when it declines, logging with logger.info_once. The new CustomAllreduce method is removed. Updated branch coming shortly.