vLLM 推理内核深度解析

第8章 ModelRunner:从 SchedulerOutput 到 GPU Kernel

作者 杨艺韬 · 16,028 字 · 发布于 · 更新于

最快的代码是根本不跑的代码。ModelRunner 的持久化 InputBatch 就是这个思路:每一步不重建整个 batch,只做差量更新。

本章要点

  • 看清 ModelRunner 在 Worker 内部的定位:把 SchedulerOutput 这一抽象的"调度指令"翻译成真正 GPU 上的 tensor 操作
  • 理解持久化 InputBatch(Persistent Batch)为什么是 V1 相对 V0 的关键跃迁:从"每步重建整个 batch"到"每步差量更新"
  • 读懂 _prepare_inputs 在 CPU 侧做什么:按 req_id 打包 block_tables、构造 slot_mapping、拼 input_ids、算 positions
  • 掌握分段 CUDA Graph (Piecewise) 的设计动因:为什么切在 attention 层而不是整图录制
  • 看懂 pad_for_cudagraph 背后的算盘:用 padding token 换 kernel launch 省下的开销,trade-off 为什么划算
  • 理解 torch.compile 和 CUDA Graph 是互补而非替代的两层优化
  • 知道多模态 / Spec Decode / LoRA 怎么接入 execute_model 这条主路径
  • 拿到 ModelRunner 调优的三个核心参数:-O / --compilation-config(编译 level)、cudagraph_capture_sizes、--enforce-eager

8.1 ModelRunner 在整个栈里的位置

前 7 章走完后,我们已经看清 vLLM 的"指令流":API Server 把请求翻成 EngineCoreRequest → EngineCore 的 Scheduler 产出 SchedulerOutput → Executor 广播给 Worker。现在 Worker 拿到 SchedulerOutput 之后,怎么把它变成真正的 GPU 计算?

答案是 ModelRunner(V1 里具体是 GPUModelRunner;v0.8.5 的 vllm/v1/worker/ 下只有它和 TPUModelRunner 两个变体)。它的职责边界非常清晰:

graph TB
    SO[SchedulerOutput<br/>调度层产物] --> MR[ModelRunner.execute_model]
    MR --> P1[_update_states<br/>增删改请求状态]
    MR --> P2[_prepare_inputs<br/>构造 GPU 输入张量]
    MR --> P3[model 前向<br/>N 层 Transformer forward]
    MR --> P4[compute_logits + sample<br/>采样出 token]
    MR --> MO[ModelRunnerOutput<br/>返回采样结果]
    style SO fill:#f59e0b,color:#333,stroke:none
    style MR fill:#2563eb,color:#fff,stroke:none
    style P3 fill:#10b981,color:#333,stroke:none
    style MO fill:#db2777,color:#fff,stroke:none

ModelRunner 是"GPU 执行的最后一公里"——上游所有的抽象(调度、KV 管理、批处理决策)到这里都必须变成具体的 torch.Tensor 和 CUDA kernel。

8.1.1 V1 ModelRunner 在 Worker 内部的调用路径

ModelRunner 是 Worker 的一个字段、Worker 再被 Executor 驱动。调用链是:

Executor (SPMD collective RPC)
    ↓ execute_model(scheduler_output)
Worker.execute_model
    ↓ 
GPUModelRunner.execute_model     ← 本章的主角
    ↓ _update_states + _prepare_inputs + model.forward + sampler
    ↓
ModelRunnerOutput 返回上游

源码里 vllm/v1/worker/gpu_worker.py 的 execute_model 方法基本就是转发——从 Executor 拿到的 scheduler_output 直接喂给 model_runner.execute_model、拿到结果后只由 driver worker 返回(其余 rank 返回 None,vllm/v1/worker/gpu_worker.py:264-269)。Worker 的其他职责(CUDA 设备初始化、模型加载、KV cache 申请、determine_available_memory)都是"启动期"的、跟每 step 热路径没关系。

这种分层让 ModelRunner 可以独立测试——不用拉起整个 Worker、不用搞多进程分布式、单独喂一个 scheduler_output 就能驱动它。仓库里的 tests/v1/worker/test_gpu_model_runner.py 就是这么做的:直接构造 GPUModelRunner、手工拼 SchedulerOutput 来测 _update_states 的各种情形(仍需一块 CUDA 设备)。

硬件变体是平行存在的——但要分清版本边界:v0.8.5 的 vllm/v1/worker/ 里只有 gpu_model_runner.py 和 tpu_model_runner.py;cpu_model_runner.py、neuron_model_runner.py、hpu_model_runner.py、xpu_model_runner.py 仍然只存在于 V0 的 vllm/worker/ 目录下,尚未迁到 V1。每个 Runner 的 _prepare_inputs 实现都不同(TPU 的 tensor 布局、Neuron 的 XLA 图),但 execute_model 的外部契约(接 scheduler_output、产 ModelRunnerOutput)是统一的。这是 vLLM 能横跨多种硬件的架构支撑——Runner 是硬件抽象的底层、上面所有调度、KV 管理、Sampling 都不感知具体硬件。

8.1.2 V0 → V1 的 ModelRunner 变革

V0 的 ModelRunner 每 step 都要重新做一遍几乎所有事情:

# V0 style(概念性)
def execute_model(scheduler_output):
    # 从零开始构造所有输入
    input_ids = [get_input_ids(req) for req in scheduler_output.requests]
    input_ids = pad_and_stack(input_ids)         # list[list] → tensor
    positions = compute_positions(scheduler_output)
    block_tables = build_block_tables(scheduler_output)
    ...  # 十几个张量,全部重新计算
    return model.forward(input_ids, ...)

这种"无状态 ModelRunner" 的问题在于——每步都重复构造 99% 完全一样的数据。decode 一步只有 1 个新 token,但 V0 要把整个 batch 的全部状态重新装一遍。

V1 改成有状态 ModelRunner:核心数据结构 InputBatch 在 Worker 生命周期内持久化,每步只做"增量更新"。这直接对应第 6 章讲过的"V1 Worker 从无状态→有状态"这个架构跃迁。

8.2 InputBatch:持久化的"当前并发全景"

vllm/v1/worker/gpu_input_batch.py 里的 InputBatch 是 V1 ModelRunner 的 state 容器。它在 Worker 初始化时预分配,之后绝大多数字段就地更新。

8.2.1 核心字段

# vllm/v1/worker/gpu_input_batch.py:55(字段名照源码,省略部分同构字段)
class InputBatch:
    # 核心槽位:下标 0..num_reqs-1 各对应一个本步在批里的请求(上限 max_num_reqs)
    _req_ids: list[Optional[str]]      # 槽位 → 请求 id(普通 list,随批增减;对外由 req_ids property 暴露)
    req_id_to_index: dict[str, int]    # 请求 id → 槽位下标

    # per-request 的运行时状态(Numpy 视图,全是预分配)
    num_computed_tokens_cpu: np.ndarray   # 每请求算了多少 token
    num_prompt_tokens: np.ndarray
    num_tokens_no_spec: np.ndarray

    # token 全量缓存(预分配到 max_num_reqs × max_model_len)
    token_ids_cpu_tensor: torch.Tensor    # [max_num_reqs, max_model_len] int32
    token_ids_cpu: np.ndarray             # 同一块内存的 numpy 视图

    # block table 是独立的小对象,自带 CPU/GPU 双份
    block_table: BlockTable               # [max_num_reqs, max_num_blocks_per_req]

    # Sampling params 按字段展开成常驻张量,不是 per-request 的 Python 对象
    temperature / temperature_cpu_tensor / temperature_cpu
    top_p / top_k / min_p / frequency_penalties / presence_penalties / ...

而每 step 真正喂给 kernel 的扁平输入缓冲(input_ids、positions、slot_mapping_cpu、query_start_loc_cpu、seq_lens_cpu)不在 InputBatch 里,而是 GPUModelRunner.__init__() 按 max_num_tokens 预分配的(vllm/v1/worker/gpu_model_runner.py:211-284)。这个分工要记清楚:InputBatch 存"每个请求的持久状态",ModelRunner 存"每一拍的扁平缓冲"。

注意两个观察:

  1. 一切皆数组——没有 List[Request] 这种 Python 对象数组。Python 对象的属性访问一次就是一次字典查找,在 per-step 要处理上百请求的热路径上会累积到可观开销。Numpy 数组的 vectorized 操作在 C 层面完成。

  2. CPU + GPU 镜像——带 _cpu 后缀的是 numpy 数组,不带后缀的是 torch.Tensor。每步在 CPU 上预先 fill 好 numpy 数组,最后一次性 .to('cuda', non_blocking=True) 搬到 GPU。避免多次小规模 H2D 传输。

8.2.2 三种变更操作

SchedulerOutput 告诉 ModelRunner 本 step 相对上一 step 有哪些变化——三类。下面三段代码是概念性简化,方法签名和字段名与源码不同(真实的 InputBatch.add_request(request: CachedRequestState, req_index) 见 vllm/v1/worker/gpu_input_batch.py:231,完整流程见 §8.2.3):

1. 新请求加入 (scheduled_new_reqs) 需要:分配一个空闲槽位、填入 prompt_token_ids、初始化 sampling params、设置 num_computed_tokens=0。

def add_request(self, new_req: NewRequestData):
    slot = self._find_empty_slot()
    self.req_id_to_index[new_req.req_id] = slot
    self.req_ids[slot] = new_req.req_id
    self.token_ids_cpu[slot, :len(new_req.prompt)] = new_req.prompt_token_ids
    self.num_prompt_tokens[slot] = len(new_req.prompt)
    self.temperature_cpu[slot] = new_req.sampling_params.temperature
    # ... 等等

2. 已有请求继续 (scheduled_cached_reqs) 需要:把新调度的 token 追加到该槽位、更新 num_computed_tokens。

def update_request(self, cached_req: CachedRequestData):
    slot = self.req_id_to_index[cached_req.req_id]
    new_start = self.num_computed_tokens_cpu_np[slot]
    new_end = new_start + cached_req.num_scheduled_tokens
    self.token_ids_cpu[slot, new_start:new_end] = cached_req.new_token_ids
    self.num_computed_tokens_cpu_np[slot] = new_end

