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_from_types, 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 feed 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. Sets the winning digit in 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 with the builtin shared-memory
   atomicAdd.
2. DSMEM merge into the leader. Every non-leader CTA walks its local histogram and reduces 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 sets the winning digit in 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/>(own hist, then merged hist)"]
    A -- "2 · DSMEM reduce" --> L
    B -- "2 · DSMEM reduce" --> L
    C -- "2 · DSMEM reduce" --> L
    L --> S["3 · BlockScan merged hist<br/>find k-th bucket"]
    S --> R["publish pass_result<br/>(kth_bucket, early_stop)"]
    R -. "read back via DSMEM,<br/>set digit in 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 pass_result · 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 reducing them into the leader
  (so the leader's own local increments are ordered ahead of the incoming reductions). We drop that
  cluster barrier by leaning on atomic scope instead: the block-private increments use the builtin
  atomicAdd, whose default .gpu scope already spans the cluster peers, so the leader's own
  increments are mutually atomic with the non-leaders' .cluster-scoped DSMEM reductions racing into
  the same hist. That saves one L2-round-trip cluster barrier per pass. (See §6 for why the local
  atomicAdd still needs a pinned 32-bit shared base to stay cheap.)
• 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 state (type state_t, nested in the agent so its counters reuse the kernel's
32-bit offset_t/out_offset_t) lives in the leader block's shared memory. It is split into two
naturally-aligned 8-byte sub-structs — size {len, k} and result {kth_bucket, early_stop} (§2.3) —
so each is fetched with a single DSMEM load. Non-leaders have no local copy; they read it from the
leader through DSMEM (and cache what they need).

2.3 The leader block and the pass_result 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 an 8-byte pass_result sub-struct with two uint32_t fields:

• kth_bucket: the splitter bucket (each CTA sets this digit in its own splitter key);
• early_stop: the early-stop flag.

Because the sub-struct is exactly 8 bytes and naturally aligned, every CTA pulls it from the leader
with a single 64-bit DSMEM load (ld.shared::cluster.u64) via the templated
leader_state_load<FieldT> accessor, which cuda::std::bit_casts the loaded word back to the
sub-struct and decodes both halves locally — no second remote load for the early-stop check, and no
broadcast of the full splitter key. (The size {len, k} sub-struct is read the same way. An earlier
version hand-packed the two halves into a uint64_t with shifts/masks and forced 16-byte state
alignment; the two 8-byte sub-structs drop both the manual packing and the over-alignment.)

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, contribute 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. The same call also adds this CTA's own back-region base num_selected into its own
back_local_cnt, 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" --> H["host: query CUDA RT APIs<br/>pick (cluster width C, dyn SMEM)"]
    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

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_algorithm, 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 API is host-only: it does not support CUDA dynamic parallelism (device-side launch), so the
cluster backend needs only one launch shape (the baseline backend has no cluster dimension and also
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.

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 launch, which is
one reason the API is host-only.

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 (size{len,k}, result{kth_bucket,early_stop}) | 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: a single segment is capped at 2^21
keys (so a segment-internal uint32_t offset cannot overflow), but cross-segment quantities are not: most
notably the launch grid's block count num_segments × cluster_blocks (which the host caps at INT_MAX)
and the residency / wave-count sizing products, so computing them in 32 bits could overflow. 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

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. 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-consume site. (It used to be
three loops: resident load, resident consume, overflow consume. 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 consume 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
consumed by 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)
   consumed by every pass + final filter      consumed by 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/consumed 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 consumes 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 consume 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).

Shared-memory atomics: builtin atomicAdd on a pinned 32-bit base

Block-private histogram increments and the local front/back placement counters use the builtin
atomicAdd on a 32-bit shared address pinned by detail::warpspeed::optimizeSmemPtr (the same
trick the dynamic-SMEM allocator uses; nvbug 4907996). Without the pin, atomicAdd(&smem[i])
re-derives the address from _TempStorage's 64-bit generic base on every update, spilling a 64-bit
base and demoting the op to a generic atomic; pinning a hoisted 32-bit base keeps the per-key
increment at pure 32-bit addressing (ATOMS.POPC.INC.32). atomicAdd keeps its default .gpu scope
(not .cta — that is atomicAdd_block), which spans the cluster peers and so is mutually atomic with
the cross-CTA DSMEM reductions (§2.1). This replaced an earlier hand-written inline-PTX red/atom
sequence keyed on the 32-bit address. The cross-CTA reductions themselves stay inline PTX
(red.relaxed.cluster.shared::cluster.add.u32): a .cluster-scoped reduction into a peer's DSMEM
window (a remapped 32-bit remote address) is not expressible through the builtin atomicAdd.

---

7. Tuning parameters

Policy layout and default per-block knobs

The policy types live in detail::batched_topk (tuning_batched_topk.cuh). The combined
topk_policy carries the selected topk_algorithm plus both sub-policies (baseline_topk_policy
and cluster_topk_policy). The selector policy_selector_from_types builds it, taking the backend
decision inline in its operator(); it lives in dispatch_batched_topk.cuh (not the tuning header)
because that decision needs baseline_can_cover_v, which is only computable from the concrete baseline
agent types. The kernel then instantiates only the arm named by backend (chosen device-side via
current_policy). A tuned selector supplied through the dispatch tuning environment (keyed on
topk_policy) replaces policy_selector_from_types wholesale.

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. Push-based broadcasting of the result
  added complexity without a clear win; the single 64-bit DSMEM load of the 8-byte pass_result
  (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). A single 64-bit load/store is still used for the
  pass_result 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/consume 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 (k >= segment_size), negative-k
  clamping (a negative runtime k selects nothing, mirroring negative segment sizes), 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, split into many %PARAM% variants. The first covers the execution-requirement
  static_asserts (e.g. non-unsorted ordering) and the argument-validation asserts for each of
  segment_sizes, k, and num_segments (bad element types, a per-segment num_segments, and a k
  whose element type is wider than 64 bits). The second is compiled for two arch sets from one source:
  a pre-SM 9.0 arch alone (89-virtual) and a mixed pre/post-Hopper pair (89-virtual;90-virtual),
  both verifying the strict unsupported-arch static_assert fires when any target arch cannot
  serve the request. It is the only top-k test built without
  CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT; all other tests define it and, rather than skipping
  outright, dispatch and assert they get cudaErrorNotSupported for unsupported configurations
  before skipping the correctness checks.

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.
• k value and type handling. k has no algorithmic maximum — a k past a segment's size selects
  the whole segment. Its element type may be at most 64 bits wide, enforced by a static_assert
  (sizeof(k's element type) <= 8): the device clamps k to the segment size through a 64-bit
  intermediate, so a wider type (e.g. __int128) is rejected to avoid a silent wrap. As with segment
  sizes, negative k is allowed under a negative static lower bound and clamped to 0 (selecting
  nothing); the clamp happens before any type cast (__get_and_clamp_param_to_nonnegative).
• 32-bit segment offsets. Segment-internal offsets are uint32_t (unsigned: offsets, ranks, and block
  counts are non-negative, and negative segment sizes — and negative k — are clamped to 0 upstream). The binding limit is not
  the offset width but the public entry's cap on the statically-known maximum segment size, currently
  2^21 keys, enforced at compile time from the segment-size argument's static upper bound: a type whose
  maximum already fits (e.g. uint8_t/int16_t/uint16_t) is accepted un-annotated, while a type whose
  maximum exceeds 2^21 (e.g. int32_t/uint32_t/int64_t) must carry a cuda::args::bounds whose upper
  end is <= 2^21 — an out-of-range static upper bound, or an un-annotated wide type, fails to compile. A
  negative static lower bound is accepted (a negative runtime size is clamped to 0, i.e. treated as empty);
  a non-negative lower bound is trusted. A per-segment runtime value outside its declared bound is a caller
  precondition violation (undefined
  behavior): the statically declared bounds are validated at compile time, while the argument values are
  bounds-checked only by assertions active in assertion-enabled (e.g. debug) builds — host-side for a
  host-known immediate value and device-side (via CUB's __assert_param_in_bounds) for values read from a
  deferred / deferred_sequence handle. 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.

[elstehle]

Flushing a few more comments on the agent.

[elstehle]

nit:
[should-defer]:
I think these could be covered by 32-bit types. As a general follow-up point: We probably want to re-assess 64-bit type usage for the cluster scenario.

[elstehle]

Flushing a few more comments. From parsing through the call stack and following through the barrier sequence diagram, things look good.

I think we need a bigger follow-up PR that will restructure the code, things are very intertwined and hard to remember right now.

[elstehle]

suggestion:
[should-defer]: For !use_block_load_to_shared, I think these two fields could be dropped.

[pauleonix]

Noted in #10023

[github-actions]

⏱️ CCCL compile-time benchmark comparison: Public headers compile-time bench

Result: 1 regression row(s), 2 improvement row(s) above threshold.

Run	Value
Config	public-headers-gcc13
Baseline	origin/main
Preset	all-dev
Targets	cub.headers.base, thrust.cpp.cuda.headers.base, libcudacxx.test.public_headers
GPU / launch args	rtx2080 / --cuda 13.3 --host gcc13

Artifacts: reports and traces

Direct file processing

-f file-processing exclusive --sort total

🔴 Direct file processing — Regressions

Rank	Regression impact	Selected Δ	Baseline	Current	Event	Matched traces
1	0.311667	0.311667	0.254034	0.565701	Processing Header File: cub/cub/device/dispatch/dispatch_batched_topk.cuh	3

🟢 Direct file processing — Improvements

Rank	Improvement impact	Selected Δ	Baseline	Current	Event	Matched traces
1	0.587114	-0.587114	4.780600	4.193486	Processing Header File: libcudacxx/include/cuda/std/__cccl/prologue.h	551
2	0.228494	-0.228494	1.653215	1.424721	Processing Header File: libcudacxx/include/cuda/std/__cccl/epilogue.h	551

[elstehle]

Congratulations on the work, @pauleonix 🥳

I've gone through this in multiple iterations now and haven't found anything blocking. I'd suggest we handle the rest in a bigger follow-up refactor. My hope is that that will bring down the mental load down in future for our future selves 🙂

[elstehle]

nit: I think this should be:

Suggested change

	// through `process_tiles` instead (as separate regions for the deterministic filter, via `nondet_consume_resident`).
	// through `process_tiles` instead (as separate regions for the deterministic filter, via `process_head_edge` and `process_tail_edge`).

[elstehle]

question: are we ever reading this on the non-deterministic path?

Suggested change

	out_offset_t num_cluster_tie_winners;
	out_offset_t num_cluster_tie_winners;

[elstehle]

nit: We're using the tid member in all other places.

Suggested change

	for (unsigned rank = threadIdx.x; rank < cluster_block_rank; rank += policy.threads_per_block)
	for (unsigned rank = tid; rank < cluster_block_rank; rank += policy.threads_per_block)

[elstehle]

nit:

Suggested change

	for (unsigned rank = cluster_block_rank + 1u + threadIdx.x; rank < layout.num_logical_cluster_blocks;
	for (unsigned rank = cluster_block_rank + 1u + tid; rank < layout.num_logical_cluster_blocks;

[elstehle]

nit:

Suggested change

	const offset_t cluster_tid = cluster_block_rank * static_cast<offset_t>(policy.threads_per_block) + threadIdx.x;
	const offset_t cluster_tid = cluster_block_rank * static_cast<offset_t>(policy.threads_per_block) + tid;

[elstehle]

suggestion: this definition appears three times and only depends on key_t, let's make it a static constexpr member var.

Suggested change

	constexpr int total_bits = int{sizeof(key_t)} * 8;
	constexpr int total_bits = int{sizeof(key_t)} * 8;

[elstehle]

suggestion: this is already a class member, so we can just read num_passes.

Suggested change

constexpr int num_passes = detail::topk::calc_num_passes<key_t>(policy.bits_per_pass);

[elstehle]

nit: This has only one read inside if constexpr (use_block_load_to_shared), I'd move it in there.

[elstehle]

suggestion: reverse_first_stage is already declared const int

Suggested change

	int streaming_stage = first_wave_is_forward ? 0 : static_cast<int>(reverse_first_stage); // reverse 1st-consumed
	// stage
	int streaming_stage = first_wave_is_forward ? 0 : reverse_first_stage; // reverse 1st-consumed stage

[elstehle]

suggestion:
[may-defer]: This snippet is duplicated. May be worth extracting into a lambda/helper.

Suggested change

	if constexpr (first_wave_is_forward)
	{
	first_wave_chunk_idx =
	(first_wave_chunk_idx + offset_t{1} == first_wave_base + num_first_wave_stages)
	? first_wave_base
	: first_wave_chunk_idx + offset_t{1};
	}
	else
	{
	first_wave_chunk_idx =
	(first_wave_chunk_idx == first_wave_base)
	? first_wave_base + num_first_wave_stages - offset_t{1}
	: first_wave_chunk_idx - offset_t{1};
	}
	if constexpr (first_wave_is_forward)
	{
	first_wave_chunk_idx =
	(first_wave_chunk_idx + offset_t{1} == first_wave_base + num_first_wave_stages)
	? first_wave_base
	: first_wave_chunk_idx + offset_t{1};
	}
	else
	{
	first_wave_chunk_idx =
	(first_wave_chunk_idx == first_wave_base)
	? first_wave_base + num_first_wave_stages - offset_t{1}
	: first_wave_chunk_idx - offset_t{1};
	}

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 12h 10m: Pass: 100%/381 | Total: 20d 14h | Max: 4h 28m | Hits: 26%/3141581

See results here.

[pauleonix]

Some diagrams showing synchronization/overlap patterns used for review:

• Whole cluster view with focus on cluster-synchronization

  • assumes two radix/histogram passes, a 4B key type would need 3 or 4 without early exit)
  

• Single CTA view focusing on bulk copies (latencies etc. are just a naive representation, nothing measured). Assumes

  • two radix/histogram passes as above
  • followed by a deterministic filter-pass that prefers smaller indices
  • 5 chunks of input data, 3 slots in smem (fitting a chunk each) and a pipeline of 2 stages/mbarriers which means that there are min(2, 5 - 3) = 2 streaming slots and 3 - 2 = 1 resident slot/chunk. So 4 chunks are re-streamed in ping-pong fashion (=> only 2 per pass).
  • No early exit from the filter pass either.