[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/cudax/test/multi_gpu/algorithms/CMakeLists.txt

@@ -0,0 +1,13 @@
+#===----------------------------------------------------------------------===##
+#
+# Part of CUDA Experimental in CUDA C++ Core Libraries,
+# under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+#
+#===----------------------------------------------------------------------===##
+
+add_subdirectory(reduce)
+add_subdirectory(exclusive_scan)
+add_subdirectory(inclusive_scan)

cccl_upstream/cudax/test/multi_gpu/algorithms/exclusive_scan/CMakeLists.txt

@@ -0,0 +1,20 @@
+#===----------------------------------------------------------------------===##
+#
+# Part of CUDA Experimental in CUDA C++ Core Libraries,
+# under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+#
+#===----------------------------------------------------------------------===##
+
+if (NOT cudax_ENABLE_NCCL)
+  return()
+endif()
+
+file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
+
+foreach (src IN LISTS test_srcs)
+  cudax_add_multi_gpu_test("algorithms.exclusive_scan" test_target "${src}")
+  target_link_libraries(${test_target} PRIVATE NCCL::nccl)
+endforeach()

cccl_upstream/cudax/test/multi_gpu/algorithms/exclusive_scan/range_basic.cu

@@ -0,0 +1,388 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/array>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/limits>
+#include <cuda/std/span>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+namespace
+{
+struct custom_plus
+{
+  template <class T>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
+  {
+    return lhs + rhs;
+  }
+};
+
+using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
+using value_types  = c2h::type_list<cuda::std::int32_t, float, custom_value>;
+using operators    = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
+
+static_assert(cudax::nccl_transportable<custom_value>);
+
+template <typename T>
+T make_value(int i)
+{
+  return static_cast<T>(i);
+}
+
+template <>
+custom_value make_value<>(int i)
+{
+  custom_value ret{};
+
+  ret.key = static_cast<std::size_t>(i);
+  ret.val = static_cast<std::size_t>(i);
+  return ret;
+};
+
+template <class T, class Op>
+[[nodiscard]] T get_identity()
+{
+  if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
+  {
+    return make_value<T>(0);
+  }
+  else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
+  {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+  else
+  {
+    static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
+  }
+}
+
+template <class T, class Op>
+[[nodiscard]] std::vector<T>
+expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
+{
+  std::vector<T> reference;
+
+  for (const auto& values : inputs_by_rank)
+  {
+    reference.insert(reference.end(), values.begin(), values.end());
+  }
+
+  std::vector<T> scan(reference.size());
+  std::exclusive_scan(reference.begin(), reference.end(), scan.begin(), init, op);
+
+  cuda::std::size_t offset = 0;
+  for (int r = 0; r < rank; ++r)
+  {
+    offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
+  }
+
+  const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
+  return {scan.begin() + offset, scan.begin() + offset + count};
+}
+
+// Run the full scan, wait for it to finish, and check that `exclusive_scan` left its argument
+// ranges untouched. This boilerplate is identical for every test regardless of how the inputs are
+// shaped.
+template <class Env, class T, class Op>
+void do_exclusive_scan(
+  cuda::std::span<cudax::nccl_communicator_ref> comms,
+  const std::vector<Env>& envs,
+  std::vector<cuda::device_buffer<T>>& in,
+  std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
+  const T& init,
+  const T& ident,
+  Op op)
+{
+  const auto envs_size    = envs.size();
+  const auto in_copy      = in;
+  const auto outputs_copy = outputs;
+
+  INFO("init = " << init);
+  INFO("ident = " << ident);
+
+  cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
+
+  // cuda::std::execution::env has no operator==, so we can only compare the sizes.
+  REQUIRE(envs.size() == envs_size);
+  // Scan call should not modify the inputs in any ways
+  REQUIRE(in.size() == in_copy.size());
+  for (cuda::std::size_t i = 0; i < in.size(); ++i)
+  {
+    INFO("device = " << i);
+    REQUIRE_THAT(in[i], Equals(in_copy[i]));
+  }
+  REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
+}
+} // namespace
+
+MULTI_GPU_TEST("exclusive_scan documentation example", c2h::type_list<int>)
+{
+  auto comms = this->communicators();
+
+  if (comms.size() < 2)
+  {
+    SKIP("The exclusive_scan documentation example requires at least two local GPUs");
+  }
+
+  auto streams_owned = nccl_test_util::make_streams();
+  // Convert to stream_ref directly, cuda::stream on their own cant be passed directly to CUB
+  auto streams = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()};
+
+  //! [exclusive_scan]
+  constexpr cuda::std::array input_values{1, 2};
+  std::vector<cuda::device_buffer<int>> inputs;
+  std::vector<cuda::device_buffer<int>> outputs;
+
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto device = comms[i].logical_device().underlying_device();
+
+    inputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values));
+    outputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values.size(), cuda::no_init));
+  }
+
+  std::vector<typename cuda::device_buffer<int>::iterator> output_iterators = make_output_iterators(outputs);
+
+  cudax::exclusive_scan(
+    cudax::distributed,
+    comms,
+    // Passing streams as the environment directly
+    streams,
+    inputs,
+    output_iterators,
+    /*__init=*/0);
+
+  constexpr cuda::std::array expected_rank_0{0, 1};
+  constexpr cuda::std::array expected_rank_1{3, 4};
+  const auto expected_0 =
+    cuda::make_buffer<int>(outputs[0].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_0);
+  const auto expected_1 =
+    cuda::make_buffer<int>(outputs[1].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_1);
+  REQUIRE_THAT(outputs[0], Equals(expected_0));
+  REQUIRE_THAT(outputs[1], Equals(expected_1));
+  //! [exclusive_scan]
+}
+
+MULTI_GPU_TEST("exclusive_scan, one element per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each scan with a few hardcoded initializers. The init participates in the fold the same
+  // way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
+  // one-element output buffer and an environment carrying its stream, so the scan is stream-ordered
+  // on the correct device. `reference` mirrors the contributions of every global rank so we can
+  // compute the host-side scan exactly like `exclusive_scan` does on the device.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(1, make_value<T>(r));
+  }
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("exclusive_scan, multiple elements per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each scan with a few hardcoded initializers. The init participates in the fold the same
+  // way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes ten copies of `rank`. `exclusive_scan` first
+  // computes a local prefix for each rank seeded by the prefix of all previous ranks. Each local
+  // rank also gets an output buffer and an environment carrying its stream. `reference` mirrors
+  // every global rank's ten contributions for the host-side scan.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
+  }
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("exclusive_scan, some ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
+  // range. Rank 0 is always non-empty. `exclusive_scan` must treat an empty rank as contributing
+  // nothing, exactly like `std::exclusive_scan` over the surviving elements. `reference` mirrors
+  // that for the host-side scan.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    if (r % 2 == 0)
+    {
+      inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
+    }
+  }
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("exclusive_scan, all ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const auto init  = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // No rank contributes any element. Scanning nothing produces no output values, exactly like
+  // `std::exclusive_scan` over an empty range.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/exclusive_scan/range_defaults.cu

@@ -0,0 +1,112 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/iterator>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+MULTI_GPU_TEST("exclusive_scan, range overloads default values", )
+{
+  using T  = cuda::std::int32_t;
+  using Op = ::cuda::std::plus<>;
+
+  constexpr auto init = T{};
+  constexpr T ident   = cuda::identity_element<Op, T>();
+  constexpr auto op   = Op{};
+
+  const auto comms   = this->communicators();
+  const auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<cuda::stream_ref> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    std::vector<T> values(values_per_rank);
+    std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
+    envs.emplace_back(streams[i]);
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  const auto expected_values = [&] {
+    std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
+
+    std::iota(reference.begin(), reference.end(), T{1});
+
+    std::vector<T> expected_values(reference.size());
+
+    std::exclusive_scan(reference.begin(), reference.end(), expected_values.begin(), init, op);
+
+    return expected_values;
+  }();
+
+  const auto check_outputs = [&] {
+    const auto exp_span = cuda::std::span{expected_values};
+
+    for (cuda::std::size_t i = 0; i < out.size(); ++i)
+    {
+      const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
+      const auto expected =
+        cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
+
+      REQUIRE_THAT(out[i], Equals(expected));
+    }
+  };
+
+  SECTION("Default init, op, ident (all)")
+  {
+    cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs);
+    check_outputs();
+  }
+
+  SECTION("Default op, ident")
+  {
+    cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init);
+    check_outputs();
+  }
+
+  SECTION("Default ident")
+  {
+    cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op);
+    check_outputs();
+  }
+
+  SECTION("Default none")
+  {
+    cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
+    check_outputs();
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/exclusive_scan/single_comm_basic.cu

@@ -0,0 +1,289 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/array>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/limits>
+#include <cuda/std/span>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+namespace
+{
+struct custom_plus
+{
+  template <class T>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
+  {
+    return lhs + rhs;
+  }
+};
+
+using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
+using value_types  = c2h::type_list<cuda::std::int32_t, float, custom_value>;
+using operators    = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
+
+static_assert(cudax::nccl_transportable<custom_value>);
+
+template <typename T>
+T make_value(int i)
+{
+  return static_cast<T>(i);
+}
+
+template <>
+custom_value make_value<>(int i)
+{
+  custom_value ret{};
+
+  ret.key = static_cast<std::size_t>(i);
+  ret.val = static_cast<std::size_t>(i);
+  return ret;
+};
+
+template <class T, class Op>
+[[nodiscard]] T get_identity()
+{
+  if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
+  {
+    return make_value<T>(0);
+  }
+  else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
+  {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+  else
+  {
+    static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
+  }
+}
+
+template <class T, class Op>
+[[nodiscard]] std::vector<T>
+expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
+{
+  std::vector<T> reference;
+
+  for (const auto& values : inputs_by_rank)
+  {
+    reference.insert(reference.end(), values.begin(), values.end());
+  }
+
+  std::vector<T> scan(reference.size());
+  std::exclusive_scan(reference.begin(), reference.end(), scan.begin(), init, op);
+
+  cuda::std::size_t offset = 0;
+  for (int r = 0; r < rank; ++r)
+  {
+    offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
+  }
+
+  const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
+  return {scan.begin() + offset, scan.begin() + offset + count};
+}
+
+// Drive the scan through the single-communicator overload, one thread per local rank. The
+// per-rank calls must rendezvous in their collectives, so issuing them serially would deadlock.
+// Catch2 assertions remain on the main thread after all worker threads have joined.
+template <class T, class Op>
+void run_case(cuda::std::span<cudax::nccl_communicator_ref> comms,
+              const std::vector<std::vector<T>>& inputs_by_rank,
+              const T& init,
+              const T& ident,
+              Op op)
+{
+  auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  const auto in_copy      = in;
+  auto outputs            = make_output_iterators(out);
+  const auto outputs_copy = outputs;
+
+  INFO("init = " << init);
+  INFO("ident = " << ident);
+
+  run_threaded(comms.size(), [&](cuda::std::size_t i) {
+    cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
+  });
+
+  REQUIRE(in.size() == in_copy.size());
+  for (cuda::std::size_t i = 0; i < in.size(); ++i)
+  {
+    INFO("device = " << i);
+    REQUIRE_THAT(in[i], Equals(in_copy[i]));
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, op);
+    const auto expected =
+      cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(expected));
+  }
+  REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
+}
+} // namespace
+
+MULTI_GPU_TEST("exclusive_scan single-comm documentation example", c2h::type_list<int>)
+{
+  auto comms = this->communicators();
+
+  if (comms.size() < 2)
+  {
+    SKIP("The exclusive_scan documentation example requires at least two local GPUs");
+  }
+
+  auto streams_owned = nccl_test_util::make_streams();
+  auto streams       = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()};
+
+  // Must be pre-allocated since it is written to by threads
+  std::vector<std::string> failed(comms.front().size());
+
+  // Every communicator rank must invoke the collective concurrently.
+  run_threaded(comms.size(), [&](cuda::std::size_t i) {
+    auto& communicator = comms[i];
+    auto environment   = streams[i];
+    const auto device  = communicator.logical_device().underlying_device();
+
+    //! [exclusive_scan_single_range]
+    constexpr cuda::std::array input_values{1, 2};
+
+    auto input  = cuda::make_device_buffer<int>(environment, device, input_values);
+    auto output = cuda::make_device_buffer<int>(environment, device, input_values.size(), cuda::no_init);
+
+    cudax::exclusive_scan(cudax::distributed, communicator, environment, input, output.begin(), /*__init=*/0);
+
+    // Every rank contributes {1, 2}, so rank r starts with a prefix of 3 * r.
+    const auto rank = communicator.rank();
+    const auto expected =
+      cuda::make_buffer<int>(output.stream(), cuda::mr::legacy_pinned_memory_resource{}, {3 * rank, 3 * rank + 1});
+
+    //! [exclusive_scan_single_range]
+
+    // catch2 isn't thread safe by default, so we can't use the usual requires expression. So
+    // we roll a hacky version of it ourselves
+    if (const auto matcher = Equals(expected); !matcher.match(output))
+    {
+      failed[rank] = matcher.describe();
+    }
+  });
+
+  for (cuda::std::size_t i = 0; i < failed.size(); ++i)
+  {
+    if (const auto& err_str = failed[i]; !err_str.empty())
+    {
+      INFO("rank: " << i);
+      REQUIRE(err_str == ""); // Should print the full error string
+    }
+  }
+}
+
+MULTI_GPU_TEST("exclusive_scan single-comm, one element per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  std::vector<std::vector<T>> inputs_by_rank;
+
+  inputs_by_rank.reserve(static_cast<cuda::std::size_t>(comms.front().size()));
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    const auto v = {make_value<T>(r)};
+
+    inputs_by_rank.emplace_back(v);
+  }
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}
+
+MULTI_GPU_TEST("exclusive_scan single-comm, multiple elements per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  constexpr auto values_per_rank = 10;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
+  }
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}
+
+MULTI_GPU_TEST("exclusive_scan single-comm, some ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  constexpr auto values_per_rank = 10;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    if (r % 2 == 0)
+    {
+      inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
+    }
+  }
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}
+
+MULTI_GPU_TEST("exclusive_scan single-comm, all ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  const std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/exclusive_scan/single_comm_defaults.cu

@@ -0,0 +1,120 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/iterator>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+MULTI_GPU_TEST("exclusive_scan single-comm, overloads default values", )
+{
+  using T  = cuda::std::int32_t;
+  using Op = ::cuda::std::plus<>;
+
+  constexpr auto init = T{};
+  constexpr T ident   = cuda::identity_element<Op, T>();
+  constexpr auto op   = Op{};
+
+  const auto comms   = this->communicators();
+  const auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<cuda::stream_ref> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    std::vector<T> values(values_per_rank);
+    std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
+    envs.emplace_back(streams[i]);
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  const auto expected_values = [&] {
+    std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
+
+    std::iota(reference.begin(), reference.end(), T{1});
+
+    std::vector<T> expected_values(reference.size());
+
+    std::exclusive_scan(reference.begin(), reference.end(), expected_values.begin(), init, op);
+
+    return expected_values;
+  }();
+
+  const auto check_outputs = [&] {
+    const auto exp_span = cuda::std::span{expected_values};
+
+    for (cuda::std::size_t i = 0; i < out.size(); ++i)
+    {
+      const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
+      const auto expected =
+        cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
+
+      REQUIRE_THAT(out[i], Equals(expected));
+    }
+  };
+
+  SECTION("Default init, op, ident (all)")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i]);
+    });
+    check_outputs();
+  }
+
+  SECTION("Default op, ident")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init);
+    });
+    check_outputs();
+  }
+
+  SECTION("Default ident")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op);
+    });
+    check_outputs();
+  }
+
+  SECTION("Default none")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
+    });
+    check_outputs();
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/inclusive_scan/CMakeLists.txt

