TrainBridge: Agentic NPU Training Optimization Guided by GPU Compilation Evidence

Yupeng Chen; Kunxiang Zhao; Xingjian Wang; Yue Zhao; Yijing Wu; Ye Liu; Kai Li; Congliang Chen

cs.CLOct 21, 2026

TrainBridge: Agentic NPU Training Optimization Guided by GPU Compilation Evidence

Authors: Yupeng Chen, Kunxiang Zhao, Xingjian Wang, Yue Zhao, Yijing Wu, Ye Liu, Kai Li, Congliang Chen,

Organizations: Shenzhen Loop Area Institute · The Chinese University of Hong Kong, Shenzhen, China · Autokernel

Abstract

We study how to improve training efficiency on Ascend NPUs. In the NPU training programs we study, framework-level compilation optimizations such as operator fusion are only partially applied, leaving fragmented computation and avoidable execution overheads. Profiling-guided tuning can address these inefficiencies, but it often requires costly trial and error. To make training on NPUs more efficient, we propose TrainBridge, an agentic approach for transferring GPU optimization guidance to NPU training. TrainBridge uses the highly optimized execution produced by mature GPU compilation stacks as a task-specific reference for NPU tuning. It extracts optimization intents from the GPU training execution to guide NPU program rewrites and further refines the rewritten program through iterative profiling. Our GPU-to-NPU migration experiments show that TrainBridge achieves an average 1.7x speedup across four workloads covering LLM pretraining, image classification, and vision-language training, and an average speedup of nearly 2.0x across 27 Modded-NanoGPT speedrun records with highly optimized GPU implementations.

Figures & tables

Implementation	Training time (s)	Relative time
GPU, compiled	128	$1.00\times$
GPU, eager	338	$2.64\times$
NPU, initial port	563	$4.40\times$

Table 1: Training times for GPU and NPU implementations of Modded-NanoGPT Record 50. Ratios use the compiled GPU program as the reference.

Implementation	Latency (ms)	Speedup
Original NPU implementation	2.604	$1.00\times$
GPU-compilation-evidence-guided agent rewrite	0.508	$5.13\times$

Table 2: Local NPU latency of KeyShift operation in Record 50 of Modded-NanoGPT speedrun before and after GPU-compilation-evidence-guided agent rewriting. The input shape is [1, 24576, 6, 128], and we use BF16 precision.

Figure 1: Overview of TrainBridge. Stage 1 aligns eager-to-compiled GPU changes with the migrated NPU program, and uses the resulting optimization evidence to guide NPU implementation optimization. Stage 2 iteratively profiles the updated NPU program and addresses residual bottlenecks.

Figure 2: Record 50 of the Modded-NanoGPT speedrun tuning results under a 12-hour budget. Results are averaged over three independent runs using the same Codex harness; error bars show standard deviation. TrainBridge achieves substantially lower end-to-end training time with lower variance across repeated runs.

Figure 3: Contribution of the two optimization stages on Record 50 under a 12-hour tuning budget. Stage 2 alone uses NPU profiling as task-specific guidance, while TrainBridge additionally applies GPU-evidence-guided Stage 1 before profiling-guided Stage 2. GPU evidence provides a substantial early gain, and subsequent NPU profiling further improves the program.

Figure 4: End-to-end training time (top) and training time relative to the GPU baseline (bottom) for ten Modded-NanoGPT records sampled approximately uniformly across the 27 evaluated records when ordered by leaderboard record number. TrainBridge consistently reduces NPU training time across the shown records; complete results for all 27 evaluated records are reported in Appendix C .

Figure 5: Training time relative to the compiled GPU baseline across four additional workloads. Lower is better; the dashed line indicates equal training time on GPU and NPU.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Record 50 of the Modded-NanoGPT speedrun tuning results under an 8-hour budget using Claude Opus 4.8 with the Claude Code harness. TrainBridge achieves substantially lower end-to-end training time than Claude Opus 4.8 alone.

