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

arXiv: https://arxiv.org/abs/2607.04395 · HTML: https://mlsanity.com/paper/2607.04395/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Software engineering coding agents, one of 16 directions in LLM agents and tool use.
- Builds on 9 of its 19 references that are in the map.
- First posted on arXiv 5 Jul 2026.

## Placed in

- LLM agents and tool use / [Software engineering coding agents](https://mlsanity.com/agents/software-engineering-coding-agents/)

## Builds on (most cited first, 9 of 9)

- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [SWE-bench: Can Language Models Resolve Real-World GitHub Issues?](https://mlsanity.com/paper/2310.06770/) (2023)
- [Unsupervised Translation of Programming Languages](https://mlsanity.com/paper/2006.03511/) (2020)
- [Is Self-Repair a Silver Bullet for Code Generation?](https://mlsanity.com/paper/2306.09896/) (2023)
- [Execution-based Code Generation using Deep Reinforcement Learning](https://mlsanity.com/paper/2301.13816/) (2023)
- [TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators](https://mlsanity.com/paper/2502.14752/) (2025)
- [CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation](https://mlsanity.com/paper/2602.24286/) (2026)
- [AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization](https://mlsanity.com/paper/2511.15915/) (2025)

## Built on by (most cited first, 2 of 2)

- [MaxKernel: Agentic Kernel Generation for TPUs](https://mlsanity.com/paper/2609.04523/) (2026)
- [From Experience to Expertise: Adoption-Aware Memory Learning for Data-Scarce NPU Kernel Synthesis](https://mlsanity.com/paper/2609.35568/) (2026)

## Abstract (arXiv)

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.