使用 Codex 进行自动研究：实现 232 倍更快的 QR 分解内核

执行摘要

通过使用 OpenAI 的 Codex 实现自动化研究循环，开发者在批量方阵紧凑-Householder QR 分解上实现了比 torch.geqrf 基线快 232 倍的性能提升。成功的关键在于从通用库调用转向分块 Householder 算法，采用“候选者束”策略避免局部最优，并充分利用 GPU 模式 popcorn CLI 和 Modal 性能分析提供的紧密反馈循环。

优化挑战：QR 分解

目标是为 FP32 CUDA 矩阵实现批量方阵紧凑-Householder QR 分解。输出需要一个紧凑表示，包含一个 H 矩阵（上三角为 R，下三角存储 Householder 向量）和一个 tau 向量（反射系数）。

串行瓶颈

标准 Householder QR 本质上是串行的：第 $j+1$ 个反射器依赖于第 $j$ 个反射器生成的矩阵。这种顺序依赖性阻止了 Tensor Cores 的有效利用，因为工作仍保持矩阵-向量形态而非矩阵-矩阵形态，导致 GPU 最强大的计算单元处于空闲状态。

分块 Householder 解决方案

为克服串行瓶颈，采用了 分块 Householder 算法。该方法将串行工作限制在宽度为 $b$ 的窄列面板内。随后，$b$ 个反射器被压缩为单个秩为 $b$ 的更新（WY 表示），从而允许使用三个连续的 GEMM（通用矩阵乘法）更新矩阵的尾部块。这将大部分计算转换为 Tensor Cores 可以以最大效率执行的格式。

自动研究方法论：“循环工程”

开发者使用高频迭代循环，在 14 天内提交超过 1,500 次。该过程依赖于 LLM（Codex 和 Claude）与结构化框架的结合。

框架设置与引导

• 工具链： 开发者使用 popcorn CLI 进行基准测试和提交排行榜，使用 Modal 进行性能分析和 nsys/NCU 分析。
• 日志记录： 一个 log.md 文件记录了每次提交、其状态（接受/拒绝）以及按形状划分的耗时，以避免重复实验。
• 目标导向提示： 在 Codex 中使用 /goal 命令设定量化目标（例如，“超越当前最佳的 n=512 时序”），使模型能够自主迭代数小时甚至数天。
• 监督机制： /btw 和 /side 命令允许开发者在不中断主优化循环的情况下查询代理的进展和当前假设。

通过束搜索摆脱局部最优

当性能达到 3,000 $μ$s 标记时，模型频繁陷入局部最优，仅进行微小的参数调优而非结构性创新。为解决此问题，开发者实施了 候选者束：

• 不再仅维护一个当前最优方案，而是保持 3-5 个活跃的“想法家族”。
• 这防止了高风险的结构性变更因初始表现略差于当前最佳而被过早丢弃。
• 该束通常包含一个“利用”束（精细调优）、一个“接近成功”束和一个“结构性/高风险”束。

内核的技术演进

内核经历了十次重大结构突破，最终达到 1,805 $μ$s 的几何平均性能：

阶段	变更	影响
1	torch.geqrf	基线起点（>108.8k $μ$s）
2	分块 WY QR（n=512）	引入面板分解和尾部更新
3	分块路径（所有形状）	将分块逻辑应用于所有矩阵尺寸
4	Triton 面板	实现自定义 panel16/32 内核
5	Cholesky-ORHR（n=4096）	对最大矩阵使用 Gram-Cholesky
6	CUDA 图重放	降低内核启动开销
7	融合 V/T 布局	消除切片复制和临时变量
8	Split16 面板	通过 Gram-Schmidt 优化尾端处理
9	形状特化	硬编码行数并融合归约
10	超面板	V256/T256 包和直接-H 返回（1.80k $μ$s）

关键洞察与权衡

领域知识 vs. 自动化

尽管代理能够推动显著的性能提升，但领域知识在引导过程中至关重要。开发者指出，前 10 名解决方案通过以下方式进一步优化：

• 数据检测： 利用特定输入分布（例如低秩情况）。
• 库移除： 将剩余的 PyTorch 函数（如三角求解）替换为自定义 CUDA/Triton 实现。
• 精度管理： 保持尾部矩阵在 FP16 中驻留，避免重复类型转换。

泛化性 vs. 特异性

社区讨论突显了这种“循环工程”方法的重大风险：对基准测试的过拟合。

“在前 10 名解决方案中，有 8 个在任何非竞赛输入下完全失效。唯一未失效的解决方案……是由专家们……在合理范围内跟进并调整其方案所实现的。”

这表明，尽管 LLM 驱动的循环在解决特定、定义明确的基准测试方面极为强大，但若缺乏大量人工引导，它们难以维持通用的鲁棒性。