# PR #39403 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[kv_offload+HMA][11/N]: Support store with multiple KV groups
- 合并时间：2026-04-26 01:00
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/39403

---

# 执行摘要

- 一句话：支持多 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 操作

```python
    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
        # ... 后续代码 ...

```

# 评论区精华

| 讨论者 | 要点 | 类别 | 状态 |
|--------|------|------|------|
| gemini-code-assist[bot] | `offload_block_ids` 切片未对齐 offloaded block 边界（`block_size_factor > 1`），且 `block_id != 0` 检查不完整 | 正确性 | 未解决 |
| gemini-code-assist[bot] | `advance_stored_idx` 无条件推进可能跳过未处理块，导致数据丢失 | 正确性 | 未解决 |
| gemini-code-assist[bot] | 循环后 `next_stored_block_idx = num_blocks` 同样可能跳过块 | 正确性 | 未解决 |
| markmc | 建议在过滤块循环前添加注释说明，便于理解 | 文档 | 已采纳（PR 已包含注释） |

- 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 支持。