NVIDIA / NVIDIA/cccl

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

Open
#7,585 0 comments 0 reactions 1 assignee Claimed by @elstehle View on GitHub
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

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.