Prhub

#50062 [Model Runner V2][Spec Decode] Add KV cache support for multi-layer MTP

原始 PR 作者 TheEpicDolphin 合并时间 2026-08-14 07:22 文件变更 10 提交数 3 评论 14 代码增减 +196 / -32

执行摘要

为多模块 MTP 新增调度与 KV cache 支持

PR body 明确指出:多模块 drafter 在 prefill 期间读取超前于已计算 token 的位置——MTP 模块 m 从 token p + m + 1 计算位置 p 的 KV,因此在每个 chunked-prefill 边界 drafter 都要消费接下来 num_speculative_tokens 个已知 prompt token 来写出精确 KV。加上 decode 期 rejection re-prefill(重写尾部 N-1 个位置),产生三个调度器与 KV cache 层必须处理的隐患:chunk 边界落在 lookahead token 不存在的位置、缓存或释放仍可能变化的 KV、prefix-cache hit 命中编码了其他请求 continuation 的 KV。

值得精读。核心设计决策包括:用一个 num_prefill_lookahead 字段驱动调度边界、编码器偏移、KV 注册延迟与 SWA 保留四个子系统,显著降低多模块 MTP 的横切复杂度;延迟哈希注册体现“先验证后共享”的安全姿态;用 raise ValueError 硬校验而非静默降级来保证 prefix-cache hit 的 soundness。建议阅读 scheduler.py_reserve_prefill_lookaheadkv_cache_coordinator.pycache_blocks,并结合 SlidingWindowSpec.extra_retained_tokens 理解 SWA 与重算窗口的交互。

讨论亮点

唯一的一条代码级 review 评论来自 depthfirst-app[bot],针对 kv_cache_coordinator.py 中防御跨请求 KV cache 数据泄漏的 assert 守卫,指出 Python -O 会剥离所有 assert,建议改用 raise ValueError 以保证安全检查始终生效。最终代码采纳该建议,并在提交 ba0d589 中同时将校验对象从 per-group block_size 改为 scheduler_block_size。WoosukKwon 的 APPROVED 评论表示高层正确性认可,认为只影响 multi-layer MTP 行为,合并安全,并要求作者补充 Inkling 端到端验收率与准确率数据;作者随后在 PR body 中补充了 GSM8K 评估与 benchmark 表。流程上 PR 经历两次 merge conflict 与多次 /ci run,rebase 后合并。

实现拆解

  1. 配置层解锁全部 MTP 深度vllm/config/speculative.py):hf_config_override 中 Inkling 分支将 n_predict 从钳制为 1 改为暴露全部 checkpoint 深度(checkpoint_depths),并在 SpeculativeConfig.__post_init__ 删除 “Inkling MTP 只支持恰好 1 个推测 token” 的校验;模块 i 草拟第 i 个推测 token,使 use_multi_module_mtp() 判定成立。配套测试 test_inkling_override_exposes_all_mtp_depthsn_predict 断言从 1 更新为 8。

  2. 调度器引入 num_prefill_lookahead 统一抽象vllm/v1/core/sched/scheduler.py):__init__ 新增字段,多模块 MTP 取 num_spec_tokens,Eagle 家族取 1,否则为 0。三个消费点由它投影:_try_schedule_encoder_inputsshift_computed_tokens 由硬编码 1 改为该字段(多模态 span 提前一个 chunk 编码);新增 _reserve_prefill_lookahead 在两个调度循环(新请求与恢复请求循环)中约束 chunk 边界,避免边界落在距 prefill 末尾 0 < remaining < N 的位置;_free_encoder_inputs 的编码器延迟释放偏移同步改为该字段。该函数在 lookahead 为 0/1 时退化为无操作,随后字段透传给 KVCacheManager

  3. KV cache 协调器延迟哈希注册与 soundness 校验kv_cache_coordinator.pykv_cache_manager.py):KVCacheCoordinator 新增 num_prefill_lookahead 参数并派生 num_reprefillable_tokens = max(0, lookahead - 1)cache_blocks 只注册 num_computed - num_reprefillable_tokens 的 token,尾部 N-1 个未定稿 token 在 rejection re-prefill 重写前不得进入前缀哈希表,防止未验证 KV 被共享;HybridKVCacheCoordinator.cache_blocks 的 EAGLE lookahead 块资格逻辑同步镜像。构造期新增硬校验:开启前缀缓存且存在 EAGLE 组时要求 scheduler_block_size >= num_prefill_lookahead,否则 raise ValueError 拒绝启动——因为现有 EAGLE last-block drop 只丢弃一个块,只有块尺寸足够大才能覆盖全部 N 个可能被污染的槽位。

  4. 滑动窗口尾沿保留扩展kv_cache_interface.pykv_cache_utils.pysingle_type_kv_cache_manager.py):SlidingWindowSpec 新增 extra_retained_tokens 字段,get_kv_cache_configs 在多模块 MTP 下向所有 SWA spec 注入 N - 1max_admission_blocks_per_requestSlidingWindowManager.get_num_skipped_tokens 同步调整,使 SWA 释放边界滞后 N-1 个 token,池子尺寸与 admission 上限按同一口径计算;replace_as 增加 drop 参数,使 SWA→FullAttention 提升路径可剔除该字段。

  5. Speculator 路由与测试配套init_speculatoruse_multi_module_mtp() 时返回 MultiModuleMTPSpeculatortests/v1/core/test_scheduler.py 为既有 eagle encoder-shift 回归测试与 test_free_encoder_inputs_defers_for_eagle_lookahead 补充 num_prefill_lookahead 参数化。

