vLLM 源码完全拆解（十二）：piecewise compilation 与 CUDA Graph——消灭 kernel 调度开销的两层机制

一次 decode 步的 CPU 账单#

自回归 decode 每生成一个 token，就要完整跑一遍模型前向。以 Llama-3-70B 为例：80 个 Transformer 层，每层至少包含 QKV 投影、注意力、o_proj、Gate/Up 投影、SiLU、Down 投影、两次 RMSNorm，再加上采样前的 lm_head。即使是融合后的 kernel，总数也在 448 个以上（80 层 × 每层 5-6 个 kernel，外加 embedding、norm、head）。

每个 CUDA kernel launch 在 CPU 侧的开销是微秒量级。PyTorch 官方在 CUDA Graphs 的技术文章里给出的典型数字是 5-10 微秒；对现代 GPU 上吞吐型的大 kernel，这个开销可以忽略——GEMM 跑 200 微秒，launch 占 3%，无所谓。但 decode 恰好是反面：batch size 小的时候，每个 token 的 GEMM 只有几百纳秒到几微秒的计算量，此时一个 decode 步的 448 个 kernel 中，绝大多数是短平快的小算子（norm、激活、逐元素加），GPU 几十微秒就算完，然后空等 CPU 发下一个 launch。

这笔账很容易算成「CPU bound」：

其中 $t_{\text{launch}}^{(i)}$ 是第 $i$ 个 kernel 的 CPU launch 开销（包括 Python 解释器逐行执行 forward、nn.Module.__call__ 钩子检查、dispatch 到 ATen、再把参数打包进 CUDA launch），$t_{\text{gpu}}^{(i)}$ 是 GPU 执行时间。串行流水线上，每一步的墙钟时间是两者的较大值之和。当 batch 小、kernel 短时，t_launch 普遍压过 t_gpu——这就是推理社区常说的「launch bound」或「CPU bound decode」。

vLLM 集成 piecewise CUDA graph 的核心 PR（#10058）给出的实测数字能说明问题：benchmark_latency.py 下，V1 引擎不带 piecewise CUDA Graph 时，OPT-125M 一次前向 0.41 秒，开启后 0.21 秒；Llama-3.1-8B 从 1.26 秒降到 1.04 秒。0.2-0.5 秒级别的收益几乎全部来自 CPU 侧——模型代码一行没改。

三段开销，三个对应的武器#

把 eager decode 的 CPU 开销拆细，其实是三段，每段武器不同：

1. Python 字节码与框架派发开销：nn.Module 逐层调用、hook 遍历、torch.autograd.Function 的记录逻辑。武器是 Dynamo 图捕获——把整个 forward 编译成一张 FX 图，Python 循环消失。
2. kernel launch 开销：CPU 把 kernel 参数写进 command buffer、提交给 GPU。武器是 CUDA Graph 回放——一次 graph.replay() 提交整段 kernel 序列。
3. kernel 自身的低效：内存搬运多、融合不足、没为当前 shape 调参。武器是 Inductor 编译——生成融合后的 Triton kernel，必要时 autotune。

上图来自 PyTorch 官方博客对 CUDA Graphs 的介绍：CPU 侧逐个 launch（上）与一次 graph launch（下）相比，GPU 时间线上的空洞明显收敛。vLLM 没有手写 CUDA Graph 逻辑，而是把这个收益通过编译子系统自动带给所有模型。

这三段分别对应 vLLM 编译子系统的三层结构：Dynamo 层（support_torch_compile）、图处理层（VllmBackend + Inductor passes）、静态执行层（CUDAGraphWrapper）。一个容易被混淆的认知是「CUDA Graph 等于 torch.compile」：不是，vLLM 的层级是 torch.compile（Dynamo+Inductor）负责消灭开销 1 和 3，CUDA Graph 负责消灭开销 2，二者是叠加关系而非替代关系——这一层叠法正是本篇的主角。

为什么不能直接 torch.compile(model, mode="reduce-overhead")#

mode="reduce-overhead" 就是 Dynamo + Inductor + CUDA Graph 的官方组合，框架自带，vLLM 为什么不走？根源在 vLLM 的推理模式与该 API 的假设冲突：

• CUDA Graph 要求 shape 与地址全静态。标准 CUDA Graph 语义下，图里所有 kernel 的 shape、数据指针在 capture 时固定，回放时一字都不能变。而 vLLM 的 token 数每个 step 都可能不同（调度器决定）、注意力元数据（block table、seq lens）每个 step 都要变。
• 注意力算子是 CUDA Graph 的天敌。PagedAttention 的执行计划依赖 ragged 的变长输入和分页 KV cache 指针，是逐 batch 动态的。强行捕获进图，要么捕获到错误的行为，要么每次 batch 变化就重捕获——后者开销巨大。
• Dynamo 的 guard 机制会产生运行时检查。每次调用编译后的函数，Dynamo 要重新评估守卫条件（输入 shape 是否变了、参数是否还是原来的对象）。在每秒数百次的 decode step 里，guard 评估本身成了新的 CPU 开销。
• 普通 torch.compile 允许反复重编译（graph break 之后遇到新 shape 就再编一次），这在长稳态服务里不可接受：一个请求触发重编译，引擎整个卡住，响应延迟出现尖刺。

vLLM 的答案是自管整个编译过程：Dynamo 只用一次，把模型变成一张图；图由 vLLM 自己切、自己编、自己缓存；CUDA Graph 只捕获注意力之外的部分；运行时彻底绕开 Dynamo guard。这套机制的每一个细节都是为了「编译产物可控 + 运行时零意外编译」。

编译等级：0/1/2/3 四级#

CompilationConfig.level 是整个子系统的总开关（vllm/config/compilation.py，V0.10.1 位于 vllm/config/__init__.py 附近的 CompilationLevel）。PR #9058 定义、PR #9715 定型的四级语义：

V1 引擎在模型不是 enforce_eager 时默认 level=3，此时又默认 cudagraph_mode=PIECEWISE。level 1 和 2 的区别只在「Dynamo 结果是否复用」：level 1 每次调用都重走 Dynamo（像 aot_eager 那样，调试用）；level 2 只允许一次追踪，运行时通过自定义 dispatcher 直接跳到编译好的字节码（下面会讲 dispatch_to_code 怎么做到）。真正生产用的是 3。

