From Experience to Expertise: Adoption-Aware Memory Learning for Data-Scarce NPU Kernel Synthesis

Longxiao Fan; Tao Zhang; Han Yan; Jiajun Li; Mingcong Song; Guoping Long; Hongjie Si; Weiwei Sun

cs.LGSep 28, 2026

From Experience to Expertise: Adoption-Aware Memory Learning for Data-Scarce NPU Kernel Synthesis

Authors: Longxiao Fan, Tao Zhang, Han Yan, Jiajun Li, Mingcong Song, Guoping Long, Hongjie Si, Weiwei Sun

Organizations: Fudan University · University of Science and Technology of China · Huawei Technologies Ltd.

Abstract

High-performance kernels underpin efficient accelerator execution but require expert tuning and lengthy manual optimization cycles. LLM coding agents promise automation, yet their CUDA knowledge transfers poorly to data-scarce domain-specific architectures (DSAs) such as NPUs, whose execution models and memory hierarchies differ substantially from those of GPUs. To address this transfer gap, post-training methods adapt LLMs to NPU programming but depend on scarce expert data and substantial training compute. Memory-learning agents instead adapt through external memory, but their uniform credit assignment gives adopted and unused experiences the same reward target, potentially biasing subsequent retrieval rankings. Moreover, when learned values guide only retrieval, high-value experiences that generalize across operators must be retrieved repeatedly rather than retained in context, thereby increasing retrieval overhead and weakening cross-task guidance. We therefore present SAGE, a persistent self-improving agent for NPU kernel synthesis. Adoption-Traced Utility estimation (ATU) combines explicit adoption records with kernel evaluation outcomes for adoption-aware credit assignment. Utility-Gated Consolidation (UGC) uses positive utility and repeated adoption across operators to select and abstract reusable rules into a bounded resident context. On NPUKernelBench, SAGE achieves a 95.5% execution rate versus 84.1% for the strongest controlled baseline, with 86.9% of solved operators outperforming torch_npu. With GLM-5.3, SAGE achieves a 43.99x speedup over the torch_npu reference on sparse flash attention. These results show that adoption-aware credit assignment and selective consolidation enable agents to accumulate and reuse hardware-specific knowledge across tasks.

Figures & tables

Figure 1: Retrieval-oriented memory learning with uniform credit and repeated retrieval (left) versus SAGE’s adoption-aware credit assignment and selective consolidation into resident context (right).

Figure 2: Overview of SAGE. (a) A frozen agent iteratively refines kernels using retrieved cold experiences, resident hot rules, and execution feedback. (b) ATU combines adoption traces with terminal outcomes to estimate item-level utility. (c) UGC uses utility and repeated adoption across operators to select reusable rules for a bounded resident context and demote underused rules.

Model	Method	Execution Rate			Overall		Performance
Model	Method	L1	L2	L3	CR	ER	$\mathbf{Fast}_{1.0}$	$\mathbf{S}_{\mathrm{self}}$
Qwen3-Coder-Next (80B)	Refinement	5.0	0.0	0.0	35.2	1.1	0.0	1.00 $\times$
	Static RAG	30.0	26.0	0.0	51.1	21.6	21.1	1.10 $\times$
	Value Memory	50.0	24.0	5.6	65.9	26.1	17.4	1.16 $\times$
	SAGE	55.0	32.0	5.6	75.0	31.8	28.6	1.27 $\times$
DeepSeek-V4-Flash (284B)	Refinement	35.0	18.0	0.0	55.7	18.2	31.3	1.20 $\times$

Table 1: Main results on the 88-operator NPUKernelBench subset. All rates are reported as percentage. The best result for each model is marked in bold.

Task schedule	L3 ER	$\mathbf{Fast}_{1.0}$
L3 Only	38.9	42.9
L2+L3 Mixed	61.1	63.6
L2 $\rightarrow$ L3	77.8	71.4

Table 2: Cross-difficulty transfer to the 18 Level-3 operators. All rates are percentages.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Level	Category	Operators
Level 1	Comparison	Equal, Less, LessEqual
	Condition	IsFinite, IsInf, NonFiniteCheckOp
	Index	GatherV3, ScatterList
	Math	Sqrt, SinMath, CastMath, AbsMath, Addcdiv, ClipByValue, CosMath, Lerp, StridesliceNegConcatV2
	TensorCreation	Arange, Eye, Fill
Level 2	Activation	DequantSwigluQuant, FastGelu, Gelu, InplaceAttnSoftmax, SwiGlu, Swish, GeGluV2, GeluGrad, MulSigmoid

Appendix

Table 3: Operators included in the 88-task NPUKernelBench evaluation subset.

Operator	Operator
lightning_indexer	sparse_flash_attention
mla_prolog_v3_x	compute_and_cache_digest
apply_rotary_pos_emb_partial	gmm_deq_swiglu_quant_gmm_deq
hc_post_x	hc_post
fused_gelu_sigmoid	fused_scatter_update_and_compute_digest
rms_norm_rope	grouped_matmul_swiglu_quant_x

Appendix

Table 4: The 16 xCCL operators used in the transfer evaluation.

Symbol	Value	Component	Description
$T$	30	Agent loop	Maximum interaction rounds
$\delta$	0.2	ATU	Unsuccessful-episode penalty
$\beta$	0.2	ATU	Unadopted-retrieval penalty
$\eta_{\min}$	0.05	ATU	Minimum utility update rate
$\gamma$	0.3	Retrieval	Utility reranking strength
$K_{\mathrm{ado}}$	3	UGC	Minimum adopted trajectories

