Implement Segmented TopK using Thread-Block-Clusters

[pauleonix]

Description

closes #9077 #9259 #9548

Cluster-based Segmented Top-K (CUB)

This PR adds a thread-block-cluster implementation of segmented top-k to CUB and unifies it
with the existing baseline backend behind CUB's public cub::DeviceBatchedTopK API. Given
num_segments segments, each of variable size, the algorithm selects the k largest (or smallest)
keys of every segment (optionally carrying a value payload, i.e. key/value pairs).

This is a production backend, not a throwaway prototype. cub::DeviceBatchedTopK now exposes two
backends behind one kernel symbol and picks between them per architecture (the compile-time,
device-side policy_selector, resolved via current_policy + if constexpr):

• baseline — the pre-existing worker-per-segment backend (one thread block per segment). It
  serves segments that fit a single thread block on all architectures, but currently only the
  fully non-deterministic request (not_guaranteed, unspecified).
• cluster — this PR's backend, for SM 9.0+. It handles segments too large for a single block
  and every supported determinism / tie-break combination.

The backend is chosen from three request facts: the architecture, the statically-known maximum
segment size, and the determinism / tie-break requirement. Deterministic (or tie-broken)
requests, and segments larger than the baseline can cover, route to the cluster backend (SM 9.0+);
among non-deterministic, baseline-coverable requests the baseline-vs-cluster choice is an
architecture / segment-size crossover. When no backend can serve the request on a target
architecture the dispatch reports it either at compile time (strict, default) or at runtime as
cudaErrorNotSupported (see §3.6, §10).

Because this is CUB's first cluster kernel, this document spends extra time on the
cluster-specific mechanics — distributed shared memory, cluster barriers (and how few of them we
use), the leader block, dynamic cluster sizing, and portability — which reviewers will not have
seen elsewhere in the codebase.

---

1. High-level overview

One cluster per segment · radix digit top-k · histogram merge · early stop · output placement · determinism

1.1 One cluster per segment, sized for portability

The kernel launches a grid of thread-block clusters and assigns exactly one cluster to one
segment. The blocks (CTAs) of a cluster co-operate through distributed shared memory (DSMEM)
and cluster barriers, so a segment is processed entirely on-chip without a global scratch
buffer or multiple kernel launches. (The non-segmented cub::DeviceTopK instead uses a multi-kernel

• global-histogram pipeline; the baseline worker-per-segment backend keeps each segment within a
  single block.)

The number of CTAs per cluster (the "cluster width") is chosen at runtime by the dispatch
layer, so it is not a template parameter; the agent reads it from
cooperative_groups::this_cluster().num_blocks(). The primary reason to spread a segment across a
runtime-sized cluster (rather than baking in a fixed size) is portability: the resident
shared-memory budget that lets a segment of a given size stay on-chip is not guaranteed on future
GPUs, so we size the cluster to whatever the device actually offers. Being able to fold the
runtime maximum-segment-size information into the dispatch decision is a welcome bonus, but not the
main motivation. For the same reason the resident key storage lives in dynamic shared memory
(Section 3.2) rather than a compile-time-sized static buffer.

1.2 Radix (digit) top-k

Selection uses the standard multi-pass radix / digit-iteration approach shared with the rest
of CUB's top-k (detail::topk): process the key one bits_per_pass-bit digit at a time, from the
most-significant digit down. Each pass:

1. Builds a histogram over the current digit, restricted to candidate keys (keys whose higher
   digits already match the running "k-th key" prefix).
2. Finds the bucket that contains the k-th element and narrows k and the candidate set to that
   bucket.
3. Folds the winning digit into a running splitter key.

After the passes converge, the splitter key (the k-th largest key) is known, and a final
filter pass writes out:

• every key strictly above the splitter (the "front"), plus
• exactly enough keys equal to the splitter (the "back"/tie region) to reach k.

flowchart LR
    P0["Pass 0 (MSD digit)<br/>candidates = all keys"] --> P1["Pass 1<br/>candidates = k-th bucket"]
    P1 --> P2["Pass 2<br/>candidate set shrinks"]
    P2 --> Pn["Pass n (LSD digit)<br/>splitter (k-th key) known"]
    Pn --> F["final filter<br/>front (&gt; splitter) + back (= splitter)"]
    P1 -. "bucket holds exactly k → early stop" .-> F

Loading

1.3 The histogram: block-private accumulation + DSMEM merge into the leader

Each pass builds the cluster-wide histogram in three steps:

1. Block-private accumulation. Every CTA lays out hist[num_buckets] at the same offset in
   its own shared memory and accumulates its own keys into it using shared-memory reduction
   atomics.
2. DSMEM merge into the leader. Every non-leader CTA walks its local histogram and folds each
   nonzero bucket into the leader CTA's hist through DSMEM. The leader's hist therefore
   does double duty: its own block-private histogram first, then the cluster-merged histogram after
   the merge.
3. Splitter identification. The leader's threads prefix-sum the merged histogram
   (cub::BlockScan), find the bucket holding the k-th key, and publish the result into a
   cluster-shared state. Every CTA reads that result back and folds the winning digit into its
   own local splitter key (so the full splitter key never has to be broadcast).

flowchart TD
    subgraph bp["1 · block-private accumulation"]
      A["CTA 0<br/>local hist"]
      B["CTA 1<br/>local hist"]
      C["CTA n<br/>local hist"]
    end
    L["leader CTA hist<br/>cluster-scope atomics<br/>(own hist, then merged hist)"]
    A -- "2 · DSMEM fold" --> L
    B -- "2 · DSMEM fold" --> L
    C -- "2 · DSMEM fold" --> L
    L --> S["3 · BlockScan merged hist<br/>find k-th bucket"]
    S --> R["publish result_pair<br/>(kth_bucket | early_stop)"]
    R -. "read back via DSMEM,<br/>fold digit into own splitter" .-> A
    R -.-> B
    R -.-> C

Loading

1.4 Early stop

When the identified splitter bucket contains exactly the remaining k candidates, every
candidate in it is part of the answer and no finer digit can change the result. The leader sets an
early_stop flag; every CTA decodes it from the same broadcast word and breaks out of the pass
loop together.

1.5 Final filter and output placement

The final pass re-reads the segment's keys and routes each through one shared placement helper
(place_one): the front (strictly-selected) and back (tie) paths are unified — a single
block-local SMEM atomicAdd on the relevant region counter yields the key's output slot, and a
uniform out < k guard drops the losing ties while always accepting a strictly-selected key. Both
regions fill forward (low index up):

• Front (strictly above the splitter) → the front region [0, num_selected).
• Back (ties equal to the splitter) → the back region [num_selected, k) (i.e. [k - num_ties, k)).

The trick is that the two region counters are pre-seeded to this CTA's absolute region base, so
the placing atomicAdd returns the final output slot with no per-key base arithmetic. Computing
those bases is the one cross-CTA step, done by a cross-CTA prefix scan: a handful of one-time
cross-CTA pushes seed every CTA's front and back counters to their region bases at once (Section 2.5),
after which per-element placement is just the bare block-local SMEM atomicAdd above, fully
decoupling the CTAs (no shared cursor).

(An earlier design instead kept a DSMEM output cursor: each output element claimed its slot with
a fetch-add into a cluster-wide cursor held in the leader via DSMEM atomics — one DSMEM atomic per
placed key. Section 2.5 keeps the comparison because it explains why the scan is the better shape.)

The scan puts its (single) cluster barrier before the filter (a post-push sync), so there is no
cluster barrier after the final filter pass and CTAs exit independently — unlike the old DSMEM
cursor, which needed a cluster barrier after the filter (whose later timing also gated early exit).

Output array (length k):

  0                                                            k-1
  |<------------- front region --------------->|<--- back / ties --->|
  [ keys strictly > splitter, packed by CTA    ][  keys == splitter   ]
                                                ^                     ^
                                           k - num_ties               k

  both regions fill forward; each CTA's front/back counters are pre-seeded to its
  region base by the cross-CTA scan

1.6 Determinism / tie-breaking

The public request carries a cuda::execution::determinism requirement and an optional tie_break
preference:

• Non-deterministic (default): both front and back use racing shared-memory atomics; among
  equal (tied) keys, which ones land in the back is arrival-order (nondeterministic), but the
  set of winning keys is always correct.
