执行摘要

• 一句话：CPU 注意力后端 KV 缓存布局标准化为 blocks-first
• 推荐动作：该 PR 是 CPU 后端标准化的一次小步重构，值得关注其设计思路：通过统一布局降低跨后端复杂度。建议阅读核心变更的代码片段以理解形状转换技巧。

功能与动机

PR #42082 提出了标准化 KV 缓存布局的需求。原始 CPU 后端使用独立的 key 和 value 缓存（unbind(0)），而其它后端采用 blocks-first 布局（key 和 value 在最后一维拼接）。此变更旨在统一布局，方便后续跨后端代码复用和 C++ 算子中的连续性假设。

实现拆解

1. 修改 get_kv_cache_shape() 方法（vllm/v1/attention/backends/cpu_attn.py 第 88-96 行）：返回值从 (2, num_blocks, num_kv_heads, block_size, head_size) 改为 (num_blocks, num_kv_heads, block_size, 2 * head_size)。
2. 更新 forward() 方法中的 KV 缓存解包逻辑（同上文件第 340-346 行）：替换原来的 kv_cache.unbind(0) 为先 view 再 chunk(2, dim=2)，将最后一维拆分为 key 和 value。
3. 同步测试用例中的 KV 缓存张量创建（tests/kernels/attention/test_cpu_attn.py 第 261-267 行）：将原先分别创建 packed_key_cache 和 packed_value_cache 改为创建一个联合张量 packed_key_value_cache，再通过 view 和 chunk 切分。

关键文件：

• vllm/v1/attention/backends/cpu_attn.py（模块 注意力后端；类别 source；类型 core-logic；符号 CPUAttentionBackend.get_kv_cache_shape, CPUAttentionBackendImpl.forward）: 核心变更文件：修改 KV 缓存形状定义和解包逻辑。
• tests/kernels/attention/test_cpu_attn.py（模块 测试；类别 test；类型 test-coverage；符号 varlen_with_paged_kv）: 测试文件同步更新，确保新布局下的功能正确。

关键符号：CPUAttentionBackend.get_kv_cache_shape, CPUAttentionBackendImpl.forward, varlen_with_paged_kv

关键源码片段

vllm/v1/attention/backends/cpu_attn.py

核心变更文件：修改 KV 缓存形状定义和解包逻辑。

# vllm/v1/attention/backends/cpu_attn.py
​
@staticmethod
def get_kv_cache_shape(
    num_blocks: int,
    block_size: int,
    num_kv_heads: int,
    head_size: int,
    cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
    # 变更后：形状变为 [num_blocks, num_kv_heads, block_size, 2 * head_size]
    # 原先是 [2, num_blocks, num_kv_heads, block_size, head_size]
    return num_blocks, num_kv_heads, block_size, 2 * head_size
​
# ... 在 forward 方法中：
​
# For decoder and cross-attention, use KV cache, size are
# [num_blocks, num_kv_heads, block_size, 2 * head_size]
# Make a view [num_blocks, num_kv_heads, block_size * 2, head_size]
# Then slice KV at dim 2
​
num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
# 先 reshape 为 [num_blocks, num_kv_heads, block_size*2, head_size]
kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
# 然后在第 2 维（block_size*2）上切分成 key 和 value 两个 chunks
key_cache, value_cache = kv_cache.chunk(2, dim=2)

tests/kernels/attention/test_cpu_attn.py

测试文件同步更新，确保新布局下的功能正确。

# tests/kernels/attention/test_cpu_attn.py
​
# KV cache for CPU attention
cache_dtype = torch.uint8 if is_fp8 else dtype
# 变更后：使用一个联合张量，形状为 [num_blocks, num_kv_heads, block_size, head_size * 2]
packed_key_value_cache = torch.empty(
    num_blocks, num_kv_heads, block_size, head_size * 2, dtype=cache_dtype
)
# 然后将最后一维拆分为 block_size*2 和 head_size，以 chunk 得到 key 和 value
packed_key_value_cache = packed_key_value_cache.view(
    (num_blocks, num_kv_heads, block_size * 2, -1)
)
packed_key_cache, packed_value_cache = packed_key_value_cache.chunk(2, dim=2)

评论区精华

Reviewer yamt 在代码行 343 处提问："is this to match the contiguity assumptions in C++ ops?"（这是为了匹配 C++ 算子的连续性假设吗？）该评论未得到作者回复，但 PR 随后被合并。这暗示变更确实出于与底层 C++ 算子的兼容性考虑。

• C++ 算子连续性假设 (design): 未直接回复，但 PR 被合并，推测确为此目的。

风险与影响

• 风险：风险较低。变更仅涉及 CPU 注意力后端，且形状变化被限制在 get_kv_cache_shape 和 forward 方法内，外部接口不变。测试已同步更新，覆盖了 FP8 和普通精度路径。需要注意的是，任何在 CPU 后端外部直接操作 KV 缓存形状的代码都需要适配，但当前仓库中外部依赖已被隔离。
• 影响：影响范围：仅限 CPU 注意力后端。影响程度：中等。任何使用 CPU_ATTN 后端的模型（如 CPU 上的 transformer 模型）都将受益于更一致的布局，便于后续优化。向后兼容性：不完全兼容—旧版序列化的 KV 缓存无法直接读取，但 KV 缓存通常不在版本间持久化，因此实际影响有限。
• 风险标记：布局变更可能影响外部直接操作 KV 缓存的代码

关联脉络

• PR #42082 注意后端 KV 缓存布局标准化: 该 Issue 是此 PR 的动机来源，要求标准化 KV 缓存布局。