Prhub

#45357 [Bugfix] Defer block freeing until in-flight steps finish under async scheduling + PD KV consumer

原始 PR 作者 llx-08 合并时间 2026-06-16 05:36 文件变更 9 提交数 9 评论 20 代码增减 +541 / -18

执行摘要

延迟块释放修复异步调度与 PD 分解竞态

异步调度下,调度器可能在步骤 N 输出被处理前就调度步骤 N+1。若请求在步骤 N 输出后完成或抢占,其 KV 块被立即释放,但步骤 N+1 仍可能写入这些块。若新 PD 消费者请求重用这些块,NIC/RDMA 写入的接收 KV 与过时 GPU 写入未排序,导致数据损坏。本 PR 延迟块归还以修复此竞态。

推荐使用 PD 分解的团队立即合并此修复。PR 的设计模式(步骤栅栏 + 延迟释放)值得学习,展示了如何在不破坏现有逻辑的前提下解决复杂的并发竞态。核心贡献者阅读可关注 pop_blocks_for_free 的设计和延迟释放队列的条件判断。

讨论亮点

Review 中主要讨论包括:NickLucche 质疑添加日志的必要性,njhill 回复保留日志是担心性能影响,但也可移除;njhill 提供了简化实现并补充了测试用例;ZJY0516 要求提供精度测试结果,llx-08 后提供了 GSM8K 基准测试结果确认正确性。后续 njhill 还追加了覆盖 PP 场景的提交。

实现拆解

  1. 引入步骤序列号栅栏:在 Scheduler 中增加 sched_step_seq 和 processed_step_seq,SchedulerOutput 携带 sched_seq。非空步骤调度时递增 sched_seq;update_from_output 处理输出时更新 processed_step_seq。
  2. 拆分缓存释放逻辑:在 KV 缓存管理器层新增 pop_blocks_for_free,仅移除请求的 bookkeeping 并返回块列表,不归还 block pool。原有的 free 改为调用 pop_blocks_for_free 后立即归还。
  3. 延迟释放队列:Scheduler 维护 _deferred_frees: deque[(fence_seq, blocks)],只有当 processed_step_seq >= fence_seq 时才能安全归还。
  4. 调度器释放路径:统一的 _free_request_blocks 方法。若延迟启用且存在在调度步骤,则调用 pop_blocks_for_free 并将块入队;否则直接 free。
  5. 排空延迟队列:update_from_output 中每次处理后调用 _drain_deferred_frees,将满足条件的块通过 block_pool.free_blocks 归还。
  6. 启用条件:仅当 max_concurrent_batches > 1(异步调度或 PP)且 KV 传输配置为消费者时设置 defer_block_free = True。
文件 模块 状态 重要度
tests/v1/core/test_deferred_block_free.py 测试覆盖 added 8.42
vllm/v1/core/single_type_kv_cache_manager.py 缓存管理器 modified 7.37
vllm/v1/core/kv_cache_coordinator.py 缓存协调器 modified 6.83
vllm/v1/core/kv_cache_manager.py 缓存管理层 modified 6.64
vllm/v1/core/sched/scheduler.py 调度器 modified 8.05

关键符号

_free_request_blocks _drain_deferred_frees pop_blocks_for_free free schedule update_from_output

关键源码片段

tests/v1/core/test_deferred_block_free.py test-coverage

新增完整测试套件,覆盖延迟释放的启用条件、正常推迟、无在调步骤立即释放、abort 场景等。

def _create_deferring_scheduler():
    """创建强制启用延迟释放的调度器(生产环境还需 KV 消费者)"""
    # 直接设置 defer_block_free 为 True,绕过生产环境的条件判断
    scheduler = create_scheduler(model=MODEL, async_scheduling=True)
    scheduler.defer_block_free = True
    return scheduler
​
​
def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5):
    """
    模拟一个请求的 prefill(步骤1)和一次 speculatively 调度的 decode(步骤2),
    用于测试延迟释放场景。返回 (request, out0, out1)。
    """
    request = create_requests(
        num_requests=1,
        num_tokens=NUM_PROMPT_TOKENS, # 33 tokens, block_size=16 时占用 3 个 block
        max_tokens=max_tokens,
        stop_token_ids=[STOP_TOKEN_ID],
    )[0]
    scheduler.add_request(request)
    out0 = scheduler.schedule() # 步骤 1:prefill
    assert out0.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS
    out1 = scheduler.schedule() # 步骤 2:decode(提前调度)
    assert out1.num_scheduled_tokens[request.request_id] == 1
    return request, out0, out1
vllm/v1/core/single_type_kv_cache_manager.py core-logic

核心拆分:新增 pop_blocks_for_free 方法,使调用方可以延迟归还 block pool。原有 free 方法基于 pop_blocks_for_free 实现。

# 在 SingleTypeKVCacheManager 中新增的方法
def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
    """
    弹出请求的 bookkeeping 并返回其 block 列表,但 *不* 归还给 block pool。
    调用方必须后续通过 block_pool.free_blocks 归还 block,顺序须为逆序。
    """
    # 默认返回 [],防止请求在被分配 block 前就被释放(abort)
    req_blocks = self.req_to_blocks.pop(request_id, [])
    # 同时清理缓存的 block 计数
    self.num_cached_block.pop(request_id, None)
    return req_blocks
​
​
def free(self, request_id: str) -> None:
    """
    立即释放请求的所有 block(逆序归还给 block pool)。
    现在基于 pop_blocks_for_free 实现,保持原有行为不变。
    """
    # 逆序释放,使尾部 block 优先回池
    self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id)))

评论区精华

延迟释放启用时的日志级别 设计

NickLucche 建议移除或降级 info 日志,因为这只是修复 bug,不应强调 ' 我们做对了 '。njhill 回复保留日志是担心性能影响,但也可移除。

结论:最终保留了 info 日志,但后续提交已移除无用日志。 · 已解决

精度测试验证 测试

ZJY0516 要求提供精度测试结果,llx-08 提供 GSM8K 对比数据确认正确性。njhill 进一步补充了测试用例并覆盖 PP 场景。

结论:精度测试通过,覆盖提升。 · 已解决

风险与影响

延迟释放可能在高 KV 缓存压力下降低池利用率,但该功能仅在有异步调度和 PD KV 消费者时启用,降低了影响范围。新引入的序列号逻辑依赖 update_from_output 按 FIFO 顺序调用,若有变更可能导致不一致。pop_blocks_for_free 与 free 的分离需要确保所有调用方都正确迁移,避免双重释放或遗漏释放。此修复不涉及模型侧变更,风险可控。

影响范围仅限于启用异步调度且使用 PD KV 消费者(分解部署)的场景。普通用户不受影响。对于使用 PD 分解的团队,此修复至关重要,能消除因竞态导致的 KV 损坏和错误输出。系统在启用后会有较小的额外计算开销(检查序列号和队列操作),但可忽略。测试覆盖包括专门的单元测试和已在 PD 环境中验证的精度测试。

核心路径变更 并发竞态修复 KV 缓存管理变更 需要精度验证 启用条件有限

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论