Teaching AI to Tune Ascend NPUs: Fudan & Huawei XiaoYi Achieve 95.5% Kernel Generation Success

Introduction

Deploying models to Ascend NPUs requires not only correct computational logic but also low-level implementations that fully utilize hardware resources. The same attention computation can run at markedly different speeds depending on tiling, data movement, parallelism, and synchronization choices. While large language models can generate code and coding agents can invoke tools, compile, and iterate, they struggle on hardware platforms with limited public training data. Fudan University and Huawei's XiaoYi team address this with SAGE, a system that connects knowledge cold-start, tool-enhanced diagnosis, and experience learning for data-scarce NPU kernel synthesis.

Why Operator Development Needs Hardware Experience

Deep learning operators (matrix multiplication, normalization, attention) are executed by operator kernels — low-level implementations that decide where data resides, how much work each core handles, and how data movement and computation overlap. Engineers must judge:

Data tiling : tensors must be split into blocks that fit on-chip storage; too large exceeds capacity, too small increases movement and scheduling overhead.

Parallel partitioning : whether tasks can be distributed across more cores, whether load is balanced, and whether some cores are underutilized.

Data movement vs. computation : reducing redundant reads/writes of intermediate results and overlapping data preparation with computation.

Synchronization dependencies : distinguishing necessary waits for correctness from overly conservative barriers that limit pipelining.

These choices interact: changing tiling may require buffer reorganization; increasing parallelism may alter synchronization; reducing a data movement may demand adjusting subsequent computation reads. A runnable implementation is only the starting point for performance optimization.

Agent Development Workflow

The Coding Agent in this work writes code, calls tools, and iterates based on feedback. Operator development comprises four stages:

Generate implementation : organize host-side configuration and kernel-side code from I/O and compute logic.

Compilation repair : fix interface, type, template parameter, and configuration issues to produce an executable candidate.

Correctness verification : run on real device, compare with reference output, check numerical error and runtime exceptions.

Performance optimization : establish latency baseline for the correct version, locate bottlenecks, attempt faster implementations.

Each modification can reintroduce errors, so optimized candidates must recompile, re-verify, and re-measure. Performance gains must be built on passing correctness.

Existing Approaches and the Gap

Prior LLM-based operator development has moved from single-shot generation to repeated on-device verification. However, on platforms with little public code, model familiarity with mainstream platforms does not transfer. A cited study shows GPT-5.2 achieves 92% CUDA correctness on KernelBench Level-1 but only 14% on Ascend C — a 78-percentage-point gap. Moving to Ascend requires not just API usage but hardware knowledge: memory hierarchy, compute unit collaboration (AIC and AIV), and synchronization constraints.

Three Core Problems SAGE Targets

Insufficient hardware knowledge; raw documentation hard to turn into development experience. Models may know an interface name but not its role in a full implementation. Reading expert code does not imply understanding applicability conditions. Post-training needs abundant high-quality data, which is scarce. SAGE constructs initial knowledge and extracts rules from implementation comparisons (differential learning).

Test feedback insufficient to locate errors or guide optimization. Compilation logs and accuracy tests indicate failure but not which compute stage or data movement/sync dependency caused it. Latency numbers tell total time but not whether slowness comes from excessive data movement, insufficient parallelism, or long waits — each requiring different optimization actions. SAGE equips the agent with debug and profiling tools to form a hypothesize → gather diagnostic evidence → targeted modification → re-verify loop.

Memory system must accurately assess experience value and continuously self-evolve. Existing memory methods convert test results into scalar feedback for experience value scores. But how should one iteration's results be distributed among experiences based on actual usage? SAGE introduces Adoption-Traced Utility Estimation (ATU) and Utility-Gated Consolidation (UGC) to track adoption, assign credit only to adopted experiences, share experiences across operators, and promote frequently reused, compact rules into a limited hot memory.

Knowledge Cold-Start: From Documentation to Differential Mining

1. Three Purpose-Built Reference Documents

At cold start, no historical repair or tuning records exist. SAGE distills and manually curates three lightweight references from official Ascend C API docs, operator development tool guides, and best practices:

Ascend C Cheat Sheet : helps locate interfaces, basic constraints, and common usage, reducing basic knowledge gaps during development.

Debugging Guide : helps understand error types, select tools that can gather further evidence, and choose inspection methods.

Implementation Examples : provide complete operator organization references so the model understands how host-side configuration and kernel-side computation connect.

These three documents support interface usage, fault diagnosis, and full implementation organization respectively.

2. Differential Mining: Extracting Patterns Documentation Misses

