用大模型生成AI芯片核心代码,成功率从0提升到95.5%
AscendKernelGen: A Systematic Study of LLM-Based Kernel Generation for Neural Processing Units
- 构建带思维链的专用数据集,训练适配昇腾芯片的代码生成模型
- 复杂核函数编译成功率从0%提升至95.5%,正确率达64.3%
- 提供完整评测框架,适合芯片研发与自动化编程研究者
为满足日益增长的计算效率需求,神经网络处理器(NPUs)已成为现代AI基础设施的关键。然而,充分发挥其性能需使用厂商特定领域语言(DSL)编写高性能计算核,这一过程要求深入硬件知识且耗时费力。尽管大语言模型(LLMs)在通用代码生成中表现良好,但在受限严格的NPU领域因训练数据稀缺而效果不佳。初步研究表明,现有先进通用LLMs在生成昇腾NPU复杂核函数时成功率接近零。为此,我们提出AscendKernelGen——一个生成与评估一体化的NPU核开发框架。引入Ascend-CoT高质量数据集,包含来自真实核实现的思维链推理;设计KernelGen-LM模型,通过监督微调与基于执行反馈的强化学习进行训练。此外,构建NPUKernelBench评测基准,全面评估编译、正确性与性能,覆盖不同复杂度层级。实验表明,该方法显著缩小了通用模型与硬件专有编程间的差距:复杂度为Level-2的核函数编译成功率由0%提升至95.5%(Pass@10),功能正确率达64.3%,而基线为完全失败。结果凸显领域专用推理与严格评估在加速器感知代码自动生成中的关键作用。AscendKernelGen已开源,可访问HuggingFace与GitHub。
原文摘要 · Abstract (English)
To meet the ever-increasing demand for computational efficiency, Neural Processing Units (NPUs) have become critical in modern AI infrastructure. However, unlocking their full potential requires developing high-performance compute kernels using vendor-specific Domain-Specific Languages (DSLs), a task that demands deep hardware expertise and is labor-intensive. While Large Language Models (LLMs) have shown promise in general code generation, they struggle with the strict constraints and scarcity of training data in the NPU domain. Our preliminary study reveals that state-of-the-art general-purpose LLMs fail to generate functional complex kernels for Ascend NPUs, yielding a near-zero success rate. To address these challenges, we propose AscendKernelGen, a generation-evaluation integrated framework for NPU kernel development. We introduce Ascend-CoT, a high-quality dataset incorporating chain-of-thought reasoning derived from real-world kernel implementations, and KernelGen-LM, a domain-adaptive model trained via supervised fine-tuning and reinforcement learning with execution feedback. Furthermore, we design NPUKernelBench, a comprehensive benchmark for assessing compilation, correctness, and performance across varying complexity levels. Experimental results demonstrate that our approach significantly bridges the gap between general LLMs and hardware-specific coding. Specifically, the compilation success rate on complex Level-2 kernels improves from 0% to 95.5% (Pass@10), while functional correctness achieves 64.3% compared to the baseline's complete failure. These results highlight the critical role of domain-specific reasoning and rigorous evaluation in automating accelerator-aware code generation. AscendKernGen is available at https://huggingface.co/AscendKernelGen and https://github.com/weich97/NPUKernelBench.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。