nano-vllm-omni · 提示 ← P16
#26 · Track E · 启动与 GPU · 书 p.109

P17 · CUDA graph bucket runner（CPU 模拟）

decode 的 kernel launch 开销靠 CUDA graph 消除；graph 只认静态 buffer——pad、清零、replay、clone，以及什么时候必须退回 eager。

⏱ 40 min 难度 易 实现 mine/p17_graph_runner.py 验收 ./check.sh p17 提示 A-26 · 参考解 C-26 solutions/p17_graph_runner.py

题面速览

启动时：为每个 batch bucket（1,2,4,8,...）在同一块静态 buffer 的前 b 行上录一张图
运行时（bs 个请求）：
  b = 最小的 ≥ bs 的 bucket；没有 → eager
  static_in[:bs] = x; static_in[bs:b] = 0     # 上一次更大 batch 留下的脏行
  graph[b].replay()                           # 结果写在 static_out[:b]
  return static_out[:bs].clone()              # 下一次 replay 会覆盖 static_out！

验收要点：每个 batch size 结果与 eager 一致 · bucket 选择与两种 eager 回退 · 返回值不别名静态 buffer · padding 行被清零 · 运行时从不重新 capture。

直觉与类比

固定尺寸的烤盘

面包房只有 1 格、3 格、6 格三种烤盘，烤箱程序都提前调好了（capture），开工后不再重新编程。来了 2 个面团：用 3 格的烤盘（向上取整），空着的那一格要先把上一炉的面包屑扫干净（清零 padding 行）。烤好之后必须把面包从烤盘上拿下来装袋再交给客人（clone）——如果直接把“烤盘上那几格”指给客人，下一炉一开烤，客人手里的面包就变成了别人的。来了 7 个面团？最大烤盘只有 6 格 → 手工一个个烤（eager）。

卡拉 OK 伴奏带

CUDA graph 像一盘录好的伴奏带：按下播放就一模一样地再放一遍。你可以换一个唱歌的人（把新数据拷进同一个话筒 = 静态 buffer），但不能在播放时让乐队改编。客人要求“用我自己的随机节奏”（自带 torch.Generator），伴奏带不会报错，只会照旧播放录制时的节奏——要求被悄悄忽略。所以这种请求必须请现场乐队（force_eager）。

框架视角 · 一次 replay 调用（bs=2、bucket 3、hidden=2）

# 张量 形状 发生了什么
1 x [2, 2] 调用方的新 batch（Python 参数，graph 看不到它）
2 选 bucket 标量 3 最小的 ≥ 2 的 bucket；没有 / force_eager → 直接 fn(x)
3 static_in[:2] ← x [2, 2] 写入 [6, 2] 原地拷贝进录制时的那块内存
4 static_in[2:3] ← 0 [1, 2] padding 行清零（之前是上一次 bs=5 留下的 5）；第 3..5 行在 bucket 外，不碰
5 graphs[3].replay() 读 static_in[:3]，写 static_out[:3] 输出 5, 5, −1（padding 行 = fn(0)）
6 static_out[:2].clone() [2, 2] 只取真实行，并且拷贝出来，和静态 buffer 脱钩

易错点

逐步图：例 1 一串调用

▶ fn = x*2+1，buckets=[16, 1, 4] → [1, 4, 16]，hidden=3；输入 x 全部填成 float(bs)
1. run(bs=10)

bucket 16 · replays 1 · eager 0 · last_bucket 16

static_in[:6,0]：101010101010　static_out[:6,0]：212121212121

2. run(bs=3)

bucket 4 · replays 2 · last_bucket 4

static_in[:6,0]：33301010　static_out[:6,0]：77712121

第 3 行被清零（上一次的 10），输出 fn(0)=1；第 4、5 行在 bucket 之外，没人碰。

3. run(bs=17)

没有 ≥17 的 bucket → eager：replays 2 · eager 1 · last_bucket 仍是 4 · 静态 buffer 不变

4. run(bs=4)

bucket 4 · replays 3 · last_bucket 4

static_in：44441010　static_out：99992121

5. run(bs=1, force_eager=True)

eager：replays 3 · eager 2 · last_bucket 4 · buffer 不变

6. run(bs=1)

bucket 1 · replays 4 · last_bucket 1

static_in：14441010　static_out：39992121

7. 为什么一定要 clone
y1 = r.run(x1)     # 若直接 return static_out[:bs]（一个 view）
y2 = r.run(x2)     # replay 把 fn(x2) 写进 static_out
y1                 # → 现在等于 fn(x2)！y1 被“悄悄改写”

