Forge-UGC：面向异构加速器的通用图编译器——FX 优化与寄存器图引擎 Forge-UGC: FX optimization and register-graph engine for universal graph compiler
FX 图 + NPUIR 四阶段编译器，Intel NPU 上 6.9–9.2× 编译加速。
前置知识
PyTorch FX 图
FX（Functional Transformations）是 PyTorch 提供的程序捕获与转换框架，它将 nn.Module 追踪为一个由节点（op、placeholder、output）和边（张量依赖）组成的有向无环图。FX 图是完全 Python 原生、可编程检查的 IR，开发者可以在节点级别插入自定义变换（pass）。
FORGE-UGC 的整个前端和中间端优化都建立在 FX 图上；如果不理解 FX 节点、target、args 的语义，就看不懂其 DCE、CSE、融合 pass 是怎么改写图的。
算子融合（Operator Fusion）
把多个相邻算子（典型如 Linear+ReLU、或注意力分解链路）合并为一个内核，避免中间张量被物化到全局内存。对于注意力场景，FlashAttention 进一步把 Q·K^T、scale、mask、softmax、·V 全部融合为 IO 感知的单次访存。
FORGE-UGC 的 attention fusion pass 和 operator fusion pass 是其性能最大来源（单 pass 减 14.6% 节点）；理解融合消除「中间张量落盘 + 额外 NPU dispatch」是看懂所有加速数字的关键。
线性扫描寄存器分配（Linear-Scan Register Allocation）
经典的多区间寄存器分配算法 Poletto & Sarkar 1999：先按区间起点排序，再贪心扫描，过期区间回收物理槽位；复杂度 O(N log N) 而非 O(N²) 的图着色。它把无限多的虚拟寄存器映射到有限的物理 buffer 槽。
FORGE-UGC 的 Phase 4 用它把 NPUIR 的虚拟寄存器减为 30–48% 的物理 buffer；这是其内存层面优于 OpenVINO 的核心算法差异。
Intel AI Boost NPU 与 NNFactory
Intel Meteor/Arrow Lake 集成的低功耗 NPU，提供 ≤11 TOPS INT8、TDP < 10W；NNFactory 是其高层编译 API，要求把整张图一次性 compile-then-run，输出是单个可调度单元，不支持 Python 端的 kernel-by-kernel 调用。
FORGE-UGC 的所有 NPU 端优化（dispatch 合并、设备切换最小化、buffer 复用）都受 NNFactory 这一执行模型约束；理解这一点才能解释为什么作者放弃 torch.compile 后端。
研究动机
随着自主 AI 智能体在边缘端普及，单次推理流水线会同时穿越 NPU（稠密 tensor）、GPU（并行计算）和 CPU（控制流）。在这一异构 SoC 中，传统部署框架 OpenVINO 与 ONNX Runtime 暴露出五大共同局限：① 必须经过 TorchScript/ONNX 中转，遇到 RoPE、GQA、SwiGLU 等 PyTorch 2.x 新算子或动态控制流就会失败，Llama-3、Mistral、Qwen2 都需要手工算子分解；② 编译管线是黑盒，开发者无法查看哪个 pass 触发了融合、贡献了多少节点削减，做不了 ablation；③ 编译时间随模型深度超线性增长——8B 模型需要 58–62 秒，无法用于迭代开发或 JIT 部署；④ 没有 NPU 友好的 liveness 分析或虚拟寄存器抽象，OpenVINO 在 EP 层做内存规划但用户不可见，CPU 与 NPU 之间的不必要数据拷贝无法消除；⑤ 没有 NPU 专用的 cost model 和 autotuning，OpenVINO 只有「LATENCY/THROUGHPUT」粗粒度 hint，ONNX Runtime 的 EP 选择是静态规则、缺乏 NPU 实际优于 CPU 的成本反馈。
本文的目标是构建一个名为 FORGE-UGC（FX Optimization & Register-Graph Engine — Universal Graph Compiler）的四阶段编译器，前端捕获、中端优化、IR 表达与任何具体后端解耦，从架构上保证「同一管线可指向任何加速器」，在本文中以 Intel AI Boost NPU 作为首个目标后端做完整验证，并预留 Qualcomm Hexagon、AMD XDNA、Apple ANE 等扩展接口。
与已有工作不同的是，现有框架的根本缺陷在于它们把编译管线当作单设备代码生成器：要么 (a) 必须走 ONNX/TorchScript 这类有损中间表示（如 OpenVINO、ONNX Runtime），要么 (b) 只支持 CPU/GPU 后端（如 torch.compile+Inductor），要么 (c) 基于 MLIR 的通用后端（IREE）不提供 Intel NPU dispatch 也不做 NPU 专用的活跃性指导分配。FORGE-UGC 的独特切入点是「直接吃 PyTorch FX 图 + 在 ATen 层做 NPU 专属优化 + 提供带虚拟寄存器的 NPUIR + 用线性扫描做 liveness-guided buffer 分配 + 设备亲和性调度」这五件事首次同时出现在同一个 NPU 编译器中，并对外暴露可独立计时的 pass 接口和 CostModel 评分，使 ablation 与 autotuning 首次成为可能。
核心方法
FORGE-UGC 的整体思路是「先捕获、再优化、再降级、再调度」，分四阶段把一个 PyTorch 模型变成可直接 dispatch 到 NPU 的扁平指令流。直觉上，它把深度学习编译器经典的前端-中端-后端三段式范式改成可分解、可度量、可单步切换的四段流水线：第一阶段用 torch.export 在 ATen 算子层捕获图，并自动解析 GPT-2 这种共享权重的「tied weight」问题；第二阶段挂六个互相独立、各自返回修改次数的优化 pass（DCE、CSE、常量折叠、注意力融合、算子融合、布局优化），按 fixed-point 迭代直到不再变化；第三阶段把优化后的图降级为带 opcode、虚拟寄存器、设备标记的 NPUIR；第四阶段做 liveness 分析、线性扫描 buffer 分配、设备亲和性调度，最终产出扁平指令流 + 预解析 callable 的 CompiledNPUExecutor。这种设计让 pass 之间的交互可视化、让 buffer 分配有形式化算法支撑、让 NPU↔CPU 切换次数成为可优化的显式目标。
与已有方案的本质区别在于三个同时存在但此前被分开实现的设计选择：第一，直接操作 PyTorch FX（at export 层），跳过 TorchScript/ONNX 转换，使 RoPE/GQA/SwiGLU 等 PyTorch 2.x 新算子零成本支持；OpenVINO/ONNX RT 需要手工分解、Inductor 不支持 NPU。第二，把 NPU 编译所需的「virtual register + liveness + linear-scan」三件套从通用 CPU/GPU 寄存器分配器搬到 NPU 专用 IR 上，并配合 NPU SR RAM 容量做了 buffer 复用比 30–48% 的形式化证明；IREE 用 MLIR 通用 buffer deallocation 代替，OpenVINO 完全不暴露。第三，引入一个**设备亲和性贪心调度器**——在拓扑序约束下优先把同设备的指令聚簇，使 NPU↔CPU 切换次数从 Llama-3.1-8B 的 264 降到 93（−64.8%），这是已有 NPU 编译器都未做的工作。
方法步骤详情
第一阶段（Phase 1, FX Graph Capture）：调用 torch.export.export(model, args=(example_input,))，并设置 torch._dynamo.config.suppress_errors = True 抑制开发期算子的 trace-time warning，得到 ExportedProgram 与 fx.GraphModule；然后 wrap_exported_for_npu() 遍历 named_modules()，用 Python id() 匹配参数 tensor，建立 tied_map[新名]=规范名，让 GPT-2 的 embedding 与 LM-head 共享同一物理张量。第二阶段（Phase 2, 6 个 pass）：每个 pass 继承 FXPassBase 并实现 run(gm)->bool，固定点循环 2 轮——(1) FXDCEPass 从 output 节点反向 BFS 标记 live，删掉不可达节点；(2) FXCSEPass 用 (target, args-name-tuple, kwargs-tuple) 哈希-consing，重复节点 replace_all_uses_with 后 erase；(3) FXConstantFoldingPass 识别 x+0、x*1、shape 计算等常量模式替换为 literal；(4) FXAttentionFusionPass 从 matmul(Q,K^T) 节点出发，链式识别 scale/mask/softmax/dropout/P@V 整段替换为单次 F.scaled_dot_product_attention，关键细节是 _unwrap_transpose 还原 K 的原始 tensor；(5) FXOperatorFusionPass 匹配 Linear+ReLU/GELU/SiLU 与 mm+add，调用 _run_npu_fused_op 编译一次并按 (op_id, activation) 缓存 NNFactory graph；(6) FXLayoutOptimizationPass 查 NPU_PREFERRED_LAYOUTS 表，在 embedding 后或首个 NPU 操作前插入 .contiguous()，并合并连续 contiguous 调用。第三阶段（Phase 3, Lowering to NPUIR）：对每个 FX 节点分配 vreg 计数器，构造 NPUIROp=(opcode, vreg_in, vreg_out, device, callable)，device 由节点 target 字符串是否包含 _npu_linear_/_npu_fused_/_npu_mm_/_npu_addmm_ 决定，所有 args 在 lowering 时替换为 _RegRef 占位符。第四阶段（Phase 4, IR Analysis & Optimization）：(a) 用 backward walk 算出每个 vreg 的 [s_i, e_i] 活跃区间与 dead_after_map；(b) 调用 BufferAllocator.allocate 做经典 linear-scan：按起点排序、用 active list + free pool 复用 buffer，复杂度 $O(N \log N)$；(c) InstructionScheduler 做设备亲和性拓扑排序——每次在「数据依赖已就绪」的指令中优先选与最近已调度同设备的指令，使 δ 减少 42–65%；(d) 输出 CompiledNPUExecutor，运行时按扁平指令流执行，预解析 callable 消除属性查找与动态分配；可选的 AutotuningCompiler 在 {$\alpha \in \{0.2, \ldots, 1.0\}$, $\lambda \in \{\mathrm{auto, channels\text{-}last, contiguous}\}$, $\pi \in \{\mathrm{fp16, int8, mixed}\}$, $\iota \in \{1, 2, 3\}$} 45 个配置上跑 CostModel $\mathrm{Score}(G) = w_1 \cdot n_{\mathrm{ops}} + w_2 \cdot n_{\mathrm{weights}} + w_3 \cdot n_{\mathrm{linear}} + w_4 \cdot d_{\mathrm{graph}} + w_5 \cdot s_{\mathrm{params}}$，<200ms 内选最优。
技术新颖性
技术新颖性体现在四个方面：(1) **首个把 PyTorch FX + NPU 编译完整栈串起来的原型**——从 torch.export 的 ATen-level capture 到 NPUIR 之间的所有 NPU 优化（融合、buffer 分配、调度）此前都散落在不同框架中，FORGE-UGC 是第一个把它们与 NPU dispatch 一体化的实现，且开发周期不到六个月由两人完成。(2) **首次为 NPU 引入 liveness-guided linear-scan buffer 分配**——传统 NPU 编译器（包括 OpenVINO/ONNX Runtime）要么不暴露 buffer 分配，要么依赖 MLIR 通用 buffer deallocation；FORGE-UGC 用 Poletto–Sarkar 算法针对 NPU 活跃区间做复用，得到 $30\text{–}48\%$ 的物理 buffer 削减。(3) **NPU 专属的设备亲和性调度器**——把 PCIe/MMIO 切换成本显式建模（每次 0.3–0.8ms），用贪心启发式在拓扑序约束下聚簇同设备指令，Llama-3.1-8B 上把切换从 264 降到 93，这一思想未见于任何现有 NPU 编译器。(4) **三套可量化的评估指标** FGR（融合增益比，cost-model 内部诊断）、CEI（编译效率指数，speedup-per-second）和 per-pass profiling，首次让 NPU 编译器的 ablation 变得可重复——这是其他 NPU 框架（OpenVINO/ONNX RT/QNN）至今缺失的能力。
实验结果
核心实验在六款 125M–8B 模型（GPT-2、Granite-350M、Qwen2-0.5B、Llama-3.2-1B、LFM2-2.6B、Llama-3.1-8B）上、以 OpenVINO 2024.4 与 ONNX Runtime 1.19 为基线，在 Intel Core Ultra 9 285HX + Intel AI Boost NPU（11 TOPS INT8）上跑 WikiText-103 与 GLUE。**编译速度**：FORGE-UGC 在 GPT-2 上 1.00s vs OpenVINO 6.93s/ONNX RT 7.27s（6.9×/7.3×），在 8B 模型上 6.72s vs 58.43s/62.15s（8.7×/9.2×），且作者指出 78% 时间花在 torch.export 这层「共享基础设施」上，FORGE-UGC 自己的 pass+backend 只有约 216ms。**图节点削减**：六个 pass 在 GPT-2 上把节点从 403 减到 333（−17.4%），其中 attention fusion 单 pass 贡献 403→344（−14.6%）最大削减，operator fusion 贡献 −3.5%；六个模型平均 −18.8%，LFM2-2.6B 最高 −21.9%。**推理延迟（WikiText-103）**：GPT-2 6.82ms vs OpenVINO 8.45ms（−19.3%）、vs ONNX RT 9.13ms（−25.3%），并随模型深度放大——Llama-3.1-8B 62.48ms vs 91.37/97.82ms（−31.6% / −36.1%）。**P99 tail**：FORGE-UGC P99/P50 稳定在 1.20，OpenVINO 与 ONNX RT 在 1.27–1.28，反映扁平指令流+预分配 buffer 的确定性。**数值保真**：fp16 派发模型 PPL 全部 bit-for-bit 一致（GPT-2 29.41→29.41、Llama-3.2-1B 9.76→9.76），max-abs $\Delta\text{logit} < 1.2 \times 10^{-5}$，KL 散度 $D_{KL} < 6.3 \times 10^{-10}$；8B 模型因 NNFactory 引入 W-int8/A-fp16 量化，max-abs 升到 $2.1 \times 10^{-5}$、KL $8.4 \times 10^{-9}$，仍在 fp16 舍入容差内。**能效**：每推理能量 30.2–40.9% 低于 OpenVINO、37.0–46.2% 低于 ONNX Runtime，节省幅度比延迟更显著（因为设备切换减少 → CPU dispatch overhead 降低，buffer 预分配 → DRAM power 下降）。**Ablation**：移除 attention fusion 让 cost score 暴增 $27.6\times$，证实其为最关键 pass；operator fusion 贡献 2.3%；DCE/CSE/常量折叠在 GPT-2 上几乎不动（trace 图干净）但在大模型上必要；fusion aggressiveness α 从 0→1.0 让 cost score 从 364.87 单调降到 8.64（FGR=42.3），$\alpha=1.0$ 始终最优（NPU 派发整个 fused subgraph，没有 GPU 上 register pressure 顾虑）。**Autotuning**：在 45 个配置上对 cost model 评分（<200ms），比默认配置再省 4.2–8.7% cost score，模型越大收益越大。**方差**：GPT-2 上 10 次独立运行所有指标 CV<2.5%（latency 2.05%、compilation 1.84%），节点削减为 0（pass 完全确定）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GPT-2 125M 编译时间 | 端到端编译耗时 (ms) | 1,000 | OpenVINO 6,930 / ONNX RT 7,271 | $6.9\times / 7.3\times$ 加速 |
| Llama-3.1-8B 编译时间 | 端到端编译耗时 (ms) | 6,720 | OpenVINO 58,430 / ONNX RT 62,150 | $8.7\times / 9.2\times$ 加速 |
| GPT-2 125M WikiText-103 推理延迟 | mean latency (ms) | 6.82 | OpenVINO 8.45 / ONNX RT 9.13 | $-19.3\% / -25.3\%$ |
| Llama-3.1-8B WikiText-103 推理延迟 | mean latency (ms) | 62.48 | OpenVINO 91.37 / ONNX RT 97.82 | $-31.6\% / -36.1\%$ |
| GPT-2 125M 每推理能量 | energy (mJ) | 69.6 | OpenVINO 99.7 / ONNX RT 110.5 | $-30.2\% / -37.0\%$ |
| Llama-3.1-8B 每推理能量 | energy (mJ) | 637.3 | OpenVINO 1,078.2 / ONNX RT 1,183.6 | $-40.9\% / -46.2\%$ |
| GPT-2 125M 数值保真 | max-abs Δlogit / KL 散度 | 6.2×10⁻⁶ / 1.8×10⁻¹⁰ | PPL 29.41→29.41 (与基线 bit-for-bit 一致) | fp16 舍入容差内，近似恒等 |
| GPT-2 125M P99 tail | P99/P50 ratio | $1.20$ | OpenVINO 1.29 / ONNX RT 1.28 | −7–9 个百分点 |
| Llama-3.1-8B 设备切换 | NPU↔CPU transition 次数 | 93（从 264 减少） | OpenVINO/ONNX RT 未直接报告但调度器无法聚簇 | −64.8% |
| Llama-3.1-8B 物理 buffer 削减 | physical buffer count / virtual reg | 468/896 | 无形式化分配（OpenVINO/ONNX RT 不可见） | −47.8% (ρ_buf) |
局限与改进
作者明确承认的限制与读者观察：(1) **验证范围仅限 Intel AI Boost NPU**——Qualcomm Hexagon、AMD XDNA、Apple ANE、Samsung NPU 等后端是「planned future extensions」，本文没有这些后端的实际编译结果，所谓的「universal」尚未在硬件层面证实。(2) **单 batch 推理**——所有延迟数字都是 batch=1；生产部署需要 batched execution，但当前 CompiledNPUExecutor 没有显式 batching 支持。(3) **Autotuning 是网格搜索**（45 个配置），作者承认 Bayesian 优化或 learned cost model 可能进一步提升。(4) **NPUIR 是固定 opcode 集**——不支持自定义算子注册，对新模型结构的覆盖受限。(5) **FGR 是 cost-model 内部诊断**——它衡量 CostModel 评分比，非真实 wall-clock 加速；67.9× FGR 对应 Llama-3.1-8B 上 29.6% 实测延迟削减，两者数量级差距大（作者明确说不应把 FGR 当作延迟比）。(6) **能耗测量用系统级 RAPL**，未做 per-component（CPU vs NPU）拆分，归因精度受限。(7) **Llama-3.1-8B 必须用 W-int8 量化**（fp16 权重 16GB 超 NPU 内存），导致数值保真指标略高于其它模型，这是 dispatch 层而非 pass 层引入的误差。(8) **作者亲自试验过 MLIR 路线**（IREE-Turbine），并坦诚 C++ pass 框架迭代太慢，间接说明这个 6 个月内建好的 FX-Python 栈能否扩展到需要复杂 dialect 转换的大模型仍是未知数。(9) **CPU/NPU 设备切换的硬编码启发式**——`is_npu = any(t in str(node.target) for t in ('_npu_linear_', ...))` 这种字符串匹配脆弱，未来加新 NPU 算子需要同步改 lowering 与调度器规则。(10) **公平性争议**——OpenVINO 仍用 ONNX 转换再 IR 化，确实有「多一道中转」的不利，但作者并未在相同上游条件下（如给 OpenVINO 也走 torch.export）做对照，难以判断 6.9–9.2× 加速里有多少来自「无 ONNX 转换」而不是「pass 实现本身更快」。
独立分析的弱点
**独立观察的弱点及改进方向**：(1) **设备路由是字符串硬匹配**——`_npu_linear_/_npu_fused_/_npu_mm_/_npu_addmm_` 出现在 target 字符串里就路由到 NPU，这是脆弱的命名耦合；改进方向是建立「算子能力表」（op_capability_table）按 dtype/layout/输入 shape 决定设备归属，而不是靠 substring。(2) **Fusion pass 仅匹配固定模式**——attention fusion 假设 QK matmul 后紧跟 scale→mask→softmax→dropout→PV，一旦用户用了 QK-norm（如 Qwen 一些变体）或 ALiBi 偏置，pattern 匹配会失败；改进方向是用 e-graph 或 ML-style sub-graph isomorphism 替代线性 forward walk。(3) **Cost model 仍是启发式加权**——Score(G) = Σwᵢ·xᵢ 没有用真实硬件 profile 校准，导致 FGR 数值（$42\text{–}68\times$）与实测延迟改善（16.6–29.6%）量级脱节；改进方向是离线 NPU profiling 收集每算子 wall-clock，再用线性回归或贝叶斯优化拟合权重。(4) **激活 fusion 缓存键是 (op_id, activation)**，没考虑 batch size、dtype 差异——当用户切到 batch=4 时会重新走一遍 _run_npu_fused_op；改进方向是扩展缓存键含 (shape, dtype, stride) 元组。(5) **buffer 分配是单层**——所有 NPU 中间张量共享一组物理 buffer，没有按 NPU SRAM 层次（HBM/L2/SRAM/寄存器）分层映射；改进方向是借鉴 Hexagon-MLIR 的 TCM-aware tiling 策略，按层数和 reuse distance 选不同存储层。(6) **无算子自动调度**——schedule 只做设备亲和性，没做算子级并行（如两路 matmul 互不依赖时可并发 dispatch）；改进方向是在 Phase 4 加入简单的 ready-queue + 同设备 batch dispatch。(7) **量化仅在 dispatch 层硬编码**——Llama-3.1-8B 用 W-int8/A-fp16 是 NNFactory 默认行为，缺少「按层 sensitivity 混合精度」策略；改进方向是加入 per-layer quantization sensitivity profiler，让 8B 模型也对敏感层保留 fp16。
未来方向
**作者提出的方向**：(1) **Triton 内核集成**——在 FORGE-UGC 流水线中支持 Triton-to-NPUIR lowering，让用户用 Triton DSL 写自定义 NPU kernel，受 Hexagon-MLIR 启发但保持 FX-native。(2) **自演化编译器**——把编译管线本身当学习系统，用运行时 telemetry（温度、内存压力、电池）逐步 refine pass 顺序与融合 aggressiveness。(3) **多后端扩展**——Qualcomm Hexagon、AMD XDNA、Apple ANE、Samsung NPU 复用 Phase 1+2+3，只在 Phase 4 写新 dispatch 模块。(4) **量化感知编译路径**——给 8B 模型加 sensitivity-driven mixed precision，让用户能 trade-off 大小 vs 精度。(5) **硬件校准 cost model**——把 CostModel $\mathrm{Score}(G) = \sum_i w_i \cdot x_i$ 系数用真实 NPU profile 拟合，使 FGR 直接可解释为 latency ratio。**基于成果可延伸的方向**：(a) 把 FORGE-UGC 作为异构 orchestrator（如作者团队同期工作 QEIL）的 NPU 后端，实现「workload-aware routing + hardware-aware compilation」闭环；(b) 把 Phase 4 的 linear-scan 与 instruction scheduling 推广到异构 GPU+NPU——同一算子有 GPU 和 NPU 两个实现时，由 scheduler 根据活跃区间长度选目标；(c) 把六 pass 暴露成 RPC 服务，让用户在远程设备上 ablation 而不必重编译；(d) 把 FGR、CEI 推广为社区公认基准，对未来 NPU 编译器做横向评测。
复现评估
**复现评估**：(1) **开源情况**——论文未在 GitHub 给具体链接，但 Intel NPU Acceleration Library (NNAL) 公开在 github.com/intel/intel-npu-acceleration-library，本文的后端 dispatch 假设用户能拿到这个；FORGE-UGC 自身代码未在正文给出仓库地址（论文提到 「The first prototype is ready to ship」 但没贴 URL），社区复现需要邮件联系作者。(2) **硬件门槛**——Intel Core Ultra 9 285HX + Intel AI Boost NPU（11 TOPS）是 Meteor Lake/Arrow Lake 才有的 SKU，普通台式机/服务器没有；想完整复现能量数据需要 RAPL 接口（Intel 平台特有）。(3) **数据**——评测用 WikiText-103（217,646 tokens 滑窗 512 步长 256）与 GLUE SST-2/MNLI，都是公开数据集，门槛低。(4) **算力与时间**——单卡工作站级别，编译 GPT-2 1s、Llama-3.1-8B 6.7s，推理延迟几 ms–60ms 级别，单次评测几小时内可完成；autotuning 额外 <200ms。(5) **难度**——中等偏低：模型即装即用（pip install + 6 款模型均可从 HuggingFace 下载），主要工作量在配 NPU 驱动（32.0.100.4514）和 NNFactory 工具链；如果只想复现节点削减与 FGR 这类静态指标，不需真机也能跑（用 mock backend 即可）。(6) **方差可控**——所有指标 CV<2.5%，3 次运行即能复现论文均值；节点削减 0% 方差说明 pass 是确定性的。(7) **隐含门槛**——若想完整复现 NPU dispatch 部分，需要 Windows 11 24H2 + Intel NPU 驱动 + NNAL 工具链，这套组合在 Linux 服务器上不通用；论文没明确说是否支持 Linux 平台。
论文图表
