菜单
主题
vLLM Logo

torch.compile 介绍及其与 vLLM 的协作方式

14 分钟阅读
Luka Govedič (Red Hat), Richard Zou (Meta), Addie Stevens (Red Hat), Kaichao You (清华大学), Michael Goin (Red Hat), Saša Zelenović (Red Hat)
#性能

注： 本博客源自我们每两周举办一次的 vLLM 办公时间 (Office Hours)，这是一个由 Red Hat 托管，邀请 vLLM 项目贡献者和加州大学伯克利分校团队参与的社区论坛。每期会议都会涵盖最新更新、特邀嘉宾深度解析以及开放问答。欢迎在每个隔周的周四美东时间下午 2:00 / 美西时间上午 11:00 通过 Google Meet 加入我们，并在我们的 YouTube 播放列表中获取录音和演示文稿。

简介

当今的大型语言模型 (LLM) 推理要求在各种硬件、工作负载和规模下尽可能高效地执行模型。高效执行需要高度优化的内核，这些内核通常需要针对不同的模型和平台进行手动调整。torch.compile 是 PyTorch 的即时 (JIT) 编译器，它能自动生成优化后的内核，使 PyTorch 代码显著提速，而无需开发人员手动针对每个支持的硬件平台进行性能调优。

对于作为便携式和高效 LLM 推理的事实标准开源引擎 vLLM 而言，torch.compile 不仅仅是一个性能增强工具。它是将优化职责从模型开发人员转移到编译器的一个核心组件。无需修改模型定义，优化在编译期间自动应用，从而实现更好的关注点分离并达到极致性能。在本文中，我们将深入探讨 torch.compile 的工作原理、它如何集成到 vLLM 中，以及 vLLM 如何利用自定义编译器处理方案来最大化性能。我们还将讨论 vLLM 中 torch.compile 集成的现状与未来工作，以进一步提升其可用性和性能。

什么是 torch.compile？

torch.compile 让您能以最小的代价优化 PyTorch 代码：使用 torch.compile 就像在函数或 torch.nn.Module 上添加装饰器一样简单。torch.compile 会自动将张量运算捕获为计算图，并为其生成优化后的代码。

在以下示例中，torch.compile 为函数 fn 中的所有点对点运算生成了一个单一的融合内核。它即时捕获并编译函数，如果任何捕获条件（例如输入形状）发生变化，它可能会进行重新编译。


图 1：torch.compile 是 PyTorch 代码的 JIT 编译器。您可以将函数、nn.Modules 和其他可调用对象包装在 torch.compile 中。

使用 torch.compile 的方法有很多。您可以将其用作内核生成器（如图 1 所示），即编译一个函数。但您也可以将 torch.compile 应用于完整的 nn.Module 模型或其子模块。根据模型的结构和您的需求（例如编译时间），我们建议在不同的地方应用 torch.compile。

为什么要使用 torch.compile？

优化模型的一种方法是编写自定义 CPU/CUDA 操作，以比模型原实现更快的方式执行相同的操作。为每个模型编写自定义内核非常耗时，且需要对性能和硬件有深刻的理解。torch.compile 几乎无需额外的工程投入就能让您达到接近峰值的性能。例如，PyTorch 的 开源 TorchBench 基准测试套件 显示在 80 多个模型上获得了 1.8-2 倍的几何平均加速。


图 2：torch.compile 为您提供快速的基准性能，从而节省您调优模型性能的开发时间。

torch.compile 的工作原理

torch.compile 流水线由两个主要阶段组成：前端 (TorchDynamo) 和后端 (TorchInductor)。我们将作简要概述，欲了解更多细节，请参阅 PyTorch 2 官方论文。

1. 前端 (TorchDynamo)：图捕获

torch.compile 的前端是一个自定义字节码解释器。它跟踪任意 Python 函数并提取仅包含张量运算的线性 torch.fx 图。torch.compile 的一个关键特性是图断点 (graph breaks)，这使它能覆盖绝大多数 Python 代码。每当 torch.compile 遇到不支持的运算时，它不会报错，而是结束当前被跟踪的图，执行该运算，然后开始跟踪一个新的图。torch.compile 将每个被跟踪的图发送到后端进行优化。

