Prhub

#50344 [BugFix] Scope divergent hybrid cache hits to capable connectors

原始 PR 作者 ivanium 合并时间 2026-08-10 06:38 文件变更 8 提交数 1 评论 5 代码增减 +130 / -27

执行摘要

发散混合缓存命中改为按 connector 能力开关,防无效状态

PR body 明确指出:按组(per-group)的混合缓存命中依赖 connector 恢复缺失的 Mamba 状态,把该策略不加区分地套用到所有 connector 上,会在恢复边界缺少有效状态时仍报告 full-attention 命中。ZJY0516 在关联 issue 请求优先处理,称运行 K3 PD(prefill-decode 分离)时在 truncate_computed_blocks 遇到断言错误,预计本 PR 可修复。作者同时声明这不是 #46455 或 #48195 的重复,而是对 #48425 原始发散处理的后续保守化,并覆盖部分(partial)发散命中。

值得精读。该 PR 展示了典型的「能力契约 + 保守默认 + 显式 opt-in」模式:先用 capability 属性把「发散命中是否需要外部状态恢复」这一策略问题与具体 connector 解耦,再由调度器统一决策,风险面小而清晰。建议重点关注 KVConnectorBase_V1.supports_divergent_local_hybrid_hits 的契约语义、MultiConnector 的 all() 聚合,以及测试中 (False, 16, 6, 32) 用例如何锁定回落行为;若团队维护第三方 KV connector,需评估是否声明该能力以恢复发散命中性能。

讨论亮点

核心 review 讨论只有一处代码级建议:njhill 针对 review 早期版本(带「hit_diverged 且 num_local 非块对齐则回退公共查找」逻辑)在 vllm/v1/core/sched/scheduler.py 上提议简化,将「未发散或块对齐时直接返回」的条件提前,使公共查找分支只保留一份。最终合并版本进一步收敛为纯 capability 分流(capable 直接返回 divergent 结果,其余走公共查找),njhill 随后 APPROVED("Thanks @ivanium!")。另外 ZJY0516 在 issue 评论请求优先合并,明确表示 K3 PD 部署在 truncate_computed_blocks 遇断言错误、预计此 PR 可修复,为该变更提供了真实用户场景支撑。

实现拆解

  1. 定义 connector 能力契约:在 vllm/distributed/kv_transfer/kv_connector/v1/base.py 的 KVConnectorBase_V1 上新增只读属性 supports_divergent_local_hybrid_hits,默认返回 False,语义为「connector 能否在发散本地命中时从外部恢复滞后的 recurrent(Mamba)状态」;NixlBaseConnector(vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py)覆写返回 True,是目前唯一开启该能力的生产实现;MultiConnector(vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py)用 all() 聚合所有子 connector,只有全部支持才返回 True,与既有 _all_support_hma 的聚合风格一致。这一步把「发散命中是否安全」从调度器下沉为 connector 自身的声明,后续新增 connector 时无需改动调度逻辑。
  2. 改造调度器本地命中查找:vllm/v1/core/sched/scheduler.py 的 schedule() 中原先「存在 connector 就走 get_computed_blocks_for_connector,否则走 get_computed_blocks」的内联分支被抽取为独立方法 _get_local_prefix_cache_hit。新逻辑为:connector 存在且具备能力时返回按组发散查找结果(含 hit_diverged 标志);其余情况(无 connector、未知或 store 风格 connector)一律走公共查找并固定 hit_diverged=False。schedule() 调用点的四元组解包保持不变,后续外部查找、CoW 截断逻辑均不受影响。
  3. 测试配套:tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py 新增两个单元测试,验证不支持发散命中的 connector 走公共查找且不调用 divergent 接口、支持者透传 divergent 结果;tests/v1/core/test_scheduler.py 的 _create_hybrid_mamba_connector_scheduler 增加 supports_divergent_hits 参数,原 test_hybrid_fa_deeper_hit_uses_external_mamba_state 重构为参数化 test_hybrid_fa_deeper_hit_respects_connector_lookup_policy,新增 (False, 16, 6, 32) 用例锁定「不支持发散时回落公共边界 32 tokens」的行为;tests/v1/kv_connector/unit/test_multi_connector.py 新增 MockDivergentHMAConnector 并注册,验证 MultiConnector 全支持 / 混合场景下的保守聚合;tests/v1/kv_connector/unit/utils.py 为 MockKVConfig / MockKVConnector 增加 supports_divergent_local_hybrid_hits 配置键与属性,使测试工具可模拟两种能力。
  4. 配置与部署:无生产配置或部署变更;能力通过 connector 类级属性 opt-in,第三方 connector 不声明即保持默认保守行为。
文件 模块 状态 重要度
vllm/v1/core/sched/scheduler.py 调度器 modified 7.41
vllm/distributed/kv_transfer/kv_connector/v1/base.py 连接器基类 modified 5.9
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py NIXL 连接器 modified 5.49
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py 多路连接器 modified 5.69
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py 前缀缓存测试 modified 6.47
tests/v1/core/test_scheduler.py 调度器测试 modified 6.34
tests/v1/kv_connector/unit/test_multi_connector.py 连接器测试 modified 6.07
tests/v1/kv_connector/unit/utils.py 测试工具 modified 4.99

关键符号

