CUDA Graph 在大模型在线推理中的学习与应用

一、为什么需要 CUDA Graph？

在大模型学习推理中，尤其是在线服务场景下，低延迟和高吞吐是关键指标。一个模型的前向计算通常会被分解成许多细粒度的 GPU 内核（kernel），例如矩阵乘法、激活函数、归一化等。

每个内核的启动都需要 CPU 向 GPU 发送一个命令，这个过程称为 CPU launch overhead（CPU 启动开销）。

当模型较小、内核数量多且每个内核执行时间很短时，CPU 启动开销可能成为性能瓶颈——GPU 大部分时间在等待命令，而不是真正计算。

一句话理解：当单个 kernel 的 GPU 执行时间 < CPU 提交它所需的时间时，GPU 就会追上 CPU，队列被抽干，开始出现气泡。

CUDA Graph 正是为了解决这一问题而设计的。

它允许将一系列 GPU 操作捕获成一个图（graph），之后只需一次调用即可重放整个图，从而大幅减少 CPU 的参与，降低延迟，提高 GPU 利用率。

在大模型在线推理中，自回归解码（autoregressive decoding）的每一个 step 都会重复执行几乎相同的计算图（只是输入 token 和内部状态变化），因此特别适合使用 CUDA Graph 进行优化。

 

先看一个demo 

import torch
import torch.nn as nn
import time

# ---------- 1. 定义一个简单的模型 ----------
class SimpleModel(nn.Module):
    def __init__(self, in_dim=128, hidden_dim=256, out_dim=128):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, out_dim)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

# ---------- 2. 关键包装器：将输出写入固定地址 ----------
class FixedOutputWrapper(nn.Module):
    """
    包装模型，使其输出写入预先分配的张量中。
    这样在 CUDA Graph 捕获时，输出地址固定不变。
    """
    def __init__(self, model, output_buffer):
        super().__init__()
        self.model = model
        self.output_buffer = output_buffer   # 预分配的输出张量

    def forward(self, x):
        # 计算模型输出，并将其复制到固定地址
        self.output_buffer.copy_(self.model(x))
        return self.output_buffer           # 返回固定地址的张量


# ---------- 3. 初始化 ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
assert device.type == 'cuda', "此示例需要 CUDA GPU"

batch_size = 64
in_dim = 128
out_dim = 128
num_iter = 200

model = SimpleModel(in_dim, 256, out_dim).to(device).eval()

# 预分配输入/输出张量（地址固定）
input_buffer = torch.randn(batch_size, in_dim, device=device)
output_buffer = torch.empty(batch_size, out_dim, device=device)

# 用包装器固定输出地址
wrapped_model = FixedOutputWrapper(model, output_buffer).to(device)





# ---------- 4. 常规执行（基线） ----------
print("=" * 40)
print("常规执行（每次调用启动多个 CUDA 内核）")

# 预热
for _ in range(10):
    _ = wrapped_model(input_buffer)


torch.cuda.synchronize()
start = time.time()
for _ in range(num_iter):
    _ = wrapped_model(input_buffer)
torch.cuda.synchronize()
end = time.time()
normal_time = (end - start) / num_iter * 1000
print(f"平均耗时: {normal_time:.4f} ms")





# ---------- 5. CUDA Graph 执行 ----------
print("\n" + "=" * 40)
print("CUDA Graph 执行（一次性提交整个计算图）")

# 创建 CUDA Graph 对象
graph = torch.cuda.CUDAGraph()

# 捕获计算图（必须使用固定地址的输入/输出）
with torch.cuda.graph(graph):
    # 在捕获上下文中，所有操作都会被记录，但不会立即执行
    out = wrapped_model(input_buffer)

# 验证输出地址是否固定
print(f"输出张量地址（重放前后应相同）: {hex(output_buffer.data_ptr())}")

# 预热重放
for _ in range(10):
    graph.replay()

torch.cuda.synchronize()
start = time.time()
for _ in range(num_iter):
    graph.replay()
torch.cuda.synchronize()
end = time.time()
graph_time = (end - start) / num_iter * 1000
print(f"平均耗时: {graph_time:.4f} ms")

# ---------- 6. 性能对比 ----------
print("\n" + "=" * 40)
print(f"加速比: {normal_time / graph_time:.2f}x")
print("说明：CUDA Graph 在小模型/小 batch 时能显著降低 CPU 启动开销。")

运行结果

root@victory-d1-tangseng-gpu-vff6l:~/self_mini_vllm# python  cuad-graph.py 
========================================
常规执行（每次调用启动多个 CUDA 内核）
平均耗时: 0.0548 ms

========================================
CUDA Graph 执行（一次性提交整个计算图）
输出张量地址（重放前后应相同）: 0x7f6c9b848600
平均耗时: 0.0114 ms

========================================
加速比: 4.80x
说明：CUDA Graph 在小模型/小 batch 时能显著降低 CPU 启动开销。

　　

 

二、CUDA Graph 基础概念

1. 什么是 CUDA Graph？

CUDA Graph 是 CUDA 10 引入的一种机制，它允许将一组 CUDA 操作（内核启动、内存拷贝等）记录为一个有向无环图（DAG）。

图中的节点表示操作，边表示依赖关系。

捕获完成后，整个图可以被实例化（instantiate）并重放（replay）。

重放时，GPU 会按照图定义的顺序和依赖关系执行所有节点，而 CPU 只需发起一次调用。

2. 捕获（Capture）与重放（Replay）

• 捕获：在某个 CUDA stream 上启动一个捕获模式，所有在该 stream 上发出的操作不会立即执行，而是被记录到图中。

• 重放：将捕获好的图提交给 GPU 执行，等效于按序执行图中的所有节点。

3. 重要限制

• 内存地址固定：图在捕获时记录了所有张量的内存地址。重放时，它假设这些地址上的数据仍然有效。因此，所有输入、输出以及中间变量都必须使用预先分配且地址不变的张量。

• 无动态控制流：图是静态的，不能包含依赖于运行时数据的条件分支或循环（但可以通过多个图或使用 CUDA Graph 的“图更新”来处理有限的变化）。

• 与 stream 的关系：图必须在某个 stream 上捕获，重放时也应在同一 stream 上（或通过事件同步）。

后续继续完善吧 ！ 知道个大概就行了先！

 

第一种浪费最反直觉：GPU 并不慢，它只是在排队等 CPU 告诉它下一步做什么。

这个问题在 Prefill 阶段几乎看不见（单个 kernel 算得久，提交开销被淹没），却会在 Decode 阶段被放大——因为 Decode 每步的计算量太小了。