Sankalp 使用 Codex loop 在 GPU Mode 的 B200 QR 竞赛中获得第12名

Sankalp在使用Codex、性能分析工具以及超过1,500次提交来优化B200 QR kernel之后,在183名参赛者中名列第12位。

By · Published

Primary source: Sankalp's blog

Why it matters

Sankalp's result shows how a strict benchmark can turn a coding agent into an experimental system. Engineers still have to choose the architecture, interpret profiles and change the search policy when the agent stalls.

The focused labor of a programmer optimizing code (Oil painting in the manner of Edward Hopper)

独立可验证的结果是 Sankalp 的 在 NVIDIA B200 上以 1,804.779 微秒在 183 名参赛者中排名第 12。他在使用 OpenAI Codex 的情况下,通过 在 14 天内进行超过 1,500 次提交 达到这一结果。该项目是 GPU Mode 的 Linear Algebra Kernels in the Age of Research 系列 及其 qr_v2 竞赛的一部分。

7 月 8 日对该项目的叙述 中,Sankalp 写道:“在 14 天里,我进行了超过 1500 次提交。” Sankalp 的博客报告称相对于一个大约为 419,000 微秒的 PyTorch 基线他实现了 232 倍的改进;该比较针对的是此基准工作负载,并未在提供的材料中独立复现。

比赛在 6 月 29 日结束,约比本文发布早六个半星期。Sankalp 的最终结果大约比 获胜者的 1,220.774 微秒提交慢 48%。他的名次仍值得关注,因为他以大约一年的 GPU 内核优化学习经验进入比赛,主要在 Triton 中,而且他说自己从未在该领域做过职业工作

Sankalp 描述了一个结合代码生成、基准测试、分析和排行榜提交的优化循环。他选择架构并解读结果。下面的工作流细节来自他的叙述;独立可测的结果是他的排行榜时间和名次。

一个为代理设计的基准

GPU Mode,一个用于 GPU 编程竞赛和内核基准的平台,以及 Core Automation 要求参赛者 实现批处理、方形紧凑 Householder QR 分解。每次提交必须接受 FP32 CUDA 矩阵并返回与 PyTorch 的 torch.geqrf 使用的相同紧凑表示:包含上三角 R 结果和存储的 Householder 向量的 H 矩阵,以及 tau 向量中的反射系数。

检查程序重建了 Q、测试了正交性和残差误差,并按在不同矩阵形状和输入条件下的几何平均运行时间对正确提交进行排名。该工作负载覆盖了高达 4,096 × 4,096 的矩阵,包括条件数较差的困难情况。参赛者可以在内部使用更低精度,但其输出仍必须通过以 FP32 风格进行的检查。

这一组合为 Codex 提供了一个快速、定量的进展定义。GPU Mode 的 popcorn command-line interface 允许代理测试、基准和提交候选方案。按形状的计时显示了每项更改是在哪些地方有所帮助或受损,而分析则提供了另一层证据。

根据 Sankalp 的叙述,他维护了一个包含操作说明、问题陈述、实验记录和带时间戳的提交日志的 AGENTS.md 文件。后来的 Codex 会话可以读取哪些方法已经失败,从而不必重新发现它们。他还给代理设定了数值目标,并让一些优化运行在夜间继续,每隔两到三小时检查一次,询问它改变了什么以及正在哪个瓶颈上努力。

最终该目录包含了 560 个命名的提交变体、119 个 Modal B200 探针和比较脚本,以及 68 份每个实验的文档。据 Sankalp 称,Modal(一家云基础设施提供商)为比赛提供了 GPU 积分。记录保存了失败的方法、分析结果和提交历史,以便后来的 Codex 会话能在早期工作基础上继续构建。

将串行工作变成矩阵乘法

他对技术工作的叙述,Sankalp 首先使用 Claude 和教学材料来理解 Householder QR。他确定了带尾部 WY 更新的分块 Householder 设计,然后使用分析、正确性检查和反复的排行榜提交来指导进一步优化。

