[论文阅读] 人工智能+ | GPU Kernel Scientist研究：LLM驱动的GPU内核迭代优化框架

zhangjiaoshou_

1352人浏览 · 2025-06-30 06:00:00

zhangjiaoshou_ · 2025-06-30 06:00:00 发布

GPU Kernel Scientist研究：LLM驱动的GPU内核迭代优化框架

论文标题：GPU Kernel Scientist: An LLM-Driven Framework for Iterative Kernel Optimization

arXiv:2506.20807 (cross-list from cs.LG)
GPU Kernel Scientist: An LLM-Driven Framework for Iterative Kernel Optimization
Martin Andrews, Sam Witteveen
Comments: 4 page paper plus Appendices. Accepted to the ES-FoMo “Efficient Systems for Foundation Models” workshop at ICML 2025
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Performance (cs.PF); Software Engineering (cs.SE)

主要作者及单位信息

Martin Andrews与Sam Witteveen均来自新加坡研究机构，通信作者为Martin Andrews
该研究成果于2025年发表在ICML的ES-FoMo工作坊，依托Google AI开发者计划支持，使用Gemini模型与Google Cloud GPU资源

研究背景：GPU优化的"荒野求生"困境

想象一下，你要在一个没有地图、没有路标甚至语言不通的陌生城市开车，而传统司机依赖的"导航系统"（丰富的文档和工具）在此完全失灵——这正是当前新GPU架构优化的真实写照。在AI加速领域，当开发者面对AMD MI300等新型GPU时，传统优化模式面临三重困境：

专家知识壁垒：优化GPU内核需要像"老司机"一样熟悉硬件架构细节，比如内存布局、指令调度等，普通开发者难以驾驭
工具链缺失：新架构缺乏成熟的profiling工具，就像开车时没有速度表和油量表，只能通过"感觉"（整体计时）判断性能
文档荒漠：AMD等非CUDA平台的开发资料如同未开垦的荒野，连专家也常需"摸着石头过河"

典型案例：某团队尝试优化MI300矩阵乘法时，因缺乏低精度操作示例，连基础正确的代码都难以写出，最终不得不让LLM通过"试错实验"自探硬件特性

在这里插入图片描述

创新点：LLM如何化身"内核优化科学家"

1. 三阶段闭环：让LLM像科学家一样做实验

突破传统"代码生成器"模式，构建完整科研闭环：

进化选择器：LLM像实验室主任，从历史代码"实验记录"中挑选最有潜力的版本作为新实验基础
实验设计器：化身研究员，提出10个"研究方向"并设计5组实验，预估性能提升幅度
内核编写器：充当工程师，根据实验方案生成可编译的HIP代码，甚至能处理矩阵核心等复杂硬件接口

2. 跨平台知识迁移：让CUDA经验"翻译"到AMD

当AMD文档不足时，LLM如同"语言翻译官"：

将Nvidia Tensor Core的优化经验（如矩阵分块策略）"翻译"为AMD Matrix Core可用的HIP代码
自动总结CUDA文档中的最佳实践，如内存对齐规则，适配到HIP编程模型

3. 黑盒优化能力：在"蒙眼"状态下精准调优

面对仅有整体计时数据的限制（无profiling工具），LLM展现独特策略：

通过对比实验推断性能影响，例如修改内存加载方式后观察整体耗时变化
设计"控制变量"实验，每次只改变一个优化因子，逐步定位关键性能瓶颈

研究方法和思路：拆解LLM优化的"三步实验法"

第一阶段：选种子（LLM进化选择器）

从"代码种群"中选取两个样本：
- “基础代码”：当前最优版本，如经52次迭代的00052号代码
- “参考代码”：来自不同优化路径，如00046号代码，在特定场景表现更优
LLM基于基准测试结果（6种矩阵尺寸的计时数据），像选种专家一样判断：“00052号平均性能最低，选作基础；00046号在首个测试场景更快，可提供对比思路”

第二阶段：做计划（LLM实验设计器）

brainstorm阶段：生成10个"研究方向"，例如：

“缓解LDS银行冲突”：调整共享内存数据布局避免访问冲突
“优化矩阵核心布局”：让数据格式匹配rocWMMA库要求

