PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.35568v1cs.LG

TL;DR

Data-scarce NPUs pose a kernel-synthesis transfer gap, while post-training depends on scarce expert data and substantial compute. SAGE uses adoption-aware utility and selective consolidation to evolve external memory; across three backbones, it improves correctness and efficiency, with larger gains on complex Level-3 operators.

  • Problem

    NPU kernel synthesis requires architecture-specific reasoning despite scarce public code, while post-training adaptation is constrained by scarce expert data and substantial compute.

    Source: §1 · pp. 1–2

  • Method

    SAGE evolves a frozen LLM’s external memory by combining adoption records with verifier outcomes for utility and consolidating positively valued experiences repeatedly adopted across operators.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    Across three backbones, SAGE improves correctness and efficiency, with larger gains on complex Level-3 operators.

    Source: §4.2 · p. 7, §5 · p. 10

  • Takeaways & Limitations

    The results support accumulating and reusing hardware-specific expertise across tasks through external memory evolution.

    Source: §5 · p. 10

    The 30-round study does not cover long-horizon synthesis of fused kernels and megakernels.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

NPU kernel synthesis faces scarce public expertise and architecture-specific execution and memory models, limiting transfer from CUDA-trained agents. SAGE addresses this through adoption-aware credit assignment and selective consolidation of reusable experience.

1 Introduction

  • GPT-5.2 correctness drops from 92% on CUDA to 14% on Ascend C, underscoring the transfer gap in data-scarce NPU kernel synthesis.Such synthesis also depends on acquiring and reusing hardware-specific expertise under scarce supervision.Source: §1 · pp. 1–2
  • Post-training for NPUs is constrained by scarce expert data and training compute, while retrieval-based memory can credit unused experiences and require repeated retrieval of reusable knowledge.Figure 1 contrasts uniform credit and repeated retrieval with SAGE’s adoption-aware credit assignment and resident-context consolidation.Source: §1 · p. 2, Figure 1 · p. 2
  • SAGE grounds credit assignment and consolidation in evidence of use: adoption plus outcomes informs utility, while repeated adoption across operators signals possible reuse.A single adoption is insufficient evidence of transferability; a bounded resident context makes consolidated rules available without another successful retrieval.Source: §1 · p. 2
  • SAGE pairs Adoption-Traced Utility estimation with Utility-Gated Consolidation to evolve a frozen LLM’s external memory and retain reusable rules under a fixed context budget.The paper’s stated contributions include evaluation across three backbones for kernel correctness and runtime performance.Source: §1 · pp. 2–3

2 Related Work

Related work spans automated kernel synthesis for mature GPU ecosystems and data-scarce accelerators, as well as memory-based agent self-improvement. SAGE is positioned against systems lacking explicit evidence of which retrieved experiences informed code revision.

2.1 LLM-based Automated Kernel Synthesis

  • GPU kernel synthesis benefits from mature CUDA/Triton ecosystems [19] [20] [3] [21] [22], while Ascend work uses SFT/RL [14] [23], DSL transcompilation [24], or retrieved knowledge and repair [10] [15] [4] [25] [26].The paper identifies abundant public kernels and documented optimization idioms as advantages of mature GPU ecosystems, unlike data-scarce accelerators.Source: §2.1 · p. 3

2.2 Memory-Augmented Agentic Self-Improvement

  • Memory agents range from reflections and reusable skills [27] to outcome-conditioned selection [16] [17] and provenance-based credit propagation [18], but do not identify which retrieved experiences informed a code revision.SAGE traces individual adoption and uses accumulated cross-task evidence to consolidate repeatedly validated knowledge into persistent memory.Source: §2.2 · p. 3

3 Methods

SAGE keeps the generator fixed while learning across an ordered task stream through persistent external memory. ATU assigns utility using adoption traces and episode outcomes, while UGC consolidates repeatedly adopted, useful experiences into bounded resident rules.

3.1 Preliminaries

  • SAGE formulates synthesis as an ordered task stream in which a fixed generator adapts only its persistent external memory.Each task specifies an operator and a correctness or optimization objective; memory transitions aim to improve subsequent execution outcomes.Source: §3.1 · pp. 3–4

3.2 SAGE Framework

  • SAGE organizes memory into cold source experiences and hot consolidated rules, with ATU updating utility and UGC selecting resident rules under a fixed token budget.New experiences enter cold memory; underused hot rules can be displaced.Source: §3.2 · p. 4, Figure 2 · p. 4, §3.4 · p. 6