3. 请求结束 (finished_req_ids) 需要:把那个槽位标记为空、后续 step 可以被新请求复用。

def remove_request(self, req_id: str):
    slot = self.req_id_to_index.pop(req_id)
    self.req_ids[slot] = None
    # 数据字段不需要清零,因为新请求会覆盖

对已在批里的请求,每步只写本步新增的 token 和新分配的块,开销与增量成正比;只有新请求、恢复的请求(以及 §8.2.3 的 condense 搬行)才按整行拷贝。V0 则每步都按序列组在 Python 里重新拼一遍输入,包括每条序列完整的 block table(vllm/worker/model_runner.py 的 ModelInputForGPUBuilder)。

8.2.3 _update_states 的真实执行顺序

§8.2.2 的三种变更操作是概念性说法、真实 _update_states 源码(vllm/v1/worker/gpu_model_runner.py:285-487)是一串严格有序的步骤——前三步腾位置,后面几步补请求、整理批:

段 1:清理 finished_req_ids(line 296-309)——从 self.requests dict(Python 状态缓存)和 self.input_batch(GPU 镜像)两处同时删除请求。删完之后把对应的 slot index 存在 removed_req_indices——这些 slot 马上就要被新请求复用。

段 2:释放 encoder cache(line 311-317)——对多模态模型特有。如果一个请求引用的图像/视频 encoder output 被全部消费完、从 encoder_cache 里删掉——否则就是显存泄漏。这段对纯文本模型是空操作。

段 3:清理 unscheduled 但还未 finished 的请求(line 319-334)——这是最容易被忽略但设计最精妙的一段。典型场景是 preemption:某个请求上一 step 是 RUNNING、本 step 被 Scheduler 决定临时挂起(可能因为 KV 缓存压力)。这种请求没 finish、但也没 scheduled——必须从 InputBatch 里先踢出去、给更紧急的请求让位。但 self.requests 里的缓存状态保留——因为它随时可能被重新调度回来。

段 4:登记 scheduled_new_reqs(line 336-397)——新请求只先写进 self.requests(建 CachedRequestState,M-RoPE 模型顺带算好整条 prompt 的位置),req_id 记进 req_ids_to_add,此时还不进 InputBatch。

段 5:更新 scheduled_cached_reqs(line 399-456)——对继续运行或刚恢复的请求,更新 num_computed_tokens、追加上一步采样出的 token 与新块(从抢占中恢复的请求则整体替换块表)。已在批里的就地写 InputBatch 对应行;不在批里的(被抢占后恢复、或上一步没被调度的)也记进 req_ids_to_add。

段 6:补进批并压实(line 458-486)——req_ids_to_add 里的请求按"小下标优先"填进段 1、3 腾出的空位(removed_req_indices),空位不够就追加到末尾;还剩空位时调 InputBatch.condense() 把末尾的请求搬进空洞,保证批里的请求始终连续占据下标 0..num_reqs-1。最后给注意力后端一次 reorder_batch 的机会(MLA 与 FlashInfer 后端用它把 decode 请求换到前面,FlashAttention 后端直接返回 False),批有变化时刷新 sampling metadata。

这几段的顺序不能乱——先删后加是硬约束、因为槽位上限 max_num_reqs 是固定的、不先腾出 slot 就没法加新请求。每段之间的约束关系被代码结构直接表达——removed_req_indices 从段 1 和段 3 累积、段 6 消费(剩下的交给 condense)——槽位的生产和消费在同一个函数内闭合。

8.2.4 "aborted 后相同 ID resubmit" 的边界处理

_update_states 代码里有一条关键注释(line 300-304):

# NOTE(woosuk): There could be an edge case where finished_req_ids and
# scheduled_req_ids overlap. This happens when a request is aborted and
# then resubmitted with the same ID. In this case, we treat them as two
# distinct requests - clearing the cached states for the first request
# and handling the second as a new request.

翻译:用户 abort 了一个请求(比如客户端断开)、然后用相同的 req_id 重新提交——finished_req_ids 里有这个 id(上一个请求被 abort)、scheduled_new_reqs 里也有这个 id(新请求)。两者看起来"是同一个 id"——代码必须把它们当成两个独立的请求处理:先清理老的(段 1)、再添加新的(段 4)。

为什么不直接"复用"?——因为老请求和新请求的 prompt、sampling_params、mm_inputs 都可能完全不同。如果偷懒"识别为同 id 就直接复用 slot、更新 token"、新的 sampling params 就被丢弃、用户看到的结果和预期完全不符。

这在线上并不罕见:OpenAI 兼容服务器会直接采用客户端传来的 X-Request-Id 请求头作为请求 id 的基础(vllm/entrypoints/openai/serving_engine.py:521-528),客户端断开后带着同一个 X-Request-Id 重试,就会产生"相同 req_id 的两个独立语义请求"。分段处理保证了语义正确。

8.2.5 槽位复用:为什么不按时间线性推

一个容易被问的问题:槽位为什么不按时间线性往后推?先看源码的实际做法:_req_ids 本身就是普通 Python list(append 追加、condense 后 del 截尾),真正按上限预分配的是那些 per-request 的数组——token_ids_cpu、num_computed_tokens_cpu、block table、各项采样参数,长度都是 max_num_reqs,取自 scheduler_config.max_num_seqs(vllm/v1/worker/gpu_model_runner.py:111;V1 下 max_num_seqs 未显式设置时默认 256,显存 ≥ 70 GiB 的卡上默认 1024,见 vllm/engine/arg_utils.py:1596-1621)。

请求结束后留下的空洞会被新请求优先填上,填不满就由 condense 把末尾的请求搬进来,所以批里的请求始终连续地排在下标 0..num_reqs-1。之后每一步都只切前 num_reqs 行(例如 block_table.commit(num_reqs)、注意力元数据里的 block_table[:num_reqs]),不会把 [max_num_seqs, ...] 整块喂给 kernel。预分配换来的是:热路径上不再分配内存,数组地址在 Worker 生命周期内固定。

这种"预分配固定尺寸"的做法,贯穿 V1 的设计(KV Block Pool、CUDA Graph 批大小、InputBatch 槽位),都是相同的工程哲学——用少量冗余 memory 换确定性的性能。

8.3 _prepare_inputs:从 InputBatch 到 GPU 张量

InputBatch 是"all request state"的全景。_prepare_inputs 做的是"把本 step 实际要计算的那些 token 挑出来,打包成 GPU kernel 能吃的张量"。

8.3.1 步骤概览

# gpu_model_runner.py 概念性简化
def _prepare_inputs(self, scheduler_output):
    # 1. 算本 step 每个请求要算几个 token
    num_scheduled_tokens_per_req = np.array([
        scheduler_output.num_scheduled_tokens[rid]
        for rid in self.input_batch.req_ids
    ])
    total_num_scheduled_tokens = num_scheduled_tokens_per_req.sum()

    # 2. 构造 cu_seqlens(每个请求的 token 起始偏移)
    # 给 FlashAttention 用的 variable-length 接口
    cu_seqlens = np.zeros(num_reqs + 1, dtype=np.int32)
    np.cumsum(num_scheduled_tokens_per_req, out=cu_seqlens[1:])

    # 3. 构造 input_ids(一维,flatten 所有请求的 new tokens)
    self._build_input_ids_flat(scheduler_output)

    # 4. 构造 positions(每个 token 在其请求里的绝对位置)
    self._build_positions(scheduler_output)

    # 5. 构造 slot_mapping(每个 new token 写到 KV Cache 的哪个 slot)
    self._build_slot_mapping(scheduler_output)

    # 6. 构造 block_tables(每个 running 请求的 KV block 列表)
    self._pack_block_tables()

    # 7. 把所有 CPU 端的 numpy 数组一次性拷到 GPU
    self._copy_to_gpu(...)

    # 8. 由注意力后端的 builder 构造元数据(FlashAttention 后端产出
    #    FlashAttentionMetadata:query_start_loc、seq_lens、block_table、slot_mapping 等)
    attn_metadata = self.attn_metadata_builder.build(
        num_reqs=num_reqs,
        num_actual_tokens=total_num_scheduled_tokens,
        max_query_len=...,
        common_prefix_len=...,
    )
    return attn_metadata, logits_indices, spec_decode_metadata

每一步都有可以细抠的工程细节。我们挑两个最关键的讲:

8.3.2 block_table.commit 的 CPU-GPU 并行化

真实的 _prepare_inputs(vllm/v1/worker/gpu_model_runner.py:500)一上来第一行、就埋了一个关键优化:

# OPTIMIZATION: Start copying the block table first.
# This way, we can overlap the copy with the following CPU operations.
self.input_batch.block_table.commit(num_reqs)

block_table.commit 内部会异步启动一次 H2D 拷贝(.copy_(..., non_blocking=True) + 在 CUDA stream 上排队)——立即返回、不等待拷贝完成。这意味着:

  • GPU 上的拷贝引擎开始干活
  • CPU 紧接着做后面的 numpy 计算(num_scheduled_tokens 统计、req_indices 构造、cu_num_tokens 累积、positions 算出、slot_mapping 计算……)
  • 等 CPU 到了需要 GPU 侧 block_table 的地方(attn_metadata_builder.build)、拷贝大概率已经完成

这是经典的CPU-GPU 流水并行——不让 GPU 等 CPU、也不让 CPU 等 GPU。block_table 的 H2D 拷贝值得提前——commit(num_reqs) 拷的是 CPU 表的前 num_reqs 行(vllm/v1/worker/block_table.py:69-71),每行 max_num_blocks_per_req = cdiv(max_model_len, block_size) 个 int32,长上下文下一行就不小(例如 max_model_len=128K、block_size=16 时一行 8192 个 int32 = 32 KiB)。CPU 侧表是 pinned 内存,拷贝可以真正异步;把它提前、overlap 在 CPU 的 numpy 计算里,传输时间就藏在了 CPU 计算后面(具体省多少取决于硬件,本机无 GPU 未实测)。

这种优化在其他框架里见得多(PyTorch 自己的 DataLoader 的 pin_memory + non_blocking)、但 vLLM 把它做到了 _prepare_inputs 这种每 step 都跑的热路径里、意识到"拷贝先发、计算后来"才是最短时间路径。

