《vLLM 内核探秘》完整目录
- 前言
- 第1章 架构总览
- 第2章 EngineCore：引擎的心脏
- 第3章 调度器：Token 的交通指挥
- 第4章 PagedAttention：虚拟内存的启示
- 第5章 KV Cache 管理：寸土寸金的显存
- 第6章 Worker 与 Executor：GPU 军团
- 第7章 模型加载与权重管理
- 第8章 前向计算与 CUDA Graph（当前）
- 第9章 采样与输出处理
- 第10章 前缀缓存：零开销的加速
- 第11章 分块预填充与混合批处理
- 第12章 投机解码：以小博大
- 第13章 量化引擎：精度与速度的平衡
- 第14章 张量并行与流水线并行
- 第15章 多模态推理
- 第16章 LoRA 适配器热切换
- 第17章 API 服务器与生产部署
- 第18章 设计模式与架构哲学
第8章 前向计算与 CUDA Graph
"The fastest code is the code that never runs." -- Robert Galanakis
:::tip 本章要点
- 理解 ModelRunner 在 Worker 内部的角色与职责
- 掌握持久化批次（Persistent Batch）模式：为什么用 NumPy 而非 Python 原生操作
- 深入分段 CUDA 图（Piecewise CUDA Graph）的设计动机与实现
- 理解 torch.compile 集成：自动优化 vs 手写内核
- 认识 GPU 前向传播的完整数据流：从 Token ID 到 Logits :::
8.1 ModelRunner 的角色
ModelRunner 是 Worker 内部的核心组件，负责驱动模型的前向传播。如果说 Worker 是士兵，ModelRunner 就是士兵手中的武器。
每一步推理，ModelRunner 的工作流程是：
Token IDs, Positions, Block Tables"] --> B["2. 执行模型前向传播"] B --> C["3. 计算 Logits"] C --> D["4. 采样下一个 Token"] D --> E["5. 更新 KV Cache 元数据"] style A fill:#f59e0b,color:#fff,stroke:none style B fill:#3b82f6,color:#fff,stroke:none style C fill:#8b5cf6,color:#fff,stroke:none style D fill:#10b981,color:#fff,stroke:none style E fill:#ec4899,color:#fff,stroke:none
步骤 2 是最耗时的------它触发了 GPU 上所有 Transformer 层的计算。但步骤 1（准备输入）和步骤 4（采样）都是 CPU 操作，在高并发场景下也可能成为瓶颈。V1 的两个关键优化------持久化批次和分段 CUDA 图------分别针对这两个 CPU 瓶颈。
8.2 GPUModelRunner 初始化：预分配的策略
源码版本 ：本节基于 vLLM v0.8.5，核心文件
vllm/v1/worker/gpu_model_runner.py。
GPUModelRunner（gpu_model_runner.py:63）在初始化时做了大量预分配工作，为后续每步推理的高效执行打下基础：
python
# gpu_model_runner.py:63-112 (关键字段)
class GPUModelRunner(LoRAModelRunnerMixin):
def __init__(self, vllm_config: VllmConfig, device: torch.device):
# ... 配置读取 ...
self.max_num_tokens = scheduler_config.max_num_batched_tokens
self.max_num_reqs = scheduler_config.max_num_seqs
self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)
# 核心：持久化批次对象
self.input_batch = InputBatch(
max_num_reqs=self.max_num_reqs,
max_model_len=self.max_model_len,
max_num_blocks_per_req=self.max_num_blocks_per_req,
device=self.device,
pin_memory=self.pin_memory,
)
# CUDA Graph 配置
self.use_cuda_graph = (
self.vllm_config.compilation_config.level
>= CompilationLevel.PIECEWISE
)
self.cudagraph_batch_sizes = list(reversed(
self.vllm_config.compilation_config.cudagraph_capture_sizes
))InputBatch 是持久化批次的核心数据结构------它在 GPU 上预分配了所有输入张量，后续每步只做差量更新。cudagraph_batch_sizes 定义了预录制 CUDA 图的批大小集合。
8.3 持久化批次模式
问题：输入准备的 CPU 开销
每一步推理前，ModelRunner 需要将调度结果转换为 GPU 可以消费的张量：
- Input IDs：本步所有请求的 Token ID 拼成一维张量
- Positions：每个 Token 的位置编码
- Block Tables：每个请求的块表（物理块 ID 数组）
- Attention metadata：序列长度、预填充/解码标记等
在 V0 中，这些张量每一步都从头构造。假设有 100 个并发请求，大部分是解码请求（每步只加 1 个 Token），但 V0 仍然要重新组装完整的 100 个请求的输入。这涉及大量 Python 列表操作、类型转换和 GPU 传输。
解法：缓存 + 差量更新
V1 的 ModelRunner 引入了**持久化批次（Persistent Batch）**模式：
具体来说，ModelRunner 在 GPU 上维护一组持久化张量：
input_ids_tensor：大小为[max_num_seqs × max_model_len]的预分配张量positions_tensor：同上大小block_table_tensor：大小为[max_num_seqs × max_num_blocks_per_seq]
这些张量在 Worker 初始化时就分配好了，之后不再重新分配。每一步，ModelRunner 只更新其中变化的部分：
- 新请求加入 → 将其 Token IDs 写入张量的对应行
- 解码请求生成新 Token → 更新一个位置
- 请求完成 → 标记该行为无效
差量更新使用 NumPy 操作 而非 Python 原生列表操作。原因是 NumPy 的批量索引操作（如 arr[indices] = values）在 C 层面执行，比 Python 循环快一到两个数量级。这看似是微优化，但在每步处理数百个请求时，Python 层面的 CPU 开销是实实在在的瓶颈。
8.3 分段 CUDA 图
标准 CUDA 图的限制
CUDA 图（CUDA Graph）是 NVIDIA 提供的一种优化技术：将一系列 CUDA 内核调用"录制"为一个图，之后每次执行只需要"回放"该图，避免了重复的内核启动开销。
对于 LLM 推理，解码阶段的每一步计算模式几乎完全相同------相同的模型层、相同的算子序列。理论上完美适合 CUDA 图。
但标准 CUDA 图有一个致命限制：图中所有张量的形状（shape）必须固定。而 LLM 推理中，每步的批大小（并发请求数）和序列长度可能变化。当新请求加入或旧请求完成时，批大小就变了。
V0 的解决方案是为每个可能的批大小预先录制一个 CUDA 图。如果批大小最大为 256，就需要录制 256 个图。这不仅消耗大量 GPU 显存（每个图都有自己的张量副本），而且预热时间很长。
V1 的分段 CUDA 图
V1 引入了分段 CUDA 图（Piecewise CUDA Graph）------将一个大图拆分为多个小段：
一个大图
形状固定"] end subgraph "分段 CUDA 图" P1["段 1
前 N 层"] P2["段 2
注意力"] P3["段 3
后 M 层"] end P1 --> P2 --> P3 style SG fill:#ef4444,color:#fff,stroke:none style P1 fill:#10b981,color:#fff,stroke:none style P2 fill:#3b82f6,color:#fff,stroke:none style P3 fill:#10b981,color:#fff,stroke:none
分段的关键洞察是：模型中大部分算子对批大小不敏感（如 LayerNorm、MLP），只有注意力层需要知道确切的序列长度和批大小。
通过在注意力层处切断 CUDA 图，V1 可以：
- 对不依赖形状的部分使用固定的 CUDA 图段
- 对注意力层动态调整参数
- 图段之间通过共享的 GPU 张量传递数据
这大幅减少了需要录制的图数量和显存占用，同时保留了 CUDA 图的大部分性能收益。
CUDA 图预热的真实流程
V1 的 CUDA 图预热在 _dummy_run（gpu_model_runner.py:1662）中完成：
python
# gpu_model_runner.py:1662-1687 (简化)
def _dummy_run(self, ...):
if not self.use_cuda_graph:
return # 未启用则跳过
# 从大到小遍历每个需要录制的批大小
for num_tokens in reversed(self.cudagraph_batch_sizes):
# 对每个批大小运行一次预热
# torch.compile 会在预热时触发 CUDA 图录制
self.model(...)预热的要点：
- 从大到小遍历：GPU 显存中先分配大图的空间，然后小图可以复用部分资源
- 显存开销可量化：预热前后测量 GPU 空闲显存的差值即为 CUDA 图总开销
- 典型配置下，CUDA 图占用 1-3 GB 额外显存
execute_model：真实的前向传播入口
execute_model（gpu_model_runner.py:1003）是每步推理的入口函数。让我们看它的核心逻辑：
python
# gpu_model_runner.py:1003-1038 (简化)
def execute_model(self, scheduler_output, ...):
self._update_states(scheduler_output) # 差量更新持久化批次
# 准备注意力元数据
attn_metadata, logits_indices, spec_decode_metadata = (
self._prepare_inputs(scheduler_output)
)
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
# 关键决策：使用 CUDA 图 or Eager 模式
if (self.use_cuda_graph
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
# CUDA 图模式：填充到最近的预录制批大小
num_input_tokens = self.vllm_config.pad_for_cudagraph(
num_scheduled_tokens)
else:
# Eager 模式：直接使用实际大小
num_input_tokens = num_scheduled_tokens注意 pad_for_cudagraph 的填充逻辑------因为 CUDA 图要求固定形状，所以实际 token 数需要被填充到预录制的某个批大小（如 1, 2, 4, 8, 16, 32, ...）。填充的 token 不参与实际计算，但会被传入 GPU 内核。这是 CUDA 图的固有 trade-off：用少量无效计算换取内核启动开销的消除。
8.5 torch.compile 集成
除了 CUDA 图，V1 还支持使用 PyTorch 的 torch.compile 进行自动优化。
torch.compile 会分析模型的计算图，自动进行算子融合、内存布局优化等变换。与手写 CUDA 内核相比，它的优势是可移植性------不需要为每种硬件平台编写专门的内核。
vLLM 在 V1 中将 torch.compile 作为一种可选的优化模式，通过编译配置控制。在某些模型和硬件组合上，torch.compile 的性能已经接近甚至超过手写内核。
8.5 完整的前向数据流
让我们把本章的所有知识串联，看一步推理的完整数据流：
Token → 向量"] EMB --> L1["Transformer Layer 1"] L1 --> L2["Transformer Layer 2"] L2 --> LN["... Layer N"] LN --> NORM["RMSNorm"] NORM --> LM["LM Head
向量 → 词表 Logits"] LM --> SAMPLE["采样
Logits → Token ID"] subgraph "每个 Transformer Layer" direction TB ATTN["注意力
（PagedAttention 内核）"] MLP_["MLP
（Gate + Up + Down）"] ATTN --> MLP_ end style INPUT fill:#f59e0b,color:#fff,stroke:none style SAMPLE fill:#10b981,color:#fff,stroke:none
- Embedding：Token ID → 隐藏状态向量（维度 = hidden_size）
- N 个 Transformer 层：每层包含注意力（读写 KV Cache）和 MLP
- RMSNorm：最后的归一化
- LM Head：将隐藏状态映射到词表大小的 Logits
- 采样：根据采样参数从 Logits 中选择下一个 Token
整个过程中，最耗时的是注意力计算（尤其是长序列时 KV Cache 的读取）和 MLP 计算（大量矩阵乘法）。PagedAttention 内核优化了前者，量化（第 13 章）优化了后者。
8.7 execute_model 完整流程
8.8 本章小结
ModelRunner 是 GPU 计算的直接执行者：
- 持久化批次------预分配张量 + NumPy 差量更新，消除每步的 Python 输入准备开销
- 分段 CUDA 图------在注意力层处切段，兼顾动态形状和 CUDA 图加速
- torch.compile------自动优化，可移植性好，性能接近手写内核
- 前向数据流------Token ID → Embedding → N × (Attention + MLP) → LM Head → Logits → 采样
下一章，我们将聚焦于前向传播的最后一步------采样，看看温度、top-p、top-k 是如何从数学定义变成高效 GPU 实现的。
源码导航
- GPU ModelRunner：
vllm/v1/worker/gpu_model_runner.py- 编译配置：
vllm/compilation/- Transformer 层实现：
vllm/model_executor/layers/- 注意力后端：
vllm/v1/attention/backends/