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

# 借助Codex实现自动化研究：我在GPU模式qr_v2问题中如何实现比基线快232倍的内核 来源：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模式与Core Automation合作，近期举办了一场以自动化研究为主题的竞赛。问题要求是`实现批量方阵紧凑型Householder QR分解`。我在183名参赛者中排名第12位，最终实现了比基线方案快232倍的加速。本文将介绍我是如何达到这一目标的，包括我的方法、学习心得以及比赛中遇到的瓶颈。这是我首次认真尝试自动化研究。有些人会称之为"循环工程"，老实说，这样称呼也没问题。 请注意，要理解本文大部分内容，无需详细研读数学原理或问题本身。我主要聚焦于自己的方法，数学和问题背景为次要内容，因为大多数读者并未参与本次竞赛。 完整竞赛页面可在此查看：问题链接与排行榜 (https://www.gpumode.com/leaderboard/774?tab=rankings) 本次竞赛是GPU模式《研究时代中的线性代数内核》(https://www.gpumode.com/news/linear-algebra-kernels-age-of-research) 系列活动的一部分。 #### 问题介绍 给定一批方阵FP32 CUDA矩阵`A`，形状为`batch x n x n`，需要返回与`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=\[12−5146167−68−424−41\]=\[6/7−69/175−58/1753/7158/1756/175−2/76/35−33/35\]⏟Q\[1421−140175−700035\]⏟R 此处`Q`是正交矩阵，意味着其列向量长度为1且相互垂直；`R`是上三角矩阵，意味着对角线以下所有元素均为零。竞赛并非要求直接输出稠密`Q`和`R`矩阵，而是要求输出紧凑的Householder表示形式，以便检查器重建`Q`并从上三角读取`R`。 对于上述3×3示例，首个反射器将第一列`(12, 6, −4)`直接映射到`(−14, 0, 0)`。其中−14成为R11。具体原理将在数学部分说明。 ### 为何此问题适合自动化研究 GPU模式为参与者提供了popcorn CLI工具，使其对代理程序友好。代理可使用该工具直接测试、基准测试和向排行榜提交结果。检查器还提供按形状分类的反馈以及整体几何平均计时。 敏锐的观察者会发现，这是设计循环的绝佳设置。代理程序渴望紧密的反馈循环，这使它们能够尽情进行爬山优化。 GPU模式竞赛通常提供某种内核迭代方式：要么直接提交，要么像Modal这样的赞助商提供积分支持。本次组织者基本允许无限次提交，只需控制提交间隔。如果提交过于频繁，队列会变长，所有人的运行都会超时。甚至有一次工作区因所有人集中提交而耗尽Modal积分。这是使学习过程更易获取的好方法。 在14天赛程中，我提交了超过1500次。 #### 学会提出更好的问题 我了解GPU内核优化的基础知识（主要在Triton框架，对CUDA有一定理解）已有一年时间，但并未在此领域全职工作。我想表达的是，在排行榜上，我相对于周围参赛者而言属于"弱势群体"。排行榜上排名在我上一位（CUDA Colonel）是NVIDIA的首席工程师。 总而言之，抛开这些不谈，由于我具备基础知识且近期阅读了GatedDeltaNet相关论文，对GPU内核通用术语保持熟悉。 你对某领域了解越深，就能越好地引导大语言模型生成提示，因为这能将未知的未知转化为已知的未知。 同时值得指出，本次竞赛即使缺乏专业知识也可参与——虽然可能无法进入前10名，但仅通过运行框架/代理循环等方式，就能实现相对基线的显著加速。 我竞赛初期的首要任务是了解QR分解及其多种实现方法（如Gram-Schmidt方法和Householder反射）。竞赛指定使用Householder反射。我通过与Claude多次讨论并观看YouTube视频建立直观理解。讨论后明确需要采用分块Householder算法作为主要架构，并配合尾部WY更新。事实证明，GPT-5.5对此也有良好见解。QR分解是相当知名的问题。 我发现此概念很有趣，因为矩阵分解在现代大语言模型训练的多种优化器变体中频繁出现，特别是在使用矩阵预条件的Shampoo风格优化器及相关方法中。Muon（Kimi使用）是另一个好例子：它不将权重更新视为巨型展平向量，而是保留矩阵结构并对动量更新进行正交化，通常通过几次近似极分解的Newton-Schulz迭代实现。 ### \(可选\) QR分解的数学原理：Householder反射 **如对数学原理感兴趣建议略读本节，否则可直接跳过。**需注意的是，Householder QR存在顺序依赖性，这使得GEMM实现变得困难。我们采用分块Householder使其更接近矩阵乘法形式。 #### 契约概述 快速回顾契约要求：输入是一批方阵FP32矩阵`A`；输出是`torch.geqrf`返回的紧凑`\(H, tau\)`格式。`H`的上三角为`R`，对角线下方存储Householder向量，`tau`每列对应一个标量。检查器从`\(H, tau\)`重建`Q`并验证`A ≈ QR`。 #### 镜子原理 暂时忘掉矩阵。在浴室镜子中，你的反射影像恰好位于玻璃后方与你在前方相同的距离处，笔直穿透玻璃。 如果x⟂是`x`垂直于玻璃面的部分，反射操作就是减去该部分两次： x_reflected = x - 2x_⟂ 因此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`表示紧凑输出矩阵。） **二维Householder反射** tau = 0.00 反射向量 x 及其垂直/平行分量。镜面规则：保留平行分量，翻转垂直分量。因此反射后x = x - 2x_⟂。 简化的二维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/(vTv)`。反射器本身为： H = I - τvv⊤，τ = 2/(v⊤v) Hx = (I - τvv⊤)x Hx = x - τv(v⊤x) HA_active = (I - τvv⊤)A_active HA_active = A_active - τv(v⊤A_active) 这将当前列变换为`\((-14, 0, 0)\)`，并相应重写其他列，使得下一个反射器基于更新后的矩阵构建。 我们对每列重复此操作。近似表示如下重复更新过程： A⁽¹⁾ = A⁽⁰⁾ - τ₁v₁(v₁⊤A⁽⁰⁾) A⁽²⁾ = A⁽¹⁾ - τ₂v₂(v₂⊤A⁽¹⁾) A⁽³⁾ = A⁽²⁾ - τ₃v₃(v₃⊤A⁽²⁾) ⋯ R = A⁽ⁿ⁾ 每行使用前一行产生的矩阵。每个反射器将当前列对角线下方元素归零，且不影响已完成列。最后一次操作后，`A`已转化为上三角矩阵`R`： A = H₁H₂⋯Hₙ⏟QR 镜面不改变长度或角度，因此每个Hj都是正交矩阵，其乘积`Q`也正交。这正是检查器验证的正交性来源。 #### 紧凑格式是什么 处理完第`j`列后，其对角线下方元素成为闲置空间。`geqrf`重用这些位置存储`v_j`的尾部（前导1是隐含的）。对角线上方为`R`，下方为反射器，`tau`作为独立向量伴随存储。这就是为何检查器需要`H`和`tau`来重建`Q`。 **紧凑geqrf存储（n = 6）**悬停某列查看 H（上三角=R，对角线下方=反射器尾部） tau（每个反射器对应一个标量） **R**的元素 R的**vj尾部****τj** 单个矩阵，双重数据。悬停（或点击）任意列j：反射器j生效后，该列对角线下方位置归零，因此`geqrf`重用这些位置存储vj的尾部。对角线上显示的前导1（高亮时显示）是隐含的。每列与对应τj配对，检查器即可重建Q。### 借助分块Householder算法减少串行计算量 Householder QR逐列归零`A`对角线下方元素。每步根据当前列构建反射器，并应用到其右侧所有元素。问题在于反射器j+1必须基于反射器j作用后的矩阵构建。因此步骤无法重排或融合。这是串行操作，串行矩阵向量运算运行在SM的慢速向量通道中，而张量核心却闲置。 **Householder QR，逐次反射器（n = 5）**数学视图：零元素出现，尾部块被重写 原始矩阵 最终由当前反射器重写的**R**元素 **0**被归零（vj存储于此） 反射器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，则： H₁H₂⋯Hᵦ = I - VTV⊤ 尾部块更新变为三个GEMM步骤： W = V⊤A_trail Z = T⊤W A_trail ← A_trail - VZ **分块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 = 32, 176, 352, 512, 1024, 2048, 4096）和批次规模的巨大差异——最大矩阵因批次过少难以填满张量核心，而n = 32时矩阵过小，需将多个矩阵打包到单次内核启动中。 ### Codex极限优化 #### 为何选择Codex 我使用ChatGPT Pro（200美元订阅）