8.3.3 np.repeat + cumsum 的向量化 req_indices 构造

_prepare_inputs 里一段漂亮的 numpy 向量化代码(line 510-523):

# vllm/v1/worker/gpu_model_runner.py:509
# E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)

# Step 1. [2, 5, 3] -> [2, 7, 10]
cu_num_tokens = np.cumsum(num_scheduled_tokens)
# Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7]
cumsums_offsets = np.repeat(cu_num_tokens - num_scheduled_tokens, num_scheduled_tokens)
# Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
arange = self.arange_np[:total_num_scheduled_tokens] - cumsums_offsets

这段在做什么:把"每请求要算几个 token"的数组(比如 [2, 5, 3]),展开成两个扁平数组:

  • req_indices:[0,0, 1,1,1,1,1, 2,2,2]——每个 flat token 属于第几个请求
  • arange:[0,1, 0,1,2,3,4, 0,1,2]——每个 flat token 在所属请求里是第几个新 token

源码注释里明确说了:

# Equivalent to but faster than:
# np.concatenate([np.arange(n) for n in num_scheduled_tokens])

为什么 np.concatenate([np.arange(n) for n in ...]) 慢?——因为它要 Python loop 每请求一次调用 np.arange、每次都分配一个小数组、最后再 concatenate——O(num_reqs) 次 Python 调用 + O(num_reqs) 次内存分配。

vectorized 版本用 3 行 numpy 操作(np.repeat × 2 + np.subtract)、零 Python loop——全部在 C 层完成。源码注释只写了 "Equivalent to but faster than",没有给出倍数;本机无 GPU 也无法复现真实 batch 下的实测,因此这里不编一个具体的微秒数——记住的是"把 O(num_reqs) 次 Python 调用换成常数次 numpy 调用"这个结构差异。

"拒绝 Python loop、用 numpy 向量化一切"是 _prepare_inputs 的核心设计原则。这不是 premature optimization——因为这段代码每 step 都跑,每步省下的 Python 开销会随 step 数线性累积。

8.3.4 torch.index_select 而不是 np.take

_prepare_inputs 里有一段显式偏离 numpy 习惯的代码(line 543-549):

# NOTE(woosuk): We use torch.index_select instead of np.take here
# because torch.index_select is much faster than np.take for large
# tensors.
torch.index_select(
    self.input_batch.token_ids_cpu_tensor.flatten(), 0,
    torch.from_numpy(token_indices),
    out=self.input_ids_cpu[:total_num_scheduled_tokens],
)

token_ids_cpu_tensor 是 torch.Tensor(不是 numpy array)——InputBatch 同时维护了一个 token_ids_cpu(numpy)和一个 token_ids_cpu_tensor(torch tensor sharing the same memory)。这里选择用 torch 版本 + torch.index_select 而不是 np.take——源码注释给出了原因:"much faster for large tensors"。

为什么 torch.index_select 更快?——源码注释没有解释原因。一个合理的推测是:torch.index_select 走 ATen 的 CPU kernel,np.take 走 NumPy 自己的索引路径,两者在大张量上的实现效率不同;这一点本章未做实测。

对每 step 都要做一次的 "从 [max_num_reqs × max_model_len] 的大张量里挑出本 step 的 tokens" 操作,这张表本身就不小:max_num_reqs=256、max_model_len=128K、int32,就是 256 × 131072 × 4 ≈ 128 MiB。源码注释只断言 "much faster for large tensors",没有给倍数,本章也不替它编一个。

vLLM 团队愿意在这里放弃"全 numpy 一致性"、混用 torch 和 numpy——这种实用主义是性能工程的标志。注释里明确写出"为什么选 torch"——让未来维护者知道这不是手滑、是有意的。

8.3.5 non_blocking=True + pin_memory 的 H2D 传输链

_prepare_inputs 末尾几行都是异步 H2D 传输(line 574-586):

# vllm/v1/worker/gpu_model_runner.py:575
self.input_ids[:total_num_scheduled_tokens].copy_(
    self.input_ids_cpu[:total_num_scheduled_tokens], non_blocking=True)
if self.uses_mrope:
    self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
        self.mrope_positions_cpu[:, :total_num_scheduled_tokens],
        non_blocking=True)
else:
    self.positions[:total_num_scheduled_tokens].copy_(
        self.positions_cpu[:total_num_scheduled_tokens],
        non_blocking=True)

non_blocking=True 告诉 PyTorch "不等待拷贝完成就返回"——CPU 立刻进入下一行代码、GPU 上的 DMA 引擎在后台搬数据。

但 non_blocking=True 只对 pin_memory 的 CPU tensor 真正生效——否则 PyTorch 内部会降级为 blocking(因为 pageable memory 不能直接 DMA、必须先拷到临时 pinned buffer 再到 GPU)。参与 H2D 的那几个 buffer 确实都带 pin_memory=self.pin_memory(is_pin_memory_available() 的结果):input_ids_cpu、positions_cpu、slot_mapping_cpu、query_start_loc_cpu、seq_lens_cpu、mrope_positions_cpu(vllm/v1/worker/gpu_model_runner.py:233-283)——这是能真正异步拷贝的前提。反过来,InputBatch.token_ids_cpu_tensor 是显式 pin_memory=False 的,源码注释写明"它不直接传 GPU,所以不需要 pin"。

这两个参数必须成对出现:pin_memory=True + non_blocking=True。缺一不可——单独用 pin_memory 还是同步、单独用 non_blocking 会 silently 退化。V0 其实也这么做:它的 ModelRunner 经 async_tensor_h2d(..., pin_memory) 把 input_tokens / positions 等异步拷上 GPU(vllm/worker/model_runner.py:897-919、vllm/utils.py:944-952),只是每步都新建一个 pinned CPU tensor;V1 的不同在于这些 pinned buffer 在启动时一次性预分配、每步复用。

这几次拷贝和 §8.3.2 的 block_table.commit 一样是异步发出的——CPU 继续算 attn_metadata_builder.build(...)、GPU 的拷贝引擎在并行搬多组数据。_prepare_inputs 返回时、所有 GPU tensor 都已经或即将就绪、紧接着 forward 开始时 CUDA stream 的顺序性会保证拷贝先于 attention kernel 完成。

这种"pinned + async"是 GPU 编程的标准姿势:需要跨 H2D 的 numpy 数组,背后都是一块 pinned torch buffer 的 .numpy() 视图(input_ids_np、positions_np、slot_mapping_np、query_start_loc_np、seq_lens_np),只在 CPU 侧用的大表则不 pin。

8.3.6 slot_mapping:新 token 写到哪个 KV 槽

每 step 新产出的 K/V 要写进 KV Cache。slot_mapping 告诉 attention kernel "第 i 个新 token 的 K/V 应该写到 KV Cache 张量的哪个 slot"。

计算逻辑:

# 对每个 new token,slot = block_id * block_size + offset_in_block
for req_id in scheduled_requests:
    start_pos = num_computed_tokens[req_id]
    for t in range(num_new_tokens):
        abs_pos = start_pos + t
        block_idx = abs_pos // block_size
        offset = abs_pos % block_size
        phys_block = block_table[req_id][block_idx]
        slot_mapping[output_idx] = phys_block * block_size + offset
        output_idx += 1

上面是逐 token 的等价写法;真实代码(vllm/v1/worker/gpu_model_runner.py:551-564)用 §8.3.3 的 req_indices / positions 把它整体向量化成几次 numpy 运算,全在 CPU 上完成。对 batch=128、decode-only step 来说只有 128 个 slot 要算,用 Python loop 也够快;但 prefill step 可能一次处理几千个 token,vectorize 起来就很关键。

8.3.7 Cascade Attention 的"取最小 num_computed_tokens、不 +1"规则

同一批里的多个请求可能共享很长的 common prefix(比如 system prompt、few-shot examples)。vLLM 的 Cascade Attention 让这段共享前缀的 KV 只被读一次——用两段 attention kernel:第一段让所有请求的 query 一起对共享前缀做 attention、第二段算各请求的独立后缀。是否真的启用,还要再过一道 attn_metadata_builder.use_cascade_attention(...) 的收益判断,不划算就返回 0。

但 _compute_cascade_attn_prefix_len(vllm/v1/worker/gpu_model_runner.py:633-715)有一个看起来反直觉的规则:

common_prefix_len = min(
    common_prefix_len,
    self.input_batch.num_computed_tokens_cpu[:num_reqs].min())
# common_prefix_len should be a multiple of the block size.
common_prefix_len = (common_prefix_len // self.block_size * self.block_size)

common_prefix_len 被 cap 到"所有请求 num_computed_tokens 的最小值"——源码注释给出了长篇推理(line 660-698):

请求 1 的 query: [D, E, X],KV cache: [A, B, C, D, E, X],num_computed_tokens: 3(即 [A, B, C]) 请求 2 的 query: [E, Y],KV cache: [A, B, C, D, E, Y],num_computed_tokens: 4(即 [A, B, C, D])

如果把 [A, B, C, D, E] 当成 common prefix、第一段 kernel 会算 [D, E, X, E, Y] vs [A, B, C, D, E] 的双向 attention。但这是错的——因为 Request 1 里的 D 不应该 attend 到 common prefix 里的 E(需要 mask)。所以 [A, B, C, D] 才是正确的 common prefix。

结论分两层,源码注释把它们分得很清楚:

  1. 理论上界是 min(num_computed_tokens) + 1——注释原文是 "capped by the minimum num_computed_tokens, and plus one to include the first token of the query"。上例里就是 [A, B, C, D]。
  2. 实现上真正取的是 min(num_computed_tokens),没有 +1(上例取 [A, B, C])。注释给的理由不是因果性,而是一个实现约束:cascade 路径必须始终跑两个 kernel。假如再来一个请求 3,query 是 [D]、num_computed_tokens=3,common prefix 一旦取到 [A, B, C, D],请求 3 就被第一段 kernel 全部吃掉、第二段 kernel 拿到空输入——当前实现不支持这种情况。

最后再做 block-size 向下对齐。所以代码里就是干净的两行:先 min(),再 // block_size * block_size;- 1 这个说法在源码里并不存在。

这种细腻的语义推理是 LLM 推理引擎的精髓——一个减号错了、模型输出就不对了、但肉眼看不出(因为 logits 数值还是"合理"的)。vLLM 把推导直接写在注释里、让后来者能 trace 为什么是这么写。

8.3.8 M-RoPE 和 1D positions 的二选一分支

_prepare_inputs 的 positions 处理有一条双分支(line 577-586):

# vllm/v1/worker/gpu_model_runner.py:577-586
# (注意 :533 还有一个同样条件的 if,那里做的是 M-RoPE 位置的**计算**
#   `self._calc_mrope_positions(...)`;这里是算完之后往 GPU 上拷)
if self.uses_mrope:
    # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
    self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
        self.mrope_positions_cpu[:, :total_num_scheduled_tokens],
        non_blocking=True)
