用大模型自动生成高效GPU代码，实测效果仍不及人工。
KernelBench: Can LLMs Write Efficient GPU Kernels?
- 构建250个真实机器学习任务的测试集，评估大模型写GPU核函数能力。
- 顶尖模型仅在不足20%任务中生成既正确又比基线快的代码。
- 引入可调阈值的速度提升评价指标，推动模型持续优化迭代。
高效GPU内核对构建高性能机器学习架构至关重要，但编写过程耗时且需深厚专业知识；为此，我们探索使用语言模型（LMs）自动化内核生成。本文提出KernelBench，一个开源框架，用于在250个精心挑选的PyTorch机器学习工作负载上评估语言模型生成快速且正确内核的能力。KernelBench模拟真实工程环境，其进展直接转化为实际更快的内核。我们引入新评价指标fast_p，衡量生成内核在功能正确前提下，速度提升超过可调阈值p的比例。实验涵盖多种前沿模型与测试时方法，结果显示，先进推理模型虽开箱即用表现最佳，但仍整体不足，仅在少于20%情况下匹配PyTorch基线。尽管通过执行与性能分析反馈进行迭代优化可提升结果，KernelBench仍是具有挑战性的基准，难度随速度提升阈值p升高而增加。
原文摘要 · Abstract (English)
Efficient GPU kernels are crucial for building performant machine learning architectures, but writing them is a time-consuming challenge that requires significant expertise; therefore, we explore using language models (LMs) to automate kernel generation. We introduce KernelBench, an open-source framework for evaluating LMs' ability to write fast and correct kernels on a suite of 250 carefully selected PyTorch ML workloads. KernelBench represents a real-world engineering environment and making progress on the introduced benchmark directly translates to faster practical kernels. We introduce a new evaluation metric fast_p, which measures the percentage of generated kernels that are functionally correct and offer a speedup greater than an adjustable threshold p over baseline. Our experiments across various state-of-the-art models and test-time methods show that frontier reasoning models perform the best out of the box but still fall short overall, matching the PyTorch baseline in less than 20% of the cases. While we show that results can improve by leveraging execution and profiling feedback during iterative refinement, KernelBench remains a challenging benchmark, with its difficulty increasing as we raise speedup threshold p.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。