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

[ 2026-09-27 · updated 2026-10-23 ] // Claude Opus 5.5 · Lean 4 · PyTorch Inductor // 13 MIN READ

An earlier post proved when one torch.compile kernel's size arguments can be passed as 32-bit integers without changing what it does. Claude Opus 5.5 agents then generalized that proof into one Lean theorem, a small checker and a search that proposes certificates. Under the captured guards and modelled semantics, the checker certifies the narrowing for 25 kernels: 15 of 35 from development, and 10 of 24 in a pre-registered held-out set run after the method was frozen. Certification is separate from speed. Once both versions were tuned, two of nine kernels got 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 beyond the kernel it was written for.

The question matters because the existing alternative is a promise. Inductor's global assume_32bit_indexing setting passes these size arguments as 32-bit everywhere on the user's assurance that sizes stay small, which is unsafe when a product of sizes exceeds 2³¹ − 1. A certificate replaces that promise, one kernel at a time, with a proof under the guards the compiler recorded.

Triton is the language Inductor writes GPU kernels in. A kernel runs as many lanes, one per element, and a mask switches a load or store off on lanes where it should not happen.

The work lives in the same fork of VeriTile, a Lean framework for Triton-style kernels by Zenan Li and collaborators (upstream), but the new theorem doesn't build on VeriTile's definitions. It defines its own small integer, mask and payload semantics and 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.

Claude Opus 5.5 agents, working in Claude Code (the agent harness that gave them tools to edit code, run Lean and launch experiments), wrote the theorem, checker, search, frontend, benchmarks and manuscript. I set the tasks, approved each GPU run and read each stage's results before starting the next. Their work was checked by Lean, by comparisons with the compiler's output and by a separate reviewing agent (below).

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 records. 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 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

Before any checker run on new code, the agents wrote a list of 26 programs using operations the 35-program development corpus did not use, hashed it and froze the checker and theorem (pre-registration). The agents that built the checker also chose 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 could not certify the kernel, not that narrowing it would change its behaviour.

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 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 first deployment path, which matched kernel text alone, had narrowed them. A cached lookup and binding took a median of 0.6 ms per kernel, and checking a new instance in Lean a median of 4.8 s.

Does it make them faster?

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. The rule 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 #189940 kernel was 1.41–1.48× faster and repeat 1.49–1.51×, but that includes effects outside the kernel that weren't separated. 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.

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' work was checked

• Compared against the compiler's output. For all 25, a separate interpreter of the compiler's Triton intermediate representation, written for the audit and independent of the frontend, matched the model at about 320 sampled lanes per kernel, at shapes near the guard limit, with 0 mismatches. A control outside the premises showed that the comparison catches differences.
• Separate-agent review. A separate agent that had not written the code reviewed it read-only in three rounds, looking for ways to hide a division from the frontend and so from the proof. The first two rounds each got past the best-effort pre-check that screens new kernels, which was extended after each; the final round found no bypass in 16 targeted probes. The manuscript discloses four frontend gaps, including one unrelated to division: a kernel that rebinds a size-argument name would be modelled with the wrong value. None occurs in the 25 certified kernels, and deployment cannot reach them because the proof lookup hashes the exact kernel body, so the frontend's contract holds for these 25 by audit, not by construction (review record).
• A claim ledger. Each of the manuscript's 26 claims is listed with its kind of support (a per-lane Lean proof, an invocation condition, a correspondence between source and model that is trusted, audited or tested, a measurement, a derivation, or a citation) and its source. The review corrected several, including one that had called runtime guard installation audited without recorded evidence (LEDGER.md).

What you can inspect and rerun

The code and recorded results behind the counts are public in scasella/certified-int-narrowing. The supported entry point is one command:

python3 repro/run_repro.py --veritile-root /path/to/VeriTile

The VeriTile root is a lake build of veritile-narrow-cat at tag stage15-patch-timing (steps in the release notes). The command checks the hash of every pinned file and of the frozen checker, reruns the frontend pre-check's regression test, regenerates all 25 instances byte for byte from the recorded compiler scans, checks each in Lean with only the standard axioms (propext, Classical.choice, Quot.sound), confirms that the 10 negative tests fail, and recomputes the deployment-identity summary from saved kernel sources. It succeeds only if every count matches the historical record. A passing reference run from 27 September 2026 is in repro/reference_run/, and the full command passed again on a copy of the release files shortly before the repository went public.

The command checks proofs, not measurements. It runs neither PyTorch nor a GPU, so the timings, the deployment-binding test (which needs the pinned Docker image), the comparison with the compiler's intermediate representation, the provenance audit and the MLIR and Alive2 baselines are recorded results to inspect. It reproduces the frontend's recorded output rather than verifying the frontend.

Interpretation and limits

Under the captured guards and the modelled Triton semantics, one theorem and one small checker certified a change Inductor does not make by default, across held-out kernels the method had not been run on, and each 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.

• The most consequential assumptions sit at the edges of the proof. 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.
• Only Inductor's 1D pointwise kernels are covered. Reductions, matmuls and multi-dimensional kernels are out of class.
• The theorem is per lane. Reading it as whole-kernel equivalence also needs disjoint stores, which hold for all 25 because each lane stores at its own index, and it assumes 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

• Record that the certified guards are enforced at run time: run each of the 25 programs at a shape just past one of its guards and confirm that torch.compile recompiles instead of reusing the narrowed kernel, as the first post's boundary test did at N = 2³¹.
• Measure whether MLIR-style partial narrowing, division and remainder only, recovers most of the full specialization's speed on the same kernels and launch configurations.
• Let the autotuner choose between the two kernels instead of always narrowing.
• Run the frozen checker on a held-out set chosen by someone other than the agents that built it.
• Test whether a bounds-based rule inside Inductor, the direction the automated PR review suggested, narrows the same 25 kernels more fully than MLIR's per-variable ranges did.

Reproducibility