# PR #56709 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Bugfix][KV Offload] Restore MTP-retained sliding-window history
- 合并时间：2026-09-16 13:14
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/56709

---

# 执行摘要

- 一句话：SWA offload 覆盖补齐 MTP 保留 token
- 推荐动作：值得精读，尤其是做 KV offload、hybrid KV cache 管理或 MTP 推测解码的工程师。建议重点关注三处设计决策：(1) `get_sliding_window_size_in_chunks` 中 `max(base, cdiv(window - 1 + extra_retained_tokens, chunk))` 的写法——为什么要用 `-1` 与外层 `max`，它决定了「保留 1 个 token 不额外多要 chunk」；对应测试 `test_swa_alignment_skip` 断言 2 与 3 的分界；(2) `_reachable_store_block_mask` 用 `dataclasses.replace` 造临时 coverage spec，把「放宽 store 掩码」与「不改变模型注意力窗口」两件事解耦，这是本 PR 最值得借鉴的模式；(3) 测试从硬编码 `boundary == 3904` 改为按 `first_retained_block` 推导 `needs_left`，把失败场景从特例变成契约。阅读时应同时打开 Issue #56701 的失败表格与 `tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py` 的参数化矩阵。若你负责性能，请补测 SWA store 体积上升在真实长上下文负载下对 CPU 缓存命中率的影响。

# 功能与动机

Issue #56701（RFC：Align sliding-window restore coverage with MTP-retained history）给出了可复现的失败路径：`get_kv_cache_configs()` 会为多模块 MTP 把 `SlidingWindowSpec.extra_retained_tokens` 设为 `num_speculative_tokens - 1`，`SlidingWindowManager.get_num_skipped_tokens()` 会因此少跳过一部分历史（即分配侧保留更老的块），但 offload 的 `get_sliding_window_size_in_chunks()` / `load_window_size_in_chunks()` 与 producer 的 `_reachable_store_block_mask()` 都仍按原始注意力窗口计算。Issue 中给出的表格显示，64 token 命中 + SWA 窗口 32 + 额外保留 3 时，lookup 只检查 chunk 2-3，分配却要恢复 chunk 1；chunk 1 缺失时报 `Chunk ... not found in cache`，仍在写入时报 `Chunk ... is not ready for reading`，且真实 `_reachable_store_block_mask()` 在该边界返回 `[False, False, True, True]`，producer 自己就漏掉了 chunk 1。PR body 明确表示实现的是 RFC 中讨论的「保守覆盖」方案，并说明替代方案（命中时只分配注意力所需区间）需要额外证明 MTP 重放与缓存不会读到被省略的历史。

# 实现拆解

1. **统一 lookup 覆盖下界（`get_sliding_window_size_in_chunks`）**：对 `SlidingWindowSpec` 由 `cdiv(sliding_window, tokens_per_chunk)` 改为 `max(cdiv(sliding_window, tokens_per_chunk), cdiv(sliding_window - 1 + extra_retained_tokens, tokens_per_chunk))`。分子刻意用 `sliding_window - 1 + extra_retained_tokens` 而不是 `sliding_window + extra_retained_tokens`，并保留原值作为下界，从而「保留 1 个额外 token 不会平白多要一个 chunk」。该值会写入 `GroupOffloadConfig.sliding_window_size_in_chunks`，同时影响 lookup 覆盖与后续 store 掩码，因此是对 SWA 组覆盖范围的统一加宽。

2. **把额外保留量并入按请求计算的窗口（`GroupOffloadConfig.load_window_size_in_chunks`）**：在原 `cdiv(sliding_window - 1 + right_padding, tokens_per_chunk)` 的分子中插入 `extra_retained_tokens`，其中 `right_padding = -num_tokens % tokens_per_chunk` 用于对齐 lookup 右边界向上取整与分配侧以真实命中边界结尾的口径差异。这样对齐命中（aligned）与非对齐命中（unaligned）都取到与 `SlidingWindowManager` 分配一致的左边界。

3. **放宽 producer 的 store 可达性掩码（`_reachable_store_block_mask`）**：新增 `from dataclasses import dataclass, field, replace`，当 `SlidingWindowSpec.extra_retained_tokens` 非 0 时，用 `replace` 造一个临时 spec：`sliding_window + extra_retained_tokens` 且 `extra_retained_tokens=0`，只把这个临时 spec 传给 `manager_cls.reachable_block_mask()`。`group_config.kv_cache_spec` 本身不被修改，所以模型真实注意力窗口、cache key 计算与 transfer 归属都不变，只让 producer 多存「cold consumer 恢复时真正需要」的块。

