[BUG]: CUB stable sort pairt fails during graph capture for certain sizes
- 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 bug and that I agree to the [Code of Conduct](CODE_OF_CONDUCT.md)
### Type of Bug
Runtime Error
### Component
CUB
### Describe the bug
Here is code to reproduce the issue. It fails with the following message:
```
CUDA Error at sort.cu:72 - operation failed due to a previous error during capture
```
Should this work? Graph capture works in other cases.. but it would be useful to know if this is supposed to work and if so presumably it is a bug.
```c++
#include
#include
#include
#include
#include
// Macro to check for CUDA errors
#define CUDA_CHECK(call) \
do { \
cudaError_t error = call; \
if (error != cudaSuccess) { \
std::cerr << "CUDA Error at " << __FILE__ << ":" << __LINE__ << " - " << cudaGetErrorString(error) << std::endl; \
exit(1); \
} \
} while (0)
int main() {
// Anything above 500 doesn't work
const int num_segments = 501;
const int segment_size = 128;
const int num_items = num_segments * segment_size;
std::vector h_offsets(num_segments + 1);
for (int i = 1; i < num_segments + 1; ++i) {
h_offsets[i] = h_offsets[i - 1] + segment_size;
}
// Keys to be sorted
std::vector h_keys_in(num_items);
std::vector h_values_in(num_items);
std::iota(h_values_in.begin(), h_values_in.end(), 0);
float* d_keys_in = nullptr;
float* d_keys_out = nullptr;
int* d_values_in = nullptr;
int* d_values_out = nullptr;
int* d_offsets = nullptr;
CUDA_CHECK(cudaMalloc(&d_keys_in, num_items * sizeof(int)));
CUDA_CHECK(cudaMalloc(&d_keys_out, num_items * sizeof(int)));
CUDA_CHECK(cudaMalloc(&d_values_in, num_items * sizeof(int)));
CUDA_CHECK(cudaMalloc(&d_values_out, num_items * sizeof(int)));
CUDA_CHECK(cudaMalloc(&d_offsets, (num_segments + 1) * sizeof(int)));
CUDA_CHECK(cudaMemcpy(d_keys_in, h_keys_in.data(), num_items * sizeof(int), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_values_in, h_values_in.data(), num_items * sizeof(int), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_offsets, h_offsets.data(), (num_segments + 1) * sizeof(int), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaDeviceSynchronize());
// --- CUB segmented sort ---
void* d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;
cudaStream_t stream;
cudaStreamCreate(&stream);
cudaGraph_t graph;
cudaGraphExec_t instance;
cub::DeviceSegmentedSort::StableSortPairs(
d_temp_storage, temp_storage_bytes,
d_keys_in, d_keys_out, d_values_in, d_values_out,
num_items, num_segments, d_offsets, d_offsets + 1);
CUDA_CHECK(cudaMalloc(&d_temp_storage, temp_storage_bytes));
CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
cub::DeviceSegmentedSort::StableSortPairs(
d_temp_storage, temp_storage_bytes,
d_keys_in, d_keys_out, d_values_in, d_values_out,
num_items, num_segments, d_offsets, d_offsets + 1);
// Raises error
CUDA_CHECK(cudaStreamEndCapture(stream, &graph));
}
```
It fails on Hopper and Ampere:
- CUDA Toolkit 12.9
- Driver 550.90.12
Build and run:
```
nvcc -arch=native sort.cu -o sort
./sort
```
### How to Reproduce
See above
### Expected behavior
It should work I think (?)
### Reproduction link
_No response_
### Operating System
_No response_
### nvidia-smi output
_No response_
### NVCC version
```
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2025 NVIDIA Corporation
Built on Tue_May_27_02:21:03_PDT_2025
Cuda compilation tools, release 12.9, V12.9.86
Build cuda_12.9.r12.9/compiler.36037853_0
```
Contributor guide
Assessment
This issue has not been assessed yet.