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