首个针对AWS神经芯片的智能核生成系统,可自动生成优化代码。
NKI-Agent: Domain-Specific Fine-Tuning and Agentic Tool Use for Neuron Kernel Generation

- 用领域微调+工具循环机制生成Neuron芯片核函数
- 在真实硬件上达77.3%通过率,工具使用提升6%性能
- 适合芯片开发、AI编译器研究者参考
针对AWS Trainium和Inferentia等新兴AI加速器,自动化核函数生成与优化仍属空白。使用Neuron Kernel Interface(NKI)编写核函数面临多引擎架构、分块编程及显式数据搬运等挑战,且缺乏公开训练数据、基准和工具增强型代理。本文提出NKI-Agent,首个结合领域特定监督微调(SFT)与编译-验证-修复代理循环的NKI核生成系统。将原有CUDA-Agent框架适配至Neuron硬件,构建包含6000个任务的训练集,并创建涵盖三个难度等级的250任务基准NKIBench。在真实Trn1硬件上,采用Claude Opus 4.8与层级感知提示词的NKI-Agent达到77.3%通过率。结果显示工具使用至关重要:无工具时单次生成仅6%通过率。在60任务子集上,经微调的Qwen3-Coder-30B-A3B以1/100成本实现25.0%通过率,优于Claude Sonnet 4的15.0%。此外,基于二进制编译奖励的组相对策略优化(GRPO)未能超越SFT,为强化学习在核生成中的奖励设计提供指导。
原文摘要 · Abstract (English)
Recent agentic approaches to LLM-based kernel generation have achieved impressive results on CUDA. For emerging AI accelerators such as AWS Trainium and Inferentia, automated kernel generation and optimization remain largely unaddressed. Writing kernels for these chips via the Neuron Kernel Interface (NKI) is particularly challenging: developers must navigate a multi-engine architecture, tile-based programming, and explicit data movement across multi-level memory hierarchy. Moreover, no publicly-available training data, benchmarks, or tool-augmented agents exist for this domain. We introduce NKI-Agent, the first system combining domain-specific supervised fine-tuning (SFT) with a compile-verify-fix agent loop for NKI kernel generation. We adapt the existing CUDA-Agent framework to Neuron hardware, curate 6,000 NKI kernel generation tasks for training, and construct NKIBench, a 250-task benchmark across three difficulty levels. Evaluated on real Trn1 hardware, NKI-Agent with Claude Opus 4.8 and a rank-aware system prompt achieves a 77.3% pass rate on the 150-task NKIBench. We show that tool use is critical: Opus 4.8 scores 6% in single-shot mode without agent tools. On a 60-task subset, we show that an SFT-trained Qwen3-Coder-30B-A3B achieves 25.0% pass rate at 1/100th the cost, outperforming Claude Sonnet 4 (15.0%). We also report that Group Relative Policy Optimization (GRPO) with binary compilation reward fails to improve over SFT, providing guidance on reward design for RL-based kernel generation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。