4. **测试配套（`tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`）**：`test_swa_offload_window_covers_unaligned_hit` 新增 `(window, extra_retained)` 参数化 `[(128, 0), (128, 3), (256, 1), (256, 3)]` 与 `boundary` 的 `3968` 取值，并用 `first_retained_block = (boundary - window + 1 - extra_retained) // 64`、`needs_left = first_retained_block // 4 < end - 1` 取代原先硬编码的 `boundary == 3904` 判定；`test_swa_alignment_skip` 新增 `extra_retained ∈ {0, 1, 3}` 参数化，断言 `sliding_window_size_in_chunks` 在 `extra_retained > 1` 时为 3、否则为 2，store 期望新增 `(1, 1), (1, 5)`、load 期望新增 `(1, 5)`，并把 `lookup.return_value = HIT` 改为按 `runner.offloaded` 实际命中的 `side_effect`，同时把请求长度 32 改为 33 以保留一个不完整尾块，避免 decode 复用 producer 已释放的块。

5. **无配置、文档或部署改动**：PR body 明确「no public option or API change」，新增 / 修改的开关语义仅存在于内部覆盖计算中。

关键文件：
- `vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py`（模块 卸载调度器；类别 source；类型 core-logic；符号 GroupOffloadConfig.load_window_size_in_chunks, get_sliding_window_size_in_chunks, _reachable_store_block_mask）: 修复的唯一生产文件：lookup 窗口覆盖公式与 store 可达性掩码都在这里，决定了「命中检查范围」与「producer 存储范围」能否覆盖分配侧保留的全部历史。
- `tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`（模块 卸载调度测试；类别 test；类型 test-coverage；符号 test_swa_offload_window_covers_unaligned_hit, test_swa_alignment_skip）: 用参数化把「分配侧保留边界」变成可验证契约：覆盖对齐 / 非对齐混合边界、EAGLE 尾块、left chunk 缺失 / 待写 / 加载中，以及 producer 存储后 cold consumer 的恢复路径。

关键符号：GroupOffloadConfig.load_window_size_in_chunks, get_sliding_window_size_in_chunks, _reachable_store_block_mask, test_swa_offload_window_covers_unaligned_hit, test_swa_alignment_skip

## 关键源码片段

### `tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`

用参数化把「分配侧保留边界」变成可验证契约：覆盖对齐 / 非对齐混合边界、EAGLE 尾块、left chunk 缺失 / 待写 / 加载中，以及 producer 存储后 cold consumer 的恢复路径。

```python
# tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

# 新增 (window, extra_retained) 组合：0 为基线，1 用于证明「多保留 1 个 token
# 不该多要一个 chunk」，3 对应 num_speculative_tokens - 1 = 3 的 MTP 场景。
@pytest.mark.parametrize(
    "window,extra_retained", [(128, 0), (128, 3), (256, 1), (256, 3)]
)
@pytest.mark.parametrize("boundary", [3840, 3904, 3968])
@pytest.mark.parametrize("eagle", [False, True])
@pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"])
def test_swa_offload_window_covers_unaligned_hit(
    boundary, eagle, left_state, window, extra_retained
):
    """A smaller SWA group can move the hit inside another group's CPU chunk."""
    # ... 构造三组 KV cache group（全注意力 / SWA / 小块 SWA）与 CPUOffloadingManager ...

    external, load_async = sched.get_num_new_matched_tokens(request, 0)

    # 分配侧真正保留的第一块：命中边界减去注意力窗口再加回 1，
    # 最后左移 MTP 额外保留量（SWA 块大小固定为 64）。
    first_retained_block = (boundary - window + 1 - extra_retained) // 64
    # 只有当这个块落在 left_key 所在的更左 chunk（每个 chunk 4 个块）时，
    # lookup 才必须额外校验 left_key；否则就是过度覆盖，不该被要求。
    needs_left = first_retained_block // 4 < end - 1

    if needs_left and left_state != "present":
        # left chunk 缺失时命中数归零；仍在写入或加载中时本轮无法命中。
        assert external == (0 if left_state == "missing" else None)
        assert not load_async
        return

    assert external == boundary and load_async

    # ... 走 add_local_computed_blocks / allocate_external_computed_blocks
    # 再调用 sched.update_state_after_alloc() ...

    job = next(iter(sched._jobs.values()))
    if needs_left:
        # 命中被接纳后，load job 必须同时包含更老的 left_key 与窗口尾部块，
        # 这正是修复前会触发 "not found in cache" / "not ready for reading" 的位置。
        assert {left_key, swa[end - 1]} <= job.keys
    manager.complete_load(job.keys, state.req_context)


@pytest.mark.parametrize("extra_retained", [0, 1, 3])
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_swa_alignment_skip(request_runner, async_scheduling: bool, extra_retained):
    """Store the window a cold consumer must restore at each FA boundary.

    With FA blocks of 16 tokens and SWA blocks of 4 tokens, a window of 8
    needs the last two blocks per segment. Retaining three extra tokens for
    MTP requires one more block; retaining just one token does not.
    """
    # ... 构造 full attention block=16 与 SWA block=4 / window=8 的混合配置 ...

    # SWA 组需要的 chunk 数：extra_retained > 1 时从 2 抬到 3，
    # extra_retained == 1 时必须保持 2（保留基线覆盖下界）。
    assert kv_group_configs[1].sliding_window_size_in_chunks == (
        3 if extra_retained > 1 else 2
    )

    # 请求 33 个 token，保留一个不完整的尾块，避免 decode 复用 producer
    # 已释放的 GPU 块，干扰 deferred store 的记录。
    num_tokens = 33

    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        # 每个对齐段保留尾部窗口，必要时再加上 MTP 需要的历史块 (1, 1)/(1, 5)。
        expected_stored=(
            (0, 0), (0, 1), (1, 2), (1, 3), (1, 6), (1, 7),
        )
        + (((1, 1), (1, 5)) if extra_retained > 1 else ()),
    )

    # cold consumer 只命中 producer 真正写过的 key，避免用固定 HIT 掩盖不一致。
    runner.manager.lookup.side_effect = lambda key, ctx: (
        LookupResult.HIT if (key, 0) in runner.offloaded else LookupResult.MISS
    )

```

