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

Satyam Kumar; Saurabh Jha

cs.ARApr 14, 2026

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

Authors: Satyam Kumar, Saurabh Jha

Abstract

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.

Figures & tables

Figure 1: FORGE-UGC four-phase architecture. Phase 1 : FX graph capture via torch.export with tied weight resolution. Phase 2 : Six composable optimization passes (DCE, CSE, constant folding, attention fusion, operator fusion, layout optimization) with optional autotuning. Phase 3 : Lowering to NPUIR with typed instructions, virtual registers, and device placement. Phase 4 : Liveness analysis, linear-scan buffer allocation, instruction scheduling, and code generation producing a CompiledNPUExecutor .

Symbol	Description
$G=(V,E)$	FX computation graph
$V$	Set of graph nodes (operations)
$E$	Set of data-dependency edges
$n_{\text{before}},n_{\text{after}}$	Node count before/after optimization
$\mathcal{P}=\{p_{1},\ldots,p_{K}\}$	Set of optimization passes
$\tau(p_{k})$	Execution time of pass $p_{k}$ (ms)

Table 1: Notation used in the FORGE-UGC methodology.

Component	Specification
CPU	Intel Core Ultra 9 285HX
NPU	Intel AI Boost (11 TOPS INT8)
NPU Driver	32.0.100.4514
NPU Memory	Shared LPDDR5, 72.7 GB
NPU Location	PCI bus 0, device 11, function 0
GPU	NVIDIA RTX PRO 5000 Blackwell

Table 2: Hardware platform specifications.

Model	Params	Hidden	Layers	Attn	Precision
GPT-2	125M	768	12	MHA	fp16
Granite-350M	350M	1024	24	MHA	fp16
Qwen2-0.5B	500M	1024	24	GQA	fp16
Llama-3.2-1B	1.0B	2048	16	GQA	fp16
LFM2-2.6B	2.6B	2560	32	MHA	fp16
Llama-3.1-8B	8.0B	4096	32	GQA	W-int8/ A-fp16

Table 3: Model specifications used in experiments. The Precision column indicates the numerical precision used during NPU inference. Models

\leq

2.6B use fp16 weights; Llama-3.1-8B uses NNFactory’s built-in symmetric int8 weight quantization with fp16 activations to fit within the NPU’s memory and compute constraints (see Section 6.2.1 ).

Figure 2: Compilation time comparison on GPT-2 (125M). FORGE-UGC compiles in 1,000ms versus 6,930ms (OpenVINO) and 7,271ms (ONNX Runtime)—a $6.9\times$ and $7.3\times$ speedup respectively.

Model	FORGE	OpenVINO	ONNX RT
GPT-2 (125M)	1,000	6,930	7,271
Granite-350M	1,420	9,840	10,320
Qwen2-0.5B	1,680	11,240	11,890
Llama-3.2-1B	2,340	18,520	19,840
LFM2-2.6B	3,850	32,160	34,280
Llama-3.1-8B	6,720	58,430	62,150

Table 4: Compilation time (ms) across model families. FORGE-UGC achieves 6.9–9.2

\times

speedup over both baselines consistently. We note that 78% of FORGE-UGC’s compilation time is spent in torch.export graph capture (Section 7.2 ), which is shared upstream infrastructure; the FORGE-UGC-specific optimization and backend phases account for only 22% of total time.

Figure 3: Compilation phase breakdown for GPT-2 (125M). FX Capture dominates at 773.5ms (78.4%), while all six optimization passes complete in 208ms (21.1%). IR lowering, buffer allocation, and scheduling together require only 8ms (0.8%).

Figure 4: Graph node count after each optimization pass for GPT-2 (125M). Attention fusion provides the largest single reduction (403 $\rightarrow$ 344, $-$ 14.6%), followed by operator fusion (344 $\rightarrow$ 332, $-$ 3.5%). Total reduction: 403 $\rightarrow$ 333 ( $-$ 17.4%).

Model	Initial	Final	$\Delta$	%
GPT-2 (125M)	403	333	$-$ 70	$-$ 17.4
Granite-350M	782	636	$-$ 146	$-$ 18.7
Qwen2-0.5B	804	652	$-$ 152	$-$ 18.9
Llama-3.2-1B	562	468	$-$ 94	$-$ 16.7
LFM2-2.6B	1,068	834	$-$ 234	$-$ 21.9
Llama-3.1-8B	1,124	896	$-$ 228	$-$ 20.3

