执行摘要

• 一句话：修复 Mooncake 连接器未卸载生产者最终部分尾部状态的问题。
• 推荐动作：该 PR 是一个有意义的 bugfix，修复了 Mooncake 连接器在特定边界条件下的数据丢失问题。实现方案设计周密，通过新增清晰的 finish-time handoff 路径而非修补现有逻辑，并配有全面的单元测试。建议关注：
  1) 新接口 register_finished_partial_tail 的返回值（是否延迟释放）的语义是否清晰；
  2) 与之前 CoW 路径中 register_boundary_state 的协调，避免重复处理。适合相关模块维护者精读，特别是 scheduler.py、single_type_kv_cache_manager.py 和 store/scheduler.py 的变更。

功能与动机

修复一个生产者请求在 prefill 完成后立即结束时，其最终部分 Mamba align 状态被丢弃的问题。原因为：部分尾部卸载仅在后续分配触发 CoW 时触发，而 prefill-only 的生产者不会再次分配，导致其最终的 Mamba 部分状态无法持久化，后续请求无法命中该前缀。

实现拆解

1. 核心调度器逻辑：在 scheduler.py 的 _connector_finished 方法中，新增 finish-time handoff 路径。当请求完成且是 KV 生产者时，先调用 kv_cache_manager.finalize_partial_tail_offloads 收集安全的生产者尾部状态，再通过新接口 connector.register_finished_partial_tail 通知连接器。该调用发生在请求块表被清理之前，以确保块引用有效。
2. KV 缓存管理器扩展：在 kv_cache_manager.py 中新增 finalize_partial_tail_offloads 方法，遍历所有单类型缓存管理器，调用各自的 finalize_partial_tail_offload 方法收集待卸载条目。
3. Mamba 缓存管理器实现：在 single_type_kv_cache_manager.py 的 MambaManager 类中，新增 finalize_partial_tail_offload 方法，检查请求是否在注册的边界恰好完成且无 in-flight token，若安全则返回组 ID、源块和边界 token 数。同时，将 _producer_partial_tail_reqs 字典值类型从 int 改为 tuple[KVCacheBlock, int] 以存储源块引用，并在 CoW 路径中断言源块匹配。
4. 连接器基类与实现：在 base.py 中定义新的 register_finished_partial_tail 虚方法；在 multi_connector.py 中实现广播至所有子连接器；在 Mooncake 的 connector.py 和 store/scheduler.py 中实现具体逻辑：验证边界合法性，为确切的源块分配一个 store job ID 并固定其引用（pool.touch），创建 ReqMeta 并暂存到 _finished_partial_tail_metas 字典，在下一元数据构建时刷新。
5. 测试覆盖：在 test_partial_prefix_cache_hits.py 中新增 test_finished_partial_tail_uses_table_source_once 和 test_connector_finish_registers_partial_tail_before_cleanup 测试，验证 finish-time 路径的正确性（包括安全/不安全场景）。在 test_mooncake_store_scheduler.py 中新增 test_finished_partial_tail_is_pre_pinned_as_store_job 测试，验证 Mooncake 调度器正确固定块并生成元数据。在 test_multi_connector.py 中新增 test_register_finished_partial_tail_notifies_every_connector 测试，验证广播机制。

关键文件：

• vllm/v1/core/sched/scheduler.py（模块 调度器；类别 source；类型 entrypoint）: 引入了请求完成时的新 handoff 路径，是变更的入口点，协调缓存管理器和连接器。
• vllm/v1/core/single_type_kv_cache_manager.py（模块 缓存管理器；类别 source；类型 core-logic；符号 finalize_partial_tail_offload）: 实现了 Mamba 缓存管理器中安全检查和状态收集的核心逻辑，是变更的关键实现。
• vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py（模块 Mooncake调度器；类别 source；类型 core-logic；符号 register_finished_partial_tail）: 实现了 Mooncake 连接器的具体处理逻辑，包括固定块引用、创建元数据和管理生命周期。
• tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py（模块 前缀缓存测试；类别 test；类型 test-coverage；符号 test_finished_partial_tail_uses_table_source_once, test_connector_finish_registers_partial_tail_before_cleanup）: 核心测试文件，新增了两个关键测试用例，验证 finish-time 路径的正确性。

