Prhub

#48481 [KV Connector] Fix PD async scheduling race condition for hybrid attn models

原始 PR 作者 arpera 合并时间 2026-07-16 18:42 文件变更 4 提交数 4 评论 5 代码增减 +158 / -7

执行摘要

修复混合模型 PD 异步调度精度回归

PR 意图修复两个已知 bug:TP>1 时 PD 异步调度精度崩溃(issue #37285)和 TP==1 时 Qwen3.5 精度塌陷(issue #42182)。同时解决了之前 PR #45096 中识别但未完全处理的竞态条件之一——分配时清零与已到达 KV 的冲突。

建议精读此 PR,特别是调度器侧跳过清零的设计模式和使用外部令牌授权推导覆盖块的方法。值得注意的设计决策:避免在异步管道中加入 CUDA 同步点,而是通过元数据告知调度器排除特定块。

讨论亮点
  1. njhill 对连接器 API 方法的担忧:最初提交通过一个 KVConnectorMetadata.get_blocks_to_skip_kv_cache_zeroing 新方法实现,njhill 指出这会给每个加载连接器增加必须实现的方法,且隐式要求连接器清空未使用的尾部,存在遗漏的失败路径。

  2. njhill 的替代方案:njhill 推送了一个额外提交,通过从 get_num_new_matched_tokens 返回的外部令牌授权直接推导需跳过的块,去除了连接器 API 方法,使逻辑集中在调度器侧。

  3. 作者立场:arpera 表示需要任何修复以发布可复现的结果,即使不覆盖所有连接器也接受。最终 njhill 的方案被合并。

实现拆解

  1. 新增 KVCacheManager 方法:在 vllm/v1/core/kv_cache_manager.py 中添加 get_zeroing_block_ids_in_range()record_blocks_for_zeroing()。前者根据 token 范围返回将被远程 KV 加载覆盖的注意力块 ID,后者用于在加载失败时重新记录需要清零的块。

  2. 调度器排除覆盖块:在 vllm/v1/core/sched/scheduler.pyschedule() 中,当处理异步 KV 加载请求时,调用 get_zeroing_block_ids_in_range() 获取被覆盖块,存入 _skip_zero_block_ids 集合。修改 _get_new_block_ids_to_zero() 方法,从新块 ID 中移除这些跳过块,并清空集合。

  3. 暴露单类型管理器属性:在 vllm/v1/core/single_type_kv_cache_manager.py 中添加 records_new_block_ids 属性,用于判断该管理器的块是否需要清零,从而区分注意力组(需清零)和 Mamba 状态组(不清零)。

  4. 单元测试:在 tests/v1/kv_connector/unit/test_nixl_connector_hma.py 中新增三个测试:验证零化块 ID 只覆盖加载的注意力块、验证调度器正确过滤连接器加载的块、验证失败加载后重新记录未写入块。

文件 模块 状态 重要度
vllm/v1/core/kv_cache_manager.py 缓存管理器 modified 7.66
vllm/v1/core/sched/scheduler.py 调度器 modified 7.58
tests/v1/kv_connector/unit/test_nixl_connector_hma.py 连接器测试 modified 6.88
vllm/v1/core/single_type_kv_cache_manager.py 缓存管理器 modified 5.57

关键符号

get_zeroing_block_ids_in_range record_blocks_for_zeroing _get_new_block_ids_to_zero records_new_block_ids

关键源码片段

vllm/v1/core/kv_cache_manager.py core-logic

核心文件:新增 get_zeroing_block_ids_in_range 和 record_blocks_for_zeroing 方法,提供被加载块的 ID 查询和失败重清零能力。

def get_zeroing_block_ids_in_range(
    self, request_id: str, start_token: int, end_token: int
) -> list[int]:
    """返回请求块中属于 [start_token, end_token) 范围的块 ID,
    仅包含那些被标记为需要清零的管理器(注意力组)。
    """
    ids: list[int] = []
    for mgr in self.coordinator.single_type_managers:
        if mgr.records_new_block_ids: # 仅注意力组需要清零
            start_idx = start_token // mgr.block_size
            end_idx = cdiv(end_token, mgr.block_size)
            blocks = mgr.req_to_blocks[request_id]
            ids.extend(blk.block_id for blk in blocks[start_idx:end_idx])
    return idsdef record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None:
    """重新记录请求从 start_token 开始的块到清零列表,
    用于 KV 加载失败后的回退。start_token 必须块对齐。
    """
    for mgr in self.coordinator.single_type_managers:
        if mgr.records_new_block_ids:
            assert start_token % mgr.block_size == 0
            start_idx = start_token // mgr.block_size
            blocks = mgr.req_to_blocks[request_id]
            mgr.new_block_ids.extend(blk.block_id for blk in blocks[start_idx:])
vllm/v1/core/sched/scheduler.py core-logic

核心调度逻辑:新增 _skip_zero_block_ids 集合,在 schedule 中收集被加载覆盖的块,在 _get_new_block_ids_to_zero 中过滤,最后传给 worker。

def _get_new_block_ids_to_zero(self) -> list[int] | None:
    # 每步清空 new_block_ids 以防止无限增长
    new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids()
    if not self.needs_kv_cache_zeroing:
        return None
​
    if self._skip_zero_block_ids:
        # 移除将被远程 KV 加载覆盖的块
        skip = self._skip_zero_block_ids
        new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip]
        skip.clear() # 清空集合,避免重复跳过
