Model Runner V2 设计文档¶

简介¶

自 vLLM V1 首次实现以来，我们发现了一些根本性的设计错误，并积累了大量的技术债务。许多功能是在最初设计未考虑的情况下强行添加的。同时，我们在采样技术（例如 Gumbel-max 采样）、工具（例如 Triton）和 CUDA 特性（例如 UVA）方面获得了宝贵的见解。基于这些知识，我们从第一性原理出发实现了 Model Runner V2 (MRV2)，使其更简洁、更高效、更模块化。

事后看来，V1 的许多设计选择并非最优。尽管 MRV2 尚未实现所有功能、未经过严苛测试，且仍有一些设计决策悬而未决，但我们认为它比 V1 有了实质性的改进。

本文档描述了 MRV2 的设计。

1. 持久化批处理 (Persistent Batch)¶

V1 中一个显著的摩擦点是其持久化批处理的实现。

背景¶

V1 引入了持久化批处理，以最小化输入准备过程中的 CPU 开销。当请求被调度进行处理时，模型运行器必须构建连续的输入张量（例如 block tables 和每个请求的 temperature 值）以馈送到模型中。在 Python 中每一步都从零开始构建这些张量通常非常缓慢，尤其是对于 block tables 这样的大型张量。

持久化批处理优化利用了以下事实：连续步骤中的请求批次大部分是相同的。每一步只有少数请求（如果有）加入或完成。通过维护持久状态张量并应用增量差异，而不是从零开始重建输入，可以显著降低 CPU 开销。

V1 版本方法的问题¶

虽然高效，但 V1 的持久化批处理设计由于将持久状态与输入张量耦合，引入了不必要的复杂性。V1 直接将持久状态张量用作模型和采样器的输入，这强制要求了严格的布局和排序。当请求加入或完成时，通常需要复杂的张量全量重排，而不是简单的行插入/删除。

V1 还必须维护 CachedRequestState，即请求状态的冗余备份副本，因为当请求仍处于活动状态时，持久张量中的行可能会被覆盖。

其结果是复杂的簿记工作，这在异步调度下变得更加困难。

MRV2 的解决方案¶

MRV2 将持久状态张量与每一步的输入张量解耦。鉴于该步骤的请求顺序（通常由注意力后端决定），MRV2 从持久状态中收集输入张量。

预分配一个具有 max_num_reqs 行的固定大小张量（大多数平台上默认为 1024）。
为每个请求分配一个永久行，用于其活跃生命周期（直到完成或抢占）。
将抢占视为完成。恢复时，将请求数据作为新状态重新添加。

这消除了对 CachedRequestState 的需求并简化了簿记。大型状态张量主要存储在 GPU 内存中，因此收集过程在 GPU 上并行运行，开销较低。

2. 异步优先 (Async-First)¶

vLLM 现在严重依赖异步调度。调度器和工作进程在 GPU 执行第 N 步时为第 N+1 步准备输入，通过重叠 CPU 和 GPU 工作来最大化利用率。

V1 最初设计时并未考虑异步调度，其支持需要进行事后补救和采取变通方案。MRV2 则假定核心模型执行循环是一个没有 CPU 同步点的 CUDA 流。CPU 入口点将工作排入流中。

3. 移除异步屏障 (Async Barrier)¶

异步执行的一个关键要求是 CPU 操作保持非阻塞。必须避免显式同步（例如 torch.accelerator.synchronize）和隐式同步（例如未锁定的 .to("cuda")）。

然而，当 CPU 和 GPU 同时访问同一内存时，异步执行可能会引入竞态条件。

示例（不安全）

class ModelRunner:
    def __init__(self, ...):
        # Pinned buffer
        self.states = torch.zeros(
            max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=True
        )

    def execute_step(self, ...):
        self.states[req_idx] = new_req.data
        states = self.states.to("cuda", non_blocking=True)

当 GPU 仍通过异步拷贝读取 self.states 时，CPU 可能会修改它。

V1 通过在临界区周围设置异步屏障来解决这个问题。这虽然避免了竞争，但有缺点：

容易漏掉需要保护的缓冲区（易出错）。
组织不灵活（所有 CPU 工作都必须保留在屏障内）。
由于同步的存在，潜在的重叠减少。

MRV2 的解决方案：消除竞争条件¶

MRV2 将持久化的 CPU 状态与拷贝的张量分离开来。

class ModelRunner:
    def __init__(self, ...):
        # Not pinned
        self.states = torch.zeros(
            max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=False
        )

    def execute_step(self, ...):
        self.states[req_idx] = new_req.data
        tmp_states = self.states.pin_memory()
        states = tmp_states.to("cuda", non_blocking=True)

