A Lean proof let agents cut a compiled PyTorch workload’s runtime by 26%

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

Claude Opus 5.5 agents found that a GPU kernel emitted by torch.compile(dynamic=True) ran faster with seven size arguments declared as 32-bit integers instead of 64-bit. A machine-checked Lean proof shows when that change is exact for this kernel, and a compiler rule applies it only then. On one NVIDIA L4, the narrowed kernel took 26% less time than the emitted kernel over a changing-shape workload, with one compiled graph and no recompilation. I sent the rule to PyTorch as an opt-in patch; a maintainer closed that PR without merging it and asked for discussion on the issue first.

Setup

After an earlier Dafny pilot that sped up deliberately naive routines, I wanted agents' proof-writing to make code from a mature compiler faster: to find work that compiled PyTorch code does needlessly, and to prove that removing it is exact.

I built on VeriTile, a project by Zenan Li and collaborators for verifying Triton kernels in the Lean proof assistant. Triton is the language Inductor writes its GPU kernels in. VeriTile provides a typed Triton-style language in Lean, its execution semantics, correctness and equivalence interfaces, and proof-checking infrastructure. I didn't build any of that. This work, in a fork, adds launch contracts, the transformation proofs and the narrowing rule.

The agents wrote the implementations, proofs, tests, measurement harnesses and the PyTorch patch. I set the tasks and spending limits, reviewed each stage before the next, and opened the PR.

First, find a baseline worth beating

The project started with simpler kernels, whose launch proofs show that checked configurations cover the whole output, every access is valid and unrelated memory is untouched.

Fusing addition and ReLU into one verified kernel ran about 1.7× faster on the device than the two-kernel pipeline at large sizes. But warmed-up torch.compile was about as fast (FUSION.md). That was a verified implementation, not an advantage over the compiler. I needed a workload where compiled code still did avoidable work that a proof could justify removing.

PyTorch issue #189940 supplied one: nested concatenation along a dynamic last dimension, with additions inside. The issue reported compiled code slower than eager (uncompiled) PyTorch on an H100. That didn't reproduce on the L4: dynamic compilation beat eager at all 7 shapes that completed. The mechanism did: compiled dynamic code was 1.52–2.38× slower than compiled static code across 14 shapes (report, §1).

def nested_cat_add(q1, k1, v1a, v1b, q2, k2, v2a, v2b):
    v1 = v1a + v1b
    v2 = v2a + v2b
    return torch.cat(
        [
            torch.cat([q1, k1, v1], dim=-1),
            torch.cat([q2, k2, v2], dim=-1),
        ],
        dim=-1,
    )

The inputs are 2-D bf16 tensors whose widths change at runtime. Inductor already fuses all of it into one kernel, so fewer launches wasn't the opportunity.

The first idea was the wrong one

The kernel maps each flat output index back to a row and column:

row = xindex // width
column = xindex % width

The size arguments arrive as int64, which makes all of that arithmetic 64-bit. That's deliberate. PyTorch's source warns that a product of size symbols can overflow even when each symbol fits.

A standalone kernel with 32-bit indexing and multiply-and-shift division, its constants computed on the host, was fast but didn't say which change mattered. So the agents rewrote Inductor's emitted kernel during code generation in three ways, keeping its masks, loads, stores and launch structure. Each variant kept the original body as a guarded fallback.

Above 1 is faster than the emitted kernel. Narrowing alone won. The division helper had a valid proof, and it made the kernel slower (controlled experiment).

The generated code showed why. With a 32-bit divisor, the compiler could reuse work that depends on the divisor instead of running the expensive 64-bit sequence per element. The recompiled emitted kernel has about 260 instructions per element, with calls to a 64-bit division slow path. The narrowed kernel has about 135 and no such calls. The custom helper added setup in every thread instead of computing its constants once on the host. So the formal work moved to the change that won.

The proof

The first narrowing variant checked its range condition inside the kernel and carried both bodies. The final one, N1, decides eligibility at compile time and declares the seven size arguments as int32 in the kernel signature. The body is unchanged, and there is no wide fallback.

That is not permission to narrow every size argument. For this kernel, call the six segment widths w₁ to w₆, the total width W = w₁ + … + w₆, the row count r and the element count N = r · W. With positive dimensions and N ≤ 2³¹ − 1, W and every segment width fit in signed 32 bits. A corollary, narrow_cat_args_fit, proves that part.