为什么 level 3 要禁用自定义算子？因为 Inductor 只认识 ATen 算子和带 fake impl 的自定义算子；vLLM 的大量手写 CUDA kernel（torch.ops._C.*）如果不注册成规范的自定义算子，Dynamo/Inductor 看不见内部结构，只能整块当黑盒——黑盒太多，融合就没戏。所以 PIECEWISE 模式下 vLLM 关闭 custom ops（custom_ops 默认从「全开」变成「全关」），让 Inductor 看到 ATen 层面的图，把 RMSNorm、激活、逐元素乘这些小算子重新融合成 Triton kernel——代价是失去手写 kernel 的极致优化，收益是整体图优化。这个取舍在 vLLM 代码里被反复权衡（custom_ops 文档字符串写明：「By default, all custom ops are enabled when running without Inductor and disabled when running with Inductor」）。

切分算子：splitting_ops 与「铰链」#

splitting_ops 决定图在哪里被剪开。V1 的默认值由 set_splitting_ops_for_v1() 设置：

1
# vllm/config/compilation.py (V0.10.1)
2
_attention_ops: ClassVar[list[str]] = [
3
    "vllm.unified_attention",
4
    "vllm.unified_attention_with_output",
5
    "vllm.mamba_mixer2",
6
]

只有注意力（和 Mamba 混合模型的 mixer 算子）被列为默认切点。为什么是它们？两个共性：一是执行计划动态——PagedAttention 每步的执行依赖 batch 的变长结构，shape 语义无法静态化；二是输出 shape 可静态描述——注意力的输出形状就是 query 的形状，这一性质让 Dynamo 可以在 fake-tensor 追踪时把注意力当黑盒「穿透」，整个模型仍然产出一张完整图，不需要 graph break。fake impl 一行：

1
def unified_attention_with_output_fake(query, key, value, output, layer_name, output_scale=None):
2
    return
3

4
direct_register_custom_op(
5
    op_name="unified_attention_with_output",
6
    op_func=unified_attention_with_output,
7
    mutates_args=["output"],   # 声明它原地写 output
8
    fake_impl=unified_attention_with_output_fake,
9
    dispatch_key=current_platform.dispatch_key,
10
)

注意这里的工程巧思：注意力计算本来是「返回一个新 tensor」，但 vLLM 把它改写成带 out 参数的 in-place 语义（mutates_args=["output"]）。为什么要多此一举？后面 CUDA Graph 捕获时，每个子图的输出 buffer 必须是地址固定的静态缓冲区。如果注意力返回新 tensor，那么注意力前后的两段子图之间就得传输一个变地址的中间量，CUDA Graph 的静态缓冲区约定就被破坏了。改成 out 风格后，注意力把结果写进一个调用方预先分配好的固定 buffer，前后子图各自回放，数据只在这个约定 buffer 里流转。

splitting_ops 为空列表时会退化成什么？代码里写得直白：PIECEWISE 编译等级 + 空 splitting_ops + PIECEWISE 图模式时，引擎会自动把 cudagraph_mode 升成 FULL——因为没有任何切点，图就是一张，piecewise 模式名不副实。这是常见的用户陷阱：显式传了 --compilation-config '{"splitting_ops": []}'，行为却和预期不同。

CUDAGraphMode：一张图还是几段图#

V1 的 CUDAGraphMode 枚举比常见的「开/关」丰富：

1
class CUDAGraphMode(enum.Enum):
2
    NONE = 0
3
    PIECEWISE = 1
4
    FULL = 2
5
    FULL_DECODE_ONLY = (FULL, NONE)
6
    FULL_AND_PIECEWISE = (FULL, PIECEWISE)

后两个复合值把 decode 和「混合 prefill-decode」两种 batch 分开路由：decode_mode() 取元组第一个元素给 decode batch，mixed_mode() 取第二个给普通 batch。这套双路由解决一个真实问题：decode batch 的 shape 每步基本不变（每请求 1 token），full graph 稳定可用；而 prefill/混批的 token 数 千差万别，full graph 命中率低，piecewise 更合适。FULL_AND_PIECEWISE（decode 用 FULL、混批用 PIECEWISE）是官方注释里「对多数模型性能最好的模式」。

capture 尺寸：为什么默认不是所有 batch size#

CUDA Graph 按 (batch size) 一图一档，vLLM 不可能给所有 token 数都捕获。_set_cudagraph_sizes() 给出默认档位:

1
# VllmConfig._set_cudagraph_sizes (核心逻辑)
2
possible_sizes = [1, 2, 4] + [8 * i for i in range(1, 1025)]
3
# 找到 >= max_num_seqs 的最小值作为 capture 上限
4
larger_sizes = [x for x in possible_sizes if x >= max_num_seqs]
5
max_batchsize_to_capture = larger_sizes[0]
6
batch_size_capture_list = [s for s in possible_sizes if s <= max_batchsize_to_capture]

默认语义：1、2、4 三个小档位单独捕获（decode 最常见的小 batch），之后按 8 为步长递增到覆盖 max_num_seqs。若 max_num_seqs=256，捕获档位是 [256, 248, ..., 16, 8, 4, 2, 1]（降序），约 35 档。

为什么小 batch 用细粒度、大 batch 用粗粒度？padding 浪费。batch size 33 来了，向上取整到 40，padding 了 7 个 token（21% 浪费）；batch 250 来了，取整到 256，只 padding 6 个（2.4% 浪费）。小 batch 恰恰是 decode 延迟敏感的场景，padding 比例最不能忍，所以 1/2/4 单独建档，8 以上统一以 8 为步——8/16/24/32…每档之间 padding 从不超 7 个 token，而 7 个 token 在 8 档大 batch 下的占比可以忽略。

运行时用一张反查表把任意 batch size 映射到档位：

1
# vllm/config/__init__.py (VllmConfig.pad_for_cudagraph)
2
def pad_for_cudagraph(self, batch_size: int) -> int:
3
    # batch_size 超过 max_capture_size 会 IndexError,
