NVIDIA / NVIDIA/cccl

[FEA]: CUB Sort Deviceless Scratch Size

Open
#9,814 10 comments 0 reactions 1 assignee Claimed by @Jacobfaib View on GitHub
cub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

### Is this a duplicate?

- [x] I confirmed there appear to be no [duplicate issues](https://github.com/NVIDIA/cccl/issues) for this request and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)

### Area

CUB

### Is your feature request related to a problem? Please describe.

We (at OpenXLA) would like to be able to obtain the amount of temporary storage needed for the CUB sorting operations, without needing a device to be attached.

In particular we'd need a deviceless API that provides us the scratch space required for the following operations:
* `cub::DeviceRadixSort::SortPairs` / `SortPairsDescending`
* `cub::DeviceSegmentedRadixSort::SortPairs` / `SortPairsDescending`
* `cub::DeviceRadixSort::SortKeys` / `SortKeysDescending`
* `cub::DeviceSegmentedRadixSort::SortKeys` / `SortKeysDescending`

It is enough if this solution only works for Hopper+ devices.

### Describe the solution you'd like

We originally created the [following doc](https://docs.google.com/document/d/1OWekzUxP7Lc3GbMrmTPnuy8AqfqaBzjw2B5YpPkPdGY/edit?usp=sharing) proposing a way to pass the properties of the device we are targeting, so that CUB doesn't need to query the actual device to obtain the information it needs to calculate the temporary storage size (very similar to the cuDNN AOT mode).

@danielfrg mentioned that CUB may already support deviceless friendly algorithms for these cases (e.g. OneSweep for `cub::DeviceRadixSort`), that we could select using the Tuning API.

### Describe alternatives you've considered

_No response_

### Additional context

XLA needs to calculate the scratch sizes during compile time, so this would allow us to compile models using CUB sort without needing a GPU to be available.

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.