CUDA Graph 与 Launch Overhead

真实收益来源

Eager 执行一段 GPU 工作时，Python/框架逐个准备并向 CUDA 提交 kernels、memcpy 和依赖。单个大 GEMM 耗时远大于 launch 时，这部分不重要；decode 小 batch 中若一层包含大量短 kernels，CPU dispatch 可能进入串行关键路径。

CUDA Graph 在一次 capture 中记录 GPU work DAG，之后 replay 复用已实例化的图。极简延迟模型：

T_{eager}\approx\sum_i(T_{launch,i}+T_{kernel,i})+T_{framework}

T_{replay}\approx T_{graph\ launch}+\sum_iT_{kernel,i}

主要减少的是许多 $T_{launch,i}$ 和 framework scheduling；图内 GEMM、attention 与 memory operation 的数学工作仍存在。只有 CPU launch 不能与 GPU 工作完全重叠并且占比足够大时，端到端收益才明显。

为什么要求静态地址与稳定控制流

Graph 节点记录 kernel 参数与内存地址。PyTorch 的常见用法是准备 static_input / static_output，capture 后只用 copy_ 把新数据写入相同 storage，再 replay：

new request tensor ─copy_→ captured static input address
                                │
                            graph replay
                                │
                         captured output address

关键约束：

shape、stride、dtype、device 和地址要满足 capture 时的节点参数；
capture 期间不能依赖 CPU 读取 GPU 值后选择动态分支；
allocator 必须提供 graph-private 或 graph-safe 内存，避免 replay 时地址失效；
随机操作需要 graph-safe RNG 状态；
collective、host callback 与第三方库是否可 capture 取决于具体支持；
capture 前要 warmup，初始化 lazy library handles、kernel modules 和 allocator 行为。

动态 batch 如何处理

生产推理的 batch、sequence length 和 scheduler 状态是动态的。常见策略不是捕获“任意 shape 的一个图”，而是：

为若干 batch-size / shape buckets 分别 capture；
padding 到已捕获 bucket，并接受额外计算；
只 capture 模型中 shape 稳定的 decode 段，动态调度和采样留在图外；
对无法 capture 的路径 fallback eager；
使用 graph pool 让多个 graphs 在已知生命周期下复用内存。

Bucket 越多，capture/warmup 时间与 graph pool 显存越高；bucket 越少，padding waste 越大。Continuous batching 还要保证 slot address 稳定，通常由固定 batch slots 和间接 metadata 配合。

可运行的静态地址实验

实验问题：固定 shape/address 时 eager 与 graph replay 输出是否一致？包含一次 input copy_ 的 replay 延迟是否低于一串小 eager kernels？动态 shape 是否会在静态 input copy 边界被拒绝？

import platform,statistics,time
import torch

torch.manual_seed(83)
assert torch.cuda.is_available(),"CUDA Graph 示例需要 CUDA"
DEVICE=torch.device("cuda"); DTYPE=torch.float32

class TinyBlock:
    def __init__(self,width=128):
        self.w1=torch.randn(width,width,device=DEVICE,dtype=DTYPE)/width**.5
        self.w2=torch.randn(width,width,device=DEVICE,dtype=DTYPE)/width**.5
    def __call__(self,x):
        return torch.relu(x@self.w1)@self.w2+x

def median_us(fn,warmup=20,rounds=200):
    for _ in range(warmup): fn()
    torch.cuda.synchronize(); samples=[]
    for _ in range(rounds):
        start=time.perf_counter(); fn(); torch.cuda.synchronize()
        samples.append((time.perf_counter()-start)*1e6)
    return statistics.median(samples)

def main():
    block=TinyBlock(); sample=torch.randn(1,128,device=DEVICE,dtype=DTYPE)
    side=torch.cuda.Stream(); side.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(side):
        for _ in range(3): block(sample)
    torch.cuda.current_stream().wait_stream(side)

    static_input=sample.clone(); graph=torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph): static_output=block(static_input)
    update=torch.randn_like(sample); eager=block(update)
    static_input.copy_(update); graph.replay()
    torch.testing.assert_close(static_output,eager,rtol=1e-5,atol=1e-6)
    input_ptr,output_ptr=static_input.data_ptr(),static_output.data_ptr()

    eager_us=median_us(lambda:block(update))
    def replay_step(): static_input.copy_(update); graph.replay()
    replay_us=median_us(replay_step)
    assert static_input.data_ptr()==input_ptr and static_output.data_ptr()==output_ptr
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={torch.cuda.get_device_name(0)} dtype={DTYPE} shape={tuple(sample.shape)}")
    print(f"input_ptr_static={input_ptr} output_ptr_static={output_ptr}")
    print(f"max_abs_diff={(static_output-eager).abs().max():.3e}")
    print(f"eager_median_us={eager_us:.3f}")
    print(f"copy_plus_graph_replay_median_us={replay_us:.3f}")
    try: static_input.copy_(torch.randn(2,128,device=DEVICE))
    except RuntimeError as error: print(f"dynamic_shape_rejected={type(error).__name__}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/cuda_graph_reference.py。2026-08-13 在 RTX 3080、PyTorch 2.13.0+cu130 上，输出最大差为 0，eager 中位数 22.252 µs，input copy + graph replay 13.916 µs，不同 shape 的 copy_ 抛出 RuntimeError。

这只是一个包含两个小 matmul、ReLU 与 residual add 的 launch-overhead 演示。同步每轮会放大 CPU 等待但保证计时边界明确；生产服务应使用 CUDA Events/端到端 trace 分别分析 GPU execution 与请求延迟。本机差值不能外推为所有模型的 graph 加速比。

何时不值得或会失败

单个大 kernel 已占主导，launch 占比很小；
shape/control flow 高度动态，bucket padding 或 graph cache 成本超过收益；
capture 内需要 CPU 同步、动态 allocation 或不支持 capture 的操作；
graph pool 提高 reserved memory，挤压 KV Cache 容量；
多个 graphs 的地址/lifetime 复用错误，导致非法访问或静默数据覆盖；
benchmark 把 capture 时间算进每次 replay，或反过来完全忽略生产中的 bucket warmup 成本。

与上下游关系

Prefill/Decode 中 decode shape 更小、kernel 更碎，常比长 prefill 更容易从减少 launch overhead 获益；但 dynamic batch 也更复杂。
GEMM/GEMV 与算术强度解释 graph 没有改变的 kernel compute/memory 主体。
FlashAttention 是 attention kernel/backend 优化；可以被 capture，但其 online softmax 与 IO 收益独立于 graph replay。
Backend Dispatch/Fallback应在 capture 前确认：graph 只会稳定重放已经选中的 kernels，不会自动把 fallback 路径升级成 fast path。
后续 Continuous Batching 页面需要说明固定 graph batch buckets 如何与请求进入/退出、padding 和 slot reuse 配合。