Prhub

#41289 [Bugfix][SimpleCPUOffloadBackend] Dedup in-flight CPU offload stores across scheduler steps

原始 PR 作者 ivanium 合并时间 2026-05-13 16:53 文件变更 3 提交数 6 评论 3 代码增减 +126 / -21

执行摘要

修复 SimpleCPUOffloadScheduler eager 模式下跨 steps 重复 offload 的问题

在 eager 模式下,SimpleCPUOffloadScheduler 通过 per-step 局部集合与 cached_block_hash_to_block 去重 store。由于 hash 只在 store 完成时注册,当某个请求在 step N+1 命中 GPU prefix cache 时,若前一步的 store 尚未完成,相同的 GPU blocks 会被重复 offload。具体描述见 PR body:'In eager mode, SimpleCPUOffloadScheduler deduplicated stores via a per-step local set plus cached_block_hash_to_block. The hash is only registered at store completion, so when req B hits the GPU prefix cache in step N+1 before req A's step-N store lands, the same GPU blocks get re-offloaded.'

建议合入。这是一个精确的 bugfix,修复了竞态条件导致的重复 offload 问题,代码改动量小,设计清晰,且有完善的测试和 CI 集成。

讨论亮点

review 中 gemini-code-assist[bot] 提出应基于 block hash 而非 block ID 去重,以避免不同物理块包含相同数据时的重复 offload。作者 ivanium 回应:不应去重不同 GPU block,即使 hash 相同,因为 KVCacheManager 将它们视为独立块,且 prefix cache 可能命中有相同 hash 的不同块,去重会导致期望的 offload 缺失。该设计讨论已达成一致,维持基于 block ID 的去重。

实现拆解

  1. 引入实例级去重集合:在 SimpleCPUOffloadScheduler.__init__ 中添加 self._in_flight_store_gpu_blocks: set[int] = set(),存储当前正在存储中的 GPU block ID。
  2. 替换 per-step 局部集合:在 _prepare_eager_store_specs 中使用 self._in_flight_store_gpu_blocks 代替原来的局部变量 gpu_blocks_this_step。当决定 offload 某个 block 时,先检查其 ID 是否已在 in_flight 中,若在则跳过,避免重复调度;否则加入集合。
  3. 清理已完成的 store:在 _process_store_event 中,当 store 事件完成后,从 _in_flight_store_gpu_blocks 中移除对应的 GPU block IDs,使这些块可以再次被调度(如果后续需要)。
  4. 测试配套:修改 test_scheduler.py,将 make_requestextra_tokens 参数化,新增 test_max_hit_len_cap_drops_last_full_block 边界测试和 test_eager_in_flight_store_dedup_across_steps 回归测试,验证跨步骤去重正确性。
  5. CI 集成:在 .buildkite/test_areas/misc.yaml 中添加 tests/v1/simple_kv_offload 作为 source file dependency 和对应的 pytest 命令,确保新测试在 CI 中自动运行。
文件 模块 状态 重要度
tests/v1/simple_kv_offload/test_scheduler.py 卸载测试 modified 7.06
vllm/v1/simple_kv_offload/manager.py KV 卸载 modified 5.99
.buildkite/test_areas/misc.yaml CI 配置 modified 2.5

关键符号

__init__ _prepare_eager_store_specs _process_store_event test_eager_in_flight_store_dedup_across_steps test_max_hit_len_cap_drops_last_full_block

关键源码片段

tests/v1/simple_kv_offload/test_scheduler.py test-coverage

新增跨步骤去重回归测试和边界测试,验证修复正确性,是质量保障的核心

def test_eager_in_flight_store_dedup_across_steps() -> None:
    """Eager mode:跨步骤去重,第二个步骤应跳过已在飞行中的 GPU block。"""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler
​
    # 步骤 1:分配并注册块,构造 store 事件但不完成
    num_blocks = 2
    req1 = make_request(num_blocks=num_blocks)
    kv_blocks1 = _alloc_and_register(fix, req1, num_blocks)
    sched.update_state_after_alloc(req1, kv_blocks1, num_external_tokens=0)
    block_ids1 = kv_blocks1.get_block_ids()
    sched_out1 = make_scheduler_output(
        {req1.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req1.request_id: block_ids1},
    )
    meta1 = sched.build_connector_meta(sched_out1)
    # store 事件已生成但未完成,block 应记录为 in_flight