4
    # 调用方需要先用 cudagraph_batch_sizes[-1] 判断,
5
    # 超出上限的 batch 走 eager
6
    return self.compilation_config.bs_to_padded_graph_size[batch_size]

bs_to_padded_graph_size 是一张预计算的定长数组: index 是真实 batch size, value 是该 batch 应 pad 到哪个档位。capture_model() 的 execute_model 路径里先判断 num_scheduled_tokens <= self.cudagraph_batch_sizes[-1](即最大捕获档位),低于上限才查表 padding,否则走 eager 并把 token 数 pad 到 TP 的倍数(为了序列并行的 all-reduce 对齐)。

影响编译产物的所有因子都要进 hash#

编译产物能跨进程复用的前提是「产物只依赖可枚举的输入」。VllmBackend.__call__ 里的 cache dir 推导列出了四大因子：

1
# vllm/compilation/backends.py (VllmBackend.__call__ 节选)
2
factors = []
3
factors.append(envs.compute_hash())          # 0. 环境变量(如 VLLM_PP_LAYER_PARTITION)
4
factors.append(vllm_config.compute_hash())   # 1. 全部 vLLM 配置
5
factors.append(code_hash)                    # 2. 被 Dynamo 追踪到的所有源文件内容
6
factors.append(compiler_hash)                # 3. Inductor/系统状态
7
cache_dir = os.path.join(envs.VLLM_CACHE_ROOT, "torch_compile_cache",
8
                         hashlib.md5(str(factors).encode()).hexdigest()[:10])

第 2 项的做法很讲究：vLLM 不只 hash 顶层模型文件，而是通过劫持 InliningInstructionTranslator.inline_call 把 Dynamo 内联过的每一个函数的源文件都记进 traced_files，再把文件内容全部读出来做 hash。为什么这么细？因为 Inductor 编译的产物融合了所有内联函数的代码，任何一个被内联的文件变了，产物都可能不同。漏掉一个，缓存就可能在代码更新后仍然命中旧的二进制——这是典型的「缓存失效域不完整」bug，vLLM 用把追踪面完整记录下来的方式把它堵住。

模型侧的接入点：support_torch_compile 装饰器#

编译的起点不在 runner，而在模型类上的一行装饰器：

1
@support_torch_compile
2
class LlamaModel(nn.Module):
3
    def forward(
4
        self,
5
        input_ids: torch.Tensor,          # 一维 token id 流
6
        positions: torch.Tensor,
7
        intermediate_tensors: Optional[IntermediateTensors] = None,
8
        inputs_embeds: Optional[torch.Tensor] = None,
9
    ) -> torch.Tensor:
10
        ...

support_torch_compile（vllm/compilation/decorators.py）做三件事。第一，把 TorchCompileWrapperWithCustomDispatcher 动态插进类的 MRO，让每个被装饰的模型实例获得编译能力。第二，从 forward 的类型注解推断哪些参数的哪些维度是动态的：注解为 torch.Tensor 或 IntermediateTensors 的参数，第 0 维默认标记为动态（token 数）。第三，重写 __call__，在首次调用时用 torch._dynamo.mark_dynamic 显式标注这些维度，然后触发一次也是唯一一次 Dynamo 编译。

动态维推断有个容易踩的细节：如果某个参数某次运行时是 None（比如非多模态模型的 inputs_embeds），那么它的整个生命周期都必须保持 None。注释原文写着「if an argument is None, it should always be passed as None during the lifetime of the model, otherwise, it cannot be captured as a single computation graph」——Dynamo 对「有时是张量有时是 None」的参数会产生两个分支，图分裂，自定义 dispatcher 直接失效。

自定义 dispatcher：绕过 Dynamo guard 的全部原理#

TorchCompileWrapperWithCustomDispatcher（vllm/compilation/wrapper.py）是理解「运行时为什么快」的钥匙。常规 torch.compile 的调用链是：调用 → 评估 guards → 查缓存 → 命中则执行编译代码，未命中则重新编译。guard 评估本身每次调用要耗费数十微秒，而且「未命中就重编译」对服务是灾难。

vLLM 的做法是抢在 guard 之前把编译产物截下来：

1
class TorchCompileWrapperWithCustomDispatcher:
2
    def __init__(self, compiled_callable=None, compilation_level=0):
3
        ...
4
        self.original_code_object = self.__class__.forward.__code__
5
        self.compiled_codes: list[CodeType] = []
6
        torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)
7
        ...
8

9
    def bytecode_hook(self, old_code: CodeType, new_code: CodeType):
10
        # Hook to save the compiled bytecode for direct execution.
11
        if old_code is not self.original_code_object:
12
            return
13
        ...
14
        self.compiled_codes.append(new_code)
15

16
    @contextmanager
17
    def dispatch_to_code(self, index: int):
18
        self.__class__.forward.__code__ = self.compiled_codes[index]
19
        yield
20
        self.__class__.forward.__code__ = self.original_code_object

bytecode_hook 是 Dynamo 编译完成时的回调——Dynamo 把 Python 字节码重写成优化后的字节码，再把新 code object 装回执行栈。vLLM 在这个钩子里把新 code object 存进 compiled_codes 列表。之后每次调用 forward，走的是装饰器重写的 __call__：

1
# vllm/compilation/decorators.py (节选)
2
# 捕获完成后直接切换 code object,不再走 Dynamo guard
3
with self.dispatch_to_code(0):
4
    model_output = self.forward(*args, **kwargs)
5
    return model_output

dispatch_to_code 是一个极其暴力的技巧：直接把类上 forward 函数的 __code__ 换成 Dynamo 产出的字节码，调用完再换回来。它成立的前提是 Dynamo 的一个保证——编译前后的 code object 具有完全相同的参数、cell 变量与自由变量布局。源码注释里给了 PyTorch 开发者论坛的讨论链接来佐证这个保证。这样每次 decode step 完全跳过 guard 评估，Python 调用栈里只剩「切换 code → 执行编译后字节码 → 切换回来」三步。

