Prhub

#52912 [Bugfix][KVOffload] P2P tier declares REQUEST_LEVEL on the producer leg

原始 PR 作者 liranschour 合并时间 2026-08-31 18:07 文件变更 5 提交数 6 评论 23 代码增减 +282 / -9

执行摘要

修复 PD 热前缀下 P2P tier 无供给致 30s 超时重算

issue #52808 明确描述了这一反直觉缺陷:「The effect is inverted from the intent: the better the producer's prefix cache works, the more often the Multi Tier path fails to deliver.」。复现数据显示:producer 冷前缀时 consumer 0.16s 内加载 256MiB(268,959,744 字节);producer 预热后 consumer 等满 32.9s、加载 0 字节再本地重算。根因有两道闸门:offloading/scheduler.py 在 BLOCK_LEVEL 策略下把 prefix-hit chunks 全部跳过(next_stored_chunk_idx = num_chunks),且所有 tier 的 on_new_request 都返回默认 BLOCK_LEVEL(grep -c REQUEST_LEVEL 为 0),request_level_tiers 恒为空,_cascade_existing_blocks_to_request_level_tiers 永不运行。

值得精读。重点关注三件事:一是 8 行修复如何解除「前缀缓存越好 PD 越慢」的反直觉失效,二是差分测试把手工镜像的 store 修剪与 load demand 两条路径钉在一起的设计思路,三是 nilig 的版本回溯分析——同一个 patch 在 v0.27.1 与 main 上表现完全不同,提醒跨版本验证的必要性。

讨论亮点

讨论围绕「这个修复在 main 上是否真的有效」展开,是本 PR 最有价值的部分:

  • nilig 最初在 v0.27.1 上打补丁验证,发现 warm 场景仍超时:「The consumer's fetch demands 112 chunks (both KV cache groups). The producer supplies 57... 55 stay demanded forever.」——55 个是滑动窗口组 out-of-window 的 block_id == 0 的 chunk,在 _build_store_jobs 中被丢弃。
  • nilig 随后在 main-based 构建上复测并撤回:「On main this PR looks good in my testing and I withdraw the sufficiency objection... the divergence is #51840, which landed after v0.27.1 and changes lookup to return HIT_PENDING instead of RETRY.」HIT_PENDING 保持连击 streak,使滑动窗口扫描止步于 trailing window;RETRY 会重置 streak,导致扫描全部 56 个 SWA chunks。
  • Etelis 发现 main 上真实存在的并发缺陷:请求 A 的 GPU→primary store 在途时,请求 B 的 REQUEST_LEVEL prepare_store 返回 keys_to_store == [],cascade 的 lookup(k) is HIT 过滤把全部 key 丢弃,_flush_pending_promotions 不处理该场景,PD consumer 也不发 LookupMsg,peer 永远不被服务。作者确认这是独立缺陷,转 issue #53062 单独修复。
  • Etelis 调侃测试规模:「Most are duplicated or not needed on an 8 line change :)」,建议参数化合并;作者采纳,将 TestOnNewRequestPolicy 类重构为参数化函数,并保留唯一覆盖 kv_request_id 分支的 producer_no_id 用例。

