PyTorch CUDA Graph 怎样固定输入形状复用推理流程

如果同一个 PyTorch 推理模型被高频重复调用，而每个请求的输入形状基本稳定，CUDA Graph 可以把已经捕获的 GPU 工作整体重放，减少 Python、C++ 和 CUDA 驱动层的调度开销。关键不是把每个新批次重新建图，而是让新数据复制进一组长期存活、形状固定的 Tensor，再调用 replay()。

官方地址：https://docs.pytorch.org/docs/stable/notes/cuda.html

这篇只讨论 CUDA Graph 的输入形状、虚拟内存地址和推理流程复用。代码是结构示例，不代表已经在本机运行；配图也是静态说明图，不是截图或性能证据。

这个问题的直接答案

把请求按固定形状分组，为每个形状准备一组静态输入和静态输出：先在旁路流上完成几次预热，再用 torch.cuda.graph 捕获模型调用。运行时只把新批次复制到 static_input，随后调用 graph.replay()，读取或复制 static_output。

CUDA Graph 不适合任意动态输入。捕获期间不能依赖会改变形状的控制流，不能在图中做 .item() 这类 CPU-GPU 同步，也不能让每次 replay 换一组输入输出地址。变长请求应使用多个形状桶，无法归桶时回退 eager 执行。

为什么 CUDA Graph 要求输入形状稳定

PyTorch 官方文档把 CUDA Graph 描述为对 CUDA stream 工作的记录：捕获后可以重复启动同一组 kernel 和参数。指针参数意味着每次 replay 使用相同的内存地址，因此图复用的是一条固定的 GPU 工作路径，而不是一段可以自由改变输入尺寸的普通 Python 函数。

捕获前至少要确认三件事：

• 模型输入、输出以及中间张量的尺寸和布局在同一张图的 replay 中保持一致。
• 控制流不能由 CPU 或不受支持的动态数据分支决定；捕获期间不能调用会把 GPU 值同步回 CPU 的操作。
• 预热、捕获和 replay 使用同一设备，并保留静态输入、静态输出和图对象的引用。

最小配方：静态输入、预热、捕获和 replay

下面的示例把批量大小和特征维度固定下来。旁路流只负责预热，捕获阶段由 torch.cuda.graph 管理流；真正的请求数据在 replay 前通过 copy_ 写入已经存在的输入 Tensor。

import torch

# 用固定形状的模型示例说明一张图对应一类输入桶
model = torch.nn.Sequential(
    torch.nn.Linear(128, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10),
).cuda().eval()

batch_shape = (8, 128)
static_input = torch.empty(batch_shape, device="cuda")
graph = torch.cuda.CUDAGraph()

# 先在旁路流上预热，让分配和库状态稳定下来
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream), torch.inference_mode():
    for _ in range(3):
        model(static_input)
torch.cuda.current_stream().wait_stream(warmup_stream)

# 捕获固定形状的前向路径；输出 Tensor 要长期保留
with torch.inference_mode(), torch.cuda.graph(graph):
    static_output = model(static_input)

def infer_one_bucket(batch: torch.Tensor) -> torch.Tensor:
    # 只替换固定地址中的内容，不替换 Tensor 本身
    if tuple(batch.shape) != batch_shape:
        raise ValueError(f"需要形状 {batch_shape}，实际得到 {tuple(batch.shape)}")
    static_input.copy_(batch)
    graph.replay()
    # 复制结果，避免下一次 replay 覆盖调用方正在使用的输出
    return static_output.clone()

这段代码的核心约束是“内容可变，地址不变”。batch 可以来自不同请求，但它的形状必须匹配；static_input 和 static_output 则由图对象持有并跨请求复用。若调用方允许直接消费固定输出，也可以不做 clone()，但必须明确下一次 replay 会覆盖它。

固定地址怎样承载不同批次数据

每次请求的真实批次不是 CUDA Graph 的新输入指针。更稳定的关系是：真实批次把内容复制到 static_input，图按照捕获时记录的地址读取它，执行固定的模型路径，再把结果写到 static_output。PyTorch 的 CUDA caching allocator 会为捕获分配图私有内存池，图和捕获期间创建的 Tensor 存活时，这组地址才有稳定的生命周期。

变长请求怎样落地

最常见的工程做法是建立形状桶，例如把序列长度分成若干档，每一档各自维护 static_input、输出 Tensor 和 CUDAGraph。请求先按形状选择图，再把有效数据复制进对应缓冲区；超出最大长度或需要完全不同控制流的请求，直接走 eager 路径。

不要为了“复用一张图”把不同长度强行塞入同一 Tensor 后再在图内动态判断。这样通常会把 padding、掩码和输出裁剪混在一起，既增加内存占用，也可能违反捕获时的静态形状约束。更合理的边界是：少量稳定形状用多张图，形状高度离散则评估图维护成本是否值得。

常见误区与速查

归纳起来，CUDA Graph 的复用单位不是“任意输入的一次函数调用”，而是“固定形状、固定地址、固定控制流的一段 GPU 工作”。先把请求形状整理成少量稳定桶，再为每桶建立长生命周期的图对象，通常比在图内追逐动态输入更容易维护。

相关问题

CUDA Graph 一定比 eager 快吗？

不一定。它主要减少 CPU 侧的参数准备和 kernel dispatch 开销；如果模型本身已经被大算子占满 GPU，收益可能有限，应该用与业务负载匹配的基准数据评估。

为什么不能每次请求重新创建 static_input？

因为图记录的是指针参数。换 Tensor 通常也就换了虚拟地址，replay 仍会读捕获时的地址，导致数据没有进入图实际读取的位置，甚至触及非法访问风险。