为什么 vLLM 能这么大胆？因为「只追一次」的前提由 vLLM 自己保证：输入结构固定（id、position、静态 buffer）、动态维固定为 token 数、注意力是黑盒算子。Dynamo 的 guard 所要防御的「输入结构变化」，在 vLLM 的调用约定下根本不会发生。guard 保护的东西被调用约定消灭了，guard 本身自然可以删掉。

split_graph 的切图算法#

Dynamo 产出图后，VllmBackend.__call__ 调用 split_graph 按 splitting_ops 切分：

1
def split_graph(graph: fx.GraphModule, ops: list[str]):
2
    subgraph_id = 0
3
    node_to_subgraph_id = {}
4
    split_op_graphs = []
5
    for node in graph.graph.nodes:
6
        if node.op in ("output", "placeholder"):
7
            continue
8
        if node.op == "call_function" and str(node.target) in ops:
9
            # 命中切分算子: 结束当前子图, 该算子独占一个新子图,
10
            # 然后再开一个新子图
11
            subgraph_id += 1
12
            node_to_subgraph_id[node] = subgraph_id
13
            split_op_graphs.append(subgraph_id)
14
            subgraph_id += 1
15
        else:
16
            node_to_subgraph_id[node] = subgraph_id
17

18
    split_gm = torch.fx.passes.split_module.split_module(
19
        graph, None, lambda node: node_to_subgraph_id[node],
20
        keep_original_order=True)
21
    ...

算法是单遍扫描：给每个节点编号其所属子图；遇到切分算子就先 +1（算子自己独占一个子图）、再 +1（切分算子后面的节点进入更新后的子图）。最关键的一行是 keep_original_order=True——注释强调「否则 pytorch 可能重排节点，图里有 mutation 时语义会变」。因为 vLLM 的图里有大量 in-place 写（KV cache 写入、out 参数），fx 的默认节点排序优化会破坏张量的写读次序，这个 flag 是用 bug 换来的经验。

切完之后，80 层的 Llama 会得到 161 个子模块：81 个编译子图（embedding + 79 个层间段 + 最终段，中间段的 Inductor 代码完全相同会命中同一份缓存）、80 个注意力单节点子图。官方设计文档记录的实际数字是 17 份不同的 Inductor handle——首段、中间段共享、尾段各一份，因为 Inductor 按「图内容」缓存而不是按「模块编号」缓存。

PiecewiseCompileInterpreter：边解释边编译#

把切好的图交回 Dynamo 之前，PiecewiseCompileInterpreter 先用 fake tensor 把图「空跑」一遍：

1
class PiecewiseCompileInterpreter(torch.fx.Interpreter):
2
    def run(self, *args):
3
        fake_args = [
4
            self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
5
            for t in args
6
        ]
7
        with self.fake_mode, enable_python_dispatcher():
8
            return super().run(*fake_args)
9

10
    def call_module(self, target, args, kwargs):
11
        output = super().call_module(target, args, kwargs)
12
        if target in self.compile_submod_names:
13
            index = self.compile_submod_names.index(target)
14
            submod = self.fetch_attr(target)
15
            sym_shape_indices = [
16
                i for i, x in enumerate(args) if isinstance(x, torch.SymInt)
17
            ]
18
            # 告诉 Inductor 为「动态 shape」编译一份
19
            compiled_graph_for_dynamic_shape = ...compiler_manager.compile(
20
                submod, args, ..., runtime_shape=None)
21
            # 为这个子图包上 PiecewiseBackend + CUDAGraphWrapper
22
            piecewise_backend = PiecewiseBackend(
23
                submod, self.vllm_config, index, ...)
24
            static_graph_wrapper_class = resolve_obj_by_qualname(
25
                current_platform.get_static_graph_wrapper_cls())
26
            self.module.__dict__[target] = static_graph_wrapper_class(
27
                runnable=piecewise_backend, ...)
28
            ...

fake mode 让整张图在不分配一行真实显存、不执行任何真实 kernel 的情况下走完，同时记录每个中间张量的符号 shape（SymInt）。sym_shape_indices 记下哪些位置的 shape 是符号——这就是后续 PiecewiseBackend 在运行时判断 batch size 的依据（runtime_shape = args[self.sym_shape_indices[0]]）。

这里还藏着一个内存管理细节：所有被包了 CUDA Graph 的子图共享同一个 graph_pool（current_platform.get_global_graph_pool()）。CUDA Graph 的显存池机制允许多个图共享中间 buffer——81 个子图如果各自持池，显存占用会爆掉；共享一个池，后捕获的图会复用先捕获的图释放出来的空间。源码注释特意提醒只在用多流时才需要重新审视这里——单流是共享全局池的前提，将来若引入多流执行，这里要重新设计。

InductorAdaptor 的「驯化」工作#

vLLM 对 Inductor 的使用远不是调一次 compile_fx 那么简单。InductorAdaptor.compile（vllm/compilation/compiler_interface.py）里塞满了 monkey patch，这层适配要解决三个只有 vLLM 才会遇到的问题。

问题一：Inductor 的缓存默认「不信任」vLLM 的调用方式。 标准 torch.compile 流程里，Inductor 的编译发生在 Dynamo 的 tracing context 之内，Dynamo 会提供 shape environment（每个符号变量的取值范围、guard 条件）。而 vLLM 的 piecewise 编译在 Dynamo 完成之后才运行，此时没有 shape env，Inductor 的代码缓存查找会直接失败。vLLM 的对策是注入一个 AlwaysHitShapeEnv——它的全部方法都是空实现（evaluate_guards_expression 恒返回 True、get_pruned_guards 返回空列表），注释坦言这些方法「obtained by trial-and-error until it works」。工程上这样做是安全的：vLLM 自己管理 shape 一致性，Inductor 缓存键里的 guard 检查是多余的防御。

问题二：需要拿到编译产物的 hash 与文件路径来做二级缓存。 Inductor 内部会为编译产物计算 compiled_fx_graph_hash，vLLM 把这个函数整个劫持掉，在它返回时顺手记下 hash；再劫持 compile_fx_inner 从编译后的 callable 里挖出生成的 Python 文件路径。这一对（hash_str, file_path）被当成 handle 存进 vLLM 自己的 CompilerManager.cache，下次启动时先查 vLLM 的缓存文件，命中就直接从 Inductor 的 FxGraphCache 里 load，跳过整条编译流水线。

