Prhub

#41228 [kv_offload+HMA][12/N]: Scheduler-side support for sliding window groups

原始 PR 作者 orozery 合并时间 2026-05-01 11:59 文件变更 1 提交数 4 评论 6 代码增减 +280 / -98

执行摘要

OffloadingConnector 调度器支持滑动窗口和 Mamba KV 缓存组

为了支持滑动窗口注意力和 Mamba 等非全注意力模型的 KV 缓存 offloading,需要调度器侧能够识别这些模型的缓存特征。本 PR 解决了调度器目前仅支持全注意力(FullAttentionSpec)的局限,通过泛化 KVCacheSpec 和引入滑动窗口相关逻辑,使调度器可以正确处理滑动窗口和 Mamba 的缓存块边界及生命周期。Issue 未提及,但从 PR 标题 [kv_offload+HMA][12/N] 可以看出这是系列工作中的一环,最终目标是使 OffloadingConnector 支持 HMA(异构内存访问)特性。

建议有相关背景的开发者精读本 PR,重点关注滑动窗口块的生命周期设计、_touch 的 LRU 更新策略,以及 _remove_pending_job 的安全性讨论。非直接涉及 KV offload 的成员可略读了解架构演化。

讨论亮点

_remove_pending_job 安全性:Gemini Code Assist 建议使用 .get().discard() 避免 KeyError。作者 orozery 回应认为不可能,宁愿捕获异常而非忽略。该讨论未改变代码。
FullAttentionSpec 命名确认:markmc 询问注释中应为 FullAttentionSpec 而非其他,作者确认,已解决。
_touch 收敛性:markmc 对排序和病态迭代提出疑问,orozery 解释与 HybridKVCacheCoordinator 类似,排序按组大小递减有助于收敛,已解决。

实现拆解

  1. 导入新类型:从 vllm.v1.kv_cache_interface 导入 FullAttentionSpecSlidingWindowSpecMambaSpec,替换原有的 KVCacheSpec 导入,为后续基于类型判断滑动窗口大小做准备。
  2. 实现 get_sliding_window_size_in_blocks:根据 spec 类型计算窗口覆盖的 offloaded 块数(全注意力返回 None,滑动窗口返回 ceiling 除法,Mamba 返回 1)。
  3. 扩展 GroupOffloadConfig:在 NamedTuple 中新增 sliding_window_size_in_blocks 字段,在 SchedulerOffloadConfig.from_spec 中动态计算该值。
  4. 拆分 TransferJobStatus.gpu_block_ids:拆为 non_sliding_window_block_idssliding_window_block_ids,前者仅在请求结束时注册 pending jobs,后者在 store 创建时立即注册,以支持滑动窗口块的提前回收。
  5. 新增命中计数RequestGroupState 增加 num_hit_blocks,由 update_num_hit_blocks 初始化,用于查找决策。
  6. 重写查找逻辑:增加 _sliding_window_lookup_touch_sliding_window_sort_key_remove_pending_job 等方法,处理后缀匹配的滑动窗口查找和 LRU 更新。
  7. 调整清理路径_on_request_finished 等适配新数据结构。
  8. 测试配套:本 PR 未包含测试文件,将在下个 PR 随 SupportsHMA 一起添加。
文件 模块 状态 重要度
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py KV 调度器 modified 9.05

关键符号

get_sliding_window_size_in_blocks update_num_hit_blocks _sliding_window_sort_key _remove_pending_job _touch _lookup

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

_remove_pending_job 使用 .get() .discard() 安全性建议 正确性

gemini-code-assist 建议改用 .get() 和 .discard() 避免 KeyError。

结论:作者认为该情况不可能发生,宁愿捕获异常而不忽略。 · 已解决

FullAttentionSpec 名称确认 question

markmc 询问注释中是否应为 FullAttentionSpec。

结论:作者确认,注释已修正。 · 已解决

_touch 排序与收敛性解释 设计

markmc 对 _touch 中排序和防止病态迭代的机制提出疑问,希望添加注释。

结论:作者解释与 HybridKVCacheCoordinator 类似,排序按组大小递减有助于收敛。 · 已解决

风险与影响

  • _remove_pending_job 使用直接字典/集合访问,若状态同步意外(如重复完成事件)可能抛出 KeyError,作者认为不可能但仍是潜在薄弱点。
  • 滑动窗口块与非滑动窗口块的生命周期管理较复杂,初始化或清理逻辑有误可能导致引用计数错误或内存泄漏。
  • _touch 的排序依赖组配置,极端情况下可能性能退化。
  • 本 PR 仅修改调度器侧,尚未与 Worker 侧 SupportsHMA 集成,中间状态可能导致不一致。
  • 用户影响:无直接可见变更,需等待后续 SupportsHMA 集成。
  • 系统影响:核心调度器数据结构与算法重构,可能影响所有 KV 传输作业的稳定性。
  • 团队影响:作为 12/N 系列的一部分,后续 PR 将直接依赖此变更,需要保持同步。
缺少测试覆盖 核心数据结构变更 滑动窗口块生命周期风险

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论