Prhub

#44490 [Bugfix][Core] Fix host memory leak from undrained new_block_ids

原始 PR 作者 Sunt-ing 合并时间 2026-07-07 15:32 文件变更 3 提交数 7 评论 10 代码增减 +18 / -15

执行摘要

修复 new_block_ids 未清空导致的主机内存泄漏

用户报告在长时间运行 max_tokens=1 负载时主机 RSS 线性增长(#44175)。根本原因是 #35219 添加的 new_block_ids 记录仅在全注意力/MLA 块分配时追加,但清空操作仅在有 Mamba 层时执行,导致非 Mamba 模型的列表无限增长。gc.freeze 隐藏了此泄漏,使其表现为内存碎片。修复是必需的,否则长时间服务会导致 OOM。

值得精读,特别是对于理解 V1 调度器与缓存管理器的交互契约以及条件记录/清空的设计权衡。展示了如何通过最小化 flag 传递来避免回归且保持正确性。

讨论亮点

主要讨论围绕设计权衡:作者最初避免添加构造函数标志,担心默认 False 会静默禁用零化路径;但 reviewer(njhill)采纳了 chaunceyjiang 的建议,添加 _record_new_block_ids 属性以减少非 Mamba 模型的不必要列表操作,确保不会引入性能开销。最终接受 flag 方案,并确认在协调器中正确传递参数。

实现拆解

  1. 在 SingleTypeKVCacheManager.init 中添加 needs_kv_cache_zeroing 参数,并根据 spec 类型设置 _record_new_block_ids 标志,控制是否记录新块 ID。
  2. 将 allocate_new_blocks 和 allocate_external_computed_blocks 中的记录条件由类型检查替换为 _record_new_block_ids,避免非 Mamba 模型的不必要记录。
  3. 在 Scheduler.schedule 方法中,无条件调用 kv_cache_manager.take_new_block_ids() 清空列表,然后根据 needs_kv_cache_zeroing 决定是否传递新块 ID 给 SchedulerOutput。
  4. 在 KVCacheCoordinator 中创建管理器时传递 needs_kv_cache_zeroing 参数,确保记录/清空行为与模型类型一致。
文件 模块 状态 重要度
vllm/v1/core/single_type_kv_cache_manager.py 缓存管理 modified 6.71
vllm/v1/core/sched/scheduler.py 调度器 modified 6.01
vllm/v1/core/kv_cache_coordinator.py 协调器 modified 4.56

关键符号

SingleTypeKVCacheManager.__init__ SingleTypeKVCacheManager.allocate_new_blocks SingleTypeKVCacheManager.allocate_external_computed_blocks Scheduler.schedule KVCacheCoordinator.__init__

关键源码片段

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

核心修复文件:添加 needs_kv_cache_zeroing 构造函数参数和 _record_new_block_ids 标志,控制 new_block_ids 的记录,避免非 Mamba 模型的不必要开销。

# 构造函数新增 needs_kv_cache_zeroing 参数,用于控制是否记录新分配的块 ID
def __init__(
    self,
    kv_cache_spec: KVCacheSpec,
    block_pool: BlockPool,
    enable_caching: bool,
    kv_cache_group_id: int,
    scheduler_block_size: int,
    dcp_world_size: int = 1,
    pcp_world_size: int = 1,
    needs_kv_cache_zeroing: bool = False, # 新增参数:是否需要 KV cache 零化
    max_admission_blocks_per_request: int | None = None,
) -> None:
    # ...
    # 仅当需要零化且 spec 类型属于需要零化的类型时,才记录新块 ID
    self._record_new_block_ids = needs_kv_cache_zeroing and type(kv_cache_spec) in (
        FullAttentionSpec,
        TQFullAttentionSpec,
        MLAAttentionSpec,
        HiddenStateCacheSpec,
    )
    self.new_block_ids: list[int] = []
    # ...# allocate_new_blocks 方法中,之前使用重复的类型检查,现在使用 _record_new_block_ids 标志
def allocate_new_blocks(self, request_id: str, num_new_blocks: int) -> list[KVCacheBlock]:
    # ...
    if self._record_new_block_ids: # 只有需要记录时才追加
        self.new_block_ids.extend(b.block_id for b in new_blocks)
    return new_blocksdef allocate_external_computed_blocks(self, request_id, ...):
    # ...
    if self._record_new_block_ids:
        self.new_block_ids.extend(b.block_id for b in allocated_blocks)
vllm/v1/core/sched/scheduler.py core-logic

修复调度器侧:无条件清空 new_block_ids,但仅当需要零化时传递给 SchedulerOutput。

# 在 schedule 方法末尾,无条件清空 new_block_ids 列表
# 只有需要零化时才将块 ID 传递给 SchedulerOutput
# 这确保了非 Mamba 模型不会积累未清空的 ID
new_attn_block_ids = self.kv_cache_manager.take_new_block_ids()
new_block_ids_to_zero = (
    (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None
)scheduler_output = SchedulerOutput(
    # ... 其他字段
    new_block_ids_to_zero=new_block_ids_to_zero,
    # ...
)

评论区精华

无条件清空 new_block_ids 的安全性 正确性

作者分析表明,由于 zeroing 路径是唯一消费 new_block_ids 的地方,无条件清空不会影响行为;Mamba 模型之前已经由 needs_kv_cache_zeroing 分支清空,因此行为无变化。

结论:接受无条件清空方案,对 Mamba 模型无影响。 · 已解决

添加 _record_new_block_ids 属性避免性能开销 设计

原作者倾向于避免添加构造函数标志,担心默认 False 会静默禁用零化;reviewer njhill 在此建议下添加 _record_new_block_ids 属性,以避免非 Mamba 模型的不必要 list 操作开销。

结论:接受 flag,确保在 KVCacheCoordinator 中正确传递参数。 · 已解决

风险与影响

变更涉及调度器和缓存管理器的核心路径,但逻辑简单审慎。主要风险是 needs_kv_cache_zeroing 参数是否在所有管理器构造路径中正确传递;由于只有 KVCacheCoordinator 一处创建管理器,风险可控。另外,无条件清空 new_block_ids 的额外开销极低(仅 pop 列表),且 _record_new_block_ids 确保非 Mamba 模型追加操作为空。

对所有 V1 标准注意力模型(如 Llama、Qwen、Gemma、DeepSeek-MLA)修复了长时间运行的内存泄漏,服务稳定性显著提升。Mamba 或混合 Mamba 模型无行为变化。影响范围覆盖大部分主流模型,但不涉及 API、前端或数据路径。

调度器核心路径变更 缺少测试覆盖

关联 Issue

#44175 [Bug]: Linear host RSS growth + step-up in E2E latency under sustained `max_tokens=1` classification load (Gemma-3-1b, V1, prefix caching)

完整报告

参与讨论