执行摘要

• 一句话：禁用 per-token-head 量化的跨层 KV 缓存布局
• 推荐动作：值得精读。本 PR 展示了如何通过一个门控修复跨组件不兼容，并体现了 Review 过程中从局部修复到统一修复的演进。对于维护者，守卫和注释是重要上下文，后续若要重新启用跨层优化需清除此守卫。

功能与动机

引用 PR body 和 Issue #48412：Per-token-head 量化将 fp32 scale 内嵌在 KV cell 尾部，注意力后端在 _ensure_scale_caches 中基于逐层连续缓冲假设创建 scale 视图。跨层统一分配共享一个交错缓冲，导致 scale 写入覆盖相邻 KV 数据，从第一次 decode 即损坏缓存。必须禁用跨层布局以恢复正确性。具体描述见 Issue #48412。

实现拆解

变更入口：vllm/v1/worker/kv_connector_model_runner_mixin.py 中的 use_uniform_kv_cache 方法。

1. 在原有检查是否支持块 stride 索引之前，新增对 kv_cache_spec.kv_quant_mode.is_per_token_head 的判断。若为 True 则直接返回 False，阻止跨层分配。
2. 该守卫置于统一分配决策点，覆盖所有 KV 连接器（OffloadingConnector、NixlConnector、MooncakeStore），而非限于单个连接器。
3. 设计依据：Per-token-head 量化需要逐层连续缓冲以正确定位 scale 内嵌区域；跨层分配破坏该假设，导致数据损坏。回退到逐层路径（由 #48411 修复传输正确性）是安全的折中。
4. 曾包含测试文件但应 reviewer 要求移除；无配置或部署变更。

关键文件：

• vllm/v1/worker/kv_connector_model_runner_mixin.py（模块 KV 缓存分配；类别 source；类型 data-contract；符号 use_uniform_kv_cache）: 核心修复文件，在 use_uniform_kv_cache 方法中添加 per-token-head 量化检查，禁用跨层 KV 缓存分配。所有 KV 连接器共享此逻辑。

关键符号：use_uniform_kv_cache

关键源码片段

vllm/v1/worker/kv_connector_model_runner_mixin.py

核心修复文件，在 use_uniform_kv_cache 方法中添加 per-token-head 量化检查，禁用跨层 KV 缓存分配。所有 KV 连接器共享此逻辑。

@staticmethod
def use_uniform_kv_cache(
    attn_groups: list[list[AttentionGroup]],
) -> bool:
    if not has_kv_transfer_group():
        return False
    if not get_kv_transfer_group().prefer_cross_layer_blocks:
        return False
​
    if len(attn_groups) != 1 or len(attn_groups[0]) != 1:
        return False
​
    attn_group = attn_groups[0][0]
    kv_cache_spec = attn_group.kv_cache_spec
    if not isinstance(kv_cache_spec, AttentionSpec):
        return False
    # Per-token-head quant carves inline-scale views that assume per-layer
    # contiguous KV buffers; the cross-layer layout breaks this and corrupts KV.
    if kv_cache_spec.kv_quant_mode.is_per_token_head:
        return False
    return kv_cache_spec.indexes_kv_by_block_stride

评论区精华

Review 中核心讨论围绕三个问题：

• 设计范围：Etelis 指出仅修改 OffloadingConnector 不够，因为 NixlConnector 和 MooncakeStore 同样有 prefer_cross_layer_blocks=True，应将检查提升到 use_uniform_kv_cache()。Achyuthan-S 采纳此建议，移除了连接器特定覆盖。
• 注释精简：Etelis 建议将冗长注释剪至两行，完整机制已在 PR 描述中。Achyuthan-S 照做。

• 测试取舍：orozery 认为不需要测试，要求移除。Achyuthan-S 移除测试文件，认为端到端已在 H100 上由 Etelis 验证。

• 注释精简 (style): Achyuthan-S 采纳，后续将注释移至 use_uniform_kv_cache 并保持精简。

• Per-token-head 检查提升到统一路径 (design): Achyuthan-S 采纳，移除了 OffloadingConnector 中的覆盖，在 mixin 中添加检查，覆盖所有连接器。
• 移除测试文件 (testing): Achyuthan-S 移除测试，PR 只保留门控变更。端到端在 H100 上验证。

风险与影响

• 风险：风险较低。主要考虑：
  1. 如果未来有人实现跨层布局对 per-token-head 的支持，此守卫会阻止跨层优化，需适时移除。
  2. 回退的逐层路径依赖 #48411 的修复；若 #48411 被回滚或重新引入传输宽度缺陷，此修复将失效。
  3. 变更仅影响 per-token-head dtype 子集，其他量化模式不受影响。
  4. 无测试覆盖，但 reviewer 进行了手动验证。
     - 影响：影响范围限于同时启用 KV offloading 和 per-token-head 量化（fp8_per_token_head / int8_per_token_head / int4_per_token_head）的用户。之前生成结果完全损坏，现在正确。非 per-token-head 用户无变化。无性能退化，无 API 改变。对其他 KV 连接器（NixlConnector、MooncakeStore）同样修复了潜在的相同问题。
     - 风险标记：跨层布局兼容性, 测试覆盖不足, 依赖上游修复

关联脉络

• PR #48411 [Bugfix][KVConnector] Fix per-token-head offload restore width: 先前修复了逐层路径的传输宽度；本 PR 依赖 #48411 确保逐层路径正确后，才安全回退到该路径。