执行摘要

• 一句话：修复 MRv2 内核块大小维度不一致
• 推荐动作：值得精读，重点关注 prepare_kernel_block_sizes 和 BlockTables 的扩展逻辑。这是 MRv2 与连接器集成的重要修复。

功能与动机

当 MRv2 默认用于密集 Qwen 模型时，发现使用 FlashInfer 后端且 block_size=128、kernel_block_size=64 时，KV 缓存张量通过 register_kv_caches API 暴露给连接器的形状不一致。MRv1 与 MRv2 的形状差异导致 CI 测试失败（问题 #42846）。本 PR 确保跨连接器的一致性。

实现拆解

1. 在 attn_utils.py 中引入 prepare_kernel_block_sizes 函数，用于计算每个 KV 缓存组的最小公共内核块大小；
2. 修改 init_attn_backend 返回 kernel_block_sizes 列表，并将正确的大小传递给元数据构建器；
3. 更新 BlockTables（block_table.py）以接受 kernel_block_sizes，计算 blocks_per_kv_block 并扩展块表尺寸，在 append_block_ids 中展开块 ID；
4. 调整 model_runner.py 的加载顺序：先初始化 attention 后端获取 kernel_block_sizes，再创建 BlockTables 并传递；
5. 移除 config/vllm.py 中 kv_transfer_config 存在时自动回退 MRv1 的防护逻辑，并删除对应测试。

关键文件：

• vllm/v1/worker/gpu/attn_utils.py（模块 注意力工具；类别 source；类型 dependency-wiring；符号 init_attn_backend, _reshape_kv_cache, prepare_kernel_block_sizes）: 核心修改：重构 init_attn_backend 返回 kernel_block_sizes，新增 prepare_kernel_block_sizes 逻辑，修改 _reshape_kv_cache 签名。
• vllm/v1/worker/gpu/block_table.py（模块 块表；类别 source；类型 core-logic；符号 BlockTables.init, BlockTables.append_block_ids, BlockTables.init_block_table_layout_tensors）: 支持内核块大小与逻辑块大小的映射，扩展块表尺寸，修改 append_block_ids 展开块 ID。
• vllm/v1/worker/gpu/model_runner.py（模块 模型运行器；类别 source；类型 data-contract；符号 ModelRunner.initialize_kv_cache）: 调用顺序调整：先初始化 attention 后端获取 kernel_block_sizes，再创建 BlockTables 和 kv cache。
• vllm/config/vllm.py（模块 配置；类别 source；类型 core-logic；符号 VllmConfig.use_v2_model_runner）: 移除 kv_transfer_config 存在时强制 MRv1 的回退逻辑，使 MRv2 在连接器存在时正常生效。
• vllm/v1/worker/gpu/spec_decode/eagle/speculator.py（模块 投机解码；类别 source；类型 core-logic；符号 EagleSpeculator.set_attn）: 适配 init_attn_backend 的新返回签名，增加占位接收 kernel_block_sizes。
• tests/test_config.py（模块 测试；类别 test；类型 test-coverage；符号 test_use_v2_model_runner_defaults_to_v1_when_kv_connector_present）: 删除测试 test_use_v2_model_runner_defaults_to_v1_when_kv_connector_present，因为回退逻辑已被移除。

关键符号：prepare_kernel_block_sizes, init_attn_backend, _reshape_kv_cache, BlockTables.init, BlockTables.append_block_ids, VllmConfig.use_v2_model_runner

关键源码片段

vllm/v1/worker/gpu/attn_utils.py

核心修改：重构 init_attn_backend 返回 kernel_block_sizes，新增 prepare_kernel_block_sizes 逻辑，修改 _reshape_kv_cache 签名。

# 从 attn_utils.py 中提取的核心修改：
# init_attn_backend 现在分三阶段执行
​
def init_attn_backend(kv_cache_config, vllm_config, device, active_layer_names=None):
    attn_backends = {}
    attn_groups = []
​
    # Phase 1: discover attention groups for each kv cache group.
    for kv_cache_group_id, kv_cache_group_spec in enumerate(kv_cache_config.kv_cache_groups):
        # ...（与原逻辑相同）
        attn_groups.append([group_map[key] for key in group_order])
