执行摘要

• 一句话：修复混合验证器 hidden-states 提取的 block size 错误
• 推荐动作：值得精读，尤其关注根因分析（block_size 不一致）和设计决策：通过 spec 类型而非 layer 名称定位缓存组更健壮；clamp 处理 padding 避免了分支同步。测试覆盖全面，可作为同类 connector 实现的参考。

功能与动机

Issue #613 报告使用 Qwen3.6-35B-A3B-NVFP4 训练后验证指标正常但实际接受率极低，根因是 hidden states 提取为全零。PR body 指出混合验证器中全局 block_size 被提升为公倍数，而 hidden-states 组保持较小 block_size，之前代码错误使用全局值，导致读取零数据和越界。此 PR 还替代了之前的修复尝试 #44328。

实现拆解

1. 新增组定位方法：在 example_hidden_states_connector.py 中添加类方法 _find_cache_kv_group_id，通过检查 HiddenStateCacheSpec 类型而非 layer 名称子串来定位 hidden-states KV 缓存组，并添加错误处理（多个 hidden 组或无 hidden 组时抛出 ValueError）。
2. 新增 block_size 获取方法：添加静态方法 _get_cache_block_size，从定位到的缓存组 spec 中读取 block_size，而非使用 cache_config.block_size；当 kv_cache_config 为 None 时回退到 cache_config.block_size。
3. 修改初始化：在 __init__ 中调用上述方法设置 self._cache_kv_group_id 和 self._block_size，替代原来的 vllm_config.cache_config.block_size 和字符串匹配逻辑。
4. 加固 register_kv_caches：添加断言确保 self._block_size 与缓存的第二维大小匹配，使不匹配在初始化阶段就失败。
5. 处理 padding 插槽：在 extract_hidden_states.py 的 basic_cache 函数中，对 slot_mapping 添加 clamp_min(0)，将 padding 插槽（值为 -1）安全地映射到 null 块（块索引 0），避免分支和同步开销，同时防止写入错误位置。
6. 单元测试：新增 tests/v1/kv_connector/unit/test_hidden_states_connector.py，覆盖 _find_cache_kv_group_id 的各种场景（None 配置、单组非 hidden、多组中定位 hidden、无 hidden 多组报错、多个 hidden 组报错）以及 _get_cache_block_size 的正确回退和 MLA 吸收场景。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py（模块 分布式；类别 source；类型 core-logic；符号 _find_cache_kv_group_id, _get_cache_block_size, init, register_kv_caches）: 核心修复文件：新增 _find_cache_kv_group_id 和 _get_cache_block_size 方法，根据 spec 类型定位 hidden-states 组并正确获取 block_size，替换原来使用全局 cache_config.block_size 和字符串匹配的错误实现。
• vllm/model_executor/models/extract_hidden_states.py（模块 模型执行器；类别 source；类型 data-contract；符号 basic_cache）: 次要修复：在 basic_cache 中对负的 slot_mapping（padding 插槽）使用 clamp_min(0)，防止写入未分配的块 0（null 块），避免之前的主机同步分支检查。
• tests/v1/kv_connector/unit/test_hidden_states_connector.py（模块 连接器；类别 test；类型 test-coverage；符号 _full, _hidden, _config, test_find_group_id_none_config_returns_zero）: 新测试文件，全面覆盖 _find_cache_kv_group_id 和 _get_cache_block_size 的各种场景（None 配置、单组、多组、混合、错误情况），确保新逻辑的正确性和鲁棒性。

关键符号：_find_cache_kv_group_id, _get_cache_block_size, basic_cache, init, register_kv_caches

关键源码片段

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

核心修复文件：新增 _find_cache_kv_group_id 和 _get_cache_block_size 方法，根据 spec 类型定位 hidden-states 组并正确获取 block_size，替换原来使用全局 cache_config.block_size 和字符串匹配的错误实现。

@classmethod
def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
    """
    定位 hidden-states 的 KV 缓存组索引。
    基于 spec 类型（HiddenStateCacheSpec）而非 layer 名称子串，
    使得在 scheduler 和 worker 两侧都能正确解析。
    """
    if kv_cache_config is None:
        return 0
    from vllm.v1.kv_cache_interface import HiddenStateCacheSpec
    groups = kv_cache_config.kv_cache_groups
    # 枚举所有组，收集类型为 HiddenStateCacheSpec 的索引
    group_ids = [
        gid
        for gid, group in enumerate(groups)
        if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
    ]
    if len(group_ids) == 1:
        return group_ids[0]
    if not group_ids and len(groups) == 1:
        # 如果没有 hidden 组但只有一个总组，默认使用组 0
        return 0
    raise ValueError(
        "Could not uniquely identify the extract-hidden-states KV cache "
        f"group among {len(groups)} groups; the hidden-states layer must be "
        "isolated in its own group (MLA verifiers are unsupported)."
    )
