数据截至 2026-08-27(上游 commit f25c580af159)

04 · 模型执行与 CUDA Graph

这一章讲什么： SchedulerOutput 到达 GPU 之后发生了什么。GPUModelRunner 怎么把「谁、几个 token、哪些块」变成一次 forward，以及 CUDA graph 这个让 decode 步 CPU 开销趋近于零的机制是怎么捕获和派发的。

1. 它要解决的小问题

调度器给的只是元数据。要在 GPU 上跑一步，得回答三串问题：

输入在哪？ batch 里每条的 token id、位置、各自该从哪些 KV 块读注意力——而且每个 step 都变。
怎么跑得快？ decode 一步可能只有几百个 token，forward 本身几毫秒；如果每个 CUDA kernel 都要 CPU 逐个发射，CPU 发射时间比 GPU 计算还长。
结果怎么出来？ 每条请求只采它最新位置的 logits，还要过温度、top-k/top-p、惩罚项。

2. 思路：持久化 batch + 展平输入 + 图重放

三个关键决策：

持久化 batch（persistent batch）。 Worker 侧维护一个常驻的 input_batch；调度器每步只发 diff（新请求全量、老请求增量、结束请求 id），runner 增量更新，而不是每步重建整个 batch 张量。
全部展平成一维。 不同长度的请求不 padding 成矩阵，而是首尾相接拼成 [total_num_tokens] 的一维序列；每条算几个 token 由 num_scheduled_tokens 切分，注意力靠「块表 + 序列长度」元数据区分边界。
CPU 发射开销用 CUDA graph 整段消掉。 把「一组固定形状的 kernel 发射序列」录下来（capture），之后同一个形状整条重放（replay），CPU 只发一次启动命令。形状（batch 内 token 数）每个 step 都不同，所以预先给一组离散尺寸各录一张图，运行时把实际尺寸向上取整到最近的已录尺寸。

图示：一步的数据流

SchedulerOutput(元数据)
   │
   ▼
① _update_states: 持久化 batch 增删改(新请求进来、结束的移除、块表追加)
   │
   ▼
② _prepare_inputs: 展平 token ids / positions,算 logits_indices,
   为每个注意力组 build 元数据(块表、seq lens)
   │
   ▼
③ cudagraph 派发: 按 token 数选模式与 padded 尺寸
   │
   ▼
④ forward: set_forward_context(元数据经 forward context 传给各层)
   │
   ▼
⑤ compute_logits(只取每条最后位置) → Sampler.sample → 回本步 token

3. 原理演示：向上取整的图派发

这是 CudagraphDispatcher 核心想法的最小版：

# 示意,非源码
CAPTURE_SIZES = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]  # 预捕获的尺寸集合

class GraphPool:
    def __init__(self):
        self.graphs = {size: capture_graph_for(size) for size in CAPTURE_SIZES}

    def padded_size(self, num_tokens):
        for s in CAPTURE_SIZES:          # 实际实现是预计算的 O(1) 查表
            if num_tokens <= s:
                return s
        return None                       # 超过最大图:退回逐 kernel 执行

def run_step(graph_pool, tokens):
    size = graph_pool.padded_size(len(tokens))
    if size is None:
        return eager_forward(tokens)
    padded = pad(tokens, size)            # 补垃圾 token,算完丢弃
    graph_pool.graphs[size].replay()      # CPU 只发一条 replay 命令
    return read_outputs(padded)

重点：用少量显存（每张图一份输入/输出缓冲）和一点浪费的计算（padding），换掉整个 CPU 发射链。decode 批量越小，这个交换越值。

4. 真实实现

4.1 `GPUModelRunner.execute_model`

在 vllm/v1/worker/gpu_model_runner.py:4238。顺着它的主干，对应上面五步：

① _update_states(scheduler_output)（:1243）：新请求登记进持久化 batch、结束的移除、req_to_new_blocks 追加进块表。
② _prepare_inputs（:2016）：构造展平的 input_ids/positions、每条请求的采样位置 logits_indices、投机解码元数据；注意力元数据由 _build_attention_metadata（:2352）按 KV cache group 逐组建。
③ CUDA graph 派发：dispatch_cudagraph 调用点在 :4104（内部转 CudagraphDispatcher.dispatch）。
④ forward：包在 set_forward_context(...) 里（:4542），注意力元数据、图模式、块映射（slot_mapping）都经这个 context 传到每一层；然后 self._model_forward(input_ids=..., positions=..., ...)（:4559）。
⑤ 只有流水并行最后一级才算 logits：sample_hidden_states = hidden_states[logits_indices] → self.model.compute_logits(...) → 采样（:4588-4594 区域）。

4.2 块表：`BlockTable`

vllm/v1/worker/block_table.py:57。它就是第一章「页表」的 GPU 侧实物：一个 [max_num_reqs, max_num_blocks] 的常驻张量，行是请求、列是逻辑块、值是物理块 id。