else:
    # Common case (1D positions)
    self.positions[:total_num_scheduled_tokens].copy_(
        self.positions_cpu[:total_num_scheduled_tokens],
        non_blocking=True)

M-RoPE (Multimodal RoPE) 是 Qwen2-VL 等模型用的旋转位置编码变体——把位置拆成 temporal / height / width 三个分量:image token 的 temporal 分量保持不变、height 与 width 随 patch 坐标变化,video token 三个分量都随帧和坐标变化。启用 M-RoPE 时文本 token 同样是 3D 位置,只是三个分量取相同值,等价于 1D RoPE(源码 vllm/v1/worker/gpu_model_runner.py:228-232 的注释)。

普通模型的 positions 是 [total_tokens] 的一维张量、而 mrope_positions 是 [3, total_tokens] 的二维张量(预分配的 buffer 是 [3, max_num_tokens + 1])(三行分别是 t/h/w)——shape 不兼容、必须分开维护。

uses_mrope 在 GPUModelRunner 初始化时就读一次 model_config.uses_mrope(vllm/v1/worker/gpu_model_runner.py:147)——判定依据是 hf_config.rope_scaling 里有没有 mrope_section 键,Qwen2.5-Omni 这类则看其 thinker 配置(vllm/transformers_utils/config.py:223-246)——一次性决定、之后每 step 走同一条路径。这避免了热路径里每 step 都做"这是不是 M-RoPE 模型"的判断。

多模态 + M-RoPE 的组合让代码里多出两段专属逻辑:新请求进来时,_update_states 里有约 35 行(vllm/v1/worker/gpu_model_runner.py:360-395)从 mm_inputs 提取 image_grid_thw、video_grid_thw、second_per_grid_ts、audio_feature_lengths,交给 MRotaryEmbedding.get_input_positions_tensor() 算出整条 prompt 的 3D 位置;每一拍再由 _prepare_inputs 调 _calc_mrope_positions()(vllm/v1/worker/gpu_model_runner.py:717)把本拍要算的那一段切出来。这部分代码对纯文本 / 纯 CLIP-style 多模态模型(比如 LLaVA)都是死代码——只在 Qwen2-VL 这类模型下才激活。

这种"一条分支服务一个特定模型家族"是 vLLM 源码里常见的模式——模型架构变化快、新的 RoPE 变体、MLA、MoE 不断出现、引擎必须能快速添加分支而不破坏主路径。M-RoPE 就是这样一个"局部添加、对主路径零影响"的典范——热路径里多一个 if、特性完全支持。

8.3.9 Block Table 打包:二维张量里的 padding

block_table 在 InputBatch 里是 [max_num_reqs, max_num_blocks_per_req] 的 pinned CPU 张量(外加 numpy 视图),GPU 上有一份同形状的镜像。_prepare_inputs 要把它搬到 GPU,同时只搬本 step 在批里的请求那几行。

这里不需要挑行:§8.2.3 的 condense 已经保证批里的请求连续排在前 num_reqs 行,所以 commit 直接拷前缀,注意力元数据也只切前缀:

# vllm/v1/worker/block_table.py:69-71
def commit(self, num_reqs: int) -> None:
    self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
                                      non_blocking=True)

# vllm/v1/attention/backends/flash_attn.py:305-306(FlashAttentionMetadataBuilder.build)
block_table = (
    self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])

第 num_reqs 行之后是历史残留数据,kernel 拿到的是 [:num_reqs] 切片,根本看不到它们。也不存在"为 CUDA Graph pad 到固定行数"的问题——V1 的 attention 不在 CUDA Graph 里(§8.4.2)。

8.3.10 hot-swap LoRA 在 _prepare_inputs 尾部的挂钩

_prepare_inputs 的最后一段(line 627-630):

# Hot-Swap lora model
if self.lora_config:
    self.set_active_loras(self.input_batch, num_scheduled_tokens)

return attn_metadata, logits_indices, spec_decode_metadata

set_active_loras 是 LoRA 多租户的核心——根据每个请求当前用的 LoRA adapter id、切换 GPU 上已加载的 LoRA 权重到生效槽位。在本 step 实际 forward 之前、必须先把 "哪些 LoRA 参与这个 batch" 搞清楚。

为什么在 _prepare_inputs 里调、而不是 execute_model 的主流程里?——因为 set_active_loras 需要 num_scheduled_tokens(知道每个请求算几个 token),这个信息在 _prepare_inputs 里刚好计算完。如果放到 execute_model 里调、要么重复算 num_scheduled_tokens、要么把变量层层传递——不如在同一个 scope 里一气呵成。

LoRA hot-swap 的开销——权重的加载 / 换入换出不是每 step 都发生:worker 侧用的是 LRUCacheWorkerLoRAManager(vllm/v1/worker/lora_model_runner_mixin.py:43),已驻留在槽位里的 adapter 不会重复加载;token→LoRA 的映射也只在与上一次不同(!= 比较)时才重新下发(vllm/adapter_commons/utils.py:27-32)。不过每步仍要在 CPU 上由 make_lora_inputs 重新拼一遍映射元组,所以稳态开销很小,但不是零。

第 16 章会详讲 LoRA 的实现细节(Punica wrapper、v0.8.5 已统一的 lora_shrink / lora_expand kernel 等)、本章这一行就是LoRA 和主 forward 接头的那一点——简单、不显眼、但承载着多租户推理的全部成本控制。

8.3.11 kv_connector 和 KV 转移的主流程挂钩

execute_model 一开头有一段KV 传输的挂钩(line 1008-1011):

# Update KVConnector with the KVConnector metadata forward().
if has_kv_transfer_group():
    get_kv_transfer_group().bind_connector_metadata(
        scheduler_output.kv_connector_metadata)

这是为跨实例 KV 传输(比如 disaggregated prefill / decode 架构)准备的——把 prefill 阶段算出的 KV 交给专门做 decode 的实例复用。KVConnectorBase_V1 接口让不同的后端接入同一个架构;v0.8.5 的 V1 目录下(vllm/distributed/kv_transfer/kv_connector/v1/)只有 SharedStorageConnector 和 LMCacheConnectorV1 两个实现。

这一行本身不搬数据——bind_connector_metadata 只是把调度器给的元数据存到 connector 上(vllm/distributed/kv_transfer/kv_connector/v1/base.py:69-80),要在 forward 之前绑好,后面才用得上。真正的传输挂在 set_forward_context 上(§8.10.5):进入时调 start_load_kv 触发(可能异步的)加载,各 attention 层按需等待自己那层读完;退出时调 wait_for_save,等各层发起的保存全部完成(vllm/forward_context.py:106-157,加载在 111-114 行、等待保存在 154-157 行)。execute_model 末尾再 clear_connector_metadata()(line 1277-1279)。逐层加载 / 保存让传输有机会与各层计算重叠。

绝大部分单机部署的用户不会触碰这个挂钩——has_kv_transfer_group() 返回 False、整段代码 short-circuit。但对 PD disaggregation 用户(比如大规模在线服务把 prefill GPU 和 decode GPU 分开)、这个挂钩是生死攸关的性能路径。

这又是一个"主路径对少数场景的开放点"——像 LoRA、spec decode、multimodal 一样、KV transfer 以最小侵入的方式嵌入到 execute_model 里、对不用它的人零开销、对用它的人提供完整功能。这种扩展点的克制是 vLLM 源码能保持可读性的重要原因——主路径始终清晰、扩展功能通过"空实现兜底 + 注入激活"融进来。这种模式和 §8.3.10 的 LoRA、§8.3.8 的 M-RoPE 一脉相承、构成了 vLLM 可扩展性的核心设计。不管是多租户、新模态、新传输方式——都按这套"零侵入接入"的路子走。

8.4 分段 CUDA Graph (Piecewise) 的动因

8.4.1 标准 CUDA Graph 的痛点

CUDA Graph 把一串 kernel launch 录制成一个图,replay 只需要一次 API 调用——省下的是每个 kernel 各自走一遍 CUDA Runtime API 的固定开销。一个 Transformer decode step 要 launch 的 kernel 数大致正比于层数 × 每层算子数,几十层的模型就是几百次 launch;单次 launch 的绝对开销随 GPU、驱动、CUDA 版本变化,本机无 GPU 无法实测,所以这里只说结构:launch 次数从几百降到 1。

但标准 CUDA Graph 有两个硬约束:

  1. 所有 shape 必须固定——录的时候多大,replay 就多大
  2. 内存地址必须固定——不能分配新 tensor

约束 1 对 LLM 特别麻烦:batch 大小可变(用户数波动)、seq len 可变(每请求不同 context)、prefill vs decode 形态差异巨大……

V0 的做法是整图录制、只用于纯 decode 批:prefill 或混合批一律走 eager;decode 批按序列数 pad 到预录的档位,档位取 [1, 2, 4] + 8 的倍数、上限是第一个 ≥ max_num_seqs 的档(vllm/config.py:4001-4030,max_num_seqs=256 时共 35 档),另外还要求本批最长序列不超过 max_seq_len_to_capture(vllm/worker/model_runner.py:762-770)。也就是说,V0 靠"只录 decode"绕开了 prefill 的形状多变,代价是 prefill 与混合批拿不到 CUDA Graph。

8.4.2 分段 (Piecewise) 的思路

观察模型结构:

graph LR
    E["Embedding<br/>shape: (T, H)"] --> L1["Layer 1<br/>Attention + MLP"]
    L1 --> L2[Layer 2]
    L2 --> DOTS[...]
    DOTS --> LN[Layer N]
    LN --> NORM[Final Norm]
    NORM --> LM["LM Head<br/>(T, vocab)"]
    subgraph "shape 敏感"
        A1[Attention<br/>依赖 cu_seqlens / block_tables]
    end
    L1 -.->|"内部"| A1
    style A1 fill:#dc2626,color:#fff,stroke:none
    style E fill:#10b981,color:#333,stroke:none
    style LM fill:#10b981,color:#333,stroke:none

只有 attention 算子对 shape 真正敏感(需要知道每个请求的 seq len、block_table)。其他算子(LayerNorm、MLP 的 Linear、激活函数)只看 (total_tokens, hidden_dim) 这一维——只要 total_tokens 能被 pad 到几个固定值,它们就能 CUDA Graph 化。

分段 CUDA Graph 的思路:

  • 把模型按 attention 切分成多个"段"
  • 每段(非 attention)用 CUDA Graph 录制,按 total_tokens 的几个离散值(cudagraph_capture_sizes)分别录
  • Attention 算子不进图,每拍按当拍的 attn_metadata 以 eager 方式调用——V1 默认的切分点就是 vllm.unified_attention 与 vllm.unified_attention_with_output 这两个自定义算子(vllm/config.py:3652-3659 的 set_splitting_ops_for_v1)
  • 段之间通过共享 GPU tensor 传递 hidden state
graph TB
    subgraph "标准 CUDA Graph"
        SG[整个模型整图录制<br/>attention 依赖每请求的 seq_len 与 block_table<br/>想覆盖 prefill 就要按组合录图<br/>组合爆炸,V0 因此只录 decode]
    end
    subgraph "Piecewise CUDA Graph"
        P1[Segment 1<br/>Embedding + Layer 1 的非 attn 部分<br/>只依赖 total_tokens<br/>按 cudagraph_capture_sizes 录图(V1 默认 67 档)]
        P2[Attention 1<br/>Eager 模式<br/>每次动态决定]
        P3[Segment 2<br/>MLP 1 + Layer 2 的非 attn 部分<br/>同样 CUDA Graph 化]
        P4[Attention 2<br/>Eager 模式]
        P5[...]
    end
    P1 --> P2 --> P3 --> P4 --> P5
    style SG fill:#dc2626,color:#fff,stroke:none
    style P1 fill:#10b981,color:#333,stroke:none
    style P2 fill:#f59e0b,color:#333,stroke:none
    style P3 fill:#10b981,color:#333,stroke:none
    style P4 fill:#f59e0b,color:#333,stroke:none

8.4.3 实际效果

vLLM 的 Piecewise 实现(vllm/compilation/backends.py + torch.compile 集成)按一组离散的 token 数分别录图。v0.8.5 的 V1 默认列表在 vllm/config.py:4031-4046(下面省略了开启 sequence parallelism 时剔除非 tp_size 倍数档位的分支):

batch_size_capture_list = [1, 2, 4] + [i for i in range(8, 513, 8)]
# 再按 max_num_batched_tokens 截断
batch_size_capture_list = [
    size for size in batch_size_capture_list if size <= max_num_tokens
]

也就是 [1, 2, 4, 8, 16, 24, 32, ..., 512],8 以上按 8 递增,未截断时共 67 个尺寸。enforce_eager=True 时这个列表直接是空的,一张图都不录。

Piecewise 相对"整图录制"的收益,源码层面能确认的是两条:要录的图只覆盖非 attention 段,attention 走 eager,因此不需要为每种 seq-len / block-table 组合各录一份;捕获顺序从大到小,小尺寸复用大尺寸的 memory pool(§6.7.2)。至于具体省了多少启动秒数和多少显存,capture_model() 自己在收尾时会把 elapsed_time 和 cuda_graph_size 打进日志——这是每个人在自己机器上唯一可信的数字来源,本章不代为编造。

8.5 pad_for_cudagraph:padding 的经济学

前面说过 CUDA Graph 需要固定 shape。实际一拍的 total_tokens 可能是任意值,比如 47。但预录制的 size 只有 [1, 2, 4, 8, 16, 24, 32, 40, 48, ...]——怎么办?

pad_for_cudagraph(47) = 48——往上 pad 到最近的预录制 size。

# gpu_model_runner.py:1022-1038(简化,省略了 eager 分支里的 SP 对齐,见 §8.8)
if self.use_cuda_graph and num_scheduled_tokens <= max_graph_size:
    num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)
else:
    num_input_tokens = num_scheduled_tokens  # Eager 回退

padding 的 token 是"无意义"的——attention 只按 num_actual_tokens 处理真实 token(slot_mapping 不含 padding),算 logits 前也会先切 hidden_states[:num_scheduled_tokens],不影响正确性,只是非 attention 段多跑了一些无用 FLOPs。

8.5.1 pad_for_cudagraph 的实现:预铺平查表,O(1)

pad_for_cudagraph(47) 怎么知道该 pad 到 48?答案比"二分查找"还要极端——它连查找都不做(vllm/config.py:3783-3788):

def pad_for_cudagraph(self, batch_size: int) -> int:
    # if batch_size > self.compilation_config.max_capture_size,
    # it should raise an IndexError.
    # the caller should make sure the batch_size is within the range,
    # i.e., batch_size <= self.compilation_config.max_capture_size
    return self.compilation_config.bs_to_padded_graph_size[batch_size]

热路径上就是一次 list 下标访问,O(1)。真正的映射是在初始化时一次性铺平的(vllm/config.py:3638-3650):

# pre-compute the mapping from batch size to padded graph size
self.bs_to_padded_graph_size = [0 for i in range(self.max_capture_size + 1)]
for end, start in zip(self.cudagraph_capture_sizes,
                      self.cudagraph_capture_sizes[1:] + [0]):
    for bs in range(start, end):
        if bs == start:
            self.bs_to_padded_graph_size[bs] = start   # 正好命中某个 size
        else:
            self.bs_to_padded_graph_size[bs] = end     # 向上取到下一个 size
self.bs_to_padded_graph_size[self.max_capture_size] = self.max_capture_size

这是典型的空间换时间 + 把分支挪到初始化期:用一个长度 max_capture_size + 1 的 int list(默认 512 时约 4 KB)换掉每 step 一次查找。注意 bs == start 那一支——如果 token 数正好等于某个 capture size,映射到它自己,零 padding 浪费。

还要注意函数的前置条件:注释明确写了调用方必须保证 batch_size <= max_capture_size,否则就是 IndexError。这个检查在调用点(vllm/v1/worker/gpu_model_runner.py:1022-1027)——num_scheduled_tokens <= self.cudagraph_batch_sizes[-1] 不成立就直接走 eager。边界检查在外、热路径里连 if 都不要,是同一种思路的两面。

8.5.2 Padding 的开销分析

按 v0.8.5 V1 的真实默认列表 [1, 2, 4] + range(8, 513, 8) 算,padding 浪费率是可以精确算出来的(这是纯算术,不是性能实测):

  • total_tokens=32 → pad 到 32,0% 浪费
  • total_tokens=33 → pad 到 40,17.5% 浪费
  • total_tokens=47 → pad 到 48,2.1% 浪费
  • total_tokens=9 → pad 到 16,43.8% 浪费(最坏的一档,正好落在 8 与 16 之间的空档)
  • total_tokens=129 → pad 到 136,5.1% 浪费

规律很清楚:8 以上每档步长固定为 8,所以绝对浪费最多 7 个 token;token 数越大,相对浪费越小。真正难看的只有最下面那几档(1/2/4/8/16 之间的跳跃),而那些档位对应的本来就是几乎空载的 batch。

这就是 trade-off 的核心:padding 浪费的是 FLOPs,CUDA Graph 省下的是 launch 开销,两者哪边大取决于模型规模、GPU 型号和 batch 形态,没有普适答案,必须在目标环境上量。但有一条和绝对数字无关的收益:CUDA Graph 让每 step 的耗时方差小得多——launch 路径被固化,不再随主机侧 Python/驱动的抖动波动,这对 P99 latency SLA 是实打实的好处。

8.5.3 不值得 CUDA Graph 化的场景

  • batch 大小经常 > 最大预录制 size → fall back 到 Eager,没有收益
  • 启动时就明确短期内不需要 steady-state 性能(比如一次性的 offline batch inference)

两个常被误列进来的场景其实不受影响:guided decoding 的 bitmask 是在模型 forward 之后、对 compute_logits 的结果施加的(vllm/v1/worker/gpu_model_runner.py:1112-1114),不在图里;开了 Spec Decode 时,目标模型的 verify 前向照样按"含 draft 的 token 总数"走 pad_for_cudagraph,只有 EAGLE 的 draft head 在 v0.8.5 不走 CUDA Graph(vllm/v1/worker/gpu_model_runner.py:1223-1226 的注释)。

这些场景下可以用 --enforce-eager 完全关掉 CUDA Graph,节省 startup 时间。

8.6 torch.compile 的角色:两层优化的协同

V1 默认启用 torch.compile。它和 CUDA Graph 不是竞争关系,是互补层次:

graph TB
    subgraph "torch.compile 层(Python 侧优化)"
        TC1[算子融合<br/>如 bias + 激活等逐元素算子合成一个 kernel]
        TC2[内存规划<br/>减少中间张量分配]
        TC3[Triton kernel 生成<br/>硬件感知的代码]
    end
    subgraph "CUDA Graph 层(Kernel launch 优化)"
        CG1[把剩下的 kernel launch<br/>打包成一个 replay 动作]
    end
    TC1 --> CG1
    TC2 --> CG1
    TC3 --> CG1
    style TC1 fill:#2563eb,color:#fff,stroke:none
    style TC2 fill:#2563eb,color:#fff,stroke:none
    style TC3 fill:#2563eb,color:#fff,stroke:none
    style CG1 fill:#10b981,color:#333,stroke:none

torch.compile 让每个 step 的算子数变少(融合掉一批逐元素算子),CUDA Graph 把剩下这些 kernel 的 launch overhead 打包省掉。两者作用在不同层次,所以可以叠加;叠加后的实际收益取决于模型、GPU 和 batch 形态,只能在目标环境实测,本章不给一个跨环境的百分比。

