执行摘要

• 一句话：将K/V打包到内容维度，统一KV缓存布局为4D
• 推荐动作：建议所有v1用户关注此PR，特别是在使用PD disagg的场景下。设计决策值得学习：通过统一布局后端接口，减少连接器耦合。合并前需确保AMD测试修复通过。

功能与动机

参照RFC #42082，标准化KV缓存布局以消除后端间的差异，减少KV-connector中is_mamba/is_mla等标志的泛滥。当前布局导致代码紧耦合，且不利于异构TP和跨层传输。本PR集中解决K/V打包问题，是整体布局标准化系列的一部分。

实现拆解

按以下步骤实现：

1. 修改注意力后端形状定义：在所有注意力后端（FlashAttention、FlashInfer、Triton、ROCm等）的get_kv_cache_shape方法中，将返回的形状从(num_blocks, 2, block_size, num_kv_heads, head_size)改为(num_blocks, num_kv_heads, block_size, 2*head_size)。量化缓存类似，但使用padded_hs。
2. 更新stride order：对应的get_kv_cache_stride_order方法适配4D形状，在NHD和HND布局下分别返回合适的permutation（例如NHD下为(0,2,1,3)）。
3. 改造K/V提取逻辑：在各个后端的forward/do_kv_cache_update中，用kv_cache.transpose(1,2).split(head_size, dim=-1)替换旧的kv_cache.unbind(1)和PagedAttention.split_kv_cache调用，获得零成本view。
4. 简化NVFP4量化缓存处理：移除nvfp4_kv_cache_split_views函数，改为调用nvfp4_split_data_scale对一侧处理。量化缓存使用separate head groups布局时，直接在dim=1上split。
5. 调整KV-connector传输拓扑：在TransferTopology中判断形状是否为4D blocks-first并断言，删除is_kv_layout_blocks_first、split_k_and_v等属性，简化虚拟split逻辑。NIXL和MoRIIO连接器均适配新的4D布局，K/V不再拆分为独立区域。
6. 更新测试和平台代码：测试文件（test_nixl_connector、test_mooncake_connector、test_minimax_m3等）更新以匹配新形状；ROCm平台代码调整以适应包装布局。

关键文件：

• vllm/utils/torch_utils.py（模块 工具函数；类别 source；类型 core-logic；符号 _nvfp4_split_data_scale, nvfp4_split_data_scale, nvfp4_kv_cache_split_views）: 修改NVFP4 KV缓存的拆分函数，移除旧的分拆视图函数，精简为单侧处理，反映布局变化对量化缓存的影响。
• vllm/distributed/kv_transfer/kv_connector/utils.py（模块 KV连接器；类别 source；类型 core-logic；符号 is_kv_layout_blocks_first, split_k_and_v）: 核心修改：在TransferTopology中识别4D packed形状，简化blocks-first判断，移除split_k_and_v属性，直接使用统一布局。
• vllm/v1/attention/backends/triton_attn.py（模块 注意力后端；类别 source；类型 core-logic；符号 to_f32_units）: 展示了核心的get_kv_cache_shape和get_kv_cache_stride_order变化，反映4D打包布局的实现。
• vllm/v1/attention/backends/flashinfer.py（模块 注意力后端；类别 source；类型 core-logic）: 修改TRTLLM去量化kernel的stride计算以支持打包布局，调整量化缓存处理路径。
• vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py（模块 NIXL连接器；类别 source；类型 core-logic）: NIXL连接器移除旧K/V拆分传输逻辑，适配单一区域传输；调整host buffer分配以匹配4D布局。
• vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py（模块 MoRIIO连接器；类别 source；类型 core-logic；符号 _content_packed_dim）: 添加对4D packed形状的传输几何计算支持。
• tests/v1/kv_connector/unit/test_nixl_connector.py（模块 NIXL测试；类别 test；类型 test-coverage；符号 test_hybrid_mamba_attention_remote_descs_use_packed_head_slices）: 更新测试以匹配新形状和传输区域计数；添加新的测试用例验证打包头部切片。

关键符号：nvfp4_split_data_scale, is_kv_layout_blocks_first, split_k_and_v, _content_packed_dim, TritonAttentionBackend.get_kv_cache_shape, TritonAttentionBackend.get_kv_cache_stride_order, FlashInferBackend.get_kv_cache_shape, TransferTopology.post_init