Appendix

Table 5: Main hyperparameters used throughout SAGE.

Correspondence	Count	Percentage
Full	44	88%
Partial	6	12%
Unrelated	0	0%

Appendix

Table 6: Expert-rated correspondence between adoption and code revisions for 50 records randomly sampled from all trace records labeled adopted.

Ranking signal	ER	L3 ER	$\mathbf{Fast}_{1.0}$
Relevance Only	84.1	61.1	70.3
Uniform Credit	89.8	72.2	79.7
ATU	95.5	88.9	86.9

Appendix

Table 7: Credit-assignment control under a shared frozen memory bank. All rates are percentages.

ID	Retrieved experience	Rel.	Adopted
181	Host-side tiling cannot dereference device-resident tensor data returned by GetData .	69.17	Yes
193	A kernel entry can read a scalar directly from GM through a gm pointer cast.	67.42	Yes
61	Scalar inputs can be exposed to tiling through ValueDepend(REQUIRED) and passed in TilingData .	57.43	No
139	Nested tiling structures require GET_TILING_DATA_WITH_STRUCT .	49.22	No
159	TilingData field types must match their kernel-side consumers.	32.47	No

Appendix

Table 8: Retrieved experiences for the GatherV3 correctness-stage runtime failure. Relevance is the retriever score; adoption is determined from the subsequent diagnosis and code revision.

Iteration	Diagnosis	Revision	Outcome
0–2	A shared Matmul object changed its $K$ dimension despite fixed tiling; the default task type also incorrectly assigned Cube work on AIV core.	Separate the QK and PV configurations and declare KERNEL_TYPE_MIX_AIC_1_0 .	Compilation succeeds; Matmul no longer hangs.
10–13	Multi-token execution exposed unmatched AIC–AIV synchronization.	Isolate a single-token path, then pair cross-core events within the token loop.	The hang becomes a diagnosable numerical mismatch.
18–23	DumpTensor showed zero softmax inputs; four host inputs were non-contiguous, and the QK path omitted the RoPE dimensions.	Synchronize the asynchronous Matmul write, apply AutoContiguous() , and materialize the full 576-dimensional Q/K representations.	First correct kernel that provides optimization baseline: $0.420$ ms ( $1.91\times$ ).
24–26	Double buffering violated the MTE3 handoff; a global barrier was also overly conservative.	Revert the unsafe buffer change and replace PIPE_ALL with an MTE3–MTE2 event.	$1.95\times$ .
27	Work remained single-core.	Partition query tokens across cores.	$3.71\times$ .
29	A relay copy materialized contiguous QK data before PV.	Split PV into nope and RoPE Matmuls that read the original GM tensors directly.	$0.155$ ms; $5.13\times$ .

Appendix

Table 9: Key transitions in the Mla task-solving trajectory on ascend910_93 .

Field	Type	Description	Example
id	integer	Stable memory-item identifier.	10035
title	string	Compact retrieval-facing rule.	Reduce algorithmic complexity and branches
type	enum	Experience type used for filtering.	algorithmic_optimization
category	string	Operator or failure category.	Scan/Reduction
error_message	string	Failure or inefficiency that motivated the item.	Quadratic polynomial evaluation with redundant writes
code_diff	object	Minimal incorrect/correct code pair.	{wrong_code, correct_code}

Appendix

Table 10: Schema of a memory-item record.

Field	Type	Description	Example
session_id	integer	Memory-evolution session identifier.	1
created_at	timestamp	Snapshot creation time.	20260811_111049
budget	integer	Maximum resident-context tokens.	10000
n_ops	integer	Completed operator trajectories at snapshot time.	18
hot_entries	list[object]	Items selected under the token budget.	[{id: 238, ...}]
Nested hot-entry fields

Appendix

Table 11: Schema of a hot-memory snapshot and its nested entry representation.

Explore similar work

NKI-Agent: Domain-Specific Fine-Tuning and Agentic Tool Use for Neuron Kernel Generation
Jul 5, 2026Junjie Tang, Jun Huan, Hao Zhou +2Graphics Processing Unit Kernels Cuda
Amazon Web Services
Hawk: Harnessing Hardware-Aware Knowledge for High-Performance NPU Kernel Generation
Jul 2, 2026Junyi Wen, Ruiyan Zhuang, Yongjia Xu +7Graphics Processing Unit Kernels Neural Processing Units
Sun Yat-sen University, Zhuhai, China · Greater Bay Area National Technology Innovation Center, Guangzhou, China · Peng Cheng Laboratory, Shenzhen, China
From Human Guidance to Autonomy: Agent Skill System for End-to-End LLM Deployment on Spatial NPUs
May 27, 2026Jiajie Li, Erwei Wang, Zhiru Zhang +1LLM Inference Optimization Agent Skill Retrieval
Cornell University · AMD

Jul 5, 2026, cs.LGJ/K move · Enter open · S save

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

Junjie Tang, Jun Huan, Hao Zhou, Yuhao Zhang+1

Amazon Web Services

Graphics Processing Unit Kernels Cuda Agentic Qwen3 Neurons

Open paper AskPDF 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.

Abstract

Explore similar work

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

Hawk: Harnessing Hardware-Aware Knowledge for High-Performance NPU Kernel Generation

From Human Guidance to Autonomy: Agent Skill System for End-to-End LLM Deployment on Spatial NPUs