Source-linked AI summary
AscendKernelGen: A Systematic Study of LLM-Based Kernel Generation for Neural Processing Units
Xinzi Cao, Jianyang Zhai, Pengfei Li, Zhiheng Hu, Cen Yan, Bingxu Mu, Guanghuan Fang, Bin She, Jiayu Li, Yihan Su, Dongyang Tao, Xiansong Huang, Fan Xu, Feidiao Yang, Yao Lu, Chang-Dong Wang, Yutong Lu, Weicheng Xue, Bin Zhou, Yonghong Tian
TL;DR
NPU kernel generation is difficult because vendor-specific DSLs require hardware expertise and general-purpose LLMs lack suitable domain knowledge. AscendKernelGen combines reasoning-oriented data, domain-adaptive training, and hardware-grounded evaluation, producing substantial improvements on complex AscendC kernels while leaving higher-level kernel support and performance optimization as future work.
Problem
NPU kernel development requires specialized DSL and hardware knowledge, while general-purpose LLMs struggle with strict constraints and scarce domain-specific training data.
Method
AscendKernelGen combines Ascend-CoT reasoning data, KernelGen-LM domain-adaptive supervised fine-tuning and execution-guided reinforcement learning, and NPUKernelBench evaluation.
Results
Domain-adaptive training substantially improves compilation success, functional correctness, and performance-related evaluation compared with general-purpose code models on hardware-specific kernels.
Takeaways & Limitations
Domain-specific reasoning supervision and rigorous execution-based evaluation support automated generation of complex accelerator-aware kernels that general-purpose models previously could not reliably produce.
Takeaways & Limitations
Level-3 kernels remain challenging, and the current study does not yet prioritize latency minimization and resource utilization through performance-aware reward models.
Abstract
from arXiv · showhide
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.
1 Introduction
NPU kernels require specialized hardware knowledge, yet general-purpose LLMs perform poorly on AscendC because vendor-specific constraints and training data are scarce. AscendKernelGen addresses this gap with domain-specific reasoning data, adaptive training, and hardware-grounded evaluation.
- Motivation: High-performance NPU kernels require vendor-specific DSL expertise, memory management, tiling, asynchronous pipelines, and explicit use of vector and matrix units.These requirements make kernel development labor-intensive and hardware-dependent.
- Motivation: Around 0% execution success on complex L2/L3 kernels shows that zero-shot general-purpose LLMs are effectively unusable for non-trivial NPU kernel development.The failures include hallucinated APIs and misuse of hardware interfaces.
- Contributions: AscendKernelGen integrates domain-specific reasoning data, domain-adaptive post-training, and a benchmark covering compilation, functional correctness, and performance.The framework targets both model adaptation and rigorous hardware-grounded assessment.
- Contributions: Ascend-CoT is curated from real-world kernels and annotates pipeline construction, synchronization logic, and arithmetic reasoning needed for low-level NPU programming.The dataset is designed to expose structured expert reasoning rather than only kernel surface form.
- Evaluation: NPUKernelBench addresses the lack of a comprehensive NPU evaluation framework by jointly assessing kernel-only optimization and full host-device integration.Its scope includes compilation, correctness, performance, and static-shape versus dynamic-shape behavior.
- Results: Domain-adaptive post-training substantially improves compilation success and functional correctness across difficulty levels, especially for complex kernels nearly unsolvable in zero-shot settings.The result supports domain-specific reasoning supervision as a key component of hardware-specific kernel generation.
3 Programming Abstractions and Constraints for NPU Kernel Generation
Low-level NPU kernels combine data partitioning, asynchronous pipelines, synchronization, and explicit memory semantics in statically structured programs. Correct generation therefore requires maintaining long-range dependencies and precise boundary, layout, and execution reasoning.
- Kernel Abstraction: A low-level NPU kernel jointly specifies global data partitioning, asynchronous pipeline stages, and explicit synchronization under a replicated execution template.Each kernel instance operates on a distinct global-data slice determined by a logical block index.
- Kernel Abstraction: Kernel programs explicitly compute memory offsets, valid ranges, and boundary conditions instead of relying on implicit indexing.These calculations are required because each processing unit handles a separate data slice.
- Execution Constraints: Asynchronous loading, computation, and write-back must be manually coordinated through producer-consumer synchronization primitives.The runtime does not manage this overlap; the kernel encodes a statically defined execution schedule.
- Execution Constraints: Low-level interfaces require explicit memory access patterns, strides, masks, and synchronization events, tightly coupling data movement, arithmetic, and control flow.These decisions are not automatically inferred or optimized.
- LLM Challenges: Correct generation depends on long-range consistency across auxiliary parameters, pipeline stages, synchronization order, boundary arithmetic, and physical layout transitions.Small errors can cause deadlocks, data hazards, incorrect outputs, or invalid memory accesses.
4 System Overview
AKGen is a unified framework that couples domain-specific data, kernel generation, and structured hardware evaluation in a closed loop. Its components are designed to expose how models acquire and apply low-level NPU programming knowledge.
- System Overview: AKGen integrates a domain-specific reasoning dataset, a kernel generation model, and a structured evaluation benchmark into a closed generation-evaluation loop.The framework supports systematic analysis under realistic hardware programming constraints.
- Generation Component: Ascend-CoT captures structured reasoning for AscendC development, including pipeline construction, synchronization, arithmetic, and memory movement.Its data comes from real-world operator implementations and documentation.
- Generation Component: KernelGen-LM is domain-adaptively post-trained to reason about data tiling, asynchronous pipeline orchestration, and low-level interface usage.The objective targets execution dependencies rather than only the surface syntax of kernels.
- Evaluation Component: NPUKernelBench evaluates tasks from simple element-wise operators to complex fused kernels through compilation, reference-based correctness, and runtime performance checks.The staged pipeline isolates distinct failure modes beyond surface-level code validity.
- Data Construction: The dataset combines documentation-based reasoning, kernel-centric CoT, and general reasoning chains; its raw corpus contains 83,916 samples.The three components balance specialized knowledge with generalization capability.
- Data Construction: Complex operators are decomposed into host-kernel pairs so supervision can jointly represent tiling computation, device execution, and cross-boundary consistency.Only samples verified against reference tiling configurations are retained.
6 LLM Fine-Tuning for NPU Kernels
KernelGen-LM uses supervised correction followed by execution-guided reinforcement learning to generate AscendC kernels that satisfy both compilation and numerical constraints. The staged design removes invalid candidates before RL and focuses later optimization on valid implementations.
- Training Strategy: KernelGen-LM combines NPU-aware supervised fine-tuning with reinforcement learning driven by execution-based correctness signals.SFT teaches syntax, APIs, and kernel structure; RL further encourages compilable and numerically correct kernels.
- Error-Derived SFT: Error-derived supervision corrects API misuse using compiler logs, kernel context, and official documentation.The model identifies failure causes and produces corrected implementations for semantically invalid API usage.
- Error-Derived SFT: Ground-truth-guided reconstruction addresses kernels that compile but fail numerical verification by pairing them with correct implementations.The resulting supervision targets memory staging, accumulation order, and host-device tiling inconsistencies.
- Training Strategy: Error-derived SFT reduces invalid candidates before reinforcement learning, stabilizing policy optimization and improving sample efficiency.It narrows the search space to reliably executable kernels.
- Reinforcement Learning: The RL stage samples executable candidates and forms preference relations from numerical accuracy and execution correctness.Unlike SFT correction, RL preferentially reinforces higher-quality implementations among kernels that already pass verification.
1 Step 1: Extract Error Logs
The procedure scans build logs, extracts error blocks, classifies failures, and identifies challenging operators. It then enriches each error sample with related code and API documentation in a structured report.
- The pipeline traverses every build log and extracts error blocks from files containing error messages.
- An LLM classifies error types, while pass/fail statistics identify challenging operators.
- The statistics can optionally filter out easy cases.
- For each error sample, the pipeline retrieves related code, variable definitions, and corresponding API documentation.
- The retrieved log excerpt, code context, and API reference are merged into a structured Markdown report.
18 Step 4: SFT Data Synthesis
The SFT synthesis pipeline turns diagnostic reports into balanced, quality-filtered training instances with expert reasoning and corrected API usage. These supervised corrections prepare reinforcement learning to optimize among valid kernels rather than recover from execution failures.
- The pipeline prompts an LLM to analyze API misuse, produce expert reasoning, and propose corrected usage with explanations.
- The resulting analyses and corrections are saved as SFT training instances.
- Data balancing uses subsampling or augmentation, and low-quality or incomplete samples are removed.
- Compilation-and-precision outcomes form preference pairs that distinguish desirable generations from kernels that compile but fail precision tests.
- Supervised correction narrows the policy search space, allowing reinforcement learning to focus on fine-grained optimization among valid kernels.
7 The Evaluation Sub-system: NPUKernelBench
NPUKernelBench is presented as an end-to-end framework for evaluating whether LLMs generate correct and efficient NPU kernels. The section introduces its overview and the construction of a reasoning- and correction-enhanced SFT dataset.
- NPUKernelBench evaluates LLM-generated NPU kernels for correctness and efficiency in an end-to-end framework.
- The section identifies a pipeline for constructing precision-corrected SFT data.
- The pipeline takes test logs and ground-truth implementations as input and outputs an SFT dataset containing reasoning and corrected kernels.
2 begin
This procedure scans log directories for kernels that compile but fail precision checks, then extracts their corresponding JSON metadata for further processing.
- The procedure traverses all log directories.
- It selects cases where compilation succeeds but the precision check fails.
- For each selected case, it extracts the corresponding JSON metadata.
8 begin
NPUKernelBench structures kernel evaluation by difficulty and interface complexity, then connects standardized generation with compilation, correctness scoring, and feedback-driven optimization.
- Hierarchical and Categorical Design: NPUKernelBench decomposes kernel difficulty into algorithmic complexity and interface complexity for systematic capability analysis.Algorithmic levels range from simple element-wise operations to globally dependent or dynamically controlled kernels; interface complexity distinguishes static- and dynamic-shape tasks.
- Standardized Generation Interface: The benchmark standardizes prompts with API descriptions, host and kernel templates, and role-based formatting constraints.These components provide semantic grounding and structural guidance for generating compilable Ascend C code.
- Comprehensive Evaluation: NPUKernelBench integrates code generation, compilation, correctness verification, and performance benchmarking in an end-to-end pipeline.Generated host-side and kernel-side code is compiled and executed on NPU hardware against reference implementations.
- Correctness Evaluation: Task-level correctness treats passing test cases as the basis for full, partial, or zero functional scores.A kernel receives full score when all test cases pass, while failures reduce the score and gate further evaluation.
- Evaluation Feedback: Detailed execution logs expose intermediate failures and partial successes that support prompt refinement, supervised fine-tuning, and reward shaping.The framework therefore functions as an evaluation-driven training environment rather than only a static benchmark.
8 Experiments and Results
Performance improves consistently across training stages and sampling settings, with the largest gains on Level 2 kernels. Domain-adaptive training also improves compilation, execution, runtime efficiency, and robustness, while complexity and model scale remain important factors.
- Main Results: Performance improves consistently from base models through supervised fine-tuning and reinforcement learning, with increasingly robust behavior on challenging kernels.Results are compared across training stages and sampling budgets on NPUKernelBench.
- Compilation and Execution: 7.92% average Pass@1 for the base Qwen-32B model rises to 26.26% after SFT and 33.46% after RL.SFT supplies API usage patterns, memory initialization, and computation templates; RL adds execution-based preference signals for semantically different implementations.
- Performance Speedup: 0.60× speedup on Level 1 tasks is achieved by the base model, while SFT reaches 1.50× on Level 2 tasks and RL reaches 1.86×.RL maintains 0.61× on Level 1 tasks and refines memory ordering, synchronization placement, and accumulation structure.
- Model Scale Sensitivity: Level 3 kernels remain highly challenging, although 32B is the only model scale showing measurable compilation success on them.Larger models improve Level 1 and Level 2 performance, while Level 3 difficulty is associated with complex control flow and memory dependencies.
- Fine-tuning Strategies: Full fine-tuning outperforms LoRA, increasing mean compilation rate from 40.29% to 55.32% and execution rate from 13.55% to 22.13%.Average speedup is 0.95× with full fine-tuning versus 0.48× with LoRA.
- Error Analysis: 51.9% of failed cases are API Signature and Overload Errors, followed by Data Type and Conversion Errors at 19.8%.Other reported categories include Variable Scope and Lifetime Errors at 16.4% and Memory and Object Usage Errors at 8.1%.
9 Conclusion
Ascend KernelGen addresses NPU kernel generation by combining domain-specific reasoning data with evaluation spanning compilation, correctness, and performance. The framework improves generation of complex operators, while the study identifies Level-3 support and performance-aware optimization as remaining challenges.
- Ascend KernelGen targets automated generation of high-performance kernels for NPU architectures.
- Ascend-CoT models expert reasoning, while NPUKernelBench evaluates compilation, numerical correctness, and runtime performance on actual hardware.
- Domain-adaptive training combines reasoning-oriented supervised fine-tuning with execution-guided reinforcement learning to generate complex operators previously out of reach.
- The study focuses on kernel-level correctness and basic optimization, leaving Level-3 kernels and performance-aware rewards for future work.