文件 模块 状态 重要度
vllm/v1/core/sched/scheduler.py 调度器 modified 7.72
vllm/v1/core/kv_cache_coordinator.py 缓存协调 modified 7.15
vllm/v1/kv_cache_interface.py 缓存接口 modified 7.04
vllm/v1/core/kv_cache_utils.py 缓存配置 modified 6.31
vllm/v1/core/single_type_kv_cache_manager.py 缓存管理 modified 6.18
vllm/config/speculative.py 推测配置 modified 6.11
vllm/v1/worker/gpu/spec_decode/__init__.py 推测路由 modified 5.35
vllm/v1/core/kv_cache_manager.py 缓存管理 modified 5.1
tests/config/test_speculative_draft_hf_overrides.py 配置测试 modified 4.57
tests/v1/core/test_scheduler.py 调度测试 modified 4.2

关键符号

_reserve_prefill_lookahead init_speculator KVCacheCoordinator.cache_blocks get_kv_cache_configs SlidingWindowManager.get_num_skipped_tokens SlidingWindowSpec.max_admission_blocks_per_request replace_as

关键源码片段

vllm/v1/core/sched/scheduler.py core-logic

核心调度逻辑入口:新增 num_prefill_lookahead 统一字段、_reserve_prefill_lookahead 边界约束,并将编码器偏移与延迟释放从硬编码改为字段驱动,两个调度循环均受影响。

# vllm/v1/core/sched/scheduler.py(关键片段)# 统一的 prefill lookahead 抽象:drafter 在 prefill 阶段要超前读取多少个
# 已知 token。它同时驱动编码器调度偏移、编码器延迟释放、KV cache 重算窗口
# 和 chunk 边界约束,避免多模块 MTP 的横切修改散落各处。
self.num_prefill_lookahead = 0
if speculative_config is not None:
    # ... 动态推测调度表构建 ...
    self.use_eagle = speculative_config.use_eagle()
    if self.use_eagle:
        # Eagle 家族(含单模块 MTP)只超前 1 个位置;多模块 MTP 中模块 i
        # 需要 token p + i + 1 的 embedding,因此要超前 num_spec_tokens 个位置。
        self.num_prefill_lookahead = (
            self.num_spec_tokens
            if speculative_config.use_multi_module_mtp()
            else 1
        )
​
​
def _reserve_prefill_lookahead(
    self,
    request: Request,
    num_computed_tokens: int,
    num_new_tokens: int,
) -> int:
    """绝不把 prefill chunk 结束在距 prefill 末尾不足 lookahead 的位置。"""
    # 在 chunked-prefill 边界上,多模块 MTP drafter 会把接下来 N 个已知
    # prefill token 当作草稿输入。若边界离末尾太近,就采不到真实 lookahead
    # token,drafter 会退回采样草稿,导致尾部模块在这些位置的 KV 被永久污染
    # (这些位置落在所有未来查询窗口之外,没有任何机制能重写它们)。因此要么
    # 本次直接完成 prefill,要么给下一 chunk 至少留下 N 个 token。
    # 对 Eagle 家族(lookahead == 1)而言该约束退化为无操作。
    remaining = request.num_tokens - num_computed_tokens - num_new_tokens
    if 0 < remaining < self.num_prefill_lookahead:
        num_new_tokens -= self.num_prefill_lookahead - remaining
    return max(num_new_tokens, 0)