That alone isn't enough, because an intermediate expression can still overflow. So the agents extracted the kernel's actual index expressions from the emitted Triton into a typed expression language. Its semantics include signed 32- and 64-bit arithmetic, type promotion, wraparound at every operation and truncating division. The kernel runs as many lanes, one per flat output index, and a mask switches a load or store off on lanes where it should not happen. The main theorem, narrow_cat_equiv, says that on every lane of the grid the 64-bit and 32-bit kernels agree on the store mask, the store's address and value, and every memory read. The tensor values stay symbolic on both sides, so the proof makes no claim about bf16 arithmetic. It shows the same data reaches the same operations (NarrowCat.lean, standard Lean axioms only).

The theorem therefore covers the modeled narrowing transformation for this kernel class. It is not a verified compiler: it assumes the extractor captured the emitted kernel correctly and that Triton's compiler implements the modeled integer semantics (full trusted list under limits).

From theorem to compiler rule

The rule checks structure and facts (RULE.md). The kernel must match the proved class exactly: its arguments, casts, masks and memory operations. Then the checker reads Inductor's own symbolic state:

• each segment width has a lower bound of at least 1;
• the total width is defined as the sum of the six segments;
• the element count is a positive row count times the total width;
• N ≤ 2³¹ − 1 is guarded, and Dynamo, the front end of torch.compile, re-checks that guard on every call.

It never uses the sizes of the example inputs as evidence. Empty inputs fail the lower bounds and compile separately. Anything the recognizer doesn't model stays 64-bit, and a three-segment version of the program is rejected. The boundary tests include inputs where every size fits in 32 bits but the total-width relation doesn't hold, so a sum of three widths overflows and the two kernels disagree. At N = 2³¹, Dynamo recompiled to a kernel with 64-bit indexing, the rule declined, and the output matched eager bit for bit.

Confirming the speedup

N1 failed its first selection rule. It ran the whole sequence faster in both rounds, but in round 1 it fell to 0.951× of guarded N at one small shape, 3000 × (1000, 120, 136), below the 0.97× floor. In round 2 it was 1.09× at that shape. That verdict stands; I didn't lower the threshold.

A separate confirmation tested whether the small-shape miss was real. Its procedure and stopping point were fixed before the GPU call (harness):

• The emitted kernel (B0), guarded N and N1 were compiled in one process as distinct functions.
• 24 paired triples ran in 2 processes, with all six orders equally often. Each block warmed up for at least 2 s under load, ran the 14-shape sequence (20 calls per shape), then made 50 calls at each of 10 small shapes.
• N1 would be adopted only if its median benefit over guarded N was at least 1.03× with a 95% CI lower bound of at least 1.00, and every small shape's lower bound was at least 0.97.
• Statistics were paired ratios with percentile-bootstrap intervals, with no interim looks and no retries.

N1 passed, including non-regression at all ten small shapes. Every variant was bitwise equal to eager at every shape, and nothing compiled during the measured blocks (CONFIRM.md). Median sequence times were 600.7 ms for B0, 482.8 ms for guarded N and 443.9 ms for N1. The speedups are medians of paired ratios, not ratios of those medians.

1.349× is 26% less time. That's warm execution time for this changing-shape workload, not 26% faster model training. The intervals describe this one experiment, not other GPUs. The first experiment, under a different protocol, showed about 30% less time; I'm not pooling the two or picking the larger.

Why sustained load mattered

Isolated kernel timings didn't match repeated execution. During sustained passes, telemetry showed the L4 at its software power cap in every sample, with the SM clock at 52–57% of idle for the dynamic kernels. Kernels ran longer inside the sequence than alone (1.39× for the emitted kernel, 1.12× for N1), while gaps between kernels were at most 1.6% of a pass. So the confirmation warmed each block under load and measured that regime.

How the agents were checked

• Frozen statements. Theorem statements were frozen before proof search. An agent could change its proof, but not pass by weakening the theorem. Later additions, like the nonempty-launch corollary, were recorded separately.
• Separate checks for separate questions. Lean checked the formal statements. Fresh builds and axiom audits checked the artifacts, and Lean's official comparator replayed the proofs and checked that they prove the original, unmodified statements. GPU tests checked the actual executable on selected inputs.
• Tested links to the source. The emitted kernel feeds an extractor, whose output feeds both the Lean definition and the recognizer (reference checker). Binding tests check those links, and mutations must be rejected: an extra size use, a changed cast, an unsupported operation.
• Timing picked the proof target. A candidate needed both a measured benefit and its own correctness argument; the losing division helper's proof didn't count toward the narrowing.

The PyTorch patch

