AscendKernelGen: A Systematic Study of LLM-Based Kernel Generation for Neural Processing Units

TL;DR

This paper tackles the bottleneck of producing high-performance, hardware-aware kernels for Huawei Ascend NPUs by addressing the limitations of general-purpose LLMs in domain-specific code generation. It introduces AscendKernelGen, a generation-evaluation framework that couples a domain-adaptive LLM (KernelGen-LM) with Ascend-CoT data and a rigorous NPU-centered benchmark (NPUKernelBench) to jointly optimize correctness and performance. The framework’s two-stage training—domain-focused supervised fine-tuning with error-derived supervision, followed by execution-guided reinforcement learning—substantially improves compilation success and functional accuracy for complex Level-2 kernels, achieving compelling speedups close to hardware roofline references. The work demonstrates that domain-specific reasoning and systematic hardware evaluation are crucial to unlocking automated accelerator-aware kernel generation, and it lays groundwork for extending the approach to more complex operators and additional accelerator platforms in the future.

Abstract

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.

Figures (7)

• Figure 1: System overview of AscendKernelGen, depicting the data construction, LLM training, and hardware-grounded evaluation pipeline for NPU kernel generation.
• Figure 2: Error-Derived Supervision for API- and Kernel-Level Error Correction.
• Figure 3: Dual-Path Evaluation Design.
• Figure 4: Pass@1 results for representative kernels across training stages.
• Figure 5: Scaling behavior of KernelGen-LM on NPU kernel generation tasks with increasing difficulty.