实现拆解

  1. 策略声明(唯一源码变更):vllm/v1/kv_offload/tiering/p2p/manager.py 新增 OffloadPolicy 导入;P2PSecondaryTierManager.on_new_request 在解析 P2PDestInfo 后,仅当请求为 remote_decoder(producer leg)且 kv_request_id 非空时返回 RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL),其余情况保持默认 BLOCK_LEVEL,把变更范围严格限定在 producer leg。

  2. 级联生效路径:REQUEST_LEVEL 使 TieringOffloadingManager.prepare_store 对已驻留 primary 的 block 执行 _cascade_existing_blocks_to_request_level_tiers,完整 block 集进入 P2PSecondaryTierManager.submit_store;调度器侧 next_stored_chunk_idx 归零,prefix-hit chunks 不再被跳过,数据从 primary tier 读取、不触碰 GPU。创建 store job 还让对端首个 FetchMsg 能绑定 session,on_request_finished 可快速失败未匹配需求,不再干等 30s。

  3. 策略单测:tests/v1/kv_offload/tiering/p2p/test_manager.py 新增参数化 test_on_new_request_policy,覆盖四种请求形态:producer(升级)、普通请求、consumer leg、空 remote_decoder 的 malformed 请求(后三者保持 BLOCK_LEVEL)。

  4. 级联差分:tests/v1/kv_offload/tiering/test_tiering_offloading.py 将原级联测试参数化为 fully_warm 与 mixed 两种形态,钉住「全部 block 已驻留时 keys_to_store == []、无 GPU→primary 传输,但 request-level tier 仍收到完整 block 集、block-level tier 不被触碰」。

  5. 调度器对称性测试:tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py(+235/-4)是三处测试中最重的,新增 _SCAN_BEHAVIOR 声明表与 test_scan_behavior_declared_for_every_lookup_result(强制每个 LookupResult 枚举成员显式声明滑动窗口扫描行为)、test_sliding_window_demand_is_store_reachable(对齐 chunk 数、窗口 chunk 数、EAGLE group、尾部 chunk 数四维网格,验证 load 路径查询的每个 chunk 都能被 store 路径保留)、test_request_level_supply_covers_consumer_demand(cold / gpu_warm / primary_warm × 是否含 SWA group),在 _build_store_jobs 真实构建 key 集的层面验证供给覆盖 demand。

  6. 配套与结果:tests/v1/kv_connector/unit/offloading_connector/utils.py 将 to_keys 参数从 list[int] 放宽为 Iterable[int](mypy 兼容,配合传入 range);PR 内全部 121 个测试通过(新增 5 个)。无配置、schema 或部署变更。

文件 模块 状态 重要度
vllm/v1/kv_offload/tiering/p2p/manager.py KV 卸载 modified 5.84
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py 调度器 modified 7.44
tests/v1/kv_offload/tiering/p2p/test_manager.py P2P 管理 modified 5.01
tests/v1/kv_offload/tiering/test_tiering_offloading.py 分级卸载 modified 4.52
tests/v1/kv_connector/unit/offloading_connector/utils.py 测试工具 modified 3.65

关键符号

P2PSecondaryTierManager.on_new_request test_on_new_request_policy test_prepare_store_cascades_existing_blocks_to_request_level_tiers test_scan_behavior_declared_for_every_lookup_result test_sliding_window_demand_is_store_reachable test_request_level_supply_covers_consumer_demand _demanded_keys

关键源码片段

vllm/v1/kv_offload/tiering/p2p/manager.py core-logic

唯一源码改动:在 on_new_request 中为 producer leg 声明 REQUEST_LEVEL,是修复的核心。

    @override
    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
        """解析 kv_transfer_params 一次并打开出站 session。
​
        consumer 侧(PD 的 remote_prefiller 或对称 P2P 的 remote_kv_source)
        按 remote_host:remote_port 打开通往 producer 的 session,使 submit_load
        触发时能立即发出 FetchMsg。prefiller(producer)侧不预先创建 session,
        等 consumer 入站连接到达 _accept_new_peers 时再创建。
​
        Producer-leg 请求(存在 remote_decoder 子字典)声明 REQUEST_LEVEL:
        对端需要该请求的每一个 block,而不只是本次新计算的 block。
        """
        _annotate_req_context(req_context)
        source = req_context.get_state(P2PSourceInfo)
        if source is not None:
            self._get_or_create_session(source.peer_id)
​
        # P2PDestInfo 非空说明本请求是 producer leg,即 remote_decoder 形态;
        # kv_request_id 非空说明确有对端消费者在等待。两者同时满足才升级为
        # REQUEST_LEVEL,使 TieringOffloadingManager 把 primary 中已驻留的
        # block(含 prefix hits)级联到 P2P store 路径。
        # 普通请求、consumer leg 与缺少 kv_request_id 的 malformed 请求
        # 仍保持默认 BLOCK_LEVEL,把变更范围限制在 producer leg。
        dest = req_context.get_state(P2PDestInfo)
        if dest is not None and dest.kv_request_id:
            return RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)
        return RequestOffloadingContext()
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py test-coverage