# 评论区精华

Review 讨论量很小，但信息点明确：

- `claude[bot]` 明确跳过自动 review：「This pull request is from a fork — automated review is disabled. A repository maintainer can comment `@claude review` to run a one-time review.」因此本 PR 没有获得自动化深度审查，最终由维护者 `orozery` 直接 APPROVED（无文字评论）。
- 设计选项在 Issue #56701 中已经定调：作者提出「make lookup and store coverage match the existing allocation contract」为首选，并指出替代方案「allocating only the attention-required interval on an external hit」需要另行证明 MTP 重放与缓存不会读取 / 发布被省略的历史。PR 选择了保守方案，且在 `_reachable_store_block_mask` 中通过 `replace` 构造临时 spec，只动 store 掩码、不动模型注意力窗口，回应了「不要扩大真实注意力窗口」的隐含约束。
- CI 线程主要由 `orozery` 的 `/ci run` 与 bot 提示构成：bot 曾报「This PR is 112 commits behind upstream `main`」，维护者两次 merge main 后重跑；作者反馈「Both V1 KV Connectors jobs hit the upstream NIXL test-stub failure already covered by #56955. The KV Offload test lanes and all KV Offload lm-eval lanes passed.」——即失败与本次改动无关，KV Offload 相关 lane 全绿。

- offload 覆盖契约：lookup/store 对齐分配侧的 retained interval (design): 采用保守覆盖：`get_sliding_window_size_in_chunks` 与 `load_window_size_in_chunks` 计入 `extra_retained_tokens`；`_reachable_store_block_mask` 用 `dataclasses.replace` 造临时 spec 只放宽 store 掩码，不改变模型注意力窗口。
- 保留 1 个额外 token 不应多要一个 chunk (correctness): 测试 `test_swa_alignment_skip` 固化该边界：`sliding_window_size_in_chunks == 3 if extra_retained > 1 else 2`，并在 `extra_retained > 1` 时才对 store/load 期望追加 `(1, 1)` / `(1, 5)`。
- CI：分支落后 112 个提交与上游 NIXL 测试桩失败 (testing): 失败与本次改动无关，等 #56955 落地后可重跑；维护者最终 APPROVED 并合入。
- fork PR 缺少自动化深度审查 (other): 该修复未经过自动化逐行审查，正确性主要依赖作者的单元测试与 tiny 模型 smoke 证据。

# 风险与影响

