超厉害！用 Codex 自动研究，GPU Mode qr_v2 内核加速 232 倍

简介

近期，GPU Mode 与 Core Automation 合作举办自动研究竞赛，问题为实现批量方阵紧凑豪斯霍尔德 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)` 相同的紧凑豪斯霍尔德 QR 表示，包括上三角为 `R`、下三角存储豪斯霍尔德向量的 `H` 矩阵，以及存储反射系数的 `tau` 向量。检查器会用 `torch.linalg.householder_product(H, tau)` 重建 `Q`，提取 `R = triu(H)`，并验证等式：A≈QR,Q⊤Q≈I,Q⊤A≈R。排行榜按不同形状和条件下几何平均运行时间排名，重要矩阵大小有批量方阵如 `512 x 512` 等。内部允许用低比特的 FP16、FP8 或 NVFP4，但返回因子要满足 FP32 风格 QR 检查。给出 `3 x 3` 小例子，`Q` 为正交矩阵，`R` 为上三角矩阵。竞赛要求输出紧凑豪斯霍尔德形式，方便检查器重建 `Q` 并读取 `R`。对于 `3×3` 例子，第一个反射器能将第一列 `(12, 6, -4)` 映射到 `(-14, 0, 0)`，`-14` 成 `R11`，原理在数学部分介绍。

为何该问题适合自动研究

GPU Mode 为参赛者提供 popcorn CLI 与智能体交互，智能体可用其测试、基准测试并提交到排行榜，检查器提供形状反馈和几何平均时间。这为编写循环提供理想环境，智能体渴望紧密反馈循环以优化。GPU Mode 竞赛提供迭代内核方式，可直接提交或由赞助商提供积分支持。此次竞赛组织者基本允许无限次提交，合理安排间隔即可。曾出现工作区耗尽 Modal 积分情况，因大家频繁提交。14 天里，作者提交超 1500 次。

积累足够知识以提出更好的问题

作者了解 GPU 内核优化基础知识一年，但未在该领域专业工作，在排行榜参赛者中处于劣势，前一位参赛者是 NVIDIA 首席工程师。不过作者掌握基础知识且读了关于 GatedDeltaNet 的文章，对 GPU 内核通用术语较熟悉。对领域了解越深，越能向大语言模型提好问题，将未知变已知。无相关领域知识也可参与竞赛，虽难进前 10 名，但靠测试框架/智能体循环可显著加速。竞赛开始，作者学习 QR 分解概念和方法，有格拉姆 - 施密特方法和豪斯霍尔德反射法，竞赛要求用豪斯霍尔德反射法。作者与 Claude 交流、看 YouTube 视频建立理解，明确用分块豪斯霍尔德算法结合尾随 WY 更新，GPT - 5.5 也有相关见解。矩阵分解在大语言模型训练优化器变体中常见，如 Shampoo 风格优化器及 Muon 方法。

（可选）QR 分解的数学原理：豪斯霍尔德反射

豪斯霍尔德 QR 分解有顺序依赖，进行通用矩阵乘法困难，用分块豪斯霍尔德方法更适合矩阵乘法。回顾问题约定，输入是一批形状为 `batch x n x n` 的方阵 FP32 矩阵 `A`，输出是 `torch.geqrf` 返回的紧凑 `(H, tau)` 格式，`H` 上三角为 `R`，对角线以下存豪斯霍尔德向量，`tau` 每列存一个标量，检查器从 `(H, tau)` 重建 `Q` 并验证 `A ≈ QR`。以浴室镜子为例说明反射原理，豪斯霍尔德反射是找到垂直分量并减去两次。豪斯霍尔德向量是反射面紧凑存储，垂直分量是 `x` 在 `v` 上投影，代入公式可得相关计算。`tau` 是 `2 / v⊤v`，`v` 确定反射面，`tau` 缩放更新量。二维豪斯霍尔德反射有相关规则和特点，在高维空间可使对角线以下元素变零。QR 分解目标是将矩阵 `A` 转上三角矩阵 `R`，豪斯霍尔德反射面可一次性转换列，以 `3×3` 例子说明找到反射面和目标向量的方法。计算 `tau` 后可得反射器公式，对每列重复操作可将 `A` 变为上三角矩阵 `R`，反射操作使 `Q` 为正交矩阵，这是检查器验证正交性来源。处理完列后，`geqrf` 复用对角线以下空间存 `v_j` 尾部，`R` 在对角线上及以上，`tau` 单独存储，检查器用 `H` 和 `tau` 重建 `Q`。

借助分块豪斯霍尔德算法减少串行工作量

豪斯霍尔德 QR 分解逐列将矩阵 `A` 对角线以下元素置零，反射器构建有数据依赖，是串行过程，矩阵 - 向量运算在慢速向量通道运行，张量核心闲置。分块算法是经典解决方案，选宽度为 `b` 的窄面板完成串行工作，因面板列少计算成本低。将面板 `b` 个反射器压缩为秩为 `b` 的更新（WY 表示），通过三次矩阵乘法应用到尾随块，串行工作限制在面板内，其余转化为 GEMM 运算。具体给出 WY 表示和尾随块更新步骤。逐列串行工作在窄面板内进行，成本低，`b` 个反射器压缩后一次性应用到尾随块，面板移动重复过程。面板是瓶颈但列少，右侧尾随块是 GEMM 运算，随面板移动缩小。想深入理解数学原理可与 Claude 探讨或查看[Mike 的文章](https://ml-mike.com/writing/qr_v2/)，他获第 5 名并分享学习经验。

其他挑战

一是内部可靠使用低精度，尤其是病态输入；二是应对不同形状（n = 32, 176, 352, 512, 1024, 2048, 4096）和批量大小变化。大矩阵批量少难利用张量核心，`n = 32` 小矩阵需打包到一次内核启动。

充分利用 Codex

作者用 ChatGPT Pro 和 Claude Pro 参赛，用 Modal 进行性能分析，Modal 每月免费提供 30 美元积分。选 Codex 原因：订阅更高版本；直觉认为 OpenAI 模型在 Triton 方面表现好；`/compaction` 在 Codex 中效果好。对问题有基本了解后，让 Codex 做基本设置，添加相关文件并维护日志。日志可记录想法有效性，达 3000 微秒性能瓶颈后更重视日志。Codex 优点是按明确指令执行，给出详细提示和目标可工作数小时。最初手动提示 Codex 在 Triton 中实现针对 `n = 512` 和 `n = 1024` 形状的优化。可用 `/goal` 引导模型循环到目标，设定合理数值目标和具体标准效果好。作者每 2 - 3 小时提供输入引导模型，有时让其夜间无人监督运行，最初督促模型多使用 Triton 等。使用 `/goal` 时可用 `/btw` 或 `/side` 不中断循环提问，作者会咨询 Claude 并反馈想法。基线 `torch.geqrf` 路径运行时间约 419 毫秒，为 `n = 512` 形状实现分块豪斯霍尔德算法后，一天内将运行时间缩到 5000 微秒。232 倍加速是基线时间与最终 1,805 微秒结果对比得出，谱系图显示从 108,803 微秒到 1,805 微秒提升。达 3000 微秒瓶颈后优化困难，作者让 Codex 进行性能分析。展示 QR v2 * B200 全表几何平均数据及内核进展突破情况。列出 QR 内核结构演变表格，包括结构变化、作用、类型和几何平均时间。在性能分析找瓶颈需反复尝试，启动开销和面板开销是主要瓶颈，精力花在减少面板开销和使 WY 更新适合 GEMM 运算上。作者反复问自己一系列问题，如开销解决方法、信息获取、性能分析看法等，还提到归约融合等优化尝试。

引入思路多样性以跳出局部最优

性能从 3000 微秒提升到 1800 微秒时，模型易陷入局部最优，表现为手动调参尝试小变体。几年前智能体常因不够智能等陷入循环，当时用不同采样方法等解决。现在模型挑战是产生新想法和有“研究品味”，作者写了相关文章。作者采用多种策略帮助模型跳出局部最优，如保留 3 - 5 个候选方案、人工干预、鼓励冒险、用强大顾问模型、指示使用子智能体、用 NCU 和 Modal 分析、清理工作等，还考虑过多智能体群策略但未尝试。认为强大顾问策略将成自动研究标准策略，Claude Code 提供 `/advisor` 命令。

实现提示

展示最终目录结构和 `AGENTS.md` 最终内容，包括提交规范和束搜索规范等相关指令。