GPU 算力没在算模型，在等 CPU 发号施令——CUDA Graph 是怎么解决这件事的
GPU 算力没在算模型，在等 CPU 发号施令——CUDA Graph 是怎么解决这件事的
1. 你大概有这样一段代码
用 transformers 加载一个模型，写了个简单的 decode 循环：
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B",
torch_dtype=torch.float16).cuda()
model.eval()
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B")
prompt = "解释一下推测解码的原理"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
output = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(output[0]))
能跑，但总觉得慢，想优化。
加量化、加 flash attention、换 vLLM……这些方法都有效，但 CUDA Graph 是其中最”无副作用”的一个——不改模型权重，不改输出质量，只改”GPU 执行任务的方式”，在 decode 阶段能有 20-50% 的提升，代码改动量也不大。
这篇文章就讲怎么给自己的推理代码加 CUDA Graph。
2. 先搞清楚为什么慢：kernel launch overhead
大模型的一次前向，不是一个大计算，而是几百个小 CUDA kernel 的串行组合：RMSNorm、RoPE、QKV projection、attention、FFN gate/up/down、residual add……每个都是独立的 kernel。
每次调用一个 kernel，CPU 都要做一次 kernel launch：打包参数、通知 driver、把任务推进 GPU 队列。这个操作本身耗时 5-20 µs，看起来不多，但一次前向几百个 kernel 累积下来就是几毫秒。
Decode 阶段每步只生成一个 token，单步耗时可能就 10-20 ms，kernel launch overhead 就占了 20-30%。GPU 的相当一部分时间在等 CPU 一遍遍地发号施令，而不是在做计算。
用 nsight systems profile 一下会很直观——timeline 上 GPU 活跃段之间有一堆小空隙，那些就是 CPU launch 的等待时间。
CUDA Graph 的解法很直接：第一次推理时把这几百个 kernel 的序列”录制”成一张图（DAG），之后每次推理直接 replay 这张图——整个 forward 只需要一次 CPU 调用，launch overhead 接近清零。
3. 先跑一个最简单的 demo，感受一下
在套到 LLM 之前，先用一个最小的可运行例子感受 CG 的效果——直接复制跑，有 GPU 就行：
import torch
import time
device = "cuda"
# 用一个小 MLP 模拟 LLM 里的一层
model = torch.nn.Sequential(
torch.nn.Linear(2048, 8192),
torch.nn.GELU(),
torch.nn.Linear(8192, 2048),
).to(device).half().eval()
# ⚠️ 关键：输入 tensor 提前分配好，地址整个生命周期不变
x = torch.randn(8, 2048, device=device, dtype=torch.float16)
# --- Step 1: warmup（让 kernel 编译完成、显存分配稳定）---
with torch.no_grad():
for _ in range(3):
_ = model(x)
# --- Step 2: capture ---
g = torch.cuda.CUDAGraph()
with torch.no_grad():
with torch.cuda.graph(g):
y = model(x) # 被"录制"进 DAG，这次不真正执行
# --- Step 3: replay ---
# 每次推理：就地更新输入（地址不变），然后 replay 整张图
x.copy_(torch.randn_like(x))
g.replay()
# y 的地址是 capture 时固定的，直接读
print(y.shape) # torch.Size([8, 2048])
# --- 对比一下速度 ---
N = 2000
torch.cuda.synchronize()
t0 = time.time()
with torch.no_grad():
for _ in range(N):
x.copy_(torch.randn_like(x))
_ = model(x)
torch.cuda.synchronize()
print(f"Eager: {(time.time() - t0) * 1000 / N:.3f} ms/iter")
torch.cuda.synchronize()
t0 = time.time()
for _ in range(N):
x.copy_(torch.randn_like(x))
g.replay()
torch.cuda.synchronize()
print(f"CG: {(time.time() - t0) * 1000 / N:.3f} ms/iter")
在 A100 上跑出来大概是：
Eager: 0.130 ms/iter
CG: 0.032 ms/iter
四倍差距，全来自消掉了几十个 kernel 的 launch overhead。这个 MLP 只有 3 个 kernel，LLM 一层有几十个，累计下来收益更大。
三个要点记住就够了：
- warmup 不能省，3 次足够
- 输入
x的地址不能变，只能就地copy_()或fill_()更新值 - replay 后直接读
y，它的地址也是 capture 时固定的
4. 套到 LLM 上：第一个坑，KV cache
理解了基本用法，来套到 LLM 推理上。
LLM 的 decode 阶段每步输入是当前 token（shape [batch, 1]），同时维护一个 KV cache 存历史计算结果。问题就出在这里——标准 transformers 的 KV cache 是一个 tuple，每步都会增长：
# step 1: past_key_values 里每个 tensor shape = [batch, heads, 1, head_dim]
# step 2: past_key_values 里每个 tensor shape = [batch, heads, 2, head_dim]
# step 3: [batch, heads, 3, head_dim]
# ...
每步 shape 不一样，就意味着每步 kernel 的参数不一样，图的拓扑在变。CUDA Graph 要求拓扑固定，这里直接冲突。
解法：StaticCache
transformers 4.38 之后提供了 StaticCache——提前分配一块固定大小的显存，KV cache 永远写在这块固定 buffer 里，shape 不变，每步只是更新写入位置：
from transformers import StaticCache
# 提前分配好，最多支持 512 个 token 的 KV cache
static_cache = StaticCache(
config=model.config,
max_batch_size=1,
max_cache_len=512,
device="cuda",
dtype=torch.float16,
)
有了 StaticCache，KV cache 的 shape 整个推理过程都不变，CUDA Graph 的”拓扑固定”约束就满足了。
5. 完整可运行代码
把所有东西拼在一起：
from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache
import torch
# --- 加载模型 ---
model_name = "Qwen/Qwen2.5-1.5B"
model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype=torch.float16
).cuda().eval()
tokenizer = AutoTokenizer.from_pretrained(model_name)
MAX_SEQ_LEN = 512
BATCH_SIZE = 1
# --- 分配固定地址的输入 buffer（地址整个生命周期不变）---
input_ids = torch.zeros(BATCH_SIZE, 1, dtype=torch.long, device="cuda")
position_ids = torch.zeros(BATCH_SIZE, 1, dtype=torch.long, device="cuda")
cache_pos = torch.zeros(1, dtype=torch.long, device="cuda")
# --- 分配 StaticCache ---
static_cache = StaticCache(
config=model.config,
max_batch_size=BATCH_SIZE,
max_cache_len=MAX_SEQ_LEN,
device="cuda",
dtype=torch.float16,
)
# --- Step 1: warmup ---
for _ in range(3):
with torch.no_grad():
model(
input_ids=input_ids,
position_ids=position_ids,
cache_position=cache_pos,
past_key_values=static_cache,
use_cache=True,
)
# --- Step 2: capture ---
static_cache.reset() # 重置 cache 状态（warmup 写进去的内容清掉）
g = torch.cuda.CUDAGraph()
with torch.no_grad():
with torch.cuda.graph(g):
logits = model(
input_ids=input_ids,
position_ids=position_ids,
cache_position=cache_pos,
past_key_values=static_cache,
use_cache=True,
).logits # shape: [batch, 1, vocab_size]
# --- Step 3: decode 循环 ---
def decode_one_token(token: int, pos: int) -> int:
"""一步 decode，返回下一个 token id"""
input_ids.fill_(token) # 就地更新，地址不变
position_ids.fill_(pos)
cache_pos.fill_(pos)
g.replay() # replay，不重新 launch kernel
return logits[0, 0].argmax().item()
def generate(prompt: str, max_new_tokens: int = 100) -> str:
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
input_token_ids = inputs["input_ids"][0].tolist()
# prefill 阶段：正常跑，不用 CUDA Graph
# （prefill 序列长度每次不同，不适合 CG）
with torch.no_grad():
static_cache.reset()
prefill_out = model(
input_ids=inputs["input_ids"],
cache_position=torch.arange(len(input_token_ids), device="cuda"),
past_key_values=static_cache,
use_cache=True,
)
# 取 prefill 最后一步的 next token
next_token = prefill_out.logits[0, -1].argmax().item()
generated = [next_token]
pos = len(input_token_ids)
# decode 阶段：每步 replay CG，飞快
for _ in range(max_new_tokens - 1):
if next_token == tokenizer.eos_token_id:
break
next_token = decode_one_token(next_token, pos)
generated.append(next_token)
pos += 1
return tokenizer.decode(generated, skip_special_tokens=True)
print(generate("解释一下推测解码的原理"))
这段代码和最开始那段的主要区别：
-
StaticCache替换了动态增长的 past_key_values - 单独的固定 buffer（
input_ids、position_ids、cache_pos），每步fill_()更新 - prefill 走正常 eager，decode 走 CUDA Graph replay
6. 第二个坑：batch size 变了怎么办
上面的代码是 batch size = 1 的版本。实际场景里，你可能想同时处理多个请求（batch inference），而且每轮 decode batch size 可能在变（有请求完成就减少，有新请求进来就增加）。
CUDA Graph 捕获的是固定 batch size 的图，batch size 变了图就不能用了。
解法：预先捕获几个常用 batch size 的图，放进一个字典，运行时按实际 batch size 选对应的图。
# 预先 capture 这几个 batch size 的图
CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32]
graphs = {} # batch_size → CUDAGraph
logits_buffers = {} # batch_size → logits tensor（capture 时固定的地址）
for bs in CAPTURE_BATCH_SIZES:
# 每个 batch size 单独分配一套 buffer
_input_ids = torch.zeros(bs, 1, dtype=torch.long, device="cuda")
_position_ids = torch.zeros(bs, 1, dtype=torch.long, device="cuda")
_cache_pos = torch.zeros(1, dtype=torch.long, device="cuda")
_cache = StaticCache(config=model.config, max_batch_size=bs,
max_cache_len=MAX_SEQ_LEN, device="cuda",
dtype=torch.float16)
# warmup
for _ in range(3):
with torch.no_grad():
model(input_ids=_input_ids, position_ids=_position_ids,
cache_position=_cache_pos, past_key_values=_cache,
use_cache=True)
# capture
_cache.reset()
g = torch.cuda.CUDAGraph()
with torch.no_grad():
with torch.cuda.graph(g):
_logits = model(input_ids=_input_ids, position_ids=_position_ids,
cache_position=_cache_pos, past_key_values=_cache,
use_cache=True).logits
graphs[bs] = {
"graph": g,
"input_ids": _input_ids,
"position_ids": _position_ids,
"cache_pos": _cache_pos,
"logits": _logits,
}
def decode_batch(tokens: list[int], pos: int) -> list[int]:
bs = len(tokens)
# 找最小的已捕获 batch size >= bs，padding 多余位置
target_bs = next(b for b in CAPTURE_BATCH_SIZES if b >= bs)
entry = graphs[target_bs]
entry["input_ids"][:bs].fill_(0) # 先清空
for i, t in enumerate(tokens):
entry["input_ids"][i].fill_(t) # 就地更新有效 token
entry["position_ids"][:bs].fill_(pos)
entry["cache_pos"].fill_(pos)
entry["graph"].replay()
return entry["logits"][:bs, 0].argmax(dim=-1).tolist() # 只取有效部分
这就是 vLLM 的 batch size bucketing 思路，只是简化版。多几张图，覆盖了实际场景里各种 batch size，运行时零额外开销。
7. 第三个坑：MoE 模型
如果你跑的是 Mixtral、DeepSeek-MoE、Qwen-MoE 这类 MoE 架构的模型，还有一个额外的问题。
Dense 模型每层所有 token 都过同样的权重矩阵，计算路径固定。MoE 不一样——router 会把 token 分配给不同的 expert，每个 expert 这一步收到多少 token 是数据依赖的，每次都可能不同：
batch=8, top_k=2, 8 个 expert → 理论上每个 expert 平均 2 个 token
实际可能：expert0 收 5 个，expert1 收 0 个，expert2 收 3 个...
expert 收到不同数量的 token，送进 expert FFN 的 tensor shape 每次都在变，CUDA Graph 直接挂。这也是为什么在一些 GitHub issue 里会看到这样的说法：
This is due to the tensor in MoE layer has dynamic shape, which is not applicable for CUDA graph. You can partially capture CUDA graph for Attention layer.
先说好消息：batch=1 的 decode 不受影响
batch=1，每步只有 1 个 token，top_k=2 的情况下恰好 2 个 expert 各收到 1 个 token，shape 固定。上面第 5 节的代码直接拿去跑 MoE 模型，decode 阶段完全没问题。
batch>1 才是麻烦的地方，下面说两种解法。
解法一：只 capture attention，MoE FFN 跑 eager（partial capture）
这是最简单、改动最小的方案，也是 Megatron-LM 对 MoE 模型的推荐做法：每个 transformer layer 里，attention 部分 capture 成 CUDA Graph，MoE FFN 部分继续走 eager。
一个 transformer layer 的结构是：
Input
→ [Attention] ← shape 固定（batch × 1 × dim），可以 capture
→ [MoE FFN] ← expert 分配动态，不 capture
→ Output
把 attention 单独包成一个 CG，MoE FFN 保持 eager，每层的调用变成：
class MoETransformerLayerWithPartialCG:
def __init__(self, attention, moe_ffn, hidden_size, batch_size):
self.moe_ffn = moe_ffn
# 给 attention 分配固定地址的输入输出 buffer
self.attn_input = torch.zeros(batch_size, 1, hidden_size, device="cuda", dtype=torch.float16)
self.attn_output = torch.zeros(batch_size, 1, hidden_size, device="cuda", dtype=torch.float16)
# warmup
for _ in range(3):
with torch.no_grad():
_ = attention(self.attn_input)
# 只 capture attention
self.attn_graph = torch.cuda.CUDAGraph()
with torch.no_grad():
with torch.cuda.graph(self.attn_graph):
self.attn_output = attention(self.attn_input)
def forward(self, hidden_states):
# attention：replay CG（无 kernel launch overhead）
self.attn_input.copy_(hidden_states)
self.attn_graph.replay()
attn_out = self.attn_output.clone()
# MoE FFN：正常 eager（shape 动态，无法 capture）
ffn_out = self.moe_ffn(attn_out)
return ffn_out
attention 部分的 kernel launch overhead 消掉了，MoE FFN 部分没动。相比整层都走 eager，能拿到一半左右的 CG 收益——对于 Mixtral 这类模型，attention 占单层耗时的 40-50%，效果还是明显的。
Megatron-LM 里通过 --cuda-graph-modules=attn 参数开启这个模式，正是这个思路。
解法二：给 MoE FFN 加 expert capacity，整层都 capture
如果想把 MoE FFN 也 capture 进去，需要消掉 shape 动态的根源：提前规定每个 expert 最多接受 capacity 个 token，无论 router 实际分了多少，tensor shape 始终是 [capacity, hidden_dim]——分配超过 capacity 的 token 丢弃，不足的用零 padding：
import math
def moe_forward_with_capacity(hidden_states, router_logits, experts,
top_k, capacity_factor=1.25):
batch, seq, dim = hidden_states.shape
hidden_flat = hidden_states.view(-1, dim) # [num_tokens, dim]
num_tokens = hidden_flat.shape[0]
num_experts = len(experts)
routing_weights = torch.softmax(router_logits, dim=-1)
topk_weights, topk_ids = routing_weights.topk(top_k, dim=-1)
# capacity 提前算死，replay 时这个值不会变
capacity = math.ceil(num_tokens * top_k / num_experts * capacity_factor)
# 固定 shape 的 buffer：[num_experts, capacity, dim]
expert_inputs = torch.zeros(num_experts, capacity, dim,
device=hidden_states.device, dtype=hidden_states.dtype)
expert_weights = torch.zeros(num_experts, capacity,
device=hidden_states.device, dtype=hidden_states.dtype)
# 把 token 填进对应 expert 的 slot（超出 capacity 就 drop）
for e in range(num_experts):
assigned = (topk_ids == e).any(dim=-1) # [num_tokens]
token_idx = assigned.nonzero(as_tuple=True)[0]
n = min(len(token_idx), capacity)
if n > 0:
w_pos = (topk_ids[token_idx[:n]] == e).float().argmax(dim=-1)
expert_inputs[e, :n] = hidden_flat[token_idx[:n]]
expert_weights[e, :n] = topk_weights[token_idx[:n], w_pos]
# 每个 expert 的输入永远是 [capacity, dim]，shape 固定，CG 可以录制
outputs = torch.stack([experts[e](expert_inputs[e]) for e in range(num_experts)])
# 把结果散回原位置
result = torch.zeros_like(hidden_flat)
for e in range(num_experts):
assigned = (topk_ids == e).any(dim=-1)
token_idx = assigned.nonzero(as_tuple=True)[0]
n = min(len(token_idx), capacity)
if n > 0:
result[token_idx[:n]] += expert_weights[e, :n, None] * outputs[e, :n]
return result.view(batch, seq, dim)
capacity_factor=1.25 给了 25% 余量，生产环境里 token drop 概率很低。代价是每个 expert 都要跑满 capacity 个位置，padding 的部分算了白算（浪费约 5-10% 算力）。
Megatron-LM 的参数是 --moe-expert-capacity-factor 1.25 --moe-pad-expert-input-to-capacity，开了这两个之后 MoE 层也能 capture。
两种方案怎么选：
- 想最快改完上线，改动最小 → 用 partial capture（解法一）
- 想把 MoE FFN 的 launch overhead 也消掉 → 用 expert capacity（解法二），但要改 MoE 层实现
- 不想自己改，直接用 vLLM →
FusedMoEkernel 内置了 capacity dispatch，CUDA Graph 开箱即用
9. prefill 为什么不加 CUDA Graph
你可能注意到上面的代码里，prefill 阶段始终走的是 eager mode，没有 CUDA Graph。
原因很简单：prefill 的输入序列长度每次都不一样。你一次处理 “你好” 和 “解释一下量子纠缠的基本原理及其在量子计算中的应用”，这两个序列的长度差了几十倍，对应的 kernel 参数完全不同，图拓扑每次都变，CG 捕获不了。
Decode 阶段每步输入永远是一个 token（长度固定为 1），天然适合 CUDA Graph，这也是为什么 vLLM 这类系统里 CG 只作用于 decode。
10. 加了 CG 之后能快多少
以 Qwen2.5-1.5B 为例，batch size = 1，decode 阶段：
| 方式 | 单步延迟 | 吞吐（tokens/s） |
|---|---|---|
| Eager (standard generate) | ~8 ms | ~125 |
| + CUDA Graph | ~5 ms | ~200 |
| 提升 | ~37% | ~60% |
小模型收益更大（kernel 数量相对 kernel 计算量的比例更高，launch overhead 占比更大）。大模型（如 70B）每个 kernel 跑的时间本来就长，launch overhead 占比低，CG 的相对收益会小一些，但绝对值（省下的毫秒数）还是有的。
11. 踩坑备忘
坑 1：capture 里不能有 Python 控制流依赖 GPU 结果
# ❌ 这会出问题：capture 时走了某个分支，replay 时永远走那个分支
with torch.cuda.graph(g):
out = model(x)
if out.max() > 0.5: # GPU 结果 → Python if → 分支被固化
out = postprocess(out)
把判断逻辑移到 graph 外面，或者改成 torch.where（GPU tensor 操作，不走 Python 分支）。
坑 2：capture 时不能有随机操作
# ❌ capture 时 dropout 用了随机数，replay 时随机种子不更新
with torch.cuda.graph(g):
out = model_with_dropout(x)
推理时 model.eval() 关掉 dropout，这个坑就不存在了。记得 eval。
坑 3：capture 里不能有 CPU-GPU 同步
tensor.item()、.cpu()、print(tensor) 这些都会触发 CPU-GPU 同步，capture 时如果有这些操作会报错或者行为异常。把这类操作移到 graph 外面。
坑 4：忘了 warmup 就 capture
第一次 capture 之前必须先把 kernel 跑热（3 次以上）。直接 capture 可能录进去的是未完全初始化的 kernel 状态，replay 时结果错误或崩溃。
坑 5：StaticCache 忘了 reset
capture 前必须调 static_cache.reset()，否则 warmup 写进去的 KV 数据还在里面，capture 的初始状态不是”空 cache”，之后的推理结果就乱了。
12. 一句话总结
给自己的 transformers 推理代码加 CUDA Graph，核心改动只有两处：一是把 past_key_values 换成 StaticCache（解决 KV cache shape 动态变化的问题），二是提前分配好固定地址的输入 buffer（解决 tensor 地址不能变的约束）。三步走：warmup → capture → 每步 replay。prefill 不动，只加速 decode。
如果这篇文章涉及的 LLM 推理效率优化你想系统深入，可以看看我之前出版的《动手学 AutoML：从 NAS 到大语言模型优化实战》，书里有专章讲 LLM 推理效率，涵盖量化、kernel 优化和推理框架的工程权衡，和本文讨论的 CUDA 层优化有直接关联。
