{"slug": "formally-verifying-ai-generated-gpu-kernels", "title": "Formally Verifying AI-Generated GPU Kernels", "summary": "Gimlet Labs has built an early research system that uses formal verification to prove semantic equivalence between reference PyTorch models and AI-generated GPU kernels, catching bugs that pass traditional numeric tests. The system addresses the growing need for trust in AI-optimized kernels as agents become more adept at generating performant code.", "body_md": "# Formally Verifying AI-Generated GPU Kernels\n\n- Published on\n- Authors\n- Name\n- Jubi Taneja\n\n- Name\n- Tom St John\n\n- Name\n- Natalie Serrino\n\nFormally Verifying AI-Generated GPU Kernels\n\n*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.*\n\nOptimizing 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](https://x.com/blelbach/status/2072138321938042971).\n\nAs 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.\n\nJason Wei [coined the Verifier's Law](https://www.jasonwei.net/blog/asymmetry-of-verification-and-verifiers-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.\n\nAt 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](https://arxiv.org/abs/2606.02963). As we've worked on this problem, we've seen firsthand how finite numerical tests aren't always sufficient in guaranteeing correctness.\n\nIn 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](https://www.youtube.com/live/5XGgqkW53FE) 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.\n\nThe Limits of Numerical Testing\n\nSuppose 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?\n\nThe 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.\n\nHowever, 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](https://arxiv.org/pdf/2507.02825) that KernelBench overestimates correctness in solutions by 31%.\n\nTo 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).\n\nThe 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.\n\nThe generated candidate fuses the manual attention into an efficient SDPA, and silently removes the intermediate clamp operation.\n\nWith 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.\n\nThis 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.\n\nThere 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!\n\nThe 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.*\n\nFormal Verification\n\nThis 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:\n\nDoes there exist any input for which these two programs produce different outputs?\n\nIf the answer is yes, we obtain a counterexample. If the answer is no, we obtain a proof.\n\nUnderneath this process are SMT solvers, which are widely used across systems software. These solvers, such as [Z3 from Microsoft Research](https://github.com/Z3Prover/z3), 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.\n\nVerifying 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](https://github.com/AliveToolkit/alive2), and [CompCert](https://compcert.org) is an entire C compiler with fully verified correctness.\n\nCompilers 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.\n\nAn 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.\n\nHow the Verifier Works\n\nSo how does a verifier actually prove that two kernels compute the same thing?\n\nA 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](https://en.wikipedia.org/wiki/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.\n\nThe 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.\n\nThe 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.\n\nRecovering the Computation\n\nBefore 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\".\n\nFor the PyTorch reference, we recover the computation from the FX graph (a captured graph of the model's ops) produced by `torch.compile`\n\n. 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.\n\nAs a result, we need to compile the candidate a second time with the Inductor backend, which emits a `.ttir`\n\nfile ([Triton's MLIR-based IR](https://pytorch.org/blog/triton-kernel-compilation-stages/)) for each kernel. We then parse the TTIR and walk backward from `tt.store`\n\ninstructions through SSA definitions, rebuilding the expression tree. This subtree gets stitched back into the main DAG by substituting back in the parsed expressions.\n\nBoth inputs are recovered in the same format, and we can see their ops in the TensorAlgebra DAG:\n\nThe recovered TensorAlgebra DAGs: (a) reference (left); (b) candidate (right).\n\nThis 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.\n\nDecomposition, Scalarization, and Bounded Verification\n\nZ3 can't reason about a high-level op like `Clamp`\n\n, and needs to be told that it's made up of a `min`\n\nand a `max`\n\n. Consequently, we need to decompose high-level tensor operators such as `Clamp`\n\n, `Softmax`\n\n, `MatMul`\n\n, and scaled dot product attention (`SDPA`\n\n) into their equivalent mathematical primitives. For our running example, these values end up mapping to the following for the reference and candidate kernels:\n\nThe 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 to the reference's intermediate clamp. For readability, only the operators relevant to our running example are expanded.\n\nAlthough the computation is now expressed using primitive mathematical operations, it still operates on whole tensors. SMT solvers, however, reason about scalar values rather than multidimensional arrays. We therefore need to transform every tensor operation into the computation performed for each output element. This process is called Scalarization.\n\nMatrix multiplication, for example, becomes an explicit summation over the reduction dimension, while each tensor index is replaced with symbolic variables representing an arbitrary output element. Instead of reasoning about an entire output tensor, the verifier now reasons about a single arbitrary element, such as `out[i, j]`\n\n.\n\nIdeally, we could simply hand these formulas to the SMT solver. In practice, however, the expanded formulas can become extremely large. Convolutions, matrix multiplications, and attention all contain reductions that expand into hundreds or even thousands of product terms. As these symbolic expressions grow, the solver's search space grows rapidly, making equivalence queries increasingly expensive to solve.\n\nOur solution is to bound the reduction: instead of fully expanding every summation, we cap the number of terms that are unrolled symbolically while keeping the output indices symbolic. The solver still reasons about an arbitrary output element, but it simply reasons over a smaller symbolic reduction. Let's illustrate this further with a simple matrix-vector multiplication example.\n\nRather than expand the full 256-term reduction, the verifier unrolls it to a small bound (`k = 2`\n\n), which has the same structure, but is feasible for the solver to execute.\n\nThis transformation changes the actual outputs, but accurately computing the exact outputs isn't the goal here. Rather, the goal is to preserve enough of the reference and candidate's mathematical structure for the solver to determine whether they remain semantically equivalent.\n\nAlthough this approach doesn't give a full proof of equality for the entire 256-term dot product in the example, it catches the types of structural differences that commonly arise with subtle kernel bugs. We don't need every one of the 256 terms to be exhaustively checked in order to find them. This technique, [bounded translation validation](https://users.cs.utah.edu/~regehr/alive2-pldi21.pdf), is used in Alive2 to catch bugs in LLVM, by e.g. unrolling loops up to some bound when performing the verification.\n\nChecking Equivalence with Z3\n\nNext, we need to hand these formulas to Z3. To reason about every possible execution, the input tensors and other arguments must remain symbolic rather than being assigned fixed values. Internally, the verifier represents each input tensor as an [uninterpreted function](https://en.wikipedia.org/wiki/Uninterpreted_function). Conceptually, this is saying \"tensor A is a function that takes indices and returns a value\", with the added guarantee that the same indices always return the same value. Z3 is free to choose those values while searching for a counterexample, whereas the arithmetic structure of the formulas stays fixed.\n\nOperations (e.g. `Clamp`\n\n, `Softmax`\n\n, etc) are handled a bit differently. Unlike inputs, which intentionally remain symbolic, the verifier generally tries to encode each operation's semantics. If an operation is left uninterpreted, the verifier is no longer reasoning about its mathematical meaning, only that identical inputs produce identical outputs. That is an intentional approximation used for certain operations like `exp`\n\n, where precise SMT reasoning is expensive or incomplete. In those cases, optimizations that depend on mathematical identities (for example, the translation invariance of softmax, `softmax(x) = softmax(x - c)`\n\n) require explicit rewrite rules or axioms, otherwise the verifier may conservatively report a false positive.\n\nWith both the reference and the candidate encoded, we then ask Z3 whether there exists any index and any input where the baseline and the candidate disagree. If it can't find any such case (`UNSAT`\n\n), we have a proof of equivalence over the entire input space. If it can (`SAT`\n\n), that assignment is a concrete counterexample that shows the candidate is not mathematically equivalent. Z3 can also return `UNKNOWN`\n\nin cases where it ran out of time or memory, although we are working to address this category of limitations in our verifier.\n\nFor our running attention example with the missing clamp, the verifier quickly finds a counterexample, returning `SAT`\n\nin about 0.01 seconds.\n\nFor this scenario, Z3 constructs attention scores of -6 and -7, which are past the clamp boundary. While the reference saturates both to -5, producing two equal values for the softmax, the candidate produces different outputs and therefore proves inequivalence.\n\nScalar parameters are also left symbolic so that we can catch bugs in configuration values that are not covered in numeric tests. It's common for AI-generated kernels to be specialized for the particular scalar values they are tested on and benchmarked with. For example, the following kernel from a KernelBench Level 2 problem illustrates that category of bug and was caught by our verifier. The reference applies a clamp, scale, clamp, and unscale using a configurable `scaling_factor`\n\n:\n\nThe candidate kernel hard-coded the scaling factor to 2.0, even though the reference supports arbitrary scaling factors:\n\nBecause our verification pipeline leaves `scaling_factor`\n\nsymbolic, the verifier finds a setting where the candidate's folded clamp disagrees with the reference, and reports the kernels as inequivalent (`SAT`\n\n).\n\nEarly Results and Limitations\n\nThis is still early, ongoing work, but as part of our [ARRAY at PLDI talk](https://www.youtube.com/live/5XGgqkW53FE), we showed that out of the 26 Triton kernels we tested from KernelBench Level 1, our verifier formally proved 16 of them to be mathematically equivalent to their references, 2 passed numeric tests yet were proven to have mathematical inequivalences, and finally, 8 were `UNKNOWN`\n\n. We believe many of these UNKNOWN cases are addressable with engineering improvements, rather than fundamental blockers.\n\nOne major limitation is kernel language - as of today, our verifier supports PyTorch and Triton, but we want to extend it to low-level languages. [VOLTA](https://arxiv.org/abs/2511.12638) is a great example of related work that uses formal verification techniques applied to PTX to prove kernel equivalence.\n\nAnother limitation in the solution is the representation of numeric values as exact reals. In kernel optimization, differing use of floating point values can lead to accuracy bugs: for example, if the candidate does more floating point operations than the reference, there will be observable numeric issues.\n\nWhile certain numeric issues cannot be caught by the verifier today, there are some categories of bug that we've seen the system be helpful in finding. For example, consider the following HardSigmoid problem from KernelBench:\n\nThe Triton candidate (abbreviated for clarity) passed testing:\n\nIt made a simple change, converting a division by `6`\n\nto a multiplication by `0.16666666666666666`\n\n. This was an fp64 literal, but Triton compiled it down to fp32, leading to a numeric change. This change was caught by the verifier, which recovered the following operations for the candidate and the reference:\n\nZ3 was able to then prove that these two implementations diverge near the upper clamp boundary. At about `x ≈ 2.9999987`\n\n, the reference and candidate produced differing values of `0.9999999`\n\nand `1.0`\n\n, a difference of `2.6e-8`\n\n. This difference would be difficult to trigger with finite numerical tests, but Z3 was able to prove structural inequivalence. Whether or not this is an acceptable level of error would be up to the kernel author, but it's better to know up front what the issues are before integrating them into a production workload expecting total equivalence.\n\nWe have focused on the verifier finding bugs, but an \"over-eager\" verifier prone to producing false positives would not be very useful in practice. It's critical that the verifier can identify that two implementations, while differing in their approach, boil down to the same fundamental math, otherwise all we can do is fall back to numeric testing. To this end, we have used the verifier to confirm that structurally different implementations of the same math (e.g. PyTorch's fused `flex_attention`\n\nversus a manual implementation) are proven equivalent rather than flagged.\n\nFuture Work\n\nFormal verification techniques are complementary to, not a replacement for, traditional numeric tests. The next steps for this work are to extend it to 1) more complex kernels and 2) lower-level backends such as PTX. Kernels pass through multiple stages of lowering, and we need verifiers at each stage of that end-to-end pipeline.\n\nAt Gimlet Labs, our goal is to optimize inference across different hardware platforms. Supporting the same workload across N different platforms requires N sets of kernels, motivating our work in kernel generation across NVIDIA, AMD, Intel, and others.\n\nKernel generation is a poster child for the asymmetry of verification [noted by Jason Wei](https://www.jasonwei.net/blog/asymmetry-of-verification-and-verifiers-law). The tasks AI will do best in are the ones with reliable verification. Stronger verifiers, and different types of verifiers that can live in the generation loop, will help improve the robustness of the optimization process and reduce much of the current oversight burden.\n\nAs AI-driven performance optimization goes mainstream, the systems that succeed will not simply generate the fastest code. They will be the ones that also establish trust in what they generate.\n\nIf you're interested in this work, you can reach us at [hello@gimletlabs.ai](mailto:hello@gimletlabs.ai), check out the [ARRAY talk](https://www.youtube.com/live/5XGgqkW53FE), and [we are also hiring](https://jobs.ashbyhq.com/gimlet)!", "url": "https://wpnews.pro/news/formally-verifying-ai-generated-gpu-kernels", "canonical_source": "https://gimletlabs.ai/blog/formally-verifying-ai-generated-kernels", "published_at": "2026-07-09 16:50:18+00:00", "updated_at": "2026-07-09 17:07:46.281316+00:00", "lang": "en", "topics": ["ai-agents", "ai-safety", "ai-research", "ai-tools", "developer-tools"], "entities": ["Gimlet Labs", "PyTorch", "Triton", "NVIDIA", "Bryce Adelstein Lelbach", "Jason Wei"], "alternates": {"html": "https://wpnews.pro/news/formally-verifying-ai-generated-gpu-kernels", "markdown": "https://wpnews.pro/news/formally-verifying-ai-generated-gpu-kernels.md", "text": "https://wpnews.pro/news/formally-verifying-ai-generated-gpu-kernels.txt", "jsonld": "https://wpnews.pro/news/formally-verifying-ai-generated-gpu-kernels.jsonld"}}