vllm/v1/core/kv_cache_coordinator.py core-logic

KV cache 协调层核心:透传 num_prefill_lookahead 并派生 num_reprefillable_tokens,cache_blocks 延迟哈希注册,构造期校验 scheduler_block_size 覆盖 lookahead 窗口。

# vllm/v1/core/kv_cache_coordinator.py(关键片段)class KVCacheCoordinator(ABC):
    def __init__(self, ..., num_prefill_lookahead: int = 0):
        # 可被 rejection re-prefill 重写的尾部 token 数:decode 阶段最多重算
        # 最后 N-1 个位置的草稿 KV,这些位置在定稿前不得进入前缀缓存。
        self.num_reprefillable_tokens = max(0, num_prefill_lookahead - 1)
​
        # 前缀缓存正确性守卫:EAGLE 家族命中前缀时只丢弃最后一个块,并用本请求
        # 自己的 lookahead token 重算。多模块 MTP 下,任何一个被缓存前缀的最后
        # N 个槽位都可能存放着写者 continuation(chunk 边界处的 lookahead token,
        # 或采样草稿)派生出的草稿 KV——这些内容不属于块哈希覆盖范围。若调度块
        # 尺寸小于 N,丢弃一个块不足以覆盖全部污染槽位,命中者会直接读到他人
        # 未经验证的 KV,此时宁可拒绝启动。
        if (
            enable_caching
            and self.eagle_group_ids
            and scheduler_block_size < num_prefill_lookahead
        ):
            raise ValueError(
                f'Multi-module MTP with prefix caching requires ' # noqa: E501
                f'scheduler_block_size ({scheduler_block_size}) >= '
                f'num_speculative_tokens ({num_prefill_lookahead}).'
            )
​
    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
        # 只把 KV 已定稿的 token 注册进哈希前缀表。末尾 N-1 个 token 在 decode 期
        # rejection re-prefill 中还会被改写:提前注册既会把未验证 KV 暴露给共享
        # 该前缀的其他请求,也会在块被共享后再次篡改其内容。
        for manager in self.single_type_managers:
            num_tokens_to_cache = max(
                0, num_computed_tokens - self.num_reprefillable_tokens
            )
            manager.cache_blocks(
                request,
                num_tokens_to_cache,
                retention_interval=self.retention_interval,
            )
vllm/v1/kv_cache_interface.py data-contract

KVCacheSpec 数据契约变更:SlidingWindowSpec 新增 extra_retained_tokens 字段并计入 admission 上限,replace_as 增加 drop 参数支持 spec 提升时剔字段。

# vllm/v1/kv_cache_interface.py(关键片段)@dataclass(frozen=True, kw_only=True)
class SlidingWindowSpec(AttentionSpec):
    sliding_window: int
    head_size_v: int = None # type: ignore[assignment]
    # 多模块 spec decode 需要重算序列末尾最多 N-1 个位置,每个重算位置都
    # 需要完整的注意力窗口;若窗口尾沿不额外滞后,这些窗口会伸进已释放的
    # null block,基于垃圾数据算出被“修正”的 KV。该字段让 SWA 释放边界
    # 滞后 N-1 个 token(保留但不参与注意力),并让 admission 上限与启动
    # 池大小按同一口径核算——单一真相源保证三者一致。
    extra_retained_tokens: int = 0
​
    def max_admission_blocks_per_request(
        self, max_in_flight_tokens: int, max_model_len: int
    ) -> int:
        # chunked prefill 期间持有最近 sliding_window - 1 个已计算 token 的 KV,
        # 外加 in-flight token;多模块 spec decode 再额外保留 extra_retained_tokens
        # 个尾沿 token。池子与 admission gate 必须按同一公式估算,否则运行期会
        # 出现请求被误拒或池子超额。
        num_tokens = min(
            self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens,
            max_model_len,
        )
        # +1:滑窗可能不从块首开始(块大小 4、窗口 6 个 token 需要 2 块:
        # [XXCD][EF] 才能装下 [CDEF])。
        return cdiv(num_tokens, self.block_size) + 1

