PassNet: 扩展大规模语言模型用于图编译器Pass生成

# PassNet: 扩展大型语言模型用于图编译器Pass生成
来源: https://arxiv.org/html/2605.29357
Yiqun Liu Yingsheng Wu¹¹Ruqi Yang Enrong Zheng Honglei Qiu Sijun He Tai Liang Jingjing Wu Yuhan Zhou Yiwei Zhang Dongyan Chen Weihan Yi Xinqi Li Siqi Bao²² Baidu, Inc.

###### 摘要

现代张量编译器（如 TorchInductor）在主流模型上实现了显著的加速，但在长尾工作上却面临系统性的性能上限——我们的性能分析显示，在默认编译下，43% 的真实世界子图经历了端到端的性能下降。虽然 LLM 为自动化优化提供了一条路径，但现有工作集中在独立的*内核生成*上。我们认为*Pass生成*——即 LLM 编写可直接集成到编译器流水线中的结构化图变换——是更合适的抽象层。我们提出 PassNet，首个用于基于 LLM 的编译器 Pass 生成的大规模生态系统，包括: (1) PassNet-Dataset，来自 100K 个真实世界模型的超过 18K 个唯一计算图；(2) PassBench，200 个精心挑选的长尾可融合任务（共包含 2,060 个子图），在*错误感知加速分数* (ES_t) 下进行评估——该指标统一了正确性、稳定性和性能——并采用分层完整性防御来抵御系统的 LLM 利用。实验表明，PassBench 既具有高度区分性，又真正未饱和: 最好的前沿模型在整体上落后 TorchInductor 37%，但在单个子图上，LLM 相对于同一编译器可实现高达 3 倍的加速——这表明瓶颈在于一致性，而非能力。在一个小模型上仅用约 4K 个 PassNet 轨迹进行微调，就获得了 2.67 倍的改进，接近前沿模型性能，证明了巨大的提升空间，并验证了 PassNet 作为推动 LLM 驱动编译器优化的在线训练基础设施的价值。所有数据、基准测试和工具均已公开提供。

## 1 引言

