I used agents and Lean to cut a compiled PyTorch workload’s runtime by 26%

[ 2026-09-26 ] // Claude Opus 5.5 · Lean 4 · PyTorch Inductor // 12 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 Lean proof shows when that change is exact, and a compiler rule applies it only then. On one NVIDIA L4, the result took 26% less time over a changing-shape workload, with one compiled graph and no recompilation. The opt-in patch is now an open PyTorch PR.

Setup

I wanted agents' proof-writing to make real software faster, not to produce more proofs or beat a deliberately slow baseline.

I built on VeriTile, a Triton-verification project by Zenan Li and collaborators. It 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.

Proof work ran on my M4 Pro MacBook Pro with 24 GB of memory, with the official Linux proof checks in a local container. CUDA runs and timing ran remotely on an L4 through Modal. Every agent was Claude Opus 5.5, directed from Claude Code. The agents wrote the implementations, proofs, tests and measurement harnesses. I set the tasks and spending limits and reviewed each stage before the next.

First, find a baseline worth beating

The project started with simpler kernels. Their launch proofs covered more than the arithmetic in one GPU program: 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, and where a proof could justify removing it.

PyTorch issue #189940 supplied one: nested concatenation along a dynamic last dimension, with additions inside. The issue reported compiled code slower than eager on an H100. That didn't reproduce on the L4: dynamic compilation beat eager at all 7 shapes that completed. The mechanism did reproduce. 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. It 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, which is not the same as 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 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).

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 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. There, 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 stays in the record. 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. The two stay separate: I'm not pooling them or picking the larger one.

Why sustained load mattered

Isolated kernel timings didn't match repeated execution. Telemetry showed the L4 at its software power cap in every sample during sustained passes, 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 and 1.12× for N1. Gaps between kernels were at most 1.6% of a pass. So the confirmation warmed each block under load and measured the regime I meant to report, rather than correcting toward isolated-call numbers.

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, axiom audits and the official comparator checked the artifacts. 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. Binding tests check those links, and mutations must be rejected: an extra size use, a changed cast, an unsupported operation. The parser and the symbolic-fact extraction are trusted, not proved (reference checker).
• Timing picked the proof target. A candidate needed both a measured benefit and its own correctness argument. The proof of the losing division helper didn't count toward the narrowing.

The upstream 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 recognizes one generated kernel, identified by a digest of its canonical form. It is not a general pass that narrows arbitrary indexing. A harmless code-generation change can make it stop matching, and the kernel then keeps 64-bit sizes. It can't make the rule misfire, provided the trusted recognizer and extractor are correct. The report asks reviewers whether the rule should work on typed index expressions earlier in the compiler instead.

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 result is not general compiler correctness. Each new kernel class would currently need its own extraction and proof.
• 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. It doesn't make static compilation obsolete.
• The GPU tests and timing ran on the PR's first commit. A second commit changed only formatting, one set type and type annotations for PyTorch's linters, and was re-checked on the CPU code-generation path only.
• PyTorch-wide CI hasn't run. Four targeted tests don't replace it.

Next experiment

• Get maintainer review on where the rule should live, and on whether it can become an automatic compiler decision instead of a flag.
• Measure it together with #193614, which fixes a separate alignment cost in the same kernel.
• Try a second program and a second GPU, including an H100.

Reproducibility