问题三：高阶算子与 Dynamo 之外的图默认禁止缓存。 Inductor 有一个 _check_can_cache 守卫：不在 Dynamo tracing context 下的图、含高阶算子的图都会拒绝缓存。vLLM 两种都占了，于是直接 patch 成恒返回 None（即「可以缓存」）。源码注释点明了理由：对 vLLM 而言这两种情况都想缓存，见 pytorch/pytorch codecache.py 的相关代码。

这层适配是「依赖内部 API」的典型样本——vLLM 的注释也承认这些 patch 与 PyTorch 版本强耦合（代码里可见 2.5/2.6/2.8 三个版本的分支处理），PyTorch 一动内部接口就断。2025 年 PyTorch 2.8 提供 standalone_compile 之后，vLLM 加了 InductorStandaloneAdaptor（用环境变量 VLLM_USE_STANDALONE_COMPILE 开启）来逐步替换这套劫持,官方注释「we’re going to replace this all with standalone_compile sometime」正是这种迁移的注脚。

set_inductor_config：动态 shape 与静态 shape 的两种编译策略#

1
def set_inductor_config(config, runtime_shape):
2
    if isinstance(runtime_shape, int):
3
        # for a specific batchsize, tuning triton kernel parameters
4
        # can be beneficial
5
        config["max_autotune"] = True
6
        config["coordinate_descent_tuning"] = True

同一个子图在 vLLM 里最多编译两套：runtime_shape=None 的动态 shape 版（token 数是符号变量）与按 compile_sizes 里每个具体尺寸编的静态版。只有静态版才开 max_autotune 与 coordinate_descent_tuning——autotune 要对每个候选 config 实测吞吐，动态 shape 下 Triton kernel 的 grid 依赖 batch，实测结果对其他 batch 不成立，白测。静态 shape 下实测出的最优 config 才可靠。

官方设计文档给出的 autotune 实例能说明这类收益：一个 8 x 2048, 2048 x 3072 的矩阵乘，Inductor 列出十余个候选 Triton config，最优者 0.0130 毫秒，比默认走 cuBLAS 的 0.0160 毫秒快约 19%。对小 batch 的小 GEMM，这类微调乘上 kernel 数量，就是「编译几分钟换 decode 提速」的来源。autotune 极费时间（文档原话「seconds to minutes」，取决于模型与 batch size），所以 vLLM 只对 compile_sizes 指定的少数档位开启，且结果进缓存。

vLLM 的自定义 Inductor pass#

PostGradPassManager 把一列自定义 pass 挂进 Inductor 的 post_grad_custom_post_pass 钩子，在 Inductor 完成自己的简化后、codegen 之前运行。默认链上有四类：

NoOpEliminationPass——删掉冗余的 reshape/slice。这个 pass 的存在理由非常具体：它是 RMSNorm-FP8 量化融合的前置条件。apply_fp8_linear 在 2D 情况下会多插一个 reshape，fused_add_rms_norm 后紧跟的 slice_scatter 在满宽时是纯 no-op，这些噪声节点让 pattern matcher 认不出「rms_norm → quant」的固定模式。NoOpEliminationPass 处理三类情形：reshape 链折叠成对 base tensor 的一次 reshape、产生原 shape 的 reshape 删除、产生原 shape 的 slice 删除。判等逻辑里专门处理了三个等价情形：两个维度都是整数值、reshape 维是 -1（推断维）、两个维是同一个 SymInt——最后一条在动态 shape 编译下必不可少，因为「batch 维」在图里是符号 s0，直接 == 比较恒为假。

FusionPass——用 Inductor 的 PatternMatcherPass 把「RMSNorm + FP8 静态量化」识别出来，替换成 vLLM 手写的融合算子 rms_norm_static_fp8_quant / fused_add_rms_norm_static_fp8_quant。这类融合的收益是双重的：norm 之后本来要写一次激活结果再到显存、量化时再读一次，融合后中间量留在寄存器/共享内存，省一次读写循环还少一个 kernel。pass_config 里 enable_noop 与 enable_fusion 是联动的——fusion 依赖 noop 先清场，代码里如果 fusion 开而 noop 关，会打警告说融合「might not work」。

AtFusionPass / AllReduceFusionPass / AsyncTPPass——注意力+量化融合、flashinfer 通信融合与异步张量并行，属于分布式/量化场景的可选 pass，默认关闭。

ActivationQuantFusionPass——激活量化的融合变体。

这一层是 vLLM「用 Inductor 又不完全受制于 Inductor」的体现：Inductor 的通用融合规则不认识 vLLM 的量化算子，vLLM 就把自己的领域知识写成 pattern 挂进去，让通用编译器产出领域特化的 kernel。

fix_functionalization：让 in-place 语义在图里活下来#

还有一块必须提到的拼图是 fix_functionalization.py。PyTorch 2.6+ 默认启用 auto-functionalization V2，它会把 auto_functionalized 高阶算子的调用重写成「复制输入 → 非原地计算 → 复制回输出」的函数式形式——这个默认行为与 vLLM 的 out 参数风格自定义算子严重冲突：KV cache 的 in-place 写、注意力 out 参数的写回，在函数化之后全变成低效的显存搬运，vLLM 论坛 RFC #14703 专门记录了这个问题。vLLM 在 CompilationConfig.__post_init__ 里对 torch 2.6+ 显式设置 enable_auto_functionalized_v2 = False，并保留 fix_functionalization pass 处理 V1 语义下的图改写。这是一个很好的例子：框架级的「正确默认」，放到推理服务的上下文里就变成了性能杀手——训练框架希望图是纯函数式的（便于求导），推理框架希望图是尽量少搬运的（in-place 是显存带宽的朋友），同一个 PyTorch 无法用一组默认值同时伺候两者。

两层封装：为什么是 PiecewiseBackend 套着 CUDAGraphWrapper#

编译完成后，每个非注意力子图的实际可执行体是一个三层洋葱：

1
CUDAGraphWrapper (capture/replay 调度, 平台无关)
2
  └── PiecewiseBackend (静态 shape 编译分派: 动态版 vs 各档位静态版)
3
        └── Inductor 编译产物 (Triton kernel 序列, 一般 shape 或具体 shape)

PiecewiseBackend 管「这个 batch size 该跑哪份编译产物」。它的两个关键数据结构：

1
class PiecewiseBackend:
2
    def __init__(self, graph, vllm_config, piecewise_compile_index,
3
                 total_piecewise_compiles, sym_shape_indices,
4
                 compiled_graph_for_general_shape, vllm_backend):
5
        ...
6
        # Inductor 按 compile_sizes 额外编译的静态 shape 档位
7
        self.compile_sizes: set[int] = set(self.compilation_config.compile_sizes)
8
        self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}
9
        self.compiled_graph_for_general_shape = compiled_graph_for_general_shape
10
        self.sym_shape_indices = sym_shape_indices
11

12
    def __call__(self, *args) -> Any:
13
        if not self.first_run_finished:
14
            # 第一次运行: 只跑「动态 shape」编译产物
15
            self.first_run_finished = True
16
            self.check_for_ending_compilation()
17
            return self.compiled_graph_for_general_shape(*args)
18

19
        runtime_shape = args[self.sym_shape_indices[0]]
20

21
        if runtime_shape not in self.concrete_size_entries:
22
            # 没为这个 size 编译过静态版, 继续用动态版
23
            return self.compiled_graph_for_general_shape(*args)
24

25
        entry = self.concrete_size_entries[runtime_shape]
26
        if not entry.compiled:
27
            # 首次命中某个静态档位: 按这个 shape 现场编译(有缓存)
28
            entry.compiled = True
29
            self.to_be_compiled_sizes.remove(runtime_shape)
30
            entry.runnable = self.vllm_backend.compiler_manager.compile(
31
                self.graph, args, ..., runtime_shape=runtime_shape)
32
            ...
33
        return entry.runnable(*args)

注意首次调用的语义：运行时遇到新档位不会立刻编译静态版，而是先保证首跑走动态版。因为编译发生在「服务已就绪」的路径上，vLLM 要确保任何代码路径都不会在请求处理中途引入长阻塞——动态版兜底，静态版后台逐步填充。to_be_compiled_sizes 清空的时机被用来判断「整个模型的编译是否全部结束」（check_for_ending_compilation 只在最后一个子图时触发收尾）。

CUDAGraphWrapper 则管「什么时候捕获、什么时候回放」。它是平台无关的抽象：构造函数里 resolve_obj_by_qualname(current_platform.get_static_graph_wrapper_cls()) 拿到平台专属实现，NVIDIA 卡是真 CUDA Graph，ROCm 卡是对应的 HIP Graph 机制，接口统一。

BatchDescriptor 与 forward context：回放 key 的设计#

CUDA Graph 回放要求「完全相同的输入形状」。vLLM 需要 scheme 回答：两个 batch 什么情况下可以共用一张已捕获的图？答案就是 BatchDescriptor（vllm/forward_context.py）：

1
class BatchDescriptor(NamedTuple):
2
    num_tokens: int
3
    uniform_decode: bool = False
4
    # False 也可用于一个均匀 decode batch, 让它复用支持非均匀 batch 的图
5

6
    @property
7
    def non_uniform(self) -> "BatchDescriptor":
8
        return BatchDescriptor(self.num_tokens, uniform_decode=False)

只有两个字段。num_tokens 是 pad 之后的 token 总数（决定所有算子的 shape），uniform_decode 标记「是否所有请求都是等长 decode」（决定注意力后端走哪条 kernel 路径）。注释里专门强调「We should keep the num of items as minimal as possible」——key 越少，能共用的场景越多：一个均匀 decode batch 如果在档位表里查不到对应的 uniform 图，可以用 non_uniform 退化成非均匀 key 去匹配一张更「通用」的图,宁可走一条稍慢的 kernel 路径也不放弃回放机会。

每 step 的回放决策由 CudagraphDispatcher.dispatch 做出：

1
# vllm/v1/cudagraph_dispatcher.py (节选)
2
def dispatch(self, batch_descriptor):
3
    if not self.keys_initialized:
4
        return CUDAGraphMode.NONE, None
5
    # 1. 精确匹配 FULL 档位
6
    if batch_descriptor in self.cudagraph_keys[CUDAGraphMode.FULL]:
7
        return CUDAGraphMode.FULL, batch_descriptor
8
    # 2. 退化成非均匀 key 再试 FULL
9
    non_uniform_key = batch_descriptor.non_uniform
10
    if non_uniform_key in self.cudagraph_keys[CUDAGraphMode.FULL]:
11
        return CUDAGraphMode.FULL, non_uniform_key
12
    # 3. 非均匀 key 再试 PIECEWISE(最通用)
13
    if non_uniform_key in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:
14
        return CUDAGraphMode.PIECEWISE, non_uniform_key
15
    # 4. 实在不行, eager
16
    return CUDAGraphMode.NONE, None

dispatcher 持有两套 key 集合（FULL 与 PIECEWISE 各一组）。选中的 mode 和 key 通过 set_forward_context 写进 forward context,所有 CUDAGraphWrapper 在被调用时从 context 里无条件信任这个决策——源码注释「blindly trust them」——包装器自己不做判断,只对比「context 里的 mode 与包装器自己的 mode 是否一致」,一致就捕获/回放,不一致就直通下层 runnable。这让 FULL 与 PIECEWISE 两类 wrapper 可以共存于同一条调用链 wrapper 只在 decode batch 时被激活,PIECEWISE wrapper 只在混批时被激活。

运行时包装：copy inputs 到静态缓冲区#

CUDA Graph 回放的硬约束是输入地址固定——回放时 kernel 从 capture 时记录的指针读数。vLLM 的 token 流每个 step 都在变（调度器写进持久化 buffer），所以捕获前必须把输入拷进静态缓冲区。这份拷贝逻辑在 VllmBackend.__call__ 的尾部：

1
# vllm/compilation/backends.py (节选)
2
# 只对 shape 是符号变量的输入(即 token 流)做拷贝;
3
# 权重和静态 buffer 的 shape 是具体的, 地址固定, 无需拷贝
4
self.sym_tensor_indices = [
5
    i for i, x in enumerate(fake_args)
6
    if isinstance(x, torch._subclasses.fake_tensor.FakeTensor)
7
    and any(is_symbolic(d) for d in x.size())
8
]
9
self.input_buffers = [example_inputs[x].clone() for x in self.sym_tensor_indices]
10

11
def copy_and_call(*args):
12
    list_args = list(args)
13
    for i, index in enumerate(self.sym_tensor_indices):
14
        runtime_tensor = list_args[index]
15
        runtime_shape = runtime_tensor.shape[0]
16
        static_tensor = self.input_buffers[i][:runtime_shape]
17
        static_tensor.copy_(runtime_tensor)
18
        list_args[index] = static_tensor
19
    return self.split_gm(*list_args)

compiler managed cudagraph input buffers 的注释说明了假设:「we assume the first run with symbolic shapes has the maximum size among all the tensors」——静态缓冲区按最大 shape 分配一次,之后每个 step 把真实输入 copy_ 到前缀位置。这个拷贝不是白给的：一个 token 流的 input_ids + positions 拷贝是两次小 memcpy，比起省下的数百次 launch,依然是大赚。

注意这段代码只在 cudagraph_copy_inputs=True 时存在。注释还注明「if the caller can guarantee that the same input buffers are always used, it can set this to False」——V1 runner 使用自己的持久化 buffer（self.input_ids、self.positions,整个生命周期地址固定）,所以主路径不需要这个编译侧拷贝。这个 flag 是给「buffer 地址无法保证」的调用路径（如某些测试或非 V1 入口）准备的兜底。

捕获一次的成本与 weak ref 技巧#

CUDAGraphWrapper.__call__ 的捕获分支有几个细节值得逐行看：

1
# vllm/compilation/cuda_graph.py (节选)
2
with ExitStack() as stack:
3
    if self.cudagraph_options.gc_disable:
4
        # piecewise 模式下一次 forward 要捕获几十个子图,
5
        # 每次都跑 gc 会把捕获拖得非常慢,
6
        # 因此只对第一个子图跑 gc, 其余全部禁掉
7
        stack.enter_context(patch("gc.collect", lambda: None))
8
        stack.enter_context(patch("torch.cuda.empty_cache", lambda: None))
9

10
    with torch.cuda.graph(cudagraph, pool=self.graph_pool):
11
        output = self.runnable(*args, **kwargs)
12
        if self.cudagraph_options.weak_ref_output:
13
            # 转 weak ref 后原 output 立即释放以省显存;
14
            # 只有最后一个子图能这么做,
15
            # 因为中间子图的输出还会被其他图引用
16
            output = weak_ref_tensors(output)
17

18
entry.output = weak_ref_tensors(output)

三个观察：其一，gc 抑制的粒度按「first graph / rest graphs」区分（CUDAGraphOptions 里 gc_disable=not is_first_graph），原因写在注释里——一次 forward 捕获几十个子图，逐个 gc 不可接受。其二，捕获时产出的 output 用 weak_ref_tensors 转成弱引用，让原始 tensor 立即归还显存池；注释专门解释为什么只有最后一个子图可以做这件事——中间子图的输出还要被下一个子图引用,提前释放会把正在使用的数据卖掉。其三,回放路径上 debug 分支会校验输入地址与捕获时一致（VLLM_LOGGING_LEVEL=DEBUG 才启用）,这是把「CUDA Graph 的地址不变式」变成可观测断言的调试手段。

还有一个容易被误解的设计：wrap 器本身不持有静态缓冲区、不拷贝输入。「Note: CUDAGraphWrapper does not store persistent buffers or copy any runtime inputs into that buffers for replay. We assume implementing them is done outside of the wrapper.」拷贝策略留在编译侧的 copy_and_call 或 runner 的持久 buffer 里,wrapper 只做「捕获/回放/校验」——这让 wrapper 对任意动态 shape 策略保持正交（源码用的词是「staying orthogonal to compilation logic」）。

capture_model：runner 侧的编排#

wrapper 是「被动」的——被调到才捕获。主动的编排发生在 GPUModelRunner.capture_model()：

1
# vllm/v1/worker/gpu_model_runner.py (V0.10.1, 节选)
2
def capture_model(self) -> None:
3
    if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:
4
        logger.warning("Skipping CUDA graph capture...")
5
        return
6
    ...
7
    start_free_gpu_memory = torch.cuda.mem_get_info()[0]
8

9
    with freeze_gc(), graph_capture(device=self.device):
10
        cudagraph_mode = self.compilation_config.cudagraph_mode
11
        if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:
12
            # 混合 prefill-decode 批: 按 capture 档位从大到小逐个捕获
13
            compilation_cases = list(reversed(self.cudagraph_batch_sizes))
14
            self._capture_cudagraphs(
15
                compilation_cases,
16
                cudagraph_runtime_mode=cudagraph_mode.mixed_mode(),
17
                uniform_decode=False)
18

19
        # decode 例行批单独按 FULL 捕获(如果启用了分离路由)
20
        if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and ...
21
            ...
22
    set_cudagraph_capturing_enabled(False)
23
    ...
24
    cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
25
    # This usually takes 5~20 seconds.
26
    logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
27
                elapsed_time, cuda_graph_size / (1 << 30))

四个值得注意的点。

从大到小捕获。compilation_cases = list(reversed(self.cudagraph_batch_sizes)) 一行加注释「Capture the large shapes first so that the smaller shapes can reuse the memory pool allocated for the large shapes」。CUDA Graph 的共享显存池只在「后捕获的图比先捕获的小」时才能完全复用——小图的中间 buffer 是大图 buffer 的前缀视角。反过来先小后大，每个大图都得新开空间，池化收益归零。这也是 CUDA Graph 工程应用的经典手法。

预热先行。_capture_cudagraphs 里对每个档位先跑 cudagraph_num_of_warmups 次 dummy run（V1 默认为 1 次）再真正捕获。预热有两个作用：触发 Triton JIT 编译、让 CUDA context 的懒初始化全部完成——否则这些一次性开销会被录进图里。预热用的 cudagraph_runtime_mode=CUDAGraphMode.NONE,即「跑热但不录」。