@@ -0,0 +1,20 @@
+#===----------------------------------------------------------------------===##
+#
+# Part of CUDA Experimental in CUDA C++ Core Libraries,
+# under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+#
+#===----------------------------------------------------------------------===##
+
+if (NOT cudax_ENABLE_NCCL)
+  return()
+endif()
+
+file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
+
+foreach (src IN LISTS test_srcs)
+  cudax_add_multi_gpu_test("algorithms.inclusive_scan" test_target "${src}")
+  target_link_libraries(${test_target} PRIVATE NCCL::nccl)
+endforeach()

cccl_upstream/cudax/test/multi_gpu/algorithms/inclusive_scan/range_basic.cu

@@ -0,0 +1,389 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/array>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/limits>
+#include <cuda/std/span>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+namespace
+{
+struct custom_plus
+{
+  template <class T>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
+  {
+    return lhs + rhs;
+  }
+};
+
+using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
+using value_types  = c2h::type_list<cuda::std::int32_t, float, custom_value>;
+using operators    = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
+
+static_assert(cudax::nccl_transportable<custom_value>);
+
+template <typename T>
+T make_value(int i)
+{
+  return static_cast<T>(i);
+}
+
+template <>
+custom_value make_value<>(int i)
+{
+  custom_value ret{};
+
+  ret.key = static_cast<std::size_t>(i);
+  ret.val = static_cast<std::size_t>(i);
+  return ret;
+};
+
+template <class T, class Op>
+[[nodiscard]] T get_identity()
+{
+  if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
+  {
+    return make_value<T>(0);
+  }
+  else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
+  {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+  else
+  {
+    static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
+  }
+}
+
+template <class T, class Op>
+[[nodiscard]] std::vector<T>
+expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
+{
+  std::vector<T> reference;
+
+  for (const auto& values : inputs_by_rank)
+  {
+    reference.insert(reference.end(), values.begin(), values.end());
+  }
+
+  std::vector<T> scan(reference.size());
+  std::inclusive_scan(reference.begin(), reference.end(), scan.begin(), op, init);
+
+  cuda::std::size_t offset = 0;
+  for (int r = 0; r < rank; ++r)
+  {
+    offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
+  }
+
+  const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
+  return {scan.begin() + offset, scan.begin() + offset + count};
+}
+
+// Run the full scan, wait for it to finish, and check that `inclusive_scan` left its argument
+// ranges untouched. This boilerplate is identical for every test regardless of how the inputs are
+// shaped.
+template <class Env, class T, class Op>
+void do_inclusive_scan(
+  cuda::std::span<cudax::nccl_communicator_ref> comms,
+  const std::vector<Env>& envs,
+  std::vector<cuda::device_buffer<T>>& in,
+  std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
+  const T& init,
+  const T& ident,
+  Op op)
+{
+  const auto envs_size    = envs.size();
+  const auto in_copy      = in;
+  const auto outputs_copy = outputs;
+
+  INFO("init = " << init);
+  INFO("ident = " << ident);
+
+  cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
+
+  // cuda::std::execution::env has no operator==, so we can only compare the sizes.
+  REQUIRE(envs.size() == envs_size);
+  // Scan call should not modify the inputs in any ways
+  REQUIRE(in.size() == in_copy.size());
+  for (cuda::std::size_t i = 0; i < in.size(); ++i)
+  {
+    INFO("device = " << i);
+    REQUIRE_THAT(in[i], Equals(in_copy[i]));
+  }
+  REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
+}
+} // namespace
+
+MULTI_GPU_TEST("inclusive_scan documentation example", c2h::type_list<int>)
+{
+  auto comms = this->communicators();
+
+  if (comms.size() < 2)
+  {
+    SKIP("The inclusive_scan documentation example requires at least two local GPUs");
+  }
+
+  auto streams_owned = nccl_test_util::make_streams();
+  // Convert to stream_ref directly, cuda::stream on their own cant be passed directly to CUB
+  auto streams = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()};
+
+  //! [inclusive_scan]
+  constexpr cuda::std::array input_values{1, 2};
+  std::vector<cuda::device_buffer<int>> inputs;
+  std::vector<cuda::device_buffer<int>> outputs;
+
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto device = comms[i].logical_device().underlying_device();
+
+    inputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values));
+    outputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values.size(), cuda::no_init));
+  }
+
+  std::vector<typename cuda::device_buffer<int>::iterator> output_iterators = make_output_iterators(outputs);
+
+  cudax::inclusive_scan(
+    cudax::distributed,
+    comms,
+    // Passing streams as the environment directly
+    streams,
+    inputs,
+    output_iterators,
+    /*__init=*/0);
+
+  constexpr cuda::std::array expected_rank_0{1, 3};
+  constexpr cuda::std::array expected_rank_1{4, 6};
+  const auto expected_0 =
+    cuda::make_buffer<int>(outputs[0].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_0);
+  const auto expected_1 =
+    cuda::make_buffer<int>(outputs[1].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_1);
+  REQUIRE_THAT(outputs[0], Equals(expected_0));
+  REQUIRE_THAT(outputs[1], Equals(expected_1));
+  //! [inclusive_scan]
+}
+
+MULTI_GPU_TEST("inclusive_scan, one element per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each scan with a few hardcoded initializers. The init participates in the fold the same
+  // way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
+  // one-element output buffer and an environment carrying its stream, so the scan is stream-ordered
+  // on the correct device. `reference` mirrors the contributions of every global rank so we can
+  // compute the host-side scan exactly like `inclusive_scan` does on the device.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(1, make_value<T>(r));
+  }
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("inclusive_scan, multiple elements per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each scan with a few hardcoded initializers. The init participates in the fold the same
+  // way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes ten copies of `rank`. `inclusive_scan` first
+  // computes a local prefix for each rank seeded by the prefix of all previous ranks. Each local
+  // rank also gets an output buffer and an environment carrying its stream. `reference` mirrors
+  // every global rank's ten contributions for the host-side scan.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    const auto v                                      = make_value<T>(r);
+    inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, v);
+  }
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("inclusive_scan, some ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
+  // range. Rank 0 is always non-empty. `inclusive_scan` must treat an empty rank as contributing
+  // nothing, exactly like `std::inclusive_scan` over the surviving elements. `reference` mirrors
+  // that for the host-side scan.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    if (r % 2 == 0)
+    {
+      inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
+    }
+  }
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("inclusive_scan, all ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const auto init  = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // No rank contributes any element. Scanning nothing produces no output values, exactly like
+  // `std::inclusive_scan` over an empty range.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
+
+  for (cuda::std::size_t i = 0; i < out.size(); ++i)
+  {
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
+    const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(exp));
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/inclusive_scan/range_defaults.cu

@@ -0,0 +1,112 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/iterator>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+MULTI_GPU_TEST("inclusive_scan, range overloads default values", )
+{
+  using T  = cuda::std::int32_t;
+  using Op = ::cuda::std::plus<>;
+
+  constexpr auto init = T{};
+  constexpr T ident   = cuda::identity_element<Op, T>();
+  constexpr auto op   = Op{};
+
+  const auto comms   = this->communicators();
+  const auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<cuda::stream_ref> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    std::vector<T> values(values_per_rank);
+    std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
+    envs.emplace_back(streams[i]);
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  const auto expected_values = [&] {
+    std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
+
+    std::iota(reference.begin(), reference.end(), T{1});
+
+    std::vector<T> expected_values(reference.size());
+
+    std::inclusive_scan(reference.begin(), reference.end(), expected_values.begin(), op, init);
+
+    return expected_values;
+  }();
+
+  const auto check_outputs = [&] {
+    const auto exp_span = cuda::std::span{expected_values};
+
+    for (cuda::std::size_t i = 0; i < out.size(); ++i)
+    {
+      const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
+      const auto expected =
+        cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
+
+      REQUIRE_THAT(out[i], Equals(expected));
+    }
+  };
+
+  SECTION("Default init, op, ident (all)")
+  {
+    cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs);
+    check_outputs();
+  }
+
+  SECTION("Default op, ident")
+  {
+    cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init);
+    check_outputs();
+  }
+
+  SECTION("Default ident")
+  {
+    cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op);
+    check_outputs();
+  }
+
+  SECTION("Default none")
+  {
+    cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
+    check_outputs();
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/inclusive_scan/single_comm_basic.cu

@@ -0,0 +1,290 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/array>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/limits>
+#include <cuda/std/span>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+namespace
+{
+struct custom_plus
+{
+  template <class T>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
+  {
+    return lhs + rhs;
+  }
+};
+
+using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
+using value_types  = c2h::type_list<cuda::std::int32_t, float, custom_value>;
+using operators    = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
+
+static_assert(cudax::nccl_transportable<custom_value>);
+
+template <typename T>
+T make_value(int i)
+{
+  return static_cast<T>(i);
+}
+
+template <>
+custom_value make_value<>(int i)
+{
+  custom_value ret{};
+
+  ret.key = static_cast<std::size_t>(i);
+  ret.val = static_cast<std::size_t>(i);
+  return ret;
+};
+
+template <class T, class Op>
+[[nodiscard]] T get_identity()
+{
+  if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
+  {
+    return make_value<T>(0);
+  }
+  else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
+  {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+  else
+  {
+    static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
+  }
+}
+
+template <class T, class Op>
+[[nodiscard]] std::vector<T>
+expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
+{
+  std::vector<T> reference;
+
+  for (const auto& values : inputs_by_rank)
+  {
+    reference.insert(reference.end(), values.begin(), values.end());
+  }
+
+  std::vector<T> scan(reference.size());
+  std::inclusive_scan(reference.begin(), reference.end(), scan.begin(), op, init);
+
+  cuda::std::size_t offset = 0;
+  for (int r = 0; r < rank; ++r)
+  {
+    offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
+  }
+
+  const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
+  return {scan.begin() + offset, scan.begin() + offset + count};
+}
+
+// Drive the scan through the single-communicator overload, one thread per local rank. The
+// per-rank calls must rendezvous in their collectives, so issuing them serially would deadlock.
+// Catch2 assertions remain on the main thread after all worker threads have joined.
+template <class T, class Op>
+void run_case(cuda::std::span<cudax::nccl_communicator_ref> comms,
+              const std::vector<std::vector<T>>& inputs_by_rank,
+              const T& init,
+              const T& ident,
+              Op op)
+{
+  auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  const auto in_copy      = in;
+  auto outputs            = make_output_iterators(out);
+  const auto outputs_copy = outputs;
+
+  INFO("init = " << init);
+  INFO("ident = " << ident);
+
+  run_threaded(comms.size(), [&](cuda::std::size_t i) {
+    cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
+  });
+
+  REQUIRE(in.size() == in_copy.size());
+  for (cuda::std::size_t i = 0; i < in.size(); ++i)
+  {
+    INFO("device = " << i);
+    REQUIRE_THAT(in[i], Equals(in_copy[i]));
+    const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, op);
+    const auto expected =
+      cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
+
+    REQUIRE_THAT(out[i], Equals(expected));
+  }
+  REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
+}
+} // namespace
+
+MULTI_GPU_TEST("inclusive_scan single-comm documentation example", c2h::type_list<int>)
+{
+  auto comms = this->communicators();
+
+  if (comms.size() < 2)
+  {
+    SKIP("The inclusive_scan documentation example requires at least two local GPUs");
+  }
+
+  auto streams_owned = nccl_test_util::make_streams();
+  auto streams       = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()};
+
+  // Must be pre-allocated since it is written to by threads
+  std::vector<std::string> failed(comms.front().size());
+
+  // Every communicator rank must invoke the collective concurrently.
+  run_threaded(comms.size(), [&](cuda::std::size_t i) {
+    auto& communicator = comms[i];
+    auto environment   = streams[i];
+    const auto device  = communicator.logical_device().underlying_device();
+
+    //! [inclusive_scan_single_range]
+    constexpr cuda::std::array input_values{1, 2};
+
+    auto input  = cuda::make_device_buffer<int>(environment, device, input_values);
+    auto output = cuda::make_device_buffer<int>(environment, device, input_values.size(), cuda::no_init);
+
+    cudax::inclusive_scan(cudax::distributed, communicator, environment, input, output.begin(), /*__init=*/0);
+
+    // Every rank contributes {1, 2}, so rank r starts with a prefix of 3 * r.
+    const auto rank = communicator.rank();
+    const auto expected =
+      cuda::make_buffer<int>(output.stream(), cuda::mr::legacy_pinned_memory_resource{}, {3 * rank + 1, 3 * rank + 3});
+
+    //! [inclusive_scan_single_range]
+
+    // catch2 isn't thread safe by default, so we can't use the usual requires expression. So
+    // we roll a hacky version of it ourselves
+    if (const auto matcher = Equals(expected); !matcher.match(output))
+    {
+      failed[rank] = matcher.describe();
+    }
+  });
+
+  for (cuda::std::size_t i = 0; i < failed.size(); ++i)
+  {
+    if (const auto& err_str = failed[i]; !err_str.empty())
+    {
+      INFO("rank: " << i);
+      REQUIRE(err_str == ""); // Should print the full error string
+    }
+  }
+}
+
+MULTI_GPU_TEST("inclusive_scan single-comm, one element per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  std::vector<std::vector<T>> inputs_by_rank;
+
+  inputs_by_rank.reserve(static_cast<cuda::std::size_t>(comms.front().size()));
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    const auto v = {make_value<T>(r)};
+
+    inputs_by_rank.emplace_back(v);
+  }
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}
+
+MULTI_GPU_TEST("inclusive_scan single-comm, multiple elements per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  constexpr auto values_per_rank = 10;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    const auto value                                  = make_value<T>(r);
+    inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, value);
+  }
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}
+
+MULTI_GPU_TEST("inclusive_scan single-comm, some ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  constexpr auto values_per_rank = 10;
+  std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+  for (int r = 0; r < comms.front().size(); ++r)
+  {
+    if (r % 2 == 0)
+    {
+      inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
+    }
+  }
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}
+
+MULTI_GPU_TEST("inclusive_scan single-comm, all ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+  auto comms       = this->communicators();
+
+  const std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
+
+  run_case(comms, inputs_by_rank, init, ident, Op{});
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/inclusive_scan/single_comm_defaults.cu

@@ -0,0 +1,120 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/iterator>
+
+#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+MULTI_GPU_TEST("inclusive_scan single-comm, overloads default values", )
+{
+  using T  = cuda::std::int32_t;
+  using Op = ::cuda::std::plus<>;
+
+  constexpr auto init = T{};
+  constexpr T ident   = cuda::identity_element<Op, T>();
+  constexpr auto op   = Op{};
+
+  const auto comms   = this->communicators();
+  const auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<cuda::stream_ref> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    std::vector<T> values(values_per_rank);
+    std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(cuda::make_device_buffer<T>(
+      streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
+    envs.emplace_back(streams[i]);
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  const auto expected_values = [&] {
+    std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
+
+    std::iota(reference.begin(), reference.end(), T{1});
+
+    std::vector<T> expected_values(reference.size());
+
+    std::inclusive_scan(reference.begin(), reference.end(), expected_values.begin(), op, init);
+
+    return expected_values;
+  }();
+
+  const auto check_outputs = [&] {
+    const auto exp_span = cuda::std::span{expected_values};
+
+    for (cuda::std::size_t i = 0; i < out.size(); ++i)
+    {
+      const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
+      const auto expected =
+        cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
+
+      REQUIRE_THAT(out[i], Equals(expected));
+    }
+  };
+
+  SECTION("Default init, op, ident (all)")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i]);
+    });
+    check_outputs();
+  }
+
+  SECTION("Default op, ident")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init);
+    });
+    check_outputs();
+  }
+
+  SECTION("Default ident")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op);
+    });
+    check_outputs();
+  }
+
+  SECTION("Default none")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
+    });
+    check_outputs();
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/reduce/CMakeLists.txt