Table 5: Graph node reduction across model families. Attention fusion consistently provides the largest reduction.

Model	PPL Pre	PPL Post	Max-Abs $\Delta$ Logit	KL Div
GPT-2 (125M)	29.41	29.41	$6.2\mathrm{e}{-6}$	$1.8\mathrm{e}{-10}$
Granite-350M	16.22	16.22	$8.4\mathrm{e}{-6}$	$3.2\mathrm{e}{-10}$
Qwen2-0.5B	14.87	14.87	$7.1\mathrm{e}{-6}$	$2.7\mathrm{e}{-10}$
Llama-3.2-1B	9.76	9.76	$9.8\mathrm{e}{-6}$	$4.1\mathrm{e}{-10}$
LFM2-2.6B	7.52	7.52	$1.2\mathrm{e}{-5}$	$6.3\mathrm{e}{-10}$
Llama-3.1-8B	6.24	6.24	$2.1\mathrm{e}{-5}$	$8.4\mathrm{e}{-9}$

Table 6: Numerical fidelity analysis. PPL Pre/Post : perplexity before and after FORGE-UGC compilation on WikiText-103. Max-Abs

\Delta

Logit : maximum absolute difference between pre- and post-compilation logits across 1,000 test sequences. KL Div : KL divergence between pre- and post-compilation output distributions. For models using fp16 dispatch (125M–2.6B), the near-zero differences confirm that the graph-level optimization passes are semantics-preserving within floating-point rounding. For Llama-3.1-8B (W-int8/A-fp16 dispatch), the slightly higher but still negligible max-abs diff reflects int8 weight quantization at the NNFactory dispatch level, not the optimization passes themselves.

	FORGE-UGC (ms)				OpenVINO (ms)				ONNX Runtime (ms)
Model	Mean	P50	P90	P99	Mean	P50	P90	P99	Mean	P50	P90	P99
GPT-2 (125M)	6.82	6.74	7.31	8.12	8.45	8.38	9.12	10.84	9.13	9.02	9.87	11.52
Granite-350M	9.41	9.32	10.08	11.24	12.67	12.54	13.81	15.93	13.28	13.12	14.55	16.87
Qwen2-0.5B	11.83	11.72	12.64	14.03	15.42	15.28	16.91	19.47	16.21	16.04	17.68	20.35
Llama-3.2-1B	18.24	18.08	19.52	21.67	24.81	24.62	27.16	31.28	26.37	26.14	28.89	33.23
LFM2-2.6B	31.56	31.28	33.74	37.48	45.23	44.87	49.52	57.01	48.14	47.72	52.69	60.64

Table 7: End-to-end inference latency (ms) on WikiText-103. FORGE-UGC achieves 18.2–35.7% lower mean latency than both baselines across all model families. Latency measured as wall-clock time from input tensor to output logits over 50 warmup-excluded iterations, averaged over 3 runs. Raw per-run data is provided in Appendix Table 25 .

	FORGE-UGC (ms)				OpenVINO (ms)				ONNX Runtime (ms)
Model	Mean	P50	P90	P99	Mean	P50	P90	P99	Mean	P50	P90	P99
GPT-2 (125M)	5.94	5.88	6.37	7.08	7.36	7.30	7.95	9.44	7.95	7.86	8.59	10.03
Granite-350M	8.21	8.13	8.79	9.80	11.04	10.93	12.04	13.88	11.57	11.44	12.68	14.71
Qwen2-0.5B	10.32	10.22	11.02	12.23	13.44	13.32	14.74	16.97	14.13	13.98	15.41	17.74
Llama-3.2-1B	15.91	15.76	17.02	18.89	21.62	21.46	23.68	27.26	22.98	22.78	25.18	28.96
LFM2-2.6B	27.52	27.28	29.42	32.68	39.43	39.10	43.15	49.67	41.97	41.60	45.93	52.86

Table 8: End-to-end inference latency (ms) on GLUE (SST-2 + MNLI, batch=1, seq=128). Results confirm that FORGE-UGC’s gains are task-agnostic.

