PaperlayerView exact revision on arXiv

Source-linked AI summary

Taming Bitwise Behavior in GPU Kernels with Tensor Core: Black-Box Reconstruction, Compiler Enforcement, and Static Verification

Ziteng Yang, Nicholas J. Riasanovsky, Warren Deng, Vivek Sarkar

arXiv:2609.11356v1cs.DCcs.PFcs.PL

TL;DR

GPU kernels can produce different bits for the same computation because reduction order and related arithmetic choices may vary, including with batch shape. This paper characterizes and controls those choices through a GEMM descriptor, library reconstruction, compiler enforcement, and static equivalence checking. It reports bit-matching cuBLAS GEMMs, improved performance under fixed-order reductions, and checkers integrated into Triton’s autotuner.

  • Problem

    GPU kernels can differ bitwise because reduction order and other floating-point choices vary, while libraries may hide their arithmetic and autotuners cannot identify bit-equivalent configurations.

    Source: §1 · p. 1

  • Method

    The paper characterizes reduction order, describes GEMM behavior with a descriptor, reconstructs cuBLAS arithmetic, enforces balanced-tree reductions, and statically checks compiled-kernel equivalence.

    Source: §1 · pp. 2–3

  • Results

    Triton GEMMs bit-match cuBLAS on tested cases, and balanced-tree enforcement with layout optimization brings 19 of 27 kernels within 10% of free-order mode.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Static equivalence checking can prune Triton autotuning to a single bit-equivalence class while preserving the search for performance.

    Source: §1 · p. 3

    For complex kernels such as flash attention, the paper provides a theoretical reduction-order description but leaves implementation open.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames bitwise reproducibility as a systems requirement because batch-dependent reduction order can change model outputs. It presents ways to describe, reconstruct, enforce, and verify the arithmetic behavior of GPU kernels.

1 Introduction

  • Batch-dependent kernel choices can change reduction order, and the resulting floating-point differences have been measured alongside accuracy changes of up to nine percentage points.The cited measurement used one 7B model in bfloat16 across batch size, GPU count, and GPU generation.Source: §1 · pp. 1–2
  • GEMMDesc records parameters that determine reduction order, including the K-axis cut in split-K GEMM, while excluding speed-only knobs.The paper characterizes factors governing reduction order in reduction and GEMM kernels.Source: §1 · p. 2
  • Triton GEMMs bit-match cuBLAS on GB300, GB200, and H100 for fp16, bf16, and fp8 e4m3 inputs.The authors report 100% bit-matches across the tested cases.Source: §1 · p. 2
  • 19 of 27 kernels across GB300 and H100 run within 10% of free-order mode under balanced-tree enforcement and data-layout optimization.Five of ten kernels on H100 even outperform free-order mode.Source: §1 · pp. 2–3
  • The static checker decides bitwise equivalence for NVIDIA PTX and AMD GCN and is integrated into Triton’s autotuner.It recovers the exact equivalence partition for GEMM and Inductor fused kernels, and comes within a factor of 2.5 on reductions and normalizations.Source: §1 · p. 3

2 Background

The background explains how GPU execution and compiler mappings shape kernel behavior, and why autotuning can select configurations with different bits. It then contrasts compiler-correctness approaches with the paper’s focus on deciding bitwise equivalence between compiled GPU kernels.

2.1 GPU kernels and the programming model

  • A GPU kernel runs thread blocks across a grid, with cooperation and data placement spanning thread, warp, and block resources.The programming model distinguishes registers, shared memory, and global memory, with block-level synchronization required for shared-memory communication.Source: §2.1 · p. 3
  • Block-level languages express tiles while compilers assign elements to lanes and registers; vendor libraries instead choose the arithmetic behind a caller’s requested result.The tile-level IR exposes reduction order, and PTX or AMDGCN is the lowest textual assembly form described here.Source: §2.1 · p. 3
  • Autotuners measure configurations because performance depends on occupancy, tile fit, and register allocation, while configurations of one template may compute the same mathematics with different bits.The configuration space for GEMM can contain thousands of choices.Source: §2.1 · p. 3