现代深度学习系统越来越多地依赖张量编译器（例如，TVM [Chen等人，2018 (https://arxiv.org/html/2605.29357#bib.bib27)]、XLA [Kaufman等人，2021 (https://arxiv.org/html/2605.29357#bib.bib7)]、TorchInductor [Ansel等人，2024 (https://arxiv.org/html/2605.29357#bib.bib10)]）将高级计算图降级为面向异构硬件的高效后端实现。这些编译器应用专家设计的、基于规则的*Pass流水线*——一系列图变换，如算子融合、分块和布局选择——已被证明非常有效：在主流架构上，TorchInductor 在超过 180 个模型上相比即时执行实现了高达 2.27 倍的推理加速 [Ansel等人，2024 (https://arxiv.org/html/2605.29357#bib.bib10)]。然而，当面对长尾的真实世界算子组合时，这些静态策略面临结构性的局限。

#### 长尾优化差距。
为了量化这一局限，我们对 TorchInductor 的默认流水线在从超过 1,000 个社区模型提取的 9,526 个子图上进行了性能分析。结果揭示了系统性的限制：34% 的子图实现了边际加速（<1.2 倍），43% 经历了端到端性能下降，8.3% 被严格降级。这种差距是结构性的，与算子覆盖范围相关，而非图复杂度（r=0.013），表明仅扩展图覆盖率无法克服由启发式方法引起的性能上限。

#### 模式集中创造了杠杆空间。
虽然长尾差距显著，但计算图模式表现出强烈的幂律集中性：对 100K 个模型去重后仅产生约 18K 个不同的图（82% 冗余），约 10,000 个子图缩减为约 1,025 个唯一的结构模式。为这一集中集合生成高质量的 Pass 足以覆盖大多数工作负载，从而能够从*人工规则*转向*自动化、数据驱动的 Pass 生成*。

#### 从内核生成到 Pass 生成。
大型语言模型（LLM）是长尾优化的有前途的方法。现有工作探索了内核生成 [Ouyang等人，2025 (https://arxiv.org/html/2605.29357#bib.bib35); Dai等人，2026 (https://arxiv.org/html/2605.29357#bib.bib37); Liao等人，2025 (https://arxiv.org/html/2605.29357#bib.bib46)]，为单个算子生成独立的 GPU 内核。然而，此类内核缺乏与编译器 Pass 的组合性，需要手动集成才能部署，并且由于代码生成不受约束而难以验证。因此，我们定义了一个新任务，**Pass 生成**：给定一个计算子图，LLM 必须编写一个结构化的编译器 Pass，该 Pass 直接与编译器的中间表示接口（形式化定义见第 3 节 (https://arxiv.org/html/2605.29357#S3)）。这种表述保留了“一行编译”的体验（例如，`torch.compile`），同时实现了可组合且可验证的优化。然而，推进这一任务需要大规模的数据和严格的评估——目前尚不存在这些资源。

#### 数据和评估瓶颈。
如今 Pass 生成与 LLM 结合是否现实？在单个长尾子图上，前沿 LLM 已经能够生成 Pass，相对于默认编译器实现高达 3 倍的加速——然而整体性能远远落后，没有模型达到高于 1.0 的几何平均加速。这一差距源于两个*基础设施*瓶颈：(1) *数据稀缺*——缺乏用于融合和布局优化的大规模专用语料库；(2) *评估盲点*——缺乏严格的基准，使得智能体能够绕过正确性而追求虚幻的收益。为弥合这些差距，我们引入了 PassNet——首个为 Pass 生成任务提供系统化数据和基准支持的生态系统。

#### 贡献。
我们形式化了*Pass 生成*，即 LLM 编写可集成到编译器流水线中的结构化图变换，并构建了一个大规模生态系统来支持这一任务。
- • **PassNet-Dataset**：我们从来自不同框架（PyTorch、PaddlePaddle）和任务类别的 100K 个真实世界模型中收集了超过 18,086 个唯一计算图。我们设计了*递归折叠*和*执行驱动的前缀分析*，以在多种粒度上构建结构多样的子图，形成了首个用于 Pass 生成的大规模开放训练集。
- • **PassBench**：我们构建了一个包含 200 个任务的基准测试，每个任务由可变数量的长尾子图组成，在*错误感知加速分数*（ES_t）下进行评估，该分数联合衡量正确性、稳定性和性能。为确保评估完整性，我们引入了分层防御，包括基于 AST 的检查、运行时分发拦截和反向评估顺序，系统地应对了开发过程中观察到的利用模式。
- • **基准验证**：通过对 6 个前沿和开源模型的广泛评估，我们表明 PassBench 既具有高度区分性（模型层级间差距 3.22 倍），又真正未饱和（最佳模型落后 TorchInductor 37%）。仅在约 4K 个 PassNet 轨迹上进行微调，就获得了 2.67 倍的改进，验证了数据集作为训练基础设施的实用性。

## 2 相关工作

#### 张量编译器。
张量编译器通过 IR 降级和调度将高级计算图转换为设备特定的内核。TVM [Chen等人，2018 (https://arxiv.org/html/2605.29357#bib.bib27)] 和 Ansor [Zheng等人，2020 (https://arxiv.org/html/2605.29357#bib.bib28)] 是基于搜索的编译器，依赖成本模型，而 XLA [Leary 和 Wang，2017 (https://arxiv.org/html/2605.29357#bib.bib33)] 应用启发式图级优化。近期系统包括 MetaSchedule [Shao等人，2022 (https://arxiv.org/html/2605.29357#bib.bib11)]、Hidet [Ding等人，2023 (https://arxiv.org/html/2605.29357#bib.bib12)] 和 BladeDISC [Zheng等人，2023a (https://arxiv.org/html/2605.29357#bib.bib8)]，行业框架如 CINN [Team，2021 (https://arxiv.org/html/2605.29357#bib.bib5)] 和 TorchInductor [Ansel等人，2024 (https://arxiv.org/html/2605.29357#bib.bib10)] 已集成到生产系统中。尽管有这些进展，现有方法仍依赖人工变换规则，并在长尾工作负载上表现不佳。

#### LLM 用于代码生成和编译器优化。
大型语言模型展现出强大的代码生成能力 [Jiang等人，2024 (https://arxiv.org/html/2605.29357#bib.bib50); Zheng等人，2023b (https://arxiv.org/html/2605.29357#bib.bib49)]，包括代码补全 [Liu等人，2025 (https://arxiv.org/html/2605.29357#bib.bib54)]、错误修复 [Huang等人，2025 (https://arxiv.org/html/2605.29357#bib.bib52); Yang等人，2026 (https://arxiv.org/html/2605.29357#bib.bib53)] 和软件工程自动化 [Jimenez等人，2024 (https://arxiv.org/html/2605.29357#bib.bib66); Yang等人，2024 (https://arxiv.org/html/2605.29357#bib.bib67); Wang等人，2024 (https://arxiv.org/html/2605.29357#bib.bib68)]。对于编译器任务，LLM Compiler [Cummins等人，2025 (https://arxiv.org/html/2605.29357#bib.bib58)] 在编译器 IR 上进行预训练，Compiler-r1 [Pan等人，2025 (https://arxiv.org/html/2605.29357#bib.bib38)] 探索了基于 RL 的自动调优，DeCOS [Cui等人，2025 (https://arxiv.org/html/2605.29357#bib.bib57)] 研究了数据高效的优化选择。TLP [Zhai等人，2023 (https://arxiv.org/html/2605.29357#bib.bib60)] 及后续工作 [Zhai等人，2024 (https://arxiv.org/html/2605.29357#bib.bib39)] 将语言模型应用于张量程序生成。然而，这些方法侧重于 Pass 选择或调度，而非合成新的变换逻辑。

#### LLM 驱动的 GPU 内核生成。
近期工作直接用 LLM 生成 GPU 内核。KernelBench [Ouyang等人，2025 (https://arxiv.org/html/2605.29357#bib.bib35)] 对 LLM 生成的内核进行基准测试，揭示了与生产编译器的差距。CUDA Agent [Dai等人，2026 (https://arxiv.org/html/2605.29357#bib.bib37)] 和 Kernelevolve [Liao等人，2025 (https://arxiv.org/html/2605.29357#bib.bib46)] 规模化智能体生成，而 STARK [Dong等人，2025 (https://arxiv.org/html/2605.29357#bib.bib40)]、Kevin [Baronio等人，2025 (https://arxiv.org/html/2605.29357#bib.bib48)] 和 Geak [Wang等人，2025 (https://arxiv.org/html/2605.29357#bib.bib30)] 探索了多智能体和基于 RL 的改进。其他工作包括 QiMeng-GEMM [Zhou等人，2025 (https://arxiv.org/html/2605.29357#bib.bib41)]、CUDA-L1 [Li等人，2025 (https://arxiv.org/html/2605.29357#bib.bib47)] 和 Autocomp [Hong等人，2025 (https://arxiv.org/html/2605.29357#bib.bib32)]。相比之下，PassNet 的目标是*Pass 生成*，即与编译器流水线集成的结构化变换。

#### 性能基准测试。
深度学习评估已从 DeepBench [Narang 和 Research，2016 (https://arxiv.org/html/2605.29357#bib.bib14)] 发展到 MLPerf [Mattson等人，2020 (https://arxiv.org/html/2605.29357#bib.bib15)]。CompilerGym [Cummins等人，2022 (https://arxiv.org/html/2605.29357#bib.bib73)] 为编译器优化提供了 RL 环境。诸如 ComPile [Grossman等人，2024 (https://arxiv.org/html/2605.29357#bib.bib16)]、TpuGraphs [Phothilimthana等人，2023 (https://arxiv.org/html/2605.29357#bib.bib18)] 和 TenSet [Zheng等人，2021 (https://arxiv.org/html/2605.29357#bib.bib19)] 等数据集为学习型编译器提供了基准。PassNet 专注于计算图级别的 Pass 生成以实现长尾优化，并采用联合衡量正确性、稳定性和加速的评估框架。

## 3 PassNet 生态系统

PassNet 生态系统弥合了第 1 节 (https://arxiv.org/html/2605.29357#S1) 中识别出的双重基础设施差距：大规模 Pass 生成语料库的稀缺以及缺乏鲁棒的多维基准测试。下面，我们在详细阐述每个生态系统组件之前，先形式化 Pass 生成任务。

### 3.1 任务形式化

在现代张量编译器中，一个*Pass* 是一个独立的图变换，它重写一个计算图，同时保留其输入输出语义 [Lattner等人，2021 (https://arxiv.org/html/2605.29357#bib.bib29); Li等人，2021 (https://arxiv.org/html/2605.29357#bib.bib22)]。遵循 MLIR [Lattner等人，2021 (https://arxiv.org/html/2605.29357#bib.bib29)] 和 TorchInductor [Ansel等人，2024 (https://arxiv.org/html/2605.29357#bib.bib10)] 采用的基于模式的重写范式，我们在下面形式化核心抽象。

###### 定义 3.1 (计算图)。
一个计算图是一个有向无环图 \(G=(V,E,\tau,\sigma)\)，其中 \(V\) 是一组算子节点，\(E\) 编码数据依赖关系，\(\tau: V \to \mathcal{T}\) 分配算子类型，\(\sigma: V \to \mathbb{Z}^+\) 分配输出形状。我们用 \(f_G: \mathcal{X} \to \mathcal{Y}\) 表示由 \(G\) 计算的函数。

###### 定义 3.2 (编译器 Pass)。
一个编译器 Pass 是一个对 \(\pi=(M,R)\)，其中 \(M: \mathcal{G} \to 2^{\mathcal{G}}\) 是一个模式匹配器，识别可优化的子图；\(R: \mathcal{G} \to \mathcal{G}\) 是一个重写器，将每个匹配的子图替换为优化的等价子图。Pass \(\pi\) 在公差 \(t\) 下对 \(G\) **有效**，如果：
\[
\forall x \in \mathcal{X}, \quad \mathrm{err}(f_G(x),\; f_{\pi(G)}(x)) \le t
\tag{1}
\]

**Pass 生成任务**是：给定一个**任务实例** \(\mathcal{T}=\{G_1,\ldots,G_k\}\)，这些子图共享相同的算子类型序列但形状和数据类型不同，生成一个有效的 Pass \(\pi\)，使其重写每个 \(G_i \in \mathcal{T}\) 并改善总体运行时间性能。这种多图形式化要求 \(\pi\) 能够在不同的形状和数据类型上泛化，排除了针对特定形状的取巧方法。与自由形式的内核生成不同，它进一步确保了与现有编译器流水线的可组合性以及通过标准编译器基础设施的可验证性。

### 3.2 数据集构建

数据集的构建包括两个阶段：**计算图收集**和**高级子图生成**，如图 1 (https://arxiv.org/html/2605.29357#S3.F1) 所示。总体目标是在保留真实世界计算模式的同时，构建结构多样且优化相关的子图。

![参见说明文字](图1: PassNet 数据集构建流水线。我们收集并过滤真实世界模型图，并通过递归图切分和元数据泛化生成层次化训练子图。)

#### 图收集与验证。
我们通过一个轻量级装饰器 (`pass_net.extract`) 从真实世界模型中提取计算图。在执行过程中，符号追踪捕获算子调用和张量依赖关系，生成标准化表示（高级 IR、权重和输入元数据）。为确保下游子图构建的高保真度，每个图都经过严格验证，满足五个关键约束：可运行、可序列化、可分解、可静态分析以及可访问自定义算子（见附录 A (https://arxiv.org/html/2605.29357#A1)）。

#### 子图生成策略。
为系统地映射优化空间，我们提出了三个子图类别，侧重于结构重复、融合潜力和原始行为。

(1) **通过递归折叠的经典子图选择**。
*经典子图* 是一个重复的结构模式，代表一个模型族内的常见计算模式。为大规模提取此类模式，我们设计了**递归折叠**：该方法