Implement Segmented TopK using Thread-Block-Clusters

.clang-format

@@ -40,7 +40,9 @@ AttributeMacros: [
                   '_CCCL_TYPE_VISIBILITY_HIDDEN',
                   '_CCCL_VISIBILITY_HIDDEN',
                   '_CCCL_LAUNCH_BOUNDS',
+                  '_CCCL_LAUNCH_BOUNDS_CLUSTER',
                   '_CCCL_BLOCK_SIZE',
+                  '_CCCL_CLUSTER_DIMS',
                   'CUB_RUNTIME_FUNCTION',
                   'THRUST_RUNTIME_FUNCTION',
                   'CCCL_DEPRECATED',

cub/benchmarks/bench/segmented_topk/fixed/keys.cu

@@ -1,25 +1,95 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+// Defer the unsupported-architecture diagnosis to the dispatch's runtime check so this benchmark compiles for the full
+// configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves)
+// across all target architectures, including pre-SM90. See CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in
+// cub/device/device_batched_topk.cuh. Must precede the CUB includes below.
+#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT
+
 #include <cub/detail/choose_offset.cuh>
+#include <cub/device/device_batched_topk.cuh>
+#include <cub/device/device_topk.cuh>
 #include <cub/device/dispatch/dispatch_batched_topk.cuh>
 
+#include <cuda/__execution/determinism.h>
+#include <cuda/__execution/output_ordering.h>
+#include <cuda/__execution/require.h>
+#include <cuda/__execution/tie_break.h>
+#include <cuda/__execution/tune.h>
 #include <cuda/argument>
 #include <cuda/iterator>
+#include <cuda/std/__execution/env.h>
+
+#include <algorithm>
+#include <vector>
 
 #include <nvbench_helper.cuh>
 
 // %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1
 // %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32
 // %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1
+// %RANGE% TUNE_BACKEND backend 0:1:1
+
+enum class topk_backend
+{
+  baseline,
+  cluster,
+  device,
+  automatic,
+};
+
+// Which backend this build benchmarks. `automatic` (the default) issues no `tune` override, leaving the choice to the
+// library's arch/size selector. `baseline`/`cluster` force one of the two DeviceBatchedTopK backends via the `tune`d
+// selector below; `device` is a reference that issues one `cub::DeviceTopK` call per segment. Autotuning sweeps the
+// baseline and cluster backends (the `%RANGE% ... 0:1:1` above); only baseline knobs are exposed here, so the cluster
+// backend uses its default sub-policy and `device` has no knobs. Override with -DTUNE_BACKEND=0/1 (force a backend),
+// =2 (device reference), or =3 (automatic).
+#ifndef TUNE_BACKEND
+#  if TUNE_BASE
+#    define TUNE_BACKEND 3 // automatic: the library's production selector, for base/benchmark builds
+#  else
+#    define TUNE_BACKEND 0 // force baseline when an actively-tuned variant does not sweep the backend
+#  endif
+#endif
+
+inline constexpr topk_backend selected_backend =
+#if TUNE_BACKEND == 0
+  topk_backend::baseline;
+#elif TUNE_BACKEND == 1
+  topk_backend::cluster;
+#elif TUNE_BACKEND == 2
+  topk_backend::device;
+#else
+  topk_backend::automatic;
+#endif
+
+// Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now).
+inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed;
+inline constexpr auto selected_tie_break   = cuda::execution::tie_break::__tie_break_t::__unspecified;
+
+// The baseline/device backends ignore these requirements, so require the defaults there to avoid a silent mismatch. The
+// cluster backend implements them, and automatic honors them via the library's selector, so both allow non-defaults.
+static_assert(selected_backend == topk_backend::cluster || selected_backend == topk_backend::automatic
+                || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed
+                    && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified),
+              "Only the cluster and automatic backends honor determinism/tie-break requirements; keep "
+              "selected_determinism and selected_tie_break at their defaults for the baseline/device backends.");
 
