第8章 前向计算与 CUDA Graph

"The fastest code is the code that never runs." -- Robert Galanakis

本章要点

• 理解 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 的工作流程是：

步骤 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）在初始化时做了大量预分配工作，为后续每步推理的高效执行打下基础：

# 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）——将一个大图拆分为多个小段：

分段的关键洞察是：模型中大部分算子对批大小不敏感（如 LayerNorm、MLP），只有注意力层需要知道确切的序列长度和批大小。

通过在注意力层处切断 CUDA 图，V1 可以：

1. 对不依赖形状的部分使用固定的 CUDA 图段
2. 对注意力层动态调整参数
3. 图段之间通过共享的 GPU 张量传递数据

这大幅减少了需要录制的图数量和显存占用，同时保留了 CUDA 图的大部分性能收益。

CUDA 图预热的真实流程

V1 的 CUDA 图预热在 _dummy_run（gpu_model_runner.py:1662）中完成：

# 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）是每步推理的入口函数。让我们看它的核心逻辑：

# 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 完整的前向数据流

让我们把本章的所有知识串联，看一步推理的完整数据流：

1. Embedding：Token ID → 隐藏状态向量（维度 = hidden_size）
2. N 个 Transformer 层：每层包含注意力（读写 KV Cache）和 MLP
3. RMSNorm：最后的归一化
4. LM Head：将隐藏状态映射到词表大小的 Logits
5. 采样：根据采样参数从 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/