@dejavucoder: 我最新的一篇博客文章 "auto-research with codex: 我如何在使用Codex的GPU Mode中实现比基线快212倍的内核…"
摘要
Sankalp的博客文章,详细描述了他如何使用Codex在GPU Mode的竞赛中为QR分解实现快232倍的GPU内核,并概述了他的自动研究方法论。
查看缓存全文
缓存时间: 2026/07/09 08:00
我最新发布的博文《auto-research with codex: how I achieved a 212x faster kernel over baseline with codex in GPU Mode’s qr_v2 problem》现已上线。在这篇文章中,我分享了自己在QR分解问题上采用自动核编程的方法。
https://t.co/tKwV9xMDC8 https://t.co/OUKzqrXjdA
Auto-research with codex: How I achieved a 232x Faster Kernel over baseline with Codex in GPU Mode’s qr_v2 problem
来源:https://sankalp.bearblog.dev/autoresearch/ 2026年7月8日
目录
- 引言 (https://sankalp.bearblog.dev/autoresearch/#intro)
- 竞赛简述 (https://sankalp.bearblog.dev/autoresearch/#contest-in-short)
- 问题简介 (https://sankalp.bearblog.dev/autoresearch/#problem-intro)
- 为何此问题适合自动研究 (https://sankalp.bearblog.dev/autoresearch/#why-this-problem-is-auto-research-able)
- 学习足够,提出更好的问题 (https://sankalp.bearblog.dev/autoresearch/#learning-enough-to-ask-better-questions)
- (可选)QR分解的数学原理:Householder反射 (https://sankalp.bearblog.dev/autoresearch/#optional-math-for-qr-decomposition-householder-reflections)
- 借助分块Householder算法减小串行工作量 (https://sankalp.bearblog.dev/autoresearch/#make-serial-work-small-with-the-help-of-the-blocked-householder-algorithm)
- 其他挑战 (https://sankalp.bearblog.dev/autoresearch/#other-challenges)
- Codex最大化 (https://sankalp.bearblog.dev/autoresearch/#codex-maxxing)
- 核进展突破 (https://sankalp.bearblog.dev/autoresearch/#kernel-progress-breakthroughs)
- 突破性思路 (https://sankalp.bearblog.dev/autoresearch/#breakthrough-ideas)
- 引入思路多样性以跳出局部最优 (https://sankalp.bearblog.dev/autoresearch/#introducing-idea-diversity-to-escape-the-local-maxima)
- 实现提示 (https://sankalp.bearblog.dev/autoresearch/#implementation-hints)
- 本可做得更好的地方 (https://sankalp.bearblog.dev/autoresearch/#what-i-could-have-done-better)
- 结论 (https://sankalp.bearblog.dev/autoresearch/#conclusion)
- 参考文献 (https://sankalp.bearblog.dev/autoresearch/#references)
- 致谢 (https://sankalp.bearblog.dev/autoresearch/#acknowledgements)
引言
竞赛简述
GPU Mode与Core Automation合作,近期举办了一场以自动研究为主题的竞赛。问题陈述是“实现批量平方紧凑Householder QR分解,即QR分解”。我在183名参赛者中排名第12,最终实现了比基准方案232倍的加速。本文讲述我如何取得这一成果。我将介绍我的方法、学习心得以及在竞赛中遇到的主要瓶颈。这是我第一次认真尝试自动研究。有些人会称之为“循环工程”,说实话这也没问题。
请注意,你无需详细了解数学或问题本身也能跟上本文的大部分内容。我着重于方法,而将数学和问题本身置于次要位置,因为大多数读者可能并未参与竞赛。
你可以在这里查看完整的竞赛页面:问题链接和排行榜 (https://www.gpumode.com/leaderboard/774?tab=rankings)
本次竞赛是GPU Mode“研究时代的线性代数核”系列 (https://www.gpumode.com/news/linear-algebra-kernels-age-of-research) 的一部分。
问题简介
我们得到一批形状为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校验。
一个小型3 x 3示例如下:
A = [\begin{bmatrix} 12 & -51 & 4 \ 6 & 167 & -68 \ -4 & 24 & -41 \end{bmatrix} = \underbrace{\begin{bmatrix} 6/7 & -69/175 & -58/175 \ 3/7 & 158/175 & 6/175 \ -2/7 & 6/35 & -33/35 \end{bmatrix}}{Q} \underbrace{\begin{bmatrix} 14 & 21 & -14 \ 0 & 175 & -70 \ 0 & 0 & 35 \end{bmatrix}}{R}]
这里Q是正交的,即其列是单位长度且互相垂直,R是上三角矩阵,即对角线以下全为零。竞赛并不要求我们直接输出稠密的Q和R;而是要求紧凑的Householder版本,让检查器能够重构Q并从上半部读取R。
对于上面的3×3示例,第一个反射器将第一列 (12, 6, -4) 直接映射到 (-14, 0, 0)。这个 -14 成为R11。其工作原理在数学部分有说明。
为何此问题适合自动研究
GPU Mode 为参与者提供了 popcorn CLI,使其对智能体友好。智能体可以使用它直接测试、基准测试和提交到排行榜。检查器还会提供形状相关的反馈以及整体几何平均时间。
敏锐的观察者会注意到,这是编写循环的绝佳设置。智能体渴望紧密的反馈循环,这让他们可以尽情地爬山式优化。
GPU Mode 的竞赛通常会提供一些迭代核的方式。要么直接提交,要么由像 Modal 这样的赞助商提供积分支持。这里的组织者基本上允许无限次提交,只要间隔足够。如果不这样做,队列会变长,所有人的运行都会超时。有一次工作区甚至耗尽 Modal 积分,因为大家都在疯狂提交。这是让学习变得可及的好方法。
在14天的竞赛中,我提交了超过1500次。
学习足够,提出更好的问题
代码图像(1)
我了解 GPU 核优化的基础知识(主要是在 Triton 中,也对 CUDA 有一定理解)已有一年,但并未在此领域专业工作。我想告诉你的是,在排行榜上我身边的人群中,我是处于劣势的选手。在我排行榜上方的那位(CUDA Colonel)是 NVIDIA 的首席工程师。
不过,除了气场展示之外,由于我掌握了基础知识,并且最近刚刚读过关于 GatedDeltaNet 的文章,我对 GPU 核的一般术语还很熟悉。
你对某件事了解得越深,就越能更好地向大语言模型发问,因为你将未知的未知转化为已知的未知。
同时,值得注意的是,这个竞赛在没有领域知识的情况下也是可以完成的——你可能不会进入前十名,但仅靠你的工具/智能体循环,就能获得比基准方案相当可观的加速。
我在竞赛中的第一步是了解什么是QR分解以及如何实现。有几种方法,比如Gram-Schmidt和Householder反射。竞赛规定使用Householder反射。我与Claude反复讨论,并观看了一些YouTube视频来建立直觉。在与Claude讨论后,很明显我们需要使用分块Householder算法作为主要架构,并配合尾随WY更新。结果发现,GPT-5.5对此也有很好的理解。QR分解是一个相当著名的问题。
我发现这个概念很有趣,因为矩阵分解出现在几种用于大语言模型训练的现代优化器变体中,特别是在使用矩阵预条件的方法中,例如Shampoo风格的优化器和相关方法。Muon(由Kimi使用)是另一个很好的例子:它不是将权重更新视为一个巨大的扁平化向量,而是保持矩阵结构,并通过几次近似极分解的Newton-Schulz迭代来正交化动量更新。
(可选)QR分解的数学原理:Householder反射
**如果您对数学感兴趣,建议快速浏览本节,否则可以跳过。**只需注意,Householder QR中存在串行依赖性,这使得进行GEMM变得困难。我们使用分块Householder使其更符合矩阵乘法的形状。
andrew
约定
快速回顾约定:输入是批量方形FP32矩阵A;输出是紧凑的(H, tau)格式,与torch.geqrf返回的相同。H的上半部是R。在对角线以下,H存储Householder向量,而tau为每列存储一个标量。检查器从(H, tau)重建Q,并验证A ≈ QR。
镜子
暂时忘记矩阵。在浴室的镜子中,你的反射在玻璃后面的距离与你在前面的距离完全相等,且与镜面成直线。
如果x⊥是x垂直于镜面的部分,反射只需减去那个部分两次:
x_reflected = x - 2 x⊥
因此,Householder反射就是找到垂直部分并减去两次。
存储镜子
Householder向量就是镜子,以紧凑形式存储。在代码中,我们并不携带整个镜面平面。我们存储一个垂直于它的向量v。镜面是垂直于v的所有东西,反射沿着v发生。
垂直部分只是x沿着v的影子,即 (v⊤x)/(v⊤v) 倍的v。将其代入上面的减法中:
Hx = x - τ v (v⊤x) ,其中 τ = 2 / (v⊤v)
因此tau只是2/(v⊤v):因子2和v的长度捆绑成一个预计算的数字。v选择镜子,tau缩放更新。(我将用数学符号Hj表示反射器,而用H表示紧凑输出矩阵。)
2D中的Householder反射 tau = 0.00
x 反射后的 x v x_parallel x_perp 镜子规则:保留x_parallel,翻转x_perp。因此反射后的 x = x - 2 x_perp。 一个简化的2D Householder步骤。拖动橙色向量:紫色镜子改变,使得绿色反射后的向量落在水平轴上。在高维空间中,落在轴上正是使对角线以下元素变为零的原因。
为什么QR需要镜子
QR希望将A变成上三角矩阵R。第1列应该变成 (∗, 0, 0)之类的形式,第2列在第2行以下应该为零,以此类推。
Householder镜子的作用是可以一次性地对一列进行这样的操作。以上面3×3示例的第一列为例:(12, 6, -4)。我们希望将其发送到x轴,使得下面的项变为零。反射只能改变方向,不能改变长度,因此目标也必须具有长度14。一个有效的目标是 (-14, 0, 0)。经过那次反射后,6和-4项就消失了,这正是我们想要的。
如何找到镜子?它位于该列及其目标之间的中间位置,因此穿过镜子的向量v就是该列减去其目标:
v = (12, 6, -4) - (-14, 0, 0) = (26, 6, -4)
Householder更新
然后计算 tau = 2 / (v⊤v)。反射器本身就是:
H = I - τ v v⊤, τ = 2 / (v⊤v) Hx = (I - τ v v⊤) x Hx = x - τ v (v⊤ x) H A_active = (I - τ v v⊤) A_active H A_active = A_active - τ v (v⊤ A_active)
这会将当前列变为 (-14, 0, 0),并一致地重写其他列,使得下一个反射器基于更新后的矩阵构建。
我们对每一列重复这个过程。用粗略的表示法,重复更新看起来像:
A(1) = A(0) - τ1 v1 (v1⊤ A(0)) A(2) = A(1) - τ2 v2 (v2⊤ A(1)) A(3) = A(2) - τ3 v3 (v3⊤ A(2)) … R = A(n)
每一行使用上一行产生的矩阵。每个反射器在不干扰已完成列的情况下将其所在列对角线以下的所有元素归零。在最后一个之后,A已演变为上三角R:
A = H1 H2 … Hn Q = H1 H2 … Hn
镜子不会改变长度或角度,因此每个Hj都是正交的,它们的乘积Q也是如此。这就是检查器验证的正交性自然得来的原因。
什么是紧凑格式
一旦处理完第j列,其对角线以下的所有空间都是死区。geqrf重用这些槽位来存储v_j的尾部(前面的1是隐式的)。对角线及其上方是R;下方是反射器;而tau作为单独的向量存在。这就是检查器需要H和tau两者来重建Q的原因。
Compact geqrf storage (n = 6) 鼠标悬停在一列上
H (上半部 = R, 对角线以下 = 反射器尾部)
tau (每个反射器一个标量)
r R的元素 v_j的尾部 τ_j
一个矩阵,两种负载。悬停(或点击)任意列 j:该列对角线以下的槽位在反射器j触发后为零,因此geqrf重用它们来存储v_j的尾部。高亮时对角线上的前导1是隐式的。将每列与其τ_j配对,检查器即可重建Q。
借助分块Householder算法减小串行工作量
Householder QR逐列将A对角线以下归零。每一步从当前列构建一个反射器,并将其应用于右侧的所有内容。问题在于反射器j+1是基于反射器j已经作用后的矩阵构建的。因此你不能重新排序步骤,也不能融合它们。它是串行的,串行的矩阵-向量工作在SM的慢速向量通道中运行,而张量核心则闲置。
Householder QR,一次一个反射器 (n = 5) 数学视图:零出现,尾部块被重写
原始 a 最终进入R 由此反射器重写 0 归零 (v_j 存放在此)
反射器j将列j对角线以下归零并确定R的第j行。但它也会重写整个橙色尾部块,而反射器j+1只能从重写后的块构建。这种数据依赖性就是分块算法要攻击的串行链。
经典的修复方法是分块算法。你选择一个窄的面板,包含b列(比如32或64),并在其内部完成所有串行工作。这没问题,因为面板只有b列宽,所以代价很低。然后,不是一次一个地应用面板的b个反射器到矩阵的其余部分,而是将它们压缩成一个单次秩b更新(“WY表示”),并用三个连续的矩阵乘法一次性击打整个尾部块。串行工作被限制在面板内,其他所有内容都变成了GEMM,这正是张量核心想要的形状。
具体来说,WY表示将一个面板的b个反射器压缩成一个单次秩b更新。将面板的Householder向量堆叠成V = [v1, v2, …, vb]的列,构建一个小的b×b上三角矩阵T,然后:
H1 H2 … Hb = I - V T V⊤
并且尾部块更新变成三个GEMM形状的步骤:
W = V⊤ A_trail Z = T⊤ W A_trail ← A_trail - V Z
分块Householder (n = 6, 面板宽度 b = 2) 将串行工作局限在面板内,其余部分用GEMM
最终进入R 存储的v (面板,串行) 由秩b更新重写 未触及
前一图中的逐列串行工作仍然发生。但只发生在狭窄的蓝色面板内部,成本很低。然后,b个反射器被压缩成 I − V T V⊤,并一次性击打整个尾部块:三个GEMM,而不是b个单独的秩1更新。然后面板滑动到橙色块上,故事重复。
面板是串行矩阵-向量工作卡住的地方,但它只有b列宽。其右侧的所有内容都是大的尾部块,那完全是GEMM。当面板沿对角线移动时,尾部块缩小。
如果你想理解数学,我建议与Claude进行头脑风暴。也可以查看Mike的博客文章 (https://ml-mike.com/writing/qr_v2/),他比我更描述性和可视化地涵盖了数学。他在竞赛中排名第5,并专注于问题分享了他的学习心得。
其他挑战
另外两件事被证明具有挑战性:在内部可靠地使用低精度(特别是对于病态输入),以及应对广泛的形状分布(n = 512, 1024, 2048, 4096)和混合条件数。
相似文章
@shiposcant: 完成阅读这篇:"你越了解某件事,就越能更好地提示LLM,因为你可以将未知的未知..."
一篇博客文章描述了使用Codex自动迭代和优化GPU内核,相较于基线实现了212倍加速。文章强调了专业知识如何放大AI的效用,通过循环实验工作流将未知的未知转化为已知的未知。
NVIDIA 工程师与研究人员如何利用 Codex 进行构建
NVIDIA 的工程师和研究人员正在使用由 GPT-5.5 驱动的 OpenAI Codex,作为处理复杂工程任务和端到端机器学习工作流的默认工具。本文重点介绍了通过在该 NVIDIA 基础设施上集成 Codex 所取得的显著生产力提升、自主系统构建以及研究自动化成果。
@levidiamode: GPU编程的第163/365天 - 今天看几个不同的agentic GPU内核优化系统。我最感兴趣的两个是…
一条推文讨论了两种agentic GPU内核优化系统:@dogacel0的Auto GPU Kernel和@songhan_mit实验室的Kernel Design Agents,两者均在MLSys Sparse Attention FlashInfer比赛中获胜。该帖子突出了使用子代理和Claude技能进行GPU编程的不同方法。
@SzymonOzog_: 周六阅读:“What happens when you run a CUDA kernel”——非常酷的博客文章,详细介绍了CPU与GPU通信的细节…
推荐一篇博客文章的推文,该文章解释了启动CUDA内核时所需的CPU-GPU通信细节。
@injaneity: https://x.com/injaneity/status/2075659478096376158
本文解释了批处理和并行操作如何改善AI计算机使用系统中的延迟和效率,重点介绍了pi-computer-use和cua-driver等开源实现,它们在Codex出现类似功能之前就取得了显著的性能提升。