在下面的代码示例中，torch.save 是一个不支持的运算：torch.compile 不知道如何执行磁盘 I/O。将 torch.compile 应用于函数 f 等同于将 torch.compile 应用于调用 torch.save 之前的计算区域和之后的计算区域。


图 3：torch.compile 捕获张量运算的线性图，并绕过像 torch.save 这样不支持的运算。

2. 后端 (TorchInductor)：优化与内核生成

torch.compile 的后端接收来自前端的图，并通过图通行和降级（Lowering）到优化后的 C++、Triton 或其他内核来优化它们。它能够：

  • 融合点对点和归约运算
  • 自动调整块大小等内核配置
  • 选择不同的矩阵乘法后端（cuBLAS、Triton、CUTLASS）并执行序言（Prologue）和后记（Epilogue）融合。
  • 使用 CUDA Graphs 来缓存和高效重放内核启动

CUDA Graphs 是编译器发挥作用的一个例子。CUDA Graphs 减少了启动开销，但对您的代码有一定假设（例如必须仅使用 CUDA 运算，输入张量必须具有静态内存地址）。torch.compile 能够自动在不支持的运算处拆分图，从而创建可以安全使用 CUDA Graph 的较小图，并自动管理静态输入缓冲区。

vLLM 集成

vLLM V1 在线和离线推理默认集成 torch.compile。您可以使用 -O0 或 --enforce-eager 来禁用它，但在大多数使用场景中，保持开启会带来性能优势。详见文档。

编译缓存

vLLM 在冷启动期间编译模型，并将产物（FX 图、Triton 内核）保存在缓存目录（默认路径为 ~/.cache/vllm/torch_compile_cache）中。在热启动时，从缓存中检索这些产物。您可以通过 VLLM_DISABLE_COMPILE_CACHE=1 或删除缓存目录来禁用缓存。

编译后的产物和缓存可以在具有相同环境的机器之间重用。如果您有自动扩缩容需求，请确保生成一次缓存目录并将其共享给各实例。


图 4：编译产物在冷启动后会被缓存，并可在机器间重用，从而确保在正确设置时实现快速、一致的启动。

动态批处理大小与特化

默认情况下，vLLM 编译一个支持所有可能批处理大小的动态批处理大小图。这意味着一个产物可以服务多种输入尺寸。然而，针对已知的批处理大小（如 1、2 或 4）进行特化可以带来性能提升。

在配置中使用 compile_sizes: [1, 2, 4] 来触发此特化。其底层逻辑是告诉 torch.compile 针对这些静态大小进行编译，并可能执行更多的自动调优以选择最佳内核。


图 5：如何指定针对特定批处理大小进行特化编译。

分段 CUDA 图 (Piecewise CUDA Graphs)

并非所有运算都与 CUDA Graphs 兼容；例如，级联注意力 (cascade attention) 就不兼容。vLLM 通过将捕获的图分解为 CUDA Graph 安全与不安全的部分，并分别执行它们来解决这个问题。这使我们能够在不损失正确性的前提下获得 CUDA Graphs 的性能优势。


图 6：vLLM 中的分段 CUDA Graphs 捕获并重放支持的 GPU 内核序列以实现低开销执行，同时跳过级联注意力等不支持的运算。

vLLM 中的自定义编译器通行 (Custom Compiler Passes)

虽然 torch.compile 包含许多内置优化，但 vLLM 添加了自定义编译器处理方案，应用额外优化以进一步提高性能。

为什么要进行自定义处理？

模型作者编写声明式、模块化的代码，专注于正确性并使用清晰的抽象，将高级运算分离为独立的子模块并按层分组。然而，实现峰值性能通常需要打破这些抽象，例如跨子模块和层融合运算。vLLM 的自定义处理方案重写了 torch.fx 图，而不是重写模型本身。

这些处理方案可以：

  • 融合内存受限的自定义算子，如激活函数和量化
  • 添加 Inductor 中不存在的优化（如删除多余的无操作符）

示例：SiLU + 量化融合