8.6.1 set_forward_context 与 CUDA Graph 的协作细节

§8.10.5 会看到 set_forward_context(attn_metadata, self.vllm_config) 把元数据注入到模型 forward。但这里有一个 CUDA Graph 相关的微妙点:

attn_metadata 并不进图。capture_model() 录图时调用的 _dummy_run 传的是 set_forward_context(None, ...)(vllm/v1/worker/gpu_model_runner.py:1494-1496),attention 本身又是切分点、在两段图之间以 eager 方式执行(§8.4.2)。所以每拍的 attn_metadata 可以是新建的对象:FlashAttentionMetadataBuilder.build() 每步都用 .to(device, non_blocking=True) 新建 query_start_loc、seq_lens、slot_mapping 的 GPU 张量(vllm/v1/attention/backends/flash_attn.py:297-308),只有 block_table 是预分配 GPU 表的 [:num_reqs] 切片。

真正必须地址固定的,是图的输入:input_ids、positions(M-RoPE 时是 mrope_positions)、多模态路径的 inputs_embeds、流水线并行下的 intermediate_tensors。它们都在 GPUModelRunner.__init__() 里按 max_num_tokens 预分配(源码注释就叫 "Persistent buffers for CUDA graphs",vllm/v1/worker/gpu_model_runner.py:211),每步只往前 num_input_tokens 个位置写新值、再切片传给模型——地址不变、内容变,replay 时读到的就是新写入的数据。

这是"预分配 + 就地更新"模式的深层价值——不只是避免分配开销、还让 CUDA Graph 能重复使用。§8.10.4 的多模态 inputs_embeds 多拷一次,也是为了把动态产物搬进这类固定地址的 buffer。

8.6.2 CompilationLevel

vLLM 用 CompilationLevel 控制优化级别:

class CompilationLevel:
    NO_COMPILATION = 0        # 全 Eager
    DYNAMO_AS_IS = 1           # 只用 torch.compile 的 Dynamo 前端
    DYNAMO_ONCE = 2            # Dynamo + 缓存
    PIECEWISE = 3              # Dynamo + 分段 CUDA Graph(V1 默认)

注意 CompilationConfig.level 这个字段本身的默认值是 0(vllm/config.py:3415);level 3 是 V1 的有效默认——VllmConfig.__post_init__ 里,只要 VLLM_USE_V1 且 enforce_eager=False,就会把它强制改写成 CompilationLevel.PIECEWISE(vllm/config.py:3884-3899)。

命令行上没有 --compilation-level 这个参数,对应的是 -O / --compilation-config(vllm/engine/arg_utils.py:873-889):传一个数字就当 level 解释(-O3 等价于 -O 3),传 JSON 串可以连 cudagraph_capture_sizes 一起指定。help 文本还写明:level 0 是不做任何优化的默认级别,level 1 和 2 仅供内部测试,level 3 才是推荐的生产级别。

8.6.3 四个 CompilationLevel 背后的具体后端选择

上面一句话带过四个 level。实际上它们对应完全不同的后端和代码路径——读懂这个映射才能判断什么情况该用哪一级。翻开 vllm/config.py:3581 的 init_backend 方法:

def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:
    if self.level == CompilationLevel.NO_COMPILATION:
        raise ValueError("No compilation level is set.")

    from torch._dynamo.backends.registry import list_backends
    torch_backends = list_backends(exclude_tags=tuple())
    if self.level in [
            CompilationLevel.DYNAMO_AS_IS, CompilationLevel.DYNAMO_ONCE
    ]:
        if self.backend == "":
            return "eager"                                  # ← 默认给 eager
        if self.backend in torch_backends:
            return self.backend
        return resolve_obj_by_qualname(self.backend)

    assert self.level == CompilationLevel.PIECEWISE
    from vllm.compilation.backends import VllmBackend
    return VllmBackend(vllm_config)                          # ← 只有 PIECEWISE 用 vLLM 自己的 backend

这揭示了四个 level 的具体分工:

Level 编号 后端 实际行为
NO_COMPILATION 0 — init_backend 直接抛 ValueError——调用方根本不会走到这里,而是走 eager 路径完全绕过 torch.compile
DYNAMO_AS_IS 1 "eager"(默认)或用户指定的 torch backend 走 torch.compile 的 Dynamo 前端做图捕获,但后端用 eager——只获得图追踪、不做任何编译优化。用途:调试 Dynamo 对 vLLM 模型的追踪行为、查哪里有 graph break
DYNAMO_ONCE 2 同 1(默认 "eager",可配) 和 level 1 的区别是缓存——只追踪一次然后复用。对于不变的模型结构省下重复追踪开销。但仍然不做真正的编译优化
PIECEWISE 3 VllmBackend(vllm_config) 走 vLLM 自家的 VllmBackend(vllm/compilation/backends.py:284),做分段图切分 + Inductor 编译 + 整合 CUDA Graph

VllmBackend 是 vLLM 独有的 torch.compile 后端,自己的 docstring(line 284-294)说得很清楚:

The major work of this backend is to split the graph into piecewise graphs, and pass them to the piecewise backend. This backend also adds the PostGradPassManager to Inductor config, which handles the post-grad passes.

它干三件 torch 默认 Inductor 不做的事:

  1. 图切分:按 compilation_config.splitting_ops 把整个模型 forward 图切开(split_graph,vllm/compilation/backends.py:440-441)——V1 默认的切分点就是两个 attention 自定义算子,attention 单独成段、不编译也不录图,其余各段独立处理
  2. PostGradPassManager 插入:在 Inductor 的 post-grad 优化点插入 vLLM 自己的 pass——v0.8.5 有 NoOp 消除、RMSNorm + 量化融合(FusionPass)、sequence parallelism 改写,以及总在最后跑的 FixFunctionalizationPass(vllm/compilation/pass_manager.py:47-58);V1 在 VllmConfig.__post_init__ 里把前两者默认关掉(vllm/config.py:3896-3897)
  3. CUDA Graph 整合:对 cudagraph_capture_sizes 里的每个尺寸,把各非 attention 段编译后的代码自动录制成 CUDA Graph——用户什么都不用写,level=3 就都有了

级别选择的实用指南:

  • 写 / 调新模型:用 level 1。看 Dynamo 是否能追踪、有多少 graph break
  • 生产推理:用 level 3(V1 的有效默认)。这是 -O 的 help 文本里唯一被标为 "recommended level for production" 的档位
  • 某些诊断场景:用 --enforce-eager(未显式指定 -O 时就等于停在 level 0、不编译)——例如你怀疑是 Inductor 编译出了错的代码、想验证裸 PyTorch 行为时

Level 2 在实践中很少用——它是"和 level 1 类似但只追踪一次",-O 的 help 文本把 level 1、2 都归为"仅供内部测试"。

这 4 层 level 和第 13 章讲的量化注册表、第 14 章讲的 GroupCoordinator 一样——表面上是个配置枚举、背后各自对应一套完全不同的代码路径。这种"统一参数接口 + 分支化实现" 是 vLLM 能快速迭代新编译特性的关键架构——新加一种编译方式就是往这里加一个 level 和对应的 backend,核心 ModelRunner 代码零改动。

8.6.4 use_spec_decode 分支里的 logits_indices

_prepare_inputs 的尾部有一个对 speculative decoding 的特殊处理(line 603-625):

# vllm/v1/worker/gpu_model_runner.py:603 起(源码里这句是折行的)
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
if not use_spec_decode:
    # NOTE: Due to chunked prefills, the batch may contain partial requests.
    # ...
    # TODO: Support prompt logprobs.
    logits_indices = attn_metadata.query_start_loc[1:] - 1
    spec_decode_metadata = None
else:
    # Get the number of draft tokens for each request.
    num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
    for req_id, draft_token_ids in scheduler_output.scheduled_spec_decode_tokens.items():
        req_idx = self.input_batch.req_id_to_index[req_id]
        num_draft_tokens[req_idx] = len(draft_token_ids)

    spec_decode_metadata = self._calc_spec_decode_metadata(
        num_draft_tokens, cu_num_tokens)
    logits_indices = spec_decode_metadata.logits_indices

无 spec decode 分支(正常路径):logits_indices = query_start_loc[1:] - 1——这一行看似简单、语义丰富。query_start_loc 是 [0, 2, 9, 12] 这样的前缀和、query_start_loc[1:] - 1 就是 [1, 8, 11]——每个请求的最后一个 token 在 flat 序列里的位置。只在这些位置计算 logits、采样出下一个 token。

为什么不对每个 token 都算 logits?——因为 prefill 阶段前 N-1 个 token 都是 "已知的输入"、不需要预测;只有第 N 个 token(最后一个)的 hidden state 用于预测"下一个 token"。这样能省掉大量 matmul(logits = hidden @ vocab_weight 是整个 forward 里最大的 matmul 之一)。

spec decode 分支:spec_decode_metadata.logits_indices 更复杂——因为投机解码要为每个 draft token 都算 logits(用 target model verify),不能只算最后一个。_calc_spec_decode_metadata 负责构造这些索引。

TODO: Support prompt logprobs——这条注释说的是 logits_indices 这条路径只覆盖每个请求的最后一个位置,不负责 prompt logprobs;它并不意味着 V1 不支持 prompt logprobs。v0.8.5 的 prompt logprobs 由 execute_model 里单独的 _get_prompt_logprobs_dict(hidden_states, scheduler_output) 从完整的 hidden_states 另算(vllm/v1/worker/gpu_model_runner.py:1172-1176、1351)。同一段注释还说明了 chunked prefill 的一个折衷:partial request(prompt 还没 prefill 完)也会顺手采样一个 token,但事后被丢弃——execute_model 里把这些请求的采样结果清空,并把随机数生成器的 offset 回拨(vllm/v1/worker/gpu_model_runner.py:1148-1164、1190-1192)。

8.7 execute_model:完整主流程

把前面所有模块串起来,一个典型 step 的完整调用栈(简化;省略了开头的 KV connector 挂钩、流水线并行的 intermediate_tensors 处理和投机解码的采样分支):

