# PR #54288 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Bugfix][KV Offload] Stop offloading the final sampled token's KV slot
- 合并时间：2026-09-06 20:22
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/54288

---

# 执行摘要

- 一句话：修复完成请求的 KV 卸载水位线错误，防止跨请求 KV 中毒。
- 推荐动作：这个 PR 值得精读，因为它修复了一个隐蔽且危险的数据一致性 bug，涉及 KV 卸载和投机解码的交互。关键设计决策包括水位线调整公式 `max()` 的防御性编程，以及测试如何基于请求状态而非步骤计数来确保断言健壮性。建议关注注释中的技术解释，它揭示了 KV 槽写入时机和投机解码下的内存状态。

# 功能与动机

当请求完成时，卸载连接器的存储水位线是 `req.num_tokens`，这包括了最终采样 token。然而，该 token 的 KV 槽从未被任何前向传递写入（它是由前一个位置的前向传递产生的），且在投机解码下会持有第一个被拒绝的 draft token 的 KV。如果生成结束恰好在块边界上，该块会被视为“完整”并按其内容哈希存储。后续请求重放相同 token 时会加载这个中毒块，导致模型使用未接受 token 的 KV 进行注意力计算。此 bug 在 `offload_prompt_only=False` 配置下被触发，且在规模运行中常见。Issue #54193 详细描述了该机制和复现步骤。

# 实现拆解

1. **核心逻辑修复**：修改 `vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py` 中的 `_build_store_jobs` 函数。在处理完成请求的分支中，将 `num_tokens_after_batch` 从 `req.num_tokens` 更改为 `max(req.num_prompt_tokens, req.num_tokens - 1)`，以确保卸载水位线不超过 GPU 前缀缓存提交的点，从而避免卸载包含未写入 KV 槽的块。
2. **测试覆盖更新**：调整 `tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py` 中的测试用例。将 `test_last_block_offloaded_at_request_finish` 重命名为 `test_final_sampled_token_does_not_complete_an_offloaded_block`，并修改其逻辑和注释以反映新契约，验证完成步骤不再存储最后一个块。同时调整了其他测试如 `test_on_request_finished_fires_after_final_block_store`，以基于请求状态而非尝试计数来确保正确性。
3. **注释优化**：根据代码审查反馈，精简了源码注释，更准确地解释了水位线调整的原因。

关键文件：
- `vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py`（模块 卸载调度器；类别 source；类型 core-logic）: 包含核心修复，修改了 `_build_store_jobs` 函数中的完成请求水位线计算，防止卸载未写入 KV 槽的块。
- `tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`（模块 调度器测试；类别 test；类型 test-coverage；符号 test_final_sampled_token_does_not_complete_an_offloaded_block, test_on_request_finished_fires_after_final_block_store）: 测试文件随源码修改，更新了测试用例以验证新契约，确保修复的正确性。

关键符号：_build_store_jobs

## 关键源码片段

### `vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py`

包含核心修复，修改了 `_build_store_jobs` 函数中的完成请求水位线计算，防止卸载未写入 KV 槽的块。

```python
# 在 _build_store_jobs 方法中，处理完成请求的分支
if req.status is RequestStatus.FINISHED_ABORTED:
    num_tokens_after_batch = req.num_computed_tokens
elif req.is_finished():
    # Clamp to the GPU prefix cache's commit point. The final sampled
    # token's slot is never committed (under spec decode it holds a
    # rejected draft's KV), so a block ending there must not be stored.
    num_tokens_after_batch = max(req.num_prompt_tokens, req.num_tokens - 1)
else:
    num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
    num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens

```

# 评论区精华

核心讨论围绕水位线调整的精确形式和测试验证方法展开。用户 `tomylin890` 建议使用 `max(req.num_prompt_tokens, req.num_tokens - 1)` 而非裸 `-1`，以覆盖边缘情况（请求无输出 token 完成），作者 `almogtavor` 同意并实施。维护者 `orozery` 提供了更简洁的注释版本，并指出测试中基于请求状态的门控比尝试计数更可靠，以防止误判。讨论确认了修复的正确性和测试的健壮性。

- 水位线公式选择 (correctness): 采用 `max()` 形式，成本为零但提供额外安全性。
- 测试门控逻辑改进 (testing): 测试改为在 `prepare_store` 中检查请求状态，并断言恰好有一次成功存储来自完成步骤。

# 风险与影响

- 风险：风险较低。修复仅调整了水位线计算，影响范围限于 KV 卸载的完成请求分支。主要风险是回归：如果逻辑错误，可能导致卸载行为异常。但 PR 已通过 241 个单元测试（含原有和修改的测试），并提供了直接的 H100 实测验证。性能影响可忽略不计；兼容性方面，水位线调整使卸载更保守，不会引入新问题。
- 影响：对用户的影响：使用 KV 卸载（特别是 `offload_prompt_only=False`）的用户将不再遇到跨请求 KV 中毒导致的生成质量下降，提升了数据一致性和多轮对话的可靠性。对系统的影响：增强了 KV 缓存管理器的正确性，防止垃圾数据进入卸载层。对团队的影响：需要理解卸载水位线逻辑和投机解码交互，但修复简化了代码逻辑（移除了之前分支的冗余）。
- 风险标记：核心路径变更 , 测试覆盖调整

# 关联脉络

- PR #54193 [Bug]: KV offloading connector: finished-request store watermark includes a KV slot never written by any forward pass: 这是本 PR 修复的关联 issue，详细描述了 bug 的机制和影响。
- PR #48596 [KV Offload] Add store path for the final block at request finish: 本 PR 中提到的原始分支添加，它引入了完成请求的存储路径，但未考虑最终采样 token 的 KV 槽未写入问题。
- PR #54362 [Bugfix][KV Offload] Fix SWA store reachability during chunked prefill: 同属 KV 卸载模块的 bugfix PR，修复了不同方面的存储可达性问题，可能相关。