mailMailgithubGithublinkedinLinkedin

Published onJuly 8, 2026<br>AuthorsNameJubi TanejaTwitterLinkedInLinkedin

NameTom St JohnTwitterLinkedInLinkedin

NameNatalie SerrinoTwitterLinkedInLinkedin

Formally Verifying AI-Generated GPU Kernels<br>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.<br>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.<br>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.<br>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.<br>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.<br>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.<br>The Limits of Numerical Testing<br>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?<br>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.<br>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%.<br>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).<br>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.<br>def forward(self, Q, K, V):<br>scale = 1.0 / math.sqrt(Q.size(-1))<br>scores = torch.matmul(Q, K.transpose(-2, -1)) * scale<br>scores = torch.clamp(scores, -5.0, 5.0)<br>weights = F.softmax(scores, dim=-1)<br>out = torch.matmul(weights, V)<br>out = torch.clamp(out, -5.0, 5.0)<br>return out<br>The generated candidate fuses the manual attention into an efficient SDPA, and silently...