-#if !TUNE_BASE
-struct tuned_policy_selector
+// Policy selector threaded through the public API's tuning environment when a concrete backend is forced (not
+// `automatic`). Its `.backend` pins the backend for this build. In a TUNE_BASE build the forced backend uses the
+// default sub-policies; otherwise the baseline knobs come from the TUNE_* macros. The cluster sub-policy is always the
+// default (no cluster knobs are exposed here).
+template <class KeyT, class ValueT, class OffsetT, cuda::std::int64_t MaxK>
+struct topk_backend_selector
 {
-  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-    -> cub::detail::batched_topk::batched_topk_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability cc) const
+    -> cub::detail::batched_topk::topk_policy
   {
-    // Single-entry policy chain driven by the tuning knobs.
+#if TUNE_BASE
+    const auto baseline =
+      cub::detail::batched_topk::baseline_policy_selector_from_types<KeyT, ValueT, OffsetT, MaxK>{}(cc);
+#else
     constexpr auto store_alg = cub::BLOCK_STORE_WARP_TRANSPOSE;
 #  if TUNE_BLOCK_LOAD_ALGORITHM == 0
     constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT;
@@ -28,17 +98,97 @@ struct tuned_policy_selector
 #  elif TUNE_BLOCK_LOAD_ALGORITHM == 2
     constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE;
 #  endif
-    return cub::detail::batched_topk::batched_topk_policy{{{
+    const auto baseline = cub::detail::batched_topk::baseline_topk_policy{{{
       cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
       cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
       cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
       cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
       cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
       cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
     }}};
+#endif // TUNE_BASE
+    const auto cluster = cub::detail::batched_topk::cluster_policy_selector{}(cc);
+    constexpr auto backend =
+      (selected_backend == topk_backend::cluster)
+        ? cub::detail::batched_topk::topk_backend::cluster
+        : cub::detail::batched_topk::topk_backend::baseline;
+    return cub::detail::batched_topk::topk_policy{backend, baseline, cluster};
   }
 };
-#endif // !TUNE_BASE
+
+// Env-based dispatch over the selected backend. `automatic`/`baseline`/`cluster` all route through the public
+// `cub::DeviceBatchedTopK` API (the latter two add a `tune`d `topk_backend_selector` that forces the backend; temp
+// storage comes from the memory resource carried by `env`); the `device` backend issues one `cub::DeviceTopK::MaxKeys`
+// per segment, reading the host-side segment sizes.
+template <typename KeyInputItItT,
+          typename KeyOutputItItT,
+          typename SegmentSizeParamT,
+          typename KParamT,
+          typename NumSegmentsParameterT,
+          typename HostSegSizeT,
+          typename EnvT>
+CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys(
+  KeyInputItItT d_keys_in,
+  KeyOutputItItT d_keys_out,
+  SegmentSizeParamT segment_sizes,
+  KParamT k,
+  NumSegmentsParameterT num_segments,
+  [[maybe_unused]] const HostSegSizeT* h_segment_sizes,
+  EnvT env)
+{
+  if constexpr (selected_backend == topk_backend::device)
+  {
+    using num_segments_val_t = typename ::cuda::args::__traits<NumSegmentsParameterT>::element_type;
+    const auto num_segs      = cub::detail::params::get_param(num_segments, num_segments_val_t{0});
+
+    // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement
+    // on top of the benchmark environment (which carries the stream and the caching memory resource).
+    const auto seg_env = cuda::std::execution::env{
+      env,
+      cuda::execution::require(cuda::execution::determinism::not_guaranteed,
+                               cuda::execution::output_ordering::unsorted)};
+
+    for (num_segments_val_t i = 0; i < num_segs; ++i)
+    {
+      const auto k_value  = cub::detail::params::get_param(k, i);
+      const auto seg_size = h_segment_sizes[i];
+      if (const auto err = cub::DeviceTopK::MaxKeys(
+            d_keys_in[i],
+            d_keys_out[i],
+            static_cast<cuda::std::int64_t>(seg_size),
+            static_cast<cuda::std::int64_t>(k_value),
+            seg_env);
+          err != cudaSuccess)
+      {
+        return err;
+      }
+    }
+    return cudaSuccess;
+  }
+  else
+  {
+    // The determinism / tie-break / ordering requirement this benchmark issues; the library honors it whether or not we
+    // additionally force a backend.
+    const auto req_env = cuda::std::execution::env{
+      env,
+      cuda::execution::require(cuda::execution::determinism::__determinism_holder_t<selected_determinism>{},
+                               cuda::execution::tie_break::__tie_break_holder_t<selected_tie_break>{},
+                               cuda::execution::output_ordering::unsorted)};
+    if constexpr (selected_backend == topk_backend::automatic)
+    {
+      // No `tune` override: the library's own selector picks the backend (arch/size crossover) -- the usual behavior.
+      return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, req_env);
+    }
+    else
+    {
+      using key_t          = cub::detail::it_value_t<cub::detail::it_value_t<KeyInputItItT>>;
+      constexpr auto max_k = ::cuda::args::__traits<KParamT>::highest;
+      const auto full_env  = cuda::std::execution::env{
+        req_env, cuda::execution::tune(topk_backend_selector<key_t, cub::NullType, cuda::std::int64_t, max_k>{})};
+      return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, full_env);
+    }
+  }
+}
 
 template <typename KeyT, int MaxSegmentSize, int MaxNumSelected>
 void fixed_seg_size_topk_keys(
@@ -51,7 +201,6 @@ void fixed_seg_size_topk_keys(
   const auto selected_elements = static_cast<::cuda::std::ptrdiff_t>(MaxNumSelected);
   const auto num_segments      = ::cuda::std::max<std::size_t>(1, (max_elements / segment_size));
   const auto elements          = num_segments * segment_size;
-  const auto total_num_items   = ::cuda::args::immediate{static_cast<::cuda::std::int64_t>(elements)};
   const bit_entropy entropy    = str_to_entropy(state.get_string("Entropy"));
 
   // Skip workloads where k exceeds the segment size
@@ -63,44 +212,36 @@ void fixed_seg_size_topk_keys(
 
   thrust::device_vector<KeyT> in_keys_buffer = generate(elements, entropy);
   thrust::device_vector<KeyT> out_keys_buffer(selected_elements * num_segments, thrust::no_init);
-  auto d_keys_in_ptr  = thrust::raw_pointer_cast(in_keys_buffer.data());
-  auto d_keys_out_ptr = thrust::raw_pointer_cast(out_keys_buffer.data());
-  auto d_keys_in      = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size);
-  auto d_keys_out     = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), selected_elements);
+  const auto d_keys_in_ptr  = thrust::raw_pointer_cast(in_keys_buffer.data());
+  const auto d_keys_out_ptr = thrust::raw_pointer_cast(out_keys_buffer.data());
+  const auto d_keys_in      = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size);
+  const auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), selected_elements);
 
