执行摘要

• 一句话：修复完成请求的 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 槽的块。

# 在 _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，修复了不同方面的存储可达性问题，可能相关。