评论区精华

prefix-cache 守卫用 assert 可能在 -O 下失效 安全

depthfirst-app[bot] 指出 `kv_cache_coordinator.py` 中防御跨请求 KV cache 数据泄漏的 `assert` 守卫在 Python `-O` 下会被静默剥离,建议改用 `raise ValueError` 以保证检查始终生效,并引用同文件其他 `raise ValueError` 先例。

结论:最终代码采纳该建议改为 `raise ValueError`,并在提交 `ba0d589` 中同时把校验对象从 per-group block_size 改为 scheduler_block_size(与命中边界对齐),后续提交 `02f7c1e` 又修复了 SWA spec 提升时的字段剔除。 · 已解决

maintainer 对正确性的认可与数据要求 question

WoosukKwon 在 APPROVED 评论中表示:“I'm probably not the best person to verify the correctness of the PR, but it seems correct to me at the high level, and only affect the behavior on multi-layer MTP. Therefore, I think it is safe to merge.” 同时要求作者补充 Inkling 端到端验收率与准确率数据。

结论:作者在 PR body 中补充了 GSM8K(1319 题,5-shot)评估表与 server benchmark 表,验收率约 43%,前缀缓存开启时吞吐翻倍以上。 · 已解决

多次 merge conflict 与 CI 重跑 other

mergify[bot] 两次提示 PR 存在 merge conflict 需要 rebase,作者多次触发 /ci run(Buildkite #82934、#83251、#83401、#83778、#83794)。

结论:作者完成 rebase 并触发最终 CI #83794 后由 WoosukKwon 合并。 · 已解决

风险与影响

  1. 核心调度路径变更scheduler.pyschedule() 两个循环都插入了 _reserve_prefill_lookahead 调用,属于调度主路径;虽然 lookahead 为 0/1 时退化为无操作,但该结论依赖 num_prefill_lookahead 初始化与 use_eagle()/use_multi_module_mtp() 判定的一致性,若未来 Eagle 家族扩展新成员而忘记同步取值,可能引入隐性行为变化。
  2. 前缀缓存正确性cache_blocks 延迟哈希注册与 Hybrid 协调器的 EAGLE lookahead 块资格逻辑是两套独立实现,hybrid 路径没有新增专门测试;若镜像逻辑漂移,可能提前注册未验证 KV 或漏缓存已定稿 KV,前者造成跨请求数据泄漏。
  3. SWA 内存池尺寸变化extra_retained_tokens 同时影响 max_memory_usage_bytes、admission cap 与 get_num_skipped_tokens,尽管以 spec 为单一真相源,但三者若不同步会导致启动池子超额或运行期请求被误拒。
  4. 依赖后续 PR:runner 侧消费早期编码 embedding 的代码尚未落地,当前用文本 embedding 兜底,多模态场景下 lookahead token 的语义精度受限。
  5. 正确性验证依赖作者:WoosukKwon 自认不是验证该 PR 正确性的最佳人选,评估数据由作者提供。

对用户:启用多模块 MTP(如 Inkling 8 深度)的用户获得 chunked prefill + 前缀缓存下的正确 KV 行为,前缀缓存开启时吞吐提升约 71%(2245→3830 tok/s),TTFT 降低约 47%;其他配置(无 spec decode、单模块 MTP/Eagle)行为不变。对系统:滑动窗口 + 多模块 MTP 组合下 KV 池启动尺寸增大(额外保留 N-1 个 token 的块),前缀缓存命中率可能轻微下降(尾部 N-1 个 token 不参与哈希注册)。对团队:num_prefill_lookahead 成为连接调度器、编码器与 KV cache 层的统一契约,为后续多模块 MTP 相关 PR 奠定基础,也是 Model Runner V2 生态的一部分。

核心调度路径变更 前缀缓存正确性敏感 多模块 MTP 专属路径,回归面有限 SWA 内存池尺寸变化 依赖后续 PR(runner 侧 embedding 消费)

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论