量化 MLP 中的一种常见模式是 SiLU 激活后跟量化降维线性层。量化线性层包括对输入进行量化操作，然后是量化矩阵乘法。单独来看，SiLU 和量化运算速度缓慢且受内存带宽限制。使用 Inductor 模式匹配工具，vLLM 中的 ActivationFusionPass 自定义处理方案将它们替换为单个融合内核，吞吐量提升高达 8%。


图 7：在 8x AMD MI300 上测试 FP8 量化的 Llama 3.1 405B 模型，融合内核（黄色 fusion）优于 default（对 RMSNorm 和 SiLU 使用 torch 算子，对量化使用自定义 FP8 内核）和 custom（未融合的自定义内核）。

图 8：对比上述 fusion 和 default 方案的详细吞吐量加速比。如果通过融合消除了所有量化开销（8%），理论上最大的吞吐量提升为 8%，我们可以看到在某些情况下达到了这一改善。

注： 自办公时间以来，我们添加了使用 torch 算子的量化实现，（当由 Inductor 编译时）它比自定义 CUDA/ROCm 内核更快。由于 Inductor 可以自动将这些 torch 算子与 SiLU torch 算子融合，因此 SiLU+量化和 RMSNorm+量化处理方案在某些情况下已过时。然而，任何涉及自定义算子（注意力、集合通信、次字节量化）的融合仍需要自定义处理方案。我们展示 SiLU+量化示例是为了与办公时间的演示文稿和录音保持一致，但其他融合处理方案的工作方式非常相似。

示例：序列并行 + 异步 TP

当使用张量并行 (TP) 时，线性层会对权重进行分片并计算不完整的矩阵乘法结果，这些结果需要在 GPU 之间同步。当计算和通信部分使用单独的内核时，我们会产生通信开销，因为 GPU 在等待通信结果的网络延迟时处于空闲状态。

相反，我们可以通过使用融合的 GEMM+集合通信内核来重叠计算和通信。此类内核的一个示例是 GEMM+reduce_scatter 和 all_gather+GEMM 内核。为了利用这些内核，我们需要将 all_reduce 集合通信操作分解为 reduce_scatter 和 all_gather，同时将 all_gather 推迟到层归一化 (layernorm) 之后，以使其能够与后续的 GEMM 融合。

如果我们要在模型定义中实现这种优化，我们就必须触及 vLLM 支持的每一个模型（有数百个！）。这不仅具有侵入性、破坏抽象、增加开发摩擦，而且最初也很难被 vLLM 采纳。相反，通过在 torch.compile 中实现该优化，它仅被限制在 2 个自定义处理方案中，并且可以通过 CLI 标志开启，从而为 vLLM 支持的所有模型提供更好的性能。

注： 此优化由社区成员 @cascade812 完全实现，感谢其杰出的贡献。有关异步 TP 的更多信息，请参阅 PyTorch 博客。

当前与即将推出的处理方案

现已可用

  • 融合处理方案：
    • RMSNorm + 量化 (FP8) 融合
    • SiLU-Mul + 量化 (FP8) 融合
    • 注意力 + 量化 (FP8) 融合（最高 7% 提升）
    • AllReduce + RMSNorm 融合（最高 15% 提升）
    • AllReduce + RMSNorm + 量化 (FP8) 融合（最高 8% 提升）
    • AllReduce + RMSNorm + 量化 (FP4) 融合（最高 10% 提升）
    • 序列并行 & 异步 TP（最高 10% 提升）
  • 其他处理方案：
    • 无操作符消除：消除或简化冗余的重塑 (reshape) 运算
    • 函数化修复：手动重新替换 auto_functionalized 运算，以避免冗余复制和内存使用

即将推出：

  • 注意力 + 量化 (FP4) 融合：#22703
  • SiLU-Mul + 量化 (FP4) 融合：#22448

处理方案可以通过 PostGradPassManager、CLI (--compilation-config) 或在离线模式下指定配置对象来添加。这允许 vLLM 用户在不修改 vLLM 源代码的情况下，执行其用例所需的自定义图转换（内核替换或其他）。

未来工作

