CUDA: Fix Integer Reduction Errors On Blackwell GPUs
Introduction
This article addresses a critical issue encountered when using CUDA integer reductions (fast_max, fast_min, fast_argmax, fast_argmin) on NVIDIA Blackwell (sm_100) GPUs and potentially newer architectures. The problem stems from the incorrect initialization of shared memory buffers with floating-point infinities (-INFINITY/INFINITY) even when dealing with integer types like int64_t, u32, and u8. This undefined behavior (UB) leads to garbage values in shared memory, causing functions like max_all(i64) to return unexpectedly large and incorrect integer values, particularly after running conv2d kernels. This, in turn, disrupts downstream logic, such as RoPE cache sizing in candle-core applications.
The root cause lies in the candle-kernels library, where the integer reduction kernels are initialized with -INFINITY/INFINITY. Casting these floating-point infinities to integers results in undefined behavior. On Blackwell GPUs, this UB manifests as seemingly random, incorrect initial values in shared memory. Consequently, operations like max_all(i64) can return huge, nonsensical integers, especially after executing a conv2d kernel. This issue has a cascading effect, breaking downstream logic in applications that rely on candle-core, such as RoPE cache sizing.
To fully grasp the issue, let's delve into the specifics. Integer reduction operations, like finding the maximum or minimum value in a set of integers, are fundamental in many CUDA-accelerated algorithms. These operations often involve shared memory to facilitate efficient communication and computation among threads within a block. When shared memory is incorrectly initialized, the results of these reduction operations become unreliable. In the context of candle-kernels, the use of floating-point infinities to initialize integer shared memory buffers introduces a subtle but significant flaw. While this might not cause issues on all GPU architectures, it becomes problematic on Blackwell GPUs due to the way the hardware handles the conversion of floating-point infinities to integers. This conversion leads to the introduction of garbage values into shared memory, which then corrupts the results of the reduction operations.
The impact of this issue extends beyond mere incorrect numerical results. It can lead to unpredictable behavior in downstream applications, making it difficult to debug and diagnose problems. For example, in the case of RoPE cache sizing, incorrect maximum values can lead to the allocation of insufficient memory, resulting in crashes or incorrect computations. Therefore, addressing this issue is crucial for ensuring the reliability and correctness of CUDA-accelerated applications running on Blackwell GPUs.
Reproducing the Issue: Minimal Examples
To demonstrate the problem, consider the following minimal example using candle-core. This example sets up a convolutional neural network (CNN) operation and then calculates the maximum value of a tensor before and after the convolution. The discrepancy between the expected maximum value and the actual value obtained on the GPU highlights the issue.
[dependencies]
anyhow = "1.0.100"
candle-core = { path = "../candle/candle-core", features = ["cuda"] }
The Cargo.toml file specifies the dependencies for the project, including anyhow for error handling and candle-core with the cuda feature enabled. This ensures that the code can utilize CUDA-enabled GPUs.
use anyhow::Result;
use candle_core::{DType, Device, Tensor};
fn main() -> Result<()> {
let device = Device::new_cuda(0).unwrap_or_else(|e| panic!("failed to init CUDA device {e:?}"));
let cpu = Device::Cpu;
// input [1, 3, 1024, 1024]，BF16
let b_size: usize = 1;
let c_in: usize = 3;
let i_h: usize = 1024;
let i_w: usize = 1024;
// [768, 3, 16, 16]，BF16
let c_out: usize = 768;
let k_h: usize = 16;
let k_w: usize = 16;
// conv2d hp
let padding: usize = 0;
let stride: usize = 16;
let dilation: usize = 1;
let groups: usize = 1;
println!(
"Input shape: [{b_size}, {c_in}, {i_h}, {i_w}], dtype=BF16, device={:?}",
device
);
println!(
"Kernel shape: [{c_out}, {c_in}, {k_h}, {k_w}], dtype=BF16, device={:?}",
device
);
println!(
"Conv2d params: padding={padding}, stride={stride}, dilation={dilation}, groups={groups}"
);
let x = Tensor::rand(0f32, 1f32, (b_size, c_in, i_h, i_w), &device)?.to_dtype(DType::BF16)?;
let k = Tensor::rand(0f32, 1f32, (c_out, c_in, k_h, k_w), &device)?.to_dtype(DType::BF16)?;
println!(
"x.shape = {:?}, x.stride = {:?}",
x.shape().dims(),
x.stride()
);
println!(
"k.shape = {:?}, k.stride = {:?}",
k.shape().dims(),
k.stride()
);
let batch: usize = 1;
let q_len: usize = 912;
let start: i64 = 0;
let end: i64 = start + q_len as i64;
println!("\n[pre] constructing ids_pre on CUDA: start={start}, end={end}, q_len={q_len}");
let ids_pre = Tensor::arange(start, end, &device)?
.reshape((1, q_len))?
.expand((batch, q_len))?
.contiguous()?
.to_dtype(DType::I64)?;
println!(
"ids_pre.shape={:?}, ids_pre.stride={:?}, device={:?}, dtype={:?}",
ids_pre.shape().dims(),
ids_pre.stride(),
ids_pre.device(),
ids_pre.dtype()
);
let max_gpu_pre = ids_pre.max_all()?.to_scalar::<i64>()?;
let max_cpu_pre = ids_pre.to_device(&cpu)?.max_all()?.to_scalar::<i64>()?;
println!("[pre] max_gpu_pre = {max_gpu_pre}, max_cpu_pre = {max_cpu_pre}");
println!("\nRunning conv2d...");
let y = x.conv2d(&k, padding, stride, dilation, groups)?;
println!(
"conv2d output shape = {:?}, stride = {:?}, dtype = {:?}",
y.shape().dims(),
y.stride(),
y.dtype()
);
println!("\n[post] constructing ids_post on CUDA: start={start}, end={end}, q_len={q_len}");
let ids_post = Tensor::arange(start, end, &device)?
.reshape((1, q_len))?
.expand((batch, q_len))?
.contiguous()?
.to_dtype(DType::I64)?;
println!(
"ids_post.shape={:?}, ids_post.stride={:?}, device={:?}, dtype={:?}",
ids_post.shape().dims(),
ids_post.stride(),
ids_post.device(),
ids_post.dtype()
);
let max_gpu_post = ids_post.max_all()?.to_scalar::<i64>()?;
let max_cpu_post = ids_post.to_device(&cpu)?.max_all()?.to_scalar::<i64>()?;
println!("[post] max_gpu_post = {max_gpu_post}, max_cpu_post = {max_cpu_post}");
println!("\n=== summary ===");
println!("before conv2d: GPU max_all(i64) = {max_gpu_pre}, CPU max_all(i64) = {max_cpu_pre}");
println!(
"after conv2d: GPU max_all(i64) = {max_gpu_post}, CPU max_all(i64) = {max_cpu_post}"
);
if max_gpu_post != max_cpu_post {
eprintln!("\n[BUG] After conv2d, GPU max_all(i64) != CPU max_all(i64).");
} else {
eprintln!("\n[INFO] This is what we want.");
}
Ok(())
}
The main.rs file demonstrates the issue. It initializes two tensors, performs a convolution operation, and then calculates the maximum value of a range of integers before and after the convolution. The key is the comparison between max_gpu_post and max_cpu_post. If they differ, it indicates the bug.
Expected vs. Actual Output
Ideally, the maximum value calculated on the GPU (max_gpu_post) should match the maximum value calculated on the CPU (max_cpu_post). However, the following output demonstrates the issue:
Input shape: [1, 3, 1024, 1024], dtype=BF16, device=Cuda(CudaDevice(DeviceId(1)))
Kernel shape: [768, 3, 16, 16], dtype=BF16, device=Cuda(CudaDevice(DeviceId(1)))
Conv2d params: padding=0, stride=16, dilation=1, groups=1
x.shape = [1, 3, 1024, 1024], x.stride = [3145728, 1048576, 1024, 1]
k.shape = [768, 3, 16, 16], k.stride = [768, 256, 16, 1]
[pre] constructing ids_pre on CUDA: start=0, end=912, q_len=912
ids_pre.shape=[1, 912], ids_pre.stride=[912, 1], device=Cuda(CudaDevice(DeviceId(1))), dtype=I64
[pre] max_gpu_pre = 911, max_cpu_pre = 911
Running conv2d...
conv2d output shape = [1, 768, 64, 64], stride = [3145728, 4096, 64, 1], dtype = BF16
[post] constructing ids_post on CUDA: start=0, end=912, q_len=912
ids_post.shape=[1, 912], ids_post.stride=[912, 1], device=Cuda(CudaDevice(DeviceId(1))), dtype=I64
[post] max_gpu_post = 4851116375995599755, max_cpu_post = 911
=== summary ===
before conv2d: GPU max_all(i64) = 911, CPU max_all(i64) = 911
after conv2d: GPU max_all(i64) = 4851116375995599755, CPU max_all(i64) = 911
[BUG] After conv2d, GPU max_all(i64) != CPU max_all(i64).
As you can see, before the conv2d operation, both the GPU and CPU calculate the correct maximum value (911). However, after the conv2d operation, the GPU calculates a drastically different and incorrect maximum value (4851116375995599755), while the CPU still calculates the correct value (911). This discrepancy highlights the bug caused by the incorrect initialization of shared memory in the CUDA integer reduction kernels.
Impact and Mitigation
The incorrect initialization of shared memory buffers with floating-point infinities in CUDA integer reduction kernels has a significant impact on the accuracy and reliability of computations performed on NVIDIA Blackwell GPUs. This issue leads to undefined behavior, resulting in garbage values in shared memory and causing functions like max_all(i64) to return incorrect results, especially after running conv2d kernels. This, in turn, can break downstream logic in applications that rely on these reduction operations, such as RoPE cache sizing in candle-core.
Mitigation Strategies
Several strategies can be employed to mitigate this issue and ensure the correctness of computations on Blackwell GPUs:
-
Correct Initialization: The most direct solution is to modify the
candle-kernelslibrary to initialize shared memory buffers with appropriate integer values instead of floating-point infinities. For example, forfast_maxkernels, the buffers should be initialized with the smallest possible integer value for the given data type (e.g.,i64::MINforint64_t). Similarly, forfast_minkernels, the buffers should be initialized with the largest possible integer value (e.g.,i64::MAXforint64_t). -
Workarounds: In cases where modifying the kernel code is not immediately feasible, workarounds can be implemented to avoid the issue. One approach is to perform the reduction operations on the CPU instead of the GPU, ensuring that the computations are performed correctly. However, this approach may introduce a performance overhead, as data needs to be transferred between the CPU and GPU.
-
Conditional Compilation: Another strategy is to use conditional compilation to selectively disable the problematic kernels on Blackwell GPUs and use alternative implementations instead. This approach allows the code to remain compatible with other GPU architectures while avoiding the issue on Blackwell GPUs.
-
Data Type Considerations: Carefully consider the data types used for reduction operations. If possible, use floating-point data types instead of integer data types, as the issue is specific to integer reduction kernels. However, this approach may not be feasible in all cases, as it may require significant changes to the code and may impact performance.
-
Testing and Validation: Thoroughly test and validate the code on Blackwell GPUs to ensure that the issue is not present. This can be done by comparing the results of computations performed on the GPU with the results of computations performed on the CPU.
Conclusion
The CUDA integer reduction issue on Blackwell GPUs highlights the importance of careful initialization and data type handling in CUDA programming. By understanding the root cause of the issue and implementing appropriate mitigation strategies, developers can ensure the accuracy and reliability of their CUDA-accelerated applications. The key takeaway is that seemingly minor details, such as the initialization of shared memory buffers, can have a significant impact on the correctness of computations, especially on newer GPU architectures like Blackwell.
For further information on CUDA best practices, refer to the NVIDIA CUDA Documentation.