Model	FORGE (mJ)	OV (mJ)	ONNX (mJ)	$\Delta$ OV	$\Delta$ ONNX
GPT-2 (125M)	69.6	99.7	110.5	$-$ 30.2%	$-$ 37.0%
Granite-350M	96.0	149.5	160.7	$-$ 35.8%	$-$ 40.3%
Qwen2-0.5B	120.7	181.9	196.1	$-$ 33.6%	$-$ 38.4%
Llama-3.2-1B	186.0	292.8	319.1	$-$ 36.5%	$-$ 41.7%
LFM2-2.6B	322.0	533.7	582.5	$-$ 39.7%	$-$ 44.7%
Llama-3.1-8B	637.3	1,078.2	1,183.6	$-$ 40.9%	$-$ 46.2%

Table 9: Energy consumption per inference (mJ) on WikiText-103. FORGE-UGC achieves 30.2–40.9% lower energy than OpenVINO and 37.0–46.2% lower energy than ONNX Runtime. Energy savings exceed latency savings because FORGE-UGC also reduces average active power through fewer device transitions and pre-allocated buffers.

Pass	Time (ms)	$\Delta$ Nodes
DCE	7	0
CSE	9	0
Constant Folding	11	0
Device Constant	21	+1
Attention Fusion	38	$-$ 59
Operator Fusion	72	$-$ 12

Table 10: Per-pass execution time (ms) on GPT-2 (125M). Operator fusion is the most expensive pass (72ms) but provides the second-largest node reduction. All six passes complete in 208ms total.

Model	Layers	Opt. Time	Attn. Fusion
GPT-2	12	208	38
Granite-350M	24	382	71
Qwen2-0.5B	24	396	74
Llama-3.2-1B	16	284	52
LFM2-2.6B	32	548	102
Llama-3.1-8B	32	572	108

Table 11: Total optimization time (ms) and attention fusion time (ms) across model families. Optimization time scales linearly with layer count.

Model	Score ( $\alpha$ =0)	Score ( $\alpha$ =1)	FGR
GPT-2 (125M)	364.87	8.64	42.3
Granite-350M	718.42	14.82	48.5
Qwen2-0.5B	742.18	15.24	48.7
Llama-3.2-1B	1,246.34	24.68	50.5
LFM2-2.6B	2,482.61	38.92	63.8
Llama-3.1-8B	4,218.47	62.14	67.9

Table 12: Fusion Gain Ratio (FGR) across model families. FGR is a cost-model-internal diagnostic : it measures the ratio of cost model scores with fusion disabled (

\alpha=0

) versus fully enabled (

\alpha=1

). Because the cost model is a heuristic weighted sum (Eq. 18 ) that is not calibrated to wall-clock time, FGR values should not be interpreted as latency speedups. The corresponding measured latency gains from fusion are 16.6–29.6% (Table 15 ).

Model	CEI ${}_{\text{OV}}$	CEI ${}_{\text{ONNX}}$	Compile (s)
GPT-2 (125M)	1.239	1.339	1.00
Granite-350M	0.948	0.994	1.42
Qwen2-0.5B	0.776	0.815	1.68
Llama-3.2-1B	0.581	0.618	2.34
LFM2-2.6B	0.372	0.396	3.85
Llama-3.1-8B	0.218	0.233	6.72

Table 13: Compilation Efficiency Index (CEI) across model families.

\text{CEI}_{\mathcal{B}}=(L_{\mathcal{B}}/L_{\text{FORGE}})/T_{\text{compile}}^{(s)}

where

T_{\text{compile}}^{(s)}

is compilation time in seconds. Higher CEI means more latency-speedup ratio per second of compilation. CEI is most informative for iterative development and JIT scenarios; for compile-once-run-millions production deployment, absolute latency (Table 7 ) is the primary metric.

Configuration	Cost Score	$\Delta$ vs. Full
All Passes	8.64	—
w/o DCE	8.69	+0.6%
w/o CSE	8.64	+0.0%
w/o Constant Folding	8.64	+0.0%
w/o Device Constant	8.64	+0.0%
w/o Attention Fusion	238.34	+2,658%