关键符号：vllm/v1/core/sched/scheduler.py#_connector_finished, vllm/v1/core/kv_cache_manager.py#finalize_partial_tail_offloads, vllm/v1/core/single_type_kv_cache_manager.py#MambaManager.finalize_partial_tail_offload, vllm/distributed/kv_transfer/kv_connector/v1/base.py#KVConnectorBase_V1.register_finished_partial_tail, vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py#MultiConnector.register_finished_partial_tail, vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py#MooncakeStoreConnector.register_finished_partial_tail, vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py#MooncakeStoreScheduler.register_finished_partial_tail

关键源码片段

vllm/v1/core/sched/scheduler.py

引入了请求完成时的新 handoff 路径，是变更的入口点，协调缓存管理器和连接器。

# vllm/v1/core/sched/scheduler.py ( 核心逻辑片段 )
def _connector_finished(
    self, request: Request
) -> tuple[bool, dict[str, Any] | None]:
    if self.connector is None:
        return False, None
​
    # 新增 : 收集完成的生产者尾部状态 ( 仅对 KV 生产者 )
    finished_partial_tails: list[tuple[int, int, int]] = []
    kv_transfer_config = self.vllm_config.kv_transfer_config
    if kv_transfer_config is not None and kv_transfer_config.is_kv_producer:
        finished_partial_tails = (
            self.kv_cache_manager.finalize_partial_tail_offloads(request)
        )
​
    # 原有逻辑 : 清理跳过的块并获取块表
    self.kv_cache_manager.remove_skipped_blocks(...)
    block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens(...)
​
    # 新增 : 通知连接器注册尾部，并可能延迟释放
    partial_tail_delay = False
    if finished_partial_tails:
        partial_tail_delay = self.connector.register_finished_partial_tail(
            request, block_ids, finished_partial_tails
        )
​
    # 原有逻辑 : 调用连接器的请求完成回调
    if not isinstance(self.connector, SupportsHMA):
        delay_free, kv_xfer_params = self.connector.request_finished(
            request, block_ids[0]
        )
    else:
        delay_free, kv_xfer_params = self.connector.request_finished_all_groups(
            request, block_ids
        )
    # 合并释放延迟标志
    return delay_free or partial_tail_delay, kv_xfer_params

vllm/v1/core/single_type_kv_cache_manager.py

实现了 Mamba 缓存管理器中安全检查和状态收集的核心逻辑，是变更的关键实现。

# vllm/v1/core/single_type_kv_cache_manager.py (MambaManager 类片段 )
class MambaManager(SingleTypeKVCacheManager):
    def __init__(self, ...):
        ...
        # 生产者尾部请求状态 : 映射到 ( 源块引用 , 边界 token 数 )
        self._producer_partial_tail_reqs: dict[str, tuple[KVCacheBlock, int]] = {}
​
    def finalize_partial_tail_offload(
        self, request_id: str, num_computed_tokens: int, num_in_flight_tokens: int
    ) -> tuple[int, KVCacheBlock, int] | None:
        """在请求完成时安全地消费生产者尾部。"""
        if self.mamba_cache_mode != "align":
            return None
        producer_tail = self._producer_partial_tail_reqs.pop(request_id, None)
        if producer_tail is None:
            return None
        source_block, boundary_tokens = producer_tail
        # 关键安全检查 : 请求必须在注册的边界恰好完成，且无 in-flight token
        if num_in_flight_tokens != 0 or num_computed_tokens != boundary_tokens:
            return None
        return self.kv_cache_group_id, source_block, boundary_tokens
​
    def _cache_partial_tail_block(self, request, num_tokens):
        ... # 缓存部分尾部时记录生产者信息
        self._producer_partial_tail_reqs[request.request_id] = (
            source_block, # 存储源块引用，而非仅 token 数
            num_tokens,
        )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

实现了 Mooncake 连接器的具体处理逻辑，包括固定块引用、创建元数据和管理生命周期。

# vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py ( 片段 )
class MooncakeStoreScheduler:
    def __init__(self, ...):
        ...
        self._finished_partial_tail_metas: dict[str, ReqMeta] = {} # 暂存完成尾部的元数据
​
    def register_finished_partial_tail(
        self, request, block_ids, partial_tail_offloads
    ) -> bool:
        """为 finish-time tail 队列并固定块。返回 False 表示不延迟释放请求块。"""
        if self.kv_role == "kv_consumer" or not partial_tail_offloads:
            return False
        tracker = self._request_trackers.get(request.request_id)
        if tracker is None or not any(block_ids):
            return False
        # 验证边界合法性和组 ID
        boundaries = {boundary for _, _, boundary in partial_tail_offloads}
        if len(boundaries) != 1:
            raise ValueError("Partial-tail offloads for one request must share a boundary")
        boundary_tokens = next(iter(boundaries))
        if boundary_tokens > tracker.prefill_end_tokens:
            return False