Record	GPU (s)	Initial NPU (s)	Tuned NPU (s)	NPU speedup	NPU/GPU before $\rightarrow$ after
50	128	563	257	$2.2\times$	$4.4\times\rightarrow 2.0\times$
49	129	569	278	$2.0\times$	$4.4\times\rightarrow 2.2\times$
48	130	574	277	$2.1\times$	$4.4\times\rightarrow 2.1\times$
47	132	542	278	$2.0\times$	$4.1\times\rightarrow 2.1\times$
46	132	574	309	$1.9\times$	$4.4\times\rightarrow 2.3\times$
45	134	579	293	$2.0\times$	$4.3\times\rightarrow 2.2\times$

Appendix

Table 3: Modded-NanoGPT per-record results. Times are end-to-end training times in seconds. NPU speedup is initial NPU time divided by optimized NPU time. Ratios and averages are computed before rounding to one decimal place. All 27 records with complete times satisfy the validation-loss requirement.

Neural Network Hardware Acceleration

See all 48 papers

Forge-UGC: FX optimization and register-graph engine for universal graph compiler
Apr 14, 2026Satyam Kumar, Saurabh JhaCompiler Optimization AI Accelerator Inference
AscendOptimizer: Episodic Agent for Ascend NPU Operator Optimization
Mar 24, 2026Jiehao Wu, Zixiao Huang, Wenhao Li +3GPU Kernel Optimization Black-Box Optimization
School of Computer Science and Technology, East China Normal University · School of Computer Science and Technology, Tongji University · Shanghai University of International Business and Economics +3
From Experience to Expertise: Adoption-Aware Memory Learning for Data-Scarce NPU Kernel Synthesis
Sep 28, 2026Longxiao Fan, Tao Zhang, Han Yan +5LLM-Based Program Synthesis Agent Skill Learning
Fudan University · University of Science and Technology of China · Huawei Technologies Ltd.

Apr 14, 2026J/K move · Enter open · S save

Forge-UGC: FX optimization and register-graph engine for universal graph compiler

Satyam Kumar, Saurabh Jha

Compiler Optimization AI Accelerator Inference Efficient Transformer Inference

Open paper AskPDF arXiv

We present Forge-UGC (FX Optimization and Register-Graph Engine for Universal Graph Compilation), a four-phase compiler for transformer deployment on heterogeneous accelerator hardware, validated on Intel AI Boost NPU. Existing frameworks such as OpenVINO and ONNX Runtime often use opaque compilation pipelines, limited pass-level visibility, and weak buffer management, which can lead to higher compilation cost and runtime overhead. Forge-UGC addresses this with a hardware-agnostic design that separates graph capture, optimization, intermediate representation lowering, and backend scheduling. Phase 1 captures graphs with torch.export at the ATen operator level, supporting modern transformer components such as rotary position embeddings, grouped-query attention, and SwiGLU without manual decomposition. Phase 2 applies six optimization passes: dead code elimination, common subexpression elimination, constant folding, attention fusion, operator fusion, and layout optimization, reducing graph node count by 14.2 to 21.9%. Phase 3 lowers the optimized graph into a typed intermediate representation with explicit virtual register assignments. Phase 4 performs liveness analysis, linear-scan buffer allocation, reducing peak buffer count by 30 to 48%, and device-affinity scheduling, reducing NPU-CPU transitions by 42 to 65%. Across six model families ranging from 125M to 8B parameters, evaluated on WikiText-103 and GLUE, Forge-UGC delivers 6.9 to 9.2x faster compilation than OpenVINO and ONNX Runtime, 18.2 to 35.7% lower inference latency, and 30.2 to 40.9% lower energy per inference. Fidelity is preserved, with max absolute logit differences below 2.1e-5 and KL divergence below 8.4e-9. We also introduce Fusion Gain Ratio, Compilation Efficiency Index, and per-pass execution profiling for systematic evaluation of NPU compilation pipelines.

Abstract

Neural Network Hardware Acceleration

Explore similar work

Forge-UGC: FX optimization and register-graph engine for universal graph compiler

AscendOptimizer: Episodic Agent for Ascend NPU Operator Optimization

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