NKI-Agent: Domain-Specific Fine-Tuning and Agentic Tool Use for Neuron Kernel Generation
July 05, 2026
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 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. We show that tools 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.
1 Introduction↩︎
The rapid scaling of deep learning models has intensified demand for hardware-specific kernel optimization. While NVIDIA GPUs have benefited from decades of CUDA ecosystem development and compiler optimizations, emerging AI accelerators such as AWS Trainium and Inferentia require new approaches to kernel generation. The Neuron Kernel Interface (NKI) [1] provides a Python-based programming model for these chips, but writing high-performance NKI kernels remains challenging due to the unique multi-engine architecture and tile-based programming model.
Recent work on agentic code generation has demonstrated that agents equipped with execution feedback can generate competitive GPU kernels. The CUDA-Agent [2] system uses Proximal Policy Optimization (PPO)-trained agents on a synthesized dataset of 6,000 CUDA tasks, achieving 92–100% faster rates over torch.compile on KernelBench [3]. However, this approach is specific to CUDA and does not address the distinct challenges of Neuron hardware.
NKI programming presents unique challenges: (1) Multi-engine architecture with four specialized compute engines (Tensor, Vector, Scalar, GpSimd); (2) Tile-based programming on 2D tiles with fixed 128-element partition dimension; (3) Explicit memory hierarchy requiring manual data movement between HBM (high-bandwidth memory), SBUF (on-chip scratchpad), and PSUM (partial-sum accumulator); (4) Evolving SDK with API changes between versions.
In this paper, we present NKI-Agent, the first system to apply domain-specific fine-tuning and agentic tool use to NKI kernel generation. Our contributions:
NKI-Agent: A multi-turn agent with compile and verify tools and a rank-aware system prompt for NKI kernel generation. It supports various models including a fine-tuned Qwen3-Coder-30B-A3B model with supervised fine-tuning (SFT) and Group Relative Policy Optimization (GRPO).
NKI-Agent-Ops-6K & NKIGen-Bench: 6,000 curated NKI tasks and a 250-task benchmark across three difficulty levels.
SFT data quality analysis: Four dataset iterations showing data curation matters more than hyperparameters.
Negative result: GRPO with binary reward fails to improve over SFT.
The NKI-Agent framework achieves 77.3% on the 150-task NKIGen-Bench with a frontier model (Opus 4.8 with a rank-aware prompt), and our domain-specialized SFT model captures a large fraction of this performance at 1/100th the cost (25.0% vs 63.3% on the 60-task subset), outperforming Sonnet 4 (15.0%) with identical tools.
2 Related work↩︎
2.0.0.1 LLM-based kernel generation.
KernelBench [3] benchmarks LLM-generated GPU kernels. CUDA-Agent [2] trains PPO-based agents on 6,000 CUDA tasks with compile-verify-profile tools. Other approaches include AlphaCode-style search [4], self-repair [5], and TritonBench [6]. Related code-translation work generates accelerator code by translating across languages: TransCoder [7], BabelTower [8], and CodeRosetta [9]. The agentic-coding evaluation lineage (SWE-bench [10], SWE-agent [11]) motivates our tool-augmented, execution-grounded setup.
2.0.0.2 RL for code generation.
CodeRL [12] and PPOCoder [13] use execution rewards; DeepSeek-R1 [14] demonstrates GRPO [15] without a critic. CUDA-Agent shows RL with graded reward outperforms SFT; we show the converse for binary reward.
2.0.0.3 Neuron hardware and NKI.
AWS Trainium/Inferentia chips [16] feature four specialized compute engines. NKI [1] provides tile-based programming; prior work is limited to reference kernels [17]. AccelOpt [18] (concurrent) addresses NKI kernel optimization, complementary to our correctness-first generation.
3 Method↩︎
Figure 1 illustrates the NKI-Agent system. It consists of four components: (1) a data synthesis pipeline producing 6,000 curated NKI tasks, (2) a multi-turn agent with compile and verify tools, (3) a binary reward function, and (4) a two-stage SFT \(\rightarrow\) GRPO training pipeline.
3.1 NKI programming model↩︎
Neuron chips feature four specialized engines: Tensor (matrix ops), Vector (elementwise/reduction), Scalar (control flow), and GpSimd (general-purpose). Programs operate on 2D tiles with a fixed 128-element partition dimension. Data must be explicitly moved between HBM (device memory), SBUF (24 MB scratchpad), and PSUM (partial sum accumulator) in a load, compute, store pattern. This differs fundamentally from CUDA’s thread-block model, requiring domain-specific knowledge for correct kernels.
3.2 Data synthesis pipeline↩︎
We synthesize a dataset NKI-Agent-Ops-6K in three stages: (1) Seed crawling from PyTorch torch.nn (\(\sim\)180 ops), nki-samples (\(\sim\)30
kernels), and KernelBench (\(\sim\)250 tasks), yielding \(\sim\)400 seeds. (2) Combinatorial expansion via shape (\(5\times\)), dtype (\(3\times\)), and fusion synthesis (\(\sim\)200 patterns), producing \(\sim\)15K candidates. (3) Filtering through compilation, execution,
stability (CV\(<\)10%), and baseline checks, balanced to 6,000 tasks. The 250 NKIGen-Bench tasks are held out. For SFT, we curate 647 high-quality episodes through four iterations
(Section 5.1).
3.3 Agent architecture↩︎
NKI-Agent operates through multi-turn interaction with two tools on real Trn1 hardware: (1) compile_kernel invokes neuronx-cc with 120s timeout and returns compilation status with error
messages; (2) verify_kernel runs the compiled kernel with random inputs, comparing against the PyTorch reference via torch.allclose. The agent iterates up to \(T{=}10\) turns: generate,
compile, verify, fix. We set \(T{=}10\) empirically: in our evaluation runs the vast majority of eventual successes converge within the first few turns, and gains beyond ten turns are negligible while the per-turn Trn1
compile/verify cost grows linearly, so \(T{=}10\) balances solution coverage against evaluation cost.
3.3.0.1 Rank-aware system prompt.
NKIBench inputs span 2-D, 3-D, and 4-D tensors (e.g.batched attention and N,C,H,W convolutions), but a naive prompt assumes a 2-D \((M, N)\) layout. We use a rank-aware prompt that instructs the model to read each task’s
input-shape block and provides templates for 2-D, 3-D, and 4-D kernels, plus a reminder that every output element must be written by an nl.store. This is a prompt-only change (no retraining or extra tool calls) and is the primary driver of
model’s performance gains on hard examples.
3.4 Training pipeline↩︎
Base model. We selected Qwen3-Coder-30B-A3B [19] after evaluating models from 1.1B to 30B parameters. For inference, we serve via vLLM with tensor parallelism and native function-calling.
Stage 1: SFT. We fine-tune Qwen3-Coder-30B-A3B on 647 curated episodes using LoRA [20] (rank 32, alpha 64, 26.7M params / 0.09%) for 2 epochs with lr=\(2{\times}10^{-4}\), batch=4, and BF16 on 8\(\times\)A100 40GB.
Stage 2: GRPO. From the SFT checkpoint, we apply GRPO [15]. For each prompt, \(G{=}4\) kernels are generated, each scored with binary reward (\(+1.0\) correct, \(-0.5\) fail), and group-relative advantages computed. We trained for 200 steps with lr=\(5{\times}10^{-7}\) and \(\beta{=}0.1\).
4 Experiments↩︎
| Model | Single-Shot | NKIAgent (10 turns) |
|---|---|---|
| Base (Qwen3-30B-A3B) | 14.0% (21/150) | 20.0% (30/150) |
| SFT | 19.3% (29/150) | 20.7% (31/150) |
| GRPO | 15.3% (23/150) | 17.3% (26/150) |
4pt
| Single-Shot | NKIAgent | |||||
|---|---|---|---|---|---|---|
| 2-4 (lr)5-7 Model | L1 | L2 | L3 | L1 | L2 | L3 |
| Base | 34% | 8% | 0% | 26% | 18% | 16% |
| SFT | 50% | 8% | 0% | 30% | 16% | 16% |
| GRPO | 44% | 2% | 0% | 22% | 18% | 12% |
| Opus 4.8 | 10% | 8% | 0% | 84% | 74% | 74% |
4.1 NKIGen-Bench↩︎
NKIGen-Bench1 contains 250 tasks across three difficulty levels: Level 1 (100 tasks: single operations such as matmul,
activations, reductions, convolutions), Level 2 (100 tasks: fused operations like conv+relu+bias, attention blocks, MLP components), and Level 3 (50 tasks: full model components including Transformer layers, ResNet
bottlenecks, Mamba blocks). Each task provides a PyTorch Model class and NKI-specific metadata. Of 250 tasks, \(\sim\)150 are adapted from KernelBench, 60 require decomposition,
and 40 are NKI-native. We evaluate on a balanced 150-task subset (50 per level) and a 60-task subset for ablations.
4.2 Evaluation setup↩︎
Models: (1) Base Qwen3-Coder-30B-A3B; (2) SFT; (3) GRPO (SFT + 200 RL steps); (4) Claude Sonnet 4; (5) Claude Opus 4.8 (frontier upper bound). All Claude models use identical NKI-Agent tools.
Modes: Single-shot (SS, one generation, no tools), MT-5 (5 turns without tools; the harness compiles the output and feeds back raw error text as a user message, but the model cannot invoke tools itself), and NKIAgent (\(\leq\)10 turns with compile/verify tools).
Metric: pass rate \(=\) compiles on neuronx-cc \(+\) numerically correct (atol=rtol=\(10^{-3}\)) vs PyTorch reference. All
on real Trn1 hardware.
4.3 Main results↩︎
Tables 1 and 2 present the main results. On the aggregate 150-task metric, SFT NKIAgent (20.7%) and Base NKIAgent (20.0%) are within statistical noise (1 task apart). The clearer signal emerges from per-level breakdown and 60-task ablations (Table 3). Key findings:
SFT + NKIAgent achieves 20.7% on 150-task and 25.0% on 60-task. The NKI-Agent framework scales to 77.3% with Opus 4.8 (Table 3), showing the framework design is sound and model capability is the primary bottleneck.
Tool use uniquely enables L2/L3. No model solves any L3 task in single-shot (0/50); NKIAgent achieves 16% via iterative compile-fix.
SFT improves L1: 34%\(\rightarrow\)50%, reflecting learned NKI patterns (tile dims,
nl.load/nl.store).L1 NKIAgent \(<\) L1 SS: a counterintuitive pattern where tool use hurts easy tasks. The agent sometimes “fixes” already-correct code after seeing unrelated compiler warnings, or exhausts turns exploring alternatives. Tool use is net-positive only for tasks that require iterative debugging (L2/L3).
GRPO degrades performance vs SFT across all settings (Section 5.2).
| Configuration | 60-task | 150-task |
|---|---|---|
| Base SS | 10.0% (6/60) | 14.0% (21/150) |
| SFT SS | 15.0% (9/60) | 19.3% (29/150) |
| SFT MT-5 (no tools) | 15.0% (9/60) | — |
| Base NKIAgent | 16.7% (10/60) | 20.0% (30/150) |
| SFT NKIAgent | 25.0% (15/60) | 20.7% (31/150) |
| GRPO NKIAgent | 16.7% (10/60) | 17.3% (26/150) |
| Sonnet NKIAgent | 15.0% (9/60) | — |
| Opus 4.8 SS | 0.0% (0/60) | 6.0% (9/150) |
| Opus 4.8 NKIAgent | 63.3% (38/60) | 77.3% (116/150) |
4.4 Single-shot vs.NKIAgent↩︎
Table 3 reveals several findings. Multi-turn without tools does not help (MT-5 \(=\) SS at 15.0%), echoing the Reflexion [21] finding that verbal self-feedback without grounded execution signal yields limited gains. SFT and NKIAgent are super-additive (+5.0pp + +6.7pp separately, +15.0pp combined). Most strikingly, Opus 4.8 achieves 0% in single-shot but 63.3% with NKIAgent tools on the 60-task subset (6% vs 77.3% on 150-task), the largest tool-use gap in our study, demonstrating that even frontier models cannot generate correct NKI kernels without iterative compiler feedback. Opus 4.8’s L3 pass rate (74% on 150-task) is particularly notable: where all other models achieve \(\leq\)16%, Opus 4.8 succeeds on nearly three-quarters of full model components.
4.4.0.1 Cost-capability tradeoff.
Our SFT model (3B active params, 7.1 tok/s) captures 40% of Opus 4.8’s performance (25.0% vs 63.3% on the 60-task subset) at \(\sim\)1/100th the per-token cost. The SFT model outperforms Sonnet (25.0% vs 15.0%) despite \(\sim\)300\(\times\) less NKI training data than CUDA, demonstrating that domain SFT on scarce data can close a large fraction of the gap to frontier models.
5 Analysis↩︎
5.1 SFT data quality progression↩︎
| Version | Issue Fixed | SS (60) | \(\Delta\) |
|---|---|---|---|
| v1 | Wrong message format | — | — |
| v2 | 64% wrong function names + overfitting | 8.3% | baseline |
| v3 | 48% used 3D shapes (should be 2D) | — | — |
| v4 | All fixed: 2D shapes, correct naming | 15.0% | +81% rel. |
| Base (no SFT) | — | 10.0% | — |
Table 4 traces our data curation journey. SFT v2 with bad data hurts performance (8.3% vs base 10.0%). Each fix targeted one issue: v1\(\rightarrow\)v2 fixed message format; v2\(\rightarrow\)v3 standardized function names (64% were wrong); v3\(\rightarrow\)v4 converted 3D tensors to 2D. This shows that each data quality fix produced larger gains than any hyperparameter change.
5.2 GRPO negative result↩︎
| Model | SS (60) | NKIAgent (60) | SS (150) | NKIAgent (150) |
|---|---|---|---|---|
| SFT v4 | 15.0% | 25.0% | 19.3% | 20.7% |
| GRPO | 13.3% | 16.7% | 15.3% | 17.3% |
| \(\Delta\) | \(-\)11% | \(-\)33% | \(-\)21% | \(-\)16% |
8pt
None
Figure 4: Illustrative NKIAgent trace (condensed from evaluation logs). The agent iteratively discovers the correct API (nisa.nc_matmul), fixes tile dimension constraints, and adds proper accumulation, errors
that are impossible to recover from in single-shot mode..
Table 5 shows GRPO consistently underperforms SFT v4 (average reward flat at \(\sim\)10–15% compile rate across 200 steps). We hypothesize: (1) binary reward is too sparse, with no signal between “almost compiles” and “completely wrong”; (2) subtle code changes flip reward discontinuously; (3) when \(\sim\)75% of generations fail, group-relative ranking selects among incorrect outputs. CUDA-Agent’s graded reward (speedup measurements) provides richer signal.
5.3 Error analysis and qualitative findings↩︎
5.3.0.1 Category-level patterns.
SFT excels at elementwise (87.5%) and activation (90.0%) operations. Convolutions, attention, and all Level 3 components achieve 0% in single-shot for our model. However, Opus 4.8 NKIAgent achieves 74% on Level 3 (37/50, 150-task), demonstrating that the NKI-Agent tool framework is sufficient for complex tasks when paired with stronger base models.
5.3.0.2 Common failures.
Incorrect tile dimensions (partition \(\neq\) 128), unsupported operation combinations, and SBUF/PSUM memory violations. Failures are systematic: entire categories fail uniformly, suggesting broader data coverage is more promising than RL. SFT reliably generates the load\(\rightarrow\)compute\(\rightarrow\)store pattern with 128-element partition but struggles with Tensor Engine matmul and multi-engine pipelining.
5.4 Example agent trace↩︎
Figure 4 shows an illustrative NKIAgent interaction on a matrix multiplication task, condensed from typical evaluation logs. This trace illustrates why tool use is essential: the agent must discover that
nl.matmul does not exist in NKI, learn the tile shape constraints of nisa.nc_matmul, and fix numerical accumulation, each requiring real compiler feedback.
6 Limitations and future work↩︎
Our work has several limitations. First, we evaluate correctness only, not runtime performance; AccelOpt [18] addresses this complementary axis. Second, binary reward was insufficient for GRPO; graded rewards remain future work. Third, all results are single-run without error bars; small absolute differences may not be statistically significant. Fourth, our SFT model uses only 3B active parameters; our Opus 4.8 results (77.3%) confirm that scaling model capability significantly improves results within the same framework. Finally, results are specific to Neuron SDK 2.24.
Conclusion↩︎
We presented NKI-Agent, an agentic system combining domain-specific supervised fine-tuning with agentic tool use for kernel generation on AWS Neuron hardware. Evaluated on NKIGen-Bench, our framework demonstrates that iterative compile-verify-fix loops are essential where even Claude Opus 4.8 scores 6% in single-shot but 77.3% with agent tools. Our SFT-trained Qwen3-Coder-30B-A3B captures 40% of frontier performance at 1/100th the cost, outperforming Claude Sonnet 4. We additionally show that data quality dominates hyperparameter tuning, and that GRPO with binary reward fails to improve over SFT, providing guidance on reward design for RL-based kernel generation.