​
        # 收集需要固定的源块
        pinned_block_ids: list[int] = []
        for group_id, block_id, _ in partial_tail_offloads:
            if group_id not in self._boundary_state_group_ids or block_id == NULL_BLOCK_ID:
                return False
            pinned_block_ids.append(block_id)
        pinned_block_ids = list(dict.fromkeys(pinned_block_ids))
​
        # 为确切块分配 store job 并固定引用
        pool = self._gpu_block_pool
        assert pool is not None
        assert request.request_id not in self._finished_partial_tail_metas
        store_job_id = self._next_store_job_id
        self._next_store_job_id += 1
        self._pinned_saves[store_job_id] = (pinned_block_ids, self._num_workers)
        pool.touch([pool.blocks[block_id] for block_id in pinned_block_ids]) # 固定引用
​
        # 创建元数据并暂存，在下一个 build_connector_meta 中刷新
        self._finished_partial_tail_metas[request.request_id] = ReqMeta(
            req_id=request.request_id,
            token_len_chunk=0, # 跳过正常保存
            block_ids=tuple(group.copy() for group in block_ids),
            block_hashes=list(request.block_hashes),
            can_save=True, # 走正常入队路径
            num_prompt_tokens=tracker.prefill_end_tokens,
            store_job_id=store_job_id,
            boundary_state_offloads=partial_tail_offloads,
        )
        tracker.has_pending_offload = True
        # store job 拥有确切块引用，请求可立即释放
        return False
​
    def build_connector_meta(self, scheduler_output):
        ...
        # 刷新暂存的完成尾部元数据
        for req_meta in self._finished_partial_tail_metas.values():
            meta.add_request(req_meta)
        self._finished_partial_tail_metas.clear()
        self._reference_save_blocks(meta) # 为所有请求固定块引用
        return meta

评论区精华

PR 讨论主要集中在 CI 运行和合并冲突解决，未涉及核心技术设计的争议。提交历史显示 PR 经历了多次向上游 rebase 以解决冲突，特别是一次合并中明确将 finish-time Mamba tail handoff 移植到基于 exact boundary-state store job 的设计（对应 commit message: Resolve the #51358 overlap by porting finish-time Mamba tail handoffs to exact boundary-state store jobs）。

• 暂无高价值评论线程

风险与影响

• 风险：回归风险：新路径仅在 is_kv_producer 且请求在注册边界恰好完成时触发，条件严格，影响范围有限。但需确保 finalize_partial_tail_offloads 在 remove_skipped_blocks 之前调用，以保证块引用未被释放。
  性能风险：每个完成的生产者请求会固定一个块引用直到所有 worker 完成传输，但 block pool 已具备引用计数，且 worker 完成后会释放，不会造成持续内存泄漏。
  安全/一致性风险：register_finished_partial_tail 返回 False（不延迟释放），依赖连接器内部固定块引用来保持生命周期，逻辑需确保与后续步骤的元数据刷新协同。
  测试覆盖：新增测试覆盖了核心逻辑，但缺少与真实 Mooncake 服务端交互的端到端测试，存在集成风险。

• 影响：用户影响：对于使用 Mooncake 连接器和 Mamba 混合模型的用户，后续请求可以重用之前请求的最终部分前缀，提升前缀缓存命中率，减少重复计算。
  系统影响：扩展了 KV 连接器的接口（新增 register_finished_partial_tail），要求所有连接器实现或继承默认空实现。Mooncake 调度器内部状态管理（新增 _finished_partial_tail_metas 和预固定 store job 逻辑）变得稍复杂。
  团队影响：为类似“请求完成时触发一次性操作”的模式提供了参考实现，但增加了连接器和缓存管理器的交互复杂度。

• 风险标记：接口变更, 分布式存储生命周期, 测试覆盖率

关联脉络

• PR #51358 [Mooncake] Implement boundary-state store jobs for Mamba CoW: PR body 提到 #51358 处理了 CoW 路径中的 Mamba 边界状态，但未处理生产者提前完成的情况。本 PR 补充了该场景。
• PR #50409 [Mooncake] Local replay-boundary scheduling: PR body 提到 #50409 改变了本地重放边界调度，但未实现连接器完成时的 handoff。本 PR 实现了该功能。