[EPIC] Add family of top-k algorithms to CUB
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 487
- Avg merge
- 2d 7h
- Merged PRs (30d)
- 296
Description
## 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:**
1. **Eliminate redundancy:** Provide a definitive family of high-performance algorithms to replace the fragmented, ad-hoc implementations currently scattered across the ecosystem.
2. **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.
3. **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:
1. [**Feature Support Matrix:**](#section-a) Detailed tracking of supported configurations (e.g., deterministic, stable, segmented) across all CUB scopes (`Device`, `Block`, `Warp`).
2. [**Implementation Roadmap:**](#section-b) The chronological plan for delivering these features.
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](https://github.com/NVIDIA/cccl/issues/6387)
**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](https://github.com/NVIDIA/cccl/issues/6290)), (β) comparison-based
### `DeviceSegmentedTopK` & `DeviceBatchedTopK`
[π Issue for tracking `DeviceSegmentedTopK` and `DeviceBatchedTopK` user requests and feature prioritization](https://github.com/NVIDIA/cccl/issues/6391)
- 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
- 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](https://github.com/NVIDIA/cccl/issues/6389)
### `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`)
- [ ] https://github.com/NVIDIA/cccl/issues/7585
### Phase 2b: Initial `DeviceBatchedTopK` for arbitrary segment sizes
- [ ] https://github.com/NVIDIA/cccl/issues/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
### Phase 2d: Improve `DeviceBatchedTopK` for multi-key scenarios
- [ ] https://github.com/NVIDIA/cccl/issues/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
Add support for sorted top-k across the family of algorithms | To be prioritized
Add support for stable sorted top-k across the family of algorithms | To be prioritized
Add support for cluster-level top-k to be used in cub::DeviceSegmentedTopK | To be prioritized
Contributor guide
Assessment
This issue has not been assessed yet.