• Deterministic: the back region is filled in segment-index order via an index-ordered
  BlockScan, so a reproducible subset of the tied keys wins. prefer_larger_index reverses the
  scan so the largest-index ties win instead of the smallest.

The determinism contract is that the set of winning elements is deterministic, not their output
positions. The user can also request weaker, valid intermediate points on the determinism /
tie-break lattice (e.g. run-to-run determinism, or GPU-to-GPU determinism with an unspecified
tie-break). For now those all map onto the two behaviors above, because we do not yet know a
cheaper way to provide them.

---

2. Cluster mechanics (the first-cluster-kernel details)

Cluster barriers · DSMEM and state · leader and result_pair · idle ranks · cross-CTA output scan

This is the section most worth close review, since these patterns are new to CUB.

2.1 Cluster barriers and how we minimize them

A cluster barrier (cluster.sync()) is expensive: unlike __syncthreads(), synchronizing
across CTAs generally requires a round-trip through L2. Keeping their count low is a first-order
concern, so several design choices exist purely to avoid cluster barriers:

• Front-load all local shared-memory initialization before the first cluster barrier. Every
  counter the later phases rely on (histogram, the output-scan accumulator, the local front/back
  counters, the state fields) is zeroed before the launch's first cluster.sync(), so that one
  cluster barrier both orders the inits and publishes them cluster-wide. No later phase needs its own
  extra init-ordering cluster barrier.