After the agent generates simple Level-1 and Level-2 operators and they compile successfully, SAGE compares the generated implementation with expert implementations. From the differences, it extracts reusable implementation patterns and writes them into cold memory. Differential mining analyzes:

Generated vs. expert implementation : provides directly comparable code objects.

Differences and hardware reasons : analyzes why implementation organization, data handling, or execution dependencies differ.

Reusable patterns : turns observations on current code into experiences future development can consult.

Applicability conditions : states when a rule is worth referencing, avoiding treating one writing style as a universal answer.

For example, tiling differences are explained with on-chip capacity, movement granularity, and tail-block handling; synchronization differences are checked against data dependencies. Extracted patterns are grounded in actual code and hardware constraints.

Tool Invocation: From "Test Failed" to "Know What to Check Next"

1. Different Tools Answer Different Questions

Final output and total latency judge correctness and performance but cannot pinpoint causes. SAGE provides the agent with autonomous debug and profiling tools:

Debug tools answer "where might the error be?"

Profiling tools answer "where is time spent?"

They supplement test results but still require interpretation with code and hardware knowledge. Checking intermediate tensors distinguishes input reading, computation, and synchronization issues; analyzing execution flow judges parallel utilization, waits, and movement overhead. The agent calls tools autonomously, forming the closed loop mentioned above.

2. MLA Case Study: How Intermediate Results Locate Errors

MLA is a fused attention operator containing matrix computation, softmax, positional encoding, and subsequent matrix multiplication. It requires collaboration across compute units and guarantees that each step reads ready data. The initial implementation hit compute configuration and execution hang issues. The agent:

Separated configurations for different matrix computations, corrected compute task types.

Isolated single-token path for checking, then handled cross-core synchronization in loops. Hang resolved, but numerical inconsistency appeared.

Critical clue from DumpTensor showed softmax input was zero . Final output test only said the whole operator was wrong; intermediate results narrowed the problem to the stage before softmax.

Subsequent fixes involved:

Waiting for async matrix computation writes to complete, ensuring data availability for subsequent reads.

Handling non-contiguous input layout with AutoContiguous().

Completing positional encoding related dimensions to build full Q/K representations.

These changes jointly produced the first correct implementation. Intermediate results narrowed the problem scope to a specific compute stage, then combined input layout and sync dependencies to locate exact causes.

3. After Correctness, Diagnosis Shifts to Performance

With a correct implementation, the agent establishes a latency baseline and iterates around bottlenecks. MLA optimization proceeded along three directions:

Reduce unnecessary waits : replace overly conservative global barriers with event synchronization that satisfies dependencies.

Increase parallel utilization : distribute work across multiple cores by query token.

Reduce intermediate data movement : adjust subsequent matrix computation read paths to avoid an extra copy.

One double-buffering attempt broke movement dependencies and was reverted. Whether an optimization trick works depends on the current implementation's dependency graph. Diagnostic evidence, knowledge rules, and correctness tests must jointly participate in judgment. Verified fix paths, optimization practices, and applicability limits enter the experience library for future tasks.

Experience Learning: Evaluate Actual Usage, Retain Reusable Rules

SAGE stores full experiences in cold memory for on-demand retrieval; verified and distilled rules go into hot memory for direct provision to subsequent tasks — analogous to "detailed development notes" and "common development guidelines."

1. ATU: Which Experiences Actually Participated in Development?

Adoption-Traced Utility Estimation (ATU) combines actual adoption records with test results to update experience value estimates.

Record adoption first, then assign stage scores : For each retrieved experience, the agent logs "adopted/not adopted" with reasoning; evaluation requires complete records. A stage may contain multiple modifications; an experience's effect may only manifest in the final implementation, so scoring happens after stage end with deduplication. Let E be the set of retrieved experiences, A ⊆ E those adopted at least once. Correctness stage passing yields stage score S_c = 1; optimization stage improvement yields S_o = (latency_best - latency_first) / latency_first (relative latency reduction, enabling cross-scale comparison). Failure or no improvement uses S = -0.1 (paper sets -0.1). Negative feedback is kept small because stage failure or non-adoption does not prove an experience is harmful.

Feedback per experience : For each adopted experience, feedback = stage score / |A| (equal split so total feedback equals stage score, avoiding inflation just because more experiences were adopted). E.g., latency drops from 1 ms to 0.6 ms → score 0.4; two adopted experiences each get 0.2; non-adopted retrieved experiences each get -0.2.

Value score stability with adaptability : Experience value score V starts at 0, updated as weighted combination of historical score and current feedback: V_new = (1 - α) * V_old + α * feedback , where α = 1 / (1 + N ), N is deduplicated cumulative retrieval count per stage. As retrievals increase, update magnitude shrinks, reducing single-result volatility; paper sets minimum α = 0.05 so new feedback always retains at least 5% weight, allowing adjustment when task type or context changes. Scores stay within feedback range [-0.1, 1], providing numeric bounds for retrieval re-ranking.