精细化设计：从10个方向中挑出5个生成具体实验方案，包含：
- 操作细则：如"修改内存加载函数，将A矩阵从行主序转为列主序"
- 性能预测：预计提升15%-40%
- 创新评分：如85分（满分100）
方案筛选：选3个最具潜力的实验：“最创新的”、“性能上限最高的”、“性能下限最稳的”

第三阶段：写代码（LLM内核编写器）

给LLM"实验手册"：
- 任务描述（如矩阵乘法）与PyTorch参考实现
- 历史"实验报告"：之前迭代的代码版本与性能数据
- “硬件说明书”：通过自探生成的"发现文档"，如矩阵核心内存布局注意事项
LLM生成HIP代码，需满足：
- 语法正确：经测试90%以上可直接编译
- 优化实现：包含共享内存双缓冲、混合精度计算（FP8输入/FP32累加/BF16输出）等高级技术
输出物：新HIP代码+优化说明，如"本版本使用32x32x16矩阵核心配置，提升算力利用率"

实验验证：在"约束考场"中测试能力

考场规则：AMD开发者挑战赛平台限制
- 仅提供整体计时数据（18种随机种子），无profiling工具
- 必须顺序测试，不能并行优化
应对策略：
- LLM通过"对比实验"推断优化效果：如修改内存加载方式后，对比前后计时差异
- 自探硬件特性：通过编译错误和运行结果，总结出"矩阵核心内存块布局规则"

主要贡献：给领域带来的三大"生产力革命"

1. 降低优化门槛：让"新手"也能做专家级优化

传统模式：需5年以上GPU优化经验的专家
新框架：普通开发者通过LLM辅助，也能生成媲美专家的HIP代码
案例：框架自动生成的代码在MI300上实现6倍于原始HIP版本的性能

2. 加速新架构落地：为硬件厂商"铺路"

当AMD等厂商推出新GPU时，开发者常因文档不足推迟适配
框架通过LLM自探和知识迁移，可在硬件发布初期快速生成优化代码，缩短生态建设周期

3. 开创"AI优化AI硬件"新范式

传统优化：人读文档→写代码→测性能→调优
新范式：LLM读文档+自探→生成代码→分析计时→自动迭代，形成闭环
意义：为AI芯片自动化优化提供可复用框架，尤其适合快速迭代的硬件环境

思维导图

在这里插入图片描述

详细总结

一、研究背景与问题

GPU内核优化挑战：传统优化需深厚架构知识、大量分析和迭代实验，新架构（如AMD MI300）因文档稀缺和工具不足，挑战更显著。
现有方案局限：进化方法依赖树结构编码，传统调优框架（如OpenTuner）侧重超参数调整，无法应对架构级优化。

二、GPU Kernel Scientist框架

三阶段LLM工作流程
- LLM进化选择器：从内核种群中选“基础代码”和“参考代码”，基于基准结果（6种MxKxN配置）判断优化路径。
- LLM实验设计器：生成10个研究方向（如LDS银行冲突缓解、矩阵核心布局优化）和5个实验计划，评估性能提升范围（如15%-40%）和创新性（最高85分）。
- LLM内核编写器：基于参考代码和实验计划生成HIP代码，支持rocWMMA库、矩阵核心（32x32x16 MFMA）、共享内存双缓冲等优化。

关键技术细节

技术点	实现方式	作用
知识迁移	LLM将CUDA文档（如Tensor Core用法）翻译为HIP实现	弥补AMD文档不足
混合精度	输入FP8、累加FP32、输出BF16	平衡精度与性能
共享内存优化	ping-pong双缓冲+动态重用途	重叠计算与数据传输，缓存缩放因子

三、实验与发现

约束条件：仅用竞赛平台的计时数据（18种子），无profiling工具，顺序测试限制吞吐量。
应对策略
- 文档缺失：LLM通过实验自探硬件特性，生成“发现文档”。
- 无profiling：LLM通过对比实验推断性能影响。
- 人力不足：全流程自动化，LLM自主决策优化方向。
定性成果：LLM生成的HIP代码可编译且性能提升显著，证明LLM在复杂优化中的可行性。