[Tracking] Top-K implementation strategies: capabilities and coverage
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
## Purpose
This issue provides an overview of the _implementation strategies_ behind the CUB top-k family of algorithms. It tracks, for every approach or implementation that we have or plan to build, _what it can do_ and _where it is expected to do well_, so we can reason about coverage, overlaps, and dispatch cross-over points at a glance.
## Terminology
We describe each strategy against the committed requirement model documented in
[Top-K: Determinism, Tie-Breaking, and Output Ordering](https://nvidia.github.io/cccl/unstable/cub/api_docs/device_topk_requirements.html#cub-topk-requirements):
- **Set membership**: `determinism` = {`not_guaranteed`, `run_to_run`, `gpu_to_gpu`}, optionally refined
by `tie_break` = {`unspecified`, `prefer_smaller_index`, `prefer_larger_index`}. A concrete
`prefer_*` tie-break implies `gpu_to_gpu`.
- **Result ordering**: `output_ordering` = {`unsorted`, `sorted`, `stable_sorted`}.
- **Comparison method**: `radix-based` (bit/digit selection) vs `comparison-based` (a comparator).
Legend for the matrix: ✅ supported / implemented, 🟡 planned or conditional (see notes), ⬜ not
supported / not applicable, ❔ to be determined.
## Implementation overview
TBD.
## Strategies
- [ ] **1. Warp bitonic-sort based**: warp-wide bitonic sort (#8391) + warp top-k (#9281). Comparison-based,
register-resident, deterministic, sorted output. Building block for a warp-per-segment device path
on small segments. Owner: @yong-wang
- [ ] **2. Warp redux based**: warp-level top-k built on warp reduction intrinsics. Targets very small k.
Owner: @yong-wang
- [ ] **4. Atomics based (small k)**: atomic-max-on-twiddled-bits approach from the GTC talk; strong for very
small k on large inputs. Owner: TBD
- [ ] **5. Cluster radix selection**: thread-block-cluster radix selection with TMA and segment streaming
for medium-sized segments. determinism/tie-break already wired. Owner: @pauleonix
- [ ] **6. Multi-CTA AIR** — multi-CTA extension of AIR radix selection for large segments, with load
balancing to route large segments to the multi-CTA path. Owner: TBD.
Contributor guide
Research direction
Start with the Top-K requirements documentation linked in the issue to understand determinism, tie-breaking, and output ordering. Then review the listed strategies—warp bitonic-sort, warp redux, atomics, cluster radix selection, and Multi-CTA AIR—and document their capabilities, coverage, and dispatch crossover points. Done means the implementation overview and strategy matrix are no longer TBD.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100