Proving GPU Kernel Correctness in Lean—Within Explicit Limits

In their preprint, “Formalizing PTX in Lean to Prove GPU Kernel Correctness,” Luk Koska, Tom Koska and Ben Koska describe a way to reason about GPU programs with asynchronous operations, specialized matrix hardware and branches driven by data. The paper, which has not been peer reviewed, reports proofs for several simplified kernels. Its contribution is a formal framework and worked demonstrations, not a general-purpose automatic verifier for production GPU code.

The researchers model PTX, an intermediate instruction language used on the path to NVIDIA GPU machine code. This puts their reasoning below higher-level languages such as CUDA, while leaving the PTX compiler and hardware outside the proof. The aim is to describe the possible executions of a kernel and establish that every permitted completed execution produces the specified result.

Why execution order matters

A GPU kernel runs many threads whose memory operations interact. A thread writing a value does not, by itself, establish that another thread will read that value at the intended point. Barriers, memory-ordering rules and completion waits provide the guarantees needed to connect one thread’s work to another’s.

The model separates each thread’s local computation from information arriving from elsewhere. An “oracle” supplies candidate values for loads and other communication operations. These are mathematical inputs, not predictions or AI-generated answers. The model then checks whether the combined execution satisfies memory-consistency and synchronization rules. Values obtained from memory can affect branches and loop lengths, allowing data-dependent control flow.

Asynchronous transfers and matrix operations require additional rules. For example, a matrix operation’s output is constrained to the intended mathematical result only when the required synchronization is present. Without it, the result remains unconstrained, preventing a proof from simply assuming that unfinished work has completed.

What the authors report proving

The examples include an elementwise affine transformation, scalar and tiled matrix multiplication, a pipelined Hopper matrix-multiplication kernel, a scalar FlashAttention implementation, a three-kernel mixture-of-experts pipeline and a small dynamically dispatched kernel. These are deliberately simple, hand-written PTX programs, not optimized compiler outputs.

The dynamic-dispatch example illustrates the approach. Two threads each increment a shared counter atomically and use the returned value as a job ticket. The proof establishes that the tickets are zero and one in some order, so exactly one thread writes each designated output. It does not require deciding in advance which thread receives which job.

Across nine PTX entry points and 454 instructions, the authors report 34,034 lines of kernel proofs, supported by additional shared lemmas. Their model contains 2,629 lines of definitions. They say Lean checked the proofs without unfinished proof placeholders or custom axiom declarations. They also report reference-based hardware tests on H100 and B200 GPUs, except that the WGMMA kernel was tested only on H100. Those tests are separate from the mathematical guarantees.

The comparison with existing tools is a tradeoff, not a demonstrated blanket advantage. As described in the paper, Volta automatically checks equivalence for a defined class of GPU kernels but requires statically determined control flow and memory-access structure. This work targets broader execution patterns at the cost of manual, LLM-assisted Lean proofs. It does not establish that this approach is faster or easier overall.

Where the guarantees stop

The largest numerical qualification is that floating-point values are modeled as exact real numbers. A proof can therefore establish that a modeled attention kernel computes its mathematical formula, but cannot establish actual rounding behavior, numerical stability, quantization error or effects on training and inference accuracy. Bit-level manipulation of floating-point representations is also outside the model.

The results establish partial correctness: completed executions have the specified outputs. They do not prove that every run finishes. For each admissible launch, the authors additionally construct a completed model execution, checking that their universal theorem is not true merely because completion is impossible. That existence result is not a guarantee against hangs.

Applying the proofs to real PTX also assumes freedom from instruction-set-level undefined behavior and correct implementation by the compiler and hardware. The authors do not prove that undefined-behavior assumption. The specification, PTX parser, formal model and Lean checker remain trusted components; machine-checking a proof does not independently validate those foundations.

Coverage is limited to 81 instruction forms with fixed qualifier combinations. Concurrent kernels, host-side memory modification, multi-GPU communication and some newer hardware features remain outside the formalized scope. The model also omits two memory-model axioms, deliberately admitting more executions rather than using those restrictions to simplify proofs.

The paper demonstrates model-level verification across several important GPU programming patterns. Scaling that result to larger, optimized kernels—and connecting it more completely to real numerical and hardware behavior—remains unfinished work.