执行摘要

• 一句话：支持多 KV 组的 offloading store 操作
• 推荐动作：该 PR 作为 kv_offload+HMA 系列第 11 部分，核心重构值得关注。建议精读 review 评论，特别是关于 block_size_factor > 1 时的切片和索引推进问题。如果计划使用多组 offloading，请确保这些潜在 bug 已被修复或理解其影响。

功能与动机

为了支持 HMA（异构内存访问）场景，需要允许 KVCacheConfig 包含多个 KV 组，每个组可以有不同的 offload 配置（如块大小）。本 PR 将 store 操作扩展到多组，与之前的 lookup 操作（PR #39401）对应。PR 说明：‘This PR extends the offloading connector to support store (via build_connector_meta) where KVCacheConfig contains multiple groups.’

实现拆解

1. 新增 advance_stored_idx 方法（+7 行）：该函数接收 num_offloadable_tokens，遍历所有 KV 组配置和状态，为每个组计算并更新 next_stored_block_idx。
2. 重构 _get_reqs_to_store 方法（+80/-42 行）：移除之前的硬编码单组断言（assert len(self.config.kv_group_configs) == 1），改为循环遍历 self.config.kv_group_configs 和 req_status.group_states。
3. 引入可 offload token 数量计算：通过 num_computed_tokens + num_scheduled_tokens 并取 min 与 req.num_tokens 对齐，处理异步调度可能缺失的 token 边界。
4. 按组切片和过滤：对每个组，从 next_stored_block_idx 开始切片 offload_keys 和 block_ids，并检查 block_id != 0 以跳过被 sliding window 或 SSM 跳过的块。收集所有有效的 key 放入统一列表，调用 manager.prepare_store 发起存储。
5. 索引更新：在存储准备完成后，调用 req_status.advance_stored_idx(num_offloadable_tokens) 推进所有组的索引。

核心代码片段（编译后整理）：

测试配套：无。本次改动未包含测试文件，建议后续补全。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py（模块 KV 组调度；类别 source；类型 core-logic；符号 advance_stored_idx, _get_reqs_to_store）: 核心调度逻辑的扩展，支持多 KV 组的 store 操作

关键符号：advance_stored_idx, _get_reqs_to_store

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

核心调度逻辑的扩展，支持多 KV 组的 store 操作

    def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
        # 遍历所有 KV 组配置和状态，更新每个组的下一步存储索引
        for group_config, group_state in zip(
            self.config.kv_group_configs, self.group_states
        ):
            # 基于可 offload 的 token 数计算块数（块大小从 group 配置获取）
            num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
            group_state.next_stored_block_idx = num_blocks
​
​
    def _get_reqs_to_store(self, scheduler_output: SchedulerOutput) -> dict[ReqId, TransferSpec]:
        # ... 前置上下文 ...
        num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
        new_offload_keys: list[OffloadKey] = []
        for group_config, group_state in zip(
            self.config.kv_group_configs, req_status.group_states
        ):
            # 根据当前已存储的索引计算要处理的块范围
            num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
            start_block_idx = group_state.next_stored_block_idx
            if num_blocks <= start_block_idx:
                continue
            # 切片：获取从 start_block_idx 到 num_blocks 的 offload keys
            offload_keys = group_state.offload_keys[start_block_idx : num_blocks]
            # 获取对应的 GPU block IDs（可能因 sliding window 而有 0 值）
            offload_block_ids = group_state.block_ids[start_block_idx : num_blocks]
            # 过滤掉 block_id 为 0 的项（表示被 sliding window 或 SSM 跳过）
            # 注意：当 block_size_factor > 1 时，此切片和过滤可能需要调整，
            # 因为一个 offloaded block 对应多个 GPU block，不应只检查第一个
            for offload_key, block_id in zip(offload_keys, offload_block_ids):
                if block_id == 0:
                    break
                new_offload_keys.append(offload_key)
        # 收集的 keys 用于后续 manager.prepare_store
        # ... 后续代码 ...

评论区精华

• offload_block_ids 切片不正确（block_size_factor > 1） (correctness): PR 作者未回复，但 PR 已合并，可能预期在后续 PR 中修复或此场景暂不使用。
• 无条件推进 next_stored_block_idx 可能跳过未处理的块 (correctness): 同前，PR 已合并，风险未被解决。
• line 480 无条件设置 next_stored_block_idx 加剧跳过问题 (correctness): 同上。

风险与影响

• 风险：
  • 回归风险：对于单组用户，行为理论上不变，但重构可能引入新 bug，特别是切片逻辑在 block_size_factor > 1 时存在已知争议。
  • 数据完整性风险：由于 advance_stored_idx 和后续无条件设置 next_stored_block_idx 可能跳过应存储的块，导致数据丢失。
  • 缺少测试：本次改动没有新增测试，增加回归概率。
  • 兼容性风险：无，因为是内部模块，不涉及 API 变化。
• 影响：
  • 用户影响：普通用户无直接影响，因为多组配置尚未在前端暴露。HMA 系列仍处于开发中。
  • 系统影响：对于配置了多个 KV 组的场景，store 操作现在可以正确按组处理。但当前仍不支持 sliding window 和 SSM，这些组的块会被跳过。
  • 团队影响：为后续 PR 奠定了基础，但需要修复 review 指出问题。
  • 风险标记：核心路径变更, 缺少测试覆盖, block_size_factor > 1 潜在 bug, 索引推进逻辑风险

关联脉络

• PR #39401 [kv_offload+HMA][9/N]: Support lookup with multiple KV groups: 同一系列 PR，扩展 lookup 功能，本 PR 扩展 store 功能，共同实现多 KV 组的 complete offloading 支持。