2.2 Related work on compiler correctness

  • Verified compilers prove preservation of source semantics, whereas translation validation checks an individual compilation using encodings and solvers.The cited approaches include CompCert, Alive2, MLIR-TV, and Alive-FP.Source: §2.2 · pp. 3–4
  • Fuzzing finds incorrect compilations but proves nothing, and existing approaches do not decide whether two compiled GPU kernels return identical bits.Test-based methods cannot establish correctness merely by finding no failing test.Source: §2.2 · p. 4

3 Understand the bit behavior under GPU parallelism

A reduction’s floating-point dependence tree determines its bitwise result, and GPU thread, warp, and shared-memory choices can change that tree. The section characterizes how reductions and GEMMs form these trees, then describes the scope of its attention analysis.

3.1 Dependence Tree of reduction

  • A reduction’s dependence tree records the order of floating-point additions, and kernels over the same inputs return the same bits when their output trees are equivalent.Swapping an operation’s inputs does not change the tree’s equivalence relation.Source: §3.1 · p. 4
  • Balanced trees reduce the longest dependence chain to ⌈log2 n⌉, versus n−1 for a left fold, while GPU layouts and thread counts can alter the tree.The tree is formed at register, warp-exchange, and shared-memory levels.Source: §3.1 · p. 4
  • Software pipelining and warp specialization leave the reduction tree unchanged when pipelining preserves each accumulator’s order and specialization does not repartition the tile.The warp-specialization claim depends on the compiler leaving the consumer tile partition unchanged.Source: §3.1 · p. 5

3.2 GEMM reduction

  • GEMM reduction order depends on matrix-instruction rounding behavior and instruction_k, while BLOCK_M, BLOCK_N, and BLOCK_K are speed knobs that do not change the accumulator chain.Plain GEMM also depends on whether a short mainloop turn runs first or last; at fp8, mma.sync and tcgen05.mma are not bitwise equivalent.Source: §3.2.2 · pp. 5–6, §3.2.1 · p. 5
  • GEMV exposes its reduction order because lanes fold their totals through butterfly shuffles, with count-down and count-up offsets producing different sums.The paper records the fold as a two-part cut per round rather than one tree over lanes.Source: §3.2.2 · p. 6
  • Chained-chunk and split-K GEMMs cut the contracted axis at different levels, and GEMMDesc records the factors that determine their reduction order.For split-K, the descriptor records the cuts, part spans, and partial and merge data types.Source: §3.2.2 · pp. 5–6

3.3 Attention

  • For attention, the paper provides a theoretical description of reduction order but leaves implementation open.Source: §3.3 · p. 6

4 Black-box reconstruction

The black-box reconstruction recovers cuBLAS’s GEMM arithmetic as a shape-to-descriptor function, using profiling and numerical probes to identify reduction-order parameters. The reconstructed Triton kernels return cuBLAS’s bytes on tested full-sum shapes, while fused kernels also compare favorably in performance.

4 Black-box reconstruction

  • A GEMMDesc records bitwise-relevant reduction factors, and reconstruction recovers a shape-to-descriptor function for each fixed cuBLAS generation.A descriptor is one point in that function; an arch profile is a table of descriptors for reachable cost-model answers.Source: §4 · p. 6
  • Profiling first narrows a shape to a kernel family, then numerical probes with +L, −L, and r distinguish its grouping and cut boundaries.Fresh draws are used to separate remaining descriptor candidates against the library.Source: §4 · pp. 6–7
  • 100% of tested full-sum shapes on GB300, GB200, and H100 match cuBLAS’s bytes, excluding a located cuBLAS bug.The reconstructed profiles cover fp16, bf16, and fp8 e4m3 against cuBLAS 12 and 13.Source: §4 · p. 7
  • On 96 realistic shapes, bit-exact fused Triton kernels run at 95–168% of cuBLAS GEMM followed by a separate epilogue kernel.Source: §4 · p. 7

5 Compiler enforcement and layout optimization

The compiler pins reduction association to a balanced tree, then changes data layouts to reduce cross-warp overhead without changing that tree. This makes layout a performance choice for pinned reductions and recovers much of the constraint’s cost.

5.1 Backend lowering

  • The pinned lane-exchange path counts shuffle offsets up so neighboring lanes pair first, unlike the default count-down order.Both modes use the same shuffle-xor step sequence; AMD uses the corresponding count-up wave64 sequence and falls back to shared shuffles for wave32 or partial warps.Source: §5.1 · p. 7