​
@staticmethod
def _get_cache_block_size(
    vllm_config: "VllmConfig",
    kv_cache_config: "KVCacheConfig | None",
    cache_kv_group_id: int,
) -> int:
    """
    返回 hidden-states 组的 block_size，从组 spec 中读取。
    cache_config.block_size 在混合验证器中被调整为公倍数，不能使用。
    """
    if kv_cache_config is None:
        return vllm_config.cache_config.block_size
    cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id]
    return cache_group.kv_cache_spec.block_size

def __init__(self, vllm_config, role, kv_cache_config):
    super().__init__(...)
    # 使用新方法替代原来的直接取 cache_config.block_size
    self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config)
    self._block_size = self._get_cache_block_size(
        vllm_config, kv_cache_config, self._cache_kv_group_id
    )
    # 后续逻辑不变 ...

vllm/model_executor/models/extract_hidden_states.py

次要修复：在 basic_cache 中对负的 slot_mapping（padding 插槽）使用 clamp_min(0)，防止写入未分配的块 0（null 块），避免之前的主机同步分支检查。

def basic_cache(
    to_cache: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
):
    # Padding slots 为 -1；将其重定向到 null 块（块 0，从不分配给请求），
    # 使得 scatter 操作保持无分支且无同步。
    block_size = kv_cache.shape[1]
    slot_mapping = slot_mapping.clamp_min(0)
    kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache

tests/v1/kv_connector/unit/test_hidden_states_connector.py

新测试文件，全面覆盖 _find_cache_kv_group_id 和 _get_cache_block_size 的各种场景（None 配置、单组、多组、混合、错误情况），确保新逻辑的正确性和鲁棒性。

def _full(block_size: int) -> FullAttentionSpec:
    """构造一个全注意力 spec 实例，用于测试组配置。"""
    return FullAttentionSpec(block_size=block_size, num_kv_heads=8, head_size=128, dtype=torch.bfloat16)
​
def _hidden(block_size: int) -> HiddenStateCacheSpec:
    """构造一个 hidden-states spec 实例。"""
    return HiddenStateCacheSpec(block_size=block_size, num_kv_heads=6, head_size=2048, dtype=torch.bfloat16)
​
def _config(*specs):
    """根据 spec 列表创建最小配置（仅暴露 kv_cache_groups）。"""
    return SimpleNamespace(
        kv_cache_groups=[
            KVCacheGroupSpec(layer_names=[f"layer.{i}"], kv_cache_spec=spec)
            for i, spec in enumerate(specs)
        ]
    )
​
def test_find_group_id_locates_hidden_group_when_not_first():
    # 混合布局：hidden-states 组不在第 0 组
    cfg = _config(_full(528), _hidden(22), _full(528))
    assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1
​
def test_get_block_size_reads_hidden_group_spec_not_global():
    # hidden 组保持 block_size = 22，全局为 528
    vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528))
    cfg = _config(_full(528), _hidden(22))
    block_size = ExampleHiddenStatesConnector._get_cache_block_size(
        vllm_config, cfg, cache_kv_group_id=1
    )
    assert block_size == 22

评论区精华

Review 评论较少，主要确认：

• fynnsu：运行测试后确认该 PR 修复了 CI 中发现的问题（原始评论：“Yes, looks like this resolved the CI issue.”）。
• mgoin：帮助清理和简化 PR，添加了 basic_cache 的 clamp 逻辑，并获批准后强制合并。

• mergify[bot]：指出 pre-commit 检查失败，建议修复后重新提交。

  • 修复 CI 隐藏状态问题 (testing): 确认 PR 修复了 CI 中观察到的隐藏状态相关问题。

风险与影响

• 风险：

  1. 非混合模型回退：对于非混合模型（仅单组或非 hidden 组），新逻辑通过 _get_cache_block_size 的 kv_cache_config is None 回退到 cache_config.block_size，行为与之前一致，风险低。
  2. 断言影响：register_kv_caches 中的断言可能暴露之前未发现的配置错误，但有助于尽早失败而非静默产生坏数据。
  3. clamp_min(0) 安全性：将 padding -1 映射到块 0，块 0 作为 null 块永不分配给请求，不会影响正常数据；且保持了 scatter 操作无分支，性能不变。
  4. MLA 吸收场景：单元测试覆盖了 HiddenStateCacheSpec 被 MLA 吸收时的报错，清晰告知用户不支持，避免无声错误。
     - 影响：用户影响：修复了混合验证器（如 Qwen3.5/3.6-A3B 系列）上 hidden-states 提取全零的问题，使 speculative decoding 训练（如 Dflash）能正常收敛并获得高接受率。
     系统影响：无性能退化，新增的查找和断言开销在初始化时，不增加推理延迟。
     团队影响：为后续添加类似 connector 提供了清晰的 group 定位模式（基于 spec 类型代替字符串匹配），降低了维护成本。

• 风险标记：核心路径变更, 依赖缓存组规范, 测试覆盖完整

关联脉络

• PR #44328 [Spec Decode] Fix hidden-state extraction block size for hybrid verifiers (previous attempt): 此 PR 在 body 中注明替代 #44328，是同一问题的前一次修复尝试，但范围更窄。