Auto-research with Codex: a 232× Faster GPU QR Kernel

The day the QR kernel stopped cooperating

The first time a “fast” QR implementation fails to be fast, it feels personal.

Our starting point was a baseline CUDA kernel that computed batched square compact-Householder QR for matrices shaped like batch x n x n (think 512×512, 1024×1024, up through 4096×4096). The checker didn’t just look at runtime. It reconstructed Q from a compact representation using torch.linalg.householder_product, extracted R as the upper triangle, and verified the relationships you’d expect from a QR decomposition: the reconstruction should be close to the original matrix, and Q should behave orthogonally.

And yet, the GPU could only deliver mediocre speed. Why? Because QR isn’t automatically GPU-friendly. Householder QR has an annoying property: each reflector depends on the updated matrix from the previous reflector. That dependency makes the computation feel serial, and serial work is exactly where GPUs underperform.

This post is the story of how an auto-research loop driven by Codex got us to a 232× speedup over baseline on the GPU kernel (the kind of acceleration that usually requires careful loop/blocking choices long before you ever start tuning micro-details). The focus here is the approach: how to turn “LLM writing code” into a reliable engineering loop.

What the problem actually asked for (compact-Householder QR)

A QR decomposition factors a matrix (A) into (A = QR), where:

• Q is orthogonal (its columns are unit-length and mutually perpendicular).
• R is upper triangular.

A common way to build QR on CPUs and GPUs is Householder reflections. A reflector is a transformation that can “zero out” parts of a column in a numerically stable way.

But the contest didn’t ask for dense Q and R explicitly. It asked for a compact Householder representation:

• You return an H matrix where the upper triangle contains R.
• The lower triangle stores Householder vectors (the data needed to reconstruct Q).
• You also return a tau vector containing reflector coefficients (scalars that control each reflection).

PyTorch’s torch.geqrf produces (a, tau) in a format described in LAPACK terms, and it pairs naturally with torch.linalg.householder_product to reconstruct Q from compact data. (docs.pytorch.org)

That compactness matters for performance too: you avoid forming full Q while still satisfying the checker.

Why this kind of contest is tailor-made for auto-research

The key isn’t that QR is special. The key is that the environment makes “iterate until it works” practical.

In GPU Mode’s auto-research contest, submissions are evaluated on a leaderboard with shape-wise feedback, and the submission harness supports rapid experimentation. There’s an agent-friendly CLI (Popcorn) that can submit, benchmark, and even run profiling workflows. (github.com)

In other words: you can run your loop while the GPU is doing real work, not while you’re waiting days between rebuilds.

That tight loop turns LLM coding from “write something and hope” into something closer to hill-climbing over design space.

The real performance villain: “serial QR” on a GPU

Here’s the part that surprised me the first time I tried to map Householder QR onto GPU thinking.

Unblocked Householder QR processes reflectors one at a time. Each new reflector is computed from a matrix that has already been modified by earlier reflectors. That creates a dependency chain.

In GPU terms, it means the computation spends too much time in matrix-vector-like behavior (think: applying one reflector to the trailing part), which doesn’t fully feed tensor cores. Meanwhile, the GPU’s fast hardware is left waiting.

So the central question becomes:

How do we keep the math stable while restructuring it so most of the work looks like big matrix multiplications?

This is a question that shows up repeatedly in GPU linear algebra: QR, LU, Cholesky, and friends all want to become BLAS level 3 operations (big dense matrix-matrix operations) as much as possible.

The fix: blocked Householder + compact WY updates

The breakthrough idea is classic in numerical linear algebra, but it becomes a weapon in kernel engineering.

Blocking (panels)

Instead of doing reflectors one column at a time, you process a panel of width (b) columns. In practice, (b) might be something like 32 or 64 depending on the target GPU and how your tiling is set up.

The panel is where the “serial-ish” work lives. Because the panel is narrow, the dependency chain is contained.

Compact WY representation

Now comes the clever part: you don’t apply each of the (b) reflectors to the trailing matrix one by one.

You convert the sequence of (b) Householder reflectors into a single structured update using a compact WY representation.

In plain language:

• Stack Householder vectors from the panel into a matrix (V).
• Build a small upper-triangular matrix (T) that encodes how the reflectors combine.
• Then the product of (b) reflectors can be expressed in a form that lets you update the trailing block with dense operations.

A common way to write it is:

[
\mathscr{H}_1 \mathscr{H}_2 \cdots \mathscr{H}_b \;=\; I - V T V^\top
]

and the trailing update becomes a sequence of matrix multiplications (often summarized as “three GEMM-shaped steps”). This is the reason blocked Householder QR maps so well to GPUs.

Conceptually, you compute intermediates like:

• (W = V^\top A_{trail})
• (Z = T^\top W)
• then update (A_{trail} \leftarrow A_{trail} - VZ)

These are matrix-matrix multiplications (often GEMM, general matrix multiply), which is where GPUs shine.

