Introduction to torch.compile and How It Works with vLLM

发表时间: 2025-08 · Blog post by Luka Govedič (Red Hat), Richard Zou (Meta), Addie Stevens (Red Hat), Kaichao You (Tsinghua University), Michael Goin (Red Hat), Saša Zelenović (Red Hat) (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2025-08-20-torch-compile

作者/机构：Luka Govedič (Red Hat), Richard Zou (Meta), Addie Stevens (Red Hat), Kaichao You (Tsinghua University), Michael Goin (Red Hat), Saša Zelenović (Red Hat)

速读

一句话结论
vLLM 深度集成了 torch.compile，通过引入分段 CUDA Graphs 和自定义编译 Pass，在不修改模型定义代码的前提下自动完成算子融合与通信掩盖，实现了大语言模型推理的高效跨平台执行。

要解决什么问题
大语言模型推理的极速执行高度依赖底层算子的极致优化。传统的做法是为不同的模型和硬件平台手动编写和调优自定义 CPU 或 CUDA 算子。这种路线卡在三个地方：首先，开发成本极高，要求开发者具备深厚的硬件和性能调优知识；其次，为了追求极致性能（例如跨层融合算子），模型作者往往被迫破坏原本清晰、模块化的代码抽象，导致代码难以维护，且 vLLM 作为支持数百个模型的推理引擎，逐一修改模型代码极具侵入性且极难被社区接受；最后，现有的通用编译加速方案存在局限，例如标准的 CUDA Graphs 无法捕获级联注意力（cascade attention）等操作，而原生的 torch.compile 缺乏针对大模型量化或分布式张量并行（TP）的特定优化逻辑，导致显存受限的激活函数、量化操作以及多卡通信延迟依然是拖慢吞吐量的瓶颈。

怎么做的
核心思路是将优化的责任从模型开发者转移给即时编译器（JIT），通过拦截并重写计算图来自动生成高性能内核。整个机制由前端捕获、后端生成以及 vLLM 的定制化增强三个部件构成。
前端 TorchDynamo 是一个自定义字节码解释器，负责图捕获。它追踪 Python 代码并提取仅包含张量操作的直线型 torch.fx 计算图。遇到不支持的操作（如磁盘 I/O）时，它不会报错，而是触发图断裂（graph breaks），将计算切分为多个安全的子图。后端 TorchInductor 接收这些图，执行逐点（pointwise）和归约（reduction）操作的融合，自动调优 block size 等内核配置，并最终生成优化的 C++ 或 Triton 内核。
为了适配大模型推理，vLLM 在此基础上做了四项关键设计：
第一，编译缓存与动态批处理。vLLM 在冷启动时编译模型，将 FX 图和 Triton 内核保存到本地缓存目录，热启动时直接复用，且默认编译支持所有批大小的动态图，同时允许用户通过配置特化已知的批大小（如 1、2、4）以触发更深度的自动调优。
第二，分段 CUDA Graphs（Piecewise CUDA Graphs）。针对级联注意力等不支持 CUDA Graphs 的操作，vLLM 将捕获到的图强行拆分为安全和不安全两部分，仅对安全部分执行低开销的内核重放，从而在保证正确性的同时绕开图捕获失效的卡点。
第三，自定义编译器 Pass。这是绕开原生编译器局限的核心。vLLM 在图层面对 torch.fx 进行重写，打破模型代码的模块化抽象以实现深度融合。例如，在处理量化 MLP 时，原有的 SiLU 激活和量化操作都是访存密集型且缓慢的，vLLM 引入 ActivationFusionPass 将其替换为单一融合内核。其图重写逻辑可抽象为：

$$ \text{FusedKernel}(x) = \text{Quantize}(\text{SiLU}(x)) $$


第四，序列并行与异步 TP。针对张量并行中的通信等待卡点，自定义 Pass 会将 all_reduce 集合通信分解为 reduce_scatter 和 all_gather，并将 all_gather 的执行时机后置到 LayerNorm 之后，使其能与后续的 GEMM 操作融合，成功让计算与跨卡通信的时间相互掩盖。此外，还引入了消除冗余 reshape 操作的无操作消除（No-op Elimination）等优化。

效果如何
实验在 8 张 AMD MI300s 显卡上运行 FP8 量化的 Llama 3.1 405B 模型。对比基线包含两种：default 基线（使用原生 PyTorch 算子执行 RMSNorm 和 SiLU，配合自定义 FP8 量化算子）以及 custom 基线（使用未融合的纯自定义算子）。
量化结果表明，vLLM 的自定义 Pass 优化效果显著。SiLU 与量化融合的 Pass 成功消除了量化带来的额外开销，使吞吐量提升了最高 8%（达到了理论上的最大提升上限）。其他自定义 Pass 也带来了可观的收益：AllReduce 与 RMSNorm 融合最高提升 15%，序列并行与异步 TP 最高提升 10%。在更广泛的 PyTorch TorchBench 测试集中，torch.compile 相比 eager 模式在 80 多个模型上实现了 1.8 到 2 倍的几何平均加速比。
但该方案当前也存在明显的代价和局限：首先是冷启动时间过长，这在需要动态扩缩容的场景下是一个巨大的痛点，为此作者计划重构 -O 命令行标志，允许用户在极速启动（-O0）和极致性能（-O3）之间做权衡；其次，为了追求极致的在线服务性能并避免重新编译，vLLM 深度依赖了 torch.compile 内部不稳定的私有 API，这导致了偶发的缓存异常，甚至在某些模型上不得不手动禁用编译缓存；最后，如果某些自定义算子（如量化）未能成功触发模式匹配并被融合，它们的回退执行速度甚至会慢于 PyTorch 的原生等价实现，从而削弱融合带来的收益。

1. 主要贡献

当前快速的大语言模型（LLM）推理要求在各种硬件、工作负载和规模下尽可能高效地执行模型。这种高效执行通常依赖于高度优化的内核，而这些内核往往需要针对不同的模型和平台进行手动调优。本文深入探讨了 PyTorch 的即时（JIT）编译器 torch.compile 如何作为核心组件集成到 vLLM（事实上的开源高效 LLM 推理引擎）中。

本文的核心创新与贡献在于：

  • 优化责任转移：通过集成 torch.compile，vLLM 将性能优化的责任从模型开发者转移到了编译器上，无需修改模型定义即可在编译阶段应用优化，实现了关注点的清晰分离。
  • 系统级集成方案：详细阐述了 vLLM 如何利用编译缓存、动态批处理大小特化以及分段 CUDA Graphs（Piecewise CUDA Graphs）来最大化推理性能。
  • 自定义编译器 Pass 机制：介绍了 vLLM 如何通过自定义编译器 Pass（如 SiLU 与量化融合、序列并行与异步张量并行）打破模型代码抽象，重写计算图，从而在不侵入数百个模型代码的前提下实现极致的底层性能优化。

2. 背景知识与工作原理

降低性能优化门槛的自动内核生成机制。优化模型的一种传统且耗时的方法是为每个模型编写自定义的 CPU/CUDA 操作。torch.compile 允许开发者以极小的额外工程代价达到接近硬件峰值的性能。根据 PyTorch 的 TorchBench 基准测试套件显示，它能在 80 多个模型上实现 1.8 到 2 倍的几何平均加速比。在实际使用中，只需将装饰器应用于函数或 torch.nn.Module，torch.compile 就能自动将张量操作捕获到计算图中，并为其生成优化的内核。它在即时（just-in-time）捕获并编译函数，如果捕获条件（例如输入形状）发生变化，则会触发重新编译。
Figure 1：torch.compile 是 PyTorch 代码的 JIT 编译器。你可以将函数、nn.Modules 和其他可调用对象包装在 torch.compile 中。
Figure 2：torch.compile 为你提供快速的基础性能，从而节省你调整模型性能的开发时间。

基于 TorchDynamo 的前端图捕获技术。torch.compile 的前端是一个自定义的字节码解释器。它追踪任意的 Python 函数，并提取仅由张量（Tensor）操作组成的直线型 torch.fx 计算图。其赋予前端良好代码覆盖率的核心特性是图中断（graph breaks）。当 torch.compile 遇到无法支持的操作（例如 torch.save 涉及磁盘 I/O）时，它不会抛出错误，而是结束当前正在追踪的图，执行该不支持的操作，然后开始追踪一个新的图。前端随后将追踪到的每个计算图发送给后端进行优化。
Figure 3：torch.compile 捕获张量操作的直线型计算图，并能绕过如 torch.save 等不支持的操作。

基于 TorchInductor 的后端优化与内核生成。torch.compile 的后端接收来自前端的计算图，并通过图传递（graph passes）和底层化（lowering）将其优化为高效的 C++、Triton 或其他内核。后端具备强大的优化能力，包括融合逐点（pointwise）和归约（reduction）操作、自动调整内核配置（如块大小）、在不同的矩阵乘法后端（如 cuBLAS、Triton、CUTLASS）之间进行选择并执行前序和后序融合。此外，它还能自动管理静态输入缓冲区，并在不支持的操作处自动分割图，从而创建可以安全使用 CUDA Graphs 的较小计算图，以有效缓存和重放内核启动并减少启动开销。

3. vLLM 集成与方法细节

vLLM 中的默认集成与编译缓存机制。vLLM V1 在在线和离线推理中默认集成了 torch.compile（可通过 -O0 禁用）。vLLM 在冷启动（cold start）期间编译模型，并将生成的工件（如 FX 图、Triton 内核）保存在默认的缓存目录 ~/.cache/vllm/torch_compile_cache 中。在热启动（warm start）时，系统会直接从缓存中检索这些工件。编译后的工件和缓存可以在具有相同环境的不同机器之间复用，对于自动扩缩容（autoscaling）的用例，只需生成一次缓存目录并在多个实例之间共享即可确保快速、一致的启动。
Figure 4：编译后的工件在冷启动后被缓存，正确设置后可跨机器重用，以确保快速一致的启动。

动态批处理大小与特化编译策略。默认情况下，vLLM 编译一个支持所有可能批处理大小的单一动态批处理大小计算图，使得一个编译工件可以服务于可变的输入大小。然而，用户可以通过在配置中设置 compile_sizes: [1, 2, 4] 来触发特化机制。在底层，该配置指示 torch.compile 针对这些静态大小进行专门的编译，并执行更多的自动调优以选择最佳内核，从而针对已知的批处理大小获得进一步的性能提升。
Figure 5：如何指定针对特定批处理大小的特化编译。
Figure 6：特化编译的代码示例。

分段 CUDA Graphs 的实现机制。由于并非所有的计算操作都兼容 CUDA Graphs（例如级联注意力机制 cascade attention 就不兼容），vLLM 采用了一种分段处理的策略。vLLM 通过将捕获的图分解为对 CUDA Graph 安全和不安全的部分，并分别执行它们。这种设计跳过了不支持的操作，同时对支持的 GPU 内核序列进行捕获和重放，从而在不损失正确性的前提下，获得了低开销执行的性能优势。
Figure 7：vLLM 中的分段 CUDA Graphs 捕获并重放支持的 GPU 内核序列以实现低开销执行，同时跳过不支持的操作（如级联注意力）。

引入自定义编译器 Pass 的动机。模型作者通常编写声明式、模块化的代码，将高级操作分离到独立的子模块中并按层分组。然而，实现峰值性能往往需要打破这些抽象界限，跨子模块和层融合操作。为了避免重写 vLLM 支持的数百个模型，vLLM 引入了自定义的编译器 Pass 来直接重写 torch.fx 图。这些 Pass 能够融合受限于内存带宽的自定义操作（如激活函数和量化），并添加 Inductor 中不存在的额外优化（如消除冗余的无用操作）。

SiLU 与量化融合的实现与演进。在量化的 MLP 层中，常规模式是 SiLU 激活函数紧跟一个量化的向下投影（down-projection）线性层（包含量化操作和量化矩阵乘法）。由于单独的 SiLU 和量化操作受限于内存带宽且速度较慢，vLLM 利用 Inductor 的模式匹配工具开发了 ActivationFusionPass。该 Pass 将这些操作替换为单一的融合内核。值得注意的是，vLLM 后续添加了基于 PyTorch 原生操作的量化实现，Inductor 可以自动将这些操作与 SiLU 融合，使得该特定 Pass 在某些情况下被原生支持取代。但任何涉及自定义操作（如注意力机制、集合通信）的融合仍严格依赖此自定义 Pass 机制。

序列并行与异步张量并行的计算通信重叠。当使用张量并行（TP）时，线性层对权重进行分片并计算出不完整的结果，这些结果需要跨 GPU 同步。如果分别使用独立的内核处理计算和通信，GPU 会因等待网络延迟而产生空闲开销。vLLM 通过使用融合的 GEMM+集合通信内核（如 GEMM+reduce_scatter 和 all_gather+GEMM）来重叠计算和通信。在具体实现上，系统将 all_reduce 集合操作分解为 reduce_scatter 和 all_gather，同时将 all_gather 推迟到 LayerNorm 之后，以便与后续的 GEMM 融合。这一复杂的优化被封装在 2 个自定义 Pass 中，通过命令行标志即可为所有模型开启。

自定义 Pass 的扩展与管理方式。vLLM 提供了灵活的接口供用户扩展，用户可以通过 PostGradPassManager、命令行参数（--compilation-config）或在离线模式下指定配置对象来添加自定义 Pass。这允许用户在不修改 vLLM 源代码的情况下，执行其特定用例所需的计算图转换或内核替换。

4. 实验环境与结果

实验环境配置：

  • 硬件配置：8x AMD MI300s GPUs。
  • 模型配置：Llama 3.1 405B（量化为 FP8）。
  • 软件环境：vLLM, torch.compile, 自定义 CUDA/ROCm 内核。

实验结果与分析：

  1. SiLU + 量化融合性能对比：

    • 实验内容：在 Llama 3.1 405B FP8 模型上，对比了三种执行模式：fusion（融合内核）、default（使用 torch ops 执行 RMSNorm 和 SiLU，加上自定义 FP8 量化内核）以及 custom（未融合的自定义内核）。
    • 实验结果：fusion（融合内核）在吞吐量上显著优于 default 和 custom 模式。
    • 分析结论：如果通过融合完全消除量化开销（约 8%），理论最大吞吐量提升为 8%，实验证实 fusion 模式在部分情况下成功达到了这一理论最大提升比例。
      Figure 8：在 8x AMD MI300s 上测试量化为 FP8 的 Llama 3.1 405B，融合内核（fusion，黄色）优于 default 和 custom（未融合的自定义内核）。
      Figure 9：详细的吞吐量加速比对比。如果通过融合消除了所有量化开销（8%），理论最大吞吐量提升将是 8%，我们可以看到在某些情况下达到了该提升。

  2. 当前可用自定义 Pass 的性能提升汇总：

    • Attention + Quant (FP8) 融合：实现了高达 7% 的性能提升。
    • AllReduce + RMSNorm 融合：实现了高达 15% 的性能提升。
    • AllReduce + RMSNorm + Quant (FP8) 融合：实现了高达 8% 的性能提升。
    • AllReduce + RMSNorm + Quant (FP4) 融合：实现了高达 10% 的性能提升。
    • 序列并行与异步张量并行 (Async TP)：实现了高达 10% 的性能提升。
    • 其他基础优化：无用操作消除（成功消除或简化了冗余的 reshape 操作）；修复函数化（避免了冗余的内存拷贝）。

5. 结论与未来工作

torch.compile 提供了一种强大且易于使用的途径来加速 PyTorch 模型，并已成为 vLLM 推理管道的核心组成部分。通过结合编译缓存、动态形状支持、CUDA Graphs 以及自定义编译器 Pass，它在任何环境下都能实现高效、可扩展的 LLM 服务。

在未来的六个月内，vLLM 与 PyTorch 编译团队将重点推进以下工作：

  1. 提升稳定性：目前 vLLM 依赖了许多私有的 torch.compile API 以满足无重编译的快速服务需求，这导致了部分缓存问题。团队正计划将推理相关功能合并到 PyTorch 主干，并迁移至 Torch 2.8 中更稳定的 API。
  2. 优化启动时间：针对自动扩缩容场景中的冷热启动痛点，计划重构命令行标志。用户可通过 -O<n>（0-3）精确控制启动时间与性能的权衡（-O0 追求极速启动，-O3 追求极致性能）。
  3. 改进自定义 Pass 机制：计划支持编译包含多个动态形状的 torch.fx 图（无需为每个静态批处理大小单独编译），并支持对自定义操作（如 RMSNorm、量化）的 PyTorch 原生实现进行模式匹配以进一步提升融合收益。
  4. 实验性编译器后端集成：探索 MPK/Mirage 编译器集成，旨在为整个模型前向传递生成一个单一的巨型内核（megakernel），从而彻底消除内核启动开销。
  5. 拓展支持范围：改进对 FlexAttention 的支持（利用 torch.compile 生成自定义 Triton 模板），并为 Flash Attention v2 和 FlashInfer 提供开销更低的完整 CUDA Graphs 支持。