def execute_model(self, scheduler_output) -> ModelRunnerOutput:
    # 1. 更新 InputBatch 状态(差量)
    self._update_states(scheduler_output)
    if not scheduler_output.total_num_scheduled_tokens:
        return EMPTY_MODEL_RUNNER_OUTPUT   # vllm/v1/outputs.py:104

    # 2. 从 InputBatch 构造本 step 的 GPU 输入
    attn_metadata, logits_indices, spec_decode_metadata = \
        self._prepare_inputs(scheduler_output)
    num_scheduled = scheduler_output.total_num_scheduled_tokens

    # 3. CUDA Graph padding 决策
    if self.use_cuda_graph and num_scheduled <= max_graph_size:
        num_input_tokens = pad_for_cudagraph(num_scheduled)
    else:
        num_input_tokens = num_scheduled

    # 4. 多模态挂钩:多模态模型一律用 embeddings 作输入,纯文本模型直接喂 token id
    if self.is_multimodal_model:
        self._execute_mm_encoder(scheduler_output)
        mm_embeds = self._gather_mm_embeddings(scheduler_output)
        inputs_embeds = ...   # get_input_embeddings 后拷进持久 buffer(§8.10.4)
        input_ids = None
    else:
        input_ids = self.input_ids[:num_input_tokens]
        inputs_embeds = None

    # 5. LoRA 挂钩(如启用)——注意真实源码里这一步在 _prepare_inputs 尾部
    #    就做完了(§8.3.10),这里只是为了让阶段顺序读起来完整

    # 6. 模型前向(真正的 GPU 计算)
    #    attn_metadata 不走函数签名,而是通过 forward context 注入(§8.10.5)
    with set_forward_context(attn_metadata, self.vllm_config):
        hidden_states = self.model(
            input_ids=input_ids,          # 多模态路径下为 None
            positions=self.positions[:num_input_tokens],
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )

    # 7. 计算 logits(先去掉 padding,再只取"需要输出 token"的位置)
    hidden_states = hidden_states[:num_scheduled]
    logits = self.model.compute_logits(hidden_states[logits_indices], None)
    if scheduler_output.grammar_bitmask is not None:   # 结构化输出的 token mask
        self.apply_grammar_bitmask(scheduler_output, logits)

    # 8. 采样(第 9 章详细讨论;投机解码时改走 rejection_sampler,§8.9)
    sampler_output = self.sampler(
        logits=logits,
        sampling_metadata=self.input_batch.sampling_metadata,
    )
    #    partial prefill 请求的采样结果会被清空;按需另算 prompt logprobs(§8.6.4)

    # 9. 投机解码挂钩(如启用)
    if self.use_spec_decode:
        spec_token_ids = ...   # ngram / EAGLE drafter 提议下一步的 draft token
    else:
        spec_token_ids = None

    # 10. 打包返回
    return ModelRunnerOutput(
        req_ids=self.input_batch.req_ids,
        req_id_to_index=self.input_batch.req_id_to_index,
        sampled_token_ids=valid_sampled_token_ids,   # list[list[int]]
        spec_token_ids=spec_token_ids,
        logprobs=logprobs_lists,
        prompt_logprobs_dict=prompt_logprobs_dict,
    )

这就是一个 step 的全部——10 个阶段,每阶段对应本专栏前面讲过的某个概念:

  • 1、2 → InputBatch(本章)
  • 3 → CUDA Graph(本章)
  • 4 → Multimodal(第 15 章)
  • 5 → LoRA(第 16 章)
  • 6 → PagedAttention(第 4 章)
  • 7 → Transformer 模型实现
  • 8 → Sampler(第 9 章)
  • 9 → Speculative Decoding(第 12 章)
  • 10 → EngineCore 回调(第 2 章)

ModelRunner 是所有章节知识的"交汇点"。把它读懂,整个 V1 引擎的工程架构就立体了。

8.8 execute_model 里"eager mode 下还要 round_up to tp_size"的冷门分支

execute_model 的 padding 决策里(vllm/v1/worker/gpu_model_runner.py:1028-1038)有一个在纯文本 decode 场景看不到的冷门分支:

else:
    # Eager mode.
    # Pad tokens to multiple of tensor_parallel_size when
    # enabled collective fusion for SP
    tp_size = self.vllm_config.parallel_config.tensor_parallel_size
    if self.vllm_config.compilation_config.pass_config. \
        enable_sequence_parallelism and tp_size > 1:
        from vllm.utils import round_up
        num_input_tokens = round_up(num_scheduled_tokens, tp_size)
    else:
        num_input_tokens = num_scheduled_tokens

只要本拍不走 CUDA Graph,且开了 sequence parallelism、tp > 1——还是要把 num_input_tokens 向上 round 到 tp_size 的倍数。注意这个 else 分支不只对应 --enforce-eager:token 数超过最大 capture size 时同样走这里,而那时模型仍是 torch.compile 编译过的。为什么要对齐?源码注释只写了"开启 SP 的 collective fusion 时按 tp_size 对齐";从实现看,SequenceParallelismPass 把 all-reduce 改写成沿第 0 维的 reduce-scatter + all-gather,并且只对能被 tp_size 整除的 shape 生效(is_applicable_for_shape 返回 shape % tp_size == 0,vllm/compilation/sequence_parallelism.py:256-260);同理,开 SP 时 CUDA Graph 的档位也会剔除不是 tp_size 倍数的尺寸(vllm/config.py:3952-3970)。

这个分支在大部分用户那里用不到——sequence parallelism 是需要显式打开的 pass 配置,默认关闭。但一旦开了,这里的 round_up 保证了无论走不走 CUDA Graph,送进模型的 token 数都满足这个 pass 的整除前提。

这体现了 vLLM 对"多配置组合"的覆盖——不是简化成"CUDA Graph 才 pad"、而是"CUDA Graph 和 TP-SP 都可能要 pad、分别判断"。配置矩阵越大、这种分支就越细。

8.9 spec_decode 的 bonus/target logits 独立存储

execute_model 的 spec decode 路径(line 1124-1146)里有一段对 PyTorch tensor 语义的深度利用:

# vllm/v1/worker/gpu_model_runner.py:1124
# When indexing with a tensor (bonus_logits_indices), PyTorch
# creates a new tensor with separate storage from the original
# logits tensor. This means any in-place operations on bonus_logits
# won't affect the original logits tensor.
bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
sampler_output = self.sampler(logits=bonus_logits, ...)
bonus_token_ids = sampler_output.sampled_token_ids

# Just like `bonus_logits`, `target_logits` is a new tensor with
# separate storage from the original `logits` tensor. Therefore,
# it is safe to update `target_logits` in place.
target_logits = logits[spec_decode_metadata.target_logits_indices]
output_token_ids = self.rejection_sampler(
    spec_decode_metadata, None, target_logits, bonus_token_ids, sampling_metadata,
)

注释解释了一个 PyTorch 的微妙行为:用 tensor 做 indexing(logits[indices_tensor])会创建 separate storage——不是 view 而是 copy。这和用 slice logits[a:b] 不同(slice 是 view、共享 storage)。

这个差异在 spec decode 里是一等大事——两个采样器都会原地修改传进来的 logits:Sampler 做 temperature 缩放用 div_、做 top-k/top-p 截断用 masked_fill_(vllm/v1/sample/sampler.py:80、vllm/v1/sample/ops/topk_topp_sampler.py),rejection_sampler 对 target_logits 同样原地 div_(vllm/v1/sample/rejection_sampler.py:270-271)。bonus 采样先跑;如果 bonus_logits 与 logits 共享 storage,这些原地修改就会写回原始 logits,随后切出来的 target_logits 拿到的就是被改过的值,验证结果随之出错。

代码利用 PyTorch "tensor indexing creates copy"的默认行为、让 target_logits 天然独立——不用显式 clone。如果未来 PyTorch 改行为(比如优化成 view)、这段代码就要加 explicit .clone()。注释里把这个语义说清楚、就是在代码里写下的文档——告诉后来者"为什么这里不 clone"。

性能上、这个差异也重要——显式 clone 是额外一次内存拷贝、而利用 indexing 语义是"免费的 copy"(无论如何都要索引)。spec decode 热路径里每一个 matmul 和每一次 tensor indexing 都是被量过的、能免费得到的都不写多余代码。

8.10 调优的三个核心参数

8.10.1 -O / --compilation-config

V1 下的有效默认是 level 3(PIECEWISE)。保持默认,除非遇到 compile bug。-O0 关掉编译,-O 1 / -O 2 按官方 help 是内部测试用途。

8.10.2 --cudagraph-capture-sizes

控制录哪几个 token 数的 CUDA Graph。它不是独立的命令行 flag,而是 -O 的 JSON 字段,写法形如 -O '{"level": 3, "cudagraph_capture_sizes": [1, 2, 4, 8]}'(这串就是 --compilation-config 的 help 文本里给的例子)。

V1 的默认值是 [1, 2, 4] + range(8, 513, 8),再按 max_num_batched_tokens 截断(vllm/config.py:4031-4046)。

  • 如果实际 workload 的 token 数永远不超过 64,把列表截到 [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64],能明显缩短 capture_model() 的时间和 CUDA Graph 占的显存
  • 如果请求高度集中在某个区间,可以在那一段加密档位、其余段稀疏化,用更少的图覆盖主要分布
  • 记住 bs_to_padded_graph_size 是按 max_capture_size + 1 长度铺平的,列表最大值直接决定这张查表的长度

8.10.3 --enforce-eager

关闭 CUDA Graph;V1 下也不再把编译 level 自动升到 3(vllm/config.py:3884-3899 的条件是 not enforce_eager),未显式指定 -O 时即不做 torch.compile。主要用于开发调试、模型 bug 排查,或 §8.5.3 那类不在乎稳态性能、只想省启动时间的场景;在线服务一般不开。

8.10.4 is_multimodal_model 场景下 inputs_embeds 的额外拷贝

execute_model 里多模态分支有一段看起来"多余"的操作(line 1050-1063):

