NKI-Agent: Domain-Specific Fine-Tuning and Agentic Tool Use for Neuron Kernel Generation

• URL: http://arxiv.org/abs/2607.04395v1
• Date: Sun, 05 Jul 2026 16:37:41 GMT
• Title: NKI-Agent: Domain-Specific Fine-Tuning and Agentic Tool Use for Neuron Kernel Generation
• Abstract summary: We present NKI-Agent, the first system combining domain-specific fine-tuning with agent loop for NKI kernel generation.<n>We show that tool use is critical: Opus 4.8 scores 6% in single-shot mode without agent tools. We show that an SFT-trained Qwen3-Coder-30B-A3B achieves 25.0% pass rate at 1/100th the cost of Sonnet 4 RL.
• Score: 18.169388901224416
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• AgenticCANN: Automated Ascend C Operator Generation via Knowledge-Augmented Agentic Evolution [16.709206818546317]
  AgenticCANN is a knowledge-augmented agentic evolution framework specifically tailored for automated Ascend C operator synthesis in low-corpus NPU environments.<n>Our method achieves 90 to 100 percent feasibility on elementwise and normalization operators, 56% on fusion operators, and up to 6.65$times$ speedup on 1B Pangu model inference kernels.
  arXiv  Detail & Related papers  (2026-07-29T09:19:17Z)
• KernelGenBench: A Multi-Source and Multi-Chip Benchmark for LLM-based Kernel Generation [13.879112914967058]
  Large language models (LLMs) have significantly increased the demand for efficient accelerator kernels, but kernel development remains a highly specialized and labor-intensive task.<n>We present KernelGenBench, a unified benchmark for evaluating LLM- and agent-generated Triton kernels across diverse operator sources and heterogeneous hardware platforms.
  arXiv  Detail & Related papers  (2026-07-22T12:40:11Z)
• Orchard: An Open-Source Agentic Modeling Framework [124.68499958175111]
  We present Orchard, an open-source framework for scalable agentic modeling.<n>At its core is Orchard Env, a lightweight environment service providing reusable primitives for sandbox lifecycle management.<n>On top of Orchard Env, we build three agentic modeling recipes.
  arXiv  Detail & Related papers  (2026-05-14T16:35:12Z)
• ChipCraftBrain: Validation-First RTL Generation via Multi-Agent Orchestration [0.0]
  Multi-agent approaches such as MAGE reach 95.9% on VerilogEval yet remain untested on harder industrial benchmarks.<n>We present ChipCraftBrain, a framework combining symbolic-neural reasoning with adaptive multi-agent orchestration for automated RTL generation.
  arXiv  Detail & Related papers  (2026-04-21T17:20:24Z)
• Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization [48.656549870801285]
  Kernel-Smith is a framework for high-performance GPU kernel and operator generation.<n>On the agent side, Kernel-Smith maintains a population of executable candidates and iteratively improves them.<n>On the training side, we convert long-horizon evolution trajectories into step-centric supervision and reinforcement learning signals.
  arXiv  Detail & Related papers  (2026-03-30T12:12:49Z)
• Towards Cold-Start Drafting and Continual Refining: A Value-Driven Memory Approach with Application to NPU Kernel Synthesis [68.7701048879757]
  EvoKernel is a self-evolving agentic framework that automates the lifecycle of kernel synthesis.<n>It learns stage-specific Q-values that prioritize experiences based on their contribution to the current objective.<n>It improves models' correctness from 11.0% to 83.0% and a median speedup of 3.60x over initial drafts.
  arXiv  Detail & Related papers  (2026-03-11T14:57:06Z)
• CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation [51.72529978689561]
  Agent is a large-scale agentic reinforcement learning system that develops kernel expertise through three components.<n>Agent delivers 100%, 100%, and 92% faster rate over torchcompile on KernelBench.
  arXiv  Detail & Related papers  (2026-02-27T18:58:05Z)
• Fine-Tuning GPT-5 for GPU Kernel Generation [5.109141377873154]
  We present Makora's environment and tools for reinforcement learning finetuning of frontier models.<n>In the single-attempt setting, our fine-tuned model improves kernel correctness from 43.7% to 77.0%.<n>When integrated into a full coding agent, it is able to solve up to 97.4% of problems in an expanded KernelBench suite.
  arXiv  Detail & Related papers  (2026-02-11T16:22:54Z)
• KernelCraft: Benchmarking for Agentic Close-to-Metal Kernel Generation on Emerging Hardware [25.808580418841718]
  New AI accelerators often require developers to manually craft low-level kernels.<n>This prevents emerging hardware platforms from reaching the market efficiently.<n>We present KernelCraft: the first benchmark to evaluate an agent's ability to generate and optimize low-level kernels for customized accelerators.
  arXiv  Detail & Related papers  (2026-02-10T14:52:02Z)
• AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization [12.269456144158783]
  We present AccelOpt, a self-improving large language model (LLM) agentic system that autonomously optimize kernels for emerging AI acclerators.<n>We build NKIBench, a new benchmark suite of AWS Trainium accelerator kernels with varying complexity extracted from real-world workloads.<n>Our evaluation confirms that AccelOpt's capability improves over time, boosting the average percentage of peak throughput from $49%$ to $61%$ on Trainium 1 and from $45%$ to $59%$ on NKIBench kernels.
  arXiv  Detail & Related papers  (2025-11-19T22:49:37Z)
• CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization [36.794824560677064]
  CudaForge is a training-free multi-agent workflow for kernel generation and optimization.<n>By leveraging base models like OpenAI-o3, CudaForge achieves 97.6% correctness generated kernels and an average 1.68$times$ speedup.
  arXiv  Detail & Related papers  (2025-10-23T22:52:00Z)
• FusionLLM: A Decentralized LLM Training System on Geo-distributed GPUs with Adaptive Compression [55.992528247880685]
  Decentralized training faces significant challenges regarding system design and efficiency. We present FusionLLM, a decentralized training system designed and implemented for training large deep neural networks (DNNs) We show that our system and method can achieve 1.45 - 9.39x speedup compared to baseline methods while ensuring convergence.
  arXiv  Detail & Related papers  (2024-10-16T16:13:19Z)
• FPGA-based AI Smart NICs for Scalable Distributed AI Training Systems [62.20308752994373]
  We propose a new smart network interface card (NIC) for distributed AI training systems using field-programmable gate arrays (FPGAs) Our proposed FPGA-based AI smart NIC enhances overall training performance by 1.6x at 6 nodes, with an estimated 2.5x performance improvement at 32 nodes, compared to the baseline system using conventional NICs.
  arXiv  Detail & Related papers  (2022-04-22T21:57:00Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.