Scheduler._get_local_prefix_cache_hit KVConnectorBase_V1.supports_divergent_local_hybrid_hits NixlBaseConnector.supports_divergent_local_hybrid_hits MultiConnector.supports_divergent_local_hybrid_hits MockKVConnector.supports_divergent_local_hybrid_hits test_connector_without_divergent_hit_support_uses_common_lookup test_capable_connector_uses_divergent_partial_hit_lookup test_hybrid_fa_deeper_hit_respects_connector_lookup_policy test_divergent_local_hybrid_hit_capability_is_conservative

关键源码片段

vllm/v1/core/sched/scheduler.py core-logic

修复主入口:新增 _get_local_prefix_cache_hit 方法,将「connector 是否支持发散命中」作为本地前缀缓存查找路径的选择条件,替换 schedule() 中原来的内联分支。

def _get_local_prefix_cache_hit(
    self, request: Request
) -> tuple[KVCacheBlocks, int, int, bool]:
    """按 connector 能力选择本地前缀缓存查找路径。    发散命中(divergent hit)指 full-attention 组与 mamba 组命中到不同
    深度;直接采用较深的 FA 命中时,恢复边界上若无有效 mamba 状态,
    会导致 truncate_computed_blocks 断言崩溃或输出错误。因此只有声明
    supports_divergent_local_hybrid_hits 的 connector(如 NIXL)才允许按组
    发散查找,并由其负责从外部恢复缺失的 recurrent state。
    """
    connector = self.connector
    if connector is not None and connector.supports_divergent_local_hybrid_hits:
        # 返回 (blocks, num_local, shared_prefix_boundary, hit_diverged)。
        return self.kv_cache_manager.get_computed_blocks_for_connector(request)
​
    # 无 connector 或未声明能力的 connector(store 风格、第三方实现):
    # 一律走公共查找,取所有 cache group 共同的保守命中边界,
    # hit_diverged 固定为 False。
    blocks, num_local, shared_prefix_boundary = (
        self.kv_cache_manager.get_computed_blocks(request)
    )
    return blocks, num_local, shared_prefix_boundary, False
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py core-logic

MultiConnector 聚合子 connector 能力:仅当全部子项都支持时才向上暴露,保证混合子连接器场景不误用发散命中。

    @property
    def supports_divergent_local_hybrid_hits(self) -> bool:
        # MultiConnector 聚合多个子 connector 的能力:
        # 只有所有子项都支持时才向上暴露。任意一个子项无法恢复
        # mamba 状态时,调度器都必须退回公共查找,避免调度层
        # 基于不完整的发散命中做决策。
        return bool(self._connectors) and all(
            c.supports_divergent_local_hybrid_hits for c in self._connectors
        )

评论区精华

_get_local_prefix_cache_hit 回退逻辑简化 style

njhill 针对 review 早期版本(含「hit_diverged 且 num_local 非块对齐则回退公共查找」分支)提出简化建议:把「未发散或 num_local % block_size == 0」的条件提前 return,使公共查找分支只保留一份。

结论:作者采纳简化方向,最终合并版本进一步收敛为纯 capability 分流:capable connector 直接返回 divergent 结果,其余一律走公共查找。njhill 随后 APPROVED。 · 已解决

K3 PD truncate_computed_blocks 断言崩溃等待此 PR 修复 question

ZJY0516 在关联 issue 评论请求优先合并:"I ran into an assertion error in truncate_computed_blocks while running K3 PD, and it appears that this PR would resolve it." 说明该修复有真实部署场景支撑。

结论:维护者触发 Buildkite CI 并合并 PR;后续 #51468 在同一测试文件上继续演进发散 FA 命中与外部 Mamba 状态的处理。 · 已解决

风险与影响

  1. 核心调度路径回归:schedule() 是每 step 必经的调度主循环,本地前缀命中查找被抽取重构,虽然逻辑等价且新增测试覆盖 capable/incapable 两条路径,但 PR body 自述 Model evaluation not run,真实 GPU 上的混合模型(K3、Llama 4 等)多组缓存行为未做模型级验证。
  2. 行为回退导致的命中率变化:未声明能力的 connector 从「总走按组发散查找」变为「公共保守查找」,在缓存组发散场景下 prefix cache 命中率可能下降;这是刻意的正确性取舍,但第三方 KV connector 生态需要感知新契约并显式 opt-in。
  3. partial divergent hit 的安全性:最终代码对 capable connector 的非块对齐发散命中直接放行(测试中 num_local=6 即透传),安全性完全依赖 NIXL 在块内边界恢复 mamba 状态的能力;review 早期版本曾包含「num_local % block_size 非零则回退」的守护分支,最终版本未保留,该边界行为建议 GPU 实测确认。
  4. MultiConnector 聚合采用 all(),混合子连接器场景能力被关闭后行为保守,可能损失部分命中率但不会出错。

用户侧:K3 等 hybrid 模型 + KV connector 部署(尤其 PD 分离)不再触发 truncate_computed_blocks 断言崩溃,推理正确性与稳定性得到保证。性能侧:未 opt-in 的 connector 场景 prefix cache 命中率可能小幅回退,属有意取舍;NIXL 场景不受影响。开发侧:KVConnectorBase_V1 新增能力契约,connector 实现方需按语义声明自身能力,MultiConnector 聚合逻辑与既有 HMA 检查风格一致,便于后续扩展。团队侧:为发散命中策略的后续演进(如按 connector 分级支持、把非块对齐回退逻辑下沉到 KV cache manager)奠定了基础。

核心调度路径变更 混合模型 GPU 验证未完成 connector 默认行为回退(性能取舍) 依赖 NIXL 恢复 partial state 待实测

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论