5.2 The data-layout optimization

  • Moving a reduction axis inside one warp avoids the cross-warp barriers and shared-memory round trip, but ordinarily changes the tree unless it is pinned.Source: §5.2 · p. 8
  • Pinned lowering fixes a balanced reduction tree independently of element placement, making valid layouts a performance knob without changing the result bits.The pass places the reduction axis within a warp where possible and weighs the layout conversion against its cost.Source: §5 · p. 7, §5.2 · p. 8
  • On H100, the layout optimization brings six of ten kernels to or past unordered-mode performance, while on GB300 it recovers most of the cost of the constraint.Source: §5.2 · p. 8

6 Static equivalence checking and partition

The checker compares compiled kernels by statically reconstructing each output’s dependence tree, then uses those signatures to identify bitwise-equivalent configurations. Its Triton autotuner integration prunes the search to configurations equivalent to a reference.

6.1 The algorithm

  • The checker walks one symbolic thread through assembly, building each output’s dependence tree from loads, arithmetic, exchanges, and loop-carried reductions.It treats coordinate-based predicates as launch facts and data-dependent predicates as data, comparing the guarded entry configuration by configuration.Source: §6.1 · p. 8
  • Symbolic affine addresses unify leaves that read the same location, while opaque address expressions match only identical expressions, conservatively separating kernels.The address evaluator admits bitwise operations only when its rules establish exactness.Source: §6.1 · p. 8
  • Canonicalized signatures compare trees up to commutativity; balanced reductions also share a signature across layouts when reduction height and butterfly structure match.A left fold retains its physical structure, and therefore remains in its own class.Source: §6.1 · pp. 8–9

6.2 Implementation

  • The PTX checker serves as Triton’s pruning predicate, retaining configurations bitwise equivalent to a reference for benchmarking.The AMDGCN checker shares the tree representation but uses instruction-set-specific address and exchange recognition.Source: §6.2 · p. 9
  • For loop-carried reductions, the checker keys folds on chunk increments because re-chunking a scalar sum regroups it; it omits the chunk key for qualifying exact tensor-core accumulations.If a reduction cannot be reconstructed, launch geometry stays in the signature; nested folds include loop trip counts to avoid an inexact summary.Source: §6.2 · p. 9
  • Across 51,152 GB300 and 4,500 AMDGCN configurations, every certified equivalent pair returned identical bytes on every input draw.On most kernels, including TorchInductor-generated ones, the checker returned 1.0–2.5 times the number of equivalence classes that actually exist.Source: §6.2 · p. 9

7 Evaluation

The evaluation tests cuBLAS reconstruction and checker soundness, then measures the performance costs of bitwise constraints and the gains from layout optimization. Results cover multiple GPU generations and kernel benchmarks.

7.1 Benchmarks

  • The evaluation uses 390 fp16 model-derived GEMM shapes and random fuzzing shapes, plus microkernels, TorchInductor kernels, and an open-source kernel zoo.Measurements include GB300, GB200, and H100; AMDGCN checker tests use gfx942.Source: §7.1 · p. 9, §7 · p. 9

7.3 Evaluation on performance cost

  • For fused epilogues, bit-exact kernels reach 95–168% of the separate-kernel baseline, versus 85–125% for torch.compile, by removing one launch.The fused kernel takes 1.8 times as long as the baseline’s kernel half but wins the combined pair.Source: Figure 5 · p. 10, §7.3 · p. 11

7.2 Evaluation on bit-level correctness & static checker soundness

  • On GB300, reconstructed Triton GEMMs match cuBLAS byte for byte on every tested fp16 shape whose cuBLAS kernel computes the full sum, across eight algorithm families.The 60 exceptions reflect a cuBLAS defect at very deep contraction lengths; shapes outside that defect were also 100% bit-identical on the two older generations.Source: §7.2 · pp. 9–10, Table 2 · p. 10
  • Across 51,152 GB300 and 4,500 AMDGCN configurations, every pair certified equivalent matched bytes on every input draw; for targeted GEMMs, checker classes match the true partition exactly.The checker over-splits some other kernels, returning 6.7–23.0 times the true classes for four kernels outside the paper’s main focus.Source: §7.2 · pp. 10–11, Table 3 · p. 11