最大测试变更(+235/-4),新增 LookupResult 扫描行为声明与 store/load 路径对称性差分测试,是本 PR 测试价值的核心。

# 每种 LookupResult 在窗口 -1 扫描中的行为:是否保持连击 streak 决定扫描终点,
# 进而决定 demanded 集合的宽度。新增枚举成员必须在此声明,防止 demand 集合
# 悄悄变宽。元组为(lookup 调用次数 , 扫描返回的 end)。
_SCAN_BEHAVIOR = {
    LookupResult.HIT: (1, 3),
    LookupResult.HIT_PENDING: (1, None),
    LookupResult.RETRY: (3, None),
    LookupResult.MISS: (3, 0),
}
​
​
@pytest.mark.parametrize("result", list(LookupResult))
def test_scan_behavior_declared_for_every_lookup_result(result: LookupResult):
    """每个 LookupResult 成员都必须有明确的滑动窗口扫描语义。"""
    assert result in _SCAN_BEHAVIOR, f"{result} has no declared scan behavior"
    expected_lookups, expected_end = _SCAN_BEHAVIOR[result]
​
    keys = to_keys([1, 2, 3])
    sched = _make_scheduler_with_lookup(dict.fromkeys([1, 2, 3], result))
​
    assert sched._sliding_window_lookup(keys, 1, _EMPTY_REQ_CTX) == expected_end
    assert len(sched.manager.lookup.call_args_list) == expected_lookups
​
​
@pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64])
@pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3])
@pytest.mark.parametrize("is_eagle_group", [False, True])
@pytest.mark.parametrize("tail_chunks", [0, 1, 5])
def test_sliding_window_demand_is_store_reachable(
    alignment_chunk_count: int,
    sliding_window_chunks: int,
    is_eagle_group: bool,
    tail_chunks: int,
):
    """load 路径查询的每个 chunk 必须被 store 路径保留。
​
    is_store_reachable_swa_chunk 会修剪 _sliding_window_lookup 永远不会查询的
    chunk,但两者是手工镜像的:本测试把关系钉死,而
    test_is_store_reachable_swa_chunk 只钉了谓词自身的取值。被查询过的 key
    就是被 demand 的 key,因为 secondary tier 命中会立即排队 promotion。
    """
    storable_chunk_count = alignment_chunk_count + tail_chunks
    # 未验证的 EAGLE group 会让两侧各放宽一个 chunk
    required_window = sliding_window_chunks + int(is_eagle_group)
​
    sched = _make_scheduler_with_lookup(
        {i: LookupResult.HIT_PENDING for i in range(storable_chunk_count)}
    )
    sched._sliding_window_lookup(
        to_keys(range(storable_chunk_count)), required_window, _EMPTY_REQ_CTX
    )
​
    # 收集扫描真实发出的 lookup 调用,即 demand 集合
    demanded = {
        int(get_offload_block_hash(lookup_call.args[0]).decode())
        for lookup_call in sched.manager.lookup.call_args_list
    }
    assert demanded, "扫描必须查询到内容"
    unsuppliable = sorted(
        chunk_idx
        for chunk_idx in demanded
        if not is_store_reachable_swa_chunk(
            chunk_idx,
            storable_chunk_count,
            alignment_chunk_count,
            sliding_window_chunks,
            is_eagle_group,
        )
    )
    assert not unsuppliable, (
        f"chunks {unsuppliable} 被 load 路径查询但被 store 路径修剪,"
        "warm producer 永远无法供应它们"
    )

评论区精华

warm producer 供给不足的反对意见及其撤回 正确性