3.3 Adoption-Traced Utility Estimation

  • SAGE records each retrieved item's adoption state and rationale, and blocks evaluation until every retrieved item has a resolved adoption state.The method also tracks retrieval and adoption counts and the distinct operators that adopted each item.Source: §3.3 · p. 5
  • ATU defers credit until episode termination, then converts the terminal score into item-level evidence for retrieved and adopted experiences.Adopted items share a fixed episode-level credit budget, while retrieved but unused items receive fixed negative credit.Source: §3.3 · p. 5
  • SAGE updates empirical utility online and uses it to rerank relevance-retrieved experiences, while resident hot knowledge bypasses retrieval.The retrieval candidate pool is first obtained from cold memory based solely on relevance; utility affects late reranking.Source: §3.3 · p. 5

3.4 Utility-Gated Consolidation

  • UGC makes cold experiences eligible for rule abstraction only after positive utility, sufficient adoption, and cross-operator validation.The abstracted rule may include a minimal code or API example, while its source and evidence are retained for provenance.Source: §3.4 · pp. 5–6
  • UGC ranks validated candidates and resident rules by expected utility per resident token, then greedily selects within the fixed token budget.Persistent non-use can lower an item's access estimate and utility, allowing demotion and restoring its source experience's retrieval eligibility.Source: §3.4 · p. 6

4 Experiments

Across controlled NPUKernelBench experiments, SAGE's memory evolution improved results across backbones and task settings. The experiments also assess transfer to xCCL, task order, cross-backbone memory portability, and component ablations.

4.1 Experimental Setup

  • The evaluation used 88 coverage-oriented NPUKernelBench operators on Ascend 910C, with three source models, controlled baselines, and up to 30 refinement rounds per operator.All methods shared the prompts, cold-start knowledge, diagnostic tools, retrieval pipeline, and hardware evaluator.Source: §4.1 · pp. 6–7, §4.2 · p. 7

4.2 Main Results

  • 84 of 88 GLM-5.2 operators were solved; versus Value Memory, Fast1.0 rose from 75.7% to 86.9% and median Sself from 3.36× to 4.47×.SAGE's gains also persisted with DeepSeek-V4-Flash and Qwen3-Coder-Next.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 Transfer Beyond the Benchmark

  • On xCCL, GLM-5.2 solved 11 of 16 operators (68.8% ER), with 81.8% of correct kernels outperforming torch_npu; GLM-5.3 reached 81.3% ER.For sparse flash attention, GLM-5.3 achieved 43.99× speedup over torch_npu.Source: Figure 3 · p. 8, §4.3 · p. 8

4.4 Effect of Task Order

  • L2→L3 improved Level-3 ER by 16.7 percentage points over mixed scheduling and 38.9 points over L3-only adaptation.Fast1.0 showed the same ordering.Source: Table 2 · p. 8, §4.4 · p. 9

4.5 Cross-Backbone Memory Transfer

  • GLM-5.2-evolved memory raised DeepSeek-V4-Flash ER from 64.8% to 72.7% and Qwen3-Coder-Next ER from 31.8% to 40.9%, without memory updates during transfer evaluation.Both recipient backbones also improved on Fast1.0, while a substantial gap to GLM-5.2 remained.Source: Figure 4 · p. 9, §4.5 · p. 9

4.6 Ablation Study

  • Uniform Credit reduced Level-3 ER from 88.9% to 66.7%; removing UGC retained much of ATU's benefit but lagged SAGE, while Random UGC underperformed utility-guided consolidation.Uniform Credit assigns the same terminal episode score to every retrieved experience regardless of adoption.Source: Figure 5 · p. 9, §4.6 · pp. 9–10

5 Conclusion

SAGE evolves external memory so frozen LLMs can accumulate hardware-specific NPU expertise. Across three backbones it improves correctness and efficiency, while the 30-round study leaves long-horizon fused-kernel synthesis untested.

5 Conclusion

  • Across three backbones, SAGE improves correctness and efficiency, with larger gains on complex Level-3 operators.Transfer results and mechanism ablations support the roles of accurate attribution and selective consolidation.Source: §5 · p. 10
  • SAGE estimates item-level utility from adoption traces and consolidates repeatedly validated knowledge within a bounded hot-memory budget.The approach lets frozen LLMs accumulate hardware-specific expertise through external memory evolution.Source: §5 · p. 10
  • The 30-round study does not cover long-horizon synthesis of fused kernels and megakernels, which requires diversity-preserving search across coupled scheduling, memory-layout, and synchronization decisions.Source: §5 · p. 10
  • Future work will explore the approach across other emerging hardware architectures and domain-specific programming languages.Source: §5 · p. 10

A.1 NPUKernelBench Evaluation Subset

The main NPUKernelBench evaluation uses an 88-operator subset, including all Level-3 tasks and coverage-oriented subsets from Levels 1 and 2.

A.1 NPUKernelBench Evaluation Subset

