# Formalizing Nvidia PTX in Lean to Prove GPU Kernel Correctness > Source: > Published: 2026-11-03 20:04:08+00:00 **Read the paper [here](https://sf-tensor.com/papers/formalizing-ptx.pdf).** LLMs have gotten quite good at optimizing GPU kernels. In large-scale training runs, however, we still want to be confident that these kernels actually do the right thing. Testing kernels can only cover a finite number of cases and miss subtle bugs such as race conditions. To solve this problem, we've spent the past few months building systems to formally prove that a GPU kernel is correct. Our [previous work](https://sf-tensor.com/news/formal-correctness) in this area can already prove many kernels correct, but still has fundamental limitations in the types of kernels that it can verify. Since then, we've explored a new approach to proving kernels correct: Formalizing their semantics in Lean, then writing a formal and computer-checked proof that they're correct. For any given kernel, we have an LLM write a proof in Lean that the kernel computes the right thing. Lean then verifies the proof. This way, we don't need to trust the proof; Lean checks that it's correct for us. The Lean kernel is a small, well-audited program. If it accepts the proof, we can be very confident that the kernel is correct, within the scope of our formalization of PTX. We work at the PTX level because PTX is what all CUDA programs compile to, and because it has a precisely specified memory consistency model that we can formalize. We use Lean because its proof checker is small and trusted, and because LLMs have gotten quite good at writing Lean proofs. This new approach supports many things that our old work did not, including data-dependent control flow, TMA, async operations, fine-grained synchronization with acquire-release patterns, proxies & proxy-fences, producer/consumer patterns, megakernels and multi-kernel pipelines. We've used this to prove the correctness of simple implementations of Hopper tensor-core GEMM, FlashAttention with online softmax, a three-kernel Mixture-of-Experts pipeline, and a megakernel with data-dependent control flow. ## GPU threads don't share a simple globally consistent view of memory. When thread A writes a value and thread B reads the same address, whether B sees A's write depends on the ordering relations between their operations. The PTX memory model specifies these rules. The model is built on relations between memory events. An event is a single memory operation (a load, a store, a barrier arrivalt) performed by a specific thread at a specific point in its execution. We describe these relations below. **Program ordewr (po)** captures sequential ordering within one thread. If thread 0 first writes to x, then arrives at a barrier, those two events are ordered by program order. **Reads-from (rf)** records where a read gets its value. If read R returns the value that write W stored, then W is R's reads-from source. **Coherence order (co)** partially orders all writes to the same address. It represents which write "happened after" which at each memory location. **Synchronizes-with (sw)** captures cross-thread ordering from synchronization primitives. When thread 0 arrives at a `bar.sync` barrier (the PTX equivalent of `__syncthreads()`) and thread 1 returns from the same barrier, thread 0's arrival synchronizes-with thread 1's return. These combine into **causality**: the transitive closure of program order and synchronizes-with (plus some additional edges). Causality is what matters for reasoning about memory visibility. If event A is causally before event B, then A's effects are visible to B. The memory model's central constraint is the causality axiom: if write W is causally before read R on the same bytes, R cannot read from any write that is coherence-before W. Put simply, if a write is guaranteed to be visible to a read, the read must see that write or something newer. A concrete example. Two threads share a global variable x, initially 0: ``` Thread 0 Thread 1 st.global.u32 [x], 1 bar.sync 0 bar.sync 0 ld.global.u32 r, [x] ``` Thread 0 writes x=1, then arrives at the barrier. Thread 1 returns from the barrier, then reads x. Let's denote the write performed by Thread 0 with "W", and the read by Thread 1 as "R". Similarly, denote the barrier arrive in thread 0 "A0", and the barrier return in thread 1 "A1". Program order relates W to A0, and A1 to R. Additionally, synchronizes-with relates A0 to A1. Combined, these imply that causality relates W to R. In other words, W has "propagated to R", meaning that R can't read any values that are older than W. Therefore, it can't read the initial value of 0. Since W is the only other write, R must read 1. This is the mechanism that makes barriers useful: they establish a causal chain between memory operations in different threads. Without the barrier, thread 1 could legally read 0, because there would be no causal ordering between the write and the read. The full PTX model includes additional relations and axioms (coherence, single-copy atomicity, SC-per-location, and others). We formalize all of these except two: the no-thin-air axiom and the Fence-SC axiom. We omit these because they are difficult to model (more details on this in our paper). Omitting these axioms may keep us from proving certain kernels correct, but does not enable us to prove kernels "correct" that are actually wrong. ## ### The PTX memory model is an after-the-fact specification. Given a complete set of events from all threads, it says which reads-from and coherence-order configurations are possible (this is called "consistent"). It doesn't define step-by-step execution. This creates a circularity for verification. What events a thread produces depends on what its loads return, because of branches and data-dependent control flow. But what loads return depends on the full set of events across all threads, because that's what the memory model reasons about. ### We break this circular dependency with a so-called oracle. We give each thread an "oracle function". Whenever the thread needs external information (a load result, an atomic return value, a warp shuffle result), it queries this oracle, which returns a specific value (more on the exact value it returns later). With an oracle, a single thread executes deterministically. Each step reads the current instruction, possibly queries the oracle, updates registers, advances the program counter, and emits zero or more events (reads, writes, barrier arrivals). Branches check a predicate register and jump or fall through. Since register values can depend on oracle answers, the thread's control flow path (including loop iterations) can depend on what values it reads from memory. An execution of a kernel assigns an oracle to every launched thread. Each thread runs deterministically under its oracle, producing a list of events. We combine all threads' events into an execution candidate and ask: is this consistent by the PTX memory model? In other words, can we choose reads-from, coherence order, and synchronization matching such that (a) all the memory model axioms are satisfied, and (b) each oracle answer matches the value from the write that reads-from says the corresponding load reads from? If yes, the execution is allowed. Correctness now means: for every consistent, complete (all threads exited) execution, the final state of global memory satisfies the specification. We separately construct a valid execution for each kernel to prove non-vacuity. This rules out the possibility that the correctness statement holds trivially because no consistent execution exists at all, which could happen if a modeling mistake made the consistency conditions unsatisfiable. ### The model constructs an execution graph rather than advancing a mutable "global memory state", so we need to separately somehow define the state of global memory after the kernel finishes. To do this, we add a synthetic read event for every byte of global memory. We causally order these reads after all kernel events. The memory model axioms now constrain what values these synthetic reads must return. This gives a concrete representation of the kernel's output. If a write W stores x=1 and no other write to x exists, the synthetic read is causally after W. The causality axiom excludes the initial value as a possible source, as R would otherwise return a value older than W, despite being causally ordered after W. ### The base model handles loads, stores, arithmetic, and data-dependent control flow. Modern GPU kernels use additional synchronization and communication primitives. We formalized many of these as an extension to the base model described above. For example, we formalized WGMMA, proxy fences, warp shuffles, and pipelines of multiple kernels launched sequentially. For pipelines, API synchronization orders the completion of one kernel before the start of the next. We model global memory passed into two kernels by equating the successor kernel's initial global memory with the predecessor's output, and chain correctness results across launches. ### We model floating-point numbers as real numbers with exact arithmetic. Our proofs establish semantic correctness (the kernel computes the right mathematical function) but not numerical correctness (the kernel's rounding behavior is acceptable). We decided not to try proving that kernels produce bitwise "correct" floating-point outputs, as two correct implementations of the same operation frequently produce different floating-point results due to different reduction orders and accumulation precision. Furthermore, functional correctness (using reals) is binary and a precondition for numerical analysis to be meaningful. An important related question that we leave for future work is the following: How do you determine whether a kernel's numerics will cause inference accuracy to decrease or training to diverge? ## We formally verified the following kernels in Lean, with no `sorry` s (unproven assumptions) and no user-added axioms. For every admissible launch, we prove that all consistent, complete, UB-free executions produce the specified output, and separately construct such an execution to prove non-vacuity. **Elementwise kernel.** Each thread computes an affine map of the input for one element. **Scalar matmul.** Each thread computes one output element by looping over the reduction dimension. The proof uses an accumulation invariant over the inner loop. **Ampere tiled matmul.** A tiled matrix multiplication using `cp.async` for asynchronous loads from global to shared memory, CTA barriers for synchronization, and `mma.sync` for the tensor-core computation. **Hopper matmul.** Uses TMA with mbarrier for loads and asynchronous WGMMA for computation. This exercises the most complex parts of our model: asynchronous agents, proxy fencing, and mbarrier phase completion. **Online FlashAttention.** Computes `softmax(alpha * Q * K^T) * V` in exact real arithmetic with online softmax. The proof use the online softmax invariant and scalar arithmetict. **Mixture of Experts.** A pipeline of three sequential kernels: expert routing, per-expert computation, and expert combination. This uses the multi-kernel composition interface: the first kernel's verified output becomes the second kernel's input. **Simple megakernel.** Two single-thread CTAs share a global work counter. Each thread atomically increments the counter and uses the returned value to choose which job to execute. ## Our correctness results assume the kernel is free of ISA-level undefined behavior, and that `ptxas` and hardware behave correctly. We don't prove UB-freedom; it's an external assumption. Proving UB-freedom is harder than it might seem. Whether a thread reaches an undefined operation can depend on values it reads from memory, which depend on what other threads wrote. And those writes can be reordered by the compiler. Consider two threads with global u32 variables x and y, both initially 0, using `.relaxed.gpu` loads and stores: ``` T0 T1 r1 = load x r2 = load y if r1 == 1: UB store x, r2 store y, 1 ``` Here UB could be a misaligned memory access. If we stop T0 before the UB, we remove its store to y. T1 then reads y=0, stores x=0, and T0 reads x=0 and never reaches UB. No partial execution reaches UB by this analysis. But the compiler can move T0's constant store (y=1) before its load of x. Then: T0 writes y=1, T1 reads y=1 and writes x=1, T0 reads x=1 and hits UB. The memory model allows this reordering because the store's value doesn't depend on the load. PTX is a virtual ISA that `ptxas` compiles to machine code. A source-level criterion for UB-freedom can't rely on the memory model axioms alone; it needs to account for compiler transformations, and it's not clear exactly which transformations are permitted. Modeling SASS (the GPU's actual machine instruction set) instead of PTX would help. GPUs execute SASS instructions in order without speculation, removing the uncertainty about compiler reordering. We believe it might not even need oracles, instead being modelable with global, nondeterministic small-step semantics and reads delayed to their first use. This model would also be closer to the hardware's execution, making it more intuitive to reason about. Unfortunately, SASS is not fully documented by NVIDIA, making a complete model impractical for now. ## We've used this model to reason about most of the techniques used in modern high-performance GPU kernels: asynchronous tensor-core matmuls, multi-kernel pipelines, dynamic work distribution via atomics, and acquire-release synchronization. We're working on improving our coverage of PTX, with the goal of being able to prove full training steps correct. We also want to extend it to multi-GPU kernels as well as other hardware like AMD and Trainium. If this kind of work interests you, come work with us at SF Tensor.