关键源码片段

vllm/utils/torch_utils.py

修改NVFP4 KV缓存的拆分函数，移除旧的分拆视图函数，精简为单侧处理，反映布局变化对量化缓存的影响。

def nvfp4_split_data_scale(kv_side: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    将NVFP4缓存的一侧（K或V）拆分为数据和缩放因子。
    输入形状为 (B, H, N, full_dim)，其中 full_dim = data_dim + scale_dim。
    布局为 [K_data | K_scale | V_data | V_scale] 每页连续存储。
    调用方需先通过 split 或切片获得单侧。
    """
    num_pages, dim_1, dim_2, full_dim = kv_side.shape
    data_dim = full_dim * 8 // 9
    scale_dim = full_dim - data_dim
    data_per_kv = dim_1 * dim_2 * data_dim
    page_bytes = kv_side.stride(0)
​
    # 从原始 stride 推导 data 和 scale 的 stride，保持物理布局（NHD 或 HND）
    s1 = kv_side.stride(1) * data_dim // full_dim
    s2 = kv_side.stride(2) * data_dim // full_dim
    data_shape = (num_pages, dim_1, dim_2, data_dim)
    data_strides = (page_bytes, s1, s2, 1)
​
    s1_s = kv_side.stride(1) * scale_dim // full_dim
    s2_s = kv_side.stride(2) * scale_dim // full_dim
    scale_shape = (num_pages, dim_1, dim_2, scale_dim)
    scale_strides = (page_bytes, s1_s, s2_s, 1)
​
    base = kv_side.storage_offset()
    data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base)
    scale = torch.as_strided(kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv).view(torch.float8_e4m3fn)
    return data, scale

vllm/distributed/kv_transfer/kv_connector/utils.py

核心修改：在TransferTopology中识别4D packed形状，简化blocks-first判断，移除split_k_and_v属性，直接使用统一布局。

# 在 TransferTopology.__post_init__ 中，识别标准 4D packed 布局
head_size = 1
# Mock 一个块形状来检查 layout
kv_cache_shape = attn_backend.get_kv_cache_shape(
    num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=head_size,
)
logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
​
# 新的注意力缓存形状应为 4D: [num_blocks, num_kv_heads, block_size, content_size]
assert kv_cache_shape[0] == 1, (
    "KV cache layout must be blocks-first; expected mocked "
    f"num_blocks=1 in leading dim, got shape {kv_cache_shape}."
)
if not self.is_mla:
    assert len(kv_cache_shape) == 4, (
        "Attention KV cache layout must be standardized as "
        "[num_blocks, num_kv_heads, block_size, content_size], "
        f"got shape {kv_cache_shape}."
    )
​
# 旧属性 is_kv_layout_blocks_first 和 split_k_and_v 被移除，
# blocks-first 已成为唯一布局，K/V 不再需要拆分传输。

vllm/v1/attention/backends/triton_attn.py

展示了核心的get_kv_cache_shape和get_kv_cache_stride_order变化，反映4D打包布局的实现。

@staticmethod
def get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size, cache_dtype_str="auto"):
    # K 和 V 打包到最后一维：逻辑形状为 (B, H, N, 2*C)
    if block_size % 16 != 0:
        raise ValueError("Block size must be a multiple of 16.")
    if kv_cache_uses_per_token_head_scales(cache_dtype_str):
        # 量化情况下额外 padding 用于存放缩放因子
        from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, get_dtype_size
        cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str]
        scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype)
        if get_kv_quant_mode(cache_dtype_str) == KVQuantMode.INT4_PER_TOKEN_HEAD:
            data_head_size = head_size // 2
        else:
            data_head_size = head_size
        padded_hs = data_head_size + scale_pad
        return (num_blocks, num_kv_heads, block_size, 2 * padded_hs)
    return (num_blocks, num_kv_heads, block_size, 2 * head_size)
​
@staticmethod
def get_kv_cache_stride_order(cache_layout=None):
    """根据cache_layout返回从逻辑形状到物理内存的permutation。
    逻辑形状 (B, H, N, 2*C)，物理排列由stride_order控制。
    NHD：物理为 (B, N, H, 2*C)
    HND：物理为 (B, H, N, 2*C)（恒等）
    """
    if cache_layout == "NHD":
        return (0, 2, 1, 3)
    elif cache_layout == "HND":
        return (0, 1, 2, 3)
    else:
        raise ValueError(f"Unknown cache layout: {cache_layout}")

评论区精华

核心讨论包括：

• AMD测试失败：NickLucche发现三个AMD侧测试失败（MLA backend、MoRIIO connector等），需要修复。
• prefix_prefill stride映射错误：depthfirst-app指出在prefix_prefill.py中，新旧布局下stride索引交换，可能导致KV数据读写出界（高严重性）。LucasWilkinson后续修复。
• 量化缓存兼容性：MatthewBonanni担心NVFP4 + 异构TP NIXL传输，LucasWilkinson回应NVFP4 + PD在主线上当前未测试/不工作，因此暂时break是允许的。
• Intel异构硬件：NickLucche担忧去掉permute逻辑会影响Intel等异构部署，需要确保backend支持HND。
• ROCm_ATTN支持：NickLucche疑问ROCM_ATTN是否支持blocks-first，LucasWilkinson确认在MI300x上测试无误。

• 文档删除：NickLucche对Claude删除NIXL文档表示不满。

• AMD测试失败 (testing): 需要在合并前修复，后续提交似乎已解决。

• prefix_prefill stride映射错误 (correctness): LucasWilkinson在后续提交中修复。
• NVFP4量化缓存兼容性 (question): LucasWilkinson回应NVFP4 + PD在主线上当前未测试/不工作，暂时break可以接受，后续再完善。
• Intel异构硬件支持 (design): 暂无明确结论，需Intel方确认。
• DOCS文档被删除 (question): 属于自动化review工具的副作用。

风险与影响

• 风险：
  1. stride计算错误（高）：在prefix_prefill和scale cache中，深度优先分析发现stride映射错误，可能导致KV数据损坏或越界读取。已在后续提交中修复，但需警惕未覆盖的kernel路径。
  2. 量化缓存回归（中）：NVFP4和FP8量化缓存布局特殊，新方案使用head-group布局，若未彻底测试可能引入精度或性能问题。当前NVFP4+PD未被测试，可能隐藏回归。
  3. 异构硬件兼容性（中）：Intel GPU、AMD等平台需要额外验证新布局，特别是NIXL连接器移除permute逻辑后对非默认布局的影响。
  4. 性能退化（低）：新增的transpose和split操作理论上为零成本view，但某些路径可能因stride不连续导致后续kernel变慢。复审评论中提到perf影响应极小。
     - 影响：影响范围广：所有使用v1 attention backends的模型都会受到影响，因为KV缓存形状改变。但经过E2E lm-eval PD disagg测试验证，覆盖了Qwen、DeepSeek、Nemotron等模型在多种backend和拓扑配置下，结果通过。用户升级到该PR后需确保模型兼容。团队需注意与外部PR（如Intel GPU、AMD ROCm）的冲突。测试配套更新完整，包括connector单元测试。
     - 风险标记：核心路径变更, 多个backend同时修改, 量化缓存路径未充分测试, AMD测试失败待修复, prefix_prefill stride隐患, 异构硬件兼容性待确认

关联脉络

• PR #42082 [RFC]: Standardize KV-cache Layouts: 本PR是RFC #42082中标准化KV缓存布局系列的第2部分，RFC提出了全局布局规范。
• PR #42374 first part of RFC #42082: 本PR基于#42374拆分为4个PR，#44455是其中的[2/N]，负责打包K/V到内容维度。
• PR #44454 [1/N][KV-Cache Layout Refactor] Refactor DSV4 KV cache config: 本PR依赖#44454作为基础，必须先合并。
• PR #44456 [3/N][KV-Cache Layout Refactor] Standardize Mamba cache; drop get_transfer_cache_regions: 与本PR同系列，后续步骤。
• PR #44458 [4/N][KV-Cache Layout Refactor] Standardize KV cache layout: 系列的最后部分。
• PR #41657 related bug PR: 本PR动机之一提到该PR导致的bug。