​
    return new_block_ids_to_zero or None # 空列表转为 None
tests/v1/kv_connector/unit/test_nixl_connector_hma.py test-coverage

新增三个单元测试,验证零化块 ID 范围过滤、调度器跳过逻辑以及失败加载后重新记录,覆盖主要代码路径。

@pytest.mark.cpu_test
def test_zeroing_block_ids_cover_only_loaded_attention_blocks():
    """仅注意力组(records_new_block_ids=True)的块被返回,
    切片到给定 token 范围。Mamba 块不被清零。"""
    manager = _make_fake_kv_cache_manager()
    # 本地缓存 token [0,16),远程加载覆盖 [16,56)
    assert manager.get_zeroing_block_ids_in_range("req-1", 16, 56) == [11, 12, 13]@pytest.mark.cpu_test
def test_scheduler_filters_connector_loaded_blocks_from_zeroing():
    """在 _skip_zero_block_ids 中的块不会出现在零化列表中。"""
    scheduler = object.__new__(Scheduler)
    scheduler.needs_kv_cache_zeroing = True
    scheduler.kv_cache_manager = FakeKVCacheManager()
    scheduler._skip_zero_block_ids = {10, 12}
    assert scheduler._get_new_block_ids_to_zero() == [9, 11]
    assert not scheduler._skip_zero_block_ids # 应已清空

评论区精华

连接器 API 方法 vs 调度器推导 设计

njhill 认为添加 `get_blocks_to_skip_kv_cache_zeroing` 连接器方法会增加每个连接器的实现负担,且存在隐式尾部清零假设和遗漏失败路径的问题。arpera 表示希望尽快修复,不介意范围窄一些。

结论:njhill 提交了替代方案,通过 `get_num_new_matched_tokens` 返回的 token 范围在调度器侧直接推导被覆盖块,移除了连接器 API 方法。 · 已解决

失败路径处理 正确性

njhill 指出如果异步加载失败,之前被跳过清零的块可能包含无效数据,需要重新记录为零化。

结论:arpera 在后续提交中添加了 `record_blocks_for_zeroing` 方法,在失败路径中被调用,重新记录未写入块。 · 已解决

风险与影响

  • 回归风险:若 _skip_zero_block_ids 集合并未正确覆盖所有被加载块,可能导致部分块被错误跳过清零,引发精度问题。测试覆盖了正常和失败路径,但未覆盖极端并发场景。
  • 性能影响:每次异步加载时需遍历管理器列表计算块 ID,但范围较小,开销可忽略。
  • 连接器兼容性:当前方案不依赖连接器实现,但需连接器正确报告 num_new_local_computed_tokensnum_computed_tokens,否则可能错误跳过清零。
  • 错误处理record_blocks_for_zeroing 在失败路径中被调用,但断言 start_token 块对齐,若传入未对齐值会导致崩溃。
  • 用户影响:使用混合模型(如 Qwen3.5、NemotronH)且启用 PD 异步调度的用户将恢复精度,无需回退到 --no-async-scheduling
  • 系统影响:仅在 KV 连接器启用异步加载时激活,不影响非 PD 或同步调度场景。
  • 团队影响:维护了调度器与连接器之间的清晰接口,降低了后续连接器实现的耦合。
核心调度路径变更 依赖连接器正确报告 token 数 失败路径断言可能崩溃

关联 Issue

#37285 [Bug]: PD disaggregation for SSM models requires `--no-async-scheduling` when TP>1
#42182 [Bug]: Qwen3.5-27B Disagg accuracy gsm8k collapses with async scheduling when TP==1
#45096 [Mamba][PD] support async scheduling for mamba PD
#48481 [KV Connector] Fix PD async scheduling race condition for hybrid attn models

完整报告

参与讨论