[EPIC] Add family of top-k algorithms to CUB

Tracking Issue: Top-K Algorithms in CUB

This issue serves as the central hub for the implementation of a family of high-performance top-k algorithms in CUB.

Mission & Goals

The current ecosystem for top-k algorithms on GPUs is fragmented, with implementations scattered across various libraries and applications. This initiative aims to centralize these efforts into a state-of-the-art family of algorithms within CUB.

By unifying top-k within CUB, we aim to:

Eliminate redundancy: Provide a definitive family of high-performance algorithms to replace the fragmented, ad-hoc implementations currently scattered across the ecosystem.
Establish a comprehensive benchmarking suite: Introduce a benchmark suite covering a broad spectrum of critical modern workloads, including Mixture of Experts (MoE), Sparse Attention, and database management.
Maximize performance: Leverage CUB’s sophisticated tuning infrastructure to automatically identify and fine-tune the optimal algorithmic strategy for every combination of workload and GPU architecture.

Scope

This issue is sub-divided into two sections:

Feature Support Matrix: Detailed tracking of supported configurations (e.g., deterministic, stable, segmented) across all CUB scopes (Device, Block, Warp).
Implementation Roadmap: The chronological plan for delivering these features.

Feature Support Matrix

Each algorithm's output requirements are tracked along two independent dimensions:

Set-membership determinism (which elements are selected, especially when keys tie at the selection boundary): nondeterministic, deterministic (reproducible set, unspecified tie-break), deterministic (prefer lower indices), deterministic (prefer higher indices).
Output ordering (how the selected elements are arranged): unordered, unstable sorted (sorted by key, equal keys in unspecified order), stable sorted (sorted by key, equal keys kept in source order).

`DeviceTopK`

🙋 Issue for tracking DeviceTopK user requests and feature prioritization

Current Status of Feature Support:

Output requirements:
- Set-membership determinism: (✓) nondeterministic, (☐) deterministic, (☐) deterministic (prefer lower indices), (☐) deterministic (prefer higher indices)
- Output ordering: (✓) unordered, (☐) unstable sorted, (☐) stable sorted
Comparison method: (✓) radix-based, (✓) radix-based for custom types (#6290), (☐) comparison-based

`DeviceSegmentedTopK` & `DeviceBatchedTopK`

🙋 Issue for tracking DeviceSegmentedTopK and DeviceBatchedTopK user requests and feature prioritization

Output requirements:
- Set-membership determinism: (✓) nondeterministic, (✓) deterministic, (✓) deterministic (prefer lower indices), (✓) deterministic (prefer higher indices)
- Output ordering: (✓) unordered, (✓) unstable sorted (small k; any k: [FEA]: support any K for non-stable sorted outputs to DeviceBatchedTopK #11328), (☐) stable sorted (Add support for stable sorted outputs to DeviceBatchedTopK #11326)
Comparison method: (✓) radix-based, (☐) radix-based for custom types, (☐) comparison-based
Segment sizes: (✓) fixed across all segments, (✓) variable per segments
Supported segment sizes: (✓) small segments (block-level), (✓) medium segments (cluster-level), (☐) large segments

`WarpTopK`, `BlockTopK`, and `ClusterTopK`

🙋 Issue for tracking WarpTopK and BlockTopK user requests and feature prioritization

`WarpTopK`

Output requirements:
- Set-membership determinism: (☐) nondeterministic, (☐) deterministic, (☐) deterministic (prefer lower indices), (☐) deterministic (prefer higher indices)
- Output ordering: (☐) unordered, (☐) unstable sorted, (☐) stable sorted
Comparison method: (☐) radix-based, (☐) radix-based for custom types, (☐) comparison-based

`BlockTopK`

Output requirements:
- Set-membership determinism: (☐) nondeterministic, (☐) deterministic, (☐) deterministic (prefer lower indices), (☐) deterministic (prefer higher indices)
- Output ordering: (☐) unordered, (☐) unstable sorted, (☐) stable sorted
Comparison method: (☐) radix-based, (☐) radix-based for custom types, (☐) comparison-based

`ClusterTopK`

Output requirements:
- Set-membership determinism: (☐) nondeterministic, (☐) deterministic, (☐) deterministic (prefer lower indices), (☐) deterministic (prefer higher indices)
- Output ordering: (☐) unordered, (☐) unstable sorted, (☐) stable sorted
Comparison method: (☐) radix-based, (☐) radix-based for custom types, (☐) comparison-based

`DeviceTopKByKey`

Output requirements:
- Set-membership determinism: (☐) nondeterministic, (☐) deterministic, (☐) deterministic (prefer lower indices), (☐) deterministic (prefer higher indices)
- Output ordering: (☐) unordered, (☐) unstable sorted, (☐) stable sorted
Comparison method: (☐) radix-based, (☐) radix-based for custom types, (☐) comparison-based

Implementation Plan and Progress

Phase 1a: Initial `DeviceTopK`

Task	Status
Determine interface for cub::DeviceTopK	✅ #5674
Add cub::DeviceTopK	✅ #5675

Phase 1b: Improve `DeviceTopK`

Task	Status
Add support for non-fundamental types to DeviceTopK	✅ #6290
Improve initial top-k implementation	Todo (#6191)

Phase 2a: Initial `DeviceBatchedTopK` for (a) small and medium segments (`<=1M`)

Add support for DeviceBatchedTopK for small and medium-sized segments of up to 1M #7585

Phase 2b: Initial `DeviceBatchedTopK` for arbitrary segment sizes

Add support for arbitrary segment sizes to DeviceBatchedTopK #9253

Phase 2c: Improve `DeviceBatchedTopK`

Task	Status
Add more cub::BlockTopK algorithm specializations/strategies	Todo
TopK optimization for small `K` using an atomics-based approach	#8274
Add family of cub::WarpTopK algorithm specializations/strategies	Todo
Tune cub::DeviceSegmentedTopK for fixed, small segment sizes	Todo
Tune cub::DeviceSegmentedTopK for larger segments	Todo
Extend and improve BlockTopK	#7799
Follow-ups to the cluster-based `DeviceBatchedTopK` backend	#10023
Support sorted output for any `k` in `DeviceBatchedTopK` (part of #9260)	#11328
Add support for stable-sorted output to `DeviceBatchedTopK` (part of #9260)	#11326

Phase 2d: Improve `DeviceBatchedTopK` for multi-key scenarios

Improve DeviceBatchedTopK for multi-key top-k #9257

Phase 3: Expose `WarpTopK`, `BlockTopK`, and `ClusterTopK`

Task	Status
Determine interface for cub::BlockTopK and cub::WarpTopK	Todo

Phase 4: Extend Feature Set for Family of Top-K Algorithms

Task	Status
Add support for custom types via decomposer_t to family of top-k algorithms	To be prioritized
Add support for deterministic top-k across the family of algorithms	To be prioritized (`DeviceBatchedTopK`: ✅ #9259)
Add support for sorted top-k across the family of algorithms	To be prioritized (`DeviceBatchedTopK`: small `k` ✅ #10811, any `k` #11328)
Add support for stable sorted top-k across the family of algorithms	To be prioritized (`DeviceBatchedTopK`: #11326)
Add support for cluster-level top-k to be used in cub::DeviceSegmentedTopK	✅ #9075

Request: Consider a stable Top-K option (defined tie policy)

Super excited to see Top-K moving into CUB. One capability that would unlock a lot of downstream value is a stable Top-K mode (i.e., when scores/keys tie, preserve input order among tied elements, or otherwise provide a documented deterministic policy).

Why stability matters

Deterministic UX & pagination. Search/recsys and analytics UIs often page by Top-K. Unstable ties cause “result flapping” between refreshes and page boundaries, degrading user trust and increasing support load.
A/B testing & metric integrity. In experimentation platforms, nondeterministic tie resolution adds variance unrelated to treatment effects, making KPIs noisier and decisions riskier.
Compliance & auditability. Finance/healthcare/public-sector workloads frequently require run-determinism for audits and reproducibility reports.
Reproducible ML pipelines.
- Feature selection / sparsification: tie instability changes which features survive, producing different models across runs.
- Vector search / quantized IR: int8/bfloat16 scores create many ties; stability prevents flaky regressions.
- Routing (e.g., MoE gating): consistent token assignment helps debugging and safety reviews.
ETL & reporting. Segmented/grouped Top-K is common in dashboards and batch jobs; stable semantics enable idempotent reruns and consistent aggregates.
Operational efficiency. Without a stable Top-K, teams often fall back to full stable sorts (wasteful) or ad-hoc secondary keys. For 64-bit keys, composing a secondary tiebreaker can imply 128-bit arithmetic or extra passes, which is costly and not always portable.
Caching & dedup. Deterministic outputs improve cache hit rates and simplify downstream deduplication logic.

Even if stability isn’t the default, having an opt-in stable mode (or clearly specified deterministic tie behavior) would let product teams choose between peak throughput and strict reproducibility, depending on their requirements.

@tedliosu

Thank you for listing out these use cases and sharing such detailed feedback regarding the importance of stability in Top-K, @tedliosu. We really appreciate your interest!

While the initial version will be unstable & unordered, we’re actively working on adding a stable, unordered top-k option to CUB. We also have plans to support ordered results in a future stage. Your input helps reinforce the broad impact of stability for downstream users, thanks again for helping guide our development.

@elstehle

@elstehle

Thanks for the quick update; it's great to hear that a stable, unordered Top-K is planned, with ordered to follow. Looking forward to trying it once it lands. Appreciate the transparency!

@tedliosu

@tedliosu, we are currently gathering user interest and requirements. Since you had already shared interest in a stable top-k algorithm, could you please help us prioritize future work items related to top-k by filling the form and/or leaving a comment on the following issue:

#6387

Thanks so much!

Navigation Menu

Repository navigation

[EPIC] Add family of top-k algorithms to CUB #5673

Tracking Issue: Top-K Algorithms in CUB

Mission & Goals

Scope

Feature Support Matrix

`DeviceTopK`

`DeviceSegmentedTopK` & `DeviceBatchedTopK`

`WarpTopK`, `BlockTopK`, and `ClusterTopK`

`WarpTopK`

`BlockTopK`

`ClusterTopK`

`DeviceTopKByKey`

Implementation Plan and Progress

Phase 1a: Initial `DeviceTopK`

Phase 1b: Improve `DeviceTopK`

Phase 2a: Initial `DeviceBatchedTopK` for (a) small and medium segments (`<=1M`)

Phase 2b: Initial `DeviceBatchedTopK` for arbitrary segment sizes

Phase 2c: Improve `DeviceBatchedTopK`

Phase 2d: Improve `DeviceBatchedTopK` for multi-key scenarios

Phase 3: Expose `WarpTopK`, `BlockTopK`, and `ClusterTopK`

Phase 4: Extend Feature Set for Family of Top-K Algorithms

Why stability matters

8 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[EPIC] Add family of top-k algorithms to CUB #5673

Description

Tracking Issue: Top-K Algorithms in CUB

Mission & Goals

Scope

Feature Support Matrix

DeviceTopK

DeviceSegmentedTopK & DeviceBatchedTopK

WarpTopK, BlockTopK, and ClusterTopK

WarpTopK

BlockTopK

ClusterTopK

DeviceTopKByKey

Implementation Plan and Progress

Phase 1a: Initial DeviceTopK

Phase 1b: Improve DeviceTopK

Phase 2a: Initial DeviceBatchedTopK for (a) small and medium segments (<=1M)

Phase 2b: Initial DeviceBatchedTopK for arbitrary segment sizes

Phase 2c: Improve DeviceBatchedTopK

Phase 2d: Improve DeviceBatchedTopK for multi-key scenarios

Phase 3: Expose WarpTopK, BlockTopK, and ClusterTopK

Phase 4: Extend Feature Set for Family of Top-K Algorithms

Activity

tedliosu commented on Aug 30, 2025

Why stability matters

elstehle commented on Sep 4, 2025

tedliosu commented on Sep 4, 2025

elstehle commented on Oct 30, 2025

8 remaining items

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`DeviceTopK`

`DeviceSegmentedTopK` & `DeviceBatchedTopK`

`WarpTopK`, `BlockTopK`, and `ClusterTopK`

`WarpTopK`

`BlockTopK`

`ClusterTopK`

`DeviceTopKByKey`

Phase 1a: Initial `DeviceTopK`

Phase 1b: Improve `DeviceTopK`

Phase 2a: Initial `DeviceBatchedTopK` for (a) small and medium segments (`<=1M`)

Phase 2b: Initial `DeviceBatchedTopK` for arbitrary segment sizes

Phase 2c: Improve `DeviceBatchedTopK`

Phase 2d: Improve `DeviceBatchedTopK` for multi-key scenarios

Phase 3: Expose `WarpTopK`, `BlockTopK`, and `ClusterTopK`