用 Codex 自动研究实现内核 232 倍加速

Auto-research with codex: How I achieved a 232x Faster Kernel over baseline with Codex in GPU Mode's qr_v2 problem

08 Jul, 2026

Table of Contents

• Intro
• Contest in short
• Problem intro
• Why this problem is auto-research-able
• Learning Enough to Ask Better Questions
• (Optional) Math for QR decomposition: Householder reflections
• Make serial work small with the help of the blocked Householder algorithm
• Other challenges
• Codex-maxxing
• Kernel progress breakthroughs
• Breakthrough ideas
• Introducing idea diversity to escape the local maxima
• Implementation Hints
• What I could have done better
• Conclusion
• References
• Acknowledgements

Intro

Contest in short

GPU Mode, in collab with Core Automation, recently hosted an auto-research themed contest. The problem statement was to implement batched square compact-Householder QR factorization aka QR decomposition. I placed 12th out of 183 participants, ending up with a 232x speedup over the baseline solution. This post is about how I got there. I will go through my approach, learnings, and bottlenecks I ran into during the contest. It was my first serious attempt at auto-research. Some people will call this "loop engineering", and honestly that is fine too.

Note that you don't need to go through the mathematics or the problem itself in detail to follow most of this blog post. I have focused on my approach while keeping the math and the problem itself secondary as most people who will read this won't have participated in the contest.

You can check out the full contest page here: Problem Link and Leaderboard

This contest was part of GPU Mode's Linear Algebra Kernels in the Age of Research series.

Problem intro

We were given a batch of square FP32 CUDA matrices A with shape batch x n x n, and had to return the same compact Householder QR representation as torch.geqrf(A): an H matrix whose upper triangle is R and whose lower triangle stores Householder vectors, plus a tau vector of reflector coefficients. The checker rebuilt Q with torch.linalg.householder_product(H, tau), took R = triu(H), and verified:

A≈QR,Q⊤Q≈I,Q⊤A≈R

Among correct submissions, the leaderboard ranked runtime by geometric mean across shapes and conditioning cases. The important sizes were batched square matrices like 512 x 512, with larger 1024, 2048, and 4096 cases too. Low-bit FP16, FP8, or NVFP4 was allowed internally, but returned factors still had to satisfy FP32-style QR checks.

A tiny 3 x 3 example is:

A=[12−5146167−68−424−41]=[6/7−69/175−58/1753/7158/1756/175−2/76/35−33/35]⏟Q[1421−140175−700035]⏟R

Here Q is orthogonal, which means its columns are unit-length and perpendicular to each other, and R is upper triangular, which means everything below the diagonal is zero. The contest was not asking us to print dense Q and R directly; it asked for the compact Householder version that lets the checker reconstruct Q and read R from the upper triangle.

For the 3×3 example above, the very first reflector maps the first column (12, 6, −4) straight onto (−14, 0, 0) in one shot. The −14 becomes R11. How that works is in the math section.

Why this problem is auto-research-able

GPU Mode provides participants with the popcorn CLI making it agent-friendly. Agents can use this to test, benchmark, and submit to the leaderboard directly. The checker also provided shape-wise feedback along with the overall geometric mean timing.

Astute observers will notice this is an apt setup for writing a loop. Agents yearn for tight feedback loops. They allow them to hill-climb to their heart's content.

GPU Mode contests usually give you some way to iterate on kernels. Either you submit directly, or a sponsor like Modal chips in credits. Here the organizers basically allowed unlimited submissions as long as you spaced them out. If you didn't, the queues got long and everybody's runs timed out. At one point the workspace even ran out of Modal credits because everyone had been hammering submissions. It's a nice way to make learning accessible.

Over the course of 14 days, I made over 1500 submissions.

Learning Enough to Ask Better Questions

I have known the basics of GPU kernel optimization (mainly in Triton with some understanding of CUDA) for a year, but haven't worked in this domain professionally. What I am trying to tell you is that I was an underdog among the people around me on the leaderboard. The person just above me on the leaderboard (CUDA Colonel) is a principal engineer at NVIDIA.

Anyway aura farming aside, since I know the basics and had recently read about GatedDeltaNet, I was fresh on the general GPU kernel lingo.

The better you know something, the better you can prompt the LLMs, because you convert unknown unknowns into known unknowns.

At the same time, it's worth noting that this contest was doable without domain knowledge - like you probably won't make it to the top 10, but you can get a respectable speedup over baseline by just relying on your harness/agent loop or whatever.

My first steps in the contest were to learn what QR decomposition is and how it can be done. There are a bunch of ways to do it - like Gram-Schmidt and Householder reflections. The contest mandated Householder reflections. I went back and forth with Claude and watched a few YouTube videos to build intuition. After my discussions with Claude, it was clear that we needed to use the blocked Householder algorithm as the main architecture with the trailing WY-update. As it turns out, GPT-5.5 also had a good idea about this. QR decomposition is a fairly well known problem.

I found the concept interesting as matrix decompositions show up in several modern optimizer variants for LLM training, especially in methods that use matrix preconditioning, such as Shampoo-style optimizers and related approaches. Muon (used by Kimi) is another good example: instead of treating a weight update as one giant flattened vector, it keeps the matrix structure around and orthogonalizes the momentum update, usually through a few Newton-Schulz iterations that approximate the polar factor.

(Optional) Math for QR decomposition: Householder reflections

I recommend skimming through this section if you are curious about the math otherwise feel free to skip. The only thing to note is that there is a sequential dependency in Householder QR which makes it problematic to do GEMM. We use blocked Householder to make it more matrix-multiplication shaped.

The contract

Quickly reviewing the contract: input is a batch of square FP32 matrices A; output is the compact (H, tau) format that torch.geqrf returns. The upper triangle of H is R. Below the diagonal, H stores the Householder vectors, and tau stores one scalar per column. The checker rebuilds Q from (H, tau) and verifies A ≈ QR.

Mirrors

Forget matrices for a second. In a bathroom mirror, your reflection is exactly as far behind the glass as you are in front of it, straight through.

If x⟂ is the part of x sticking out perpendicular to the glass, reflection just subtracts that part twice:

xreflected=x−2x⟂

So a Householder reflection is about finding the perpendicular part and subtracting it twice.

Storing the mirror

A Householder vector is the mirror, stored compactly. In code, we don't carry around the whole mirror plane. We store one vector v sticking straight out of it. The mirror is everything perpendicular to v, and the reflection moves along v.

The perpendicular part is just the shadow of x along v, which is v⊤xv⊤v copies of v. Plug that into the subtraction above:

ℋx=x−τv(v⊤x),τ=2v⊤v

So tau is just 2v⊤v: the factor of 2 and the length of v bundled into one precomputed number. v picks the mirror, tau scales the update. (I'll write the mathematical reflector as ℋj and reserve H for the compact output matrix.)

Householder reflection in 2D tau = 0.00

x reflected x v x_parallel x_perp Mirror rule: keep x_parallel, flip x_perp. So reflected x = x - 2x_perp.