执行摘要

• 一句话：修复 V2 warmup 未预留推测解码 lookahead 块
• 推荐动作：建议 MRV2 和 speculative decoding 相关开发者精读，尤其关注 _reserved_block_count 的注释和对照测试；该 PR 体现了「单一数据源 + 与真实 allocator 对拍」的测试思路，值得在类似的资源预留类逻辑中复用。

功能与动机

V2 warmup 手工构建 SchedulerOutput，每个请求的预留按 cdiv(num_computed + num_scheduled, block_size) 计算。但调度器实际通过 KVCacheManager.allocate_slots 预留了 num_computed + num_scheduled + num_lookahead_tokens 个槽位，因为推测解码器为其草稿生成 KV，而这些 KV 位于目标模型调度范围之外。若 warmup 预留不足，后续推理可能面临块不足，且推测解码器直接基于持久化块表构建槽位映射而不约束列索引，可能导致越界访问内存。

实现拆解

1. 集中 lookahead 口径：在 vllm/config/vllm.py 新增 VllmConfig.num_lookahead_tokens 属性，统一 EAGLE、draft model、DFlash、DSpark 等方法的 lookahead 计算（DFlash 为 num_speculative_tokens + 1，其余为 num_speculative_tokens，无推测时为 0），避免调度器与 warmup 各自推导导致漂移。
2. 重构调度器取值：vllm/v1/core/sched/scheduler.py 的 __init__ 改用 vllm_config.num_lookahead_tokens 初始化 self.num_lookahead_tokens，删除原先按方法重复设置的分支，仅保留 use_eagle 标志的设定。
3. 修正 warmup 预留逻辑：vllm/v1/worker/gpu/warmup.py 将原有的 _warmup_block_count 替换为 _reserved_block_count，精确模拟 KVCacheManager.allocate_slots 的行为：CrossAttention 只算编码器长度；Mamba 在 align 模式下基于不含 lookahead 的 token 数并追加 speculative 块，其他模式叠加 lookahead 后再按块大小上取整；普通 attention/mamba 都受 max_model_len 截断。同时新增 _warmup_block_counter 绑定 runner 配置，避免重复传参。
4. 配套测试：新增 tests/v1/worker/test_gpu_warmup_blocks.py，用 _StepRecorder 记录 warmup 产生的每次分配，通过 _assert_covers_lookahead 校验每个请求至少持有 cdiv(num_computed + num_scheduled + lookahead, block_size) 块；另有对照测试直接驱动真实 KVCacheManager.allocate_slots，在混合 prefill/spec/non-spec 轨迹上逐组逐个断言与 _reserved_block_count 完全一致。测试需要 CUDA 环境。

关键文件：

• vllm/v1/worker/gpu/warmup.py（模块 GPU 预热；类别 source；类型 core-logic；符号 _reserved_block_count, _warmup_block_counter, block_count, _warmup_block_count）: 核心修复所在：新增 _reserved_block_count 模拟 KVCacheManager.allocate_slots，并通过 _warmup_block_counter 统一应用于 run_mixed_prefill_decode_warmup 与 warmup_kernels，确保 warmup 预留块数与调度器一致。
• tests/v1/worker/test_gpu_warmup_blocks.py（模块 预热测试；类别 test；类型 test-coverage；符号 _attention_group, _mamba_group, _make_runner, _StepRecorder）: 新增完整测试套件，覆盖 attention、Mamba（三种缓存模式）以及真实 KVCacheManager 对拍测试，验证 warmup 预留与调度器一致。
• vllm/config/vllm.py（模块 配置；类别 source；类型 core-logic；符号 num_lookahead_tokens）: 新增 VllmConfig.num_lookahead_tokens 属性，作为 lookahead 预留的唯一数据源，供调度器和 warmup 共同使用。
• vllm/v1/core/sched/scheduler.py（模块 调度器；类别 source；类型 core-logic；符号 num_lookahead_tokens, use_eagle）: 重构 Scheduler.__init__，使用 VllmConfig.num_lookahead_tokens 替代重复的按方法分支，减少配置口径扩散。

关键符号：VllmConfig.num_lookahead_tokens, _reserved_block_count, _warmup_block_counter, run_mixed_prefill_decode_warmup, warmup_kernels

关键源码片段

vllm/v1/worker/gpu/warmup.py

核心修复所在：新增 _reserved_block_count 模拟 KVCacheManager.allocate_slots，并通过 _warmup_block_counter 统一应用于 run_mixed_prefill_decode_warmup 与 warmup_kernels，确保 warmup 预留块数与调度器一致。

# vllm/v1/worker/gpu/warmup.py
​
def _reserved_block_count(
    num_tokens: int,
    kvcache_spec: KVCacheSpec,
    *,
    num_lookahead_tokens: int,
    max_model_len: int,
    max_encoder_len: int,
) -> int:
    """模拟 `KVCacheManager.allocate_slots` 的预留逻辑。
​
    Warmup 手工构造 `SchedulerOutput`，必须按调度器的口径预留块数，
    否则正式推理会因块不足而失败或越界访问。
    """
    if isinstance(kvcache_spec, CrossAttentionSpec):
        # 交叉注意力只覆盖编码器序列，与 lookahead 无关
        return cdiv(max_encoder_len, kvcache_spec.block_size)
