Prhub

#41445 [kv_offload+HMA][13/N]: Enable HMA support

原始 PR 作者 orozery 合并时间 2026-05-01 19:30 文件变更 6 提交数 1 评论 3 代码增减 +662 / -251

执行摘要

启用 OffloadingConnector HMA 支持

原 offloading connector 仅支持单一 KV 缓存组,无法兼容使用滑动窗口或混合注意力状态的 HMA 模型。用户 naveline67 在评论中反馈在 Qwen3.6 上使用 ——kv-offloading-size 16 工作正常,表明实际部署效果良好。

值得精读,特别是测试框架如何通过 GPUBlock 和 kv_cache_groups 参数模拟多组场景,以及 SupportsHMA 接口的设计模式。

讨论亮点

gemini-code-assist[bot] 在 review 中指出测试工具中 offload_addresses_to_skip 的计算可能错误,应使用 logical_offset % block_size_factor 而非 -logical_offset % block_size_factor,但该问题未在代码中看到修复。markmc 审核并批准。

实现拆解

  1. 在 OffloadingConnector 上实现 SupportsHMA 接口:继承 SupportsHMA,新增 request_finished_all_groups 方法,委托给 scheduler.request_finished(忽略 block_ids 参数)。
  2. 调整 Scheduler 的 request_finished 签名:移除 block_ids 参数,简化接口。
  3. 重构测试框架:引入 GPUBlock dataclass 表示带组索引的 GPU 块;重构 RequestRunner 支持可选的 kv_cache_groups 和 block_size_factor,模拟多组不同 block size 场景。
  4. 增加多组测试用例:test_two_groups_full_and_sliding_window 和 test_two_groups_different_block_sizes 验证调度逻辑。
  5. 增强端到端测试:用 MODEL_PARAMS 参数化替换硬编码列表,加入 google/gemma-3-1b-it 等 HMA 模型,测试 CPU offloading 的延迟与准确性。
文件 模块 状态 重要度
tests/v1/kv_connector/unit/offloading_connector/utils.py 测试工具 modified 7.52
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py 调度器测试 modified 7.28
tests/v1/kv_connector/unit/test_offloading_connector.py 端到端测试 modified 6.86
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py 连接器 modified 7.03
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py 调度器 modified 4.96

关键符号

OffloadingConnector.request_finished_all_groups OffloadingConnectorScheduler.request_finished GPUBlock _mock_get_layers _update_gpu_block_idx test_two_groups_full_and_sliding_window test_cpu_offloading

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

评论区精华

offload_addresses_to_skip 计算可能存在逻辑错误 正确性

gemini-code-assist[bot] 在 review 中指出,在 utils.py 中计算 offload_addresses_to_skip 时,应使用 logical_offset % block_size_factor 而不是 -logical_offset % block_size_factor。当前实现可能导致跳过错误数量的子块。

结论:尚未在代码中看到对应的修复,但 PR 已合并。该问题可能影响测试准确性,但端到端测试结果正常。 · unresolved

风险与影响

测试工具中潜在的 offload_addresses_to_skip 计算错误可能导致测试跳过错误数量的子块,影响某些场景的覆盖率可信度,但端到端测试通过。多 KV 组路径首次启用可能在极端配置下触发未预期行为,但已覆盖滑动窗口和不同块大小组合。

对使用 offloading connector 的用户,现在可以支持 HMA 模型(如 Gemma-3、Mamba、Falcon-H1),提升兼容性。非 HMA 模型行为不变。测试框架重构可能影响其他基于此测试工具单元的测试,但概率低。

测试工具潜在 Bug 多 KV 组路径首次启用

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论