ML Compilers on KKKZOZ's Blog

vllm & torch.compile

Tue, 11 Aug 2026 00:00:00 +0000

vLLM 中同时存在多条与 FX 和 torch.compile 相关的执行路径，它们虽然共享 PyTorch 的 graph infrastructure，但解决的问题并不相同。理解这些路径的边界，是区分“模型结构改写”“Inductor 编译优化”和“CUDA Graph replay”的前提。

进入正文前，可以先区分三类用法：

HF Transformers modeling backend：直接使用 fx.Tracer 分析 Hugging Face forward，识别可融合结构并通过 AST 改写模型源码
vLLM 自定义 torch.compile 主路径：由 Dynamo 产生完整 FX GraphModule，再由 vLLM 控制 graph partition、shape specialization、custom passes、cache 和 CUDA Graph integration
普通 @torch.compile 小函数：FX 主要作为 Dynamo、AOTAutograd 和 Inductor 的内部 IR，vLLM 通常不直接处理这些 graph

Piecewise CUDA Graph 不是第四条独立的 tracing 路径，而是建立在 vLLM Piecewise Compilation 之上的运行时优化：编译阶段先拆分并编译 graph regions，随后只对兼容区域执行 CUDA Graph capture/replay。

这篇文章的核心问题是：vLLM 如何围绕标准 torch.compile 增加一层 LLM-serving-aware 的 graph processing、compilation policy 和 runtime integration。

Piecewise CUDA Graph

Mon, 10 Aug 2026 00:00:00 +0000

核心结论

Piecewise CUDA Graph 的核心思想是：不要求整个 Transformer forward 都满足 CUDA Graph capture 条件，而是将 CUDA Graph 不兼容的算子作为边界，只 capture 其余 CUDA-Graph-safe 区域。

在 vLLM 中，最典型的边界是 Attention。因此，PIECEWISE 模式并不是“把一个 CUDA Graph 切开执行”，而是将模型划分为多个独立的 CUDA Graph，并在它们之间以 eager 模式执行 CUDA-Graph-unsafe 操作：

CUDA Graph -> Attention (eager) -> CUDA Graph -> Attention (eager) -> CUDA Graph

CUDA Graph 会提前 capture 一组 GPU operations，后续通过 graph replay 一次性提交，从而降低 CPU 逐个启动 kernel 的开销。Transformer 中的 RMSNorm、Linear/GEMM、activation、FFN 和 residual 等计算，在 capture size 确定时通常具有稳定的执行结构；Attention 则需要处理 KV Cache、动态 shape 和运行时 metadata，对 CUDA Graph 的兼容要求更高。

ML Compilers Overview

Sun, 09 Aug 2026 00:00:00 +0000

ML Compiler Overview

ML Compiler 位于深度学习框架和底层硬件之间，作用是把 TensorFlow、PyTorch、JAX 等框架描述的高层张量计算转换成能在 CPU/GPU/TPU 等硬件上高效执行的程序。

它的典型架构是 Frontend → IR → Optimization → Backend:

Frontend 接收框架已经构造或捕获出来的计算图/算子程序（例如 tf.function 得到的 TF Graph、PyTorch Dynamo 得到的 FX Graph），转换成统一的中间表示 IR
中间层进行算子融合、常量折叠、布局变换、内存规划、并行化等与机器学习计算相关的优化
Backend 再根据目标硬件进行 lowering、调度和代码生成

因而 ML compiler 输入通常是带有 tensor shape、dtype、算子及数据依赖关系的计算图或 IR，输出则是面向特定设备的低层 IR、kernel 或可执行程序。

需要特别区分的是：graph capture 不一定属于 compiler 本身——例如 TensorFlow 由 tf.function tracing 捕获图，PyTorch 由 TorchDynamo 捕获图，然后 ML compiler 才接手，对这个图做优化并最终生成高效机器码。

IR Pipeline

下面这套 IR 划分描述的是一个典型 ML Compiler 从高层模型计算逐渐 lowering 到硬件可执行 kernel 的过程。不同编译器实际使用的 IR 名称和数据结构可能不同，但抽象层次通常可以归纳为：