perf(custom-ar): add graph-safe Kimi-K3 dual-input all-reduce

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

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,

csrc/include/custom_all_reduce.cuh

@@ -480,6 +480,116 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_1stage(RankData* _
     end_sync<ngpus, true>(sg, self_sg, rank);
 }
 
+template <typename T, int ngpus>
+__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<T, pack_size>;
+    using A                 = typename opus::vector_t<opus::fp32_t, pack_size>;
+
+    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<const P*>(left);
+        P* staging_packs    = reinterpret_cast<P*>(left_staging);
+        for(int copy_idx = start; copy_idx < left_size; copy_idx += step)
+            staging_packs[copy_idx] = left_packs[copy_idx];
+    }
+    end_sync<ngpus>(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<const P*>(
+            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<T>(accumulator[element]);
+
+            if(current_idx < left_size)
+                reinterpret_cast<P*>(left_result)[current_idx] = output;
+            else
+                reinterpret_cast<P*>(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<const P*>(
+                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<ngpus, true>(sg, self_sg, rank);
+}
+
 template <typename T, int ngpus, bool is_broadcast_reg_outptr = false>
 __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 <typename T>
+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<int64_t>(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<int>(left_size / pack_size);
+    const int right_packs = static_cast<int>(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<T, ngpus>                        \
+            <<<blocks, threads, 0, stream>>>(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

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

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"),                                                                      \