LLM2D
KernelBench: LLMs能编写高效的GPU内核吗?
KernelBench: Can LLMs Write Efficient GPU Kernels?
作者: Anne Ouyang, Simon Guo, Simran Arora, Alex L. Zhang, William Hu, Christopher R\'e, Azalia Mirhoseini
发布日期: 2/18/2025
arXiv ID: oai:arXiv.org:2502.10517v1

摘要

arXiv:2502.10517v1 交叉类型:跨领域 摘要:高效的GPU内核对于构建高性能的机器学习架构至关重要,但编写它们是一项耗时的挑战,需要大量的专业知识;因此,我们探索使用语言模型(LMs)来自动化内核生成。我们介绍了KernelBench,这是一个开源框架,用于评估LMs在250个精心选择的PyTorch机器学习工作负载上的能力,以编写快速且正确的内核。KernelBench代表了一个实际的工程环境,通过对引入的基准的改进可以直接转化为更快的实际内核。我们引入了一个新的评估指标fast_p,该指标衡量生成内核中功能正确且比基线快于可调整阈值p的百分比。我们在各种最先进的模型和测试时间方法上进行的实验显示,前沿推理模型在一开始表现最佳,但总体上仍然不尽如人意,在不到20%的情况下能够匹配PyTorch基线。尽管我们展示了通过在迭代细化过程中利用执行和剖析反馈可以提高结果,但KernelBench仍然是一个具有挑战性的基准,随着加速阈值p的提高,其难度也会增加。