返回 clone 后，y1 保持 fn(x1)。

三级提示

L1I/O 示例与关键问题I/O 示例 + 关键问题 + 细节（轻度剧透）

例 1、例 2 见逐步图。

关键问题

  1. CpuGraph.replay() 不接受参数。新的输入是怎么“进入”这张图的？输出又在哪里拿？
  2. 为什么所有 bucket 的图都建在同一块 static buffer 的前 b 行，而不是每个 bucket 各自分配？
  3. bs=3、bucket=4 时第 3 行里可能是什么？不清零会导致什么后果？
  4. static_in[bs:] 全部清零 vs 只清 static_in[bs:b]：哪个是必要的？多清有什么代价？
  5. force_eager 的典型场景是请求带了自己的 torch.Generator。为什么 replay 会“悄悄忽略”它而不是报错？
  6. last_bucket 在走 eager 时要不要更新？
  7. run 发现某个 bucket 还没 capture，能不能现场 capture？为什么 spec 说 never？

细节

  • CUDA graph 记录的是 kernel 以及它们读写的设备地址；replay 时地址不变，所以必须把新数据 copy_ 进同一块内存。
  • static_in[:b] 是共享存储的 view；所有图共享同一块内存，所以上一次大 batch 的数据会残留在后面的行里。
  • “largest bucket first”：CPU 模拟里顺序不影响结果，但请按 spec 来。
  • bucket 选择：最小的 b ≥ bs；buckets 在 __init__ 里已排好序，线性找即可（也可 bisect）。
  • stats：replays 与 eager 是计数器；last_bucket 只在 replay 时记录。
  • 返回值只含前 bs 行（padding 行的结果是“对 0 做计算”的垃圾）。
L2思路与常见坑思路拆分与常见坑（中度剧透）
  • __init__：max_bs = buckets[-1]；分配 static_in、static_out [max_bs, hidden]，用传入的 dtype。初始化成 0 比 empty 更安全（capture 时就会跑一次 fn）。
  • capture：遍历 sorted(buckets, reverse=True)，用 static_in[:b] / static_out[:b] 构造 CpuGraph，.capture()，存进 self.graphs[b]。
  • run 的决策顺序：① 算 bs；② force_eager 或找不到 bucket（或该 bucket 不在 graphs 里）→ eager：计数 +1，return self.fn(x)；③ 否则 replay 分支。eager 分支不碰静态 buffer。
  • replay 分支四步：copy 输入 → 清零 [bs, b) → replay → 更新 stats → 返回 clone 的前 bs 行。清零在 replay 之前；clone 在下一次 replay 之前（即返回前）。
  • self.static_in[:bs] = x 其实是原地写（等价于 copy_），没问题；但 self.static_in = x 是错的。
L3详解：答案详解：答案、逐步追踪、骨架提示（重度剧透）
关键问题答案
  1. 输入：先 copy_ 进 static_in；replay 读 static_in、写 static_out；输出从 static_out 读。
  2. 省内存，而且 run 时只需一份“写入位置”：无论选哪个 bucket，都写 static_in 的前 bs 行。
  3. 可能是上一次更大 batch 的真实数据（例 1 中的 10）。本题不清零只让 padding 输出变垃圾；真实模型中可能 NaN 扩散、影响 MoE 路由与 expert 负载统计、影响整 batch 归约的算子。
  4. 只需清 [bs, b)：这是这张图会读的范围；b 之后的行这张图根本不读，清它们只是多花带宽。
  5. replay 只是重放 capture 时录下的 kernel，不重新执行 Python；per-row generator 的逻辑根本不会被执行，所以不报错，只会结果不可复现/不符合 seed。
  6. 不更新；它的语义是“最近一次 replay 用的 bucket”。
  7. 不能：capture 很慢（要 warmup、同步），在 CUDA 上还会和正在运行的 stream / 内存池冲突；运行时捕获会造成延迟尖刺。

自检清单：① 连续两次 run，第一次的返回值不变；② 先跑一次大 batch 再跑小 batch，检查 static_in 的 [bs, b) 行全 0；③ capture 之后 self.graphs 里每个对象的 id() 在多次 run 后不变；④ bs > max bucket 和 force_eager 时 static buffer 完全不被修改。

面试官追问 & 对照源码

对照源码：worker/gpu_ar_model_runner.py:634-690 · worker/gpu_model_runner.py:243-272 · worker/base.py:140-147