dummy run 的输入也要仿真。捕获用的是 _dummy_run(num_tokens, ...),它构造假的 input_ids、positions 与注意力元数据。这里有个注释明确指出的细节:「When setting max_query_len = 1, we switch to and capture the optimized routine of FA2 for pure decode」——decode 批的 dummy 输入故意把 query 长度设为 1,让 FA2 在捕获时就走 FlashDecode 路径,CUDA Graph 里固化下来的就是这条优化路径而不是通用 prefill 路径。

捕获完成即上锁。set_cudagraph_capturing_enabled(False) 关掉全局捕获开关,此后任何意外的捕获尝试都会被检测并报错(而非静默走 eager)。捕获总耗时与占用显存都被记录——源码注释「This usually takes 5~20 seconds」——这笔一次性成本就是 V1 引擎启动比 eager 慢几十秒的主要来源。

显存账：CUDA Graph 的池化与每档成本#

CUDA Graph 的主要显存开销有两块：每张图录制的 kernel 参数与中间 buffer。中间 buffer 大头靠 graph pool 共享摊薄,但每个 batch 档位仍要一份专属的输入/输出静态区。V0.10.1 的 capture_model 会打印一笔账：

1
Graph capturing finished in 12 secs, took 2.17 GiB

档位多、模型层深时,这 2 GiB 往往就是几十上百兆 × 档位数的累积。两个常见的调优动作因此而来：一是砍档位——--compilation-config '{"cudagraph_capture_sizes": [1, 2, 4, 8]}' 只保留最常用的几档,牺牲命中率换显存;二是 FULL_DECODE_ONLY——decode 档位用整图、prefill 混批完全走带编译的 eager,只为解码一侧付捕获成本（CompilationConfig.cudagraph_mode 的 docstring 明说这个模式适合 P/D 分离架构里的 decode 实例,prefill 地位不重要,省下的显存可以换成 KV cache)。

另一个隐性成本是捕获期间的临时峰值：dummy run 本身会触发中间张量分配，与模型加载后的常驻显存叠加，显存水位在捕获阶段最容易 OOM。vLLM 用「先捕获最大档位」减轻这个问题（最大档位一次性吃掉峰值，小档位复用空间），但启动脚本的 显存余量仍要为最大档位的 dummy run 留出。

哪些代码写在 CUDA Graph 下会出事#

TorchCompileWrapperWithCustomDispatcher.bytecode_hook 有一段狠代码：

1
if self.vllm_config.compilation_config.cudagraph_mode != \
2
    CUDAGraphMode.NONE and "update" in new_code.co_names:
3
    import depyf
4
    src = depyf.decompile(new_code)
5
    msg = "Assigning / modifying buffers of nn.Module during forward pass
6
           is not allowed when using cudagraph inside the compiler because
7
           it will cause silent errors. ..."
8
    raise RuntimeError(msg)

编译后的字节码里只要出现名为 update 的调用（buffer 改写),直接抛错并反编译出源码方便定位。为什么这么凶?CUDA Graph 回放期间 CPU 不参与执行,如果图里录了「写某个 buffer」的操作而 buffer 地址在两次回放之间被 PyTorch allocator 复用/移动,错误的结果会被静默接受——不是崩溃,是算错。这个检查把「张量地址不变」从约定升级成了可检测的失败。

同文件里 use_custom_dispatcher 与 guard 的互斥、support_torch_compile 对 None 参数的终生一致性要求,前面已提。还有三个实际使用中常见的坑:

1. graph break 让自定义 dispatcher 静默失效。模型 forward 里有 Dynamo 无法追踪的语句（打印日志、断言张量值、依赖 .item() 的分支）会让 Dynamo graph break,产出多段编译代码。level 3 的 dispatcher 只在「第一段」上工作,后续段落直接走 Dynamo 默认派发,guard 开销回归——性能不报错地劣化。排查方法是把 VLLM_LOGGING_LEVEL=DEBUG 打开,数 Dynamo 的重编译次数（源码: Dynamo bytecode transform time 只应出现一次）。
2. cache 命中失败的典型来源。cache dir hash 包含「被追踪文件的内容」,装了 torch 版本不同、改过模型代码、甚至改了会被内联的工具函数,hash 就变,整个编译+捕获流程重来。VLLM_DISABLE_COMPILE_CACHE=1 可以关缓存排查,但注意它同时关掉了 Inductor 的 FX graph cache（源码里 VLLM_DISABLE_COMPILE_CACHE 是「总开关」）。
3. 编译时间预算。capture_model 的 5-20 秒加上 Dynamo+Inductor 编译的分钟级成本,让「频繁重启的服务」必须依赖编译缓存目录。官方文档建议部署时把 ~/.cache/vllm/torch_compile_cache 目录直接拷到新机器复用——hash 因子设计保证了只要配置与代码不变,缓存命中。

调试与观测： 两个官方入口#

depyf 图导出。VllmBackend.__call__ 里写入 computation_graph.py 的代码改自 depyf 项目（源码注释给了行级链接）,把切分后的 fx 图完整落盘。加上 --debug-dump-path 或读 ~/.cache/vllm/torch_compile_cache/<hash>/rank_0_<dp>/,能看到两份关键产物:computation_graph.py（切分后的图,直观看到子图边界与注意力单节点）与 transformed_code.py(Dynamo 改写的 Python 字节码反编译)。官方文档的读图方法:先看子模数数量是否符合层数预期,再确认每个 submod 是被 CUDAGraphWrapper 包住还是直通。

compilation_counter。vllm/compilation/counter.py 的全局计数器暴露了编译过程的每一步: num_graphs_seen(Dynamo 产图数,应=1）、num_piecewise_graphs_seen(切图数,应≈2×层数+1)、num_inductor_compiles(Inductor 编译次数,应≈子图内容种类数)、num_cudagraph_captured(捕获图数,应≈档位数×子图数)。测试文件里大量用这些计数器做断言——比起 benchmark,计数器更能精确定位「哪一层的图没被编译」「哪个图没被捕获」。