@@ -0,0 +1,20 @@
+#===----------------------------------------------------------------------===##
+#
+# Part of CUDA Experimental in CUDA C++ Core Libraries,
+# under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+#
+#===----------------------------------------------------------------------===##
+
+if (NOT cudax_ENABLE_NCCL)
+  return()
+endif()
+
+file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
+
+foreach (src IN LISTS test_srcs)
+  cudax_add_multi_gpu_test("algorithms.reduce" test_target "${src}")
+  target_link_libraries(${test_target} PRIVATE NCCL::nccl)
+endforeach()

cccl_upstream/cudax/test/multi_gpu/algorithms/reduce/range_basic.cu

@@ -0,0 +1,341 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+namespace
+{
+struct custom_plus
+{
+  template <class T>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
+  {
+    return lhs + rhs;
+  }
+};
+
+using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
+using value_types  = c2h::type_list<cuda::std::int32_t, float, custom_value>;
+using operators    = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
+
+static_assert(cudax::nccl_transportable<custom_value>);
+
+template <typename T>
+T make_value(int i)
+{
+  return static_cast<T>(i);
+}
+
+template <>
+custom_value make_value<>(int i)
+{
+  custom_value ret{};
+
+  ret.key = static_cast<std::size_t>(i);
+  ret.val = static_cast<std::size_t>(i);
+  return ret;
+};
+
+template <class T, class Op>
+[[nodiscard]] T get_identity()
+{
+  if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
+  {
+    return make_value<T>(0);
+  }
+  else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
+  {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+  else
+  {
+    static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
+  }
+}
+
+// Run the full reduction, wait for it to finish, and check that `reduce` left its argument ranges
+// untouched. This boilerplate is identical for every test regardless of how the inputs are shaped.
+template <class Env, class T, class Op>
+void do_reduce(cuda::std::span<cudax::nccl_communicator_ref> comms,
+               const std::vector<Env>& envs,
+               std::vector<cuda::device_buffer<T>>& in,
+               std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
+               const T& init,
+               const T& ident,
+               Op op)
+{
+  const auto envs_size    = envs.size();
+  const auto in_copy      = in;
+  const auto outputs_copy = outputs;
+
+  INFO("init = " << init);
+  INFO("ident = " << ident);
+
+  cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init, op, ident);
+
+  // cuda::std::execution::env has no operator==, so we can only compare the sizes.
+  REQUIRE(envs.size() == envs_size);
+  // Reduction call should not modify the inputs in any ways
+  REQUIRE(in.size() == in_copy.size());
+  for (cuda::std::size_t i = 0; i < in.size(); ++i)
+  {
+    INFO("device = " << i);
+    REQUIRE_THAT(in[i], Equals(in_copy[i]));
+  }
+  REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
+}
+} // namespace
+
+MULTI_GPU_TEST("reduce, one element per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each reduction with a few hardcoded initializers. The init participates in the fold the
+  // same way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
+  // one-element output buffer and an environment carrying its stream, so the reduction is
+  // stream-ordered on the correct device. `reference` mirrors the contributions of every global
+  // rank so we can fold them on the host exactly like `reduce` does on the device.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto values = {make_value<T>(comms[i].rank())};
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce(comms, envs, in, outputs, init, ident, Op{});
+
+  const T expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size());
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      reference.push_back(make_value<T>(r));
+    }
+
+    return std::accumulate(reference.begin(), reference.end(), init, Op{});
+  }();
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("reduce, multiple elements per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each reduction with a few hardcoded initializers. The init participates in the fold the
+  // same way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes ten copies of `rank`. `reduce` first does a local
+  // CUB reduction of each rank's range, then combines the partials across ranks. Each local rank
+  // also gets a one-element output buffer and an environment carrying its stream. `reference`
+  // mirrors every global rank's ten contributions for the host-side fold.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto v = make_value<T>(comms[i].rank());
+    const std::vector<T> values(values_per_rank, v);
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce(comms, envs, in, outputs, init, ident, Op{});
+
+  const T expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size() * values_per_rank);
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      const auto v = make_value<T>(r);
+
+      reference.insert(reference.end(), values_per_rank, v);
+    }
+
+    return std::accumulate(reference.begin(), reference.end(), init, Op{});
+  }();
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("reduce, some ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
+  // range. Rank 0 (the reduction root) is always non-empty. `reduce` must treat an empty rank as
+  // contributing nothing, exactly like `std::accumulate` over the surviving elements. `reference`
+  // mirrors that for the host-side fold.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto rank = comms[i].rank();
+    if (rank % 2 == 0)
+    {
+      const std::vector<T> values(values_per_rank, make_value<T>(rank));
+      in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    }
+    else
+    {
+      in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
+    }
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce(comms, envs, in, outputs, init, ident, Op{});
+
+  const T expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size() * values_per_rank);
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      if (r % 2 == 0)
+      {
+        reference.insert(reference.end(), values_per_rank, make_value<T>(r));
+      }
+    }
+
+    return std::accumulate(reference.begin(), reference.end(), init, Op{});
+  }();
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("reduce, all ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const auto init  = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // No rank contributes any element. Reducing nothing seeded by `init` is just `init`, so every
+  // output must equal `init` regardless of the operator.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce(comms, envs, in, outputs, init, ident, Op{});
+
+  // Reducing nothing seeded by `init` yields `init`, exactly like `std::accumulate` over an empty
+  // range.
+  const T expected = init;
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/reduce/range_defaults.cu

@@ -0,0 +1,112 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/memory_resource>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
+
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+MULTI_GPU_TEST("reduce, range overloads default values", )
+{
+  using T  = cuda::std::int32_t;
+  using Op = ::cuda::std::plus<>;
+
+  constexpr auto init = T{};
+  constexpr T ident   = cuda::identity_element<Op, T>();
+  constexpr auto op   = Op{};
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const std::vector<T> values(values_per_rank, static_cast<T>(comms[i].rank()));
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  const auto expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size() * values_per_rank);
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      reference.insert(reference.end(), values_per_rank, static_cast<T>(r));
+    }
+
+    const auto val = std::accumulate(reference.begin(), reference.end(), init, op);
+
+    return cuda::make_buffer(cuda::stream_ref{::CUstream{}}, cuda::mr::legacy_pinned_memory_resource{}, 1, val);
+  }();
+
+  SECTION("Default init, op, ident (all)")
+  {
+    cudax::reduce(cudax::broadcasted, comms, envs, in, outputs);
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+
+  SECTION("Default op, ident")
+  {
+    cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init);
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+
+  SECTION("Default ident")
+  {
+    cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init, op);
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+
+  SECTION("Default none")
+  {
+    cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init, op, ident);
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/reduce/single_comm_basic.cu

@@ -0,0 +1,347 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/functional>
+#include <cuda/memory_resource>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/functional>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
+
+#include <exception>
+#include <future>
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+namespace
+{
+struct custom_plus
+{
+  template <class T>
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
+  {
+    return lhs + rhs;
+  }
+};
+
+using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
+using value_types  = c2h::type_list<cuda::std::int32_t, float, custom_value>;
+using operators    = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
+
+static_assert(cudax::nccl_transportable<custom_value>);
+
+template <typename T>
+T make_value(int i)
+{
+  return static_cast<T>(i);
+}
+
+template <>
+custom_value make_value<>(int i)
+{
+  custom_value ret{};
+
+  ret.key = static_cast<std::size_t>(i);
+  ret.val = static_cast<std::size_t>(i);
+  return ret;
+};
+
+template <class T, class Op>
+[[nodiscard]] T get_identity()
+{
+  if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
+  {
+    return make_value<T>(0);
+  }
+  else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
+  {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+  else
+  {
+    static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
+  }
+}
+
+// Drive the reduction through the single-communicator overload of `reduce`, one thread per
+// rank. That overload opens its own NCCL group on a single communicator, so issuing the
+// per-rank calls serially on one thread would deadlock at `ncclGroupEnd`. Running each rank on
+// its own thread lets the per-thread groups rendezvous across ranks. Only the `reduce` call
+// happens on the worker threads; every Catch2 assertion runs on the main thread after the
+// join, since the assertion macros are not safe to fire concurrently.
+template <class Env, class T, class Op>
+void do_reduce_threaded(
+  cuda::std::span<cudax::nccl_communicator_ref> comms,
+  std::vector<Env>& envs,
+  std::vector<cuda::device_buffer<T>>& in,
+  std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
+  const T& init,
+  const T& ident,
+  Op op)
+{
+  const auto in_copy      = in;
+  const auto outputs_copy = outputs;
+
+  INFO("init = " << init);
+  INFO("ident = " << ident);
+
+  run_threaded(comms.size(), [&](cuda::std::size_t i) {
+    cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init, op, ident);
+  });
+
+  // Reduction call should not modify the inputs in any ways
+  REQUIRE(in.size() == in_copy.size());
+  for (cuda::std::size_t i = 0; i < in.size(); ++i)
+  {
+    INFO("device = " << i);
+    REQUIRE_THAT(in[i], Equals(in_copy[i]));
+  }
+  REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
+}
+} // namespace
+
+MULTI_GPU_TEST("reduce single-comm, one element per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each reduction with a few hardcoded initializers. The init participates in the fold the
+  // same way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
+  // one-element output buffer and an environment carrying its stream, so the reduction is
+  // stream-ordered on the correct device. `reference` mirrors the contributions of every global
+  // rank so we can fold them on the host exactly like `reduce` does on the device.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto values = {make_value<T>(comms[i].rank())};
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
+
+  const T expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size());
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      reference.push_back(make_value<T>(r));
+    }
+
+    return std::accumulate(reference.begin(), reference.end(), init, Op{});
+  }();
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("reduce single-comm, multiple elements per rank", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  // Seed each reduction with a few hardcoded initializers. The init participates in the fold the
+  // same way on host and device, so any value works for every operator under test.
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Global rank `comms[i].rank()` contributes ten copies of `rank`. `reduce` first does a local
+  // CUB reduction of each rank's range, then combines the partials across ranks. Each local rank
+  // also gets a one-element output buffer and an environment carrying its stream. `reference`
+  // mirrors every global rank's ten contributions for the host-side fold.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto v = make_value<T>(comms[i].rank());
+    const std::vector<T> values(values_per_rank, v);
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
+
+  const T expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size() * values_per_rank);
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      const auto v = make_value<T>(r);
+
+      reference.insert(reference.end(), values_per_rank, v);
+    }
+
+    return std::accumulate(reference.begin(), reference.end(), init, Op{});
+  }();
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("reduce single-comm, some ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const T init     = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
+  // range. Rank 0 (the reduction root) is always non-empty. `reduce` must treat an empty rank as
+  // contributing nothing, exactly like `std::accumulate` over the surviving elements. `reference`
+  // mirrors that for the host-side fold.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const auto rank = comms[i].rank();
+    if (rank % 2 == 0)
+    {
+      const std::vector<T> values(values_per_rank, make_value<T>(rank));
+      in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    }
+    else
+    {
+      in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
+    }
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
+
+  const T expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size() * values_per_rank);
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      if (r % 2 == 0)
+      {
+        reference.insert(reference.end(), values_per_rank, make_value<T>(r));
+      }
+    }
+
+    return std::accumulate(reference.begin(), reference.end(), init, Op{});
+  }();
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}
+
+MULTI_GPU_TEST("reduce single-comm, all ranks empty", value_types, operators)
+{
+  using T  = c2h::get<0, TestType>;
+  using Op = c2h::get<1, TestType>;
+
+  const auto init  = make_value<T>(GENERATE(0, 1, -1, 5));
+  const auto ident = get_identity<T, Op>();
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  // No rank contributes any element. Reducing nothing seeded by `init` is just `init`, so every
+  // output must equal `init` regardless of the operator.
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
+
+  // Reducing nothing seeded by `init` yields `init`, exactly like `std::accumulate` over an empty
+  // range.
+  const T expected = init;
+
+  for (const auto& buf : out)
+  {
+    const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
+
+    REQUIRE_THAT(buf, Equals(exp));
+  }
+}

