AscendKernelGen：基于LLM的神经处理单元算子生成系统性研究

为了满足对计算效率不断增长的需求，神经处理单元 (NPU) 已成为现代人工智能基础设施的关键。然而，要释放它们的全部潜力，需要使用供应商特定的领域特定语言 (DSL) 开发高性能计算内核，这项任务需要深厚的硬件专业知识并且是劳动密集型的。虽然大语言模型 (LLM) 在一般代码生成方面显示出了前景，但它们却面临着 NPU 领域训练数据的严格约束和稀缺问题。我们的初步研究表明，最先进的通用 LLM 无法为 Ascend NPU 生成功能复杂的内核，成功率接近于零。为了应对这些挑战，我们提出了 AscendKernelGen，一种用于 NPU 内核开发的生成评估集成框架。我们引入了 Ascend-CoT（一个包含源自现实世界内核实现的思想链推理的高质量数据集）和 KernelGen-LM（一个通过监督微调和带有执行反馈的强化学习进行训练的领域自适应模型）。此外，我们还设计了 NPUKernelBench，这是一个用于评估不同复杂程度的编译、正确性和性能的综合基准。实验结果表明，我们的方法显着弥合了通用大语言模型和特定于硬件的编码之间的差距。具体来说，与完全失败的基线相比，复杂的 Level-2 内核的编译成功率从 0% 提高到 95.5% (Pass@10)，而功能正确性达到 64.3%。这些结果凸显了特定领域推理和严格评估在自动化加速器感知代码生成中的关键作用。