Formally Verifying AI-Generated GPU Kernels

Published onJuly 8, 2026

Authors

NameJubi Taneja

TwitterLinkedIn

NameTom St John

TwitterLinkedIn

NameNatalie Serrino

TwitterLinkedIn

Formally Verifying AI-Generated GPU Kernels

tl;dr: As AI agents get better at generating performant GPU kernels, building trust in their outputs becomes the bottleneck. At Gimlet, we've built an early research system that uses formal verification to complement traditional numeric tests in catching bugs in AI-generated (and human-written!) kernels.

Optimizing GPU kernels with AI agents has quickly gone from a cool concept to a viable technique in a seemingly short period of time. In a recent GPU Mode kernel competition, NVIDIA's Bryce Adelstein Lelbach commented that he didn't even need to look at the code his agent generated.

As generating optimized kernels becomes more tractable, building enough confidence in them to put them into production becomes the new bottleneck. Agents given overly narrow targets can come up with solutions that pass numerical validation tests, but don't actually share the same functionality needed in a real workload, which can span many different shapes and inputs. End-to-end workload issues arising from kernel bugs can be very difficult to diagnose, especially when they are subtle or only occur with very specific inputs.

Jason Wei coined the Verifier's Law, stating that "The ease of training AI to solve a task is proportional to how verifiable the task is. All tasks that are possible to solve and easy to verify will be solved by AI." This applies to any type of agent-driven solution, including kernel generation - underscoring the importance of good verification techniques.

At Gimlet Labs, optimizing kernel performance with AI agents is especially important for us, since we run inference workloads across heterogeneous hardware, and therefore need to support the same set of kernels across a variety of different platforms. You can read about some of our recent work here. As we've worked on this problem, we've seen firsthand how finite numerical tests aren't always sufficient in guaranteeing correctness.

In this post, we'll discuss our early work building a tensor algebra equivalence checker that uses formal verification to prove semantic equivalence between reference PyTorch models and their optimized implementations, including Triton kernels. Some of our work was presented recently at ARRAY 2026 at PLDI, and you can catch the full presentation here if you're interested. This is still an early research prototype, and the examples we walk through will be intentionally simple, but each is a real case the verifier caught that had passed numeric testing.

The Limits of Numerical Testing

Suppose you start with a PyTorch implementation of a model. You ask an AI agent to make it faster, and a few seconds later it produces an optimized Triton kernel. The kernel compiles successfully, benchmarks show a performance improvement, and all of your tests pass. Should you trust it?

The answer is "maybe". Today, traditional testing techniques for verifying kernel correctness are the standard. The engineer will define a set of tests that cover all of the cases that the kernel needs to support, and any performance or correctness check is done on these cases.

However, we've started to see that approach break down with the emergence of AI agents for kernel optimization. Kernels written by AI agents are specifically generated to maximize performance, which means that the optimization pressure is focused on whatever validation harness you use. There are many examples of reward hacking, cheating, and other correctness issues already well-documented by the community for the solutions generated by these systems. Researchers found that KernelBench overestimates correctness in solutions by 31%.

