Forge-UGC: FX optimization and register-graph engine for universal graph compiler
| Authors | Satyam Kumar & Saurabh Jha |
| Year | 2026 |
| HF Upvotes | 3 |
| arXiv | 2604.16498 |
| Download | |
| HF Page | View on Hugging Face |
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.
Engineering Breakdown
Plain English
Forge-UGC is a four-phase compiler framework designed to deploy transformer models efficiently on heterogeneous accelerator hardware, specifically validated on Intel AI Boost NPU. The authors identify that existing frameworks like OpenVINO and ONNX Runtime suffer from opaque compilation pipelines, limited visibility into individual optimization passes, and weak buffer management strategies that increase both compilation overhead and runtime latency. Forge-UGC solves this by introducing a hardware-agnostic architecture that explicitly separates graph capture, optimization, intermediate representation lowering, and backend scheduling into distinct phases. The framework captures computation graphs at the ATen operator level using torch.export and natively supports modern transformer components like rotary position embeddings, grouped-query attention, and SwiGLU activations without requiring manual decomposition.
Core Technical Contribution
The core novelty is a transparent, multi-phase compilation pipeline that makes the black-box nature of existing compilers into an inspectable, modular system where each optimization pass can be understood and debugged independently. Unlike ONNX Runtime or OpenVINO which bundle passes together, Forge-UGC explicitly decouples graph capture (Phase 1), optimization (Phase 2), IR lowering (Phase 3), and backend scheduling (Phase 4), enabling fine-grained control over which optimizations fire and when. The register-graph engine (the 'RG' in the name) appears to be a novel intermediate representation designed specifically for heterogeneous accelerators, managing data movement and register allocation across different hardware units without opaque heuristics. The hardware-agnostic design allows a single compiled artifact to potentially target multiple backends by swapping only the final scheduling phase, rather than recompiling the entire graph for each new hardware target.
How It Works
Phase 1 takes a PyTorch model and uses torch.export to capture the computation graph at the ATen operator granularity, preserving high-level semantics of transformer operations (rotary embeddings, grouped-query attention, etc.) rather than decomposing them into primitive ops, which keeps the graph semantically rich for later optimization. Phase 2 performs graph-level optimizations such as operator fusion, dead code elimination, and layout optimization, with each optimization as a discrete pass that can be toggled or inspected, providing engineers with visibility into what changed and why. Phase 3 lowers the optimized graph into an intermediate representation (the register-graph IR) that is hardware-agnostic but accounts for realistic hardware constraints like register pressure, memory hierarchy, and data movement costs. Phase 4 takes the lowered IR and performs backend-specific scheduling and code generation for the target accelerator (e.g., Intel AI Boost NPU), mapping logical operations to physical hardware execution units and managing buffer allocation. The separated phases mean that improvements to the optimization pass (Phase 2) or the IR design (Phase 3) do not require changes to the torch.export logic or backend scheduling, enabling iterative improvement without cascading rewrites.
Production Impact
For teams deploying transformers on edge or embedded accelerators, Forge-UGC provides a critical advantage: visible, debuggable compilation that reduces surprise performance drops between model export and runtime execution. Production systems often spend 30-50% of total latency in compilation or in suboptimal buffer management; the explicit buffer management and pass-level visibility could eliminate much of that waste by allowing engineers to identify and fix compilation bottlenecks without vendor black boxes. The hardware-agnostic design means you could train and optimize once, then target multiple accelerator platforms (CPU, NPU, GPU) by only changing the final backend, significantly reducing the time and complexity of multi-platform deployment. However, adoption requires understanding the four-phase model and potentially retraining teams on a new compiler abstraction; this has non-zero onboarding cost. The production win is substantial for long-running inference systems where compilation overhead compounds, but for one-shot inference or latency-critical services, the compilation cost itself may be negligible.
Limitations and When Not to Use This
The paper does not demonstrate end-to-end performance numbers or comparisons to OpenVINO/ONNX Runtime on real transformer models, so the actual speedup or latency improvement is unknown from the abstract alone. The framework is validated only on Intel AI Boost NPU; it is unclear how well the hardware-agnostic IR generalizes to fundamentally different accelerator architectures (e.g., TPUs, mobile NPUs, or GPUs with very different memory hierarchies). The automatic buffer management strategy is not detailed, leaving ambiguity about whether Forge-UGC's approach is truly better or simply exposes the same heuristics in a more visible way. The paper abstracts modern transformer components (rotary embeddings, grouped-query attention) without manual decomposition, but does not address custom operators or user-defined kernels, which real production models often include.
Research Context
Forge-UGC builds on the torch.export standardization effort (PyTorch 2.x) that aims to create portable, stable computation graphs, and extends it with a principled compilation framework that improves upon earlier work in compiler design for deep learning (e.g., TVM, Glow). It directly addresses known pain points in existing production-grade compilers like ONNX Runtime and OpenVINO, which have been widely deployed but are known to have opaque optimization pipelines and suboptimal memory management. The register-graph IR appears to be positioned as an alternative to existing deep learning IRs (MLIR, TVM Relay) with a focus on explicit buffer management and hardware heterogeneity. This work opens a research direction toward compiler transparency and debuggability in ML, where future work might extend the framework to support dynamic shapes, sparse operations, or quantization-aware compilation without losing the visibility guarantees.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::