Add support for `DeviceBatchedTopK` for small and medium-sized segments of up to 1M
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
The goal of this issue is to establish the functionality for `cub::DeviceBatchedTopK` that works for segment sizes of up to `1M` items. This issue can be closed when feature matrix (see [Scope](#Scope)) is covered and once the algorithm is integrated into FlashInfer.
# Scope
- **supported segment sizes**: `[1, 1M]`
- **supported k values**: `[1, 8192]`
- **batch sizes**: `[1, 64]`
- **supported devices:** B200, B300
- **support for cuda graph launches**
## Output requirements feature matrix:
| | unordered | ordered |
|:--------------------------------------------------------------------------------------------------------------------------:|:---------:|:-------:|
| non-deterministic
`determinism::not_guaranteed` | ⏳ | 🟡 |
| deterministic; implementation-defined tie-break
`determinism::run_to_run` | ⏳ | 🟡 |
| deterministic; ties resolved toward smaller source index
`determinism::gpu_to_gpu, tie_break::prefer_smaller_index` | ⏳ | 🟡 |
| deterministic; ties resolved toward larger source index
`determinism::gpu_to_gpu, tie_break::prefer_larger_index` | 🟡 | 🟡 |
### On output requirements
We have documented the behaviour for various output requirements in this issue:
- https://github.com/NVIDIA/cccl/issues/9354
# Progress
## Implementation for small (CTA-sized), variable-size segments
Task | Status
-- | --
Gather users & libraries and plan cub::DeviceBatchededTopK | ✅ Done
Add initial benchmark suite for cub::DeviceBatchededTopK | ✅ #6508
Add `DispatchSegmentedTopK` that's flexible enough to cover arbitrary
per-segment and uniform parameter combinations | ✅ #6864
Add tests for `DeviceBatchededTopK` with fixed-size segments | ✅ #6863
Add cub::BlockTopK based on AIR Top-K | ✅ #7092
Add implementation for small, variable-size segments to DeviceBatchedTopK | ✅ #7839
## Interface for small (CTA-sized), variable-size segments
Status
- [x] #7495
- [x] #9254
- [x] #7616
Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Determine Argument Annotation Framework**
Determine argument annotation framework for `cub::DeviceBatchededTopK` | ✅ #7495 | 2 weeks | 2 weeks
**Expose Argument Annotation Framework**
Expose the argument annotation publicly | ✅ #9254 | 1 week | 1 week
**DeviceBatchededTopK API:**
Expose `cub::DeviceBatchededTopK` for small segment sizes | ✅ #7616 | 1.5 weeks | 1.5 weeks
## Support for medium-sized segment sizes (`<=1M`)
Status
- [x] #8364
- [x] #9075
- [x] #8360
Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Load Balancing Top-K**
Delegate large segments to multi-cta implementation | ✅ #8364 | 2 weeks | 1.5 weeks
**Cluster-Optimized Top-K**
Optimize segmented Top-K with clusters for sm100f | ⏳ #9075 | 3 weeks | 1.5 weeks
**DeviceBatchededTopK API:**
Expose `cub::DeviceBatchededTopK` | Todo #8360 | 0.5 week | 0.5 weeks
## Support for deterministic, tie-breaking, & ordered output requirements
Status
- [x] #9255
- [x] #9259
- [ ] #9260
Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Introduce tie-breaking to requirements API**
Introduce an option to specify a `tie_breaking` behavior | ✅ #9255 | 0.5 weeks | 0.5 weeks
**Add support for deterministic outputs**
Add support for deterministic outputs (respecting tie-breaking criterion) | ⏳ #9259 | 1 week | 1 week
**Add support for ordered outputs**
Add support for ordered outputs | #9260 | 0.5 weeks | 0.5 weeks
## Integration into FlashInfer
Status
- [x] #8370
- [x] #8371
Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Transition FlashInfer to use CCCL from GitHub**
Make FlashInfer use CCCL from Github | ✅ #8370 | 2 weeks | 2 weeks
**FlashInfer integration MVP**
Integrate `DeviceBatchededTopK` into FlashInfer | Todo #8371 | 2 weeks | 2 weeks
Contributor guide
Assessment
This issue has not been assessed yet.