​
    # Phase 2: pick a kernel block size per kv cache group that is supported
    # by all backends within that group.
    kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups)
​
    # Phase 3: create metadata builders and determine cudagraph support.
    attn_backend_workspace = None
    min_cg_support = AttentionCGSupport.ALWAYS
    min_cg_attn_backend = None
    for kv_cache_group_id, groups in enumerate(attn_groups):
        kv_cache_group_spec = kv_cache_config.kv_cache_groups[kv_cache_group_id]
        kernel_block_size = None
        if kv_cache_group_id < len(kernel_block_sizes):
            kernel_block_size = kernel_block_sizes[kv_cache_group_id]
        for group in groups:
            group.create_metadata_builders(
                vllm_config=vllm_config,
                device=device,
                kernel_block_size=kernel_block_size, # 之前为 None
                num_metadata_builders=1,
            )
            # ...（cudagraph support 检查）
​
    return (
        attn_backends,
        attn_groups,
        AttentionCGSupportInfo(...),
        kernel_block_sizes, # 新增返回
    )

vllm/v1/worker/gpu/block_table.py

支持内核块大小与逻辑块大小的映射，扩展块表尺寸，修改 append_block_ids 展开块 ID。

# BlockTables 构造函数片段
class BlockTables:
    def __init__(self, block_sizes, max_num_reqs, max_num_batched_tokens,
                 max_num_blocks_per_group, device, kernel_block_sizes, ...):
        self.kernel_block_sizes = kernel_block_sizes
        # 计算每个逻辑块包含的内核块数
        self.blocks_per_kv_block = [
            bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
        ]
        # 注意：内核块的总数可能大于逻辑块数
        # 使用 max_num_blocks_per_group * blocks_per_kv_block 来分配
        self.block_tables = []
        for i in range(self.num_kv_cache_groups):
            max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i]
            self.block_tables.append(StagedWriteTensor(
                (self.max_num_reqs, max_num_blocks), ...))
        # ...
​
    def append_block_ids(self, req_index, new_block_ids, overwrite):
        for i in range(self.num_kv_cache_groups):
            start = ...
            block_ids = new_block_ids[i]
            bpk = self.blocks_per_kv_block[i]
            if bpk > 1:
                # 将逻辑 block_id 展开为多个内核 block_id
                # 例如逻辑块 0 对应内核块 0、1、...
                block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
            self.block_tables[i].stage_write(req_index, start, block_ids)

评论区精华

AI 审查指出对 MLA 等具有 storage_block_size != block_size 的模型，虚拟块拆分应被禁用，否则会导致运行时错误。然而最终实现未采纳此建议，仅针对通用情况修复。njhill 提供了简化修改并批准了 PR。

• MLA 模型需要禁用虚拟块拆分 (design): 最终 PR 未采纳该建议，仅处理通用情况。
• kernel_num_blocks 计算对 MLA 模型不安全 (correctness): 未解决，但 PR 已合并，未来可能需要改进。
• prepare_kernel_block_sizes 应检查 storage_block_size (correctness): 未处理，当前实现假设所有模型都兼容。

风险与影响

• 风险：修改了 KV 缓存管道的核心逻辑，可能影响所有使用 MRv2 和连接器的模型。对于 MLA 等特殊存储布局的模型，仍存在潜在不兼容风险（未处理 storage_block_size 差异）。需要额外测试覆盖。
• 影响：对用户：修复了 MRv2 与 FlashInfer 配合使用时的崩溃问题；移除回退意味着连接器必须适配 MRv2 的块表布局，否则可能出错。对开发者：引入内核块大小选择机制，简化了跨后端兼容性维护。影响面中等。
• 风险标记：核心路径变更, 可能影响特殊存储布局模型

关联脉络

• PR #42846 [Bug][CI] NIXL + FlashInfer fails with Qwen3 MRV2 and --block-size 128: 本 PR 直接修复该 Issue 报告的 bug。
• PR #42955 [MRv2] Default to MRv1 when a connector is present: 该 PR 曾引入连接器时回退 MRv1 的临时方案，本 PR 在修复底层问题后将其回退。