vLLM TorchDynamo 编译：FX Graph 分割原理与实践
背景 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 的模型计算中包含两种性质不同的算子： ...