执行摘要

• 一句话：修复 DCP 下统一 KV 缓存规格的 block table 宽度不一致
• 推荐动作：值得精读。该 PR 改动小，但揭示了 DCP 下 KV cache spec 机制的设计缺陷，修复思路（委托给合并 specs 并断言一致性）具有借鉴意义，且测试覆盖了多种场景。

功能与动机

Issue #50825 报告 GLM-5.2 在 TP8/DCP2/fp8_ds_mla 下解码崩溃，报错 RuntimeError: The expanded size of the tensor (8192) must match the existing size (16384)，发生在 mla/indexer.py。根因是 UniformTypeKVCacheSpecs 未覆盖 max_num_blocks_per_req，导致 runner 从 group spec 分配的 block table 宽度与 metadata builders 从 per-layer spec 构造的宽度不一致。PR 描述指出该问题在 #50302 之前就已存在，且不是其对齐部分的重复。

实现拆解

1. 在 vllm/v1/kv_cache_interface.py 中为 UniformTypeKVCacheSpecs 新增 max_num_blocks_per_req 方法。该方法遍历 kv_cache_specs 中的每个 spec，收集各自的 max_num_blocks_per_req 值，并通过 assert len(widths) == 1 确保所有 layer 宽度一致，然后返回该唯一值。这样 group spec 与 per-layer spec 始终报告同一宽度。
2. 在 tests/v1/core/test_kv_cache_utils.py 中新增参数化测试 test_uniform_type_spec_block_table_width_matches_layer_spec，覆盖 mla 和 mamba 两种 layer 类型，以及 dcp_size = 1 和 2 的情形，验证 layer_spec 与 uniform_spec 的 max_num_blocks_per_req 一致。测试还特别注明 mamba 状态是复制而非 DCP 分片。
3. 提交历史中，MatthewBonanni 补充了宽度一致性断言，并将测试参数化，进一步增强覆盖。

关键文件：

• vllm/v1/kv_cache_interface.py（模块 缓存规格；类别 source；类型 core-logic；符号 max_num_blocks_per_req）: 核心修复：为 UniformTypeKVCacheSpecs 新增 max_num_blocks_per_req 覆盖，委托给内部各 layer spec，确保 DCP 下 block table 宽度一致。
• tests/v1/core/test_kv_cache_utils.py（模块 缓存测试；类别 test；类型 test-coverage；符号 test_uniform_type_spec_block_table_width_matches_layer_spec）: 新增参数化测试，验证 UniformTypeKVCacheSpecs 与内部 layer spec 在 DCP 分片下返回相同的 block table 宽度。

关键符号：max_num_blocks_per_req, test_uniform_type_spec_block_table_width_matches_layer_spec

关键源码片段

vllm/v1/kv_cache_interface.py

核心修复：为 UniformTypeKVCacheSpecs 新增 max_num_blocks_per_req 覆盖，委托给内部各 layer spec，确保 DCP 下 block table 宽度一致。

在 UniformTypeKVCacheSpecs 中新增 max_num_blocks_per_req 方法，将宽度计算委托给内部各 layer spec，并断言一致性：

@dataclass(frozen=True)
class UniformTypeKVCacheSpecs(KVCacheSpec):
    # 多个相同 type 的 layer 共享的 KV cache 规格组（例如 DSA 模型中的所有 MLA layer）。
    kv_cache_specs: dict[str, KVCacheSpec]
​
    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
        # Metadata builders 从 per-layer spec 构造，而 runner 从 group spec 分配
        # block table，所以 group spec 必须与内部 layer spec 返回相同的宽度，
        # 否则 DCP 下两者相差 dcp_world_size 倍，导致 sparse indexer 崩溃。
        widths = {
            spec.max_num_blocks_per_req(vllm_config, max_len)
            for spec in self.kv_cache_specs.values()
        }
        # 所有 layer 的宽度必须一致，否则说明分组不合理，提前暴露配置错误。
        assert len(widths) == 1, (
            'All layers in the same KV cache group must need the same number '
            f'of block table entries, got {sorted(widths)}.'
        )
        return next(iter(widths))

tests/v1/core/test_kv_cache_utils.py

新增参数化测试，验证 UniformTypeKVCacheSpecs 与内部 layer spec 在 DCP 分片下返回相同的 block table 宽度。

新增参数化测试，验证 group spec 与 per-layer spec 在 DCP 分片下宽度一致：

@pytest.mark.parametrize(
    'layer_type,dcp_size,expected_width',
    [
        ('mla', 1, 64),
        ('mla', 2, 32),
        # Mamba 状态是复制的，不进行 DCP 分片，其宽度是常驻状态块数而非 cdiv。
        ('mamba', 2, 3),
    ],
)
def test_uniform_type_spec_block_table_width_matches_layer_spec(
    layer_type, dcp_size, expected_width
):
    # runner 从 group spec 分配 block table，metadata builders 则从 per-layer spec
    # 构造，因此聚合规格必须与其包装的 layer 返回相同的宽度。
    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
    vllm_config.parallel_config.decode_context_parallel_size = dcp_size
    layer_spec = new_mla_spec() if layer_type == 'mla' else new_mamba_spec()
    uniform_spec = UniformTypeKVCacheSpecs(
        block_size=layer_spec.block_size,
        kv_cache_specs={'layer1': layer_spec, 'layer2': layer_spec},
    )
​
    assert layer_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width
    assert uniform_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width

评论区精华

MatthewBonanni 在 review 中批准并做了小幅调整：LGTM, thanks! Just made a couple tweaks，补充了断言和参数化测试。PR 描述中作者详细澄清了与 #48404、#46794、#43970 的区别，强调 DCP 分片是独立根因，避免重复修复。claude[bot] 因 fork 自动 review 被禁用。

• 修复方案与测试覆盖 (correctness): 已解决：合并时加入了断言和参数化测试。
• 与既有修复 PR 的关系 (question): 澄清了非重复，确认独立根因。

风险与影响

• 风险：修改了 block table 宽度的计算，可能影响其他使用 UniformTypeKVCacheSpecs 的场景（非 DCP 下宽度不变）。如果内部 layer 的宽度计算不一致，断言会失败并中断启动，但这能提前暴露配置错误。测试仅在 CPU 环境运行，未覆盖真实 GPU DCP 集群，但 PR 描述提及在 4×H200 上验证过 fail→pass。
• 影响：用户影响：修复了 DCP 下 DSA 模型（如 GLM-5.2）的解码崩溃，使这些配置可用。系统影响：避免 block table 分配 dcp_world_size 倍宽，节省显存。团队影响：为后续类似 block table 宽度不一致问题提供了统一的处理模式。
• 风险标记：核心路径变更, 依赖 DCP 配置, 测试覆盖有限

关联脉络

• PR #50302 [Bugfix] Universally align block table width to 128 tokens: 统一 block table 宽度计算，本 PR 修复其未覆盖的 DCP 分片场景。
• PR #48404 [Bugfix] Size the sparse-indexer expanded block table from the runner's block-table width: 之前通过重分配 buffer 治疗症状，本 PR 从根因修复 DCP 分片导致的宽度不一致。
• PR #43970 [Bugfix] Align MLA indexer block table with MTP speculative decode: 解决 MTP 对齐问题，但与本 PR 的 DCP 分片问题不同源。