​
    num_speculative_blocks = 0
    if isinstance(kvcache_spec, MambaSpec):
        # MambaManager 在任何缓存模式下都会追加 speculative 运行态块
        num_speculative_blocks = kvcache_spec.num_speculative_blocks
        if kvcache_spec.mamba_cache_mode == "align":
            # align 模式按不含 lookahead 的原始 token 数计算，保持块对齐
            return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
​
    # 非 align 的 Mamba 与普通 Attention 都要在 token 数上叠加 lookahead，
    # 并受 max_model_len 截断，与 `KVCacheManager` 一致
    num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len)
    return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks

tests/v1/worker/test_gpu_warmup_blocks.py

新增完整测试套件，覆盖 attention、Mamba（三种缓存模式）以及真实 KVCacheManager 对拍测试，验证 warmup 预留与调度器一致。

# tests/v1/worker/test_gpu_warmup_blocks.py
​
class _StepRecorder:
    """从 warmup 发出的 `SchedulerOutput` 重建每次 step 的块持有情况。"""
​
    def __init__(self) -> None:
        # 记录 ( 每组已持块数 , num_computed_tokens, num_scheduled_tokens)
        self.steps: list[tuple[list[int], int, int]] = []
        self._held: dict[str, list[int]] = {}
​
    def execute_model(self, scheduler_output) -> None:
        # 处理 prefill 请求：直接记录新分配的块数
        for new_req in scheduler_output.scheduled_new_reqs:
            self._held[new_req.req_id] = [len(ids) for ids in new_req.block_ids]
            self._record(new_req.req_id, new_req.num_computed_tokens, scheduler_output)
​
        # 处理 decode 请求：累加新增块数
        cached = scheduler_output.scheduled_cached_reqs
        for i, req_id in enumerate(cached.req_ids):
            new_block_ids = cached.new_block_ids[i]
            if new_block_ids is not None:
                self._held[req_id] = [
                    held + len(ids)
                    for held, ids in zip(self._held[req_id], new_block_ids)
                ]
            self._record(req_id, cached.num_computed_tokens[i], scheduler_output)
​
    def _record(self, req_id: str, num_computed: int, scheduler_output) -> None:
        self.steps.append(
            (
                list(self._held[req_id]),
                num_computed,
                scheduler_output.num_scheduled_tokens[req_id],
            )
        )
​
​
def _assert_covers_lookahead(
    steps: list[tuple[list[int], int, int]], num_lookahead_tokens: int
) -> None:
    """断言每个 step 持有的 attention 块数至少覆盖 token 范围加 lookahead。"""
    assert steps, "warmup ran no steps"
    for num_blocks, num_computed, num_scheduled in steps:
        num_tokens = min(
            num_computed + num_scheduled + num_lookahead_tokens, MAX_MODEL_LEN
        )
        assert num_blocks[0] >= cdiv(num_tokens, BLOCK_SIZE), (
            f"{num_blocks[0]} blocks for {num_computed}+{num_scheduled} tokens "
            f"and {num_lookahead_tokens} lookahead tokens"
        )

评论区精华

本 PR 的 review 讨论很少：主要审查者 WoosukKwon 直接批准，没有提出修改意见；claude[bot] 提示来自 fork 的 PR 自动审查被禁用。提交历史显示 pre-commit 两次失败后由 njhill 和 Benjamin Chislett 修复。没有记录关于设计权衡的实质性 review 讨论，但代码内的详细注释（如 Mamba align 模式下 lookahead 与 speculative blocks 的处理）体现了与 KVCacheManager 的精确对齐考量。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. vllm/v1/worker/gpu/warmup.py 中的 _reserved_block_count 必须与 KVCacheManager.allocate_slots 保持同步，任何 allocator 逻辑变化都可能使 warmup 预留出现偏差。
  2. VllmConfig.num_lookahead_tokens 是单点数据源，未来新增推测方法必须同步更新该属性，否则调度器与 warmup 会同时受影响。
  3. 调度器重构后 use_eagle 只由 speculative_config.use_eagle() 设置，若 DSpark 等方法的 use_eagle() 返回 False 且依赖旧逻辑，则 use_eagle 标志行为可能变化；当前配置属性对 DSpark 使用 use_eagle() or uses_draft_model()，但调度器中的 use_eagle 标志本身可能影响其他逻辑，需关注。
  4. 新增测试仅限 CUDA 环境（skipif(not current_platform.is_cuda())），在 ROCm、CPU 等平台不会执行，存在覆盖盲区。
  5. 对 vllm/config/vllm.py 的修改属于配置语义变更，使用 VllmConfig 的第三方扩展可能受属性新增影响。
     - 影响：影响所有使用 V2 model runner 且开启推测解码（EAGLE、MTP、DFlash、draft model 等）的用户，修复了因 warmup 块预留不足导致的潜在运行时失败或越界访问；同时统一了调度器与 warmup 的 lookahead 口径，降低后续维护成本。对团队而言，新增的测试模式（用真实 KV cache manager 对拍预测函数）为类似 warmup 及资源预留逻辑提供了可借鉴的验证思路。
     - 风险标记：核心路径变更, 配置口径集中化, CUDA-only 测试覆盖, 调度器行为重构

关联脉络

• PR #50531 [Bugfix][MRV2] Reserve spec-decode lookahead blocks in V2 warmup: 本 PR 的原始版本，由 rchalamala 提交，PR body 声明当前 PR 是该 PR 的复制品，因此两者在功能演进上完全一致。