[EPIC]: Support device-resident problem sizes in CUB device-level algorithms
- 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.
## Context
A typical CUB device-level interface looks roughly like this:
```
cub::DeviceAlg::Alg(int* d_in, int* d_out, int num_items, stream);
```
Because the interface is asynchronous, the memory referenced by `d_in` and `d_out` is accessed later in stream order. However, scalar parameters such as `num_items` are consumed immediately by the host-side dispatch path because they affect temporary storage sizing and grid configuration.
In some workflows, these scalar parameters are themselves produced on the device. For example, the `num_items` passed to a device scan may come from `d_num_selected`, which was produced by a preceding device select operation. Copying such values from device memory back to the host introduces a host dependency, breaks asynchrony, and prevents capturing the full workflow as a reusable CUDA Graph.
## Existing Recommendation
In some cases, we can work around device-resident scalar parameters by adding algorithm overloads that accept `cub::FutureValue`. For example, this is already done for the initial value in [`DeviceScan::ExclusiveScan`](https://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceScan.html#_CPPv4I000000EN3cub10DeviceScan13ExclusiveScanE11cudaError_tPvR6size_t14InputIteratorT15OutputIteratorT7ScanOpT11FutureValueI10InitValueT14InitValueIterTE9NumItemsT12cudaStream_t).
For problem size, the required changes are more involved because `num_items` affects temporary storage sizing and launch configuration. For use cases where the problem size is device-resident and CUDA Graph capture is required, the existing recommendation has generally been to launch CUB from a parent kernel using CUDA Dynamic Parallelism (CDP):
```c++
__global__ void kernel(int* d_in, int* d_out, const int* d_num_items) {
cub::DeviceAlg::Alg(d_in, d_out, *d_num_items);
}
```
## Problem
In some cases, a CDP-based solution for device-resident problem size is not efficient. The table below reports normalized elapsed times for four use cases. Lower is better. Source code for these benchmarks is available [here](https://github.com/gevtushenko/cubcdp).
Each variant means:
- **base** - host-side launch in whole-program mode (`-rdc=false`)
- **rdc** - host-side launch in separable compilation mode (`-rdc=true`)
- **cdp** - device-side launch in separable compilation mode
- **copy** - copy `num_items` from device memory, then perform a host-side launch in whole-program mode
Here, we use **base** as the speed-of-light (SOL) baseline. RDC can affect generated SASS. It is a prerequisite for CDP, which is the focus of this analysis, but some projects enable RDC for other reasons. Therefore, we include **rdc** as a separate data point because it may be the relevant baseline for those projects.
### `cub::DeviceRadixSort`
For `cub::DeviceRadixSort`, the cost of the suggested workarounds depends heavily on problem size.
On A6000 Ada, CDP and copy-based dispatch are close to the **base** elapsed time for large problem sizes. At `2^28` elements, both are within roughly 1-2% of **base**. For smaller problem sizes, however, fixed overhead dominates: at `2^16` elements, CDP takes more than 2x as long as **base**.
On B200, the trend is similar, but the problem size required to amortize CDP and copy overheads is much larger. Even at `2^28` elements, CDP remains substantially slower than **base**, while the copy-based workaround is closer but still not equivalent.
### `cub::DeviceScan`
`cub::DeviceScan` has two relevant prefix-sum implementations. On A6000 Ada, CUB uses the decoupled look-back implementation, and the trend is similar to `cub::DeviceRadixSort`: the CDP and copy-based workarounds become more competitive as the problem size grows, but fixed overheads dominate at smaller sizes.
On B200, CUB uses warpspeed scan. This implementation depends heavily on occupancy, which is affected by CDP. As a result, the CDP workaround is roughly 4x slower than **base** across the measured problem sizes, independent of `num_items`.
The copy-based workaround avoids the CDP occupancy issue, but it still breaks asynchrony and CUDA Graph capture, so it is not a viable solution for graph-captured workflows.
### Describe the solution you'd like
Instead of recommending CDP as the primary workaround, CUB should support device-resident problem sizes in device-level algorithm interfaces for motivated use cases.
Users should be able to pass `num_items` as a device-accessible scalar, while also providing any host-accessible bounds required for temporary storage sizing and launch configuration. This could be modeled with the existing `cub::FutureValue` mechanism, or through the more general parameter/guarantees framework proposed in #7495.
The important property is that the active problem size can be produced by earlier GPU work, consumed by a later CUB algorithm in stream order, and used inside CUDA Graph capture without a device-to-host synchronization or CDP launch.
### Describe alternatives you've considered
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.