zgba 站群
Auto-research with codex: How I achieved a 232x Faster Kernel

Auto-research with codex: How I achieved a 232x Faster Kernel

使用 Codex 进行自动研究:我如何实现 232 倍更快的内核

08 Jul, 2026 2026 年 7 月 8 日

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. GPU Mode 与 Core Automation 合作,最近举办了一场以自动研究为主题的比赛。题目要求实现批量方形紧凑 Householder QR 分解,即 QR 分解。我在 183 名参赛者中排名第 12,最终实现了比基准解决方案快 232 倍的速度提升。本文将介绍我是如何达成这一成绩的。我将回顾我的方法、学到的经验以及比赛中遇到的瓶颈。这是我第一次认真尝试自动研究。有些人会称这为“循环工程”,老实说,这也无妨。请注意,你不需要详细了解数学原理或问题本身就能理解本文的大部分内容。我侧重于我的方法,而将数学和问题本身置于次要地位,因为大多数阅读本文的人并没有参加比赛。你可以在这里查看完整的比赛页面:问题链接和排行榜。本次比赛是 GPU Mode“研究时代的线性代数内核”系列的一部分。

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. 我们获得了一批形状为 batch x n x n 的方形 FP32 CUDA 矩阵 A,必须返回与 torch.geqrf(A) 相同的紧凑 Householder QR 表示:一个 H 矩阵,其上三角是 R,下三角存储 Householder 向量,外加一个反射系数 tau 向量。检查器使用 torch.linalg.householder_product(H, tau) 重建 Q,取 R = triu(H),并验证:A≈QR, Q⊤Q≈I, Q⊤A≈R。在正确的提交中,排行榜根据形状和条件情况下的几何平均运行时间进行排名。重要的尺寸包括 512 x 512 的批量方形矩阵,以及更大的 1024、2048 和 4096 案例。内部允许使用低位的 FP16、FP8 或 NVFP4,但返回的因子仍必须满足 FP32 风格的 QR 检查。

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 一个微小的 3x3 示例是: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. 这里 Q 是正交的,这意味着它的列是单位长度且彼此垂直,而 R 是上三角的,这意味着对角线以下的所有元素都是零。比赛并没有要求我们直接打印稠密的 Q 和 R;它要求的是紧凑 Householder 版本,让检查器能够重建 Q 并从上三角读取 R。对于上述 3x3 示例,第一个反射器直接将第一列 (12, 6, −4) 映射到 (−14, 0, 0)。−14 变成了 R11。其工作原理在数学部分中有说明。

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. GPU Mode 为参与者提供了 popcorn CLI,使其对智能体友好。智能体可以使用它进行测试、基准测试并直接提交到排行榜。检查器还提供了按形状反馈以及整体几何平均时间。敏锐的观察者会注意到,这是一个适合编写循环的设置。智能体渴望紧密的反馈循环。这允许它们随心所欲地进行爬山算法。GPU Mode 比赛通常提供某种迭代内核的方式。要么你直接提交,要么像 Modal 这样的赞助商提供积分。在这里,组织者基本上允许无限次提交,只要你间隔开时间。如果你不这样做,队列会变长,每个人的运行都会超时。有一次,工作区甚至因为每个人都疯狂提交而耗尽了 Modal 积分。这是一种让学习变得可及的好方法。在 14 天的过程中,我提交了超过 1500 次。

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 我了解 GPU 内核优化的基础知识(主要在 Triton 中,对 CUDA 有一定了解)已经一年了,但并没有在这个领域专业工作。我想告诉你的是,在排行榜上我周围的人中,我是一个弱者。排行榜上排在我前面的人(CUDA Colonel)是 NVIDIA 的首席工程师。无论如何,抛开光环效应不谈,既然我了解基础知识,并且最近阅读了关于 GatedDeltaNet 的内容,我对通用的 GPU 内核术语很熟悉。你对某事了解得越好,你就能越好地提示 LLM,因为你将未知的未知转化为已知的未知。同时,值得注意的是,这场比赛在没有领域知识的情况下也是可行的——比如你可能进不了前 10 名,但你可以仅仅依靠你的测试框架/智能体循环或其他方式,获得比基准值得尊敬的速度提升。我在比赛中的第一步是学习什么是 QR 分解以及如何完成它。有很多方法可以做到这一点——比如 Gram-Schmidt 和 Householder 反射。比赛强制要求使用 Householder 反射。我与 Claude 来回讨论,并观看了一些 YouTube 视频来建立直觉。在与 Claude 讨论之后,很明显我们需要使用分块 Householder 算法作为主要架构,并带有尾随 WY 更新。事实证明,GPT-5.5 对此也有很好的想法。QR 分解是一个相当知名的问题。我觉得这个概念很有趣,因为矩阵分解出现在 LLM 训练的几种现代优化器变体中,特别是在使用矩阵预条件的方法中,例如 Shampoo 风格的优化器和相关方法。Muon(Kimi 使用)是另一个很好的例子:它没有将权重更新视为一个巨大的扁平向量,而是保留了矩阵结构并对动量更新进行正交化,通常通过几次近似 Newton-Schulz 迭代来实现。