​
    # 步骤 2:使用相同 token 的请求,预期命中 GPU prefix cache
    req2 = Request(
        request_id="req2-dedup",
        prompt_token_ids=req1.prompt_token_ids,
        sampling_params=req1.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req1._block_hasher,
    )
    kv_blocks2 = _alloc_and_register(fix, req2, num_blocks)
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=num_blocks * BLOCK_SIZE)
    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: 1},
        new_reqs={req2.request_id: block_ids2},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    # 由于步骤 1 的 store 尚未完成,步骤 2 的相同 GPU block IDs
    # 应被 in_flight 集合排除,验证无新 store 事件
    # 实际断言:store_gpu_blocks 为空(已由 in_flight 过滤)
    assert len(meta2.store_gpu_blocks) == 0, "Should not re-offload in-flight blocks"
​
    # 完成步骤 1 的 store 后,步骤 3 应能正常 store
    simulate_store_completion(sched, meta1.store_event)
    req3 = Request(
        request_id="req3-after",
        prompt_token_ids=req1.prompt_token_ids,
        sampling_params=req1.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req1._block_hasher,
    )
    kv_blocks3 = _alloc_and_register(fix, req3, num_blocks)
    sched.update_state_after_alloc(req3, kv_blocks3, num_external_tokens=num_blocks * BLOCK_SIZE)
    block_ids3 = kv_blocks3.get_block_ids()
    sched_out3 = make_scheduler_output(
        {req3.request_id: 1},
        new_reqs={req3.request_id: block_ids3},
    )
    meta3 = sched.build_connector_meta(sched_out3)
    # store 事件正常生成
    assert meta3.store_event >= 0
vllm/v1/simple_kv_offload/manager.py core-logic

核心变更文件:引入 _in_flight_store_gpu_blocks 实例变量,修改 _prepare_eager_store_specs 和 _process_store_event 实现跨步骤去重

# 在 __init__ 中新增实例变量
self._in_flight_store_gpu_blocks: set[int] = set()# _prepare_eager_store_specs 中的去重逻辑
# 使用实例集合代替 per-step 局部集合
in_flight = self._in_flight_store_gpu_blocks# 遍历请求数据
for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
    # ... 计算块和状态
    # 如果该 GPU block ID 已经在 in_flight 中,则跳过
    if gpu_block_id in in_flight:
        continue
    # ... 分配 CPU block 并准备 store
    # 将新调度的 block 加入 in_flight
    in_flight.update(gpu_block_ids)
    gpu_block_pool.touch(gpu_block_ids) # 防止被释放# _process_store_event 中的清理
def _process_store_event(self, event_idx: int) -> None:
    """处理完成的 store 事件,从 in_flight 集合中移除对应的 GPU block。"""
    transfer = self._store_event_to_blocks.pop(event_idx)
    if not self._lazy_mode:
        # 仅 eager 模式需要管理 in_flight 集合
        self._in_flight_store_gpu_blocks.difference_update(transfer.gpu_block_ids)
    self._process_store_completion(transfer.gpu_block_ids, transfer.cpu_block_ids)

评论区精华

去重应基于 block hash 还是 block ID 设计

gemini-code-assist[bot] 提出应基于 hash 去重,以避免不同物理块包含相同数据时的重复 offload。ivanium 回应:不应去重不同 GPU block,即使 hash 相同,因为 KVCacheManager 将它们视为独立块,且 prefix cache 可能命中有相同 hash 的不同块,去重会导致期望的 offload 缺失。

结论:维持基于 block ID 的去重,作者解释合理,reviewer 未进一步反对。 · 已解决

风险与影响

主要风险在于 _process_store_event 中清理集合的逻辑:如果 store 事件丢失或异常,对应的 GPU block IDs 不会被移除,可能导致这些块永久跳过 offload,但随着 KV cache 释放这些块会被重用。该风险较低,因为事件丢失是小概率事件,且系统有重试机制。此外,该变更为 eager 模式专用,不影响 lazy 模式。测试覆盖了单机场景,但多 worker(TP/PP)场景下的并发安全未显式测试,需依赖已有的事件计数机制。

影响范围限于使用 SimpleCPUOffloadBackend 的用户(CPU offload 场景),修复后可减少不必要的 CPU 写入,提升 CPU cache 利用率和整体吞吐。对其他 offload 后端无影响。新增测试纳入 CI,提高了回归防护。

核心路径变更 竞态条件修复 并发安全 TP/PP 场景未显式测试

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论