NKI-Agent: Domain-specific fine-tuning and agentic tool use for neuron kernel generation

Recent agentic approaches to LLM-based kernel generation have achieved strong results on CUDA, yet emerging AI accelerators such as AWS Trainium and Inferentia remain unaddressed. Writing kernels for these chips via the Neuron Kernel Interface (NKI) is particularly challenging due to a multi-engine architecture, tile-based programming with a fixed 128-element partition dimension, and explicit memory management. Moreover, no 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 NKIGEN-BENCH, 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 NKIGEN-BENCH (63.3% on a 60-task subset), while our SFT-trained Qwen3-Coder-30B-A3B achieves 25.0% pass rate (60-task) at 1/100th the cost, outperforming Claude Sonnet 4 (15.0%) with identical tools. Tool use proves essential: even Opus 4.8 scores 0% in single-shot mode on the 60-task subset (6% on 150-task) versus 77.3% with agent tools. 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.