Value influences retrieval but relevance remains primary : System first obtains candidate pool by relevance, then re-ranks by value score. Relevance answers "is this related to the current problem?"; value estimation answers "does actual usage evidence support prioritizing this?" Adoption records show an experience participated in diagnosis or modification but cannot alone prove independent contribution; they must combine with stage results for evaluation basis.

2. UGC: Which Experiences Deserve to Become Standing Guidelines?

Utility-Gated Consolidation (UGC) decides which experiences enter the resident context and when to yield space.

Require cross-operator evidence, then measure value per context length : An experience must have positive value score, be adopted at least 3 times , and cover at least 3 distinct operators to become a consolidation candidate. Repeated adoption reduces chance; cross-operator coverage avoids treating multiple rounds within one task as generality evidence. The agent then distills a concise rule; source experience and rule share identifier and value score.

Compete for limited resident budget : Candidates ranked by value score × estimated access frequency / rule token length . Frequency estimated as "retrieval count / eligible stages" during cold memory phase. Multiplying by frequency prioritizes knowledge often needed; dividing by length accounts for context cost of continuously providing the rule. Equally effective and frequently used, more compact rules are better suited for residence.

Fixed capacity, evict long-unused rules : Resident budget set to 10K tokens , limiting hot memory growth. Each update merges existing rules with new candidates, sorts by the above metric, rebuilds from empty set: add if fits, skip if not, continue checking later candidates. Thus new and old rules compete on same standard; shorter candidates also get a chance. Resident rules bypass retrieval; access frequency updates via actual usage: freq_new = (1 - β) * freq_old + β * I_used , where I_used is 1 if rule used in current stage else 0. β = 0.1 retains 90% history, adds 10% current signal, smoothing recent changes. After 7 consecutive updates without use, frequency estimate drops to ~48% of original; with unchanged value score and length, resident priority drops proportionally.

Replaced rules return to retrievable state; source experiences and history remain.

ATU evaluates usage effect; UGC combines reuse frequency and text cost to allocate resident space, forming bidirectional updates between cold and hot memory. All parameters set uniformly based on their numeric roles and preliminary system stability; same config for all base models and operators, no per-model or per-operator tuning.

Experimental Setup and Results: Where Performance Gains Appear

1. Hardware, Budget, and Measurement

Experiments run on Ascend 910C ; candidate operators undergo real hardware compilation, correctness, and performance verification. Key config:

Device resources : 8 accelerators per server, 16 accelerator dies total; each operator evaluation worker gets one card.

Software environment : ARM64, Huawei Cloud EulerOS 2.0, CANN 8.5.1, driver 25.5.1, Python 3.11.14, torch_npu 2.9.0.

Model service : vllm-ascend 0.18.0rc1 local deployment; DeepSeek-V4-Flash uses DP=4, TP=4; GLM-5.2/5.3 use TP=16.

Interaction budget : Claude Code orchestrates development flow; max 30 rounds per operator, shared between correctness repair and performance optimization.

Latency measured via device-side NPU Event: warm-up 5 runs, timed 100 runs, report average ; candidates and torch_npu use same inputs and timing boundaries. NPUKernelBench correctness requires max absolute error and max relative error both ≤ 10⁻³ .

2. MLA: From 0.420 ms to 0.155 ms

In MLA's development trajectory, the agent obtained the first correct version after multiple repair rounds, then used remaining budget to optimize. Self-acceleration from first correct to best version: 2.71× , latency reduction 63.1% . This trajectory embodies combined gains from reducing waits, increasing parallelism, and reducing data movement.

3. Main Experiment: More Correct Implementations, More Implementations Faster Than Reference

Main experiment covers 88 operators from NPUKernelBench. Four configurations share base model, initial knowledge, tools, interaction budget, and evaluation environment; comparison focuses on memory mechanism:

Refinement : iterates within current operator, no cross-task experience retention.

Static RAG : retrieves fixed knowledge base, no further experience writes or value updates.

Value Memory : references EvoKernel (ICML 2026) value-driven memory, continuously accumulates experiences, uses same test score target for same-round retrieved experiences, then adjusts retrieval ranking with value scores.

SAGE : combines actual adoption to evaluate experiences, and selects/updates resident rules.

Metrics answer three questions: ER (Execution Rate) — how many operators solved correctly within budget; Speedup vs torch_npu — fraction of solved operators faster than torch_npu; Self-optimization acceleration — median speedup from first correct to best version.

