vLLM TorchDynamo 编译：FX Graph 分割原理与实践

Mon, 01 Jun 2026 10:40:00 +0800

背景

vLLM 利用 PyTorch 2.x 的 torch.compile 路径，通过 TorchDynamo 捕获模型计算图，再经过图分割（graph splitting）和分段编译（piecewise compilation）来优化 GPU kernel 执行效率。本文将深入剖析 FX Graph 的分割原理及其背后的设计思想。

一、TorchDynamo 编译概览

TorchDynamo 是一个 Python 级别的 JIT 编译器，它通过 PEP 523 的 frame evaluation callback 在 Python 字节码执行之前捕获计算图。vLLM 利用这一机制，将模型 forward 函数中的计算捕获为 fx.GraphModule，然后送入自定义后端 VllmBackend 进行编译。

编译流程大致如下：

model.forward()
└── TorchDynamo 捕获计算图
└── fx.GraphModule (原始完整图)
└── VllmBackend.__call__()
├── split_graph() → 图分割
├── PiecewiseCompileInterpreter → 分段编译
└── codegen → 生成胶水函数
└── 返回可调用对象给 Dynamo

二、FX Graph 分割：split_graph 的工作原理

2.1 为什么需要分割

vLLM 的模型计算中包含两种性质不同的算子：

编译 on Richelieu's Blog

vLLM TorchDynamo 编译：FX Graph 分割原理与实践

背景

一、TorchDynamo 编译概览

二、FX Graph 分割：split_graph 的工作原理

2.1 为什么需要分割