cccl_upstream/cudax/test/multi_gpu/algorithms/reduce/single_comm_defaults.cu

@@ -0,0 +1,122 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/buffer>
+#include <cuda/memory_resource>
+#include <cuda/std/cstdint>
+#include <cuda/std/execution>
+#include <cuda/std/type_traits>
+
+#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
+
+#include <exception>
+#include <future>
+#include <numeric>
+#include <vector>
+
+#include <algorithm_common.h>
+#include <nccl_test_common.h>
+#include <testing.cuh>
+
+MULTI_GPU_TEST("reduce single-comm, overloads default values", )
+{
+  using T  = cuda::std::int32_t;
+  using Op = ::cuda::std::plus<>;
+
+  constexpr auto init = T{};
+  constexpr T ident   = cuda::identity_element<Op, T>();
+  constexpr auto op   = Op{};
+
+  auto comms   = this->communicators();
+  auto streams = nccl_test_util::make_streams();
+
+  std::vector<cuda::device_buffer<T>> in;
+  std::vector<cuda::device_buffer<T>> out;
+  std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
+
+  in.reserve(comms.size());
+  out.reserve(comms.size());
+  envs.reserve(comms.size());
+
+  constexpr auto values_per_rank = 10;
+  for (cuda::std::size_t i = 0; i < comms.size(); ++i)
+  {
+    const std::vector<T> values(values_per_rank, static_cast<T>(comms[i].rank()));
+
+    in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
+    out.emplace_back(
+      cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
+    envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
+  }
+
+  auto outputs = make_output_iterators(out);
+
+  const auto expected = [&] {
+    std::vector<T> reference;
+
+    reference.reserve(comms.front().size() * values_per_rank);
+    for (int r = 0; r < comms.front().size(); ++r)
+    {
+      reference.insert(reference.end(), values_per_rank, static_cast<T>(r));
+    }
+
+    const auto val = std::accumulate(reference.begin(), reference.end(), init, op);
+
+    return cuda::make_buffer(cuda::stream_ref{::CUstream{}}, cuda::mr::legacy_pinned_memory_resource{}, 1, val);
+  }();
+
+  SECTION("Default init, op, ident (all)")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i]);
+    });
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+
+  SECTION("Default op, ident")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init);
+    });
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+
+  SECTION("Default ident")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init, op);
+    });
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+
+  SECTION("Default none")
+  {
+    run_threaded(comms.size(), [&](cuda::std::size_t i) {
+      cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init, op, ident);
+    });
+
+    for (const auto& buf : out)
+    {
+      REQUIRE_THAT(buf, Equals(expected));
+    }
+  }
+}