-  auto segment_sizes    = ::cuda::args::constant<MaxSegmentSize>{};
-  auto k                = ::cuda::args::constant<MaxNumSelected>{};
-  auto select_direction = ::cuda::args::constant<cub::detail::topk::select::max>{};
+  const auto segment_sizes = ::cuda::args::constant<MaxSegmentSize>{};
+  const auto k             = ::cuda::args::constant<MaxNumSelected>{};
 
   state.add_element_count(elements, "NumElements");
   state.add_element_count(segment_size, "SegmentSize");
   state.add_element_count(selected_elements, "NumSelectedElements");
   state.add_global_memory_reads<KeyT>(elements, "InputKeys");
   state.add_global_memory_writes<KeyT>(selected_elements * num_segments, "OutputKeys");
 
+  // Host copy of segment sizes — all entries equal MaxSegmentSize for fixed-size segments. Consumed only by the
+  // per-segment device backend. Segment sizes fit in a signed 32-bit integer (the library caps them at 2^21).
+  const std::vector<cuda::std::int32_t> h_segment_sizes(num_segments, static_cast<cuda::std::int32_t>(MaxSegmentSize));
+
   caching_allocator_t alloc;
   state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
-    auto env = cub_bench_env(
-      alloc,
-      launch
-#if !TUNE_BASE
-      ,
-      cuda::execution::tune(tuned_policy_selector{})
-#endif // !TUNE_BASE
-    );
-    // TODO(bgruber): call the public API once available
+    const auto env = cub_bench_env(alloc, launch);
     _CCCL_TRY_CUDA_API(
-      cub::detail::batched_topk::dispatch_with_env,
+      batched_topk_keys,
       "batched topk failed",
       d_keys_in,
       d_keys_out,
-      static_cast<cub::NullType**>(nullptr),
-      static_cast<cub::NullType**>(nullptr),
       segment_sizes,
       k,
-      select_direction,
       ::cuda::args::immediate{static_cast<::cuda::std::int64_t>(num_segments)},
-      total_num_items,
+      h_segment_sizes.data(),
       env);
   });
 }

