NVIDIA / NVIDIA/cccl

[Tracking] Top-K implementation strategies: capabilities and coverage

Open
#9,747 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.