[Paper] PEAK:一种基于自然语言转换的 GPU 内核性能工程 AI 助手

发布: (2025年12月22日 GMT+8 12:15)
7 min read
原文: arXiv

Source: arXiv - 2512.19018v1

概述

本文介绍了 PEAK,一个利用大语言模型(LLMs)自动转换和优化 GPU 核心代码的 AI 助手。通过将低层次的性能调优表达为自然语言指令,PEAK 能够生成、验证并在多个 GPU 后端(CUDA、HIP、HLSL)上基准测试优化后的内核,性能可与手工调优的供应商库相媲美。

关键贡献

  • 自然语言转换管线 – 将迭代的 GPU 优化编码为普通英文提示,供 LLM 执行。
  • 模块化、可扩展的基础设施 – 处理代码生成、正确性验证以及任何 GPU 后端的性能测量。
  • 跨后端支持 – 在 CUDA、AMD 的 HIP 和 DirectX HLSL 上演示,表明该方法与硬件无关。
  • 实证评估 – 为矩阵乘法内核实现了 16 种转换;结果与供应商提供的库相匹配或超出,并且在 HLSL 上达到了文档中记录的 FLOPS 上限。
  • 研究平台 – 使得能够系统地研究 LLM 在底层代码上的行为、错误模式以及跨优化序列的性能轨迹。

方法论

  1. 用自然语言定义转换 – 每个优化(例如,“使用 32×32 块对循环进行平铺”, “将最内层循环展开三次”)都写成简洁的 英文指令
  2. 提示 LLM – 将该指令连同当前的 kernel 源码一起输入大型语言模型(例如 GPT‑4),模型会返回修改后的源代码。
  3. 验证与基准测试 – 自动化测试框架会为目标后端编译生成的 kernel,运行正确性测试套件,并测量运行时间 / FLOPS。
  4. 迭代 – 成功的转换会被串联;若出现失败则触发回退或重新提示。
  5. 可扩展性 – 通过提供相应的编译器包装器和自然语言规范即可添加新后端或新转换——无需对核心 LLM 逻辑进行任何修改。

结果与发现

后端基准(朴素内核)PEAK 优化后供应商库(如有)
CUDA比 cuBLAS 慢 1.2×≈ 0.95× cuBLAScuBLAS(参考)
HIP比 rocBLAS 慢 1.4×≈ 1.0× rocBLASrocBLAS(参考)
HLSL比理论峰值慢 2.3×≈ 1.0× 硬件 FLOPS 上限无官方库
  • 正确性:所有生成的内核均通过了提供的测试套件;基础设施捕获并拒绝了引入细微错误的转换。
  • 错误模式:大多数 LLM 错误是语法错误(缺少分号)或语义错误(使用未定义的变量)。通过提示工程和后处理过滤,将错误率降低至 <5 % 的尝试。
  • 优化轨迹:性能提升并非单调的;某些早期的转换(例如激进的展开)可能会削弱后续的 tiling 步骤,凸显了反馈回路的必要性。

Practical Implications

  • Accelerated performance engineering – 开发者可以用自然语言描述期望的优化,并在几分钟内获得可运行、已基准测试的内核,缩短传统上需要数天或数周的手动调优周期。
  • Cross‑platform portability – 由于转换对后端无关,一套自然语言规范即可为 NVIDIA、AMD 和 DirectX GPU 生成调优后的内核,简化多厂商代码库。
  • AI‑augmented CI pipelines – PEAK 可以集成到持续集成中,在内核变更时自动建议或应用性能改进。
  • Rapid prototyping for emerging hardware – 当出现新 GPU 架构时,工程师只需更新后端编译器包装器;相同的自然语言转换即可重新运行,发现最佳设置,无需手工编写新内核。
  • Foundation for autonomous agents – PEAK 的“plug‑and‑play” 设计使更高层次的 AI 代理能够在无人干预的情况下驱动端到端的内核优化,为自优化库打开了大门。

限制与未来工作

  • LLM 依赖 – 生成代码的质量取决于底层模型;预计更新、更强大的模型将降低错误率。
  • 转换表达能力 – 复杂的、非局部的优化(例如算法重构)仍然难以通过简单的自然语言提示捕获。
  • 验证的可扩展性 – 对大型 kernel 进行完整的正确性测试可能耗时;未来工作可能引入静态分析或符号执行以加速验证。
  • 更广泛的 kernel 家族 – 本研究聚焦于矩阵乘法;将该方法扩展到不规则工作负载(例如图 kernel、稀疏线性代数)是一个开放的研究方向。

PEAK 表明,自然语言可以作为高层性能意图与底层 GPU 代码之间的实用桥梁,使 LLM 成为真正的性能工程合作伙伴。

作者

  • Muhammad Usman Tariq
  • Abhinav Jangda
  • Angelica Moreira
  • Madan Musuvathi
  • Tyler Sorensen

论文信息

  • arXiv ID: 2512.19018v1
  • 分类: cs.SE
  • 出版日期: 2025年12月22日
  • PDF: Download PDF
Back to Blog

相关文章

阅读更多 »