- 风险：
 1. **CPU 缓存占用与传输量上升（最实质的代价）**：放宽 store 掩码后，带 `extra_retained_tokens > 1` 的 SWA 组会把每个对齐段的可达块从 2 个增加到 3 个（测试中显式新增 `(1, 1), (1, 5)`），即每段 SWA store 体积约 +50%。`test_swa_alignment_skip` 原文档曾强调该机制为真实 DeepSeek V4（约 100K token）减少约 78% 的 SWA store，本次修复会把其中一部分收益让回去；PR body 也承认「Restoring the required older blocks can increase CPU cache usage and transfer volume」，但没有给出量化性能数据。

2. **覆盖公式需要与分配侧长期保持一致**：`sliding_window - 1 + extra_retained_tokens` 与 `SlidingWindowManager.get_num_skipped_tokens()` 的取整口径必须继续对齐，任何一侧单独调整（例如把 `-1` 语义改成窗口包含式，或新增其他保留来源）都会重新引入同类漏算。当前没有跨模块的共享常量或断言来固化这一契约。

3. **命中提前返回的语义变化**：`get_num_new_matched_tokens()` 现在在 `needs_left` 且左块 missing / pending / loading 时返回 0 或 `None`，即主动放弃本次外部命中。缺块时命中数归零是安全的，但会让「本来能命中 64 token」的请求退化为本地重算，属于命中率上的行为变更；测试只验证了返回值，没有度量端到端影响。

4. **验证覆盖有限**：PR body 自述 GPU 验证使用 tiny dummy 权重、BF16、eager、关闭 JIT warmup，并明确「Trained-model quality, acceptance rate, and performance have not been evaluated; no throughput improvement is claimed.」真实权重下 MTP 拒绝 draft 后的 re-prefill 路径、长上下文下的 eviction 交互（CPU cache 竞争）都未被端到端覆盖。

5. **回归基线可解释**：作者给出把生产文件回退到 base 后「14 failed, 88 passed」，说明新测试确实能捕获该缺陷；但相关套件在 H20 上需要预置 `facebook/opt-125m` 配置 / 分词器缓存并开启 offline 模式才能跑完（378 + 96），环境依赖仍偏重。
- 影响：
 - **用户侧**：使用 `OffloadingConnector` + 多模块 MTP（`num_speculative_tokens >= 2`，即 `extra_retained_tokens > 0`）+ 混合 SWA 架构（DeepSeek V4 / GLM 类混合 KV cache）的部署，不再因为「命中已接纳但 chunk 未存储 / 未就绪」的断言而出现请求级故障；代价是 SWA 组在 CPU 侧多存少量块，长上下文高并发场景下 CPU 缓存压力与传输量上升。
 - **系统侧**：offload 的 lookup / store 覆盖现在是「注意力窗口 + extra_retained_tokens」的统一口径，与 `SlidingWindowManager` 的分配契约对齐，减少了 hybrid KV cache 管理中「分配范围 > 查找范围」这类结构性不一致；对 `extra_retained_tokens == 0` 的配置行为完全不变（`max` 退化、store 掩码分支不进入），因此非 MTP 用户零影响。
 - **团队侧**：把「offload 覆盖必须覆盖整个 retained interval」这一隐性契约显式化，并给出了可复用的验证范式（按分配侧真实保留边界参数化测试、store 与 load 期望互相对应），对后续 hybrid/offload 相关改动有参考价值。
 - 风险标记：CPU 缓存与传输量上升 , 命中率行为变更 , 覆盖公式需与分配侧长期对齐 , 缺少真实权重端到端验证 , CI 依赖上游 #56955, fork PR 无自动深审

# 关联脉络

- PR #55818 KV Offload 分配缩放相关改动（标题依据 PR body 表述推断）: PR body 明确说明 `#55818` 处理 allocation scaling，与本 PR 的分配 / 覆盖边界相邻，但不修复 lookup/store 不一致。
- PR #53499 KV Offload page-size 归一化（标题依据 PR body 表述推断）: 同属 offload 的页 / 块尺寸口径调整，与本 PR 的 chunk 覆盖取整逻辑共享同一类边界条件。
- PR #56955 上游 NIXL test-stub 失败修复（标题依据 issue 评论推断）: 本 PR 的 V1 KV Connectors CI 作业因该 NIXL 测试桩问题失败，作者的 CI 结论以此为前提。
- PR #55055 [BugFix][Model Runner V2][Spec Decode] Fix decode instance's multi-layer MTP kv caches during P/D: 同属「多模块 MTP + KV cache 正确性」脉络：前者修正 P/D 下未验证 draft 污染 KV cache，本 PR 修正 MTP 额外保留历史在 offload 路径上的覆盖缺口。