Agents turned a one-kernel Lean proof into a checker that certified 25 PyTorch kernels

[ 2026-09-27 ] // Claude Opus 5.5 · Lean 4 · PyTorch Inductor // 11 MIN READ

My last post proved, for one torch.compile kernel, when 32-bit size arguments are exact. Claude Opus 5.5 agents then built a general version: one Lean theorem, a small checker, and a search that proposes certificates. Frozen before a pre-registered test, it certified 10 of 24 kernels it had never seen, and 25 in all. The narrowed kernels are safe to use. Most of them aren't faster.

From one proof to a checker

The first post ended with a limit: each new kernel class would need its own extraction and proof. That proof, NarrowCat.lean, covered one concatenation kernel under four hand-stated hypotheses. I wanted to know whether the idea generalized, or whether it only worked for the kernel it was written for.

The work lives in the same fork of VeriTile, a Triton-verification project by Zenan Li and collaborators, but it doesn't build on VeriTile's definitions. The new theorem defines its own small integer, mask and payload semantics, and it checks with import Mathlib alone. It reuses the first proof's modelling choices: Triton's type promotion and wraparound, symbolic tensor values, and the grid lemma.

The agents were Claude Opus 5.5, directed from Claude Code, with occasional second opinions from Claude Fable 5.1. They wrote the theorem, checker, search, frontend, benchmarks and manuscript. I set the tasks, approved each GPU run and reviewed each stage before the next. Proofs ran on my M4 Pro MacBook Pro, and timing on one NVIDIA L4 through Modal.

How the checker works

An Inductor kernel reaches the checker in two parts. A frontend turns the printed Triton kernel into a typed term. An exporter reads Inductor's symbolic shape state: each size argument and the element count as a polynomial in shape symbols, plus the guards Inductor installs, such as a product of sizes being at most 2³¹ − 1. Both of those steps are trusted, not proved.

A Python search then proposes a certificate. The certificate cannot assert a bound. It can only choose which guard bounds each subexpression, which quotient bound each division uses, and whether a load is checked only where its own mask holds. Lean recomputes every interval from those choices and emits side conditions: polynomials in the shape symbols that must be non-negative. One tactic, the same for every kernel, proves them.

For the held-out torch.tile kernel, the key step is 2ab − a(b − 1) = ab + a ≥ 0. It uses Inductor's guard on the product 2ab, which a per-symbol range would lose. During development, the search once proposed a side condition that held on sampled values but not on large ones. Lean rejected it, and that case is now a regression test.

Every accepted kernel instantiates the same theorem, proved once. On every lane of the launch grid:

• on active lanes, the 64-bit and 32-bit kernels read the same addresses and store the same value to the same address;
• on inactive lanes, neither one reads or stores, though their intermediates may differ;
• every division and remainder has a divisor of at least 1, in both widths;
• every size argument fits in 32 bits, so passing it as int32 is exact;
• no unspecified masked-off value reaches a store.

Tensor values stay symbolic, so this says the same data reaches the same operations. It says nothing about floating-point arithmetic.

Requiring agreement only where it can be observed

The obvious safety condition is that every integer intermediate fits in 32 bits on every lane. It's too strong. Five of the 25 certified kernels fail it, each with a concrete launched lane where an intermediate leaves int32 but nothing observable changes.

The sharpest case is the kernel Inductor emits for a rotary embedding. Under Inductor's own guards, there's a shape and a lane where one load's offset is exactly 2³¹:

The two kernels compute different offsets for a load that neither of them performs. The checker requires a load's offset to fit only where the load's own mask holds, and it derives those bounds from the mask itself, never from the certificate. Under that condition rotary is certified. Under the strict one it is rejected. Running both widths through the compiler's own intermediate representation at that lane reproduced the table.

That rule adds exactly one kernel. Treating inactive lanes separately matters for at least five.

The held-out test

The development corpus had 35 programs. Before any checker run on new code, the agents wrote a list of 26 programs using operations the development corpus did not use, hashed it and froze the checker (pre-registration). The same party wrote the list, so this guards against editing afterward, not against choosing easy programs.

The ten cover patterns absent from development: triangular masking, movedim, expand then reshape, tile, multi-way cat and stack, and permute then flatten. Neither search failure is a real overflow. One came from a guard Inductor had rewritten into a form the exporter did not pass through, and the other from a negative symbolic factor the checker cannot handle. A rejection means the checker couldn't show safety, not that the kernel is unsafe.

Binding the proof to the live compilation

A proof about a kernel's text isn't enough, because the same kernel body can be compiled under different guards. The first deployment path had exactly that gap: it matched kernel text alone.

The research integration hooks Inductor's kernel definition. It narrows a kernel only if the kernel's hash matches a certified one, and the live size polynomials and guards match the proof's premises up to renaming of shape symbols. On real Inductor compilations, 178 cases over 25 programs, every certified kernel bound. With the body unchanged, all 128 variants with a certified guard dropped or weakened were refused. The old text-only matcher had narrowed them. A cached lookup and binding took a median of 0.6 ms per kernel. Checking a new instance in Lean took a median of 4.8 s.