append_row（:157）：调度器发来新块，追加到该请求行的已用区；
commit_block_table（:231）：CPU 侧累积的改动一次性 H2D 拷到设备张量——每步一次小拷贝，而不是每个 kernel 各自同步；
注意力后端拿整行块表去做 gather（如 FlashAttention 后端元数据字段 block_table，vllm/v1/attention/backends/flash_attn.py:262）。

4.3 CUDA graph 的三档模式

CUDAGraphMode（vllm/config/compilation.py:53）：

模式	录进图里的范围	适用
`NONE`	不用图	调试、不支持的后端
`PIECEWISE`	除注意力外的分段（注意力算子保持 eager）	配合 torch.compile 的 splitting_ops，通用
`FULL`	整个 forward（含注意力）	支持「图内注意力」的后端 + 形状满足时最快

另有两个组合档：FULL_DECODE_ONLY（decode 用 FULL、prefill 走 eager）和 FULL_AND_PIECEWISE（decode FULL、prefill PIECEWISE）——它们是「配置态」，运行时被解析成具体模式（decode_mode() / mixed_mode()，同文件 :66-70）。

4.4 派发：`CudagraphDispatcher.dispatch`

在 vllm/v1/cudagraph_dispatcher.py:235。决策链：

图未初始化 / 模式为 NONE / token 数超过最大捕获尺寸 → 直接 NONE 退回 eager（:271-281）；
uniform_decode（整批都是等长 decode，含投机 token 数）是单独的 dispatch key 维度——等长批可以用更特化的图；
尺寸向上取整查表：_compute_bs_to_padded_graph_size（:72）预计算 batch size → padded graph size 的完整映射数组，运行时是 O(1) 数组下标，不做二分。

信任链值得注意：dispatcher 把 (runtime_mode, BatchDescriptor) 写进 forward context；CUDAGraphWrapper（vllm/compilation/cuda_graph.py:145）收到后「盲目信任」（docstring 原话 "blindly trust"）：模式匹配就 replay（没捕获过就现场捕获），不匹配就直接调底层 runnable。决策全部集中在 dispatcher 一处，wrapper 不做判断——单一事实来源，避免各层各猜一套。

4.5 捕获时机

不是懒捕获为主：启动时 compile_or_warm_up_model（vllm/v1/worker/gpu_worker.py:749）阶段就对 cudagraph_capture_sizes 里的每个尺寸录好图（_capture_cudagraphs，vllm/v1/worker/gpu_model_runner.py:7045）；runner 持有的捕获尺寸集在 :797-806 排序备用。代价是启动慢，换来的是第一张生产 token 就在图上跑。

4.6 采样：`Sampler`

vllm/v1/sample/sampler.py:21，一个 nn.Module。forward（:73）的处理顺序写在类 docstring 里，要点：

先转 float32 再处理（数值稳定）；
依次过：bad words、logit bias、min tokens → 三种 penalty（repetition/frequency/presence）→ 温度 → top-k/top-p → 采样或 argmax；
有个反直觉设计：返回的 top-k logprobs 用原始 logits（未加惩罚/温度）计算（docstring 第 84-87 行的 NOTE）——展示给用户的概率反映模型本来面貌，不反映采样时的人为扭曲。

5. 关键细节与坑

padding 的 token 是真实参与计算的。它们读 null_block（块 0）的 KV，产出被丢弃；块 0 永不入缓存（第一章）正因此。
图尺寸集是显存换 CPU 的汇率。cudagraph_capture_sizes 每多一个尺寸多一份图缓冲；尺寸太稀则 padding 浪费的计算变多。
FULL 图对注意力后端有要求（后端须支持图内捕获），否则只能 PIECEWISE；这是换后端时性能突变的一个常见来源。
KV cache 大小是 profiling 量出来的：先按 gpu_memory_utilization（默认 0.92，vllm/config/cache.py:111）跑 dummy forward 测峰值，再扣除 CUDA graph 缓冲估计（determine_available_memory，vllm/v1/worker/gpu_worker.py:512）。同样模型在不同卡上块数不同，属预期行为。
跨层共享 KV 的模型（如 YOCO 类）在 initialize_kv_cache_tensors 里直接让多层指向同一张量（vllm/v1/worker/gpu_model_runner.py:7432-7435），调度侧无感。

1. 它要解决的小问题​

2. 思路：持久化 batch + 展平输入 + 图重放​

图示：一步的数据流​

3. 原理演示：向上取整的图派发​

4. 真实实现​

4.1 GPUModelRunner.execute_model​

4.2 块表：BlockTable​

4.3 CUDA graph 的三档模式​

4.4 派发：CudagraphDispatcher.dispatch​

4.5 捕获时机​

4.6 采样：Sampler​

5. 关键细节与坑​