• No cluster barrier between the block-private histogram and the merge. Naively, all CTAs would
  need a cluster barrier after building their local histograms before folding them into the leader
  (so the leader's own local increments are ordered ahead of the incoming folds). We drop that
  cluster barrier and instead make the leader's own histogram increments cluster-scoped atomics,
  so they are mutually atomic with the non-leaders' cluster-scoped folds. The leader pays wider-scope
  atomics on its own histogram, but the whole cluster saves one L2-round-trip cluster barrier per
  pass. (Non-leaders keep cheap CTA-scope atomics on their private histograms.)
• The cross-CTA output scan needs only one cluster barrier. Because the front/back placement
  counters (front_local_cnt/back_local_cnt) are pre-zeroed with the front-loaded inits above, the
  scan is a set of fire-and-forget pushes followed by a single post-push cluster barrier — there
  is no separate "start of scan" cluster barrier. (This is the only cluster barrier the scan adds;
  earlier notes about a "cluster-prefix barrier" refer to exactly this one post-push cluster barrier.)
• No cluster barrier after the final filter pass. The last cluster barrier is the post-push scan
  barrier that precedes the filter; the filter itself writes output through block-local atomics
  into gmem, touching no peer, so the last cross-CTA access is the already-fenced output-seeding step.
  Every value a CTA still needs from a peer afterward (notably the leader's early_stop) is cached
  into registers before that last cluster barrier, so no CTA reads a peer's shared memory during or
  after the filter and each can return independently. This co-residency safety matters: reading a peer
  that has already returned faults the launch ("cluster target block not present"). (This exact shape
  may change if the output seeding changes; see Section 2.5.)

The net per-pass cost is small: a local __syncthreads() plus two cluster barriers (one before
the leader reads the merged histogram, one that publishes the pass result / orders the histogram
reset), and one post-push cluster barrier for the output scan.

A single-CTA "cluster" (Section 3.4) has no peers, so cluster_or_block_sync(single_cta) routes
every one of these cluster barriers to a plain __syncthreads() and all atomics drop to CTA scope.

2.2 Distributed shared memory (DSMEM) and state

Every CTA allocates the identical _TempStorage layout, so any field can be reached in a peer CTA
at a known offset. Cross-CTA reads/writes are done through DSMEM (a CTA's shared memory mapped into
another rank's address space); the histogram merge and the output scan write into a peer's fields,
and the leader's state is read back by every CTA.

The cluster-shared cluster_topk_state lives in the leader block's shared memory and holds
len, k, and a packed result_pair (below). Non-leaders have no local copy — they need to read
it from the leader through DSMEM (and cache what they need).

2.3 The leader block and the result_pair broadcast

One CTA is the leader. It owns the merged histogram and the shared state, does the
BlockScan-based splitter identification, and publishes the per-pass result.

The per-pass result is packed into a single 64-bit result_pair:

• low 32 bits: the splitter kth_bucket (each CTA folds this digit into its own splitter key);
• high 32 bits: the early_stop flag.

Packing both into one naturally-aligned 64-bit word means every CTA pulls the pass result with a
single DSMEM load and decodes both halves locally — no second remote load for the early-stop
check, and no broadcast of the full splitter key.

The leader is always the CTA that is last in the cross-CTA scan order. This is required, not
arbitrary: the leader's histogram has been merged with everyone else's, so it cannot recover its
own local front/candidate counts directly. Sitting last in the scan, it receives the sum over all
other CTAs and derives its own counts by subtracting that sum from the known totals. (For the
deterministic-prefer-smaller case "last" is the highest effective rank; otherwise it is rank 0.)

2.4 Idle ranks

The launch is sized for the largest segment a cluster may have to handle, so a smaller segment may
not have enough chunks to give every CTA of its cluster work. Those CTAs are idle ranks: they
own no chunks, fold nothing, and never lead. They cannot simply return, because they still have to
arrive at every cluster.sync() — a CTA that returned early would make the cluster barrier
unreachable and hang the cluster. Each cluster-barrier site notes (TODO(cccl)) that a sub-cluster mbarrier
over just the working ranks would let idle ranks exit and free their SM slots (see Section 10).

2.5 Seeding output positions: the cross-CTA scan (and the cursor it replaced)

Output positions are seeded through DSMEM atomics. The chosen design is a cross-CTA scan; an
earlier DSMEM output cursor is contrasted throughout because the contrast is exactly what motivates
the scan: how many atomics (cluster-bounded and paid once vs. one per placed element), whether they
are fire-and-forget or fetch-add, and — consequently — where the cross-CTA cluster barrier lands.

As introduced in Section 1.5, each CTA needs its front/back region bases. The chosen design is a
cross-CTA prefix scan fused with priming the placement counters
(prime_placement_counters): each CTA adds its selected count and its candidate count into every
successor's front_local_cnt and back_local_cnt — the very counters the final filter later places
into. In parallel each CTA also folds its own back-region base num_selected into its own
back_local_cnt (seed_local_back), so after the scan the counters already hold the absolute
region bases the filter needs: front_local_cnt = sel_prefix and
back_local_cnt = num_selected + cand_prefix. The leader is last, so it ends up holding the full
predecessor sum and derives its own counts by subtraction. As a bonus, after the scan every CTA also
knows the totals it needs for the early-exit bookkeeping without any further DSMEM traffic.

Two 32-bit pushes, not one 64-bit push. An earlier version packed both counts into a single
(front_count << 32) | cand_count word and scanned them with one 64-bit add. That was dropped
because red.add.u64 on shared memory is not a native atomic — ptxas emulates it with a
CAS-spin loop (ATOMS.CAST.SPIN.64 in SASS) — whereas two red.add.u32s are native
warp-aggregatable shared atomics. So each CTA now issues two independent 32-bit DSMEM reductions
per successor (front lane and back lane); they were never going to carry into each other anyway.

Every rank adds its counts into all higher ranks via DSMEM atomics — fire-and-forget, all in parallel:

  from \ into  │   r1       r2       r3 = leader
  ─────────────┼─────────────────────────────────────
  r0 (c0)      │   c0       c0       c0
  r1 (c1)      │   ·        c1       c1
  r2 (c2)      │   ·        ·        c2
  ─────────────┼─────────────────────────────────────
  prefix held  │   c0       c0+c1    c0+c1+c2   ← leader; own = totals − sum

  (cₓ = one CTA's counts, pushed as two 32-bit red.add.u32 — front lane + back lane — per cell)

These pushes are fire-and-forget atomics (they do not return a value), so, unlike a fetch-add
into a shared cursor, they don't force a CTA-to-CTA round-trip to read a result back. They are also
lane-parallel (each thread owns a strided slice of the successor ranks), which measurably beat
the original single-threaded loop.

Crucially, the scan's DSMEM traffic is bounded by the cluster width and paid once (O(cluster²)
pushes), independent of k — the per-element placement that follows uses only cheap block-local
SMEM atomics. The shared cursor instead issues one DSMEM fetch-add per placed element, so its
cross-CTA atomic count scales with the number of placed keys: normally ≈ k, which is usually
≫ cluster width, and it can even exceed k when candidates are placed before the CTA has
observed early exit. (In the rare k < cluster_size regime the cursor could issue fewer, so this is
a "typically", not a guarantee.)

The two schemes also differ in where the cross-CTA cluster barrier lands, not whether one exists:
the scan's cluster barrier sits before the filter (a post-push sync), so there is no cluster
barrier after the final filter pass and CTAs exit independently; the old shared cursor instead
needed a cluster barrier after the filter, whose later timing was precisely what stopped CTAs
from early-exiting. The front-loaded inits (Section 2.1) hold either way.

---

3. Launch, dynamic cluster sizing, and portability

One kernel symbol · dynamic-SMEM block tile · wave-aware sizing · single-CTA path · portability · strict vs lenient

The dispatch first picks a backend (baseline vs cluster) from the request facts, then — for the
cluster backend — sizes the launch. Both steps fit together as follows:

flowchart TD
    Req["request:<br/>arch, max seg size,<br/>determinism/tie-break"] --> BE{"needs cluster?<br/>deterministic/tie-break,<br/>seg &gt; baseline coverage,<br/>or crossover prefers it"}
    BE -- "no" --> BL["baseline backend<br/>(worker-per-segment,<br/>1 block/segment)"]
    BE -- "yes" --> Arch{"SM 9.0+?"}
    Arch -- "no" --> NS["unsupported<br/>(strict: compile error;<br/>lenient: cudaErrorNotSupported)"]
    Arch -- "yes" --> Sel{"k ≥ segment size?"}
    Sel -- "yes" --> All["select-all fast path<br/>(copy only, no radix)"]
    Sel -- "no" --> Origin{"launch origin"}
    Origin -- "host" --> H["query CUDA RT APIs<br/>pick (cluster width C, dyn SMEM)"]
    Origin -- "CDP / device" --> D["portable width 8 + 48 KiB<br/>stream overflow"]
    H --> One{"fits 1 CTA and<br/>≤ single_block_max?"}
    One -- "yes" --> SC["single-CTA path<br/>(cluster-barrier-free)"]
    One -- "no" --> W["wave-aware search:<br/>min waves, tie-break largest C"]
    W --> Fit{"max segment<br/>fully resident?"}
    Fit -- "yes" --> RES["cluster launch, all resident"]
    Fit -- "no" --> STR["widest cluster + max SMEM,<br/>stream overflow"]

Loading

3.1 One kernel symbol, both backends (plus a CDP cluster variant)

Both backends live behind a single kernel symbol (device_batched_topk_kernel). The device-side
current_policy + if constexpr instantiate only the agent named by the selected topk_backend, so
a build compiles exactly one backend per architecture — no wasted kernels — while host code stays
agnostic to the architecture actually chosen (a TU may target several). This follows DeviceScan's
pattern rather than emitting a distinct kernel per backend.

The cluster backend needs two launch shapes (the baseline backend has no cluster dimension and always
uses the plain host launch):

• Host launch: a kernel with no __cluster_dims__, launched via cudaLaunchKernelEx with a
  runtime cluster dimension. This lets the host pick the cluster width per problem.
• CDP (device-side) launch: a kernel with a static __cluster_dims__(max_portable_cluster_blocks, 1, 1),
  gated on CUB_RDC_ENABLED (same pattern as dispatch_segmented_sort.cuh).

The two differ because the host uses a whole battery of CUDA runtime APIs to find the ideal
(cluster width, dynamic-SMEM) combination that keeps the runtime maximum segment resident —
cudaDeviceGetAttribute, cudaFuncGetAttributes, cudaOccupancyMaxActiveClusters,
cudaFuncSetAttribute, cudaOccupancyMaxPotentialClusterSize. None of these are available from a
device-side (CDP) launch. So, to stay safely portable, the CDP path just uses the portable cluster
width (max_portable_cluster_blocks == 8) and the portable shared-memory budget (48 KiB), and lets
the agent stream any overflow.

3.2 Dynamic shared memory and the resident "block tile"

The segment's keys are staged in dynamic shared memory partitioned into fixed-size chunks (one
chunk = one slot = ChunkBytes). smem_block_tile_layout computes, from a dynamic-SMEM byte budget,
how many whole chunks fit and thus the per-CTA block_tile_capacity (resident key capacity).

The base is rounded up to slot_alignment so every bulk-copy destination has the same alignment its
gmem source has. slot_alignment is max(LoadAlignBytes, alignof(key_t)):

• LoadAlignBytes only needs to be the TMA minimum (16 B) for bulk-copy correctness, but the
  policy currently uses 128 B so every chunk starts on a cache-line boundary.
• It is bumped above LoadAlignBytes only for over-aligned key types whose alignof exceeds the
  load alignment.

Dynamic SMEM (per CTA), offset increasing ─────────────────────────────►

 |pad| chunk slot 0 | chunk slot 1 | ... | chunk slot m-1 |
  ^    ChunkBytes      ChunkBytes          ChunkBytes
  |    (aligned)       (aligned)           (aligned)
  base rounded up to slot_alignment = max(LoadAlignBytes, alignof(key_t))
       \____ resident slots ____/ \__ streaming slots (p_eff) __/

Static _TempStorage (compile-time sized):
 | edge_keys[] (peeled head/tail) | hist[num_buckets] | state (len, k, result_pair) | scan storage |

3.3 Wave-aware cluster-size search (host)

The host dispatch chooses (cluster_blocks, dynamic_smem_bytes) analytically. The free variable is
the cluster width C; each C is paired with the smallest dynamic SMEM that keeps the max segment
fully resident. It enumerates the feasible C range, queries cudaOccupancyMaxActiveClusters for
clusters-per-wave, and picks the C that minimizes the number of waves, tie-breaking toward the
largest C.

The largest-C tie-break is chosen for parallelism, not L1 behavior: the bulk copies land keys
straight into shared memory and largely bypass L1, so spreading a segment over more CTAs mostly buys
more SMs working it in parallel. We scale that parallelism via the cluster width rather than a
bigger block because threads_per_block has to be a compile-time constant for cub::BlockScan, so
the block size is fixed and the cluster width is the free knob.

All the size arithmetic is done in 64-bit to stay overflow-safe: the segment-size bound is capped at
2^31 - 1 but can still be that large (e.g. an explicit cuda::args::bounds<1, 2^31 - 1>), and the
downstream products (segment size × cluster/tile counts, total grid blocks) would otherwise overflow 32
bits. If full residency is impossible
(segment larger than the widest cluster can hold), it maximizes residency with the widest launchable
cluster at the largest SMEM and lets the agent stream the overflow (Section 4).

3.4 Single-CTA path and the runtime collapse

Small segments that fit resident in one CTA and are at/below a tuning threshold
(single_block_max_seg_size) take a dedicated single-CTA path: it is cluster-barrier-free
(all cluster barriers become __syncthreads(), all atomics CTA-scope) and cheaper than spreading a tiny
segment across CTAs.

This happens two ways:

• Host-sized launch: the dispatch launches such a segment with clusterDim.x = 1.
• Runtime collapse (per-segment sizes): when the size is per-segment/deferred, the launch is
  sized for the largest segment, so a small segment collapses at runtime — ranks other than 0
  return immediately (freeing their SM slots) and rank 0 runs the cluster-barrier-free path.

The eligibility math (single_cta_eligible) is shared between host and device so both sides agree
exactly. A related knob, min_chunks_per_block, controls how many chunks a CTA must own to join the
effective cluster; it is currently 1, i.e. it just implements the "drop zero-chunk CTAs" idle-
rank avoidance and nothing more aggressive.

3.5 Portability

The CDP path (and any device without opt-in SMEM) runs within the portable 48 KiB total SMEM
budget. This works because the agent peels the unaligned boundary edges into a tiny static buffer
(Section 4.3), so streaming needs only a single resident-or-streaming slot; the only hard requirement
is that at least one load-aligned chunk fits. Segments that exceed the small portable block tile are
still correct — the agent re-streams overflow from gmem. On the host path the kernel raises its
dynamic-SMEM opt-in lazily (cudaFuncSetAttribute) only up to the selected size.

3.6 Unsupported requests: strict vs. lenient

A request may have no viable backend on a given architecture — e.g. a deterministic / tie-break
request, or a segment larger than the baseline can cover, while a pre-SM 9.0 target is present (the
cluster backend requires SM 9.0+). Because host code compiles for a list of architectures, this is
diagnosed two ways:

• Strict (default). A static_assert fires at compile time if the request cannot be served on
  any architecture in CMAKE_CUDA_ARCHITECTURES. This is the least-surprising UX for callers
  building the default multi-arch preset.
• Lenient. Defining _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT defers the diagnosis to runtime:
  the dispatch returns cudaErrorNotSupported on devices whose architecture cannot serve the
  request. CUB's own tests and benchmarks define this macro so they compile the full configuration
  space across all target architectures and skip at runtime where unsupported.

---

4. Data movement: resident vs. streaming, and the TMA pipeline

Chunking and the resident/streaming split · first pass · raw TMA pipeline and ping-pong · boundary edges · latency hiding · straddling reuse

4.1 Chunking and the resident/streaming split

Each CTA's assigned keys are partitioned into ChunkBytes chunks. If a CTA's chunks fit its resident
slots, they all stay resident and are re-read from SMEM every pass. If a CTA has more chunks than
slots (an "overflow"), it reserves a small round-robin streaming region at the tail of its block
tile and re-streams the overflow chunks from gmem each pass. The split is decided independently
per CTA — CTAs need not agree, because all cross-CTA traffic and cluster barriers are still reached uniformly.

CTA's assigned chunks in gmem:  [c0][c1][c2][c3][c4][c5][c6][c7]
                                 \____________/ \_______________/
                                 fit in slots    overflow
                                 (stay resident) (re-streamed each pass)

Block tile slots:  [ slot0 ][ slot1 ][ slot2 ]                        | [ stream slots (p_eff), round-robin ]
                    \____ resident, re-read from SMEM every pass ____/

4.2 The first pass: filling the resident tile through the pipeline

The very first radix pass has to load the resident keys (later passes just re-read them from SMEM).
That initial load runs through the same async bulk-copy pipeline, using up to PipelineStages
in-flight copies (prologue = min(PipelineStages, resident_chunks)) — the pipeline depth is used for
the initial resident load whether or not any streaming happens afterward. So "pipeline depth"
(PipelineStages) and "number of streaming slots" (p_eff) are separate things: the first-pass
resident load can use the full depth even with zero overflow.

The whole first pass is one fused loop over this rank's chunks — resident chunks first, then the
overflow chunks — behind a single TMA-issue site and a single histogram-fold site. (It used to be
three loops: resident load, resident fold, overflow fold. Merging them removed the replicated
cp.async.bulk issue sequences that had bloated the first pass's SASS.)

The first overflow wave is primed as part of this same pass rather than strictly after it: an
up-front loop issues the resident-chunk loads plus — when the stream is wider than the resident window
— the early-consumed overflow copies onto stages the resident load never occupies; then, as resident
stages drain in the fold loop, they are re-armed with the remaining first-wave overflow copies
(interleaved priming). The two pipelines overlap seamlessly on one shared set of mbarriers, so no TMA
idle gap opens between "resident loaded" and "streaming primed" (Section 4.5).

4.3 The overflow stream and the raw TMA pipeline

Overflow chunks move through a fixed set of p_eff (<= PipelineStages) streaming slots driven by an
elected-thread cp.async.bulk (TMA) copy against raw per-stage mbarriers. This inlines the
internals of BlockLoadToShared (one mbarrier per stage, one elected thread issuing the copy +
mbarrier_arrive_expect_tx, and per-stage phase-parity waits); BlockLoadToShared itself was tried
first but did not behave well as a software pipeline, so it was replaced with direct PTX. The
mbarriers are initialized once up front and reused across all radix passes and the final filter. (The
streaming state — inflight mask, ping-pong direction, slot base — lives directly in the agent; it was
previously a separate overflow_streamer helper, now inlined.)

The ping-pong pattern (why it matters). Reusing mbarriers across passes is ordinary; the
non-obvious part is that the stream reverses ("ping-pongs") its chunk-visit direction every pass to
cut memory traffic two ways:

1. Skip re-loads via resident reuse. The last p_eff chunks visited in a pass are still sitting
   in the streaming slots when the pass ends. By reversing direction, those same chunks are the
   first ones the next pass needs, so they are read straight from SMEM instead of being re-fetched
   from gmem (the "priming" copies are skipped).
2. Warm L2 for the chunks that must be re-loaded. Even chunks that do have to be re-fetched are
   re-fetched most-recently-touched-first, so they are more likely to still be hot in L2.

The same ping-pong direction is carried into the final filter, which is where the straddling-CTA
nuance in Section 4.6 applies.

Pass N   (forward):   c0 c1 c2 c3 c4 c5 [c6 c7]
                                         ^^^^^  last p_eff chunks still in slots at pass end
Pass N+1 (reverse):  [c7 c6] c5 c4 c3 c2 c1 c0
                      ^^^^^  first chunks needed = already resident → read SMEM, skip priming
                            remaining reloads run most-recently-touched-first → still warm in L2

4.4 Boundary edges

The bulk-copy path requires load-aligned chunks, but a segment's start (and possibly its end) are
generally unaligned. Both unaligned boundaries — the head prefix and the tail suffix — are
always peeled into a tiny persistent static buffer (edge_keys), loaded once in the first pass and
folded into every subsequent pass and the final filter. Peeling the tail leaves the tail chunk's
aligned bulk as a normal (possibly partial) aligned chunk that resides or streams like any other; no
partial tail chunk is ever kept resident. Peeling both boundaries removes dual-boundary pressure so
streaming needs only one full slot — this is what enables the portable 48 KiB path.

(An earlier design kept the tail suffix resident inside its own partial tail chunk and peeled it only
in the single-slot portable case. Always peeling it is simpler — the stream already had to handle a
tail chunk shorn of its peeled suffix for that portable case — and removes the "reserve a resident slot
for the tail" bookkeeping and its forced-resident-tail special cases.)

Segment in gmem:
  start (unaligned)                                        end (unaligned)
   │                                                         │
   ▼                                                         ▼
   ┌──────┬─────────────┬─────────────┬────  ────┬───────────┐
   │ head │  chunk 0    │  chunk 1    │   ...    │   tail    │
   │prefix│ (aligned)   │ (aligned)   │          │  suffix   │
   └──────┴─────────────┴─────────────┴────  ────┴───────────┘
      │                                                │
      ▼ peeled once into edge_keys (static)            ▼ peeled once into edge_keys (static)
   folded into every pass + final filter      folded into every pass + final filter

4.5 Latency hiding

Later passes overlap the first streaming wave with useful work: the stream issues its first p_eff
loads, then runs the caller's resident-chunk work (overlap_work) — histogramming or filtering the
already-resident keys — before waiting on the in-flight copies. Further waves rely on pipeline
depth.

The first pass hides even more: its first overflow wave is interleaved into the resident load
(Section 4.2), so the initial TMA copies for streaming are already in flight while the boundary edges
are staged/folded and the resident keys are histogrammed. The stream ring and the resident window
share the same mbarriers (stage_cycle = max(stream_stages, prologue) distinct stages), and the wave
is always primed in consumption order — the earliest-consumed copy issues first — so the first
wait_stage finds its copy the most in-flight:

• Forward first wave (first_wave_is_forward): the overflow is consumed ascending from stage 0.
  The resident window sits at [stage_base, stage_base + prologue) walked +1, and a stage_rot
  correction realigns a misaligned reload tail so the last resident folds free the shared stages in
  consume order.
• Reverse first wave (32-bit deterministic configs need three radix passes, which makes the first
  wave descend): the overflow is consumed descending from reverse_first_stage. The resident window
  is placed as a cyclic block starting at (reverse_first_stage + 1) mod stage_cycle, with a
  descending chunk→stage map, so the descending resident re-arms free the shared stream stages in
  consume order too. This replaced an earlier scheme that primed reverse ascending while consuming it
  descending, where the first-consumed 16 KiB copy could be issued last and erase the overlap.

A single first_wave_chunk ring counter is seeded once in consume order and handed from the up-front
loop to the fold loop's re-arms; every cursor steps unit-stride with a conditional-select wrap, so no
hot path pays a runtime % by the non-power-of-two prologue/stage_cycle.

t ────────────────────────────────────────────────────────────►
 issue    │██ p_eff async bulk copies (first wave) ██│
 overlap   │        ██ histogram / filter already-resident keys ██│   (overlaps copies)
 wait+use │                                      │wait s0│use│wait s1│use│ …

4.6 Straddling-CTA reuse in the final filter

Only the single straddling CTA — the one whose candidates cross the k-boundary while ties are
still unresolved — actually needs the final stream in strict scan order. Every other CTA is
order-independent:

• an "all-below" CTA wins every one of its candidates,
• an "above" CTA places none, and
• a CTA that already resolved its ties in its resident region has no back keys left.

The stream's initial ping-pong direction is preselected at run entry from the compile-time
pass count: a streaming rank flips direction once per histogram pass, so starting at
(!tie_reversed) ^ (num_passes & 1) makes the leftover direction after all num_passes passes equal
exactly the filter's required !tie_reversed. The straddling CTA therefore enters the final filter
already in scan order and — like every order-independent CTA — always reuses the chunks already
resident in its streaming slots, skipping the re-prime copies entirely. Under early stop fewer passes
run so the leftover direction is arbitrary, but then every CTA is all-below (no straddling CTA), so the
direction is order-independent anyway. Preselecting the direction at compile time removes the old
runtime "force the scan direction and conditionally re-prime" branch from the straddling CTA's critical
path.

---

5. Determinism and tie-breaking internals

Shared tile engine · lazy scan · blocked vs striped loads · reverse residency

The deterministic and non-deterministic filters share one tile engine (process_tiles, gated by a
Deterministic template parameter and a POD det_filter_state/nondet_filter_state); the
deterministic-only tie machinery is compiled in only where needed. On the deterministic path the
front (strictly-selected) keys are always order-independent, so they go through the same place_one
atomic as everything else. Only the back (ties) needs ordering, and only on the straddling CTA:

• CTAs whose candidates are entirely at/below the k-boundary ("all-below") win every candidate, so
  they place the back with arrival-order atomics (order-independent set) and skip the scan.
• The straddling CTA needs an index-ordered BlockScan (emit_indexed) so the num_back
  smallest-index (or largest-index, for prefer_larger_index) ties win reproducibly.

Lazy scan: rather than running the BlockScan on every tile, the deterministic filter runs the
cheap atomics by default and only falls back to the BlockScan on the single tile that actually
crosses the k-boundary (plus a forced "last-tile" scan when needed). All other tiles run at nearly the
speed of the non-deterministic path. Single-tile CTAs skip the laziness and scan directly to avoid the
extra latency.

Blocked vs. striped loads. A BlockScan over a tile needs the keys in a blocked thread
arrangement (thread t owns a contiguous run), but blocked loads bank-conflict on SMEM and don't
coalesce as well as striped ones (consecutive threads read consecutive keys). Only the straddling
CTA, and only while its ties are unresolved, actually needs blocked. So process_tiles runs the
straddling deterministic CTA in two phases: phase A loads blocked while is_tie_active, then
switches once to striped for the remaining tiles (past the boundary no tie scan is possible, so
only strictly-selected keys are placed). Every other deterministic CTA and the entire
non-deterministic path load striped throughout; a static_assert enforces that a striped
deterministic tile never reaches the index-ordered scan.

reverse_residency: for prefer_larger_index the final scan visits high indices first, so the
resident/streaming split is flipped to keep the high-index chunks resident, restoring the same
"first-visited chunks stay resident, skip re-reading the overflow" symmetry the ascending path enjoys.

---

6. Lower-level optimizations (worth knowing while reviewing)

Split load/atomic loops · register scans · unroll clamping · select-all · coarse early-exit · forward-fill placement · elected leader · generic fallback

• Split load/atomic loops in the histogram (and filter). The per-CTA loop over resident keys is a
  non-unrolled block-stride loop over tiles of Unroll * threads_per_block keys, but each tile is
  processed by two separate fully-unrolled inner loops: the first reads the whole tile of keys
  from SMEM into registers, the second feeds those registers to apply (the histogram SMEM atomic).
  Merely (partially) unrolling the block-stride loop was not enough. The compiler cannot prove the
  SMEM key-read range and the SMEM histogram range are disjoint, so a single fused unrolled loop
  serializes each load against its dependent atomic. Splitting them lets the whole wave of LDS
  loads issue and be in flight ahead of the atomics, which is what recovered histogram throughput.
• Non-leader register scans. During splitter identification only the leader does useful work with
  the merged histogram. Non-leaders exclusive-scan their own (un-merged) histogram into registers
  in parallel, so once the leader publishes the bucket, the owning lane can read its strictly-selected
  prefix and its splitter-bucket count without any extra SMEM traffic; this feeds the output scan and
  the early-exit counters and lets hist reset on its normal schedule.
• Unroll clamping for single-CTA-eligible segments. When the static segment-size bound is small
  enough that the segment will take the single-CTA path — the only path that knows its per-CTA chunk
  count at compile time — the per-thread unroll factors are clamped (both floor and ceil flavors
  for the two loop styles) so sub-tile segments don't pay for predication/registers they can't use.
  Larger/unbounded types keep the full unroll.
• Select-all fast path. When k >= segment_size, every element wins; the kernel skips the radix
  passes, histogram, and ordering entirely and just copies keys (and values) across the full cluster.
• Coarse early-exit checking. The final filter checks "am I done?" only at region boundaries and
  before each streaming refill copy — never per tile — using block-local placement counters behind one
  __syncthreads(); the atomic front/back paths run without a __syncthreads() between those points.
  Both filters exploit this: the non-deterministic filter also stops streaming between chunks once
  this CTA's front is full (front_local_cnt reached sel_prefix + my_front) and its back is full
  (back_local_cnt >= k), so a CTA whose whole contribution is already placed skips re-fetching its
  remaining overflow chunks.
• Unified, forward-fill placement. Front (selected) and back (tie) placement share one leaf
  (place_one): a single SMEM atomicAdd on the region counter — already primed to the absolute
  region base — plus a uniform out < k guard, both regions filling forward. This keeps the hot loop
  branch-uniform (only the target counter differs), lets the shared increment warp-aggregate, and is
  what let the deterministic and non-deterministic filters merge into one tile engine.
• Elected block leader beyond TMA. The single thread elected via __block_elect_one()
  (is_block_leader) is reused for every "one thread does it" site (mbarrier driving, the scan's
  local seed, the single-CTA back-base write), not just the TMA issue, so the compiler can keep those
  guards in uniform registers.
• Generic (non-TMA) fallback. Key types that aren't bulk-tileable (over-aligned, padded, or
  non-contiguous iterators) fall back to plain per-element gmem loads/stores; that path reserves no
  streaming slots and re-reads overflow from gmem each pass (still ping-ponging direction for L2
  locality).

Open question: hand-written shared-memory atomics

Histogram increments and the local output counters use shared-space red/atom PTX keyed on the
32-bit shared address rather than the builtin atomicAdd(&smem). The motivation was codegen: when
a shared-memory base address spills to local memory it becomes a generic address, which in turn
demotes the atomic to a generic atomic with a spilled 64-bit base reloaded on every update. Addressing
directly with a hoisted 32-bit shared address side-steps that. This is worth re-investigating — it
may be that fixing the underlying address spill (so the compiler keeps the shared address in a
register) removes the need for the hand-written PTX entirely.

---

7. Tuning parameters

Policy layout and default per-block knobs

All policy types live in detail::batched_topk (tuning_batched_topk.cuh). The combined
topk_policy carries the selected topk_backend plus both sub-policies (baseline_topk_policy
and cluster_topk_policy); policy_selector builds it and makes the backend decision, and the kernel
instantiates only the arm named by backend (chosen device-side via current_policy).

cluster_topk_policy carries only the per-block knobs; the cluster width and dynamic-SMEM capacity
are chosen at runtime by the dispatch. Current defaults:

Knob	Default	Role
threads_per_block	512	block size (compile-time; required by BlockScan)
histogram_items_per_thread	8	histogram-loop unroll
pipeline_stages	8	max in-flight bulk copies / mbarriers
chunk_bytes	16 KiB	slot size / stream granularity
load_align_bytes	128	chunk alignment (TMA needs 16; 128 = cache line)
bits_per_pass	11	radix digit width (2048 buckets)
min_blocks_per_sm	1	launch bound
tie_break_items_per_thread	8	filter-loop unroll
single_block_max_seg_size	8192	single-CTA fast-path threshold
min_chunks_per_block	1	effective-cluster join threshold (drop zero-chunk CTAs)
copy_items_per_thread	8	select-all copy unroll

The tuning is currently identical across all cluster-capable CCs (SM 9.0+).

---

8. Things we tried and dropped

BlockLoadToShared · histogram unroll · push broadcast · chunk subdivision · 64-bit scan · forced nounroll

• BlockLoadToShared for pipelining → replaced with raw mbarrier + cp.async.bulk PTX. It didn't
  behave well as a software pipeline; the hand-rolled pipeline is what ships.
• Cascading/halving histogram unroll → reverted. It helped small segments but regressed
  medium/large ones too much; the simpler clamped unroll won.
• Pull-based → push-based leader broadcast → reverted (kept the 64-bit packing). Push-based
  broadcasting of the result added complexity without a clear win; the single 64-bit result_pair
  packing (fewer DSMEM loads) was kept.
• Subdividing the first resident chunk for ramp-up → reverted. It did not pan out at all — it was
  meant to help small segments but measured worse there.
• 64-bit packed cross-CTA scan → split into two 32-bit pushes. Packing both counts into one
  (front << 32) | cand word and scanning them with a single red.add.u64 looked cheaper, but a
  64-bit red.add on shared memory is CAS-spin-emulated (ATOMS.CAST.SPIN.64); two native
  red.add.u32s are faster (Section 2.5). The 64-bit packing is still used for the result_pair
  broadcast (a plain load/store, not an atomic).
• Forcing nounroll on every runtime-bound loop → partially reverted. After annotating all loops
  with explicit unroll/nounroll pragmas, the histogram/filter sub-tile remainders and the radix pass
  loop regressed under forced nounroll and were restored to compiler-decided (partial/full)
  auto-unroll. The bulk-copy priming/fold loops went the other way: letting the compiler
  auto-unroll them replicated the whole cp.async.bulk issue sequence per iteration and bloated the
  first pass's SASS (it more than doubled the UBLKCP count and hurt even non-streaming kernels, which
  still carry the cold streaming machinery), so they are pinned to nounroll and were merged into the
  single fused first-pass loop (Section 4.2).

---

9. Testing

Functional, API, layout, and compile-fail tests · benchmarks

All functional tests drive the public cub::DeviceBatchedTopK API (not the low-level dispatch),
so they exercise the same backend-selection path production callers hit.

• catch2_test_device_segmented_topk_{keys,pairs}.cu: functional coverage across key types (including
  half/bf16 and small/large index types), keys-only and pairs, non-deterministic and deterministic
  (both tie-break directions), a range of segment sizes spanning single-CTA, fully-resident multi-CTA,
  and streaming/overflow regimes, plus the select-all path and a positive size-based baseline→cluster
  crossover case. The pairs test additionally pins first-pass scheduling corner cases with tiny custom
  tunings (forcing a single CTA and specific stream_slots/stages): forward stage_base priming
  with a resident chunk below a wider stream, a wrapping reverse resident window, a misaligned reload
  tail (stage_rot), and a partial final overflow wave (overflow_chunks % stream_stages != 0) —
  paths end-to-end tests otherwise reach only nondeterministically.
• catch2_test_device_batched_topk_{api,env_api}.cu: smoke tests of the public API surface (with and
  without an explicit execution environment).
• catch2_test_device_segmented_topk_cluster_layout.cu: a host-only unit test that pins down the
  smem_block_tile_layout math (capacity, padding, per-rank worst-case chunk counts) independent of a
  GPU.
• test_device_batched_topk_requirements_fail.cu / test_device_batched_topk_unsupported_arch_fail.cu:
  compile-fail tests. The first checks the requirement static_asserts (e.g. non-unsorted ordering);
  the second is pinned to a pre-SM 9.0 arch (89-virtual) and verifies the strict unsupported-arch
  static_assert fires. It is the only top-k test built without
  _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT; all other tests define it and skip unsupported
  configurations at runtime.

Benchmarks under cub/benchmarks/bench/segmented_topk/ cover fixed and variable (keys / indexed)
segment sizes. A TUNE_BACKEND knob selects which backend a build measures (baseline / cluster /
device-reference / automatic), so each backend can be tuned independently while automatic benchmarks
the production selector.

All keys+pairs test binaries pass on a Blackwell (B200) GPU.

---

10. Known limitations and future work

Idle-rank spin · 32-bit segment offsets · SM 9.0+ only

• Idle ranks spin on cluster barriers. They cannot exit early without making the cluster barrier
  unreachable. A sub-cluster mbarrier over just the working ranks (the "L2-free sync" plan) would let
  them return and free their SM slots.
• 32-bit segment offsets. Segment-internal offsets are signed int32_t, so the supported segment
  size range is [0, 2^31 - 1] keys. This is enforced at compile time from the statically-known range
  (both ends): a type whose whole range already fits (a narrow unsigned type such as uint8_t/uint16_t)
  is accepted un-annotated, while a signed type (negative min, e.g. int32_t) or an unsigned type whose
  max exceeds 2^31 - 1 (e.g. uint32_t) must carry a cuda::args::bounds within [0, 2^31 - 1] — an
  out-of-range static bound, or an un-annotated type that could leave the range, fails to compile. A
  per-segment runtime value outside its declared bound is a caller precondition violation (undefined
  behavior): only host-known values are checked (by the argument wrapper's assertions); there is no
  device-side guard. Cross-segment counts use 64-bit where needed; widening the internal offset to 64-bit
  is future work (the cross-CTA scan already avoids 64-bit shared atomics — it pushes the two counts
  as separate 32-bit red.add.u32s because a 64-bit shared red.add is CAS-spin-emulated — so a
  wider offset would need native 64-bit shared atomics or a different scan).
• Cluster backend is SM 9.0+ only. Clusters require Hopper or newer. On pre-SM 9.0 targets the
  baseline (worker-per-segment) backend serves single-block, non-deterministic requests; deterministic
  / tie-break requests and segments too large for a single block have no backend there and are reported
  strictly (compile-time) or leniently (runtime cudaErrorNotSupported) per §3.6. A multi-block
  baseline backend for large segments on older GPUs is future work.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 67925474-bfd5-49eb-b0ef-67488a8688b7

📥 Commits

Reviewing files that changed from the base of the PR and between 563d246 and fc46a4a.

📒 Files selected for processing (12)

• cub/benchmarks/bench/segmented_topk/fixed/keys.cu
• cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh
• cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh
• cub/cub/device/device_batched_topk.cuh
• cub/cub/device/dispatch/dispatch_batched_topk.cuh
• cub/test/CMakeLists.txt
• cub/test/catch2_test_device_segmented_topk_keys.cu
• cub/test/catch2_test_device_segmented_topk_pairs.cu
• cub/test/catch2_test_device_topk_common.cuh
• cub/test/test_device_batched_topk_requirements_fail.cu
• cub/test/test_device_batched_topk_unsupported_arch_fail.cu
• docs/cub/api_docs/device_topk_requirements.rst

🚧 Files skipped from review as they are similar to previous changes (12)

• cub/test/CMakeLists.txt
• cub/test/test_device_batched_topk_unsupported_arch_fail.cu
• cub/test/test_device_batched_topk_requirements_fail.cu
• cub/test/catch2_test_device_topk_common.cuh
• docs/cub/api_docs/device_topk_requirements.rst
• cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh
• cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh
• cub/benchmarks/bench/segmented_topk/fixed/keys.cu
• cub/test/catch2_test_device_segmented_topk_keys.cu
• cub/cub/device/device_batched_topk.cuh
• cub/test/catch2_test_device_segmented_topk_pairs.cu
• cub/cub/device/dispatch/dispatch_batched_topk.cuh

---

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added automatic backend selection for segmented/batched Top‑K, with clearer support for baseline, device, and SM90+ cluster execution.
  • Introduced new benchmark variants that sweep cluster-related tuning knobs and requirements.
• Bug Fixes
  • Improved handling for k==0, empty/no-work cases, and negative segment sizes (treated as empty).
  • Tightened validation and dispatch behavior for unsupported configurations and segment-size argument types/ranges.
• Documentation
  • Updated Top‑K requirements to require unsorted, and documented determinism/tie-break and architecture constraints.
• Tests
  • Expanded Top‑K benchmarks and device tests, including shared-memory layout checks and unsupported-architecture failure coverage.

Walkthrough

Changes

Adds cluster-based segmented TopK execution with unified backend dispatch, CUDA cluster launch support, parameter validation, benchmark integration, and expanded deterministic, streaming, layout, and unsupported-architecture tests.

Segmented TopK cluster optimization

Layer / File(s)	Summary
TopK contracts and backend policies cub/cub/detail/segmented_params.cuh, cub/cub/device/device_batched_topk.cuh, cub/cub/device/dispatch/tuning/..., docs/cub/api_docs/device_topk_requirements.rst	Adds deferred-argument bounds checks, segment-size validation, backend-aware policies, and updated execution-requirement documentation.
Unified baseline and cluster dispatch cub/cub/device/dispatch/..., cub/cub/detail/launcher/cuda_runtime.cuh, thrust/thrust/system/cuda/detail/core/triple_chevron_launch.h	Routes requests between baseline and cluster kernels and adds cluster launch configuration.
Cluster TopK agent implementation cub/cub/agent/agent_batched_topk_cluster.cuh, cub/cub/agent/agent_batched_topk.cuh	Implements cluster histogram merging, radix refinement, shared-memory streaming, deterministic placement, select-all handling, effective cluster widths, and zero-k handling.
Benchmark backend configurations cub/benchmarks/bench/segmented_topk/{fixed,variable}/...	Centralizes benchmark drivers and adds automatic, baseline, cluster, and device backend configurations with variable segment-size support.
TopK behavior and contract coverage cub/test/..., cub/test/CMakeLists.txt	Adds coverage for cluster layouts, deterministic tie-breaking, streaming paths, segment-size handling, invalid contracts, and unsupported architectures.

Assessment against linked issues

Objective	Addressed	Explanation
Prototype cluster optimization for segmented TopK [#9077]	✅
Provide benchmark coverage comparing segmented cluster TopK with DeviceTopK [#9077]	✅
Support deterministic selection and configurable tie-breaking [#9077]	✅
Provide sorted output [#9077]	❌	The implementation requires output_ordering::unsorted; sorted and stable-sorted output are explicitly rejected.

Possibly related issues

• [#8360] — The changes expand DeviceBatchedTopK across bounded variable segment sizes and update related validation and dispatch behavior.

Suggested reviewers: elstehle, griwes, oleksandr-pavlyk, naderalawar

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 11

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

cub/benchmarks/bench/segmented_topk/fixed/keys.cu (1)

146-235: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

important: Apply the repository’s const/constexpr rule consistently to the new immutable environments, argument wrappers, iterators, input buffers, and host data.

• cub/benchmarks/bench/segmented_topk/fixed/keys.cu#L146-L235: qualify the environment objects, fixed argument wrappers, host segment sizes, and launch environment.
• cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh#L177-L253: qualify the environment objects, input buffer, argument wrappers, iterator handles, and launch environment.
• cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh#L189-L282: qualify the environment objects, input buffer, argument wrappers, iterator handles, and launch environment.

As per coding guidelines, “All variables that are not modified must be declared const” and compile-time variables must be constexpr.

Source: Coding guidelines

cub/cub/device/device_batched_topk.cuh (1)

162-165: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

important: Validate k and num_segments beyond their wrapper form.

These checks accept every cuda::args wrapper. For example, a floating-point immediate, a bool, or an invalid scalar deferred handle passes the gate. A deferred num_segments also passes because it is single-valued, despite the host-resolved contract.

Add integral non-bool value checks, validate deferred handles for k, and restrict num_segments to constant, immediate, or plain integral forms.

#!/bin/bash
fd 'test_device_batched_topk.*fail\.cu' cub/test \
  -x rg -n -C3 'segment_sizes|num_segments|\bk\b|deferred|immediate' {}

Also applies to: 232-248

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ff40ff92-60e5-4a2c-880d-140f90939995

📥 Commits

Reviewing files that changed from the base of the PR and between 22e4e68 and 240dcd3.

📒 Files selected for processing (26)

• cub/benchmarks/bench/segmented_topk/fixed/keys.cu
• cub/benchmarks/bench/segmented_topk/variable/common.cuh
• cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu
• cub/benchmarks/bench/segmented_topk/variable/indexed.cu
• cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh
• cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu
• cub/benchmarks/bench/segmented_topk/variable/keys.cu
• cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh
• cub/cub/agent/agent_batched_topk.cuh
• cub/cub/agent/agent_batched_topk_cluster.cuh
• cub/cub/detail/launcher/cuda_runtime.cuh
• cub/cub/detail/segmented_params.cuh
• cub/cub/device/device_batched_topk.cuh
• cub/cub/device/dispatch/dispatch_batched_topk.cuh
• cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh
• cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh
• cub/cub/util_device.cuh
• cub/test/CMakeLists.txt
• cub/test/catch2_test_device_segmented_topk_cluster_layout.cu
• cub/test/catch2_test_device_segmented_topk_keys.cu
• cub/test/catch2_test_device_segmented_topk_pairs.cu
• cub/test/catch2_test_device_topk_common.cuh
• cub/test/test_device_batched_topk_requirements_fail.cu
• cub/test/test_device_batched_topk_unsupported_arch_fail.cu
• docs/cub/api_docs/device_topk_requirements.rst
• thrust/thrust/system/cuda/detail/core/triple_chevron_launch.h

[coderabbitai]

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

important: Convert the new local header inclusions to configured repository-relative angle-bracket paths.

• cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu#L32-L32: convert indexed_common.cuh.
• cub/benchmarks/bench/segmented_topk/variable/indexed.cu#L25-L25: convert indexed_common.cuh.
• cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh#L37-L37: convert common.cuh.
• cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu#L32-L32: convert keys_common.cuh.
• cub/benchmarks/bench/segmented_topk/variable/keys.cu#L25-L25: convert keys_common.cuh.
• cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh#L42-L42: convert common.cuh.

As per coding guidelines, “All header inclusions must use angle-bracket syntax.”

📍 Affects 6 files

• cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu#L32-L32 (this comment)
• cub/benchmarks/bench/segmented_topk/variable/indexed.cu#L25-L25
• cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh#L37-L37
• cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu#L32-L32
• cub/benchmarks/bench/segmented_topk/variable/keys.cu#L25-L25
• cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh#L42-L42

Source: Coding guidelines

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
rg -n -C4 'deferred(_sequence)?\s*\{|get_segment_size\s*\(' cub/test cub/cub

Repository: NVIDIA/cccl

Length of output: 42582

---

🏁 Script executed:

sed -n '1,180p' cub/cub/detail/segmented_params.cuh

Repository: NVIDIA/cccl

Length of output: 8354

---

🏁 Script executed:

rg -n -C3 'element_type|struct __traits|deferred_sequence|deferred\{' cub/cub cub/cuda

Repository: NVIDIA/cccl

Length of output: 31378

---

🏁 Script executed:

rg -n -C3 'deferred_sequence\{|deferred\{' cub/test cub/cub | sed -n '1,220p'

Repository: NVIDIA/cccl

Length of output: 20232

---

🏁 Script executed:

sed -n '120,150p' cub/cub/detail/segmented_params.cuh

Repository: NVIDIA/cccl

Length of output: 1542

---

🏁 Script executed:

sed -n '1,220p' cub/cuda/argument

Repository: NVIDIA/cccl

Length of output: 210

---

🏁 Script executed:

rg -n -C4 'struct __traits|using element_type|using value_type|deferred_sequence|deferred<' cub/cuda cub/cub

Repository: NVIDIA/cccl

Length of output: 30603

---

important: Materialize both deferred reads as typename ::cuda::args::__traits<_Arg>::element_type before returning. auto __value can preserve a proxy reference from *handle / handle[index], and get_segment_size then instantiates max against that proxy type. Make __value and __size const too.

Source: Coding guidelines

[Jacobfaib]

cmake lgtm

[github-actions]

😬 CI Workflow Results

🟥 Finished in 1h 22m: Pass: 71%/344 | Total: 2d 16h | Max: 52m 07s | Hits: 90%/338131

See results here.

[elstehle]

discussion: Let's discuss whether we want to make the macro public.

[elstehle]

just a personal note for posterity/myself:

Generated outline on the compilation & dispatch behaviour

note, I found this paragraph very helpful:

Think of it as a pipeline of gates. Gate 1–3 (contract, annotation form, 2²¹ limit) are unconditional compile errors — the macro can't touch them; fix them by changing your call. Gate 4 is the only conditional one: your request kind + segment regime + the archs in your fatbin decide whether the request needs the SM90+ cluster backend, and if any target arch can't provide it, the macro chooses when you find out — compile time (strict, default) or runtime cudaErrorNotSupported (relaxed). Deterministic-anything and segments-too-big-for-one-block are the two things that demand SM90+; segment size ≤ 2²¹ is a hard wall; the 8 Ki / 16 Ki numbers only steer which backend runs, never whether it compiles.

I now have the full picture. The key insight is that there are four independent failure mechanisms, and the macro only gates one of them. Let me lay it out.

The one macro: _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT

It flips a single decision — what happens when the request resolves to the unsupported backend for at least one architecture — between compile-time error (default/"strict") and runtime cudaErrorNotSupported ("relaxed"):

#if _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \
  && !defined(_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT)
  static_assert(
    !any_target_cc_unsupported<selector_t>(),
    "cub::DeviceBatchedTopK: the requested top-k configuration is not supported on at least one architecture ...");
#endif

• Strict (default, macro undefined): the static_assert scans every arch in CMAKE_CUDA_ARCHITECTURES. If any one would be unsupported, it's a compile error. Guarantee: if it compiles, every target arch works, so a runtime cudaErrorNotSupported for arch reasons essentially can't happen.
• Relaxed (macro defined): compiles regardless; at runtime the dispatch picks the policy for the actual device's CC, and if that resolves to unsupported it returns cudaErrorNotSupported (with the two-phase temp-storage protocol handled first). CUB's own tests/benchmarks define this so one binary spans the whole config space and skips at runtime.

Crucially, the macro does not affect the other three mechanisms below — those are always hard compile errors regardless of the macro.

The four mechanisms (only #4 is macro-gated)

#	Check	When it fires	Failure mode	Macro-affected?
1	Output / determinism contract	bad (determinism, tie_break, output_ordering) combo	compile error (static_assert)	No
2	Argument annotation form	bad segment_sizes/k/num_segments wrapper or element type	compile error	No
3	Static max segment size > 2²¹	seg_traits::highest > 2^21	compile error	No
4	Backend unsupported on a target arch	request resolves to unsupported for some CC	compile error or runtime cudaErrorNotSupported	Yes

Matrix A — the backend selector (drives mechanism #4)

This is policy_selector::operator()(cc). "unsupported" here is what the macro turns into either a compile error or a runtime error. Everything else launches fine.

Let me define the segment-size regimes first, since they feed the selector:

Regime	Meaning
S₁ static max ≤ baseline coverage (~16 Ki, largest worker tile)	baseline can cover
S₂ baseline coverage < static max ≤ 2²¹	baseline cannot cover → cluster-only
S₃ static max > 2²¹	rejected by mechanism #3 (never reaches selector)

Now the decision, per request kind × arch (automatic mode):

Request	seg regime	cc < 9.0	cc 9.0–9.x	cc ≥ 10.0
Non-deterministic (not_guaranteed + unspecified)	S₁ (baseline covers)	baseline	baseline	cluster if static max ≥ 8 Ki, else baseline
Non-deterministic	S₂ (too big for baseline)	unsupported	cluster	cluster
Deterministic (run_to_run/gpu_to_gpu, or any concrete tie-break)	S₁ or S₂	unsupported	cluster	cluster

    else if (deterministic)
    {
      backend = cluster_capable(cc) ? topk_backend::cluster : topk_backend::unsupported;
    }
    else if (!baseline_can_cover)
    {
      backend = cluster_capable(cc) ? topk_backend::cluster : topk_backend::unsupported;
    }
    else
    {
      const bool beneficial = cc >= ...{cluster_beneficial_min_cc_major, 0}
                           && static_max_segment_size >= cluster_beneficial_min_segment_size;
      backend = (cluster_capable(cc) && beneficial) ? topk_backend::cluster : topk_backend::baseline;
    }

The only two ways to land in unsupported (→ macro-gated) under automatic:

• deterministic request while a pre-SM90 arch is in the build, or
• segment too big for the baseline (S₂) while a pre-SM90 arch is in the build.

force_baseline/force_cluster (via a tune()d selector) add more unsupported cases: force_baseline + (deterministic OR S₂) → unsupported on all archs; force_cluster + pre-SM90 → unsupported.

Matrix B — segment-size thresholds (there are four distinct ones)

These are easy to conflate; they're genuinely different numbers doing different jobs:

Threshold	Value	Role	Effect if exceeded
Baseline worker tile	~16 Ki	max a single-block baseline can tile	forces cluster-only (S₂) → SM90+ required (macro-gated on old archs)
Cluster perf crossover	8 Ki	cluster_beneficial_min_segment_size	below it on SM100+, baseline is kept even though cluster could run — pure perf, never a failure
Cluster single-CTA fast path	8 Ki	single_block_max_seg_size (policy)	above it, cluster uses the barrier-full multi-CTA path — perf/path choice, never a failure
Hard API limit	2²¹ (~2 M)	max supported static segment size	compile error, always (mechanism #3)

Note the 8 Ki crossover and the 8 Ki single-CTA threshold are numerically equal but conceptually separate (one is in the selector and non-tunable; the other is a policy field). Also: within the cluster backend, exceeding resident SMEM capacity (up to ~2 M) just triggers gmem re-streaming — slower, not a failure.

  // Only the statically-known *maximum* segment size is constrained: it must not exceed 2^21 (about 2 million).
  ...
    static_assert(
      ::cuda::std::cmp_less_equal(seg_traits::highest, ::cuda::std::int64_t{1} << 21),
      "cub::DeviceBatchedTopK: the statically-known maximum segment size exceeds the maximum currently supported ...");

This is why an un-annotated int segment size is a compile error: its type-max (2³¹) exceeds 2²¹, so you must attach cuda::args::bounds<…> ≤ 2²¹.

Matrix C — output / determinism requirements (mechanism #1, always compile-time)

Three layered static_asserts admit exactly five (determinism, tie_break) pairs, and only output_ordering::unsorted:

determinism	tie_break	output	Result
specified but tie_break omitted (or vice versa)	—	—	compile error (must be paired)
not_guaranteed or run_to_run	prefer_smaller/larger_index	—	compile error (concrete tie-break needs gpu_to_gpu)
any valid pair	any valid	sorted/stable_sorted	compile error (not implemented)
empty env	—	(defaults to stable_sorted)	compile error (must request unsorted)
not_guaranteed	unspecified	unsorted	OK (any arch, subject to Matrix A)
run_to_run	unspecified	unsorted	OK → deterministic → cluster → SM90+
gpu_to_gpu	unspecified / prefer_smaller_index / prefer_larger_index	unsorted	OK → deterministic → cluster → SM90+

The mental model, in one paragraph

Think of it as a pipeline of gates. Gate 1–3 (contract, annotation form, 2²¹ limit) are unconditional compile errors — the macro can't touch them; fix them by changing your call. Gate 4 is the only conditional one: your request kind + segment regime + the archs in your fatbin decide whether the request needs the SM90+ cluster backend, and if any target arch can't provide it, the macro chooses when you find out — compile time (strict, default) or runtime cudaErrorNotSupported (relaxed). Deterministic-anything and segments-too-big-for-one-block are the two things that demand SM90+; segment size ≤ 2²¹ is a hard wall; the 8 Ki / 16 Ki numbers only steer which backend runs, never whether it compiles.

Want me to turn this into a decision-flow diagram, or add it as a section to the review guide?

[elstehle]

question: we are guarding PDL through _CCCL_HAS_PDL, which just seems to be a user switch. Do we expect there to be any reason in future for providing the same for cluster support?

[elstehle]

question: Should we document the 2^21 segment size restriction here?

[elstehle]

question: Can you remind me, would we accept an upper bound that's inferred, solely from the type's range coverage, of ~65K from passing a cuda::args::immediate{cuda::std::int16_t{8}};?

[elstehle]

question: Should we document the rationale why we reject spans? Is that something we generally agreed on for the argument annotation framework? I just think that I would appreciate understanding this when I revisit this in a year from now and have forgotten most of the context.

[elstehle]

question: Should we extend coverage to a list that contains a mixture of architectures supporting clusters and not supporting architectures?

[elstehle]

suggestion: based on offline discussion, this is something we'd want to have

[elstehle]

suggestion: we may want to call this something like needs_set_determinism, to indicate that we're talking st membership determinism at this stage.