7.3 Evaluation on performance cost

  • Above 5 GFLOP, standalone Triton GEMMs reach 56–93% of cuBLAS performance when bit-exact and 69–91% with large-tensor acceleration.The free-order Triton GEMM reaches 61–88% of cuBLAS over the layer groups.Source: §7.3 · p. 11
  • 19 of 27 kernels finish within 10% of free-order performance; on H100, layout optimization brings six of ten kernels to or above it, while GB300 reaches 57–99% on all but one.On H100, the column-sum epilogue improves by as much as 42%.Source: Figure 6 · p. 11, §7.3 · p. 11

8 Discussion and future work

The discussion identifies limits in evaluation scope and verification level, and points to composing kernels and autotuning within equivalence classes as future directions.

8 Discussion and future work

  • Cross-machine robustness and the contribution to whole ML/RL training jobs were not measured because of resource limitations.The descriptor is designed to hold across machines and library versions, but that robustness remains unmeasured.Source: §8 · p. 11
  • Checker soundness is relative to PTX and AMDGCN assembly; checking emitted SASS remains an open direction, but that format is undocumented.Source: §8 · p. 11
  • Attention and GEMM fusions whose epilogues affect reduction order receive theoretical analysis in an appendix, not deep coverage.Source: §8 · p. 11
  • A future generator such as TorchInductor could pin the emitted reduction order and tune within one bitwise-equivalence class.Source: §8 · p. 12

9 Conclusion

The paper makes floating-point accumulation structure explicit: it can be described, reconstructed from a closed library, enforced by a compiler, and checked statically. The work reports cuBLAS byte matching across three GPU generations and static equivalence checkers spanning NVIDIA and AMD instruction sets.

9 Conclusion

  • A GEMM descriptor records the structure that determines bit-level semantics, which are shaped by compiler and hardware mapping rather than only by kernel authors.The paper says equal descriptors owe each other the same bits.Source: §9 · p. 12
  • Black-box reconstruction recovers closed-library arithmetic, and the paper reports matching cuBLAS bytes across three GPU generations.The recovered structure can also be required during compiler lowering, with performance costs described as recoverable at an acceptable rate.Source: §9 · p. 12
  • Static checkers decide bitwise equivalence between compiled kernels on NVIDIA and AMD instruction sets, allowing autotuning to search within one equivalence class.The paper reports applying the checkers to tens of thousands of autotuner configurations.Source: §9 · p. 12

A High level graph

The section outlines the work’s mechanisms and describes a layout optimization for pinned reductions, then explains how model configurations determine the GEMM shapes used in evaluation.

A High level graph

  • Figure 7 places four mechanisms along the path from a kernel template to returned bytes, beside the closed library whose arithmetic is reconstructed.The figure is described as a high-level overview of the work.Source: A High level graph · p. 12

B The layout optimization

  • The layout optimization considers pinned reductions, skipping cases that cross a thread block, are already warp-synchronous, or would spill after relayout.It also requires more lanes on the reduction axis and skips axes already spread enough; eligible operands are converted to an idealized layout and results converted back.Source: B The layout optimization · p. 12

C The models the shapes are read from

  • Model configuration keys determine GEMM dimensions: token count gives M, weight output width gives N, and weight input width gives K.LoRA uses adapter rank for K, reaching K=8, while its N remains the hidden or FFN width.Source: C The models the shapes are read from · p. 12

D The cuBLAS bug on split-K tail

The section identifies a cuBLASLt split-K path that can omit the K tail and gives a condition for predicting when the loss occurs. Its measurements show that the affected shapes vary by architecture and library version.

