AscendKernelGen: 신경망 처리 장치(NPU)를 위한 LLM 기반 커널 생성에 관한 체계적 연구

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

계산 효율성에 대한 수요가 급증함에 따라 신경망 처리 장치(NPU)는 현대 AI 인프라의 핵심 요소가 되었습니다. 그러나 NPU의 잠재력을 온전히 활용하려면 벤더별 도메인 특화 언어(DSL)를 사용하여 고성능 연산 커널을 개발해야 하는데, 이는 깊은 하드웨어 전문 지식과 많은 노동력을 필요로 하는 작업입니다. 대규모 언어 모델(LLM)이 일반적인 코드 생성 분야에서 유망한 모습을 보여주었지만, NPU 도메인의 엄격한 제약 사항과 훈련 데이터 부족으로 인해 어려움을 겪고 있습니다. 사전 연구 결과, 최신 범용 LLM은 Ascend NPU를 위한 복잡한 기능적 커널을 생성하지 못해 성공률이 거의 0에 가까운 것으로 나타났습니다. 이러한 문제를 해결하기 위해 우리는 NPU 커널 개발을 위한 생성-평가 통합 프레임워크인 AscendKernelGen을 제안합니다. 우리는 실제 커널 구현에서 도출된 사고 연쇄(Chain-of-Thought) 추론을 포함하는 고품질 데이터셋인 Ascend-CoT와, 지도 미세 조정 및 실행 피드백 기반 강화 학습을 통해 훈련된 도메인 적응형 모델인 KernelGen-LM을 소개합니다. 또한 다양한 복잡도 수준에서 컴파일, 정확성, 성능을 평가하기 위한 포괄적인 벤치마크인 NPUKernelBench를 설계했습니다. 실험 결과, 우리의 접근 방식은 일반 LLM과 하드웨어 특화 코딩 간의 격차를 크게 줄이는 것으로 나타났습니다. 구체적으로 복잡한 레벨-2 커널의 컴파일 성공률은 0%에서 95.5%(Pass@10)로 향상되었으며, 기능적 정확성은 베이스라인의 전면적인 실패에 비해 64.3%를 달성했습니다. 이러한 결과는 가속기 인식 코드 생성 자동화에 있어 도메인 특화 추론과 엄격한 평가가 중요한 역할을 함을 강조합니다.

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.