执行摘要

• 一句话：修复 DecodeBenchConnector 在前缀缓存命中时错误选择块的问题。
• 推荐动作：这个 PR 值得精读，因为它修复了一个关键 bug 并展示了良好的测试实践。设计决策基于块大小计算区间，考虑了多种缓存组类型，值得借鉴。

功能与动机

根据 PR body，DecodeBenchConnector 当前选择每个 KV 缓存组的前 ceil(num_external_tokens / block_size) 个块，但当请求有本地前缀缓存命中时，外部令牌区间从该前缀之后开始，而连接器仍然从块零填充，可能覆盖共享前缀缓存块并留下新分配的后缀块未填充。需要选择实际覆盖外部计算令牌的块区间。

实现拆解

1. 修改 DecodeBenchConnectorScheduler 类的 __init__ 方法，存储 kv_cache_groups 引用，以便在后续方法中使用。
2. 重写 update_state_after_alloc 方法：计算总计算令牌数和本地计算令牌数；对于每个 KV 缓存组，区分非循环缓冲区和循环缓冲区；对于非循环缓冲区，基于块大小计算外部块的起始位置和结束位置，选择对应的块；对于循环缓冲区，保留所有块；过滤掉空块；存储待填充的块 ID。
3. 添加新的单元测试覆盖前缀命中、非对齐块、空填充跳过等场景，确保修复正确。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py（模块 KV 连接器；类别 source；类型 core-logic；符号 update_state_after_alloc）: 核心逻辑变更文件，修复了块选择算法。
• tests/v1/kv_connector/unit/test_decode_bench_connector.py（模块 连接器测试；类别 test；类型 test-coverage；符号 test_decode_bench_connector_prefix_hit_fills_only_suffix_blocks, test_decode_bench_connector_selects_external_blocks_per_group, test_decode_bench_connector_selects_unaligned_external_block, test_decode_bench_connector_skips_null_padding）: 测试覆盖文件，验证修复的正确性。

关键符号：update_state_after_alloc

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py

核心逻辑变更文件，修复了块选择算法。

def update_state_after_alloc(
    self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
):
    """更新分配后的状态，为每个组选择覆盖外部令牌的块。"""
    req_id = request.request_id
    if num_external_tokens == 0:
        return
​
    # 计算总计算令牌数和本地计算令牌数
    total_computed_tokens = request.num_tokens - 1
    num_local_computed_tokens = total_computed_tokens - num_external_tokens
​
    block_ids_per_group_list: list[list[int]] = []
    for group_idx, (group, group_blocks, group_block_size) in enumerate(
        zip(self.kv_cache_groups, blocks.blocks, self.group_block_sizes, strict=True)
    ):
        # 检查是否为循环缓冲区组
        is_circular_buffer = all(
            isinstance(spec, CircularBufferSpec)
            for spec in iter_layer_specs(group.kv_cache_spec)
        )
        if not is_circular_buffer:
            # 计算外部块的起始和结束位置
            num_computed_blocks = cdiv(total_computed_tokens, group_block_size)
            external_block_start = num_local_computed_tokens // group_block_size
            assert 0 <= external_block_start <= num_computed_blocks <= len(group_blocks), (
                "DecodeBenchConnector block range exceeds allocated blocks: "
                f"request={req_id}, group={group_idx}, "
                f"range=[{external_block_start}, {num_computed_blocks}), "
                f"allocated={len(group_blocks)}"
            )
            selected_blocks = group_blocks[external_block_start:num_computed_blocks]
        else:
            # 循环缓冲区组保留所有块
            selected_blocks = group_blocks
​
        # 过滤空块并获取块 ID
        block_ids = [block.block_id for block in selected_blocks if not block.is_null]
        if not block_ids:
            logger.warning(
                "DecodeBenchConnector: No blocks selected for KV cache group "
                "%d with %d external tokens for request %s",
                group_idx, num_external_tokens, req_id,
            )
        block_ids_per_group_list.append(block_ids)
​
    block_ids_per_group = tuple(block_ids_per_group_list)
    # 存储待填充的块 ID
    self._pending_fills[req_id] = (block_ids_per_group, num_external_tokens)
    self._filled_requests.add(req_id)
​
    block_counts = tuple(len(group) for group in block_ids_per_group)
    logger.debug(
        "DecodeBenchConnector: Selected %d total blocks across %d KV cache "
        "groups for request %s (per-group counts: %s)",
        sum(block_counts), len(blocks.blocks), req_id,
        ", ".join(map(str, block_counts)),
    )

评论区精华

PR body 中提供了详细的测试计划和结果，包括单元测试和 E2E 测试，验证了修复的正确性。Review 评论中 ivanium 批准了 PR，没有实质性争议。

• 修复块选择逻辑以覆盖前缀缓存命中场景 (correctness): 修复正确，测试覆盖全面，已合并。

风险与影响

• 风险：风险较低，因为变更仅影响 DecodeBenchConnector 的块选择逻辑，且测试覆盖全面。但需要确保在非循环缓冲区组中，块索引计算正确，避免越界。代码中的断言提供了运行时检查。
• 影响：影响 DecodeBenchConnector 的 KV 缓存填充行为，确保在前缀缓存命中时正确填充后缀块，避免数据覆盖和未填充块。对用户无直接影响，但提升了基准测试的准确性。
• 风险标记：核心路径变更, 测试覆盖良好

关联脉络

• PR #54879 [Bugfix][KV Connector] Safely fill circular buffers in DecodeBench: 修复 DecodeBench 在填充含 CircularBufferSpec 的 KV 缓存组时损坏元数据的 bug，与本 PR 互补，本 PR 修复块选择逻辑。