执行摘要

• 一句话：SimpleCPUOffload 添加 PCP/DCP 支持
• 推荐动作：推荐开发者和技术管理者精读此 PR，特别是理解如何利用已有 KVCacheCoordinator 的 CP 支持通过简单缩放块大小来扩展 offload 功能。这是一次增量扩展的典型例子，展示了最小入侵式修改。同时，测试辅助函数的组织方式值得参考。

功能与动机

根据 PR 描述，目的为“Add PCP + DCP support by scaling kv block sizes by pcp_world_size * dcp_world_size”。SimpleCPUOffloadConnector 依赖 KVCacheCoordinator，而 KVCacheCoordinator 已支持上下文并行，因此需要让 offload 模块也能感知 CP 并正确缩放逻辑块大小。

实现拆解

实现分为以下几个步骤：

1. 计算 CP world size 并移除断言：在 SimpleCPUOffloadScheduler.__init__ 中从 parallel_config 读取 decode_context_parallel_size 和 prefill_context_parallel_size，计算乘积 self.cp_world_size，并移除原有的 assert dcp_world_size == 1 and pcp_world_size == 1 断言。
2. 缩放 lazy target 估计：在 _estimate_lazy_target_blocks 静态方法中添加 cp_world_size 参数，遍历 KV cache groups 时将每个 spec 的 block_size 乘以 cp_world_size，确保 lazy offload 模式下保留足够的 GPU 空闲块。同时调用处传入 self.cp_world_size。
3. 缩放状态更新中的块大小：在 update_state_after_alloc 中，从 kv_cache_groups 中获取 block_size 后乘以 self.cp_world_size，使得 computed tokens 到 blocks 的映射考虑 CP 扩展。
4. 增加测试覆盖：在 tests/v1/simple_kv_offload/test_scheduler.py 中新增辅助函数 (_make_cp_vllm_config, _make_cp_scheduler, _make_cp_request, _allocate_cp_gpu_blocks) 和三个测试用例：test_cp_block_size_scaling 验证不同 CP 组合下块大小缩放是否正确；test_cp_eager_store_and_load_roundtrip 验证 eager 模式在 CP 下 store 和 load 的完整性；test_cp_lazy_target_blocks_scaling 验证 lazy 模式 target blocks 数量符合期望。这些测试依赖现有的 KVCacheCoordinator 和 BlockPool，仅需设置并行配置和虚拟块大小。

关键文件：

• vllm/v1/simple_kv_offload/manager.py（模块 卸载器；类别 source；类型 core-logic）: 核心逻辑变更：计算 cp_world_size 并移除 CP 断言，在块大小计算和 lazy target 估计中缩放 cp_world_size。
• tests/v1/simple_kv_offload/test_scheduler.py（模块 测试；类别 test；类型 test-coverage；符号 _make_cp_vllm_config, _make_cp_scheduler, _make_cp_request, _allocate_cp_gpu_blocks）: 新增 CP 场景的测试覆盖，包括辅助函数和三个测试用例验证块大小缩放、eager store-load 往返及 lazy target 计算。

关键符号：SimpleCPUOffloadScheduler.init, SimpleCPUOffloadScheduler._estimate_lazy_target_blocks, SimpleCPUOffloadScheduler.update_state_after_alloc, _make_cp_vllm_config, _make_cp_scheduler, _make_cp_request, _allocate_cp_gpu_blocks, test_cp_block_size_scaling, test_cp_eager_store_and_load_roundtrip, test_cp_lazy_target_blocks_scaling

关键源码片段

vllm/v1/simple_kv_offload/manager.py

核心逻辑变更：计算 cp_world_size 并移除 CP 断言，在块大小计算和 lazy target 估计中缩放 cp_world_size。

# __init__ 中：从 parallel_config 读取 dcp/pcp world size 并计算乘积
class SimpleCPUOffloadScheduler:
    def __init__(self, vllm_config, kv_cache_config, ...):
        # ...
        # 新增：计算 CP world size
        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
        self.cp_world_size = dcp_world_size * pcp_world_size
        # 移除了 assert dcp_world_size == 1 and pcp_world_size == 1
        # ...
        # lazy target 估计中缩放 block_size
        self._target_free = self._estimate_lazy_target_blocks(
            kv_cache_config,
            vllm_config.scheduler_config.max_num_batched_tokens,
            self.cp_world_size,
        )
        # ...
​
    @staticmethod
    def _estimate_lazy_target_blocks(
        kv_cache_config: "KVCacheConfig",
        max_num_batched_tokens: int,
        cp_world_size: int = 1,
    ) -> int:
        """GPU blocks to keep available per step in lazy mode."""
        WATERMARK_RATIO = 1.0
        target = 0
        for g in kv_cache_config.kv_cache_groups:
            spec = g.kv_cache_spec
            # 关键：将 block_size 乘以 cp_world_size 以匹配 CP 缩放
            block_size = spec.block_size * cp_world_size
            if isinstance(spec, MambaSpec):
                target += 2
            elif isinstance(spec, SlidingWindowSpec):
                target += cdiv(spec.sliding_window, block_size) + 1
            else:
                target += cdiv(max_num_batched_tokens, block_size)
        return int(target * (1 + WATERMARK_RATIO))

评论区精华

PR 作者在一条 review 评论中说明了测试调整原因：“Preexisting tests were failing due to caches being exactly block sizes. + 1 here to balance out the - 1 logic that is included for needing to recompute the last token to create logits.” 这表明先前的 num_tokens = num_blocks * BLOCK_SIZE 在计算请求 token 数时与缓存分配逻辑存在边界偏移，通过加 1 解决。其余 review 仅由 gemini-code-assist 自动审查且无具体反馈，ivanium 给予了 LGTM 批准。整体讨论较少，说明变更清晰且风险可控。

• 测试中 num_tokens 计算调整 (testing): 调整被接受，作为预存在问题的修复。

风险与影响

• 风险：核心风险在于 块大小缩放的正确性：若 cp_world_size 计算错误或与其他模块（如 GPU block pool）的块大小不一致，可能导致缓存索引越界或内存浪费。但由于底层 KVCacheCoordinator 已处理 CP 逻辑，且测试验证了不同 CP 组合下的行为，风险可控。另一个潜在风险：update_state_after_alloc 中乘以 cp_world_size 可能影响 prefix caching 匹配，因为 block hash 也依赖于虚拟块大小；测试要求 hash_block_size 等于 scaled virtual block size 来模拟真实行为，确保了匹配。非 CP 场景 (dcp=1, pcp=1) 行为不变（因为 cp_world_size=1）。总体风险较低。
• 影响：用户影响：启用上下文并行（PCP 和/或 DCP）的用户现在可以同时使用 SimpleCPUOffloadConnector 进行 KV 缓存卸载，之前该配置会被断言拒绝。系统影响：修改仅影响 SimpleCPUOffloadScheduler 的初始化及两个内部方法，不涉及其他模块。团队影响：无，变更已在批准后合并。
• 风险标记：核心路径变更, 块大小缩放兼容性

关联脉络

• 暂无明显关联 PR