D The cuBLAS bug on split-K tail

  • Filling both inputs with ones makes each correct output equal K, so a shortfall directly counts missing K terms.The reported defect is in the nvjet split-K path cuBLASLt reaches at ALGO_ID 66 and reproduces on GB200, GB300, and H100 under cuBLAS 12.8.5 and 13.1.1.Source: D The cuBLAS bug on split-K tail · p. 12
  • The tail is lost exactly for ALGO_ID 66 when t≠0, q mod s=0, and s>t; the four-clause condition predicted every H100 loss with no false positives or negatives.Across 1,542,555 real GEMMs per H100 library, omitting the ALGO_ID clause misclassified 3,557 shapes under each library.Source: D The cuBLAS bug on split-K tail · p. 13
  • The three Table 5 shape examples distinguish which architecture-and-library pairings lose the tail; on GB300, the two library versions’ affected shapes were nearly disjoint, while on H100 they nearly coincided.The smallest losing K on H100 was 11,528 for both libraries.Source: Table 5 · p. 14, D The cuBLAS bug on split-K tail · p. 14
  • Across random-shape sweeps, 1,062 of 496,906 GB300 GEMMs under cuBLAS 13.1.1 differed from the full-K reference, and 1,059 met the tail-loss condition.The passage reports 2,714 and 2,995 flagged real GEMMs on H100 under versions 13.1.1 and 12.8.5, respectively.Source: D The cuBLAS bug on split-K tail · p. 14
  • The loss appears in ordinary as well as skinny shapes: 363 of 885 (M, N) pairs under 13.1.1 and 416 of 885 under 12.8.5 lose K somewhere.For 32 × 32, 64 × 64, and 128 × 128, K=11,528 loses 8 values.Source: D The cuBLAS bug on split-K tail · p. 14

E The large-tensor acceleration on GB300

The GB300 acceleration arm changes scheduling and tiling while retaining the descriptor’s arithmetic. Its validation reports zero byte differences, and its descriptor path adds host-time overhead that is visible for short GEMMs.

E The large-tensor acceleration on GB300

  • The GB300 branch adds persistent-grid scheduling, warp specialization, deeper pipelining, and plan-mode-specific tile rules while leaving other generations on the original kernel body.The branch also includes a GROUP_M tile swizzle, tensor-memory descriptors, 32-bit indexing where needed, and copy-free operand handling.Source: E The large-tensor acceleration on GB300 · p. 14
  • Grouped tl.dot operations preserve the reference’s increasing-K accumulator updates because a wide dot lowers to chained K=16 MMAs on one fp32 accumulator.For KPD 8, masked zeros stand in for each group’s short tail; the residue tile keeps one group per dot.Source: E The large-tensor acceleration on GB300 · p. 14
  • 154,904 byte comparisons returned zero differences across the acceleration arm.The evaluation included 7,868 plan-parameter combinations over nine modes at ten input draws each, using ordinary and wide exponent distributions.Source: E The large-tensor acceleration on GB300 · p. 15
  • The descriptor path costs about 100 μs of host time per call versus about 20 μs for the pointer launch, becoming visible on GEMMs under roughly 100 μs.The host cost is for building tensor descriptors; device time does not see it.Source: E The large-tensor acceleration on GB300 · p. 15

F The reduction order of attention

Attention combines two matrix multiplications with a scaled accumulation over keys; its reduction descriptor captures the choices that determine the resulting arithmetic. Block boundaries, masking, probability narrowing, and within-block layout all affect that arithmetic.

F The reduction order of attention

  • Attention applies two matrix multiplications around a scaled key-axis accumulation, which maintains running maximum, denominator, and output quantities across key blocks.The final answer is the accumulated output divided by the denominator.Source: F The reduction order of attention · p. 15
  • The key block width sets where output rescaling rounds occur, while BLOCK_K in the tensor-core mainloop changes speed without changing the instruction results’ order.A row of N keys incurs ⌈N/w⌉ block-boundary roundings.Source: F The reduction order of attention · p. 15
  • For causal masking, BLOCK_M determines where the straddling block begins and how many keys are live; unmasked configurations share boundaries regardless of BLOCK_M.Source: F The reduction order of attention · p. 15
  • Within a block, the row maximum selects an input without rounding, while the row-sum tree depends on thread count and lane layout.Source: F The reduction order of attention · p. 15
  • Probabilities are rounded from fp32 to fp16, bf16, or fp8 before the second matrix multiplication, making operand type part of the arithmetic order.Source: F The reduction order of attention · p. 15
  • The attention descriptor records key block width, masking and query block width when masked, row-sum tree, probability type, both matrix-multiplication descriptors, and final-division placement.Source: F The reduction order of attention · p. 15
Loading 2609.11356v1…