Add support for ordered outputs to `DeviceBatchedTopK`
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
Currently ordering requirements can't be fulfilled meaning that the least-surprising default (deterministic, prefer small index, ordered output) does not compile and users need to opt out of ordered output.
Given the difference in parallelism between the TopK selection and the output sorting (segments size is typically much bigger than K) and the fact that selected keys are distributed among cluster CTAs, it is not clear if kernel fusion is actually worth it. To make the interface usable ASAP we want to start by just launching a segmented sort following the TopK selection kernel. Allowing CUDA Graph-capture to work is a priority, so we need to use `DeviceSegmentedRadixSort` instead of `DeviceSegmentedSort`.
This issue can be closed once the full requirements interface is functional for segment sizes and GPU architectures supported by the cluster backend.
Contributor guide
Assessment
This issue has not been assessed yet.