NVIDIA / NVIDIA/cccl

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

Open
#5,673 4 comments 2 reactions 1 assignee Claimed by @elstehle View on GitHub
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.

## 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](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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.