Table 14: Pass ablation on GPT-2 (125M). Each row removes one pass from the full pipeline. Attention fusion removal causes a 27.6

\times

cost model degradation, confirming it as the most critical pass.

Model	Layers	w/ Fusion	w/o Fusion	$\Delta$ %
GPT-2	12	6.82	8.18	$-$ 16.6%
Granite-350M	24	9.41	12.32	$-$ 23.6%
Qwen2-0.5B	24	11.83	15.48	$-$ 23.6%
Llama-3.2-1B	16	18.24	22.71	$-$ 19.7%
LFM2-2.6B	32	31.56	44.18	$-$ 28.6%
Llama-3.1-8B	32	62.48	88.72	$-$ 29.6%

Table 15: Attention fusion impact across model families. Measured wall-clock latency reduction scales with model depth as more attention blocks are fused. These measured values complement the cost-model-based FGR diagnostic (Table 12 ).

Model	V-Regs	Phys.	$\rho_{\text{buf}}$	Trans. $\downarrow$
GPT-2	333	218	34.5%	42.1%
Granite-350M	636	412	35.2%	48.3%
Qwen2-0.5B	652	418	35.9%	49.1%
Llama-3.2-1B	468	324	30.8%	44.7%
LFM2-2.6B	834	478	42.7%	58.2%
Llama-3.1-8B	896	468	47.8%	64.8%

Table 16: Buffer allocation statistics across model families. Linear-scan allocation reduces physical buffer count by 30–48% through liveness-guided reuse.

$\alpha$	Cost Score	Nodes	Fused Ops
0.0	364.87	403	0
0.2	142.31	392	4
0.4	58.42	378	8
0.6	22.18	358	12
0.8	10.84	342	18
1.0	8.64	333	24

Table 17: Fusion aggressiveness

\alpha

sensitivity on GPT-2 (125M).

\alpha=1.0

yields the best cost model score; aggressive fusion consistently helps.

Model	Default	Autotuned	$\Delta$ %
GPT-2	8.64	8.28	$-$ 4.2%
Granite-350M	14.82	13.92	$-$ 6.1%
Qwen2-0.5B	15.24	14.18	$-$ 7.0%
Llama-3.2-1B	24.68	22.84	$-$ 7.5%
LFM2-2.6B	38.92	35.74	$-$ 8.2%
Llama-3.1-8B	62.14	56.72	$-$ 8.7%

Table 18: Autotuning vs. default configuration across model families. Autotuning improves cost model score by 4.2–8.7% while adding

<

200ms to compilation.

Metric	Mean	Std Dev	CV (%)
Compilation Time (ms)	1,000	18.4	1.84
Inference Latency (ms)	6.82	0.14	2.05
P99 Latency (ms)	8.12	0.19	2.34
Node Reduction (%)	17.4	0.0	0.00
FGR	42.3	0.82	1.94

Table 19: Variance across 10 independent runs on GPT-2 (125M) with FORGE-UGC. CV

<

2.5% across all metrics confirms high reproducibility. The low latency variance reflects the pre-allocated buffers and deterministic scheduling of the compiled executor, which eliminates runtime allocation jitter. We acknowledge that CV

<2.5\%

is tighter than typical for shared-memory NPU systems; raw per-run data supporting these statistics is provided in Appendix Table 25 .

	Mean Latency (ms)			$\Delta$ vs. Best Baseline		FORGE-UGC Metrics
Model	FORGE	OV	ONNX	Latency	Compile	FGR	CEI ${}_{\text{OV}}$	$\rho_{\text{buf}}$
GPT-2 (125M)	6.82	8.45	9.13	$-$ 19.3%	$-$ 85.6%	42.3	1.24	34.5%
Granite-350M	9.41	12.67	13.28	$-$ 25.7%	$-$ 85.6%	48.5	0.95	35.2%
Qwen2-0.5B	11.83	15.42	16.21	$-$ 23.3%	$-$ 85.1%	48.7	0.78	35.9%
Llama-3.2-1B	18.24	24.81	26.37	$-$ 26.5%	$-$ 87.4%	50.5	0.58	30.8%
LFM2-2.6B	31.56	45.23	48.14	$-$ 30.2%	$-$ 88.0%	63.8	0.37	42.7%

