在线大模型服务中，每次生成一个 token 都需要执行数十个甚至上百个 GPU kernel。在 eager 模式下，CPU 需要逐个启动这些 kernel，每个 kernel 的启动开销（包括参数校验、命令提交等）虽然只有微秒级，但当 kernel 数量众多且单个 kernel 执行时间很短时，CPU 启动开销可能成为瓶颈，导致 GPU 利用率下降。例如，在 decode 阶段，每个 token 的生成往往只涉及少量计算，但 kernel 启动开销却占据了大量时间。CUDA Graph 正是为了解决这一问题而设计的，它通过将一段 GPU kernel 序列捕获为静态图，并一次性提交执行，从而显著降低 CPU 开销。本文将以一个在线大模型服务场景为例，深入探讨 CUDA Graph 的工作原理、与动态形状和批处理大小变化的兼容性、与 PagedAttention 的协同，以及其捕获成本、内存占用和适用边界。
问题背景：kernel 启动开销为何成为瓶颈
在大模型推理中，一个典型的 decode 步骤可能包含多个 Transformer 层的计算，每层又包含矩阵乘法、归一化、激活函数等操作，总计可能涉及数百个 kernel。在 eager 模式下，CPU 需要逐个启动这些 kernel，每个 kernel 的启动开销包括参数解析、命令缓冲区提交等，通常在 5～10 微秒左右。当 kernel 执行时间很短（例如几微秒）时，启动开销占比极高，导致 GPU 空闲等待。
更严重的是，CPU 启动 kernel 的速度可能跟不上 GPU 的执行速度，导致 GPU 利用率下降。例如，在 batch size 较小的情况下，kernel 执行时间可能只有几微秒，而 CPU 启动一个 kernel 需要约 10 微秒，此时 GPU 大部分时间在等待 CPU 提交命令。这种 CPU 与 GPU 之间的速度不匹配，使得单次请求的延迟增加，也限制了系统的吞吐量。
传统的优化手段包括增大 batch size 以摊薄启动开销，但 batch size 受限于显存大小和请求到达率。另一种手段是使用异步执行，但异步执行只能减少 CPU 等待时间，并不能减少 CPU 启动 kernel 的总次数。CUDA Graph 则从根本上去除了逐个 kernel 启动的开销，将整个计算流作为一个整体提交。
CUDA Graph 的核心机制：捕获、实例化与重放
CUDA Graph 的工作流程分为三个阶段：捕获（Capture）、实例化（Instantiate）和重放（Replay）。
捕获阶段，调用 cudaStreamBeginCapture() 后，CUDA runtime 进入录制模式，后续提交到该 stream 的所有操作（kernel launch、memcpy、memset 等）都不会真正执行，而是被记录为图节点。每个节点保存了 kernel 的入口地址、grid/block 维度以及所有参数值（对于 tensor 而言是 GPU 虚拟地址）。节点之间的依赖关系由 stream 上的提交顺序和跨 stream 的 event 同步自动推断。捕获结束时，调用 cudaStreamEndCapture() 返回一个 cudaGraph_t 对象，它是对计算流的静态描述。
实例化阶段，调用 cudaGraphInstantiate() 将 cudaGraph_t 编译为可执行的 cudaGraphExec_t。这一阶段会进行依赖分析，确定哪些 kernel 可以并发执行，并将所有参数绑定到固定的内存地址。实例化过程可能耗时较长，但只执行一次。
重放阶段，调用 cudaGraphLaunch(exec, stream) 将整个图一次性提交到 GPU。CPU 只发出一次 launch 指令，GPU 端的调度器按照预定的执行计划依次或并发地执行所有 kernel。由于重放不经过 Python/PyTorch 的 dispatcher，也没有 CPU 端的逐 operation 调度，CPU 开销几乎降为零。
在 PyTorch 中，CUDA Graph 被封装为 torch.cuda.CUDAGraph 类，提供了 capture_begin()、capture_end() 和 replay() 方法，分别对应捕获开始、捕获结束和重放。
动态形状与批处理大小变化：CUDA Graph 的约束与应对
CUDA Graph 的一个关键约束是：捕获时录制的参数（如 grid 维度、tensor 形状和地址）被固化在 cudaGraphExec_t 中，因此一份图只能服务一种 batch size。如果 batch size 改变，kernel 的 grid 维度、中间 tensor 的形状和内存布局都会发生变化，导致图失效。
例如，在 vLLM 的 decode 阶段，每个请求每步只产生一个 token，batch size 为 4 时，输入 tensor 的第一维为 4，对应的 kernel grid 大小、中间 tensor 形状都是按 4 设计的，与 batch size 为 8 时的布局完全不同。因此，一份图只能用于一种 batch size。
为了应对这种动态变化，常见的做法是为一组离散的 batch size（如 1、2、4、8、16 等）各捕获一份图。当实际请求数恰好命中某个预录的 batch size 时，走图重放；否则回退到 eager 模式。这种策略在 vLLM 和 SGLang 等系统中广泛使用。
此外，CUDA Graph 还要求捕获期间不能有动态内存分配，因为动态分配可能导致地址变化。因此，所有 tensor 必须预先分配，这通常通过内存池来实现。捕获期间分配的内存区域会被 CUDA Graph 锁定，不能用于其他用途，这增加了显存占用。
与 PagedAttention 的兼容性：内存管理的挑战与解决
PagedAttention 是 vLLM 中提出的注意力算法，灵感来自操作系统的虚拟内存分页技术。它将 KV cache 划分为固定大小的块，通过块表管理，实现了 KV cache 内存的按需分配和共享，从而减少了内存碎片和冗余复制，提高了批处理大小。
然而，PagedAttention 的动态内存分配特性与 CUDA Graph 的静态图约束存在冲突。在捕获阶段，PagedAttention 可能会动态分配新的 KV cache 块，这会导致地址变化，破坏图的稳定性。因此，直接使用 CUDA Graph 时，需要预先分配足够的 KV cache 块，并确保捕获期间不发生新的分配。
vLLM 的做法是：在捕获图之前，为每个请求预先分配固定数量的 KV cache 块，并确保这些块在捕获期间保持不变。此外，vLLM 还利用了 PagedAttention 的块表机制，在重放时通过更新块表来改变 KV cache 的物理位置，而不改变图结构，从而实现了动态批处理和内存共享。
工程实现：多图复用与内存池
在实际系统中，CUDA Graph 的工程实现涉及多图复用和内存池管理。
多图复用是指为不同的 batch size 或不同的模型配置捕获多个图，并在运行时选择合适的图进行重放。例如，SGLang 会为一组离散的 decode batch size 各捕获一份图，当实际请求数命中某个预录的 batch size 时，走图重放；否则回退到 eager 模式。
内存池管理是 CUDA Graph 捕获的关键。捕获期间，所有中间 tensor 的内存地址会被固化在图中，因此这些内存必须保持稳定。PyTorch 的缓存分配器在捕获期间会锁定内存池，防止地址变化。此外，捕获期间分配的内存区域会被 CUDA Graph 整体锁定，不能用于其他用途，这增加了显存占用。
在 SGLang-Omni 框架中，为 Fish Audio 的 S2-Pro TTS 模型添加 CUDA Graph 支持时，作者发现需要为慢 AR 和快 AR 两个不对称的自回归过程统一捕获一个图，并使用了 deferred graph capture 和 persistent buffer 等技术。通过统一捕获，TPS 从 55.6 提高到 88，效果显著。
何时使用 CUDA Graph：收益与代价的权衡
CUDA Graph 并非适用于所有场景，其收益与代价需要权衡。
收益：当计算流固定且重复执行时（如 decode 阶段的固定 kernel 序列），CUDA Graph 可以显著降低 CPU 启动开销，提高 GPU 利用率，从而提升吞吐量和降低延迟。对于 kernel 数量多且单个 kernel 执行时间短的场景，收益尤为明显。
代价：CUDA Graph 的捕获和实例化需要额外的时间，且捕获期间的内存池被锁定，增加了显存占用。此外，由于图是静态的，无法适应动态形状和动态控制流，需要为不同的 batch size 捕获多份图，增加了内存开销和实现复杂度。
以下表格对比了 CUDA Graph 与 eager 模式在关键维度上的差异：
| 维度 | Eager 模式 | CUDA Graph |
|---|---|---|
| CPU 启动开销 | 每个 kernel 单独启动，开销高 | 一次启动整个图，开销极低 |
| 动态形状支持 | 完全支持 | 仅支持预捕获的形状 |
| 批处理大小变化 | 灵活调整 | 需要为每个 batch size 单独捕获图 |
| 内存占用 | 动态分配，碎片较多 | 内存池锁定，占用较高 |
| 实现复杂度 | 简单 | 需要处理捕获约束和内存管理 |
| 适用场景 | 原型开发、动态形状、小模型 | 固定形状、高吞吐、大模型推理 |
常见失败模式与调试策略
CUDA Graph 的失败模式主要源于其静态特性。常见的失败包括：
- 捕获失败：捕获期间如果出现 host-device 同步（如调用
.item()或torch.multinomial），会导致捕获失败。这是因为同步操作会中断 stream capture。 - 数值错误：如果捕获期间使用的内存地址在重放时发生变化，可能导致 kernel 读写错误内存，产生数值错误。这通常是由于内存池管理不当或动态分配未预分配所致。
- 性能问题：如果图结构不合理，例如存在不必要的依赖边，可能导致 GPU 利用率下降。此外，如果 batch size 频繁变化，导致经常回退到 eager 模式，性能提升可能不明显。
调试 CUDA Graph 时，可以使用 CUDA 提供的工具（如 cuda-gdb 和 nsys）来检查图的捕获和重放过程。常见策略包括：确保捕获期间没有动态内存分配，使用固定大小的输入，以及在捕获前预热内存池。
结论
CUDA Graph 通过捕获并重放 GPU kernel 序列，有效消除了大模型推理中的 kernel 启动开销，显著提升了吞吐量和降低了延迟。然而，它并非万能，其静态特性与动态形状和批处理大小变化存在冲突，需要结合 PagedAttention 等内存管理技术进行适配。在实际应用中，应根据具体场景权衡收益与代价，选择是否使用 CUDA Graph。
以下流程图展示了 CUDA Graph 在在线服务中的典型工作流程：
flowchart TD
A[接收请求] --> B{是否命中预录batch size?}
B -- 是 --> C[选择对应CUDA Graph]
C --> D[重放图]
D --> E[生成token]
B -- 否 --> F[回退到eager模式]
F --> E
E --> G{是否结束?}
G -- 否 --> A
G -- 是 --> H[返回响应]
该流程展示了在线服务中如何根据请求的 batch size 选择图重放或 eager 模式，以平衡性能和灵活性。