Proving a CUDA kernel matches its PyTorch specification

We built deterministic translations from PyTorch and CUDA into Lean, allowing us to prove that an optimised GPU kernel computes the same function as its PyTorch specification, or identify a concrete input on which it does not. Unlike testing, the comparison includes real Float32 behaviour such as NaNs.

This is a case study of the same pattern behind the Logos Harness: start from an approved reference specification capturing intended behaviour, let an agent transform or optimise the implementation, and use deterministic semantic translation plus formal proof to decide whether the new program preserves that meaning. CUDA is an unusually demanding test because optimisation changes both execution structure and numerical behaviour.

The problem: tests cannot see the inputs nobody sent

Writing GPU kernels by hand is slow and requires unusual expertise, so a great deal of effort now goes into having models write them instead. The practical appeal is obvious: a practitioner describes the computation once, at a high level, in something like PyTorch, and gets a fast CUDA kernel back without writing any CUDA. However, there are no guarantees on whether the translation is sound or if the new, paralell program is thread or memory safe.

KernelBench (Ouyang et al., arXiv:2502.10517) is the standard benchmark for this task. It provides 250 PyTorch programs, each an nn.Module called Model, and asks for a ModelNew that computes the same thing with hand-written kernels. Its scoring metric, fast_p, counts a task as solved when the kernel is correct and at least p times faster than PyTorch. Correctness is decided by running the reference and the candidate on randomly generated input tensors and comparing the outputs.

However, this is not enough to guarantee the semantic faithfullness over all of Float32 and indeed the benchmark itself names the gap:

The failures that survive testing are not exotic. ProofWright (Chatterjee et al., arXiv:2511.12294) records a kernel that a model wrote for a simple activation function, which looks correct, compiles, and passes repeated unit tests, but violates a safety property whenever the input length is not a multiple of four. Two threads end up writing the same output element.

The usual defence against that class of error is NVIDIA's Compute Sanitizer, a correctness-checking suite shipped with the CUDA toolkit. It runs a kernel and watches the execution, reporting out-of-bounds and misaligned memory accesses, shared-memory data races, reads of uninitialised memory, and misuse of synchronisation primitives. It is the standard tool for exactly this job, and the authors report that it did not find the bug.

That is a safety concern rather than a semantic one and this is taken care of not in Lean but VerCors annotations. VerCors is a flexible tool used for reasoning about the data-race freedom and memory safety of paralellised code. It can be used not only for CUDA code but also programs written in C and Java. Code does not need to be executed in order to establish these proofs.

What we built

The approach is to stop comparing outputs and compare meanings. Both sides of the comparison are translated into Lean, where they become mathematical objects that can be related by a theorem.

• A PyTorch to Lean translator. The reference program is captured as a computation graph and imported into Lean as an ordinary value with a denotational semantics. The graph is generic over the scalar type: it names no number system of its own, so the same definition and the same theorem can be read over the real numbers or over IEEE 32-bit floats as Lean defines them.
• A CUDA to Lean translator. The kernel is compiled from the instantiated clang AST into an imperative Lean program: the kernel body becomes nested loops over the thread indices, with the launch configuration as ordinary parameters, shared arrays as array variables, and a barrier as the point where one loop ends and the next begins.
• A comparison step in Velvet. The translated kernel becomes the body of a Velvet method and the PyTorch-derived semantics becomes its postcondition.
• A VerCors gate. Semantics are not enough: in order for a translated CUDA kernel to be accepted, it must come with formal safety guarantees.

There are two key points about this setup that make it practical.

The first is that both translations are deterministic compilers. Neither is a language model being asked what a kernel means. A compiler can be read, reviewed and held still, and the accepted input fragment can be declared: kernel launches, atomics, warp shuffles, goto and unbounded while loops are rejected by name rather than approximated. It is impossible to translate somthing which the compiler does not support. Bugs, meanwhile, survive translation intact and that is the point. We are compiling exactly what the model wrote.

The second is that the two deterministic compilers reduce the amount of code a human has to trust. Whatever an agent writes, whether that is an invariant, a lemma or a counterexample, is a proposal, and Lean checks it.

Velvet: reasoning about imperative code in Lean

The comparison happens in Velvet (Gladshtein et al., CAV 2026), a Dafny-style verifier for imperative programs implemented as a Lean library. It supports require, ensures, invariant and decreasing annotations on a method; a weakest-precondition calculus turns the annotated program into Lean goals, which its automation attacks with grind and hands to an SMT solver like cvc5 or Z3.

Velvet is powerful because it allows for both automatic and interactive proof modes. Whilst a standalone verifier might just report a failure, Velvet is embedded in Lean, so the solver is one tactic among several. If it fails, the goal is still there, ready to be proved in interactive mode.

