NVIDIA / NVIDIA/cccl

[BUG]: CUB stable sort pairt fails during graph capture for certain sizes

Open
#5,661 5 comments 0 reactions 1 assignee Claimed by @nanan-nvidia View on GitHub
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

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.