GPU kernel reductions get predictable bitwise behavior with new techniques
Taming Bitwise Behavior in GPU Kernels with Tensor Core: Black-Box Reconstruction, Compiler Enforcement, and Static Verification
Distributed, Parallel, and Cluster ComputingPerformanceProgramming Languages
Summary
When running math on GPUs, small differences in the order and way numbers are combined can cause tiny but important differences in results. The authors studied what causes these differences, especially in popular GPU math routines. They created a way to understand and replicate the exact behavior of closed-source GPU libraries and made improvements to an open-source compiler to control these differences. They also built tools to check if two GPU programs will produce exactly the same bit-level results. This helps make GPU calculations more consistent and reliable.
What this means in practice
- •For machine learning engineers: Ensure exact reproducibility of GPU computations in large language model training and inference by controlling bitwise behavior of matrix multiplications and reductions.
- •For gpu compiler developers: Integrate static bit-equivalence checking and optimized reduction scheduling to improve performance and correctness guarantees in GPU kernel autotuning.
Authors
Ziteng Yang, Nicholas J. Riasanovsky, Warren Deng, Vivek Sarkar
Abstract
Determinism and numerical reproducibility are increasingly required of GPU kernels in machine learning systems, yet deterministic implementations of the same kernel can still differ bit for bit. Floating-point reduction order is the primary cause, alongside partial-sum precision, fused multiply-add operations, and rounding placement. These choices may be hand-coded, selected by a block-level language such as Triton, or hidden inside a closed-source library such as cuBLAS or rocBLAS. A tile shape chosen for performance therefore also determines the arithmetic, potentially breaking batch invariance. Preserving a fixed order can cost up to 20 percent, while an autotuner cannot identify which configurations are bitwise equivalent. We characterize the factors determining the bitwise behavior of reductions and general matrix multiplication (GEMM). First, we introduce a descriptor of GEMM reduction order, including the partitioning of K in split-K GEMM. Using it, we perform the first black-box reconstruction of a closed-source library's arithmetic for bit-level correctness. Our family of Triton GEMMs matches NVIDIA cuBLAS in all tested cases on Blackwell and Hopper. For realistic LLM shapes with fused epilogues, it matches or exceeds torch.compile performance. Second, we enforce balanced-tree reduction during Triton lowering and introduce a data-layout optimization that brings 19 of 27 kernels on GB300 and H100 within 10 percent of free-order performance. Third, we develop sound static checkers for bitwise equivalence between compiled GPU kernels, including the first checker spanning NVIDIA PTX and AMD GCN. Integrated into Triton's autotuner, the checker restricts search to a single bit-equivalence class.