💻 computer science

Equivalence Checking of ML GPU Kernels

This paper introduces Volta, the first sound and complete equivalence checker for GPU kernels, which formally verifies the correctness of machine learning computations optimized by hand, compilers, or LLMs.

Original authors: Benjamin Driscoll, Kshitij Dubey, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin Driscoll, Kshitij Dubey, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast, invisible machinery of modern artificial intelligence, the most critical work happens not in the cloud, but on specialized computer chips called GPUs. These chips are designed to perform millions of tiny calculations simultaneously, a necessity for training the large language models that now write code, translate languages, and generate art. To make these models run fast enough to be useful, engineers must write highly specialized instructions, known as kernels, that tell the GPU exactly how to move data and perform math. Over the last few years, companies have begun using artificial intelligence itself to write these kernels, hoping to find faster ways to do the work than human engineers can. However, this speed comes with a risk: when an AI or a compiler rewrites a piece of code to be faster, it can accidentally introduce subtle errors. These errors might cause the computer to produce the wrong answer, or worse, to crash silently in ways that are nearly impossible to find through standard testing. The core challenge is that these chips execute thousands of threads of work at the same time, and if they do not coordinate perfectly, they can step on each other's toes, creating a race condition where the final result depends on the unpredictable order in which things happen.

A team of researchers has developed a new tool called Volta to solve this problem. Instead of guessing whether a new, faster version of a kernel is correct, Volta acts as a formal verifier that mathematically proves the two versions produce identical results. The researchers built a system that takes the low-level instructions of a reference kernel—the original, trusted version—and the optimized kernel—the new, faster version—and runs them through a symbolic engine. Rather than feeding the code specific numbers and seeing what comes out, the engine treats the inputs as abstract symbols. It traces every possible path the code could take, tracking how data moves through the thousands of parallel threads and how they synchronize with one another. If the code attempts to access memory in a way that could cause a conflict, or if the threads get stuck waiting for each other forever, the tool immediately flags the error. If the code runs cleanly, the tool translates the final output of both kernels into complex mathematical expressions and checks if those expressions are fundamentally the same, regardless of the specific numbers fed into them.

The researchers tested Volta on a wide variety of real-world machine learning tasks, including matrix multiplications, convolutions, and the attention mechanisms that power large language models. They found that the tool could successfully verify kernels that had been optimized by hand, by compilers, and even by large language models. In one instance, they examined a kernel generated by an AI that had been optimized through thirteen rounds of automated improvement. Volta confirmed that this AI-generated code was mathematically equivalent to the original human-written reference, proving that the aggressive optimizations had not broken the logic. The tool also proved its value by catching errors that other methods missed. For example, it detected data races in a popular, widely cited tutorial for GPU programming that had been used by thousands of developers for years. These errors were hidden because they only appeared under very specific timing conditions that standard testing rarely catches. The tool also identified a bug in an AI-generated kernel where the code attempted to read data from a memory location that did not exist; while the current hardware happened to ignore this mistake, the researchers showed that the code was fundamentally unsafe and could fail on future machines.

The strength of this approach lies in its ability to handle the unique complexity of GPU programming, where thousands of threads must coordinate their actions. Previous tools could check single-threaded programs or high-level mathematical operations, but they struggled to decompose the massive parallelism of a GPU into manageable pieces. Volta overcomes this by assuming that the kernels it analyzes follow a specific, structured pattern common in machine learning, where the number of threads and the size of the data are known in advance. Within this framework, the tool can prove with certainty that a race condition exists if the threads are not properly synchronized, and it can prove that two different versions of a program are equivalent if they produce the same symbolic result. The researchers demonstrated that their tool could verify these properties in a matter of seconds or minutes, even for kernels containing hundreds of thousands of instructions. They also proved that the mathematical logic behind their tool is sound, meaning that if the tool says two programs are equal, they are truly equal for all possible inputs.

This work represents a significant step toward making the development of artificial intelligence safer and more reliable. As companies increasingly rely on automated systems to generate the code that powers their models, the need for a rigorous way to check that code becomes critical. The researchers showed that it is possible to move beyond simple testing, which can only check a limited number of scenarios, to a method that provides a formal guarantee of correctness. By verifying the equivalence of optimized kernels, Volta gives developers the confidence to use faster, more aggressive optimizations without fear of introducing silent bugs. The tool is currently available for use, and the researchers have made their code and the proofs behind it open to the public, allowing others to build upon this foundation. While the tool does not yet cover every possible type of GPU code, it successfully handles the vast majority of the kernels that drive modern machine learning, offering a new standard for trust in the automated generation of high-performance computing code.

Technical Summary: Equivalence Checking of ML GPU Kernels

Problem Statement

The rapid advancement of deep learning and large language models (LLMs) has led to massive expenditures on executing GPU kernels, which are now prime targets for aggressive optimization. These optimizations are performed manually, by compilers, and increasingly by Large Language Models (LLMs). While empirical testing and program analysis are used to validate these kernels, they lack formal guarantees. This is particularly critical given the unique challenges of GPU programming:

Concurrency Errors: Fine-grained synchronization among thousands of threads introduces subtle data races and deadlocks that are difficult to detect via testing alone. These errors may only manifest in rare execution schedules or rely on implicit synchronization that is not guaranteed across hardware versions.
Floating-Point Semantics: Optimizations often reorder arithmetic operations, causing small numerical differences between reference and optimized implementations. While tests often tolerate these, they can mask genuine semantic bugs (e.g., omitted clipping operations).
Lack of Formal Tools: Existing equivalence checking techniques support single-threaded integer programs or high-level tensor operations but fail to handle the parallelism, synchronization, and non-deterministic scheduling inherent in GPU kernels. There is no prior equivalence checker for GPU programs.