Table 20: Comprehensive performance summary across all model families on WikiText-103. FORGE-UGC achieves the best latency, fastest compilation, highest FGR, and highest CEI across all models. The “

\Delta

vs. Best Baseline” columns report improvement over whichever baseline (OpenVINO or ONNX RT) performs best on each model.

Model	$\delta_{\text{before}}$	$\delta_{\text{after}}$	$\Delta$ %	Lat. $\Delta$ %
GPT-2	86	50	$-$ 41.9%	$-$ 4.2%
Granite-350M	168	87	$-$ 48.2%	$-$ 5.8%
Qwen2-0.5B	174	89	$-$ 48.9%	$-$ 6.1%
Llama-3.2-1B	126	70	$-$ 44.4%	$-$ 5.1%
LFM2-2.6B	248	104	$-$ 58.1%	$-$ 8.4%
Llama-3.1-8B	264	93	$-$ 64.8%	$-$ 11.2%

Table 21: Instruction scheduling impact across model families. Device transitions (

\delta

) are reduced by 42–65%, directly translating to latency improvement.

Model	FORGE	OV	ONNX RT
GPT-2	1.20	1.29	1.28
Granite-350M	1.21	1.27	1.29
Qwen2-0.5B	1.20	1.27	1.27
Llama-3.2-1B	1.20	1.27	1.27
LFM2-2.6B	1.20	1.27	1.27
Llama-3.1-8B	1.20	1.27	1.27

Table 22: P99/P50 latency ratio across frameworks. FORGE-UGC exhibits the tightest tail distribution (ratio closest to 1.0), critical for SLA-bound deployments.

Feature	TVM	XLA	IREE	torch.compile	OpenVINO	ONNX RT	NNAL	FORGE-UGC
Direct PyTorch FX input	No	No	No	Yes	No	No	No	Yes
No lossy export required	No	Partial	No	Yes	No	No	N/A	Yes
Pass-level visibility	Partial	No	Yes	Partial	No	No	N/A	Yes
Attention fusion (NPU)	No	Yes	Partial	No	No	No	No	Yes
Operator fusion (NPU)	Partial	Yes	Partial	GPU/CPU only	Partial	Partial	No	Yes
NPU-aware autotuning	Yes	No	No	No	No	No	No	Yes

Table 23: FORGE-UGC vs. existing frameworks across key compiler capabilities. FORGE-UGC is the only framework providing pass-level visibility, NPU-aware autotuning, explicit buffer allocation, and direct PyTorch FX integration for Intel NPU targets. torch.compile (Inductor backend) targets CPU/GPU and lacks NPU dispatch, liveness-aware allocation, and NNFactory integration. IREE provides composable MLIR passes and multi-backend support but requires torch-mlir conversion and has no Intel NPU backend.

Metric	WikiText-103	GLUE	Std Dev
$\Delta$ vs. OpenVINO	$-$ 26.4%	$-$ 26.1%	0.21%
$\Delta$ vs. ONNX RT	$-$ 30.2%	$-$ 29.8%	0.28%
$\Delta$ Compile Time	$-$ 85.6%	$-$ 85.6%	0.00%

Table 24: Cross-dataset robustness: mean latency improvement (%) of FORGE-UGC over baselines. Standard deviation

<1.2\%

confirms task-agnostic improvements.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Model	Run 1	Run 2	Run 3	Mean
GPT-2 (125M)	6.78	6.84	6.83	6.82
Granite-350M	9.38	9.44	9.40	9.41
Qwen2-0.5B	11.79	11.86	11.84	11.83
Llama-3.2-1B	18.19	18.28	18.24	18.24
LFM2-2.6B	31.48	31.62	31.57	31.56
Llama-3.1-8B	62.31	62.58	62.54	62.48

Appendix

Table 25: Raw per-run mean inference latency (ms) on WikiText-103 across 3 independent runs for FORGE-UGC. These raw values support the variance statistics reported in Table 19

Compiler Optimization

See all 27 papers

Jun 4, 2026

AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference

Transformer inference increasingly depends on specialized compiler and runtime support, but real model graphs still require semantic decisions about which regions are worth specializing and which CUDA implementation families are plausible. We present AgentCompile, an LLM-guided CUDA inference compiler that uses LLM outputs only as advisory search metadata. Given compiler-derived region summaries and bounded candidate spaces, the LLM proposes semantic labels, candidate priorities, parameter hints, and risk annotations; the compiler materializes CUDA candidates through templates, checks interface and hardware constraints, validates candidates empirically, selects implementations by measured latency, and falls back when specialization is unsupported or unprofitable. In end-to-end autoregressive generation, AgentCompile averages 5.66x, 4.05x, and 4.26x speedup over PyTorch eager on Qwen3-1.7B, Qwen3-4B, and Llama-3.2-1B-Instruct, respectively, across five representative workloads. We will open-source the project.

Xuanzhe Li, Ziyan Weng, Zhiyu Zhu +1

City University of Hong Kong (Dongguan) · City University of Hong Kong

High-Performance Computing GPU Kernel Generation GPU Acceleration Efficient Transformer Inference Large Language Model-Guided Optimization LLM Inference Acceleration+1

Ask PDF

Sep 24, 2026

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Deep learning inference and training performance depends critically on GPU kernel efficiency. Modern compilers such as PyTorch Inductor automatically generate GPU kernels from high-level model code, but frequently underperform expert-written implementations by wide margins. Recent LLM-assisted kernel optimizers can close this gap for standalone kernels, yet treat compiled models as black boxes, generally optimizing individual standalone kernels without respecting the compiler's structural decisions or verifying the model end-to-end. We present KernelOPT, a multi-agent system that treats compiled models as structured artifacts. It preserves vendor library calls (cuBLAS, cuDNN) and exclusively targets generated Triton sub-kernels using five profiling-guided LLM agents. A four-gate verification cascade applies static validation, multi-seed correctness checking, model-level float64-fallback verification, and performance gating (

γ{=}1.03

) to filter candidates and verify the re-stitched model end-to-end. When candidates fail verification, the system preserves the compiler baseline. The system accepts PyTorch nn Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems (100 Level 1, 100 Level 2 and 50 Level 3) on NVIDIA H200, KernelOPT achieves geometric mean speedups over torch compile of 1.40

\times

(L1), 1.15

\times

(L2), and 1.07

\times

(L3) across all kernels, including fallback cases. Optimized-only geomeans (excluding cases where verification gates preserve the compiler baseline) are substantially higher: 2.54

\times

(L1: 36/100), 1.84

\times

(L2: 23/100), and 1.37

\times

(L3: 11/50), reflecting where the optimizer achieves meaningful leverage.

Aheli Poddar, Sanskar Prasad, Arindam Samanta +3

Red Hat

GPU Kernel Optimization Agentic Search Large Language Model-Guided Optimization Compiler Optimization Neural Network Hardware Acceleration

Ask PDF

Apr 16, 2026

Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels

We present Nautilus, a novel tensor compiler that moves toward fully automated math-to-kernel optimization. Nautilus compiles a high-level algebraic specification of tensor operators into efficient tiled GPU kernels. Nautilus's successive lowering design allows high-level optimizations, expression rewrites, and tile optimizations to be jointly applied in a single end-to-end system. Nautilus presents a novel auto-scheduler that discovers sequences of high-level optimizations, while preserving the regular program structure needed by tile optimizers. Nautilus's auto-scheduler captures complex interactions and trade-offs in the high-level optimizations, including aggressive global transformations like advanced reduction fusion. Nautilus is the first end-to-end tensor compiler capable of starting from a math-like description of attention and automatically discovering FlashAttention-3-like kernels, offloading the entire burden of optimization from the programmer to the compiler. Across five transformer-based models and 150 evaluation configurations on NVIDIA GH200 and RTX 5090 GPUs, Nautilus achieves up to 23% higher throughput than state-of-the-art compilers on GH200 and up to 42% on RTX 5090, while matching or exceeding manually written cuDNN kernels on many long-sequence configurations.

Yifan Zhao, Yuchen Yang, Matei Budiu +1

University of Illinois Urbana-Champaign USA

GPU Kernel Optimization GPU Acceleration High-Performance Computing GPU Kernel Generation

Ask PDF

Abstract

Compiler Optimization

Explore similar work

AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels