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

📅 2026/8/17 2:59:49
超厉害!用 Codex 自动研究,GPU Mode qr_v2 内核加速 232 倍
简介近期GPU Mode 与 Core Automation 合作举办自动研究竞赛问题为实现批量方阵紧凑豪斯霍尔德 QR 分解。183 名参赛者中作者获第 12 名实现比基线快 232 倍加速。文章将分享方法、知识及瓶颈。读博客无需了解相关数学知识作者会重点介绍方法。可点击[问题链接和排行榜](https://www.gpumode.com/leaderboard/774?tabrankings)查看完整竞赛页面此次竞赛是 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⊤vv 确定反射面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 最终内容包括提交规范和束搜索规范等相关指令。