For a kernel the arrangement looks like this. The specification comes from the PyTorch graph, the body is compiler output, and the invariant is the only line an agent wrote:

-- from the PyTorch graph:
def spec (x : Scalar) : Scalar := max x 0

method relu_kernel (in_ : Array Scalar) (mut out : Array Scalar) (n : Nat)
  require n = in_.size                                 -- sizes from the launch
  ensures ∀ i < out.size, out[i]! = spec in_[i]!       -- the PyTorch meaning
  do
  let mut gid := 0
  while gid < n
    invariant ∀ k < gid, out[k]! = spec in_[k]!        -- fixed invariant
  do
    let v := in_[gid]!
    out := out.set! gid (if v > 0 then v else 0)
    gid := gid + 1

prove_correct relu_kernel by loom_solve                -- automatic proof discharge

The body and the ensures are never edited. If loom_solve fails, then either the agent may choose to continue the proof in interactive mode, or it can add more Lean theorems to the database of theories available to the SMT solver by including the attribute @[solverHint] to proven results.

Coverage: stating what a PyTorch program means

Before anything can be proved, the reference program has to exist in Lean. For that we build on TorchLean, a Lean 4 library of neural-network definitions with shape-indexed tensors, dependently typed operators, an operator-tagged graph, an evaluator and an independent shape checker. Its design is the one described above: operators are written against an abstract scalar.

Running our translator over all of KernelBench showed that the library as it stood could express very little of the benchmark. After extending it, all 200 Level 1 and Level 2 programs have a formal semantics in Lean: the benchmark's entire single-operator and operator-fusion tiers, covering matrix multiplies, convolutions, activations, normalisations, reductions, pooling, losses and the fused sequences built from them.

That is the step which makes the rest possible. For any of those 200 programs we can now write down, as a Lean statement, what the PyTorch reference means, which is what a candidate kernel has to be proved against. A problem counts here only when its graph passes structural validation and shape inference and every operator it uses has defined semantics. KernelBench is a proxy rather than a target. We are not trying to win KernelBench; what we want is to be able to state, in Lean, what a wide range of real PyTorch programs mean, and KernelBench is a convenient stand-in. Levels 3 and 4 of KernelBench are particularly challenging, representing entire architectures.

Example one: ReLU, and why the number system decides the verdict

The simplest problem in the benchmark is return torch.relu(x), and a model wrote the obvious kernel for it:

int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
  float v = in[idx];
  out[idx] = v > 0.0f ? v : 0.0f;
}

It compiles, and it passes the benchmark's five random inputs. The specification it has to match is infered from the provided PyTorch code.

def spec (x : Scalar) : Scalar := max x 0

The proof attempt does not go through, and the shape of the failure is interesting. Everything reduces to one question about a single float: is if 0 < x then x else 0 the same as max x 0? Splitting both comparisons gives four cases. Two close immediately but two remain open, and the reason is that a NaN compares false against everything. The case "not greater than zero and not less than or equal to zero" is inhabited, by exactly the NaNs.

Over the real numbers both of those cases are empty. A real number is negative, zero or positive, so there is nothing left to prove. If we were reasoning over the reals, then we would be able to prove this kernel 'correct'. Over Float32 the case survives and Velvet stops there.

Stopping is not yet a verdict. When the automation cannot close a goal, that on its own says nothing about the kernel. It may be wrong, or the argument may simply need a fact about floating-point arithmetic that nobody has supplied yet. The two are told apart by trying to settle the open case directly, and here the open case says exactly what to look for: a value for which both comparisons come out false. The agent supplies the following code, providing a witness:

def nan : Float32 := Float32.ofBits 0x7FC00000

theorem kernel_ne_spec_at_nan :
    (if 0 < nan then nan else 0) != spec nan := by decide

The kernel sends NaN to +0.0 whilst torch.relu persists it. The verdict is that the two are not equivalent, and the input that separates them is written down.

A second model, asked independently, made the same mistake by a different route, using fmaxf instead of the conditional.

It's worth catching NaN's

One objection is that nobody sends a NaN to an activation function, so a kernel that mishandles them is not a big deal. However, it becomes important when combining kernels.

If a NaN arises during a computation, it should be carried forward to future layers in order that it is loud and catchable. A kernel which incorrectly maps a NaN to +0.0 silences this which can hide numerical errors when the integrated into a full network.

Example two: LeakyReLU, and the lemma pool

The next problem is LeakyReLU with a slope of 0.01. After the agent has translated both the kernel and the PyTorch semantics into Lean, the comparison reduces to a single question:

(if v >= 0 then v else slope * v)  =?=  (if v > 0 then v else v * slope)

Splitting both guards gives four cases. One closes. The other three each name a fact about Float32 that aren't available at the time.

