Prhub

#51766 [Bugfix][Core] Preserve Mamba running CoW after external hits

原始 PR 作者 Dao007forever 合并时间 2026-08-11 18:15 文件变更 2 提交数 1 评论 3 代码增减 +73 / -0

执行摘要

修复外部命中后 Mamba running CoW 语义被误判的问题

PR body 指出 Kimi-K3 实测几何为 external state@24,960 → prefill to partial tail@25,472(无 Mamba 块增长)→ continue to prompt end@25,525,生产环境在受影响几何上观察到 3/12 的逐字重复(verbatim repeats)。根因是 allocate_external_computed_blocks() 填充了请求的 Mamba 块表,但当首个续写停留在同一 Mamba 块内时,allocate_new_blocks() 提前返回而未将请求加入 _allocated_block_reqs;随后 partial-tail CoW 被误分类为首次 prefill CoW,Core 用替换块更新表,Worker 却按 running 请求追加返回块 ID,导致 Core 与 Worker 对活 Mamba 状态列产生分歧。

值得精读,尤其是对 vLLM V1 调度器块表所有权和 CoW 语义感兴趣的同学。1 行核心改动背后是 Core/Worker 状态一致性的深刻问题,配套测试完整复现了外部加载序列。建议关注 ywang96 提出的 last_state_block_idx 在投机解码下的分支差异,以及是否补充 num_speculative_blocks > 0 的回归测试。

讨论亮点

