LMDeploy PyTorch 后端 CUDA Graph Runner 设计解析：三段式分发、分段捕获（PCG）与失效重置机制

• 人工智能
• 大模型
• 模型推理服务
• 推理引擎
• 本地部署
• 模型量化

本篇技术指南以 LMDeploy 仓库中 graph_runner 设计文档 为核心骨架，结合 runner.py、full_graph.py、piecewise.py、standard.py 等源码实现，完整解析 PyTorch CUDA 后端中 CUDA Graph 的选择、捕获、回放与失效四大生命周期。读完本文，你将掌握 LMDeploy 解码（decode）全图路径、分段 CUDA 图（Piecewise CUDA Graph，PCG）预填充路径与 eager 回退路径的三段式调度模型，理解 eager 边界（eager boundary）输出适配、bridge 存储复用、标准解码器描述符（descriptor）等核心设计，并能在源码层面定位每条设计决策的落点。

概览：一条 forward 的三段式出路

CUDA Graph 是减少 kernel 启动开销、提升 LLM 推理吞吐的关键技术。LMDeploy 的 PyTorch CUDA 后端在CUDAGraphRunner中将其抽象为三个并行的执行路径，设计文档开篇给出了清晰的拓扑：

CUDAGraphRunner | +-- supported decode --------> full CUDA graph +-- prepared PCG prefill ----> piecewise CUDA graph plan +-- everything else ---------> eager model forward

• supported decode（解码）：走传统的完整 CUDA 图路径，由CUDASingleGraphRunner负责一整张模型图的捕获与回放；
• prepared PCG prefill（预填充）：走分段 CUDA 图路径，由PiecewiseGraphManager管理的分段执行计划（plan）驱动；
• everything else（其余情况）：包括不支持的请求、未准备（unprepared）的请求、启动热身被禁用等场景，一律回退到eager 模型 forward。

三条路径共享外层分发器与模型输出契约，但刻意不共用同一个"重模式"执行器——它们的 graph key、缓冲、捕获生命周期与失败规则各不相同（README.md）。

在源码层面，这一分发逻辑位于 runner.py 的CUDAGraphRunner.__call__：

if self._should_use_full_graph(step_context, kwargs): return self._forward_full_graph(**kwargs) descriptor = self._get_piecewise_graph_descriptor(step_context, kwargs) if descriptor is not None: manager = self._piecewise_graph_manager if manager.has_plan(descriptor): return manager.replay(descriptor, kwargs) if is_preparing_prefill(): return manager.prepare(descriptor, kwargs) # Serving never captures... return self._forward_eager(**kwargs)

对应设计文档中的三层判定顺序：① 支持的 decode 走完整 CUDA 图；② 支持的 prefill 推导分段描述符，已有计划则回放、启动期 dummy 请求则准备计划；③ 不支持或未准备的请求直接 eager 执行。

文件与职责地图

设计文档用一张表格界定了各文件的所有权，结合源码可将其职责进一步落实：

一个重要的边界约定是：模型与后端无关的nn模块禁止 import 本包。模型通过继承PiecewiseCudaGraphMixin选择共享 prefill 运行时（cudagraph.py）；选定 CUDA 算子实现则负责自身语义所需的任何 eager 边界。

外层分发与整体重置

分发判定的细节

设计文档强调"serving 从不捕获分段计划"：如果启动热身被禁用，或者没有匹配的计划，则在 PCG 有机会改写 KV 或状态缓存之前，就先行选择 eager 执行。这一安全性约束在 runner.py__call__中有明确注释：

# Serving never captures. If graph-runner warmup was skipped or this call is # unsupported, eager execution is selected before PCG can mutate state. return self._forward_eager(**kwargs)

失效的所有者：reset()

设计文档规定CUDAGraphRunner.reset()是失效（invalidation）的唯一所有者：它丢弃每一张全图、每一个分段计划、共享的分段 bridge 存储，以及 DeepEP 图缓冲。任何改变地址的生命周期操作都必须通过这个完整重置，而不是逐个修补捕获的指针。

源码中 reset 的实现与此完全一致：

def reset(self): super().reset() self._full_graph_runners.clear() if self._piecewise_graph_manager is not None: self._piecewise_graph_manager.reset() _destroy_deepep_buffer()

其中_destroy_deepep_buffer()在 DeepEP 启用时会销毁进程级缓冲并执行dist.barrier()，确保重置屏障内所有 rank 同步（runner.py）。

完整 CUDA 图路径：解码的经典方案

CUDASingleGraphRunner拥有一个 graph key与一张完整模型图。设计文档给出的生命周期如下：

allocate model-owned static buffers -> fill buffers and redirect StepContext -> eager warmup -> capture the complete model call -> retain output buffers -> refill and replay for later matching decode steps

源码中 capture 方法 严格按此顺序执行：make_buffers_cudagraph分配静态输入缓冲 →_bind_inputs填缓冲并重定向 StepContext → 先跑一次 eager warmup →_capture_model用torch.cuda.graph捕获完整模型调用 → 保存输出缓冲 → 首次调用直接返回 warmup 输出。

关于"首次捕获返回 warmup 输出"这一点，设计文档给出的理由非常关键：这防止 SSM/状态型模型将捕获时的状态更新应用两次。在 runner.py_capture_full_graph中同样有对应注释：

# SSM capture warmup updates state, so the first call returns that # warmup output instead of replaying and applying the update twice. return output

外层 runner 以旧版解码 key保留这些执行器，并沿用共享图池（graph pool）策略。graph key 由(batch_size, is_decoding, enable_microbatch, query_len)加上模型的get_cudagraph_extra_key组成（runner.py），batch size 则通过_get_capture_tokens向上取到已捕获的桶大小。

解码 torch.compile：可选优化

设计文档明确指出，decode 的torch.compile是 full_graph.py 内部的可选优化：它不用于划分分段执行，也不是正确性依赖。源码中build_decode_model_forward仅在enable_decode_torch_compile为真且 PyTorch >= 2.8 时才启用torch.compile，并会调整 Dynamo 的recompile_limit、加入triton等 trace 规则，同时强制triton.cudagraphs=False（full_graph.py）。

分段 CUDA 图路径：预填充的流水线式捕获

与 FX/Dynamo 无关的直接追踪

PCG（Piecewise CUDA Graph）直接对普通 forward 使用 CUDA stream capture 进行追踪，不构建 FX/Dynamo 图。一个 eager 边界结束当前捕获、以 eager 方式运行原函数、绑定其结果，然后开始下一次捕获：

GraphStep -> EagerStep -> GraphStep -> ...

设计文档强调，刻意只存在两种步骤类型：

• GraphStep：拥有一张可回放的torch.cuda.CUDAGraph；
• EagerStep：拥有一份边界声明、一份回放实参模板和一个结果绑定。

在 piecewise.py 中，GraphStep.run用piecewise::graph:{index}的 profiler 范围包裹self.graph.replay()，EagerStep.run则先通过_resolve_eager_argument把_FrameValueRef解析为当前请求帧的值，再执行原函数并绑定输出。

启动期构造：warmup → build → 捕获 → 发布

PiecewiseGraphManager.prepare()只在启动 dummy prefill 时运行，六步流程对应 piecewise.pyprepare与 standard.py：

1. StandardDecoderPiecewiseGraphRuntime.warmup()以桶形状（bucket-shaped）eager 运行 forward，完成懒 kernel 初始化；
2. build()分配桶形状静态输入并调用trace_piecewise_cuda_graph()；
3. _CaptureBuilder捕获模型直到第一个 eager 边界；
4. 结束并回放该前缀图，使 eager 函数能消费其真实输出；
5. 运行原始 eager 函数，记录实参与结果绑定，然后恢复捕获；
6. 只有最后一段成功后才发布完整计划。

为什么构造过程中要回放捕获的前缀？设计文档解释：结束 CUDA 捕获只记录 kernel，并不执行它们。因此构造过程本身已经物化了启动请求的输出——管理器直接返回该结果，而不会立即回放新计划。这一点在 trace_piecewise_cuda_graph 中得到体现：_CaptureBuilder.build中_finish_piece_and_replay()每次capture_end()后都会执行graph.replay()。

Serving 回放：不重跑 Python forward

对已准备的描述符，回放分为四步（standard.pyreplay）：

1. 填充计划拥有的顶层图输入缓冲；
2. 绑定 live 请求帧（frame_inputs）；
3. 在引擎主 forward 流上按记录顺序执行每个 graph/eager 步骤；
4. 返回可复用输出缓冲的逻辑 raw-token 视图（由get_outputs_cudagraph按真实 token 数切片）。

模型 Pythonforward在 serving 回放期间不会被重新运行——这是 PCG 收益的根本来源。

流与 DP 相位：无需专用 PCG 流

设计文档特别说明了几点值得注意的工程细节：

• 捕获与回放都运行在引擎主 forward 流（ModelAgent.stream）上，与完整 decode 图使用的非默认流相同；管理器在prepare()时从当前流惰性解析该流，不分配专用 PCG 流。由于计划流就是调用者流，稳态路径上的跨流等待（cross-stream waits）退化为 no-op。
• CUDA 图私有池保持 plan 局部性，因为不同 token 桶可能以任意请求顺序回放。
• DP rank 处于不同本地相位时：全局 step 仍是 prefill，本地解码 rank 复用同一 token 桶计划，eager 的 attention/state/routed-expert 边界根据请求元数据选择 live 相位；投机注意力（speculative attention）会把[batch, query_len, heads, dim]结果归一化回被捕获投影使用的展平 token 形状。

外层forward_piecewise_cudagraphprofiler 范围覆盖输入绑定、完整有序计划与输出投影；嵌套的piecewise::graph:*与piecewise::eager:*范围暴露单个步骤（见 piecewise.py 中record_function('forward_piecewise_cudagraph')）。

图输入与请求帧输入：两种生命周期问题的解法

设计文档强调，graph inputs与frame_inputs解决的是不同的生命周期问题，不应混为一谈。

图输入（Graph Inputs）

图输入在回放前被复制进计划拥有的固定地址张量。当前标准运行时支持三类，token 轴均为 1（见 standard.py_GRAPH_TOKEN_INPUT_AXES）：

• input_ids，token 轴 1；
• position_ids，token 轴 1；
• 可选的mrope_position_ids，token 轴 1。

它们的物理 token 范围就是选中的 token 桶；每次回放都会先清零填充尾部，再复制逻辑请求前缀（standard.py_fill_bucket_inputs）。

frame_inputs：命名的请求帧查找表

设计文档特别澄清了一个常见误解：frame_inputs并不是"一切不进 CUDA 图的东西"的通用集合，而是当前请求的命名查找表，用于那些可能作为 eager 函数直接实参出现的值。当前标准运行时提供三类：

• past_key_values；
• attn_metadata；
• state_ids。

（见 standard.py_FRAME_INPUT_NAMES。）

构造期间，_bind_eager_argument()将直接 eager 实参与构造帧按**身份（identity）**比较：匹配的变成_FrameValueRef(name)而不是保留 dummy 请求对象；在回放某个 eager 步骤之前，_resolve_eager_argument()用 live 请求的frame_inputs[name]替换该引用（piecewise.py）。

设计文档同时给出了两条安全红线：

• 这个映射不会让任意被捕获的使用变安全：如果图内工作消费了这些值之一，其张量地址必须在计划生命周期内稳定，或者必须由复制的图输入表示并被描述符覆盖。例如，cache arena 拥有 runner 生命周期的存储所有权，尽管请求本地的 cache 元数据是 eager 的。
• 绑定刻意保持浅层：只解析直接 eager 调用实参；在算子证明每个包含值的所有权之前，不支持嵌套请求对象。

其他 eager 实参

一个 eager 实参模板还可以包含：先前 eager-only 的结果槽、稳定的图 bridge、由计划描述符表示的不可变常量、runner 拥有的模型/后端对象，以及非拥有型 CUDA 张量视图（其存储由图池、静态输入、模型/cache 分配或 bridge 维持存活）。已发布的步骤不得保留启动请求元数据，也不得意外拥有捕获的图池存储。

非拥有视图的机制在 piecewise.py_make_weak_cuda_view中实现：它通过torch._C._construct_storage_from_data_pointer（PyTorch 2.1 起可用的私有 API）构造一个只持有元数据、不拥有图池存储的弱视图。这正对应设计文档"Deliberately Unsupported"中"PCG 不支持 PyTorch 2.0"的限制来源。

eager 边界输出策略：四种适配器

@eager_boundary装饰器在非活动分段构造期间是透明的，嵌套的被装饰调用归属于最外层 eager 边界（piecewise.py）。一个边界有且仅有一种输出策略：

桥接张量可以安全地被后续 eager 步骤和后续图分段同时消费——因为两者看到的是同一个稳定 bridge。但单个边界目前不能返回混合树（部分叶子 eager-only、部分叶子进图）；设计文档要求把该操作拆成多个边界，直到真实算子证明需要算子拥有的复合适配器。

通用适配器会拒绝非 strided 张量、不支持的视图以及与边界输入发生别名（alias）的输出。保留别名的行为属于算子拥有的适配器，而不是通用追踪器——这一点在_allocate_bridge的严格校验中得到印证（piecewise.py）。

Bridge 生命周期与复用

每个图可见的 eager 输出都需要稳定地址，普通 bridge 在计划生命周期内保持保留。算子只有在确认该输出在紧随其后的 graph/eager 步骤之后不再被使用时，才可以设置reuse_bridge_after_next_step=True——此时池将该存储逻辑上开放给后续边界，但不释放物理分配，被捕获的图可以继续持有该地址。

ReusableBridgePool的实现支持这一契约（piecewise.py）：

• release(bridge)只把池内张量槽位标记为available=True，物理存储保持存活；
• allocate_padded_tensor允许用更大的 token 轴回填同一槽位（narrow）；
• 同一个池可以为互斥的 token 桶计划提供后备存储，因为模型 forward 是串行的且每个计划常驻；
• 并发计划回放不受支持——那需要独立槽位或显式租约。

标准解码器描述符：_StandardGraphDescriptor

_StandardGraphDescriptor用一个完整计划标识三个要素（standard.py）：

• 聚合 token 桶（token_bucket）；
• 图可见输入名（graph_input_names）；
• 不可变的额外 forward 实参（forward_constants）。

设计文档特别说明：精确请求数被刻意排除在描述符之外——对当前支持的 Qwen3 与 Qwen3.5 分区，所有请求分区相关的消费者都是 eager 的，而被捕获的张量代数只依赖展平 token 桶。这是这些算子分区已验证的属性，而不是通用解码器假设；未来任何被捕获的 batch 形状算子，都必须向描述符添加容量/布局事实，或将该算子移出 eager。

Token 桶策略与静态资格

初始运行时使用固定 512-token 步长，直至配置的piecewise_cudagraph_max_tokens上限（standard.py_DEFAULT_TOKEN_STRIDE）。所有候选桶都在启动时捕获并常驻，没有运行时捕获或驱逐。该上限是执行策略，不是调度器或分页容量。在配置层面，它对应 config.pyBackendConfig.piecewise_cudagraph_max_tokens，默认None（即默认关闭 PCG）。

当前静态资格要求（runner.py_make_piecewise_graph_manager）：

• 模型实现PiecewiseCudaGraphMixin；
• CUDA 执行且eager 模式被禁用（backend_config.eager_mode为假）；
• 每个选定的 CUDA 实现都支持 PCG（step_meta_plan.enable_piecewise_cuda_graph()）。

并行策略方面：TP 与默认模式 DP/EP 支持每个 rank 一个独立计划；DeepEP routed experts 保持 eager 边界；DP_TPlayer 模式在边界安装之前就被拒绝，因为其 live per-rank 切分与集合通信尚未被计划表示——但这不拒绝DP2/EP4 部署形成的默认 attention-TP 组。

请求时选择还会拒绝microbatch prefill、chunked prefill 与 live LoRA 适配器（当前描述符无法表示这些模式）；动态额外 forward 实参被拒绝，不可变的标量类 extras 则纳入描述符。

数值正确性提示

设计文档给出了一条重要警告：填充会改变 GEMM 形状，可能改变浮点结果。验证机制正确性时应将 PCG 与等价的桶形状 eager 执行比较；同时单独测量其对用户可见的数值与性能影响（与原始 eager 执行对比）。512-token 步长是临时性的，不是通用的盈利阈值。

失败与重置规则

设计文档列出七条硬性规则，可对照源码逐条验证：

1. 不支持的或缺失的 serving 计划在 PCG 开始之前选择 eager；
2. warmup、build、replay 失败都会传播——PCG 可能已写入 KV/状态缓存后，绝不重试整个模型 forward；
3. 失败的 build不发布任何部分计划；活跃捕获清理是 best-effort 的，同时保留原始异常（piecewise.py_discard_active_capture）；
4. serving 期间从不驱逐计划；
5. CUDAGraphRunner.reset()在权重、cache arena 或其他被捕获地址可能变化之前使完整计划失效；
6. 稳态回放使用流排序而非设备级同步。

新增分段支持的八步方法论

设计文档给出了模型/算子集成 PCG 的完整清单，核心原则是保持通用 runner 对模型和算子无感知：

1. 判定哪些工作真正依赖 live 请求布局、元数据、状态、宿主控制流或不支持的集合通信，把其余张量代数保留为被捕获；
2. 若模型符合标准解码器输入/输出契约，用PiecewiseCudaGraphMixin选择接入，不要每个模型加一个运行时文件；
3. 让每个选定 CUDA 实现上报supports_piecewise_cuda_graph()并通过enable_piecewise_cuda_graph()安装自己的包装器；能力安装通过CudaStepMetaPlan保持原子性（相关抽象见 step_metadata.py）；
4. 通过模型、层、算子调用链保持语义元数据显式——不要藏在 graph runner 里、不要作为通用算子缓存挂在StepContext上、也不要向后端传入模型所有者；
5. 选择一种输出策略：固定图可见张量用普通 bridge，不同语义用窄的算子拥有的适配器，仅由后续 eager 步骤消费的值用 eager-only 槽；
6. 把每个会改变捕获形状或边界顺序的事实放进描述符，不要以只被 eager 边界消费的元数据作为 key；
7. 确保启动时 warmup 并捕获 serving 可能选择的每一种尺寸；
8. 验证：相同聚合 token 数不同请求分区、prefix/history 缓存写入、状态更新、reset/rebuild、步骤顺序、保留内存与真实延迟。

同时设计文档划出了piecewise.py与standard.py的禁止事项：不要添加模型/算子分支、attention 重分发、合成元数据、调度器页、运行时 warmup、回退恢复或签名/切片 DSL。

刻意不支持的设计边界

设计文档列出九项明确不支持的能力，它们是显式设计边界，添加任何一项都需要具体的算子、清晰的所有权、聚焦的正确性证据与实测性能：

• 运行时捕获与计划驱逐；
• 单模型 runner 上的并发 build 或回放；
• 单个边界结果中混合 eager-only 与图可见叶子；
• 嵌套请求对象的递归绑定；
• 通用 alias/view 保持型输出 bridge；
• DP_TPlayer 执行；
• 动态额外标准解码器 forward 实参；
• 边界顺序的任意数据依赖变化；
• PCG 在 PyTorch 2.0 上不可用（非拥有 CUDA 视图依赖 PyTorch 2.1 起的私有 API，见 piecewise.py_make_weak_cuda_view）。

评审清单：graph runner 变更的十条自检

设计文档附带的评审清单适用于任何 graph-runner 相关改动，可在代码评审时逐条对照：

1. 是否有一个组件清晰拥有每个图、池、流、bridge、缓冲与重置转换？
2. 是否可能有已发布步骤保留 dummy/startup 请求对象？
3. 每个被捕获的张量地址在计划生命周期内是否保持有效？
4. 输入值能否在不改变描述符的情况下改变边界顺序？
5. 不支持/未准备的请求是否在产生副作用之前选择 eager？
6. 异常是否可能导致完整 forward 被执行两次？
7. 算子特有的切片、元数据、状态或集合逻辑是否留在拥有它的 CUDA 实现中？
8. 新的输出形状是否需要 bridge 适配器、描述符事实，或两者都要？
9. bridge 复用生命周期是否在声明处得到证明？
10. 连续批处理布局、cache/状态变更、reset、内存与延迟，是否有与改动规模相称的证据覆盖？

文档末尾还提到，早期独立的 PCG 测试与实验存放在/home/yaoqian/space/tmp/lmdeploy_test/develop/pcg（接口仍在稳定中，该路径为开发环境路径，仓库内并未包含这些文件）。

总结

LMDeploy 的 CUDA Graph Runner 通过"完整 decode 图 + 分段 prefill 图 + eager 回退"的三段式架构，把 CUDA Graph 的收益覆盖到 decode 与 prefill 两个阶段：decode 沿用成熟的单图方案，prefill 则以 PCG 的GraphStep -> EagerStep流水线把请求相关、状态相关或集合通信相关的计算留在 eager 边界，其余张量代数全部捕获进 CUDA 图。所有缓冲、bridge、流、池与失效转换都由单一组件（CUDAGraphRunner/PiecewiseGraphManager）统一所有，并通过描述符机制把"哪些事实决定捕获形状与边界顺序"显式化。对希望深入 LMDeploy PyTorch 后端推理加速的读者，建议从 README.md 通读设计，再对照 runner.py、piecewise.py 与 standard.py 逐一落点验证，即可完整掌握这套机制的实现脉络。

• 人工智能
• 大模型
• 模型推理服务
• 推理引擎
• 本地部署
• 模型量化