# vllm/v1/worker/gpu_model_runner.py:1050 起(省略了开头的 NOTE 注释)
if self.is_multimodal_model:
    input_ids = self.input_ids[:num_scheduled_tokens]
    if mm_embeds:
        inputs_embeds = self.model.get_input_embeddings(input_ids, mm_embeds)
    else:
        inputs_embeds = self.model.get_input_embeddings(input_ids)
    # TODO(woosuk): Avoid the copy. Optimize.
    self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
    inputs_embeds = self.inputs_embeds[:num_input_tokens]
    input_ids = None

看起来是:先算出 inputs_embeds、又把它 copy 到 self.inputs_embeds、再重新 slice 出来——这 copy_ 一步似乎多余。但源码的 TODO 注释告诉我们作者自己也知道这是技术债、但有原因:

原因是 CUDA Graph。self.inputs_embeds 是预分配的 persistent buffer、地址固定;而每次 get_input_embeddings 返回的 tensor 是动态分配的、地址每次不同。CUDA Graph 要求"所有 tensor 地址固定"、如果直接用动态分配的 inputs_embeds、CUDA Graph replay 就会错。

于是不得不把动态产物 copy 到固定地址 buffer——牺牲一次内存拷贝(几 MB)、换来下游所有 attention/MLP 都能 CUDA Graph 化。这条 TODO 指向未来改进方向——或许可以让 get_input_embeddings 直接写到预分配 buffer、避免这次 copy。

为什么纯文本模型不需要这一步?——因为纯文本模型的 inputs_embeds 是 embedding 层直接产生的、embedding 层本身就在 CUDA Graph 内、输出地址天然固定。多模态模型的 inputs_embeds 需要合并文本 token 的 embedding 和视觉/音频 encoder 的 embedding、不在一个简单 kernel 里、所以要额外走这个固定 buffer 的"桥"。

注释的另一行 For text-only models, we use token ids as input. While it is possible to use embeddings as input just like the multimodal models, it is not desirable for performance since then the embedding layer is not included in the CUDA graph.(line 1064-1068)——显式说明了"纯文本模型保留 token_ids 输入是为了让 embedding 层进 CUDA Graph"。这是两种代码路径各自的优化点、合并看似能简化代码但会损失性能。

这个小段揭示了 CUDA Graph 对代码结构的约束如何渗透到业务逻辑里——不只是"录制一下就完事"、而是整个 tensor 分配、数据流动的每一步都要考虑"地址是否固定"。这是高性能 LLM 推理的真实复杂性。

8.10.5 set_forward_context:把 attn_metadata 注入模型

前向计算那一行(line 1091-1097)看似平凡:

# vllm/v1/worker/gpu_model_runner.py:1091
with set_forward_context(attn_metadata, self.vllm_config):
    output = self.model(
        input_ids=input_ids,
        positions=positions,
        intermediate_tensors=intermediate_tensors,
        inputs_embeds=inputs_embeds,
    )

但 set_forward_context(attn_metadata, self.vllm_config) 是一个关键的依赖注入机制。模型 forward 接口里没有 attn_metadata 参数——attention 层怎么知道每个请求的 seq_len、block_table?

答案是 set_forward_context 把 attn_metadata 包进一个 ForwardContext dataclass、赋给 vllm/forward_context.py 里的模块级全局变量 _forward_context(进入时先把旧值存到 prev_context、finally 里再还原),attention 层内部调用 get_forward_context() 取出。注意它不是 threading.local() 也不是 contextvars——因为一个 worker 进程里 forward 是单线程串行的,全局变量 + 上下文管理器就够了。这就避免了"把 attn_metadata 透穿给每一层模型代码"——否则模型定义里每一层 forward 都要 attn_metadata: AttentionMetadata | None = None 参数、极其侵入。

这种"上下文变量 + 吸收点"模式,和《LangGraph 设计与实现》专栏里的 Runtime 注入、《Hyper 与 Tower:工业级 HTTP 栈》专栏里 Tower 的 layer 信息传递是同构的——跨层共享的运行时依赖不该走函数签名、该走上下文容器。

with 块保证 attn_metadata 的生命周期被精确管理——forward 完就从 context 里 pop 出去、下一 step 的 metadata 不会污染这一 step。这和 Python 的资源管理哲学(如 open(...) as f、torch.no_grad())一脉相承——用上下文管理器表达"这段时间有效"。

8.10.6 与本专栏其他章节的呼应

与第 4 章(PagedAttention)的呼应——本章的 slot_mapping、block_tables、cu_seqlens 都是给第 4 章的 Paged KV cache 和 FlashAttention 用的。第 4 章讲"KV 怎么分页存储"、本章讲"怎么告诉 attention kernel 每个请求对应哪些 page"。两章合起来是 vLLM 最核心的内存-计算协同。

与第 6 章(Worker 架构)的呼应——第 6 章讲 Worker 从 V0 无状态到 V1 有状态的变革、本章的 InputBatch 持久化正是这个变革的具体载体。理解了 Worker 为什么要有状态、再看 InputBatch 的 numpy 预分配策略就水到渠成。

与第 9 章(Sampler)的呼应——本章的 execute_model 最后产出 logits、然后立刻走 self.sampler(...)——第 9 章展开 Sampler 的内部实现。logits_indices 的构造(§8.6.4)是本章和第 9 章的接口——本章决定"算哪些位置的 logits"、第 9 章决定"怎么从 logits 采样 token"。

与《Hyper 与 Tower:工业级 HTTP 栈》专栏的呼应——vLLM 的 CompilationLevel 4 个分支和该专栏里 Steer 的 Picker 是同一个模式——"配置参数 → 选择完全不同的代码路径"。虽然在两个语言、两个领域、但设计哲学相同:用枚举 + 分支 builder 让参数控制能以零运行时开销映射到不同实现。

8.11 本章小结

ModelRunner 是 vLLM 从"调度指令"到"真正 GPU 计算"的最后一公里:

  • InputBatch 持久化:预分配槽位 + numpy 差量更新,避免每 step 重建;配合 ModelRunner 里那组供 CUDA Graph 使用的固定地址输入 buffer
  • _prepare_inputs:在 CPU 侧高效打包所有 GPU 输入;slot_mapping 决定新 token 写 KV 的位置;condense 让批里的请求连续排列,block_table 只拷、只切前 num_reqs 行
  • 分段 CUDA Graph (Piecewise):在 attention 边界切段,非 attention 段 CUDA Graph 化、attention 段 Eager;只需按 token 数录一组图,不必为每种 seq-len / block-table 组合各录一份
  • pad_for_cudagraph:bs_to_padded_graph_size 查表 O(1);8 以上档位步长 8,绝对 padding 浪费最多 7 个 token,换来每 step 耗时方差大幅收敛
  • torch.compile + CUDA Graph 是互补的两层:前者减少 kernel 数,后者省 launch 开销
  • execute_model 是所有章节的交汇点:10 个阶段覆盖 InputBatch、CUDA Graph、Multimodal、LoRA、PagedAttention、Sampler、SpecDecode 等所有主题
  • 三个调优参数:-O / --compilation-config (保持默认 level 3)、cudagraph_capture_sizes (作为 -O 的 JSON 字段定制)、--enforce-eager (主要用于调试)

下一章我们会进入"采样与输出处理"——第 9 章。execute_model 走完最后一步把 logits 张量交给 Sampler,从 [batch_size, vocab_size] 的浮点矩阵到一串 token id 之间还有一段精巧的工程:top-k / top-p 截断、temperature 缩放、speculative decoding 的接收-拒绝逻辑、stop token 判定。这一步发生在 GPU 上但又有强烈的 CPU-GPU 协同特征,是"从张量到产品"最容易被低估的环节。

8.12 ModelRunner 源码拆解的七个高阶设计原则

把本章所有源码级观察归纳成框架设计的原则:

① CPU-GPU 流水并行(§8.3.2)——能异步启动的 H2D 拷贝立刻发起、不阻塞 CPU 计算。"让两种硬件同时工作"是异构计算的第一定律。

② 向量化一切(§8.3.3)——热路径里的 Python loop 是性能之敌。np.repeat + cumsum 组合把 O(num_reqs) 次 Python 调用换成常数次 numpy 调用,源码注释只写 "faster than"、不承诺倍数。

③ 混用 torch 和 numpy(§8.3.4)——实用主义胜过风格一致性。torch.index_select 比 np.take 快就用 torch、对维护性的担心小于对性能的真实需求。

④ 分阶段状态更新(§8.2.3)——删除、释放、登记、更新、补入与压实顺序严格、互不耦合、保证"槽位的生产和消费在一个函数内闭合"。

⑤ 边界语义显式化(§8.2.4 的 abort-resubmit、§8.3.7 的 cascade prefix 为何不 +1)——把"看似相同但语义不同"的场景在代码里写清楚、注释里完整论证、让维护者看得懂"为什么这么写"。

⑥ 多配置组合的全覆盖(§8.8 的非 CUDA Graph 路径 + TP-SP 分支)——不简化"只考虑主流场景"、把每一种组合的正确性都算到位。

⑦ 利用语言语义避免显式操作(§8.9 的 tensor indexing 即 copy)——PyTorch 的 "tensor indexing creates new storage" 就是免费的 clone、不写多余代码、注释把语义依赖写清楚。

这七条原则在 _prepare_inputs 和 execute_model 两个函数里反复出现——两个函数合计 430 行左右(_prepare_inputs 488-631、execute_model 1003-1288)、承载了 vLLM V1 架构的全部热路径精髓。读懂它们、再看任何其他高性能推理引擎(SGLang、TGI、TensorRT-LLM)都能快速识别出"这家在优化什么"——不同团队的选择不同、但面对的问题空间是一样的。


源码导航

  • GPU ModelRunner:vllm/v1/worker/gpu_model_runner.py
  • InputBatch:vllm/v1/worker/gpu_input_batch.py
  • 编译配置:vllm/compilation/
  • Piecewise backend:vllm/compilation/backends.py
  • Attention Metadata 实现:vllm/v1/attention/backends/flash_attn.py(FlashAttentionMetadata / FlashAttentionMetadataBuilder)
  • V1 的 TPU 变体:vllm/v1/worker/tpu_model_runner.py(CPU / Neuron / HPU / XPU 变体在 v0.8.5 仍只有 V0 版本,位于 vllm/worker/)
  • forward context 注入:vllm/forward_context.py