核心性能问题是顺序依赖。传统的 Householder 实现按列顺序处理,每个反射器依赖于前一步产生的矩阵。这会把大量工作留在较慢的矩阵-向量操作中,而 B200 的张量核则处于等待状态。

分块方法将串行工作限制在狭窄的面板内,并将更大的尾部更新转换为矩阵乘法,在那里 GPU 有更多并行工作可做。Sankalp 报告说,他在采用该架构的第一天就在权重较重的 512 × 512 情况上达到了大约 5,000 微秒的成绩。

进一步的提升需要跨栈的更改。他的提交历史 经过了自定义 Triton 面板、分组 WY 更新、CUDA 图重放、融合布局组装、固定形状内核专用化以及针对最大矩阵的自定义 Cholesky 路径。Sankalp 报告称被跟踪的整表结果从 108,803 微秒下降到大约 1,805 微秒。这一过程与他用于 232 倍计算的大约 419,000 微秒的 PyTorch 基线是分开的。

Sankalp 对后期分析的描述 表示启动开销和面板处理占主导地位,而内核很少受原始内存带宽或计算能力的限制。他的优化表记录了减少启动、融合归约、固定形状专用化、合并 V/T 组装以及移除拷贝、连接和临时表示的工作。

Sankalp 写道他给予 Codex 对 Modal 分析的访问权限,并使用了 Torch profiling、NVIDIA Nsight Systems 和后来的 NCU。这些都是自我报告的方法细节。在他的叙述中,Codex 实现并测量了更改,而他则检查结果、设定目标并重新定向搜索。

人类改变了搜索过程

根据 Sankalp 的叙述,当结果降到 3,000 微秒以下后,优化变得更难。他描述 Codex 花更多时间在参数调整和它已经尝试过的想法的小变体上。Sankalp 通过改变实验争夺注意力的方式作出回应。

指示 Codex 保持三到五个候选族的光束,而不是保留单一的现任方案并拒绝每一个较慢的尝试。候选者包括接近当前最好结果的更改、有前途的接近未中选者以及风险较高的结构性想法。

同一叙述中,Sankalp 描述了使用无头的 claude -p 调用作为顾问,指派子代理去寻找优化想法,并在新运行前清除积累的上下文。他的 AGENTS.md 指示将超时视为不确定、保留带时间戳的日志、要求有完整输出作为证据并记录为何某些候选被提升或拒绝。该文件还警告在没有实质性变化的情况下不要重复已被拒绝的想法。

Codex 提供了持久性、代码生成和庞大的搜索预算。Sankalp 选择了架构、修改了反馈循环、识别出重复行为并改变了实验策略。

他对前10名参赛作品的评述 中,Sankalp 表示,更快的竞争者使用了输入分布检测器、删除了更多库调用、实现了自定义的三角矩阵求逆,并更积极地处理低精度数据。这些是 Sankalp 的观察结论,而非独立审计的发现。在 同一篇文章的结论 中,他写道他“无法在 B200 上使用 tcgen05 指令”来进一步利用其张量核心。

GPU Mode 将竞赛变成代理测试平台

在一篇关于 AI 编写系统代码的 Core Automation 文章 中,GPU Mode 联合创始人 Mark Saroufim 写道,他与联合创始人 Andreas Kopf 在 2023 年末发起了一个 GPU 编程读书会。该读书会后来扩展为一个 YouTube 频道、内核竞赛和黑客马拉松。

Saroufim 用短语 "数据匮乏" 来形容在网上可用于训练代码模型的 GPU 内核材料的短缺。GPU Mode 的竞赛生成了可衡量的示例,展示了代理和人类在这一稀缺的系统问题类别上进行工作的情况。

Sankalp 获得第12名的成绩为这项活动提供了一个扎实的视角。一名新人使用代理在一个狭窄且可测量的任务上冲击竞争榜单的前列。该工作仍然要求他学习问题、构建实验、挑战重复性行为并解释性能分析结果。他的结果显示了一个能够在每次运行后产出可靠证据的评估者的价值。

Reader comments

Conversation for this story loads after sign-in.