Papers With

← All papers
First page of Formalizing PTX in Lean to Prove GPU Kernel Correctness

Formalizing PTX in Lean to Prove GPU Kernel Correctness

Luk Koska, Tom Koska, Ben Koska

cs.AI Nov 3, 2026 · v1
verification tooling mathlib ai-agents
Read PDF arXiv abstract
TL;DR
Formalizes PTX semantics and memory model in Lean 4 with Mathlib and proves correctness of several GPU kernels, using LLM-assisted proofs.
Abstract
We present an approach to mathematically formalizing PTX, an approach that allows modeling and reasoning about patterns common in GPU kernels for Deep Learning, including acquire-release, tensor cores, proxies, async operations, multi-kernel pipelines and data-dependent control flow. We implement it in Lean and use it to prove the correctness of simple implementations of a WGMMA GEMM kernel, FlashAttention, and a 3-kernel MoE pipeline.
Problem

GPU kernels for deep learning, increasingly written by LLMs, need correctness guarantees without extensive human review. Existing verifiers cover only restricted kernel classes and exclude features such as asynchronous operations and data-dependent control flow.

Approach

The authors build a formal model of PTX in Lean 4, combining an interpreter covering 81 instruction forms with the axiomatic PTX memory model, including mixed-proxy semantics. The model supports acquire-release, tensor cores, proxies, async operations, and multi-kernel pipelines. Kernel correctness theorems state that every consistent, complete execution produces the specified output, and a constructed execution establishes non-vacuity. LLM coding agents assisted in writing the model and the proofs.

Results

Correctness is proved for 9 hand-written PTX entry points totaling 454 instructions, including elementwise, scalar, Ampere and Hopper (WGMMA) matmuls, FlashAttention, and a 3-kernel MoE pipeline. The proofs total 34,034 lines, average about 75 proof lines per instruction, and contain no sorry or custom axioms. They are checked by Lean 4 v4.35.0 with Mathlib.