A.2 xCCL Transfer Evaluation Set

The xCCL transfer evaluation uses 16 operators.

A.2 xCCL Transfer Evaluation Set

A.3 Level-3/GMM Per-Operator Speedups

The appendix reports per-operator speedups for Level-3/GMM and xCCL tasks, gives SAGE’s hyperparameter settings, and describes UGC’s memory replacement and agent runtime.

A.3 Level-3/GMM Per-Operator Speedups

  • 1.54× is the median speedup over torch_npu across 16 correctness-passing Level-3/GMM operators; individual speedups range from 0.38× to 71.41×.Figure 6 uses log2-transformed horizontal distances and marks parity; green and red distinguish speedups above and below 1×.Source: §A.3 · p. 13, Figure 6 · p. 13

A.4 xCCL Per-Operator Speedups

  • 7.32× versus 3.32× is the median speedup over torch_npu on 11 commonly solved xCCL operators for GLM-5.3 and GLM-5.2, respectively.GLM-5.3 also solves two additional operators at 11.81× and 10.86×, with a 10.86× median across all 13 correctness-passing operators; both inherit memory evolved over 88 NPUKernelBench tasks.Source: §A.4 · p. 14, Figure 7 · p. 14

B Hyperparameter Configuration

  • The same hyperparameters are used across backbones and evaluation sets, without per-backbone or per-operator tuning.The settings include δ = β = 0.2, γ = 0.3, Kado = Kop = 3, and a 10K-token resident hot-memory budget.Source: B Hyperparameter Configuration · p. 14, Table 5 · p. 15

C.1 Utility-Gated Consolidation Procedure

  • After each task episode, UGC pools hot-memory items with validated candidates, ranks them by expected utility per token, and rebuilds hot memory under budget B.Items that do not fit are skipped while scanning continues; displaced residents return to validated status and regain retrieval eligibility.Source: §C.1 · p. 15

C.2 Infrastructure and Agent Runtime

  • The agent runtime orchestrates iterative generation, compilation, execution, diagnosis, and memory updates, with callable tools for runtime and performance diagnosis.Experiments use Ascend 910C infrastructure; diagnostics include memory and synchronization checks and kernel profiling.Source: §C.2 · pp. 15–16

C.3 Initialization and Experimental Integrity

The evaluation controls memory mechanisms while holding the backbone, agent harness, prompts, tools, interaction budget, and evaluation harness fixed. It also specifies safeguards, correctness validation, and device-side latency measurement.

C.3 Initialization and Experimental Integrity

  • Cold-start knowledge comes from distilled AscendC references and patterns mined by comparing generated kernels with expert solutions on operators disjoint from the 88 evaluation operators.The references comprise a quick reference, debugging guide, and add_custom implementation example.Source: §C.3 · p. 16
  • All memory-based methods start from identical memory, and the comparison holds the backbone, harness, prompts, tools, interaction budget, and evaluation harness fixed so only memory mechanism differs.The interaction budget is T = 30.Source: §C.3 · p. 16
  • Evaluation results are recorded before memory updates, while cross-backbone transfer uses fixed evolved memories without further updates.The 88 operators form the learning stream.Source: §C.3 · p. 16
  • Safeguards prohibit high-level substitutes, expert-implementation leakage, semantic shortcuts, and reduced-precision transformations that fail the evaluation correctness criteria.The constraints apply to all methods during evaluation-operator task solving.Source: §C.3 · p. 16

C.4 Correctness and Performance Measurement

  • Correctness validation compares generated NPUKernelBench kernels with torch_npu using average and maximum absolute and relative error, plus average accuracy.The passage lists these benchmark error metrics but does not provide the full correctness threshold.Source: §C.4 · p. 16
  • Latency uses device-side Ascend NPU events, with five warm-ups and 100 timed iterations under identical inputs and synchronization boundaries for generated kernels and torch_npu references.The reported latency is mean per-execution latency over the timed iterations.Source: §C.4 · pp. 16–17

D.1 Expert Audit of Adoption Records

The audit checks whether trace records labeled adopted correspond to observable code revisions. An Ascend C expert rated all 50 sampled records as at least partially corresponding.

D.1 Expert Audit of Adoption Records

  • The audit randomly sampled 50 records from all records labeled adopted, using each adoption record as the sampling unit.The saved trajectory provided the code evidence for expert assessment.Source: §D.1 · p. 17
  • 44 of 50 sampled adopted records (88%) showed full correspondence with code revisions, six (12%) showed partial correspondence, and none were unrelated.An Ascend C expert rated each record against its associated experience, adoption rationale, and code revision.Source: §D.1 · p. 17, Table 6 · p. 17

D.2 Fixed-Bank Credit Assignment

