AI 智能体在一个下午内将 CUDA C++ 注意力内核移植并优化至 CuTeDSL，比原版快 1.2 倍，比 cuDNN 版本快 2.8 倍。
AI translation, not an official translation. Refer to the original for technical details.
Adapted from @maharshii# 用智能体进行 GPU 内核移植 几周前，我的 AI 智能体在一个下午内将一个 CUDA C++ 注意力内核移植到了 CuTeDSL，并完成了优化。最终的内核比原始版本快 1.2 倍，比 cuDNN 版本快 2.8 倍。我认为"移植"是智能体内核开发的一个极具价值的应用场景。 我所说的"移植"包含两层含义： - 不同 GPU 版本之间的移植：当你需要将一个在某种 GPU 上运行的内核转换为在另一种 GPU 上运行，并利用其特有的架构功能时。例如，从 H100（Hopper 架构）转换到 B200（Blackwell 架构）。 - 不同抽象层级之间的移植：当你需要将一个内核从某一抽象层级转换到另一层级时。例如，从底层的 NVIDIA CuTe C++ 转换为像 Triton 这样的高层领域特定语言（DSL）。 智能体内核开发在上述两种场景中似乎都表现出色。下面我来谈谈将 Sage Attention 3 内核从 CuTe C++ 移植并优化到 CuTeDSL 的亲身经历。 # Sage Attention 3 为了在 NVIDIA 消费级 Blackwell GPU 上提升注意力操作的效率，Sage Attention 3 利用量化技术和低位张量核心，加速注意力操作中的 GEMM（矩阵乘法）运算。RTX 5090 和 6000 PRO 等 GPU 上的 FP4 张量核心，相比 FP16 张量核心性能大幅提升。因此，我们必须在注意力内核的 GEMM 运算中巧妙地加以利用。论文链接：https://arxiv.org/pdf/2505.11594 Sage Attention 3 还提出了一种方法，用以缓解将数值量化为 FP4 e2m1 时造成的严重精度损失——该格式仅有 15 个可表示的值。 FP4 注意力机制面临两大主要挑战： - 逐张量（per-tensor）和逐词元（per-token）量化方案均不足以保留模型精度。因此需要将量化分组大小限制为 1x16，即 NVFP4 格式。 - 注意力图 P 的值主要集中在 [0, 1] 区间内的较小数值。直接量化为 FP4 时，这些值会迫使缩放因子落入极窄的动态范围内。硬件要求量化因子采用 FP8 e4m3 数据类型，这导致以 FP8 表示这些缩放因子时出现显著的精度损失。 针对第二个挑战，论文提出了一种两级量化方法。该方法首先通过逐词元（即每行）量化，将每个词元的取值范围归一化到 [0, 448 x 6]，以充分利用 FP8 e4m3 缩放因子的表示范围，再进一步应用 FP4 微缩放量化。 # 移植过程 Sage Attention 3 的原始实现使用 CUDA C++ 编写，涉及约 20 个文件，充斥着模板代码和 NVIDIA 的 CuTe 抽象。仅仅理解整个代码库就需要耗费我数天时间。代码地址：https://github.com/thu-ml/SageAttention 我希望将其移植到 CuTeDSL，以获得快速 JIT 编译的优势，并实现更便捷的维护，无需再深陷 C++ 模板的泥沼。不仅如此，我还希望通过迭代优化内核，榨取 C++ 实现可能遗漏的最后一点性能。 经过优化的 CuTeDSL 移植版本只有 4 个主要文件。其余文件包括实际的测试框架（用于验证正确性和基准测试），以及一些由 Claude 生成的探针文件，用于检查 CuTeDSL 特有而原始 C++ 实现中不存在的 API 行为。 移植的整体流程与我之前一篇文章中所描述的大体相似：https://x.com/maharshii/status/2086442755748970889 对于上下文目录，我只给了 Claude 访问 NVIDIA 的 cutlass 和 SageAttention 代码库的权限，以及一个本地安装的 CuTeDSL 版本。我还要求它维护一个 INDEX.md 文件，用于说明上下文中每个目录的用途，大致如下所示： 在移植过程中，我要求我的智能体将遇到的任何 CuTeDSL 注意事项和陷阱记录在 README.md 文件中。这有助于它保留与 DSL 相关的知识，避免遗忘。 # 测试 测试工具文件沿用了我上一篇文章中的方案，并根据移植过程的阶段进行了迭代改进。它运行一组按顺序排列的测试，使得一旦出现失败，便能快速定位问题。测试从主机端代码开始，随后进入内核端测试，最终与参考实现进行对比。 # PTX/SASS 转储 工具文件还会让智能体转储 PTX、SASS 和 CUBIN 文件，从而允许它检查寄存器溢出、指令数量、向量化以及其他优化机会。它会创建 asm_info.json 文件，将从转储文件中提取的内核相关信息写入其中，大致如下所示： # 优化日志 智能体还在同一个 README.md 文件中持续记录了它尝试过的优化操作。观察到收益最大的优化来自于调整 TMA/MMA warp 的寄存器数量，以及智能体相较于 C++ 实现所做的一些刻意改动： README 文件中的优化日志表格如下： 这些实现上的差异使得 CuTeDSL 内核比原始 C++ 实现快了 1.2 倍。 # 结语 总体而言，我认为快速移植内核是智能体内核开发的一个非常好的应用场景。在整个移植过程中，我没有亲自编写一行 CuTeDSL 代码，却依然获得了比原始实现更快的内核，将原本需要数天的工作压缩到了几个小时之内。过程中仍需要一些来回沟通，以引导智能体走向更好的路径，这使得开发效率高于让智能体完全自主探索的方式。不过，我相信这一差距——即人类仍需介入的部分——会随着模型不断进步而持续缩小。
