From 861bc77230ac989e4c74d4aa8acf3a32d5cccbc0 Mon Sep 17 00:00:00 2001 From: Yanyuan Qin Date: Tue, 28 Jul 2026 21:01:08 +0000 Subject: [PATCH] perf(custom-ar): add graph-safe dual-input all-reduce --- .../device_communicators/custom_all_reduce.py | 77 +++++++- aiter/ops/custom_all_reduce.py | 12 ++ csrc/include/custom_all_reduce.cuh | 185 ++++++++++++++++++ csrc/include/custom_all_reduce.h | 7 + csrc/include/rocm_ops.hpp | 9 + csrc/kernels/custom_all_reduce.cu | 103 ++++++++++ op_tests/test_custom_all_reduce_dual.py | 120 ++++++++++++ 7 files changed, 509 insertions(+), 4 deletions(-) create mode 100644 op_tests/test_custom_all_reduce_dual.py diff --git a/aiter/dist/device_communicators/custom_all_reduce.py b/aiter/dist/device_communicators/custom_all_reduce.py index e225164171..d2c6e7a4b6 100644 --- a/aiter/dist/device_communicators/custom_all_reduce.py +++ b/aiter/dist/device_communicators/custom_all_reduce.py @@ -642,7 +642,6 @@ def get_external_ipc_meta(self, tensor): class CustomAllreduce: - _SUPPORTED_WORLD_SIZES: ClassVar[list[Any]] = [2, 4, 6, 8] def _select_ops(self): @@ -664,6 +663,7 @@ def _select_ops(self): # kernel ops (arch) self._ops_meta_size = ops.meta_size_gfx1250 self._ops_all_reduce = ops.all_reduce_gfx1250 + self._ops_all_reduce_dual = None self._ops_all_gather = ops.all_gather_gfx1250 self._ops_reduce_scatter = ops.reduce_scatter_gfx1250 self._ops_dispose = ops.dispose_gfx1250 @@ -685,6 +685,7 @@ def _select_ops(self): self._ops_meta_size = ops.meta_size self._ops_init_custom_ar = ops.init_custom_ar self._ops_all_reduce = ops.all_reduce + self._ops_all_reduce_dual = ops.all_reduce_dual self._ops_all_gather = None self._ops_reduce_scatter = ops.reduce_scatter self._ops_dispose = ops.dispose @@ -725,9 +726,9 @@ def __init__( self.group = group - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), "CustomAllreduce should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "CustomAllreduce should be attached to a non-NCCL group." + ) if not all(in_the_same_node_as(group, source_rank=0)): # No need to initialize custom allreduce for multi-node case. @@ -1062,6 +1063,32 @@ def should_custom_ar(self, inp: torch.Tensor, prefill_support: bool = False): inp ) + def should_custom_ar_dual( + self, + left: torch.Tensor, + right: torch.Tensor, + ) -> bool: + """Return whether two registered inputs fit the one-stage dual kernel.""" + if self.disabled or self._ops_all_reduce_dual is None: + return False + if left.numel() == 0 or right.numel() == 0: + return False + if left.device != right.device or left.dtype != right.dtype: + return False + if left.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return False + if not left.is_contiguous() or not right.is_contiguous(): + return False + left_bytes = left.numel() * left.element_size() + right_bytes = right.numel() * right.element_size() + if left_bytes % 16 != 0 or right_bytes % 16 != 0: + return False + total_bytes = left_bytes + right_bytes + if not self._car_min_size < total_bytes <= self._car_max_size: + return False + one_stage_limit = 160 * 1024 if self.world_size <= 4 else 80 * 1024 + return self.fully_connected and total_bytes < one_stage_limit + def should_custom_ar_bytes(self, inp: torch.Tensor, prefill_support: bool = False): """Return whether the tensor size fits custom AR even if it is strided. @@ -1145,6 +1172,48 @@ def custom_all_reduce( registered_input=False, ) + def all_reduce_dual( + self, + left: torch.Tensor, + right: torch.Tensor, + *, + left_out: torch.Tensor | None = None, + right_out: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Reduce two separately registered inputs in one one-stage launch.""" + if left_out is None: + left_out = torch.empty_like(left) + if right_out is None: + right_out = torch.empty_like(right) + if not left_out.is_contiguous() or not right_out.is_contiguous(): + raise ValueError("dual-input custom allreduce outputs must be contiguous") + if self._ops_all_reduce_dual is None: + raise RuntimeError("dual-input custom allreduce is unavailable") + self._ops_all_reduce_dual( + self._ptr, + left, + right, + left_out, + right_out, + self._pool["input"].data_ptr, + self._pool["input"].max_size, + ) + return left_out, right_out + + def custom_all_reduce_dual( + self, + left: torch.Tensor, + right: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Use the graph-registered dual kernel or return ``None`` for fallback.""" + if not self.should_custom_ar_dual(left, right): + return None + if not self._IS_CAPTURING: + return None + if torch.cuda.is_current_stream_capturing(): + return self.all_reduce_dual(left, right) + return torch.zeros_like(left), torch.zeros_like(right) + # reduce_scatter split_dim enum — must match `aiter::ReduceScatterSplitDim` # in csrc/include/custom_all_reduce.cuh. _RS_SPLIT_FIRST = 0 diff --git a/aiter/ops/custom_all_reduce.py b/aiter/ops/custom_all_reduce.py index 11a72c0b8d..a6e97ad644 100644 --- a/aiter/ops/custom_all_reduce.py +++ b/aiter/ops/custom_all_reduce.py @@ -35,6 +35,18 @@ def all_reduce( ) -> None: ... +@compile_ops("module_custom_all_reduce", develop=True) +def all_reduce_dual( + _fa: int, + left: torch.Tensor, + right: torch.Tensor, + left_out: torch.Tensor, + right_out: torch.Tensor, + staging_ptr: int, + staging_bytes: int, +) -> None: ... + + @compile_ops("module_custom_all_reduce", develop=True) def reduce_scatter( _fa: int, diff --git a/csrc/include/custom_all_reduce.cuh b/csrc/include/custom_all_reduce.cuh index ee304c132c..f59128ca03 100644 --- a/csrc/include/custom_all_reduce.cuh +++ b/csrc/include/custom_all_reduce.cuh @@ -480,6 +480,116 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_1stage(RankData* _ end_sync(sg, self_sg, rank); } +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_1stage_dual(RankData* _left_staging_dp, + RankData* _right_dp, + RankSignals sg, +#ifndef USE_ROCM + volatile +#endif + Signal* self_sg, + const T* __restrict__ left, + T* __restrict__ left_staging, + T* __restrict__ left_result, + T* __restrict__ right_result, + int rank, + int left_size, + int right_size) +{ + constexpr int pack_size = 16 / sizeof(T); + constexpr int tnum_gpu = THREAD_NUM / ngpus; + using P = typename opus::vector_t; + using A = typename opus::vector_t; + + auto left_dp = *_left_staging_dp; + auto right_dp = *_right_dp; + int warp_id = threadIdx.x / tnum_gpu; + int lane_id = threadIdx.x % tnum_gpu; + const int size = left_size + right_size; + const int step = gridDim.x * tnum_gpu; + const int start = blockIdx.x * tnum_gpu + lane_id; + __shared__ P tmp_smem[2][tnum_gpu * ngpus]; + + // Stage the first local input into the existing pre-registered pool. + // Each block copies exactly the logical packs it will later reduce. The + // release/acquire system barrier makes those stores visible to the + // corresponding block on every peer before any peer load. + if(threadIdx.x < tnum_gpu) + { + const P* left_packs = reinterpret_cast(left); + P* staging_packs = reinterpret_cast(left_staging); + for(int copy_idx = start; copy_idx < left_size; copy_idx += step) + staging_packs[copy_idx] = left_packs[copy_idx]; + } + end_sync(sg, self_sg, rank); + + const int first = blockIdx.x * tnum_gpu; + int iters = 0; + { + int remaining = size - first; + iters = remaining > 0 ? (remaining + step - 1) / step : 0; + } + + int buffer = 0; + int idx0 = start; + if(idx0 < size) + { + const bool use_left = idx0 < left_size; + const int input_idx = use_left ? idx0 : idx0 - left_size; + const P* input = reinterpret_cast( + use_left ? left_dp.ptrs[warp_id] : right_dp.ptrs[warp_id]); + tmp_smem[buffer][warp_id * tnum_gpu + lane_id] = input[input_idx]; + } + __syncthreads(); + + for(int iteration = 0; iteration < iters; ++iteration) + { + const int current_idx = idx0 + iteration * step; + const int next_idx = current_idx + step; + const int next_buffer = buffer ^ 1; + + if(warp_id == 0 && current_idx < size) + { + P first_value = tmp_smem[buffer][lane_id]; + A accumulator; +#pragma unroll + for(int element = 0; element < pack_size; ++element) + accumulator[element] = upcast_s(first_value[element]); +#pragma unroll + for(int gpu = 1; gpu < ngpus; ++gpu) + { + P value = tmp_smem[buffer][gpu * tnum_gpu + lane_id]; +#pragma unroll + for(int element = 0; element < pack_size; ++element) + accumulator[element] += upcast_s(value[element]); + } + + P output; +#pragma unroll + for(int element = 0; element < pack_size; ++element) + output[element] = downcast_s(accumulator[element]); + + if(current_idx < left_size) + reinterpret_cast(left_result)[current_idx] = output; + else + reinterpret_cast(right_result)[current_idx - left_size] = output; + } + + if(next_idx < size) + { + const bool use_left = next_idx < left_size; + const int input_idx = use_left ? next_idx : next_idx - left_size; + const P* input = reinterpret_cast( + use_left ? left_dp.ptrs[warp_id] : right_dp.ptrs[warp_id]); + tmp_smem[next_buffer][warp_id * tnum_gpu + lane_id] = input[input_idx]; + } + __syncthreads(); + buffer = next_buffer; + } + end_sync(sg, self_sg, rank); +} + template __global__ void __launch_bounds__(512, 1) cross_device_reduce_2stage(RankData* _input_dp, RankData* _output_dp, @@ -3896,6 +4006,81 @@ class CustomAllreduce #undef KL } +template +void allreduceDual(hipStream_t stream, + T* left, + T* right, + T* left_output, + T* right_output, + int64_t left_size, + int64_t right_size, + T* left_staging, + int64_t left_staging_bytes) +{ + constexpr int pack_size = 16 / sizeof(T); + if(left_size <= 0 || right_size <= 0) + throw std::runtime_error("dual-input custom allreduce requires non-empty inputs"); + if(left_size % pack_size != 0 || right_size % pack_size != 0) + throw std::runtime_error( + "dual-input custom allreduce requires each input byte size to be a multiple of 16"); + if(!full_nvlink_) + throw std::runtime_error( + "dual-input custom allreduce requires fully connected peer GPUs"); + if(left_staging == nullptr || + left_staging_bytes < static_cast(left_size) * sizeof(T)) + throw std::runtime_error( + "dual-input custom allreduce staging buffer is too small"); + + const int64_t total_bytes = (left_size + right_size) * sizeof(T); + const int64_t one_stage_limit = + world_size_ <= 4 ? 160 * 1024 : (world_size_ <= 8 ? 80 * 1024 : 0); + if(total_bytes >= one_stage_limit) + throw std::runtime_error( + "dual-input custom allreduce currently supports only the one-stage size range"); + const int left_packs = static_cast(left_size / pack_size); + const int right_packs = static_cast(right_size / pack_size); + + auto staging_it = input_buffer.find(left_staging); + if(staging_it == input_buffer.end()) + throw std::runtime_error( + "dual-input custom allreduce staging metadata is unavailable"); + RankData* left_staging_ptrs = staging_it->second; + RankData* right_ptrs = get_buffer_RD(stream, right); + constexpr int threads = 512; + const int total_packs = left_packs + right_packs; + const int threads_per_gpu = threads / world_size_; + const int blocks = std::min( + kMaxBlocks, (total_packs + threads_per_gpu - 1) / threads_per_gpu); + +#define DUAL_REDUCE_CASE(ngpus) \ + case ngpus: \ + cross_device_reduce_1stage_dual \ + <<>>(left_staging_ptrs, \ + right_ptrs, \ + sg_, \ + self_sg_, \ + left, \ + left_staging, \ + left_output, \ + right_output, \ + rank_, \ + left_packs, \ + right_packs); \ + break + + switch(world_size_) + { + DUAL_REDUCE_CASE(2); + DUAL_REDUCE_CASE(4); + DUAL_REDUCE_CASE(6); + DUAL_REDUCE_CASE(8); + default: + throw std::runtime_error( + "dual-input custom allreduce supports world sizes 2, 4, 6, and 8"); + } +#undef DUAL_REDUCE_CASE +} + // reduce_scatter dispatch. Python wrapper is responsible for: // - normalizing dim < 0 // - rejecting shapes where n (or k, for kFirst) % ngpus != 0 diff --git a/csrc/include/custom_all_reduce.h b/csrc/include/custom_all_reduce.h index b5d01b38b1..63935caa4b 100644 --- a/csrc/include/custom_all_reduce.h +++ b/csrc/include/custom_all_reduce.h @@ -38,6 +38,13 @@ void all_reduce(fptr_t _fa, bool open_fp8_quant, int64_t reg_inp_ptr, int64_t reg_inp_bytes); +void all_reduce_dual(fptr_t _fa, + const aiter_tensor_t& left, + const aiter_tensor_t& right, + const aiter_tensor_t& left_out, + const aiter_tensor_t& right_out, + int64_t staging_ptr, + int64_t staging_bytes); // reduce_scatter dispatcher. (m, n, k, split_dim) describe the canonical // shape the Python wrapper collapsed the input to: // split_dim = 0 (kFirst): only `k` (= numel) used diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index cfdc18f3e0..b390883adf 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp @@ -494,6 +494,15 @@ namespace py = pybind11; py::arg("open_fp8_quant"), \ py::arg("reg_inp_ptr"), \ py::arg("reg_inp_bytes")); \ + m.def("all_reduce_dual", \ + &aiter::all_reduce_dual, \ + py::arg("_fa"), \ + py::arg("left"), \ + py::arg("right"), \ + py::arg("left_out"), \ + py::arg("right_out"), \ + py::arg("staging_ptr"), \ + py::arg("staging_bytes")); \ m.def("reduce_scatter", \ &aiter::reduce_scatter, \ py::arg("_fa"), \ diff --git a/csrc/kernels/custom_all_reduce.cu b/csrc/kernels/custom_all_reduce.cu index e536f97414..c69f545708 100644 --- a/csrc/kernels/custom_all_reduce.cu +++ b/csrc/kernels/custom_all_reduce.cu @@ -117,6 +117,62 @@ static void _all_reduce(fptr_t _fa, void* inp, void* out, } } +static void _all_reduce_dual(fptr_t _fa, + void* left, + void* right, + void* left_out, + void* right_out, + int64_t left_numel, + int64_t right_numel, + AiterDtype dtype, + void* staging, + int64_t staging_bytes) +{ + hipStream_t stream = aiter::getCurrentHIPStream(); + auto fa = reinterpret_cast(_fa); + switch(dtype) + { + case AITER_DTYPE_fp32: + fa->allreduceDual(stream, + reinterpret_cast(left), + reinterpret_cast(right), + reinterpret_cast(left_out), + reinterpret_cast(right_out), + left_numel, + right_numel, + reinterpret_cast(staging), + staging_bytes); + break; + case AITER_DTYPE_fp16: + fa->allreduceDual(stream, + reinterpret_cast(left), + reinterpret_cast(right), + reinterpret_cast(left_out), + reinterpret_cast(right_out), + left_numel, + right_numel, + reinterpret_cast(staging), + staging_bytes); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case AITER_DTYPE_bf16: + fa->allreduceDual(stream, + reinterpret_cast(left), + reinterpret_cast(right), + reinterpret_cast(left_out), + reinterpret_cast(right_out), + left_numel, + right_numel, + reinterpret_cast(staging), + staging_bytes); + break; +#endif + default: + throw std::runtime_error( + "dual-input custom allreduce only supports float32, float16, and bfloat16"); + } +} + static void _reduce_scatter(fptr_t _fa, void* inp, void* out, int m, int n, int k, aiter::ReduceScatterSplitDim split_dim, @@ -440,6 +496,53 @@ void all_reduce(fptr_t _fa, use_new, open_fp8_quant, is_broadcast_reg_outptr); } +void all_reduce_dual(fptr_t _fa, + const aiter_tensor_t& left, + const aiter_tensor_t& right, + const aiter_tensor_t& left_out, + const aiter_tensor_t& right_out, + int64_t staging_ptr, + int64_t staging_bytes) +{ + if(left.device_id != right.device_id || left.device_id != left_out.device_id || + left.device_id != right_out.device_id) + throw std::runtime_error( + "dual-input custom allreduce tensors must be on the same device"); + if(left.dtype() != right.dtype() || left.dtype() != left_out.dtype() || + left.dtype() != right_out.dtype()) + throw std::runtime_error( + "dual-input custom allreduce tensors must have the same dtype"); + if(left.numel() != left_out.numel() || right.numel() != right_out.numel()) + throw std::runtime_error( + "dual-input custom allreduce output sizes must match their inputs"); + if(!left.is_contiguous() || !right.is_contiguous() || !left_out.is_contiguous() || + !right_out.is_contiguous()) + throw std::runtime_error( + "dual-input custom allreduce tensors must be contiguous"); + if(left_out.data_ptr() == left.data_ptr() || left_out.data_ptr() == right.data_ptr() || + right_out.data_ptr() == left.data_ptr() || + right_out.data_ptr() == right.data_ptr() || + left_out.data_ptr() == right_out.data_ptr()) + throw std::runtime_error( + "dual-input custom allreduce outputs must not alias inputs or each other"); + if(staging_ptr == 0 || + staging_bytes < static_cast(left.numel() * left.element_size())) + throw std::runtime_error( + "dual-input custom allreduce staging buffer is too small"); + + HipDeviceGuard device_guard(left.device_id); + _all_reduce_dual(_fa, + left.data_ptr(), + right.data_ptr(), + left_out.data_ptr(), + right_out.data_ptr(), + left.numel(), + right.numel(), + left.dtype(), + reinterpret_cast(staging_ptr), + staging_bytes); +} + void reduce_scatter(fptr_t _fa, const aiter_tensor_t& inp, const aiter_tensor_t& out, diff --git a/op_tests/test_custom_all_reduce_dual.py b/op_tests/test_custom_all_reduce_dual.py new file mode 100644 index 0000000000..6b025eb214 --- /dev/null +++ b/op_tests/test_custom_all_reduce_dual.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +from aiter.dist.device_communicators.custom_all_reduce import CustomAllreduce + + +def _communicator(**overrides): + values = { + "disabled": False, + "_ops_all_reduce_dual": object(), + "_car_min_size": 0, + "_car_max_size": 64 * 1024 * 1024, + "world_size": 8, + "fully_connected": True, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_should_custom_ar_dual_accepts_kimi_k3_boundary(): + communicator = _communicator() + left = torch.empty((1, 3584), dtype=torch.bfloat16) + right = torch.empty((1, 7168), dtype=torch.bfloat16) + + assert CustomAllreduce.should_custom_ar_dual(communicator, left, right) + + +@pytest.mark.parametrize( + "left,right", + [ + ( + torch.empty(0, dtype=torch.bfloat16), + torch.empty(8, dtype=torch.bfloat16), + ), + ( + torch.empty(8, dtype=torch.bfloat16), + torch.empty(8, dtype=torch.float16), + ), + ( + torch.empty(8, dtype=torch.int32), + torch.empty(8, dtype=torch.int32), + ), + ( + torch.empty(7, dtype=torch.bfloat16), + torch.empty(8, dtype=torch.bfloat16), + ), + ( + torch.empty((8, 16), dtype=torch.bfloat16)[:, ::2], + torch.empty(8, dtype=torch.bfloat16), + ), + ], +) +def test_should_custom_ar_dual_rejects_unsupported_inputs(left, right): + communicator = _communicator() + + assert not CustomAllreduce.should_custom_ar_dual(communicator, left, right) + + +def test_should_custom_ar_dual_rejects_transport_and_size_mismatches(): + left = torch.empty(8, dtype=torch.bfloat16) + right = torch.empty(8, dtype=torch.bfloat16) + + assert not CustomAllreduce.should_custom_ar_dual( + _communicator(disabled=True), left, right + ) + assert not CustomAllreduce.should_custom_ar_dual( + _communicator(_ops_all_reduce_dual=None), left, right + ) + assert not CustomAllreduce.should_custom_ar_dual( + _communicator(fully_connected=False), left, right + ) + assert not CustomAllreduce.should_custom_ar_dual( + _communicator(_car_min_size=32), left, right + ) + + left_at_limit = torch.empty(20 * 1024, dtype=torch.float32) + assert not CustomAllreduce.should_custom_ar_dual( + _communicator(), left_at_limit, right + ) + + +def test_all_reduce_dual_passes_explicit_staging_contract(): + calls = [] + + def op(*args): + calls.append(args) + + staging = SimpleNamespace(data_ptr=0x1234, max_size=64 * 1024 * 1024) + communicator = _communicator( + _ops_all_reduce_dual=op, + _ptr=0x5678, + _pool={"input": staging}, + ) + left = torch.empty((1, 3584), dtype=torch.bfloat16) + right = torch.empty((1, 7168), dtype=torch.bfloat16) + + left_out, right_out = CustomAllreduce.all_reduce_dual( + communicator, + left, + right, + ) + + assert left_out.shape == left.shape + assert right_out.shape == right.shape + assert calls == [ + ( + communicator._ptr, + left, + right, + left_out, + right_out, + staging.data_ptr, + staging.max_size, + ) + ]