The agents ported the rule to PyTorch as an opt-in setting, config.triton.narrow_proven_size_args, on nightly 2.15.0.dev20260926 (git 6aa9e2fc, built from main ef166fb2). That build emits the same Triton source for this program, byte for byte. The patch applied, its four new GPU tests passed, and with the flag on the 14-shape sequence was bitwise equal to eager, with one graph, one kernel and all seven sizes as int32 (patch check).

Identical source doesn't guarantee identical compiled kernels or autotuning, so the patch was timed on its own revision. The flag was off in one callable and on in another, in the same patched process, under the same paired, pre-registered design, run once.

That's about 28% less time (report, §8). It comes from a different stack from the 26% result, so the two aren't combined. For four unrelated programs, first-call latency with the flag on was 2–4% higher, more than the check's own cost. With three processes per cell, that isn't separated from noise.

The patch recognized one generated kernel, identified by a SHA-256 digest of its canonical form; it was not a general pass that narrows arbitrary indexing. A harmless code-generation change could stop the match, leaving 64-bit sizes, but could not make the rule misfire, provided the trusted recognizer and extractor are correct. The report asked reviewers whether the rule should instead work on typed index expressions earlier in the compiler.

What happened to the PR

I opened pytorch/pytorch#198733 on 26 September 2026, marked it ready for review the same day and linked it from issue #189940. It had three commits: the patch tested and timed above, a lint pass with no change in behavior, and a change limiting the two recognition tests to CUDA. On the final commit, PyTorch's Inductor test jobs passed apart from two already marked unstable; several general build jobs stopped at a runner-setup step that PyTorch's CI bot judged an infrastructure failure, so the tests depending on them were skipped.

No human code review was submitted. On 29 September, Jason Ansel, a TorchInductor maintainer, asked PyTorch's automated Claude reviewer to review the changes. The bot review marked the PR as needing discussion, judged the measured speedup plausible and worth pursuing, and answered the report's question: the rule should work on Inductor's symbolic index expressions, not on generated source. Its main objection was that matching a hash of one kernel's source can't be maintained in PyTorch's tree: an unrelated code-generation change would silently switch the optimization off, and each new kernel class would need another recognizer. It recommended a generic decision instead: choose each size argument's type from value bounds, where Inductor already decides those types, reusing the 32-bit indexing guards Inductor already installs. That, it argued, would also keep what must be trusted inside Inductor's existing symbolic reasoning, instead of a hand-written source parser and an external Lean proof whose premises must be kept in step with the recognizer by hand. It also flagged broad exception handling, guards compared as strings, and test gaps: the two-shape test did not assert that the second shape reused the narrowed kernel without recompiling, and no test checked that the rule declines when the size guard or the width relation fails. The fork's boundary tests exercised N = 2³¹; the PR's tests did not.

On 30 September, Alban Desmaison, PyTorch's lead core maintainer, closed the PR without merging it. He wrote that PyTorch's usual process is to discuss a feature on its issue before anyone prepares a PR, to avoid work on “a feature we're not sure we want to add”, and to make sure the issue is marked “actionable” first. As of 23 October 2026, issue #189940 is open and not marked actionable.

Interpretation and limits

Agents improved code from a mature compiler, found the mechanism with a controlled experiment, and wrote a machine-checked argument for the change that was kept. The gain isn't a new division algorithm. The proof lets the compiler use cheaper arithmetic it already supports.

• Everything was measured on one L4, for one program. The issue's H100 is untested.
• The source extractor, the recognizer's parse, SymPy and Inductor's symbolic facts, and the Triton-to-machine-code toolchain are trusted, not proved. Masked-off lanes are assumed never to be dereferenced, as the emitted kernel already assumes.
• The patch and proof cover one kernel class; under this approach, each new class needed its own extraction and proof. The next stage, Agents turned a one-kernel Lean proof into a checker that certified 25 PyTorch kernels, replaced the per-kernel proof with one general theorem and a checker; it remains outside PyTorch.
• Whether PyTorch wants this optimization, and in what form, is unsettled (details above).
• Static compilation is still about 12% faster at warm steady state, at about 20 s more cold compilation over this shape sequence (29 s against 9 s). Narrowing improves that trade-off but doesn't make static compilation obsolete.

Next experiment

• Ask on issue #189940 whether maintainers want the 64-bit size-argument cost addressed before sending another PR.
• Prototype, inside Inductor, the bounds-based decision the bot review recommended, and test whether it narrows this kernel, keeps its speedup on the same L4, and agrees with the Lean checker on the next stage's development and held-out kernels.
• Measure the change on an H100, alone and together with #193614, which addresses a separate alignment cost in the same kernel, and try a second program.

Reproducibility