To understand where traditional testing helps and where it begins to struggle, let's take a single running example and break down where bugs arise and where naive test coverage may have failed. This is adapted for simplicity from a real workload that passed a suite of traditional numerical tests, and later we'll show how our research system was able to catch issues in the implementation. We show it in PyTorch for clarity, though our system verifies both PyTorch and Triton kernels (we'll look at a real generated Triton kernel later on).

The original reference implementation computes scaled dot-product attention (SDPA) manually. After computing the attention scores, it applies an intermediate clamp that restricts values to the range [-5, 5] before the softmax operation is executed.

def forward(self, Q, K, V): scale = 1.0 / math.sqrt(Q.size(-1)) scores = torch.matmul(Q, K.transpose(-2, -1)) * scale scores = torch.clamp(scores, -5.0, 5.0) weights = F.softmax(scores, dim=-1) out = torch.matmul(weights, V) out = torch.clamp(out, -5.0, 5.0) return out

The generated candidate fuses the manual attention into an efficient SDPA, and silently removes the intermediate clamp operation.

def forward(self, Q, K, V): scale = 1.0 / math.sqrt(Q.size(-1)) out = F.scaled_dot_product_attention(Q, K, V, scale=scale) out = torch.clamp(out, -5.0, 5.0) return out

With random testing, both the baseline and the optimized candidate return the same output. The optimization improved performance, but it actually did not preserve the contract of computation. While the original implementation guaranteed that all attention scores entering the softmax operation would remain within a bounded range, the optimized version removed that guarantee.

This difference was not immediately visible because typical inputs rarely exercised the clamp boundary. Most randomly generated test cases produced attention scores that already fell within the valid range, and hence treat clamp as an identity function. As a result, both implementations generated nearly identical outputs.

There is an obvious solution here: we need to add more configurations, shapes, and ranges to the inputs. However, this quickly breaks down. It's difficult to predict the entire range of inputs that a kernel may see in production. Even if you sweep different configurations, what if you don't sweep every combination of configurations? Similarly, perhaps randomized inputs can be produced over a larger range, but what if conditional branching or other logic in the kernel requires specific combinations of values in order to trigger the bug? This type of issue may not be discovered over hundreds of tests, but in production with large scale workloads, such edge cases will be triggered much more often!

The common issue here is that no matter how large the test suite becomes, testing can only examine a finite subset of an effectively infinite space of behaviors. Testing is sampling.

Formal Verification

This is precisely the problem that formal verification was designed to solve. Rather than evaluating a program on a finite set of examples, formal verification reasons about all possible inputs simultaneously. Instead of generating random tensors, we treat inputs symbolically. The verification question becomes:

Does there exist any input for which these two programs produce different outputs?

If the answer is yes, we obtain a counterexample. If the answer is no, we obtain a proof.

Underneath this process are SMT solvers, which are widely used across systems software. These solvers, such as Z3 from Microsoft Research, take a logical formula over symbolic values and either find concrete values satisfying it (SAT) or prove no such values exist (UNSAT). In formal verification, we ask it whether there exists a case where "baseline != candidate", such that SAT hands us a bug-triggering input, and UNSAT is proof of equivalence over the whole input space.

Verifying optimizations in this way has a long history in compilers. By formally verifying each individual transformation, we can build confidence in the entire set of transformations applied. LLVM does this in production with Alive/Alive2, and CompCert is an entire C compiler with fully verified correctness.

Compilers adopted this technique because optimizer bugs are the exact kind of issue that testing misses: subtle, input-specific, and very difficult to diagnose downstream. These issues are similar to the ones we see when optimizing AI kernels for better performance.

An AI agent generating candidate kernels has many degrees of freedom: it can restructure loops, fuse operations, change memory layout, reorder reductions. But that freedom isn't unlimited — the reference PyTorch model is the mathematical contract, defining what must be computed, not how. An optimized kernel is free to compute that specification any way it likes, fused, tiled, reordered, vectorized differently, as long as the function it computes matches the function the reference defines, for every input.

How the Verifier Works

So how does a verifier actually prove that two kernels compute the same thing?

A theorem prover cannot reason about Python, PyTorch, or Triton directly. We need to translate both the reference and the candidate kernels into mathematical formulas that the solver can understand. While proving equivalence of two completely arbitrary programs is not feasible (see the Halting problem), kernels for ML workloads are typically expressible as a DAG of arithmetic ops over concrete tensor shapes, which makes them more amenable to formal verification.

The overall process is shown below. We will use the missing clamp example from earlier as a running example to illustrate each stage of the pipeline. The goal is to determine whether the two models are mathematically equivalent for every possible input.

The verification pipeline. Both implementations are lowered to a common TensorAlgebra DAG, decomposed into primitive math, and scalarized to a formula for one symbolic output element. Z3 then compares the two formulas and returns a proof of equivalence, a concrete counterexample, or unknown.

Recovering the Computation

Before we can compare the two implementations, we first need to capture the computation they perform, using a common intermediate representation so the reference and the candidate can be compared. We call this the "TensorAlgebra DAG".

For the PyTorch reference, we recover the computation from the FX graph (a captured graph of the model's ops) produced by torch.compile. Triton kernels have to be handled differently, because Triton JIT kernels are opaque in the FX graph. You can't see their computational semantics at that level.

As a result, we need to compile the candidate a second time with the Inductor backend, which emits a .ttir file (Triton's MLIR-based IR) for each kernel. We then parse the TTIR and walk backward from tt.store instructions through SSA definitions, rebuilding the expression tree. This subtree gets stitched back into the main DAG by substituting back in the parsed expressions.

Both inputs are recovered in the same format, and we can see their ops in the TensorAlgebra DAG:

The recovered TensorAlgebra DAGs: (a) reference (left); (b) candidate (right).

This recovery process works because Triton emits an inspectable IR (TTIR). Low-level kernel implementations, such as CUDA, would require a different extraction pipeline to recover the computation before it can be verified for mathematical equivalence.

Decomposition, Scalarization, and Bounded Verification

Z3 can't reason about a high-level op like Clamp, and needs to be told that it's made up of a min and a max. Consequently, we need to decompose high-level tensor operators such as Clamp, Softmax, MatMul, and scaled dot product attention (SDPA) into their equivalent mathematical primitives. For our running example, these values end up mapping to the following for the reference and candidate kernels:

The same DAGs after expansion, with each high-level op replaced by primitive math: (a) reference (left); (b) candidate (right). Note the candidate has nothing corresponding

[truncated for AI cost control]