SAGE correctly solves 84/88 operators (95.5%). Relative to Value Memory, ER improves by 11.2 percentage points ; self-optimization acceleration median rises from 3.36× to 4.47× . This compares memory mechanisms within a unified framework; cold-start knowledge and diagnostic tools not ablated separately.

4. Backbone Capability: Can Experience Translate to Performance Gains?

External memory supplies hardware knowledge; the base model (backbone) still handles diagnostic understanding, code modification, and optimization selection. With identical SAGE mechanism and interaction budget, three models show significant differences:

Level-3 operators solved: 1/18, 8/18, 16/18

Self-optimization acceleration median: 1.27×, 2.08×, 4.47×

Agent must understand tool feedback, judge experience applicability conditions, and coordinate layout, movement, and sync modifications — all still dependent on model capability. Earlier correctness leaves more optimization rounds within the shared 30-round budget, but experiment does not isolate this factor's contribution.

Cross-model memory transfer further separates knowledge from model capability. In read-only re-evaluation on the same 88 operators, receiving models use GLM-5.2's accumulated memory without further memory updates:

DeepSeek-V4-Flash : ER from 64.8% → 72.7% ; speedup-vs-torch_npu from 64.9% → 71.9% .

Qwen3-Coder-Next : ER from 31.8% → 40.9% ; speedup-vs-torch_npu from 28.6% → 38.9% .

Transfer gains show stronger models' accumulated experiences have reuse value; but receiving models still fall short of GLM-5.2, indicating experience quality and ability to utilize experience jointly determine outcomes . Engineering selection should weigh model serving cost, correctness coverage, and optimization magnitude. This reflects specific models' combined differences, not disentangled parameter scale or training data effects.

5. Task Ordering: Is Accumulating Simpler Task Experience First More Effective?

Paper fixes 25 Level-2 operators and all 18 Level-3 operators , compares three task schedules. All groups start from same initial memory, use same per-operator budget, update memory after each task without re-solving; results only count Level-3, recorded before corresponding experience writes.

Mixed execution and sequential execution use same 43 tasks , so latter two directly reflect ordering effect: completing Level-2 first yields Level-3 ER 16.7 percentage points higher than mixed, and 7.8 percentage points higher than Level-3-only. Level-3-only group lacks extra Level-2 development process. Results support accumulating lower-difficulty experience first, then tackling complex operators; this is an independent task flow, statistical scope differs from main experiment.

6. More Complex Tasks and Mechanism Validation

On 16 more complex xCCL operators , system inherits existing memory and continues updating. With GLM-5.3, correctly solved count rises from GLM-5.2's 11/16 to 13/16 . On the 11 operators both solve , median speedup vs torch_npu improves from 3.32× to 7.32× ; sparse flash attention operator jumps from 23.35× to 43.99× . This shows differences in both coverage and optimization quality on same tasks. GLM-5.3 median acceleration across all 13 solved operators is 10.86× (different statistical set from paired result).

Ablation under GLM-5.2 on 88 operators: full SAGE ER = 86.9% ; without UGC = 82.7% ; random resident rule selection = 78.8% ; uniform stage score for all retrieved experiences = 79.2% . These support evaluating experience value via actual adoption and evidence-based resident rule selection.

Gains still vary by task. GLM-5.2 on correctly passing Level-3/GMM tasks shows speedup vs torch_npu ranging 0.38×–71.41× , including some slower than reference. Thus effectiveness should be judged by coverage ratio and typical optimization magnitude; single-operator results need further validation for end-to-end model benefits.

Trade-offs and Future Work

SAGE combines initial hardware knowledge, tool-enhanced diagnosis, and experience updates, enabling agents to continuously optimize Ascend operators and accumulate reusable knowledge through self-iteration. For production deployment, several directions need advancement:

More complex fused operators and longer tuning pipelines : joint exploration of layout, scheduling, movement, and sync; retain multiple candidate paths to avoid premature fixation. Current 30-round setting does not cover long-horizon complex fused kernel development.

Traceable shared experience library : store rule provenance, code changes, and measurement records for multi-person collaboration; explicit applicability conditions; handle software version changes, rule conflicts, and stale knowledge.

More input and hardware configurations : extend optimization across shapes, dtypes, and chip architectures; verify which old experiences reuse and which need re-validation.

Full development cost accounting : simultaneously measure model calls, compilation, device evaluation, and memory maintenance time and resources; assess performance gain vs. development investment.

End-to-end effect and toolchain integration : integrate generated operators into real models and existing workflows; measure overall latency, throughput, tail latency; retain performance fallback paths.