nilig 在 v0.27.1 上打补丁验证发现 warm 场景仍超时:consumer demand 112 chunks,producer 只供给 57,55 个 out-of-window 且 block_id == 0 的滑动窗口 chunk 被 _build_store_jobs 丢弃,故质疑修复不充分。后在 main-based 构建上复测:warm-GPU 0.075s、warm-CPU-only 0.072s 均加载 268MiB,与 cold 一致,正式撤回反对。

结论:确认是 v0.27.1 与 main 的行为差异(#51840 的 HIT_PENDING vs RETRY)所致;main 上修复有效。 · 已解决

并发共享前缀下 HIT_PENDING 级联丢弃全部 key 正确性

Etelis 发现请求 A 的 GPU→primary store 在途时,请求 B 的 REQUEST_LEVEL prepare_store 返回 keys_to_store == [],cascade 的 lookup(k) is HIT 过滤把全部 key 丢弃且 ready=0;_flush_pending_promotions 只处理 promotions,PD consumer 从不发 LookupMsg,server 侧 re-poll 机制不生效,B 的 peer 永远不被服务。

结论:确认为独立缺陷,作者承诺以 issue #53062 单独修复,不并入本 PR。 · 已解决

store 路径修剪与 load 路径 demand 的对称性测试缺口 测试

is_store_reachable_swa_chunk 修剪 _sliding_window_lookup 永不查询的 chunk,两者手工镜像但无测试绑定;test_is_store_reachable_swa_chunk 只钉谓词取值,lookup 测试只钉扫描行为,任一侧可单独漂移。PR 以 _SCAN_BEHAVIOR 声明表加网格差分测试补齐该缺口。

结论:已通过新增差分测试修复缺口。 · 已解决

8 行变更配大量测试是否冗余 style

Etelis 评论「most are duplicated or not needed on an 8 line change :)」,建议参数化合并 TestOnNewRequestPolicy 的四个断言与 tiering 级联测试;作者采纳并重构为参数化测试,同时保留 producer_no_id 用例(唯一覆盖 kv_request_id 分支的测试)。

结论:参数化重构完成,测试保留但去重。 · 已解决

风险与影响

  • P2P store 带宽上升:producer leg 每次 PD 轮询都会全量进入 store 路径(含 prefix hits),P2P 传输量从「仅新 block」变为「全部 block」。8K token 场景约 256MiB / 0.07s,远小于修复前 30s 超时重算,净收益明确,但多副本部署下带宽成本需关注。
  • 首次实际启用 request-level 级联:REQUEST_LEVEL 路径在代码中早已存在但从未被任何 tier 声明(grep 为 0),本 PR 是首次启用,TieringOffloadingManager._cascade_existing_blocks_to_request_level_tiers 的并发行为未经大规模生产验证。
  • 依赖 #51840 的 HIT_PENDING 语义:修复有效性依赖 lookup 在 promotion 时返回 HIT_PENDING 并保持扫描 streak 的行为;_SCAN_BEHAVIOR 表把该行为固化进测试,未来改变 LookupResult 语义或扫描逻辑会第一时间暴露,但也意味着改动成本显式化。
  • issue 中提示的 PYTHONHASHSEED 排查建议对本次失败模式具有误导性,排查同路径问题时需注意区分。
  • 用户侧:PD 分离部署 + 多级 tiering(primary + P2P secondary)的用户,warm 前缀场景从 30s 超时本地重算变为百毫秒级直接供给,冷热场景行为拉齐。
  • 系统侧:P2P secondary tier 的 store 路径现在承载 producer leg 的完整 block 集,带宽占用上升;on_request_finished 的快速失败路径开始生效,session 绑定时序更早。
  • 团队侧:kv_offload 调度器测试基础设施获得两类可复用的新模式——LookupResult 扫描行为声明表和 store/load 路径对称性差分测试;后续扩展 LookupResult 枚举或调整 SWA 修剪逻辑时会被强制审视。
核心路径策略变更 依赖上游 HIT_PENDING 语义 首次启用 request-level 级联

关联 Issue

#52808 [Bug]: PD Multi Tier supplies nothing when the producer already holds the blocks

完整报告

参与讨论