Does it make them faster?

The narrowed code is always different. In every certified kernel, all of the 64-bit division call sites (8 to 43 per kernel) disappear, and the machine code is 21–62% smaller. Whether that shows up as time depends on the launch configuration.

At one fixed configuration, 7 of the 25 were at least 1.10× faster, 16 were within 2% and 2 were slightly slower. Some of those gains came from a configuration that suited the wide kernel badly. So the agents tuned both versions of nine kernels over the same four configurations, chose on half the rounds and timed on the other half, with L2 flushed:

Above 1 means the narrowed option helped. The ranges span two builds, with and without pointer-alignment hints, and are not confidence intervals. head_merge's narrowed kernel was about 2% slower than its 64-bit one at their best configurations, so the rule kept the original. That selection rule was evaluated on the same run's other half. It is not an implemented tuner, and the integration as built narrows whenever a proof binds.

The #189940 figure is not a correction of the first post. That post timed a 14-shape sequence of whole calls under sustained load and pre-registered the comparison, which gave 1.349×. Table 03 is kernel time with L2 flushed and both variants tuned. Timed as whole programs on repeated inputs, the same two kernels were 1.41–1.48× and 1.49–1.52× faster, but that includes effects outside the kernel that weren't separated. Each number answers a different question. I’m keeping them separate rather than picking the largest.

The gains show up where size-dependent division or remainder sits in a hot kernel. A certified control with no division was 1.00× as a kernel and 0.99× as a program.

Compared with MLIR and Alive2

MLIR's integer-range narrowing pass, run unchanged on all 25 kernels, narrowed none of 233 64-bit operations. Given the checker's proven bounds as clamps, it narrowed all 65 divisions and remainders but left 161 operations wide, mostly address arithmetic, and made only 3 kernels fully 32-bit. Its ranges are per variable, so it can't use a guard on a product of sizes.

That leaves an open question: MLIR-style partial narrowing, which narrows only division and remainder, might recover most of the speed. I didn't measure it.

Alive2, a translation validator for LLVM, verified the division-free control in 0.1 s. It timed out at 60 s on all 25 certified kernels, each of which contains 64-bit division. That's one encoding and one budget, not a speed comparison between the tools. It does show why the proof avoids bit-level division reasoning.

How the agents were checked

• Frozen before testing. The checker and theorem were hash-frozen before the held-out run, and held-out failures stayed failures.
• Compared against the compiler's output. For all 25, an independent interpreter of the compiler's intermediate representation matched the model at about 320 sampled lanes per kernel per build, at shapes near the guard limit, with 0 mismatches. A control outside the premises showed that the comparison catches differences.
• Outside review. A fresh read-only reviewer, then two confirmation passes, looked for ways to sneak a division past the frontend. The first two rounds found bypasses: reassigning a variable after the check, and an expression form the extractor drops. A new pre-check that rejects them was added, and the final round found no bypass in 16 targeted probes. The same review found two latent frontend defects. Neither construct occurs in any certified kernel, so the frontend's contract holds for these 25 by audit, not by construction.
• A claim ledger. Each of 26 claims in the manuscript is tagged as proved, tested, measured, derived or assumed, with its source. The review corrected several, including one that had called runtime guard installation audited when there was no recorded evidence for it (LEDGER.md).
• One reproduction command. It verifies every pinned file, regenerates all 25 instances byte for byte, checks them in Lean with only the standard axioms, and confirms the 10 negative tests fail. I re-ran it on the public copy before publishing.

Interpretation and limits

The positive result is about safety and reach. One theorem and one small checker certified a change the compiler cannot make on its own, across kernels the method hadn't seen, and the proof is tied to the compilation it came from. The speed result is narrow. Certification makes the narrowed kernel an eligible alternative, and tuning should decide whether to use it.

• Only Inductor's 1D pointwise kernels are covered. Reductions, matmuls and multi-dimensional kernels are out of class.
• The frontend, the shape exporter, the Triton-to-machine-code toolchain and floating point are trusted. That Inductor's shape guards become runtime guards is assumed, with no recorded evidence.
• The theorem is per lane. Reading it as whole-kernel equivalence also assumes disjoint stores, no aliasing between inputs and outputs, and in-bounds accesses in the original kernel.
• One GPU, one compiler version (a PyTorch 2.15 nightly, Triton 3.8.0), and nine kernels with tuned timings. The other 16 have fixed-configuration data only.
• One whole-program output mismatch in an early run (attention, forced pointwise cat) didn't recur in eight controlled repeats. Its cause is unknown, and it is excluded from the validated results.

Next experiment

• Measure partial narrowing, division and remainder only, against the full specialization.
• Let the autotuner choose between the two kernels instead of always narrowing.

Reproducibility