The blocked Householder algorithm with WY-style updates is a well-studied approach in both CPU and GPU contexts. (cs.utexas.edu)

Turning the math into a kernel (what actually broke)

Even once the algorithm is correct on paper, kernel engineering adds a second layer of pain.

1) Layout and tiling

A QR kernel lives or dies by memory behavior. On GPUs, each streaming multiprocessor (SM) has limited fast memory (like registers and shared memory). You can’t afford to fetch the same matrix elements repeatedly from global memory.

So the implementation needs:

• A clear data layout strategy for A, H, and V.
• Tiling so that sub-blocks of the trailing matrix stay hot while you apply the WY update.
• Careful boundary handling near the panel edges and at the “tail” end of the matrix.

2) Mixed precision and accuracy traps

The contest allowed lower-bit arithmetic internally (FP16/FP8/NVFP4), but the returned representation had to satisfy FP32-style checks.

That means:

• Accumulations often need to happen in higher precision (for example, FP32 accumulation even if inputs are low-bit).
• Errors can compound across many panel steps.

This is where you learn to treat “passes checker sometimes” as a near-miss, not a victory.

3) The panel-vs-trailing pipeline

The blocked approach naturally suggests a pipeline:

• Serial-ish work inside the panel: build reflectors and pack Householder vectors.
• Dense work on the trailing block: compute WY intermediates and apply the update.

If your kernel accidentally reintroduces fine-grained dependencies during the trailing update, performance collapses back toward baseline.

Where Codex fits: not magic, just better iteration

The main role of Codex wasn’t to “discover QR.” It was to speed up the loop where humans normally spend hours writing variants, then running, then debugging.

The agent loop idea

A practical auto-research loop looks like this:

1. Start from a correct reference kernel (even if it’s slow).
2. Generate one focused change: a tiling strategy, a different way to form intermediates, a different micro-kernel for a GEMM-like step, or a refined panel packing routine.
3. Submit to the checker quickly.
4. Keep changes that improve runtime and preserve correctness.

In a contest setting, you can do this repeatedly—hundreds or thousands of times—because the infrastructure supports frequent submissions and shape-wise evaluation. (sankalp.bearblog.dev)

Profiling as part of the loop

A surprisingly effective strategy is to make profiling “another signal” for the loop. Popcorn’s tooling includes hosted profiling support (including Nsight Compute profiling flows). ()

In other words: Codex writes code variants; profiling tells you where performance went missing.

“Idea diversity” beats local maxima

Kernel optimization has a nasty failure mode: you hit a version that seems faster, then every nearby tweak causes regressions.

A strong auto-research strategy intentionally explores multiple algorithmic directions—new blocking sizes, alternative packing layouts, different ways to map WY intermediates to GPU tiles—rather than refining one strategy forever.

That diversity is what helps escape local maxima.

A skeleton of the blocked WY QR flow

This isn’t meant as copy-paste code. It’s a map of what a fast GPU QR kernel ends up doing structurally:

for j in 0..n-1 step b:
 # Panel factorization (serial-ish inside width b)
 factorize panel A[j:, j:j+b]
 store Householder vectors into H
 compute tau for reflectors

 # Compress reflectors into WY form
 V = pack(panel Householder vectors)
 T = build_T_from_tau_and_V

 # Trailing update (dense, GEMM-shaped)
 A_trail = A[j:, j+b:]
 W = V^T * A_trail
 Z = T^T * W
 A_trail -= V * Z

 # Move to next panel

The performance win comes from making sure the heavy part (the trailing update) stays GEMM-friendly.

Implementation hints that matter more than they seem

A few recurring “gotchas” show up when you implement this for real:

• Panel width (b) isn’t just a tuning knob. It determines how much work is serial-ish and how large the GEMM-like updates get.
• Numerical stability must be designed in. If low-bit math is used internally, ensure accumulations are high enough and the returned compact representation matches the checker’s expectations.
• Don’t rebuild Q. Work in the compact representation (H and tau) because forming dense Q is a performance trap.
• Respect data dependencies. The dependency chain is unavoidable, but blocked QR contains it. The moment your trailing update depends on per-reflector results, you’ve lost the advantage.

Conclusion: speed came from restructuring, not from brute-force tuning

The headline 232× speedup didn’t come from one clever warp-level trick. It came from a deeper shift in how the algorithm is organized for the GPU.

Blocked Householder QR with a compact WY update turns the serial dependency chain into something you can quarantine inside a narrow panel. Then it lets the GPU spend most of its time on dense matrix-matrix work—the kind of computation tensor cores are built to chew through.

Codex and auto-research mostly helped the engineering part: exploring kernel variants quickly, validating them against a strict QR checker, and using profiling feedback to steer toward designs that were actually faster.

That’s a story worth remembering: when a numerical kernel feels “inherently serial,” the winning move is often to restructure the math so the GPU sees large, regular operations instead.