From 130ff00e6984b927852cb0f70dfbbb91a6015360 Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:30:01 +0200 Subject: [PATCH] [Kernel] Two-rank FP16 all-reduce through pinned host memory for GPUs without P2P For two-rank tensor-parallel groups whose GPUs NCCL cannot connect peer to peer (P2P disabled or unavailable), one kernel exchanges the FP16 payload through a pinned host buffer both ranks map: each block writes its chunk, fences, raises a flag, waits for the peer's flag and adds the peer's chunk. One FP16 rounding and commutative, so the result has NCCL's bits. No CPU work, capturable in CUDA graphs. It sits after the SM70 ring in the all-reduce chain and takes FP16 payloads up to kernel_config.host_reduce.max_bytes (256 KiB); larger ones stay on NCCL. kernel_config.host_reduce.enabled turns it off. The exchange kernel is adapted from kernels/skinny_ar.cu of mzen17/v100-skinny-unify (MIT, notice kept in the source), turned into handle-based channels so a process can own more than one two-rank group. Co-authored-by: Mike Zeng <92705460+mzen17@users.noreply.github.com> Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014ZLMmGB8mjc43VpwPxLQNn (cherry picked from commit f53a668ebd02d28600b778a46005bec2addd4b8f) --- CMakeLists.txt | 4 +- csrc/host_reduce.cu | 271 ++++++++++++++++++ docs/design/architecture/runtime_reference.md | 3 +- tests/distributed/test_host_reduce.py | 74 +++++ vllm/config/kernel.py | 19 ++ .../device_communicators/cuda_communicator.py | 21 ++ .../device_communicators/host_reduce.py | 120 ++++++++ 7 files changed, 509 insertions(+), 3 deletions(-) create mode 100644 csrc/host_reduce.cu create mode 100644 tests/distributed/test_host_reduce.py create mode 100644 vllm/distributed/device_communicators/host_reduce.py diff --git a/CMakeLists.txt b/CMakeLists.txt index f69b44c2be..67452129bf 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -367,9 +367,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") set(VLLM_SM70_TURBOMIND_INCLUDE_DIRS "") cuda_archs_loose_intersection(SM70_TURBOMIND_ARCHS "7.0" "${CUDA_ARCHS}") if (SM70_TURBOMIND_ARCHS) - list(APPEND VLLM_EXT_SRC "csrc/sm70_ring_collectives.cu") + list(APPEND VLLM_EXT_SRC "csrc/sm70_ring_collectives.cu" "csrc/host_reduce.cu") set_gencode_flags_for_srcs( - SRCS "csrc/sm70_ring_collectives.cu" + SRCS "csrc/sm70_ring_collectives.cu" "csrc/host_reduce.cu" CUDA_ARCHS "${SM70_TURBOMIND_ARCHS}") set(VLLM_SM70_TURBOMIND_ROOT "${CMAKE_CURRENT_LIST_DIR}/csrc/sm70_turbomind") diff --git a/csrc/host_reduce.cu b/csrc/host_reduce.cu new file mode 100644 index 0000000000..0d84c58b88 --- /dev/null +++ b/csrc/host_reduce.cu @@ -0,0 +1,271 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Two-rank FP16 all-reduce through pinned host memory, for GPU pairs without +// peer access (PCIe boxes with P2P disabled, KVM guests, mismatched boards). +// NCCL's LL protocol costs tens of microseconds for the small decode +// all-reduces there, and the custom all-reduce needs P2P/IPC. One kernel does +// the whole exchange: each block writes its chunk to a host slot, fences, +// raises a per-block flag, spins on the peer's flag, then reads the peer's +// chunk straight from host memory and adds it. No CPU work and no stream +// synchronization, so it can be captured in CUDA graphs (the call counter +// lives in device memory and the last block advances it). +// +// Numerics: out = mine + peer in FP16, one rounding and commutative, so both +// ranks produce identical bits, the same as NCCL's two-rank FP16 sum. +// +// Slot reuse: double buffering by call parity. Before call n+2 overwrites +// slot (n & 1), this rank has observed the peer's flag for call n+1, which +// the peer raises only after it finished call n on its own stream. +// +// The exchange kernel is adapted from kernels/skinny_ar.cu of +// mzen17/v100-skinny-unify (https://github.com/mzen17/v100-skinny-unify), +// which carries this notice: +// +// MIT License. Copyright (c) 2026 v100-skinny contributors +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the +// "Software"), to deal in the Software without restriction, including +// without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit +// persons to whom the Software is furnished to do so, subject to the +// following conditions: +// +// The above copyright notice and this permission notice shall be included +// in all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN +// NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE +// USE OR OTHER DEALINGS IN THE SOFTWARE. +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kMaxBlocks = 16; +constexpr int kMaxChunks = 64; +constexpr int kThreads = 256; +constexpr int kFlagStride = 32; // one 128-byte cache line per flag +// flags: [parity 2][rank 2][kMaxChunks] * kFlagStride uints +constexpr size_t kFlagsBytes = + 2ull * 2 * kMaxChunks * kFlagStride * sizeof(unsigned); + +struct Channel { + void* host = nullptr; + void* dev = nullptr; + size_t map_bytes = 0; + size_t slot_bytes = 0; + unsigned* counters = nullptr; // [0] call sequence, [1] finished blocks + int rank = -1; + int device = -1; +}; + +std::mutex channels_mutex; +std::vector> channels; + +Channel& channel(int64_t handle) { + std::lock_guard lock(channels_mutex); + TORCH_CHECK(handle >= 0 && handle < static_cast(channels.size()) && + channels[handle], + "host_reduce: unknown handle ", handle); + return *channels[handle]; +} + +__device__ __forceinline__ int4 load_volatile(const int4* p) { + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} + +__device__ __forceinline__ int4 add_half8(int4 a, int4 b) { + const half2* pa = reinterpret_cast(&a); + const half2* pb = reinterpret_cast(&b); + int4 o; + half2* po = reinterpret_cast(&o); +#pragma unroll + for (int i = 0; i < 4; i++) po[i] = __hadd2(pa[i], pb[i]); + return o; +} + +// x -> out over n8 groups of eight halves. +__global__ void __launch_bounds__(kThreads) + host_reduce_kernel(const int4* __restrict__ x, int4* __restrict__ out, + int n8, unsigned* flags, int4* data, int rank, + int slot_int4, unsigned* counters) { + __shared__ unsigned s_seq; + if (threadIdx.x == 0) { + s_seq = *reinterpret_cast(counters) + 1u; + } + __syncthreads(); + const unsigned seq = s_seq; + const int parity = seq & 1u; + const int peer = rank ^ 1; + int4* my_slot = data + static_cast(parity * 2 + rank) * slot_int4; + const int4* peer_slot = + data + static_cast(parity * 2 + peer) * slot_int4; + volatile unsigned* my_flag = + flags + + (static_cast(parity * 2 + rank) * kMaxChunks + blockIdx.x) * + kFlagStride; + volatile unsigned* peer_flag = + flags + + (static_cast(parity * 2 + peer) * kMaxChunks + blockIdx.x) * + kFlagStride; + + const int per_block = (n8 + gridDim.x - 1) / gridDim.x; + const int begin = blockIdx.x * per_block; + const int end = min(n8, begin + per_block); + + // 1. Publish this rank's chunk (posted PCIe writes), fence, raise the flag. + for (int i = begin + threadIdx.x; i < end; i += kThreads) my_slot[i] = x[i]; + __threadfence_system(); + __syncthreads(); + if (threadIdx.x == 0) { + *my_flag = seq; + // 2. Wait for the peer's chunk. + while (*peer_flag != seq) { + } + } + __syncthreads(); + // 3. Read the peer's chunk from host memory and add. + for (int i = begin + threadIdx.x; i < end; i += kThreads) { + out[i] = add_half8(x[i], load_volatile(peer_slot + i)); + } + // 4. The last block to finish advances the call counter. + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) { + const unsigned finished = atomicAdd(counters + 1, 1u); + if (finished == gridDim.x - 1) { + counters[1] = 0; + __threadfence(); + *reinterpret_cast(counters) = seq; + } + } +} + +// Map the shared file both ranks opened, clear this rank's flags and slots so +// a stale mapping cannot satisfy a wait, and register it with CUDA. The caller +// barriers on both ranks before the first all-reduce. +int64_t open_channel(const std::string& path, int64_t rank, int64_t max_bytes) { + TORCH_CHECK(rank == 0 || rank == 1, "host_reduce supports exactly two ranks"); + TORCH_CHECK(max_bytes > 0 && max_bytes % 16 == 0 && max_bytes <= (1 << 22), + "host_reduce: invalid max_bytes ", max_bytes); + auto ch = std::make_unique(); + ch->rank = static_cast(rank); + ch->slot_bytes = static_cast(max_bytes); + const size_t page = static_cast(sysconf(_SC_PAGESIZE)); + ch->map_bytes = (kFlagsBytes + 4 * ch->slot_bytes + page - 1) / page * page; + const int fd = open(path.c_str(), O_RDWR | O_CREAT, 0600); + TORCH_CHECK(fd >= 0, "host_reduce: cannot open ", path, ": ", + std::strerror(errno)); + const bool sized = ftruncate(fd, static_cast(ch->map_bytes)) == 0; + void* host = sized ? mmap(nullptr, ch->map_bytes, PROT_READ | PROT_WRITE, + MAP_SHARED, fd, 0) + : MAP_FAILED; + const int error = errno; + close(fd); + TORCH_CHECK(sized, "host_reduce: ftruncate failed: ", std::strerror(error)); + TORCH_CHECK(host != MAP_FAILED, + "host_reduce: mmap failed: ", std::strerror(error)); + ch->host = host; + unsigned* flags = static_cast(host); + for (int parity = 0; parity < 2; parity++) { + std::memset( + flags + + static_cast(parity * 2 + rank) * kMaxChunks * kFlagStride, + 0, static_cast(kMaxChunks) * kFlagStride * sizeof(unsigned)); + std::memset(static_cast(host) + kFlagsBytes + + static_cast(parity * 2 + rank) * ch->slot_bytes, + 0, ch->slot_bytes); + } + C10_CUDA_CHECK(cudaGetDevice(&ch->device)); + C10_CUDA_CHECK(cudaHostRegister(host, ch->map_bytes, cudaHostRegisterMapped)); + C10_CUDA_CHECK(cudaHostGetDevicePointer(&ch->dev, host, 0)); + C10_CUDA_CHECK(cudaMalloc(&ch->counters, 2 * sizeof(unsigned))); + C10_CUDA_CHECK(cudaMemset(ch->counters, 0, 2 * sizeof(unsigned))); + std::lock_guard lock(channels_mutex); + channels.push_back(std::move(ch)); + return static_cast(channels.size()) - 1; +} + +void close_channel(int64_t handle) { + std::unique_ptr ch; + { + std::lock_guard lock(channels_mutex); + TORCH_CHECK(handle >= 0 && handle < static_cast(channels.size()), + "host_reduce: unknown handle ", handle); + ch = std::move(channels[handle]); + } + if (!ch) return; + const c10::cuda::CUDAGuard guard(ch->device); + C10_CUDA_CHECK(cudaDeviceSynchronize()); + C10_CUDA_CHECK(cudaHostUnregister(ch->host)); + munmap(ch->host, ch->map_bytes); + C10_CUDA_CHECK(cudaFree(ch->counters)); +} + +void host_reduce_out(torch::Tensor& output, const torch::Tensor& input, + int64_t handle) { + Channel& ch = channel(handle); + TORCH_CHECK(input.is_cuda() && input.scalar_type() == at::kHalf && + input.is_contiguous(), + "host_reduce: contiguous FP16 CUDA input required"); + TORCH_CHECK(output.is_contiguous() && output.sizes() == input.sizes() && + output.scalar_type() == at::kHalf && + output.device() == input.device(), + "host_reduce: output must match input"); + const int64_t bytes = input.numel() * 2; + TORCH_CHECK(bytes > 0 && bytes % 16 == 0 && + bytes <= static_cast(ch.slot_bytes), + "host_reduce: payload of ", bytes, " bytes outside the channel"); + TORCH_CHECK((reinterpret_cast(input.data_ptr()) & 15) == 0 && + (reinterpret_cast(output.data_ptr()) & 15) == 0, + "host_reduce: 16-byte alignment required"); + const c10::cuda::CUDAGuard guard(input.device()); + const int n8 = static_cast(bytes / 16); + const int blocks = + std::max(1, std::min(kMaxBlocks, (n8 + kThreads - 1) / kThreads)); + unsigned* flags = static_cast(ch.dev); + int4* data = + reinterpret_cast(static_cast(ch.dev) + kFlagsBytes); + host_reduce_kernel<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(output.data_ptr()), n8, flags, data, ch.rank, + static_cast(ch.slot_bytes / 16), ch.counters); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(_C, m) { + m.def("host_reduce_open(str path, int rank, int max_bytes) -> int", + &open_channel); + m.def("host_reduce_close(int handle) -> ()", &close_channel); + m.def( + "host_reduce_out(Tensor(a!) output, Tensor input, int handle) -> " + "()"); +} +TORCH_LIBRARY_IMPL(_C, CUDA, m) { m.impl("host_reduce_out", &host_reduce_out); } diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index 66c5bc04ff..0524024cb9 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -37,8 +37,9 @@ identify initialization/hash code; they do not evaluate its semantics. | `vllm.config.gdn_schedule.GdnScheduleConfig` | kernel_config.gdn.schedule | resolve: [vllm/config/gdn_schedule.py](../../../vllm/config/gdn_schedule.py#L129); graph_options: [vllm/config/gdn_schedule.py](../../../vllm/config/gdn_schedule.py#L160) | [vllm/config/gdn_schedule.py](../../../vllm/config/gdn_schedule.py#L49) | | `vllm.config.gdn_state.GdnStateConfig` | kernel_config.gdn.state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L39); graph_options: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L52) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L23) | | `vllm.config.gdn_state.GdnStateTraceConfig` | observability_config.gdn_state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L99); compile_ignored_aliases: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L142) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L59) | +| `vllm.config.kernel.HostReduceConfig` | kernel_config.host_reduce | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) | | `vllm.config.kernel.IrOpPriorityConfig` | kernel_config.ir_op_priority | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L84) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L67) | -| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1044) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) | +| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1061) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L632) | | `vllm.config.kernel.Sm70AwqConfig` | kernel_config.sm70_awq | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L338); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L317) | | `vllm.config.kernel.Sm70Fp8Config` | kernel_config.sm70_fp8 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L433); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L369) | | `vllm.config.kernel.Sm70GgufConfig` | kernel_config.sm70_gguf | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L498) | diff --git a/tests/distributed/test_host_reduce.py b/tests/distributed/test_host_reduce.py new file mode 100644 index 0000000000..82a86dd73c --- /dev/null +++ b/tests/distributed/test_host_reduce.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The host-memory two-rank all-reduce gives NCCL's bits, eagerly and in a graph.""" + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from vllm.distributed.device_communicators.host_reduce import ( + HostReduceCommunicator, +) + +SIZES = [8, 4096, 5120, 3 * 5120, 8 * 5120, 131072] + + +def _worker(rank: int, port: int, failures) -> None: + os.environ["NCCL_P2P_DISABLE"] = "1" + torch.accelerator.set_device_index(rank) + device = torch.device("cuda", rank) + init = f"tcp://127.0.0.1:{port}" + dist.init_process_group("nccl", init_method=init, rank=rank, world_size=2) + cpu_group = dist.new_group(backend="gloo") + comm = HostReduceCommunicator(cpu_group, device, "tp:0") + try: + assert comm.status["enabled"], comm.status + generator = torch.Generator(device=device).manual_seed(1234 + rank) + for size in SIZES: + for _ in range(3): + x = torch.randn(size, generator=generator, device=device) + x = x.half() + expected = x.clone() + dist.all_reduce(expected) + got = comm.all_reduce(x) + assert got is not None + assert torch.equal(got, expected), size + static = torch.randn(5120, generator=generator, device=device).half() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = comm.all_reduce(static) + for _ in range(4): + static.copy_(torch.randn(5120, device=device).half()) + expected = static.clone() + dist.all_reduce(expected) + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(graph_out, expected) + assert comm.all_reduce(torch.ones(3, device=device).half()) is None + assert comm.all_reduce(torch.ones(8, device=device)) is None + except Exception as exc: # noqa: BLE001 - reported to the parent + failures.put(f"rank {rank}: {exc!r}") + raise + finally: + comm.close() + dist.destroy_process_group() + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 2 + or not hasattr(torch.ops._C, "host_reduce_open"), + reason="needs two GPUs and the native host all-reduce", +) +def test_host_reduce_matches_nccl(): + import socket + + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + ctx = mp.get_context("spawn") + failures = ctx.Queue() + mp.spawn(_worker, args=(port, failures), nprocs=2, join=True) + assert failures.empty() diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index f83e58c965..afc805b0f7 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -615,6 +615,19 @@ class Sm70RingConfig: """Largest calibrated input payload; larger messages retain NCCL.""" +@config +class HostReduceConfig: + """Two-rank FP16 all-reduce through pinned host memory for GPU pairs without + peer access; exact, the same bits as NCCL's two-rank sum.""" + + enabled: bool = True + """Admit the operator for two-rank groups whose GPUs NCCL cannot connect + peer to peer.""" + + max_bytes: int = Field(default=262144, gt=0, le=4194304, multiple_of=16) + """Largest payload; larger all-reduces stay on NCCL.""" + + @config class KernelConfig: """Configuration for kernel selection and warmup behavior.""" @@ -741,6 +754,10 @@ def resolve_sm70_rmsnorm_gated(self, *, qualified: bool) -> None: """Keep only losslessly packed TP4 HC shards; large M uses gathered GEMM.""" sm70_router_weight_storage: Literal["dual", "row_major"] = "dual" """Read the FP16 checkpoint router directly or retain its packed duplicate.""" + + host_reduce: HostReduceConfig = Field(default_factory=HostReduceConfig) + """Two-rank host-memory all-reduce policy, resolved from peer access.""" + collective_kernel_selections: dict[str, Any] = Field( default_factory=dict, init=False ) @@ -1073,6 +1090,8 @@ def compute_hash(self) -> str: "ple_disk_row_readers", "qsa_auto_e4m3_reason", "qsa_host_kv_reason", + # Collective transport chosen outside the compiled graph. + "host_reduce", } if not self.sm70_skinny_moe_applicable: ignored_factors.add("sm70_skinny_moe") diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 1c02b54d3f..d15f79664c 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -196,6 +196,13 @@ def __init__( self.hc_ll_comm = Sm70HcLLCommunicator( self.cpu_group, self.device, self.unique_name ) + self.host_reduce_comm = None + if "tp" in self.unique_name and self.world_size == 2: + from .host_reduce import HostReduceCommunicator + + self.host_reduce_comm = HostReduceCommunicator( + self.cpu_group, self.device, self.unique_name + ) if self.world_size > 1: self._log_all_reduce_backend_selection() @@ -274,6 +281,7 @@ def _log_all_reduce_backend_selection(self) -> None: all_potential_ar_backends = [ "NCCL_SYMM_MEM", "SM70_RING", + "HOST_REDUCE", "QUICK_REDUCE", "FLASHINFER", "CUSTOM", @@ -299,6 +307,11 @@ def _log_all_reduce_backend_selection(self) -> None: ) if self.ring_comm.status["enabled"]: enabled_ar_backends.append("SM70_RING") + if ( + self.host_reduce_comm is not None + and self.host_reduce_comm.status["enabled"] + ): + enabled_ar_backends.append("HOST_REDUCE") if ( self.pynccl_comm is not None and not self.pynccl_comm.disabled @@ -336,6 +349,11 @@ def all_reduce(self, input_): if ring_out is not None: self._collective_trace.record("sm70_ring", input_) return ring_out + if self.host_reduce_comm is not None: + host_out = self.host_reduce_comm.all_reduce(input_) + if host_out is not None: + self._collective_trace.record("host_reduce", input_) + return host_out # since currently we perform copy input -> symm_input -> out-of-place AR # return symm_output, we don't need to check if input is symmetric if self.pynccl_comm is not None and should_nccl_symm_mem_allreduce( @@ -654,6 +672,9 @@ def destroy(self): ring_comm = getattr(self, "ring_comm", None) if ring_comm is not None: ring_comm.close() + host_reduce_comm = getattr(self, "host_reduce_comm", None) + if host_reduce_comm is not None: + host_reduce_comm.close() if self.pynccl_comm is not None: self.pynccl_comm.destroy() self.pynccl_comm = None diff --git a/vllm/distributed/device_communicators/host_reduce.py b/vllm/distributed/device_communicators/host_reduce.py new file mode 100644 index 0000000000..d4f64e9efc --- /dev/null +++ b/vllm/distributed/device_communicators/host_reduce.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Two-rank FP16 all-reduce through pinned host memory (csrc/host_reduce.cu) +for GPU pairs that NCCL cannot connect peer to peer.""" + +import os +import uuid +from typing import Any + +import torch +import torch.distributed as dist + +from vllm.config import get_current_vllm_config_or_none +from vllm.config.kernel import HostReduceConfig +from vllm.platforms import current_platform + +_OPS = ("host_reduce_open", "host_reduce_close", "host_reduce_out") + + +def _nccl_peer_access(group, device: torch.device) -> bool: + """Whether NCCL may connect the two ranks peer to peer; NCCL_P2P_DISABLE + turns it off regardless of the hardware.""" + if os.environ.get("NCCL_P2P_DISABLE", "0") not in ("", "0"): + return False + local = device.index + if local is None: + local = torch.accelerator.current_device_index() + uuids = [""] * 2 + dist.all_gather_object(uuids, current_platform.get_device_uuid(local), group=group) + visible = { + current_platform.get_device_uuid(i): i + for i in range(torch.accelerator.device_count()) + } + peer = visible.get(uuids[1 - dist.get_rank(group)]) + # A peer outside this worker's visible devices cannot be reached directly. + return peer is not None and torch.cuda.can_device_access_peer(local, peer) + + +class HostReduceCommunicator: + def __init__(self, group, device: torch.device, name: str): + self.group = group + self.device = device + self.rank = dist.get_rank(group) + self.handle: int | None = None + self.status: dict[str, Any] = { + "enabled": False, + "reason": None, + "scope": "collective_capability", + } + cfg = get_current_vllm_config_or_none() + self.policy = cfg.kernel_config.host_reduce if cfg else HostReduceConfig() + reason = None + if not self.policy.enabled: + reason = "disabled_by_configuration" + elif dist.get_world_size(group) != 2: + reason = "requires_two_ranks" + elif not current_platform.is_cuda(): + reason = "requires_cuda" + elif missing := [op for op in _OPS if not hasattr(torch.ops._C, op)]: + reason = "operator_missing:" + ",".join(missing) + if dist.get_world_size(group) == 2: + reasons: list[str | None] = [None] * 2 + dist.all_gather_object(reasons, reason, group=group) + reason = next((r for r in reasons if r is not None), None) + if reason is None: + access = [False] * 2 + dist.all_gather_object( + access, _nccl_peer_access(group, device), group=group + ) + if any(access): + reason = "peer_access_available" + if reason is None: + self._open() + self.status.update( + enabled=reason is None, reason=reason, max_bytes=self.policy.max_bytes + ) + if cfg: + cfg.kernel_config.collective_kernel_selections[name] = self.status + + def _open(self) -> None: + # Both ranks map one shared file; it is unlinked once both hold it, so + # nothing outlives the processes. + paths = [f"/dev/shm/vllm_host_reduce_{uuid.uuid4().hex}"] + dist.broadcast_object_list( + paths, src=dist.get_global_rank(self.group, 0), group=self.group + ) + with torch.accelerator.device_index(self.device.index): + self.handle = torch.ops._C.host_reduce_open( + paths[0], self.rank, self.policy.max_bytes + ) + dist.barrier(group=self.group) + if self.rank == 0: + os.unlink(paths[0]) + self.status.update(dtype="float16", transport="pinned_host_memory") + + def rejection_reason(self, tensor: torch.Tensor) -> str | None: + if not self.status["enabled"]: + return self.status["reason"] + if tensor.dtype != torch.float16: + return "fp16_input_required" + if tensor.device != self.device or not tensor.is_contiguous(): + return "requires_contiguous_local_cuda_input" + nbytes = tensor.numel() * tensor.element_size() + if not 0 < nbytes <= self.policy.max_bytes or nbytes % 16: + return "payload_outside_channel" + if tensor.data_ptr() % 16: + return "requires_16_byte_alignment" + return None + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor | None: + if self.rejection_reason(tensor) is not None: + return None + output = torch.empty_like(tensor) + torch.ops._C.host_reduce_out(output, tensor, self.handle) + return output + + def close(self) -> None: + if self.handle is not None: + torch.ops._C.host_reduce_close(self.handle) + self.handle = None