Auto-research with Codex: Achieving a 232x Faster QR Decomposition Kernel

Executive Summary

By implementing an automated research loop using OpenAI's Codex, a developer achieved a 232x speedup over the torch.geqrf baseline for batched square compact-Householder QR factorization. The success was driven by transitioning from generic library calls to a blocked Householder algorithm, utilizing a "beam of candidates" to avoid local maxima, and leveraging tight feedback loops provided by the GPU Mode popcorn CLI and Modal profiling.

The Optimization Challenge: QR Decomposition

The goal was to implement batched square compact-Householder QR factorization for FP32 CUDA matrices. The output required a compact representation consisting of an H matrix (where the upper triangle is R and the lower triangle stores Householder vectors) and a tau vector of reflector coefficients.

The Serial Bottleneck

Standard Householder QR is inherently serial: reflector $j+1$ depends on the matrix produced by reflector $j$. This sequential dependency prevents effective use of Tensor Cores, as the work remains matrix-vector shaped rather than matrix-matrix shaped, leaving the GPU's most powerful compute units idle.

The Blocked Householder Solution

To overcome the serial bottleneck, the blocked Householder algorithm was employed. This approach confines serial work to a narrow panel of $b$ columns. The $b$ reflectors are then compressed into a single rank-$b$ update (the WY representation), allowing the trailing block of the matrix to be updated using three back-to-back GEMMs (General Matrix Multiplications). This transforms the bulk of the computation into a format that Tensor Cores can execute with maximum efficiency.

Auto-Research Methodology: "Loop Engineering"

The developer utilized a high-frequency iteration loop, making over 1,500 submissions over 14 days. The process relied on a combination of LLMs (Codex and Claude) and a structured harness.

Harness Setup and Steering

• Tooling: The developer used the popcorn CLI for benchmarking and submitting to the leaderboard, and Modal for profiling and nsys/NCU analysis.
• Logging: A log.md file tracked every submission, its status (accept/reject), and shape-wise timings to prevent redundant experiments.
• Goal-Oriented Prompting: The /goal command in Codex was used to set quantitative targets (e.g., "beat our active best's n=512 timings"), allowing the model to iterate autonomously for hours or days.
• Oversight: The /btw and /side commands allowed the developer to query the agent's progress and current hypotheses without pausing the main optimization loop.

Escaping Local Maxima with Beam Search

As performance reached the 3,000 $\mu$s mark, the model frequently stalled in local maxima, performing minor parameter tuning rather than structural innovation. To solve this, the developer implemented a beam of candidates:

• Instead of a single incumbent, the agent maintained 3-5 active "idea families."
• This prevented high-risk structural changes from being prematurely discarded just because they initially performed worse than the current best.
• The beam typically included one "exploit" beam (fine-tuning), one "near-miss" beam, and one "structural/high-risk" beam.

Technical Evolution of the Kernel

The kernel evolved through ten major structural breakthroughs to reach the final 1,805 $\mu$s geomean:

Stage	Change	Impact
1	torch.geqrf	Baseline starting point (>108.8k $\mu$s)
2	Blocked WY QR (n=512)	Introduced panel factorization and trailing updates
3	Blocked route (all shapes)	Applied blocked logic across all matrix sizes
4	Triton panels	Implemented custom panel16/32 kernels
5	Cholesky-ORHR (n=4096)	Used Gram-Cholesky for the largest matrices
6	CUDA graph replay	Reduced kernel launch overhead
7	Fused V/T layout	Eliminated slice copies and temporaries
8	Split16 panels	Optimized tail-end processing via Gram-Schmidt
9	Shape specialization	Hardcoded rows and fused reductions
10	Superpanels	V256/T256 packs and direct-H returns (1.80k $\mu$s)

Critical Insights and Trade-offs

Domain Expertise vs. Automation

While the agent could drive significant speedups, domain expertise was critical for steering. The developer noted that the top 10 solutions further optimized by:

• Data Detection: Exploiting specific input distributions (e.g., low-rank cases).
• Library Removal: Replacing remaining PyTorch functions (like triangular solve) with custom CUDA/Triton implementations.
• Precision Management: Keeping the trailing matrix resident in FP16 to avoid repeated type conversions.

Generalization vs. Specificity

Community discussion highlighted a significant risk in this "loop engineering" approach: overfitting to the benchmark.

"8 out of the 10 top solutions... completely broke at any other input than the competition ones. The only solutions that did not break... were made by experts who... followed and adjusted their solution in reasonable bounds."

This suggests that while LLM-driven loops are exceptionally powerful for solving specific, well-defined benchmarks, they struggle to maintain general-purpose robustness without heavy human steering.