NVIDIA / NVIDIA/cccl

[BUG]: `cub::DeviceHistogram::HistogramEven` produces garbage bins for `__nv_bfloat16` samples

Open
#10,940 0 comments 0 reactions 1 assignee Claimed by @NaderAlAwar 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

Silent Failure

### Component

CUB

### Describe the bug

`DeviceHistogram::HistogramEven` with `__nv_bfloat16` samples/levels compiles cleanly but computes garbage bin indices: in practice every valid sample lands in a single bin (bin 0 or 1) regardless of its value. `__half` and `float` work correctly.

Root cause, in `cub/device/dispatch/kernels/kernel_histogram.cuh` (`Transforms::ScaleTransform`):

- `ComputeScale` and `SampleIsValid` have dedicated `__nv_bfloat16` overloads, but **`ComputeBin` only has a `__half` overload**. For `__nv_bfloat16`, the generic `ComputeBin` template is selected, and its tag dispatch uses `cuda::std::is_floating_point` — which is `false` for `__nv_bfloat16` — so it takes the non-floating-point path:

```cpp
return static_cast(((sample - min_level) * scale.fraction.bins) / scale.fraction.range);
```

However, the `__nv_bfloat16` `ComputeScale` overload initialized the *other* member of the `ScaleT` union (`scale.reciprocal`), so `ComputeBin` reads uninitialized/aliased union bytes and every sample maps to the same bin.

- Additionally, the `__nv_bfloat16` `SampleIsValid` overload is missing the `const` qualifier, so the (const) `BinSelect` caller silently falls back to the generic template instead of the intended overload. This happens to be functionally harmless today, but it means the dedicated overload is dead code.

The fix is to add a `ComputeBin(__nv_bfloat16, __nv_bfloat16, ScaleT)` overload mirroring the `__half` one (using `scale.reciprocal`, with an `NV_PROVIDES_SM_80` native path and a float fallback), and to add `const` to the `__nv_bfloat16` `SampleIsValid` overload.

Scope: the bug is in `ScaleTransform`, which only the *Even* entry points use, so it affects `HistogramEven` and `MultiHistogramEven` (all overloads funnel into the same dispatch). `HistogramRange` uses `SearchTransform` (an `UpperBound` search over the levels array) and was verified to produce correct results for `__nv_bfloat16` with the same data.

### How to Reproduce

```cpp
// bf16_hist_repro.cu
#include

#include
#include

#include

#include
#include
#include

int main()
{
constexpr int num_samples = 1000;
constexpr int num_levels = 9; // 8 bins

std::mt19937 rng(42);
std::uniform_real_distribution dist(0.0f, 8.0f);
std::vector<__nv_bfloat16> h_samples(num_samples);
std::vector expected(num_levels - 1, 0);
for (int i = 0; i < num_samples; ++i)
{
h_samples[i] = __float2bfloat16(dist(rng));
const float s = __bfloat162float(h_samples[i]);
if (s >= 0.0f && s < 8.0f)
{
++expected[static_cast(s)];
}
}

thrust::device_vector<__nv_bfloat16> d_samples(h_samples);
thrust::device_vector d_histogram(num_levels - 1, 0);

const __nv_bfloat16 lower = __float2bfloat16(0.0f);
const __nv_bfloat16 upper = __float2bfloat16(8.0f);

size_t temp_bytes = 0;
cub::DeviceHistogram::HistogramEven(
nullptr,
temp_bytes,
thrust::raw_pointer_cast(d_samples.data()),
thrust::raw_pointer_cast(d_histogram.data()),
num_levels,
lower,
upper,
num_samples);
thrust::device_vector d_temp(temp_bytes);
cub::DeviceHistogram::HistogramEven(
thrust::raw_pointer_cast(d_temp.data()),
temp_bytes,
thrust::raw_pointer_cast(d_samples.data()),
thrust::raw_pointer_cast(d_histogram.data()),
num_levels,
lower,
upper,
num_samples);
cudaDeviceSynchronize();

thrust::host_vector h_histogram = d_histogram;
printf("bin: got expected\n");
bool ok = true;
for (int i = 0; i < num_levels - 1; ++i)
{
printf("%3d: %8d %8d\n", i, h_histogram[i], expected[i]);
ok &= (h_histogram[i] == expected[i]);
}
printf(ok ? "PASS\n" : "FAIL\n");
return ok ? 0 : 1;
}
```

From a CCCL checkout:

```bash
nvcc -std=c++17 -arch=sm_89 -Icub -Ithrust -Ilibcudacxx/include bf16_hist_repro.cu -o bf16_hist_repro
./bf16_hist_repro
```

### Expected behavior

Observed output — every sample lands in one bin:

```
bin: got expected
0: 1000 138
1: 0 109
2: 0 132
3: 0 117
4: 0 127
5: 0 135
6: 0 104
7: 0 138
FAIL
```

Replacing `__nv_bfloat16` with `__half` (and the conversion intrinsics accordingly) in the same program produces the correct histogram.

### Reproduction link

_No response_

### Operating System

_No response_

### nvidia-smi output

_No response_

### NVCC version

_No response_

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.