课程 06 · 执行优化

把整个 decode 录成一张图

decode 每步只算一个 token，却要启动几百个 kernel。CUDA Graph 的做法是：把这一整套启动序列录下来，之后一次回放。

约 45 分钟 源码：model_runner.py L195–257 · layers/sampler.py 前置：04

1先算一笔启动开销

Qwen3-0.6B 有 28 层，每层十来个算子（线性、RMSNorm、RoPE、attention、silu…）。一次 decode 前向大约要启动 300–500 个 CUDA kernel。

• 每个 kernel 的 CPU 侧启动开销约 3–8 µs（Python + PyTorch dispatcher + driver）。
• 400 个 kernel × 5 µs ≈ 2 ms。纯启动，还没算 GPU 真正干活的时间。

而 0.6B 模型在 4070 Ti 上 decode 一步的计算可能只有 1–2 ms。也就是说，一半以上的时间花在「告诉 GPU 要干什么」上，而不是「GPU 在干」。大模型更夸张：kernel 更多，但单个 kernel 的利用率也不高。

2录图：固定 shape，固定地址

model_runner.py · L222–248max_bs = min(self.config.max_num_seqs, 512)
self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
self.graphs = {}
self.graph_pool = None

for bs in reversed(self.graph_bs):                 # 从大到小录
    graph = torch.cuda.CUDAGraph()
    set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
    outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
    with torch.cuda.graph(graph, self.graph_pool):
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
    if self.graph_pool is None:
        self.graph_pool = graph.pool()                 # 复用显存池
    self.graphs[bs] = graph
    torch.cuda.synchronize()

几个必须理解的设计：

1. 只录固定的 batch size：[1, 2, 4, 8] + 16 的倍数直到 512。CUDA Graph 要求 shape 和显存地址完全固定，所以只能为每个可能的 bs 录一张图。
2. 为什么先 warmup 再 capture：第一次跑会触发显存分配、cuBLAS 算法选择、kernel 编译缓存。这些必须在录制前完成，否则会被录进图里（或者直接报错）。
3. 为什么 reversed：先录最大的 bs，显存池一次分配到位；后面的小图复用同一个 pool，不会因为碎片而失败。
4. slot_mapping 预填 -1：实际 bs 小于录制 bs 时，多出来的槽位保持 -1，Triton kernel 第一行就 if slot == -1: return——padding 不产生任何写入。

3回放：把输入拷进固定 buffer

model_runner.py · L195–212def run_model(self, input_ids, positions, is_prefill):
    if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
        return self.model.compute_logits(self.model(input_ids, positions))   # 正常路径
    bs = input_ids.size(0)
    graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]   # 找最小够用的图
    graph_vars = self.graph_vars
    graph_vars["input_ids"][:bs] = input_ids          # 拷进固定 buffer
    graph_vars["positions"][:bs] = positions
    graph_vars["slot_mapping"].fill_(-1)
    graph_vars["slot_mapping"][:bs] = context.slot_mapping
    graph_vars["context_lens"].zero_()
    graph_vars["context_lens"][:bs] = context.context_lens
    graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
    graph.replay()                                     # 一次回放
    return self.model.compute_logits(graph_vars["outputs"][:bs])

注意最后一行：compute_logits 在图外面。因为 logits 计算的 shape 随 batch 变化，而且后面紧跟采样，没必要录进去。

回放的代价是几次 copy_：把新输入搬进固定的 buffer。相比省下的几百次 kernel launch，这点拷贝可以忽略。

路径	条件	为什么
eager	prefill	shape 每次都不同，录图无意义
eager	enforce_eager=True	调试 / 兼容性 / 显存紧张
eager	bs > 512	超过录制上限
graph	其余 decode	shape 稳定，启动开销占比最高

4torch.compile：只编译采样

layers/sampler.py · L7–12@torch.compile
def forward(self, logits, temperatures):
    logits = logits.float().div_(temperatures.unsqueeze(dim=1))
    probs = torch.softmax(logits, dim=-1)
    sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)
    return sample_tokens

为什么只编译这一步？三个原因：

• shape 稳定：[bs, vocab]，bs 只有几十种取值，编译缓存命中率高。
• 算子链长：div → softmax → exp → div → argmax，每一步都是独立的 kernel，编译后可以融合。
• 不碰模型：整模型 torch.compile 会遇到 attention 的自定义 kernel、动态 shape、CUDA Graph 冲突，风险远大于收益（vLLM 走的是另一条路，见第 10 课）。

5enforce_eager 该不该开

	enforce_eager=True	False（默认）
启动速度	快（不录图）	慢（要为每个 bs 录一遍）
显存占用	低	高（图池 + 静态 buffer）
decode 速度	慢（每步几百次 launch）	快
调试	可以用 pdb / print	几乎无法调试（图内不可断点）

所以 example.py 用 enforce_eager=True 是为了让第一次跑的人别被录图卡住；bench.py 用 False 是为了跑出真实性能。面试里能说清「为什么默认开着、什么时候关掉」，比记住参数名有用得多。

6练习 A · 估算收益

假设 400 个 kernel、每个启动 5 µs，decode 一步的计算时间 1.5 ms。开图后能省多少？加速比是多少？

展开：答案

启动开销 400 × 5 µs = 2.0 ms。开图前一步约 1.5 + 2.0 = 3.5 ms；开图后只剩回放开销（约 0.05 ms）+ 计算 = 1.55 ms。

加速比约 2.3×。这就是为什么「decode 开图」几乎是所有推理框架的标配，而 prefill 不需要——prefill 的计算量大，启动开销被掩盖了。

7练习 B · 闭卷测验

8练习 C · 代码寻宝

1. 为什么 graph_bs 是 [1,2,4,8] + range(16, ..., 16) 这种非均匀分布？

小 batch 是延迟敏感场景的高频取值（1、2、4、8 很常见），必须精确匹配；大 batch 之间差异不大，用 16 的步长向上取整就够了。next(x for x in graph_bs if x >= bs) 总是找最小够用的图，最多浪费十几个空槽。

2. 如果 decode 时 batch 突然从 4 涨到 5，会发生什么？

会选 bs=8 的那张图，把 5 条序列的输入拷进前 5 行，剩下 3 行保持 -1 / 0 padding。Triton 的 KV 写入 kernel 会因为 slot == -1 跳过它们；注意力的 varlen 不适用（这是 decode 的 with_kvcache 路径），padding 的 context_lens 为 0，不会影响有效行。

3. 为什么 compute_logits 不放进图里？

它的输出 shape 是 [bs, vocab]，且 bs 会变；更重要的是它后面紧跟采样，采样在 CPU 侧做（.tolist()），把图断开反而更灵活。图内只放「shape 固定、纯 GPU」的部分。

9面试复述卡

「CUDA Graph 为什么能加速 decode，prefill 为什么不用？」

展开：30 秒版本

「decode 每步的计算量小，但 kernel 数量不变，所以 CPU 侧的启动开销占了很大比例。CUDA Graph 把固定的静态图和显存地址录下来，之后一次 replay 就把几百次 launch 变成一次。prefill 的 shape 每次都不同（prompt 长度不定），录不了图；而且 prefill 计算量大，启动开销本来就被掩盖了。」