Methodology

The authors present Volta, the first equivalence checker designed specifically for a practical class of ML GPU kernels. Volta operates as a black-box verifier, analyzing PTX assembly (the lowest documented level of the NVIDIA GPU stack) without requiring knowledge of the optimization process.

1. Target Class: Structured-CTAs

Volta targets Structured-CTAs (Cooperative Thread Arrays). This class assumes:

The number of threads, tensor sizes, and pointer targets are statically known.
Branch targets and memory access addresses are statically resolvable based on thread IDs and loop indices.
Control flow does not depend on runtime data values (e.g., efficient ML kernels do not make data-dependent runtime decisions).
These assumptions hold for many high-performance ML kernels (e.g., in JAX and XLA) and are checked by Volta; if violated, an exception is raised.

2. Symbolic Execution and Confluence

The core of Volta is a symbolic execution engine that treats tensor values as real numbers (modeling the convention used by optimizers like --ffast-math).

Symbolic State: The engine executes the reference and optimized kernels symbolically, deriving expressions for output tensor elements as functions of input tensors.
Scheduling: It employs a round-robin scheduling strategy. Threads execute until they hit a barrier (e.g., syncthreads, syncwarp).
Confluence: A key technical result is the proof of confluence: the symbolic evaluator guarantees that all possible execution schedules lead to the same symbolic expression (or report a race/deadlock). This property, formally verified in the Agda proof assistant, ensures that the checker's result is independent of the specific schedule chosen during analysis.

3. Race and Deadlock Detection

Volta extends symbolic execution to detect data races and deadlocks utilizing barrier synchronization:

Context Tracking: It maintains a context X tracking memory events (reads/writes) and synchronization status for each thread.
Race Detection: A race is detected if a thread reads/writes a memory location without having synchronized with the last writer/previous readers.
Deadlock Detection: A deadlock is reported if threads become blocked waiting for a barrier that cannot be satisfied.
Guarantees: The system is proven to be sound (if a race is reported, a concrete race exists) and satisfies a "true positives" property.

4. Decision Procedure

Once symbolic expressions are generated, Volta must decide if the reference and optimized outputs are mathematically equivalent.

Expression Form: The expressions involve multivariate polynomials, addition, multiplication, and exponentiation (e.g., $p(\bar{x})e^{h(\bar{x})}$ ).
Decidability: The authors prove the decidability of equalities of the form $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ . The proof relies on the fact that if the exponents $h_i$ are distinct polynomials, the coefficients $p_i$ must be zero for the sum to be identically zero. This avoids heavy number-theoretic machinery (like Lindemann–Weierstrass) and handles arbitrary real coefficients.
Canonicalization: The implementation uses a canonicalizer to normalize rational functions and exponential terms, handling cases like softmax rescaling.

Key Contributions

Linear Time Race/Deadlock Algorithm: A linear time algorithm for detecting races and deadlocks in structured-CTAs utilizing barrier synchronization.
Decidability Proof: A proof of decidability for equalities involving multivariate polynomials and exponentials over the reals, specifically for the form $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ .
Volta Implementation: The first equivalence checker for a practical class of GPU kernels. It verifies correctness for convolutions, matrix multiplications, reductions, and attention mechanisms.
Formal Verification: The key confluence lemma is formally verified in Agda, and the system provides soundness and completeness guarantees for the defined class of programs.

Results

Volta was evaluated on a diverse set of benchmarks, including:

Human-Optimized Kernels: Verified equivalence for matrix multiplication and reduction kernels from standard tutorials. It correctly rejected versions of reduction kernels (Red-5, Red-6, Red-7) that contained data races due to deprecated warp-synchronous execution patterns.
LLM-Generated Kernels: Verified kernels generated by LLMs (e.g., for 2D convolution and matrix multiplication). In one case, Volta detected out-of-bounds shared memory reads in an LLM-generated kernel that were not flagged by NVIDIA's Compute Sanitizer or the LLM's own testing.
Compiler-Generated Kernels: Verified equivalence for kernels generated by the TileLang compiler, including those using tensor cores.
Attention Mechanisms: Successfully verified the equivalence of standard softmax implementations against optimized online softmax formulations (used in FlashAttention), despite the complex symbolic expressions involving exponentials and max operations.
Performance: Volta verified complex attention kernels (e.g., FlashAttention-2) in under 3 minutes. The decision procedure significantly outperformed general-purpose SMT solvers (like Z3) on benchmarks containing exponentials, particularly when amortizing canonicalization costs across multiple output elements.

Significance and Claims

The paper claims to provide the first formal equivalence checker for GPU kernels. Its significance lies in:

Formal Guarantees: Offering stronger confidence than empirical testing, which is crucial for production deployment where bugs are costly.
Handling Concurrency: Addressing the unique challenges of GPU synchronization (barriers, warp-level sync) that previous equivalence checkers could not handle.
Robustness to Optimization: Verifying correctness even when optimizations reorder floating-point operations, by modeling values as reals and checking for mathematical identity rather than bit-wise equality.
Practical Applicability: Demonstrating that the "Structured-CTA" assumption covers a wide range of real-world ML workloads (convolutions, matmuls, attention) and that the approach scales to realistic problem sizes.

The authors note that while the work focuses on NVIDIA GPUs (via PTX), the techniques are applicable to other runtimes. They also acknowledge limitations: the approach does not support kernels with dynamic control flow (e.g., top-k/argsort) or asynchronous primitives (e.g., pipeline, tma), which are left for future work.