decide cannot help: it evaluates closed expressions, and v is a variable. Congruence cannot help either, because the two sides are different terms. This is the point at which an automatic verifier hands you nothing. Velvet, being embedded in Lean, leaves the goal open in front of you, so the missing facts can simply be proved.

Because Lean 4.33 defines Float32 down to the bits, we can provide the required theorems. The agent proves them, and each one is then tagged so that the automation may use it from that point on:

@[grind] theorem float32_mul_comm (a b : Float32) : a * b = b * a
@[grind] theorem zero_lt_imp_not_le_zero (x : Float32) : 0 < x → ¬ (x ≤ 0)
@[grind] theorem zero_classify (v : Float32) (s) (h : v.toModel.unpack = .zero s) :
    v = Float32.ofBits 0x00000000 ∨ v = Float32.ofBits 0x80000000

With those three, LeakyReLU closes: equivalent for every input, NaN included. As the pool of lemmas available to the automatic solver grows, the automation becomes more and more powerful, motivating the need for a floating point library in Lean.

Floats are no longer opaque

Until Lean 4.33, Float and every operation on it were opaque to the kernel: there was no definition to unfold and therefore nothing to prove. In 4.33 Float became a wrapper over Float.Model, a subtype of UInt64, with a real definition (Float32 over Float32.Model, a subtype of UInt32). Operations unpack a bit pattern into a signed zero, a finite significand and exponent, an infinity or a NaN; compute on exact integers; round once; repack. Closed expressions can therefore be evaluated.

However, the definitions do not define everything. Basic operations like: + - * /, sqrt, abs, neg, and propositions, isNaN, isInf, isFinite, are given meaning. Not defined are more complicated functions like: sin, cos, tan, exp, log, pow, which remain calls into the platform's C library. The reason is as follows:

So a kernel calling expf or tanhf cannot be given a bit-exact verdict, and that is not a gap in our tooling: the answer really does depend on the machine. What we can do is prove everything around the call and name the operation left undefined, so that it gets tested on the device rather than assumed.

A verdict today describes IEEE arithmetic on its own, and real hardware does more: nvcc fuses a*b+c into a single multiply-add by default, and flush-to-zero changes behaviour near zero. Because Lean's model computes exact integer significands before rounding, these are definable rather than assumed. A fused multiply-add is the exact a * b + c with one rounding step. The same kernel could then carry a verdict per device, one for an H100 and another for a part that rounds differently.

How this differs from existing work

ProofWright (Chatterjee et al.) is the closest published system and the one this work is built against. It has two independent halves. An agent writes VerCors annotations for memory and thread safety, establishing those properties for 74 of 100 level-1 kernels, at an average overhead of about three minutes per kernel. Separately, a static analyser captures the PyTorch program into MLRocq, a Rocq library implementing 99 of the operators appearing in level 1, roughly 92% coverage. Semantic equivalence end to end lands on 14 of 100, all element-wise kernels with one output element per thread.

The safety half is real work on a property we do not address at all. On the semantic half, three differences account for most of the distance.

The scalars. MLRocq's tensors use the integers, and its ReLU is (x + |x|) / 2, which is ReLU over the integers because 2x / 2 = x there. But that's not what a float kernel computes, and more complex functions like softmax, GELU and normalisation have no faithful representation. In order to provide meaningful equivalence claims, the semantic representation of the kernels must be as close as possible to the actual implementation.

The translation. ProofWright's PyTorch front end is deterministic and its authors say so. The CUDA side is not: an agent synthesises the Rocq reading of the kernel, and the lowering to VerCors annotations is generated by a model and checked by hand. The paper is direct about it, and about where it wants to end up:

That procedural compiler is what we built, and it is why the trusted surface here is two pieces of reviewable code rather than a review of each generated artefact. Hand-checking scales with the number of kernels; a compiler does not.

Where this goes

There is a great deal more that could be built on in this direction.

• Widening the specification languages Our methodology is could in principlebe extended to other languages for defining the high level specification. All we extract from PyTorch is the computation graph which is available in other formats. For example it would be interesting to extend our appraoch to a JAX jaxpr, an ONNX model or a Triton program.
• Float semantics for a particular device. As above: fused multiply-add and flush-to-zero are definable in Lean's model, so the same kernel can carry a verdict per target rather than one verdict for idealised IEEE arithmetic.
• A floating-point library for Lean. We proved what we needed one stuck goal at a time. That does not scale, and now that floats are no longer opaque a shared library is the obvious thing to build. Work has started elsewhere, relating a real-indexed float type to core Float32.

The larger goal goes beyond CUDA verification. Logos is building toward making large software transformations safe: an agent should be able to migrate, optimise or retarget a program while a deterministic semantic layer turns “preserve the approved behaviour” into a machine-checkable obligation. GPU kernels make that problem unusually unforgiving. The fact that a five-line ReLU optimisation already turns on NaN semantics is exactly why the numerical model has to be part of the proof.