AscendKernelGen：面向神经处理器的基于大语言模型的内核生成系统性研究
人工智能
2026-04-20 v2 机器学习
摘要
为满足计算效率日益增长的需求，神经处理器（NPU）已成为现代人工智能基础设施中的关键组件。然而，要实现其潜力的充分释放，需开发针对特定硬件的域特定语言（DSL）的高性能计算内核，这需要深厚的硬件专业知识且工作量大。虽然大型语言模型（LLM）在通用代码生成方面显示出前景，但在NPU领域面临严格约束和训练数据稀缺的挑战。我们的初步研究表明，最先进的通用LLM难以为Ascend NPU生成功能复杂的内核，几乎成功率为零。为解决这些挑战，我们提出AscendKernelGen，一个集成生成-评估的NPU内核开发框架。我们引入Ascend-CoT，包含从真实世界内核实现中推导的链步理由（chain-of-thought reasoning）的高质量数据集，并通过监督微调和强化学习训练KernelGen-LM模型。此外，我们设计了NPUKernelBench，一个综合性基准，用于评估不同复杂度水平下的编译、正确性和性能。实验结果表明，我们的方法显著缩小了通用LLM与硬件特定编码之间的差距。具体而言，复杂等级-2内核的编译成功率从0%提升至95.5%（Pass@10），功能正确性达64.3%，而基线模型完全失败。这些结果凸显了领域特定推理和严格评估在加速器感知代码生成中的关键作用。AscendKernelGen 可在 https://huggingface.co/AscendKernelGen 和 https://github.com/weich97/NPUKernelBench 上获取。
引用
@article{arxiv.2601.07160,
title = {AscendKernelGen: A Systematic Study of LLM-Based Kernel Generation for Neural Processing Units},
author = {Xinzi Cao and Jianyang Zhai and Pengfei Li and Zhiheng Hu and Cen Yan and Bingxu Mu and Guanghuan Fang and Bin She and Jiayu Li and Yihan Su and Dongyang Tao and Xiansong Huang and Fan Xu and Feidiao Yang and Yao Lu and Chang-Dong Wang and Yutong Lu and Weicheng Xue and Bin Zhou and Yonghong Tian},
journal= {arXiv preprint arXiv:2601.07160},
year = {2026}
}
备注
33 pages,7 figures,16 tables