vLLM 与 torch.compile 的集成已经走过了很长一段路。以下是我们未来六个月关注的一些领域。

提高稳定性
vLLM 与 torch.compile 的集成使用了许多私有（以单下划线开头）的 torch.compile API，并依赖于不稳定的实现细节。我们这样做是因为使用公共 torch.compile API 不足以满足我们的要求——vLLM 需要快速的服务性能，并且在模型服务期间不能进行重新编译。这导致了一些问题，如缓存混乱，或者在某些模型上需要禁用 vLLM 的 torch.compile 缓存。PyTorch 编译器团队正致力于将 vLLM（及通用推理）相关功能从 vLLM 上游合并到 torch.compile，并将 vLLM 迁移到使用更稳定的 API。许多这些特性已经存在于 torch 2.8 中，即将引入 vLLM！

缩短启动时间
我们了解到，启动时间是 vLLM torch.compile 和 CUDA Graphs 的一大痛点，特别是在根据需求动态启动新机器的自动扩缩容场景中。我们计划显著减少 vLLM 的冷启动（首次）和热启动（第二次及以后）时间，特别是与 Dynamo 和 Inductor 编译相关的部分。请关注 GitHub 上的 startup-ux 标签，或加入 vLLM Slack 上的 #feat-startup-ux 频道以了解最新进展！

一个重要的用户体验改进是计划改进的 -O 命令行标志。通过在 vLLM CLI 上指定 -O<n>（其中 n 是 0-3 之间的整数），用户可以更容易地直接控制在启动时间和性能之间进行权衡。虽然 -O0 几乎不执行任何优化并尽快启动，但 -O3 将花费更多时间但提供最佳性能。

自定义处理方案改进
我们计划对自定义处理机制进行几项广泛的改进，以提高它们的灵活性并使其更易于编写，同时提升已应用优化的最终性能：

  • 编译多个动态形状的 torch.fx 图。这将让我们根据批处理大小特化前向传播图，而无需为每个静态大小单独编译。更多信息详见 RFC。
  • 启用自定义算子的 torch 实现匹配。目前，自定义算子（rms_norm、quant 等）需要被启用以允许模式匹配和融合，但可能存在未被融合的自定义算子（特别是每层发生 4 次的量化）。这些算子比它们的 torch 等效实现慢，这降低了融合的好处。我们有一个原型可以模式匹配自定义算子的 torch 实现，预示着进一步的性能提升。

实验性 torch.compile 后端集成
我们也在探索实验性的 MPK/Mirage 编译器集成。MPK 是一种精度调度超大规模内核 (megakernel) 编译器，这意味着它为整个模型的前向传播生成一个单一内核，与 CUDA Graphs 相比，它可以进一步降低 CPU 开销并消除内核启动开销。关于提议集成的更多信息详见 RFC。

其他性能改进
vLLM 的 torch.compile 集成目标是提供良好的基准性能，以避免需要编写和维护大量自定义内核。我们将继续维护和提升性能。正在进行的工作亮点包括：

  • 改进 FlexAttention 支持。FlexAttention 是一种 API，允许使用不同的注意力变体，而无需为每种变体编写自定义注意力内核。其底层使用 torch.compile 来生成自定义 Triton 模板。
  • 完整 CUDA Graphs 支持 Flash Attention v2 和 FlashInfer。完整 CUDA Graphs 比分段 CUDA Graphs 的开销更小，应该在高开销设置中提高性能。

总结

torch.compile 为加速 PyTorch 模型提供了一种强大且易用的方式。在 vLLM 中，它是推理流水线的核心部分。结合缓存、动态形状支持、CUDA Graphs 和自定义处理方案，它能在任何环境中实现高效、可扩展的 LLM 服务。

随着编译器栈的成熟和对新硬件支持的扩展，torch.compile 和 vLLM 将继续突破推理性能的极限——同时保持模型开发的清晰和模块化。在 PyTorch 文档和 vLLM 文档中阅读有关 torch.compile 的更多信息，并加入 vLLM Slack 上的 #sig-torch-compile 频道以提问、分享反馈并贡献您自己的自定义处理方案！