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

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.

• The paper demonstrates that incorporating iterative compile-verify loops and rank-aware prompting significantly boosts performance in generating complex Neuron kernels.
• The approach leverages supervised fine-tuning on 6,000 curated tasks, achieving up to 74% pass rate on complex Level 3 tasks compared to single-shot failures.
• The study highlights that data quality and systematic prompt engineering outperform reinforcement learning methods for domain-specific kernel synthesis.

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

Overview and Motivation

This paper introduces NKI-Agent, an agentic code generation system architected specifically for kernel development targeting AWS’s Neuron hardware via the Neuron Kernel Interface (NKI) (2607.04395). Unlike the extensively studied CUDA ecosystem, the NKI domain presents pronounced challenges: a multi-engine architecture, explicit and tile-based memory management, and rapid SDK/API evolution. The absence of publicly available high-quality data, native benchmarks, and suitable LLM-based agentic frameworks for NKI accentuates the need for bespoke solutions.

The framework combines supervised fine-tuning (SFT) on carefully curated NKI data, agentic iterative compilation and verification, and exploration of reinforcement learning via Group Relative Policy Optimization (GRPO). The methodology synthesizes a large dataset, introduces an evaluation suite (NKIGen-Bench), and systematically studies model and system design choices, focusing both on correctness and data/agentic workflow efficacy.

Figure 1: NKI-Agent system architecture, showing iterative code generation, compilation, verification, and feedback on AWS Trn1 hardware, driven by a data pipeline producing 6,000 tasks.

NKI-Specific Challenges and Architecture

NKI mandates explicit programming paradigms not found in traditional GPU programming, such as the orchestration of Tensor, Vector, Scalar, and GpSimd engines, and rigorous data movement across HBM, SBUF, and PSUM. The agent framework comprises multiple components:

• Data Synthesis Pipeline: Yields 6,000 curated NKI kernel generation tasks for SFT, from diverse sources, including PyTorch ops, public NKI samples, and existing benchmarks.
• Multi-Turn Agent: Implements a compile-verify loop (≤10 turns), with real hardware feedback underpinning error correction, a necessity for the domain’s brittle APIs and architectural constraints.
• Prompt Engineering: Crucially, a rank-aware prompt adapts code synthesis to input tensor rank and NKI’s strict template needs, substantially improving model success on high-difficulty tasks.
• Two-Stage Learning Pipeline: SFT followed by optional RL (GRPO).

Data Pipeline and Benchmarking

6,000 NKI tasks were generated via a multi-phase process: initial seed crawling, combinatorial shape/type/fusion expansion, and robust filtering (compilation, execution, stability, baseline correctness). Importantly, the NKIGen-Bench is a new, publicly defined suite with 250 held-out evaluation tasks across three difficulty levels:

• Level 1: Isolated ops (matmul, conv, reductions)
• Level 2: Fused primitives and function blocks
• Level 3: Complex model components, including full layers (Transformers, ResNet, Mamba blocks)

Model performance is measured on both an aggregate 150-task balanced subset and a focused 60-task ablation set, with all results on physical Trn1 hardware.

Numerical Results and Key Findings

Figure 2: Pass rate by difficulty level showing the sharp drop for single-shot generation on higher complexity (L3); NKI-Agent tool use uniquely enables nontrivial L2 and L3 performance, with Opus 4.8 (rank-aware prompt) achieving up to 74% on L3.

Figure 3: Ablation study on the 60-task subset; SFT-trained model captures over 40% of the frontier model's performance at substantially reduced computational cost and data footprint.

The most critical empirical claims, supported by extensive evaluation, include:

• Tool Use is Essential: Single-shot code generation (without compile-verify turns) fails completely on Level 3 tasks (0% across all models), and even frontier LLMs (Opus 4.8) achieve only 6% in single-shot versus 77.3% in full agentic mode. This is the largest tool-use delta reported for code generation agents on any architecture or domain.
• Supervised Fine-Tuning: SFT on just 647 curated episodes boosts performance from baseline (14–16.7%) to 20.7% (150-task). On the focused 60-task subset, the fine-tuned Qwen3-Coder-30B-A3B model achieves 25%, outperforming Claude Sonnet 4 (15.0%) at 1/100th the inference cost.
• Prompt Engineering: The rank-aware prompting provides a non-trivial jump for hard tasks, particularly on Opus 4.8, with L3 pass rate reaching 74%.
• RL Negative Result: GRPO with a binary correctness reward consistently degrades or at best fails to surpass SFT, attributed to sparse/discontinuous reward signals. Thus, reward granularity and domain-specific metric design remain open challenges; graded (e.g., speedup) rewards found effective in CUDA are required here.
• Data Quality Over Hyperparameters: Ablations across multiple SFT data versions reveal that systematic data curation (e.g., correcting message formats and input ranks) yields larger gains than any hyperparameter tuning, and bad data can drastically reduce performance (e.g., 8.3% vs. 10.0% base in Table 3).

Agentic Error Recovery and Qualitative Insights

A qualitative trace underscores why agentic compilation/emulation feedback is indispensable. The agent typically requires several turns to resolve non-existent APIs, incompatible tile shapes, or incorrect data movement—a process infeasible for any LLM in a single generation due to the highly specific and under-documented nature of the NKI APIs.

Failure analysis reveals systematic coverage gaps by operation category (e.g., activations are handled well; convolutions and engine-pipelined ops are not), suggesting that future progress is better served by dataset and template expansion rather than further model scaling or RL under the current reward scheme.

Broader Implications and Future Directions

Practical Implications: For hardware-specific LLM code generation—particularly for emerging, non-commodity accelerators—iterative, tool-driven agentic frameworks are categorically superior to single-shot paradigms. SFT on high-quality, niche data allows strong baseline performance at low cost, informing the design of accessible, efficient developer tools for future specialized silicon.

Methodological Implications: The negative result for RL with binary reward motivates exploration of more expressive, graded reward functions (e.g., performance, energy profiling) for kernel synthesis. Moreover, scaling SFT-trained models and integrating richer prompt engineering can further close the gap to proprietary, frontier models even in very low data regimes.

Future Work: Next steps include extending evaluation to cover runtime performance, developing reward functions for optimization RL, and automated data expansion for underrepresented NKI ops, as well as robustness studies across SDK/API versions. The potential for hybrid approaches—combining SFT with inference-time retrieval or code-repair, and cross-domain knowledge transfer—is substantial.

Conclusion

NKI-Agent establishes the necessity and effectiveness of agentic, tool-augmented workflows and domain-specialized fine-tuning for hardware kernel generation in new accelerator domains. Models augmented with multi-turn compilation and prompt strategies can generate and correct highly idiosyncratic code that neither SFT nor even large frontier LLMs can synthesize in single-pass, while careful data curation outpaces additional parameter scaling or naive RL. These results inform future AI toolchain designs for rapidly evolving, heterogeneous AI hardware platforms.