With a shared frozen memory bank, ATU outperforms uniform-credit and relevance-only controls across reported metrics. The accompanying properties characterize its credit and reranking behavior and UGC’s usage-aware decay, illustrated by adoption and kernel-synthesis trajectories.

D.2 Fixed-Bank Credit Assignment

  • ATU consistently outperforms uniform credit and relevance-only ranking across all reported correctness and optimization metrics when memory contents are held identical.The replay varies cold-memory ranking while keeping hot memory identical; verifier outcomes supervise generated kernels, while adoption tracing attributes evidence to experiences used during synthesis.Source: §D.2 · p. 17, Table 7 · p. 17

E.1 Property 1: Bounded and normalized adoption credit.

  • Adopted experiences share the terminal episode score across the adoption set, so larger adoption sets do not receive more total credit; with δ = β = 0.2, negative item credit is bounded below by −0.2.The credit normalization and lower bound are described for the default δ and β values.Source: §E.1 · p. 18

E.2 Property 2: Stable but continually adaptive utility.

  • ATU’s nonzero step-size floor prevents utility estimates from freezing, while old evidence is forgotten geometrically; with ηmin = 0.05, the stated bound gives an approximately 13.5-update half-life.This lets retrieval-conditioned utility track changing task streams, memory, and generator context rather than converge to a fixed stationary value.Source: §E.2 · p. 18

E.3 Property 3: Controlled utility-aware reranking.

  • Utility modulates relevance during reranking, allowing comparable candidates to change order while preserving sufficiently large relevance gaps; negative evidence is treated conservatively because failure need not mean knowledge is harmful.Historically useful experience may still be inapplicable to a structurally different operator.Source: §E.3 · pp. 18–19

E.4 Property 4: Resident priority decays under persistent non-use.

  • Persistently unused resident knowledge loses priority relative to a persistently useful alternative; at the default μ = 0.1, the usage contribution’s half-life is approximately 6.6 memory updates.UGC’s consolidation is not irreversible promotion, allowing stale expertise to lose resident priority as operator patterns or optimization regimes change.Source: §E.4 · p. 19

F.1 Adoption-Aware Memory Use

  • For GatherV3, the agent adopted two of five retrieved experiences: one identified the host–device boundary failure, and the other supplied the applicable repair.The repair moved scalar reading and dependent work partitioning to the device side; the example argues that retrieval alone does not establish use.Source: §F.1 · p. 20, Table 8 · p. 20
  • For the Level-3 Mla operator, SAGE required 23 generation attempts to obtain its first correct kernel, then used the remaining budget for optimization.The operator fuses two QK paths, softmax, RoPE, and PV matrix multiplication.Source: §F.1 · p. 20

F.2 Memory Record Schemas

This appendix documents SAGE’s memory-record and hot-memory schemas, then describes its skill-based kernel-generation workflow and a generated sparse flash attention result. The recorded memory fields preserve experience history, while resident rules are selected under a token budget.

F.2 Memory Record Schemas

  • Each memory-item record stores a retrieval-facing rule, diagnostic evidence, ATU utility, retrieval and adoption counts, cross-operator adoption history, and lifecycle state.These fields remain in the cold store after promotion, preserving the item’s history; the example values are synthetic.Source: §F.2 · p. 20
  • A shared Matmul object changed its K dimension despite fixed tiling, so the trajectory records separating QK and PV configurations and declaring KERNEL_TYPE_MIX_AIC_1_0.The default task type also incorrectly assigned Cube work on an AIV core.Source: §F.2 · p. 21
  • Hot-memory snapshots select resident items under a token budget using utility per token, which combines estimated benefit, access frequency, and context cost.Before promotion, access frequency is estimated from retrievals over eligible episodes; resident items bypass retrieval, and their access estimates are updated from observed context usage.Source: §F.2 · p. 21, Table 11 · p. 22

F.3 Generated Kernel and Skill-Based Task Instruction

  • The generator skill builds and validates a complete Ascend C operator project; after correctness, the optimizer continues from the best feasible kernel using on-device performance feedback.The task instruction limits runs to 30 iterations, restricts code references to the skill’s examples, and directs debugging through skill references.Source: §F.3 · p. 22
  • 43.99× speedup over the reference torch_npu implementation was achieved by SAGE with GLM-5.3 on sparse flash attention during xCCL evaluation.The kernel passed correctness validation.Source: §F.3 · p. 22
  • The generated sparse-attention kernel gathers selected KV pages, computes query-key scores, applies softmax, and multiplies the probabilities by gathered values to produce output.The implementation batches page relays and pairs MTE2–MTE3 and MTE3–MTE2 events to order transfers and drain writes before reusing the buffer.Source: §F.3 · pp. 27–30
Loading 2609.35568v1…