核心 review 讨论来自 ywang96(#51763 原修复作者):

  • 肯定方向:ywang96 认为本 PR 方向优于 #51763——#51763 通过抑制 mid-block resume 后的 partial-tail stop 来避开该过渡,但牺牲了一个有效缓存条目;本 PR 保留细粒度 partial-tail 缓存,并修复了所有权簿记,对 Core/Worker 活 Mamba 状态列分歧的解释也比他原来的“切割落在块网格外”更完整。

  • _allocated_block_reqs 其他读者的质疑:ywang96 提出该一行改动会影响 last_state_block_idx 在投机解码下的分支选择(single_type_kv_cache_manager.py:1565-1576):标记后 blocks_allocated 为真时走 prev_block_len - 1 - num_speculative_blocks,否则走 prev_block_len - 1。两个公式仅在 num_speculative_blocks == 0 时一致,因此投机解码场景需要确认语义是否仍然正确。该讨论线程未在 PR 中闭合(review_comments 为空),属于未完全解决的疑虑。

实现拆解

实现拆解如下:

  1. 定位提前返回路径:在 vllm/v1/core/single_type_kv_cache_manager.pyallocate_new_blocks() 中,align 模式分支存在 if num_required_blocks <= len(req_blocks) and not has_partial_hit: return [] 的提前返回。该路径原本不登记 _allocated_block_reqs,导致后续 partial-tail CoW 走错分支。

  2. 单行修复:在提前返回前加入 self._allocated_block_reqs.add(request_id),将请求标记为“已分配”(running)。这样后续 partial-tail CoW 会走 running-request 分支:活块保留在请求表中,缓存快照移到 CoW 目标块,Worker 不再追加替换块,Core 与 Worker 的状态列保持一致。

  3. 新增回归测试:在 tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py 中新增 test_external_mamba_hit_same_block_uses_running_cow_on_continue,用较小等价几何模拟外部加载序列:external state@10 → prefill to partial tail@14 → continue to prompt end@15。测试断言:续写时 continuation_blocks.get_block_ids()[1] == [](不返回新 Mamba 块)、活块仍留在原表位、partial-tail 哈希迁移到排队 CoW 拷贝目标块。

  4. 配套验证:PR 作者在无源码改动时运行新测试失败(断言 [10] == [] 不满足),应用修复后该套件 23 个测试全部通过;pre-commit 对两个变更文件通过。额外说明 macOS 上因 PyTorch MPS teardown 的无关断言而本地禁用了全局 accelerator cleanup,该改动未包含在 PR 中。

文件 模块 状态 重要度
vllm/v1/core/single_type_kv_cache_manager.py 调度器 modified 5.48
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py 前缀缓存 modified 6.34

关键符号

allocate_new_blocks

关键源码片段

vllm/v1/core/single_type_kv_cache_manager.py core-logic

核心修复文件。在 `allocate_new_blocks()` 的 no-new-block 提前返回路径上新增 `self._allocated_block_reqs.add(request_id)`,将外部命中但无需新块的 Mamba 请求标记为已分配,使后续 partial-tail CoW 走 running 分支,修复 Core 与 Worker 的活块表分歧。

def allocate_new_blocks(
    self, request_id: str, num_tokens: int, num_tokens_main_model: int
) -> list[KVCacheBlock]:
    assert isinstance(self.kv_cache_spec, MambaSpec)
    if self.mamba_cache_mode != "align":
        # 非 align 模式:为投机解码(MTP/EAGLE)多分配 speculative blocks
        if self.num_speculative_blocks > 0:
            num_tokens += self.block_size * self.num_speculative_blocks
        return super().allocate_new_blocks(
            request_id, num_tokens, num_tokens_main_model
        )
    else:
        # align 模式:忽略 lookahead tokens,避免破坏块对齐
        num_tokens = num_tokens_main_model
        req_blocks: list[KVCacheBlock] = self.req_to_blocks[request_id]
        # num_tokens 可能包含后续被拒绝的 draft tokens,因此这是过估
        num_required_blocks = (
            cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
        )
        partial_hit = self._partial_hit_reqs.get(request_id)
        has_partial_hit = partial_hit is not None
        # 若所需块数不超过已有块数且无 partial hit,无需分配新块。
        # 修复:即使不分配新块,也要把请求标记为 allocated,
        # 这样后续 partial-tail CoW 会走 running 分支,
        # 活块保留在请求表内,缓存快照移到 CoW 目标块,
        # Core 与 Worker 的活列保持一致。
        if num_required_blocks <= len(req_blocks) and not has_partial_hit:
            self._allocated_block_reqs.add(request_id)
            return []
        else:
            prev_block_len = len(req_blocks)
            blocks_allocated = request_id in self._allocated_block_reqs
            # 记录最后状态块:running 时取倒数第 (1+num_speculative_blocks) 块,
            # 首次命中 prefix cache 时取最后一块。
            if blocks_allocated:
                self.last_state_block_idx[request_id] = (
                    prev_block_len - 1 - self.num_speculative_blocks
                )
            elif prev_block_len > 0:
                self.last_state_block_idx[request_id] = prev_block_len - 1
            # ... 后续分配逻辑(null blocks、speculative block 复用等)
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py test-coverage

新增回归测试 `test_external_mamba_hit_same_block_uses_running_cow_on_continue`,完整复现外部加载后首个续写停留在同一 Mamba 块的序列,断言不返回新块、活块位置不变、partial-tail 哈希迁移到 CoW 目标。该测试在无修复时失败,有效锁定回归。

def test_external_mamba_hit_same_block_uses_running_cow_on_continue():
    """外部 mid-block 命中后,即使首个续写不需要新 Mamba 块,
    请求也必须成为 running 请求,以保证后续 partial-tail CoW 走正确的分支。"""
    hash_block_size = 2
    mamba_block_size = 4 * hash_block_size
    kv_cache_config = KVCacheConfig(
        num_blocks=32,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                FullAttentionSpec(
                    block_size=hash_block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            ),
            KVCacheGroupSpec(
                ["mamba"],
                MambaSpec(
                    block_size=mamba_block_size,
                    shapes=(1, 1),
                    dtypes=(torch.float32,),
                    mamba_cache_mode="align",
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=hash_block_size,
    )
    # 模拟外部加载:状态 @10,随后 prefill 到 partial tail@14,
    # 再续写 1 个 token 到 @15。
    request = make_request("0", [0] * 15, hash_block_size, sha256)
    loaded_blocks = manager.allocate_slots(
        request, num_new_tokens=0,
        num_external_computed_tokens=10, delay_cache_blocks=True,
    )
    assert loaded_blocks is not None
​
    request.num_computed_tokens = 10
    first_step_blocks = manager.allocate_slots(request, num_new_tokens=4)
    assert first_step_blocks is not None
​
    source_block_id = manager.get_blocks("0").get_block_ids()[1][1]
    partial_hash = request.block_hashes[14 // hash_block_size - 1]
    partial_block = manager.block_pool.get_cached_block(
        partial_hash, kv_cache_group_ids=[1]
    )
    assert partial_block is not None
    assert partial_block[0].block_id == source_block_id
​
    # 续写阶段:停留在同一 Mamba 块,不返回新块,活块留在原表位。
    request.num_computed_tokens = 14
    continuation_blocks = manager.allocate_slots(request, num_new_tokens=1)
    assert continuation_blocks is not None
    assert continuation_blocks.get_block_ids()[1] == []
    assert manager.get_blocks("0").get_block_ids()[1][1] == source_block_id
​
    # partial-tail 缓存快照应迁移到排队 CoW 拷贝的目标块,
    # 而不是覆盖活块本身。
    copies, _ = manager.take_kv_cache_block_copies()
    cow_copy = next(c for c in copies if c.src_block_id == source_block_id)
    assert cow_copy.dst_block_id != source_block_id
​
    moved = manager.block_pool.get_cached_block(
        partial_hash, kv_cache_group_ids=[1]
    )
    assert moved is not None
    assert moved[0].block_id == cow_copy.dst_block_id

评论区精华

标记 allocated 对 last_state_block_idx 在投机解码下的影响 设计

ywang96 指出该改动会改变 `_allocated_block_reqs` 的读取分支:`blocks_allocated` 为真时 `last_state_block_idx` 取 `prev_block_len - 1 - num_speculative_blocks`,为假时取 `prev_block_len - 1`。两个公式仅在 `num_speculative_blocks == 0` 时一致,需确认投机解码场景下语义正确。

结论:该问题被提出后未被进一步讨论或显式解决,PR 已由 ZJY0516 批准合并,但 review_comments 为空,属于遗留的未闭合疑虑。 · 待处理

与 #51763 方案对比:保留 partial-tail 缓存 vs 抑制 stop 设计

ywang96 认为本 PR 方向优于 #51763:抑制 partial-tail stop 会牺牲有效缓存条目;修复所有权簿记保留缓存并解释了续写不安全的根因。

结论:结论明确:采纳本 PR 方向,保留 partial-tail 缓存并修复 running CoW 所有权。 · 已解决

风险与影响

主要风险集中在 last_state_block_idx 的语义变化:

  • 投机解码场景:在 vllm/v1/core/single_type_kv_cache_manager.py 中,blocks_allocated 由 False 变 True 后,last_state_block_idx 的公式从 prev_block_len - 1 变为 prev_block_len - 1 - num_speculative_blocks。当 num_speculative_blocks > 0 时,两者不再一致,可能影响 MTP/EAGLE 草稿块的状态保存位置。虽然后续 else 分支的块复用逻辑(reuse previous speculative blocks in this step)依赖 blocks_allocated,但新增测试仅覆盖 num_speculative_blocks == 0 的几何,未覆盖投机解码组合。

  • 外部加载与 partial-tail 缓存交互:改动把外部命中后但未新增块的请求标记为 allocated,可能影响后续 free 或 offload 路径对 _allocated_block_reqs 的读取(如 #51358 涉及的 connector handoff 生命周期),PR body 已声明与 #51358 范围不同,但未做交叉验证。

  • 回归覆盖面:源码仅 1 行,但行为分支变化影响面主要在 Mamba align 模式 + 外部前缀加载场景;常规模型不受影响。

影响范围限定在 vLLM V1 调度器核心的 Mamba align 缓存模式:

  • 用户侧:修复了 Kimi-K3 等 Mamba 模型在外部 KV 加载 + partial-tail 缓存场景下的状态列错乱,可消除 3/12 的逐字重复生成;对普通 Transformer 模型无影响。

  • 系统侧:Core 与 Worker 的块表所有权语义被修正,后续 CoW 行为与 running 请求一致;保留了 partial-tail 缓存命中,避免 #51763 那样的缓存收益损失。

  • 团队侧:该 PR 是 Mamba 外部命中正确性链条中的一环,与 #51763、#51358 形成系列修复;commits 仅 1 个,代码量小,review 已通过。

投机解码分支未覆盖 1 行核心逻辑变更 外部加载与 CoW 交互复杂 讨论遗留未闭合问题

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论