现在 CPU 写入 self.states，而 GPU 从 tmp_states 读取，消除了竞争，无需显式同步。

4. StagedWriteTensor¶

对于像 block tables 这样的大型张量，MRV2 通过使用 StagedWriteTensor 避免了每一步的全量 CPU 到 GPU 拷贝：

将基准张量保留在 GPU 上。
在 CPU 上暂存差异。
将差异打包到连续的缓冲区中。
将打包的差异拷贝到 GPU。
启动一个内核来应用差异。

使用示例

# Initialize state on GPU
state = StagedWriteTensor(size=(1024, 1000), dtype=torch.int32, device="cuda")

# Write [3, 1, 2] into row 2, starting at index 3
state.stage_write(row=2, start=3, value=[3, 1, 2])

# Write [-1, -2, -5] into row 0, starting at index 1
state.stage_write(row=0, start=1, value=[-1, -2, -5])

# Apply staged changes
state.apply_write()

这支持不规则更新，无需 CPU-GPU 同步，且内核启动次数极少。它对于 block tables 以及 num_computed_tokens 等混合 CPU/GPU 写入的状态特别有用。

5. GPU 原生输入元数据准备与输出处理¶

MRV2 使用 Triton 内核来准备诸如 input_ids、positions、query_start_loc 和 seq_lens 等输入。

优势

更好的异步行为：GPU 可以推导 CPU 尚不知道的值（例如在推测解码中）。
更低的 CPU 开销：输入准备在 GPU 上非常廉价，避免了 Python 瓶颈。

通用虚拟寻址 (UVA)¶

MRV2 在某些路径中使用 UVA，让 GPU 内核直接访问大型常驻 CPU 的张量（例如 prefill_token_ids），而无需将这些张量复制到 GPU 内存中。

6. Triton 原生采样器¶

MRV2 大部分采样逻辑都在 Triton 中重新实现，以实现更好的数值/内存控制和优化。

Gumbel 采样内核¶

MRV2 引入了一个 Triton Gumbel 采样内核，避免了显式的 softmax 实例化，并使用来自种子输入的无状态核内 RNG（随机数生成）。

高效 Top-K Logprobs 计算¶

V1 在 top-k 之前实例化全词汇表 logprobs。MRV2 先从 logits 中识别 top-k token，然后仅对选定的 token 计算 logprobs。这降低了 GPU 峰值内存使用率。

内存高效的 Prompt Logprobs¶

MRV2 支持更细粒度的分块，包括在单个 prompt 内部进行分块，以避免长 prompt 导致的内存激增。

与推测解码更好的兼容性¶

MRV2 不会将每个请求的采样状态扩展以匹配每个 logit 的形状，而是在内核中使用间接寻址（idx_mapping）将每个 logits 向量映射到正确的请求状态。这简化了对复杂采样参数和 logits 处理器的支持。

7. 模块化¶

MRV2 强调模块化。与 V1 中庞大且纠缠不清的 gpu_model_runner.py 相比，MRV2 将功能逻辑拆分到专门的文件中（例如 mrope_utils.py、penalties.py 等）。

它还将模型输入整合到 InputBatch 类中，并减少了直接对模型运行器属性的耦合。

8. 不滥用 `dummy_run`¶

在 V1 中，dummy_run 承担了过多的职责：

初始内存分析和 torch.compile
CUDA Graph 捕获
预热 (Warmups)
EP+DP 的空 DP 前向传播

MRV2 简化了这一点：

execute_model 支持 dummy run，且不影响状态。
dummy_run 将分析、预热和空 DP 前向传播的任务委托给 execute_model。
CUDA Graph 捕获使用单独的专用路径。

这降低了复杂性，消除了因 execute_model 和 dummy_run 行为不一致而导致的错误。

9. 显式 CUDA Graph 管理¶

V1 的 CUDA Graph 处理是隐式的，难以推理。MRV2 使用 CUDAGraphManager，通过标准 PyTorch API 显式捕获并启动完整的 CUDA Graph。

这使得 Graph 生命周期和执行模式的决策更易理解且更易扩展。例如：MRV2 可以将多个草稿模型 (draft-model) 的前向传播捕获到一个 CUDA Graph 中。

开发理念¶

MRV2 的变更应符合更高的代码质量标准。随着与 V1 功能差距的填补，应在 MRV2 的设计背景下从第一性原理重新考虑功能，而不是简单地快速移植 V1 的行为。

一个关键要求是保持模块化和清晰的抽象边界，即使这需要更多的预先设计迭代。