vLLM 推理内核深度解析
第8章 ModelRunner:从 SchedulerOutput 到 GPU Kernel
第8章 ModelRunner:从 SchedulerOutput 到 GPU Kernel
“The fastest code is the code that never runs.” — Robert Galanakis
本章要点
- 看清 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 调优的三个核心参数:
cudagraph_capture_sizes、compilation_level、enforce_eager
8.1 ModelRunner 在整个栈里的位置
前 7 章走完后,我们已经看清 vLLM 的”指令流”:API Server 把请求翻成 EngineCoreRequest → EngineCore 的 Scheduler 产出 SchedulerOutput → Executor 广播给 Worker。现在 Worker 拿到 SchedulerOutput 之后,怎么把它变成真正的 GPU 计算?
答案是 ModelRunner(V1 里具体是 GPUModelRunner,另有 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:#fff,stroke:none
style MR fill:#3b82f6,color:#fff,stroke:none
style P3 fill:#10b981,color:#fff,stroke:none
style MO fill:#ec4899,color:#fff,stroke:none
ModelRunner 是”GPU 执行的最后一公里”——上游所有的抽象(调度、KV 管理、批处理决策)到这里都必须变成具体的 torch.Tensor 和 CUDA kernel。
8.1.0 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、拿到结果原样返回。Worker 的其他职责(CUDA 设备初始化、模型加载、KV cache 申请、determine_available_memory)都是”启动期”的、跟每 step 热路径没关系。
这种分层让 ModelRunner 可以独立测试——不用拉起整个 Worker、不用搞多进程分布式、单独喂一个 scheduler_output 就能跑 execute_model。这对 vLLM 团队的开发效率是巨大的——调 attention kernel 的时候不用起 200G 显存的集群、小模型本地就能测。
TPUModelRunner / CPUModelRunner / NeuronModelRunner 是平行存在的硬件变体——每个 Runner 的 _prepare_inputs 实现都不同(TPU 的 tensor 布局、CPU 的 numa 感知、Neuron 的 XLA 图),但 execute_model 的外部契约(接 scheduler_output、产 ModelRunnerOutput)是统一的。这是 vLLM 能横跨多种硬件的架构支撑——Runner 是硬件抽象的底层、上面所有调度、KV 管理、Sampling 都不感知具体硬件。
8.1.1 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 核心字段
# 概念性简化
class InputBatch:
# 核心槽位:max_num_seqs 个 "逻辑槽",每个对应一个 RUNNING 请求
req_id_to_index: dict[str, int] # 请求 id → 槽位下标
req_ids: list[Optional[str]] # 反向:槽位 → 请求 id
# per-request 的运行时状态(Numpy 数组,全是预分配)
num_computed_tokens_cpu_np: np.ndarray # 每请求算了多少 token
num_prompt_tokens: np.ndarray
num_tokens_no_spec: np.ndarray
# 输入张量(预分配到 max_num_seqs × max_model_len)
token_ids_cpu: np.ndarray # [max_num_seqs, max_model_len] int32
input_ids: torch.Tensor # GPU 侧 flat 输入
positions: torch.Tensor
slot_mapping_cpu: torch.Tensor
block_table_cpu: np.ndarray # [max_num_seqs, max_blocks_per_seq]
# Sampling params 也是 per-request 预分配
temperature_cpu: np.ndarray
top_p_cpu: np.ndarray
...
注意两个观察:
-
一切皆数组——没有
List[Request]这种 Python 对象数组。Python 对象的属性访问一次就是一次字典查找,在 per-step 要处理上百请求的热路径上会累积到可观开销。Numpy 数组的 vectorized 操作在 C 层面完成。 -
CPU + GPU 镜像——带
_cpu后缀的是 numpy 数组,不带后缀的是 torch.Tensor。每步在 CPU 上预先 fill 好 numpy 数组,最后一次性.to('cuda', non_blocking=True)搬到 GPU。避免多次小规模 H2D 传输。
8.2.2 三种变更操作
SchedulerOutput 告诉 ModelRunner 本 step 相对上一 step 有哪些变化——三类:
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
# 数据字段不需要清零,因为新请求会覆盖
三种操作都是 O(1) 或 O(k)(k = 本 step 新增 token 数),远低于 V0 的 O(batch_size × avg_seq_len)。
8.2.2.5 _update_states 的真实四段执行
§8.2.2 的三种变更操作是概念性说法、真实 _update_states 源码(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 及后)——新请求填进 self.requests 和 self.input_batch、优先用段 1、3 腾出的空 slot(removed_req_indices 里的)。这就是为什么 InputBatch 不需要”垃圾回收”——每步都在做紧凑化。
这四段的顺序不能乱——先删后加是硬约束、因为 max_num_seqs 是固定的、不先腾出 slot 就没法加新请求。每段之间的约束关系被代码结构直接表达——removed_req_indices 从段 1 和段 3 累积、段 4 消费——槽位的生产和消费在同一个函数内闭合。
8.2.2.6 “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 就被丢弃、用户看到的结果和预期完全不符。
这是真实生产系统里的用户可见问题——长连接断开重连、React 客户端 StrictMode 双触发、前端快速重试——都可能产生”相同 req_id 的两个独立语义请求”。分段处理保证了语义正确。
8.2.3 槽位复用:为什么不按时间线性推
一个容易被问的问题:为什么不让 req_ids 是一个动态 list(append / remove)?答案是:CUDA Graph 要求固定 shape。
如果槽位数量动态变化,每次 batch 数都可能不同,CUDA Graph 无法有效复用。V1 的做法是固定最大槽位数(max_num_seqs,典型 128-256),用”空槽位”占位。即使本 step 只有 10 个 RUNNING 请求,输入张量还是会 [max_num_seqs, ...]——浪费了点 memory,换来的是 CUDA Graph 可用性。
这种”预分配固定尺寸”的做法,贯穿 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 if rid is not None
])
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. 构造 AttentionMetadata(注意力 kernel 需要的所有元信息)
attn_metadata = AttentionMetadata(
num_prefills=...,
num_decode_tokens=...,
slot_mapping=self.slot_mapping,
block_tables=self.block_tables,
context_lens=...,
cu_seqlens_q=...,
cu_seqlens_k=...,
)
return attn_metadata, logits_indices, ...
每一步都有可以细抠的工程细节。我们挑两个最关键的讲:
8.3.1.5 block_table.commit 的 CPU-GPU 并行化
真实的 _prepare_inputs(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 拷贝特别重要——它是 [max_num_seqs, max_blocks_per_seq] 的大张量(典型 128 × 1024 int32 = 512KB)、同步拷贝会阻塞几十微秒。把它提前、overlap 在 CPU 的 numpy 计算里、每 step 节省这几十微秒。
这种优化在其他框架里见得多(PyTorch 自己的 DataLoader 的 pin_memory + non_blocking)、但 vLLM 把它做到了 _prepare_inputs 这种每 step 都跑的热路径里、意识到”拷贝先发、计算后来”才是最短时间路径。
8.3.1.6 np.repeat + cumsum 的向量化 req_indices 构造
_prepare_inputs 里一段漂亮的 numpy 向量化代码(line 510-523):
# 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 层完成。对 batch=128 的 decode step、这段从 ~200 μs 压到 ~5 μs、省下 40 倍。
“拒绝 Python loop、用 numpy 向量化一切”是 _prepare_inputs 的核心设计原则。这不是 premature optimization——因为这段代码每 step 都跑、一次推理 1000 step 的话就是 1000 次调用、省下的总时间是秒级的。
8.3.1.7 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 在 CPU 上走的是 ATen 的高度优化实现、对内存连续的大张量用了 SIMD(AVX2 / AVX512)。np.take 虽然也优化、但走的是 NumPy 自己的索引路径、对大张量的 bandwidth 没那么极致。
对每 step 都要做一次的 “从 [max_num_seqs × max_model_len] 的大张量里挑出本 step 的 tokens” 操作、max_model_len 可能到 32K/128K、max_num_seqs 到 256——这是一个几百 MB 的张量、SIMD vs 普通索引的差距能有 2-3 倍。
vLLM 团队愿意在这里放弃”全 numpy 一致性”、混用 torch 和 numpy——这种实用主义是性能工程的标志。注释里明确写出”为什么选 torch”——让未来维护者知道这不是手滑、是有意的。
8.3.1.8 non_blocking=True + pin_memory 的 H2D 传输链
_prepare_inputs 末尾几行都是异步 H2D 传输(line 574-586):
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)。InputBatch 里所有 _cpu 前缀的 tensor 都是 pinned memory(构造时 torch.zeros(..., pin_memory=True))——这是能真正异步拷贝的前提。
这两个参数必须成对出现:pin_memory=True + non_blocking=True。缺一不可——单独用 pin_memory 还是同步、单独用 non_blocking 会 silently 退化。vLLM 团队把所有热路径的 CPU buffer 都改成 pinned——这是从 V0 到 V1 的重要工程整改。
这段代码和 §8.3.1.5 的 block_table.commit 同步发生——CPU 继续算 attn_metadata_builder.build(...)、GPU 的拷贝引擎在并行搬多组数据。_prepare_inputs 返回时、所有 GPU tensor 都已经或即将就绪、紧接着 forward 开始时 CUDA stream 的顺序性会保证拷贝先于 attention kernel 完成。
这种”pinned + async”是 GPU 编程的标准姿势、vLLM 把它用到了极致——所有 numpy 数组背后都有对应的 pinned torch buffer、所有 H2D 都是 non_blocking。每个微小的细节都被测过、都在热路径上。
8.3.2 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
这是纯 CPU 的 numpy 向量化操作。对 batch=128、decode-only step 来说只有 128 个 slot 要算,用 Python loop 也够快;但 prefill step 可能一次处理几千个 token,vectorize 起来就很关键。
8.3.2.5 Cascade Attention 的 “最小 num_computed_tokens - 1” 规则
prefill 场景下、多个请求可能共享很长的 common prefix(比如 system prompt、few-shot examples)。vLLM 的 Cascade Attention 把这段共享前缀的 KV只算一次——用两段 attention kernel:第一段算共享前缀、第二段算各请求的独立后缀。
但 _compute_cascade_attn_prefix_len(gpu_model_runner.py:633-708)有一个看起来反直觉的规则:
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 668-694):
请求 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。
结论:common_prefix_len 要 cap 到 min(num_computed_tokens) - 1、再加上 block-size 对齐。这是对”哪些 KV 已经在 cache 里”的严格推导——不能把”某些请求还在 prefill、某些已经 decode”的不对齐 KV 错当成 common。
为什么是 -1 而不是 +1(包含)?——因为 attention 的因果性:query 的第一个 token 必须能 attend 到自己对应的 KV、而不是只 attend 到前面的 common prefix。如果 common prefix 覆盖了 query 的第一个 token 的 KV、第一段 kernel(无 mask)就会错乱因果顺序。
这种细腻的语义推理是 LLM 推理引擎的精髓——一个减号错了、模型输出就不对了、但肉眼看不出(因为 logits 数值还是”合理”的)。vLLM 把推导直接写在注释里、让后来者能 trace 为什么是这么写。
8.3.2.6 M-RoPE 和 1D positions 的二选一分支
_prepare_inputs 的 positions 处理有一条双分支(line 577-586):
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 等模型用的旋转位置编码变体——对 image token 使用2D 坐标(height, width)、对 video token 使用3D 坐标(time, height, width)。普通文本 token 只用 1D 绝对位置。
普通模型的 positions 是 [total_tokens] 的一维张量、而 mrope_positions 是 [3, total_tokens] 的二维张量(三行分别是 t/h/w)——shape 不兼容、必须分开维护。
uses_mrope 在 Worker 初始化时就根据 model_config.hf_config.rope_scaling 是否有 “mrope” 类型来决定——一次性决定、之后每 step 走同一条路径。这避免了热路径里每 step 都做”这是不是 M-RoPE 模型”的判断。
多模态 + M-RoPE 的组合让 _prepare_inputs 有一整段 _calc_mrope_positions 的额外逻辑(§_update_states 里 line 361 起约 25 行)——从 mm_inputs 里提取 image_grid_thw、video_grid_thw、second_per_grid_ts 等、拼成 3D 坐标张量。这部分代码对纯文本 / 纯 CLIP-style 多模态模型(比如 LLaVA)都是死代码——只在 Qwen2-VL 这类模型下才激活。
这种”一条分支服务一个特定模型家族”是 vLLM 源码里常见的模式——模型架构变化快、新的 RoPE 变体、MLA、MoE 不断出现、引擎必须能快速添加分支而不破坏主路径。M-RoPE 就是这样一个”局部添加、对主路径零影响”的典范——热路径里多一个 if、特性完全支持。
8.3.3 Block Table 打包:二维张量里的 padding
block_table 在 InputBatch 里是 [max_num_seqs, max_blocks_per_seq] 的 numpy 数组。_prepare_inputs 要把它搬到 GPU,同时只保留本 step scheduled 的请求那几行。
一个 subtlety:本 step 可能只有 12 个请求 scheduled(out of 128 max),但为了 CUDA Graph 固定 shape,block_table 张量还是要 pad 到固定 num_seqs。具体做法:
# block_tables tensor shape 永远是 [max_num_seqs, max_blocks_per_seq]
# 未 scheduled 的行用 0 填充(kernel 里会跳过)
valid_rows = [req_id_to_index[rid] for rid in scheduled_req_ids]
block_tables_gpu[:len(valid_rows)] = block_table_cpu[valid_rows]
# 剩下的行是历史残留数据,kernel 不看
Kernel 通过 cu_seqlens 知道 batch 的真实边界,跳过 padding 行。
8.3.3.5 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 都发生、只在”batch 里 LoRA 组合变化”时才真正切换权重(内部有 LRU 缓存)。如果连续几 step 都是同一批 LoRA、set_active_loras 几乎是 no-op(只做哈希比较)。这就让 LoRA 在稳态下零额外开销、只在 LoRA 变动时付一次拷贝成本。
第 16 章会详讲 LoRA 的实现细节(Punica CUDA kernel、SGMV 打包调度等)、本章这一行就是LoRA 和主 forward 接头的那一点——简单、不显眼、但承载着多租户推理的全部成本控制。
8.3.4 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 的机器。KVConnector 接口让不同的传输后端(TCP、RDMA、NCCL)都能接入同一个架构。
为什么在 execute_model 最开始而不是 _update_states 之后?——因为 KV 传输可能要和 forward 并行——在 CPU 这边做 _update_states 和 _prepare_inputs 的时候、底层的 RDMA 引擎已经在搬上一 step 的 KV 了。这样”计算和通信重叠”、网络延迟被掩盖在计算时间里。
绝大部分单机部署的用户不会触碰这个挂钩——has_kv_transfer_group() 返回 False、整段代码 short-circuit。但对 PD disaggregation 用户(比如大规模在线服务把 prefill GPU 和 decode GPU 分开)、这个挂钩是生死攸关的性能路径。
这又是一个”主路径对少数场景的开放点”——像 LoRA、spec decode、multimodal 一样、KV transfer 以最小侵入的方式嵌入到 execute_model 里、对不用它的人零开销、对用它的人提供完整功能。这种扩展点的克制是 vLLM 源码能保持可读性的重要原因——主路径始终清晰、扩展功能通过”空实现兜底 + 注入激活”融进来。这种模式和 §8.3.3.5 的 LoRA、§8.3.2.6 的 M-RoPE 一脉相承、构成了 vLLM 可扩展性的核心设计。不管是多租户、新模态、新传输方式——都按这套”零侵入接入”的路子走。
8.4 分段 CUDA Graph (Piecewise) 的动因
8.4.1 标准 CUDA Graph 的痛点
CUDA Graph 把一串 kernel launch 录制成一个图,replay 只需要一次 API 调用——省下每 kernel 5-10 μs 的 launch overhead。对每 step 要 launch 500+ kernel 的 LLM decode 来说,这省的是几毫秒。
但标准 CUDA Graph 有两个硬约束:
- 所有 shape 必须固定——录的时候多大,replay 就多大
- 内存地址必须固定——不能分配新 tensor
约束 1 对 LLM 特别麻烦:batch 大小可变(用户数波动)、seq len 可变(每请求不同 context)、prefill vs decode 形态差异巨大……
V0 的做法是为每个可能的批大小录一张图。假设 max_num_seqs=128,要录 128 张图——显存大量浪费、启动慢。
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:#ef4444,color:#fff,stroke:none
style E fill:#10b981,color:#fff,stroke:none
style LM fill:#10b981,color:#fff,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 算子走 Eager 模式——FlashAttention 本身已经高度优化,额外的 kernel launch 不显著
- 段之间通过共享 GPU tensor 传递 hidden state
graph TB
subgraph "标准 CUDA Graph"
SG[整个模型<br/>对每个 batch_size × seq_len 组合<br/>都要录一张图<br/>组合爆炸]
end
subgraph "Piecewise CUDA Graph"
P1[Segment 1<br/>Embedding + Layer 1 的非 attn 部分<br/>只依赖 total_tokens<br/>录 13 张图 × 13 个 size]
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:#ef4444,color:#fff,stroke:none
style P1 fill:#10b981,color:#fff,stroke:none
style P2 fill:#f59e0b,color:#fff,stroke:none
style P3 fill:#10b981,color:#fff,stroke:none
style P4 fill:#f59e0b,color:#fff,stroke:none
8.4.3 实际效果
vLLM 的 Piecewise 实现(vllm/compilation/backends.py + torch.compile 集成)对 13 个 batch_size(默认 [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256])分别录图。
性能对比(Llama-8B,A100):
| 模式 | single-request decode | batch=32 decode | startup 时间 |
|---|---|---|---|
| Eager (no CUDA Graph) | 25 ms/step | 45 ms/step | 15 s |
| Full CUDA Graph (V0) | 18 ms/step | 30 ms/step | 120 s(录 32 张图) |
| Piecewise (V1) | 19 ms/step | 32 ms/step | 45 s |
Piecewise 牺牲一点点 steady-state 性能,换来 2-3× 的 startup 加速 + 大幅减少的显存占用(CUDA Graph 本身要存 captured tensor copies)。在生产部署里这个 trade-off 极度划算——startup 快意味着 K8s 副本拉起快,故障恢复快。
8.5 pad_for_cudagraph:padding 的经济学
前面说过 CUDA Graph 需要固定 shape。实际 batch 的 total_tokens 可能是任意值,比如 47。但预录制的 size 只有 [1, 2, 4, 8, 16, 24, 32, 48, 64, ...]——怎么办?
pad_for_cudagraph(47) = 48——往上 pad 到最近的预录制 size。
# gpu_model_runner.py
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 mask / valid_tokens index 跳过,不影响正确性,只是多跑了一些无用 FLOPs。
8.5.0.5 pad_for_cudagraph 的实现:二分查找最接近的上界
pad_for_cudagraph(47) 怎么知道该 pad 到 48?实现在 vllm/config.py 里的 VllmConfig.pad_for_cudagraph:
def pad_for_cudagraph(self, batch_size: int) -> int:
# cudagraph_capture_sizes is a sorted list
sizes = self.compilation_config.cudagraph_capture_sizes
# find smallest size >= batch_size
idx = bisect.bisect_left(sizes, batch_size)
return sizes[idx]
关键是 bisect.bisect_left——对有序列表做二分查找、找到”第一个 >= batch_size 的元素”的位置。cudagraph_capture_sizes 是构造时就排序好的列表(默认 [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256])——所以每次调用都是 O(log n) 而不是线性扫描。
为什么用 bisect 而不是 min(s for s in sizes if s >= batch_size)?——后者是 O(n)、前者是 O(log n)。对 13 个元素的列表这个差别是 4 步 vs 13 步 Python 比较——在每 step 都调用一次的热路径里、省下 9 次 Python 比较就是 ~1 μs 级别的节省。十万次调用就是秒级的差异。
bisect 是 Python 标准库里最被低估的高性能工具——很多**“显然应该 loop 一下**“的场景、bisect 能把它变成 O(log n)。vLLM 这里用得很得当:capture_sizes 是不变的、初始化完就是只读列表、二分查找最合适。
bisect_left vs bisect_right 的选择——bisect_left(sizes, 32) 返回的是”能放下 32 的最小 size 的 idx”、当 32 in sizes 时返回的是 32 自己那个位置(而不是下一个)。这保证 pad_for_cudagraph(32) == 32——如果 batch 正好等于某个 capture size、零 padding 浪费、直接用。
8.5.1 Padding 的开销分析
假设预录制 size = [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256]。实际 token 数的 padding 浪费率:
- total_tokens=32 → pad 到 32, 0% 浪费
- total_tokens=33 → pad 到 48, 31% 浪费
- total_tokens=47 → pad 到 48, 2% 浪费
- total_tokens=97 → pad 到 128, 24% 浪费
- total_tokens=129 → pad 到 192, 33% 浪费
最坏浪费 ~33%,平均约 15-20%。听起来挺高?来算 trade-off:
- 不用 CUDA Graph:每 kernel launch ~10μs × 500 kernel = 5 ms 开销
- 用 CUDA Graph pad 20%:多算 20% 的 FLOPs ≈ 多 3-5 ms
两者级别相当。但 CUDA Graph 让总 latency 的方差小得多——每 step 时间稳定可预测,对 P99 latency SLA 极友好。而不用 CUDA Graph,launch overhead 随系统负载波动明显。
8.5.2 不值得 CUDA Graph 化的场景
- batch 大小经常 > 最大预录制 size → fall back 到 Eager,没有收益
- 用了 guided decoding → logit mask 每步不同,CUDA Graph 失效,走 Eager
- 开了 Spec Decode → verify 阶段的 token 数高度动态
- 启动时就明确短期内不需要 steady-state 性能(比如一次性的 offline batch inference)
这些场景下可以用 --enforce-eager 完全关掉 CUDA Graph,节省 startup 时间。
8.6 torch.compile 的角色:两层优化的协同
V1 默认启用 torch.compile。它和 CUDA Graph 不是竞争关系,是互补层次:
graph TB
subgraph "torch.compile 层(Python 侧优化)"
TC1[算子融合<br/>linear + bias + gelu → 一个 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:#3b82f6,color:#fff,stroke:none
style TC2 fill:#3b82f6,color:#fff,stroke:none
style TC3 fill:#3b82f6,color:#fff,stroke:none
style CG1 fill:#10b981,color:#fff,stroke:none
torch.compile 让每个 step 的算子数变少(融合后从 500+ 降到 200+),CUDA Graph 把剩下这 200+ kernel 的 launch overhead 打包省掉。两者叠加后 steady-state 时间可以再降 30-50%。
8.6.0.5 set_forward_context 与 CUDA Graph 的协作细节
§8.8.5 会看到 set_forward_context(attn_metadata, self.vllm_config) 把元数据注入到模型 forward。但这里有一个 CUDA Graph 相关的微妙点:
CUDA Graph 录制时——attn_metadata 是一个具体的 AttentionMetadata 实例、里面的 block_tables、cu_seqlens、slot_mapping 都是具体的 GPU tensor 地址。CUDA Graph 会把这些地址编译到图里、replay 时直接用。
replay 时不同 shape 的 attn_metadata 怎么办?——答案是 attn_metadata 内部的 tensor 本身是持久化的固定地址 buffer(在 InputBatch 里预分配)、只是每 step 往这些 buffer 里写不同的值。地址不变、内容变——CUDA Graph 看到的是同一块内存、replay 时读取的是新写入的数据。
这是”预分配 + 就地更新”模式的深层价值——不只是避免分配开销、还让 CUDA Graph 能重复使用。每个 AttentionMetadata 字段都是 InputBatch 的一个 slice、slice 的底层 storage 在 Worker 生命周期内不变。
对比如果每 step 都 torch.zeros(...) 重建 attn_metadata——每次新的 tensor 地址、CUDA Graph 会 silently 用到错误地址、产生不可预测的行为(或者 assert、或者更糟——读到脏数据算出错答案)。vLLM 在整个架构设计上强制所有热路径 tensor 都是预分配 buffer 的 slice——这是 CUDA Graph 能安全 replay 的前提。
8.6.1 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 默认)
--compilation-level 3 是默认、也是推荐。只有调试模型 bug / 想看原生 PyTorch 行为时才降级到 0 或 1。
8.6.2 四个 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 不做的事:
- 图切分:把整个模型 forward 图按”能进 CUDA Graph 的 piece”切开——遇到动态 shape、非法控制流(比如 attention 里基于 seq_len 的分支)就在那里切断、每段独立处理
- PostGradPassManager 插入:在 Inductor 的 post-grad 优化点插入 vLLM 自己的 pass(比如针对 LoRA 的融合、针对特定 attention kernel 的替换)
- CUDA Graph 整合:把每段编译后的代码自动录制成 CUDA Graph——用户什么都不用写,level=3 就都有了
级别选择的实用指南:
- 写 / 调新模型:用 level 1。看 Dynamo 是否能追踪、有多少 graph break
- 生产推理:用 level 3(默认)。前面表格说 padding + torch.compile 两层协同能再降 30-50% 延迟
- 某些诊断场景:用
--enforce-eager直接跳到 level 0——例如你怀疑是 Inductor 编译出了错的代码、想验证裸 PyTorch 行为时
Level 2 在实践中很少用——它是”和 level 1 类似但缓存一次”,对 vLLM 这种稳态运行的服务用处不大。它存在主要是为了让 torch.compile 开发者能对比 “每次追踪” vs “只追踪一次” 的行为差异、做 debug。
这 4 层 level 和第 13 章讲的量化注册表、第 14 章讲的 GroupCoordinator 一样——表面上是个配置枚举、背后各自对应一套完全不同的代码路径。这种”统一参数接口 + 分支化实现” 是 vLLM 能快速迭代新编译特性的关键架构——新加一种编译方式就是往这里加一个 level 和对应的 backend,核心 ModelRunner 代码零改动。
8.6.3 use_spec_decode 分支里的 logits_indices
_prepare_inputs 的尾部有一个对 speculative decoding 的特殊处理(line 603-625):
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.
# ...
# Maintenance note: 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 负责构造这些索引。
维护注释:Support prompt logprobs——说明当前实现不支持”对 prompt 里每个 token 都计算 logprob”(用户调试用的功能)。代码在 non-spec-decode 分支里显式丢弃了 partial request 的 sample(“we will ignore the sampled tokens from the partial requests”)——这是工程折衷:支持 chunked prefill 让内存稳定、但牺牲了 prompt logprobs。这条维护标记留给未来版本补齐。
8.7 execute_model:完整主流程
把前面所有模块串起来,一个典型 step 的完整调用栈:
def execute_model(self, scheduler_output) -> ModelRunnerOutput:
# 1. 更新 InputBatch 状态(差量)
self._update_states(scheduler_output)
if self.input_batch.num_reqs == 0:
return ModelRunnerOutput.empty()
# 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. 多模态挂钩
if self.is_multimodal_model:
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
inputs_embeds = self.model.get_input_embeddings(
self.input_ids[:num_input_tokens], mm_embeds,
)
else:
inputs_embeds = self.model.get_input_embeddings(
self.input_ids[:num_input_tokens],
)
# 5. LoRA 挂钩(如启用)
if self.lora_config is not None:
self._set_active_loras(scheduler_output.lora_requests)
# 6. 模型前向(真正的 GPU 计算)
with set_forward_context(attn_metadata, ...):
hidden_states = self.model(
input_ids=None, # 已经转成 embeddings
positions=self.positions[:num_input_tokens],
inputs_embeds=inputs_embeds,
kv_caches=self.kv_caches,
attn_metadata=attn_metadata,
)
# 7. 计算 logits(只对"需要输出 token"的位置)
# logits_indices 是 [num_sampled_tokens] 的下标
logits = self.model.compute_logits(
hidden_states[logits_indices],
sampling_metadata=None,
)
# 8. 采样(第 9 章详细讨论)
sampler_output = self.sampler(
logits=logits,
sampling_metadata=self.input_batch.sampling_metadata,
)
# 9. 投机解码挂钩(如启用)
if self.use_spec_decode:
spec_token_ids = self.drafter.propose(...)
else:
spec_token_ids = None
# 10. 打包返回
return ModelRunnerOutput(
sampled_token_ids=sampler_output.sampled_token_ids.cpu().tolist(),
spec_token_ids=spec_token_ids,
logprobs=sampler_output.logprobs_tensors,
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.7.1 execute_model 里”eager mode 下还要 round_up to tp_size”的冷门分支
execute_model 的 padding 决策里(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
当 enforce_eager + sequence parallelism + tp > 1 时——还是要把 num_input_tokens 向上 round 到 tp_size 的倍数。为什么?因为 sequence parallelism 的 collective fusion(把 all-reduce/reduce-scatter 和 compute 融合)要求 token 数能整除 tp_size——否则切分不均匀、fusion pass 会 assert fail。
这个分支在大部分用户那里用不到——enforce_eager 是罕见的调试 flag、sequence parallelism 又是 advanced 优化——但一旦开了、不做这个 round_up 就会silent 错误(最后一个 rank 拿到的 token 数不对、但不 assert、直接算错)。代码在enforce_eager 路径里仍然做一小部分 padding、保证多机多卡场景的正确性。
这体现了 vLLM 对”多配置组合”的覆盖——不是简化成”CUDA Graph 才 pad”、而是”CUDA Graph 和 TP-SP 都可能要 pad、分别判断”。配置矩阵越大、这种分支就越细。
8.7.2 spec_decode 的 bonus/target logits 独立存储
execute_model 的 spec decode 路径(line 1124-1146)里有一段对 PyTorch tensor 语义的深度利用:
# 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 里是一等大事——rejection_sampler 会原地修改 target_logits(比如做 top-k mask、temperature scaling)。如果 target_logits 和 logits 共享 storage、这些原地修改就会污染原始 logits、让 bonus 采样结果错乱。
代码利用 PyTorch “tensor indexing creates copy”的默认行为、让 target_logits 天然独立——不用显式 clone。如果未来 PyTorch 改行为(比如优化成 view)、这段代码就要加 explicit .clone()。注释里把这个语义说清楚、就是在代码里写下的文档——告诉后来者”为什么这里不 clone”。
性能上、这个差异也重要——显式 clone 是额外一次内存拷贝、而利用 indexing 语义是”免费的 copy”(无论如何都要索引)。spec decode 热路径里每一个 matmul 和每一次 tensor indexing 都是被量过的、能免费得到的都不写多余代码。
8.8 调优的三个核心参数
8.8.1 --compilation-level
默认 3 (PIECEWISE)。保持默认除非遇到 compile bug。
8.8.2 --cudagraph-capture-sizes
控制录哪几个 size 的 CUDA Graph。默认 [1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128, 192, 256]。
- 如果实际 workload 的 batch 永远不超过 64,可以改成
[1, 2, 4, 8, 16, 24, 32, 48, 64],省 startup 时间 - 如果大量请求聚集在某个特定 size(比如 48),把相邻 size 加密
[1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, ...],降低 padding 浪费率
8.8.3 --enforce-eager
关闭所有 CUDA Graph 和 torch.compile。只在开发调试、模型 bug 排查时用。生产永远不要开。
8.8.4 is_multimodal_model 场景下 inputs_embeds 的额外拷贝
execute_model 里多模态分支有一段看起来”多余”的操作(line 1050-1063):
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)
# Maintenance note: avoid the copy and optimize this path.
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_ 一步似乎多余。但源码维护注释告诉我们作者自己也知道这是技术债、但有原因:
原因是 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 化。注释里的维护标记指向未来改进方向——或许可以让 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.8.5 set_forward_context:把 attn_metadata 注入模型
前向计算那一行(line 1091-1097)看似平凡:
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 存进线程局部变量(Python 的 threading.local() 或 contextvars)、attention 层内部调用 get_forward_context() 取出。这就避免了”把 attn_metadata 透穿给每一层模型代码”——否则模型定义里每一层 forward 都要 attn_metadata: AttentionMetadata | None = None 参数、极其侵入。
这种”上下文变量 + 吸收点”模式和 LangGraph 的 Runtime(第 14 章讲过)、Tower 的 layer 信息传递是同构的——跨层共享的运行时依赖不该走函数签名、该走上下文容器。三种不同语言/场景的框架、解决同一个问题、都是这套机制——说明这是跨语言的通用模式。
with 块保证 attn_metadata 的生命周期被精确管理——forward 完就从 context 里 pop 出去、下一 step 的 metadata 不会污染这一 step。这和 Python 的资源管理哲学(如 open(...) as f、torch.no_grad())一脉相承——用上下文管理器表达”这段时间有效”。
8.8.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.3)是本章和第 9 章的接口——本章决定”算哪些位置的 logits”、第 9 章决定”怎么从 logits 采样 token”。
与 hyper-tower 第 8 章(Filter/Steer)的呼应——vLLM 的 CompilationLevel 4 个分支和 Steer 的 Picker 是同一个模式——“配置参数 → 选择完全不同的代码路径”。虽然在两个语言、两个领域、但设计哲学相同:用枚举 + 分支 builder 让参数控制能以零运行时开销映射到不同实现。
8.9 本章小结
ModelRunner 是 vLLM 从”调度指令”到”真正 GPU 计算”的最后一公里:
- InputBatch 持久化:预分配槽位 + numpy 差量更新,避免每 step 重建;配合 CUDA Graph 要求的固定 shape
- _prepare_inputs:在 CPU 侧高效打包所有 GPU 输入;slot_mapping 决定新 token 写 KV 的位置;block_tables 用固定 shape 张量 + kernel 侧跳过 padding
- 分段 CUDA Graph (Piecewise):在 attention 边界切段,非 attention 段 CUDA Graph 化、attention 段 Eager;用 13 张图代替 V0 的 N 张图
- pad_for_cudagraph:平均 15-20% padding 浪费,换 P99 latency 稳定——划算
- torch.compile + CUDA Graph 是互补的两层:前者减少 kernel 数,后者省 launch 开销
- execute_model 是所有章节的交汇点:10 个阶段覆盖 InputBatch、CUDA Graph、Multimodal、LoRA、PagedAttention、Sampler、SpecDecode 等所有主题
- 三个调优参数:
compilation-level(默认)、cudagraph-capture-sizes(定制)、enforce-eager(只调试用)
下一章我们会进入”模型从磁盘到 GPU”——第 7 章 模型加载。这是启动时间的关键路径,也是一个很少被讨论但工程上非常精巧的领域。
8.10 ModelRunner 源码拆解的七个高阶设计原则
把本章所有源码级观察归纳成框架设计的原则:
① CPU-GPU 流水并行(§8.3.1.5)——能异步启动的 H2D 拷贝立刻发起、不阻塞 CPU 计算。“让两种硬件同时工作”是异构计算的第一定律。
② 向量化一切(§8.3.1.6)——热路径里的 Python loop 是性能之敌。np.repeat + cumsum 组合比”for 循环 + concat”快 40×。
③ 混用 torch 和 numpy(§8.3.1.7)——实用主义胜过风格一致性。torch.index_select 比 np.take 快就用 torch、对维护性的担心小于对性能的真实需求。
④ 分阶段状态更新(§8.2.2.5)——删除、释放、补充、添加四段顺序严格、互不耦合、保证”槽位的生产和消费在一个函数内闭合”。
⑤ 边界语义显式化(§8.2.2.6 的 abort-resubmit、§8.3.2.5 的 cascade prefix -1)——把”看似相同但语义不同”的场景在代码里写清楚、注释里完整论证、让维护者看得懂”为什么这么写”。
⑥ 多配置组合的全覆盖(§8.7.1 的 enforce_eager+TP-SP 分支)——不简化”只考虑主流场景”、把每一种组合的正确性都算到位。
⑦ 利用语言语义避免显式操作(§8.7.2 的 tensor indexing 即 copy)——PyTorch 的 “tensor indexing creates new storage” 就是免费的 clone、不写多余代码、注释把语义依赖写清楚。
这七条原则在 _prepare_inputs 和 execute_model 两个函数里反复出现——两个函数合计 600 多行、承载了 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/__init__.py- TPU / CPU / Neuron ModelRunner 变体:
vllm/v1/worker/tpu_model_runner.py/cpu_model_runner.py/neuron_model_runner.py