cub/benchmarks/bench/segmented_topk/variable/common.cuh

@@ -19,6 +19,11 @@
 
 #include <nvbench_helper.cuh>
 
+// Segment sizes are stored as signed 32-bit integers: the library caps supported segment sizes at 2^21 (about 2
+// million), so a wider type buys nothing here. The benchmarks always pass an explicit `cuda::args::bounds` (a bare
+// int32 would be rejected, as its type maximum exceeds the cap).
+using segment_size_t = cuda::std::int32_t;
+
 namespace
 {
 enum class pattern_kind : int
@@ -57,7 +62,7 @@ enum class pattern_kind : int
 
 template <int MaxSegmentSize, int K>
 [[nodiscard]] thrust::device_vector<float>
-gen_data(int num_segments, pattern_kind pattern, const cuda::std::int64_t* d_seg_sizes)
+gen_data(int num_segments, pattern_kind pattern, const segment_size_t* d_seg_sizes)
 {
   const auto num_keys = static_cast<std::size_t>(num_segments) * static_cast<std::size_t>(MaxSegmentSize);
   auto d_keys         = thrust::device_vector<float>{num_keys, thrust::no_init};

cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu

@@ -0,0 +1,32 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+// Cluster-backend variable-segment indexed (arg-top-k) benchmark. A base build runs the library's production
+// (automatic) selector; a tuning variant forces the cluster backend and sweeps its per-block knobs plus the
+// determinism / tie-break requirement. The baseline backend is benchmarked separately in the plain `indexed.cu`.
+
+// %RANGE% TUNE_CLUSTER_THREADS_PER_BLOCK tpb 128:1024:32
+// %RANGE% TUNE_CLUSTER_MIN_BLOCKS_PER_SM mbs 1:2:1
+// %RANGE% TUNE_CLUSTER_MIN_CHUNKS_PER_BLOCK mcb 1:2:1
+// %RANGE% TUNE_CLUSTER_CHUNK_KIB ckib 8:32:1
+// %RANGE% TUNE_CLUSTER_LOAD_ALIGN_BYTES_POW2 la 4:7:1
+// %RANGE% TUNE_CLUSTER_PIPELINE_STAGES ps 2:16:1
+// %RANGE% TUNE_CLUSTER_BITS_PER_PASS bpp 8:11:1
+// %RANGE% TUNE_CLUSTER_HIST_IPT hipt 1:24:1
+// %RANGE% TUNE_CLUSTER_TIEBREAK_IPT tipt 1:24:1
+// %RANGE% TUNE_CLUSTER_COPY_IPT cipt 1:24:1
+// %RANGE% TUNE_REQUIREMENT req 0:2:1
+
+#ifndef TUNE_BACKEND
+#  if TUNE_BASE
+#    define TUNE_BACKEND 2 // automatic: the library's production selector, for base/benchmark builds
+#  else
+#    define TUNE_BACKEND 1 // cluster: the backend this file tunes
+#  endif
+#endif
+
+#ifndef TUNE_REQUIREMENT
+#  define TUNE_REQUIREMENT 1 // deterministic + prefer-smaller-index (forward): the safe, least-surprising default
+#endif
+
+#include "indexed_common.cuh"