LMCache GPU KV Cache Layout 单一事实来源：`normalize_kv_and_discover

LMCache GPU KV Cache Layout 单一事实来源：`normalize_kv_and_discover_format`不变量深度解析

【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache

本篇技术指南围绕 LMCache 中 GPU KV Cache 布局管理的一条核心设计不变量展开：normalize_kv_and_discover_format是唯一解析 KV-cache 布局的地方。文中将完整剖析这一不变量背后的规范类型、kv_format包结构、格式探测与几何描述（spec）机制、Helper 接口面，以及新增一个 KV 格式的完整操作步骤，帮助你理解 LMCache 如何在同一套代码库中同时支撑 vLLM（flash-attn / flash-infer / MLA / cross-layer）、TRT-LLM、SGLang 等多种引擎的 KV Cache 布局，并避免下游模块各自"猜"形状导致的漂移。

核心不变量：布局解析只有一个入口

normalize_kv_and_discover_format是唯一解析 KV-cache 布局的地方。它在返回格式的同时，也返回 kv_caches 的规范（permute 之后）形式，因此调用方无需再做一次独立的归一化。其他所有模块都通过lmcache/v1/gpu_connector/utils.py中接受EngineKVFormat参数的 Helper 来查询 KV-cache 信息。

这里所说的"布局解析"包括：列表嵌套深度、张量维度顺序、HND 与 NHD、MLA 与 MHA、按层（per-layer）与跨层（cross-layer）。这些信息全部编码在EngineKVFormat中，下游代码绝不能从原始 shape 重新推导。

在源码中，这个不变量由 detection.py 的detect_format落实，而utils.py中的normalize_kv_and_discover_format只是它的一个薄封装：

def normalize_kv_and_discover_format( kv_caches: DiscoverableKVCache, serving_engine: EngineType, layout_hints: "LayoutHints | None" = None, ) -> tuple["lmcache_native.EngineKVFormat", DiscoverableKVCache]: return detect_format(kv_caches, serving_engine, layout_hints)

其完整执行流程为（见 detection.py）：

先做引擎无关的连续视图恢复：attempt_permute_to_contiguous_view(kv_caches)；
依据serving_engine查表得到对应的EngineDetector（查不到则抛出ValueError）；
由 detector 的discover()把引擎的原始布局重塑为规范形式并识别格式，一次返回(engine_kv_format, kv_caches)；
识别失败（返回None）同样抛出ValueError。

规范类型：`DiscoverableKVCache`

所有 KV-cache 值在 LMCache 中都属于以下递归联合类型（定义于 types.py）：

DiscoverableKVCache = Union[torch.Tensor, list["DiscoverableKVCache"]]

实际形态只有三种：

单个torch.Tensor：vLLM cross-layer、TRT-LLM；
扁平的list[torch.Tensor]：vLLM per-layer、SGLang MLA；
嵌套的list[list[torch.Tensor]]：SGLang MHA 的[K_list, V_list]。

引擎适配器如果递过来其他容器（例如 vLLM 的dict[str, Tensor]），有责任在调用任何 Helper 之前解包成该形式。也就是说，类型转换的责任在边界处（adapter），而不是散布在消费方。

与规范类型配套的还有LayoutHints（定义于 types.py），这是服务引擎在注册 KV Cache（REGISTER_KV_CACHE）时传给 LMCache 的提示信息：

Hint 键	含义
`kv_layout`	维度物理顺序：`"NHD"`（大多数 vLLM 构建的默认值）、`"HND"`（`VLLM_KV_CACHE_LAYOUT=HND`）、`"BLHNC"`/`"BLNHC"`（vLLM 标准化布局，block 最外层）
`num_kv_heads`	每层 KV head 数；TRT-LLM 用它把 4-D pool 张量重塑为规范的 6-D 形式
`tokens_per_block`	每个分页块承载的 token 数；TRT-LLM 与 SGLang MHA 都会用到
`kv_list_layout`	SGLang 外层 KV 列表组织方式，`"k_v"`表示扁平注册中包含等长的 K、V 两半
`head_dim`	每 head 维度；TRT-LLM 使用

包结构：`kv_format/`与门面模式

布局逻辑全部位于lmcache/v1/gpu_connector/kv_format/目录，utils.py中的公开 Helper 只是薄门面（facade），委托给它内部实现——对调用方来说，单一事实来源的接口面保持不变。

lmcache/v1/gpu_connector/kv_format/ ├── types.py # DiscoverableKVCache, LayoutHints（基础类型） ├── contiguity.py # attempt_permute_to_contiguous_view（零拷贝视图恢复） ├── detection.py # detect_format() 编排 ├── specs/ # 几何层 │ ├── base.py # KVFormatSpec ABC + shape_desc/concrete_shape 渲染 │ ├── registry.py # 自动发现 spec 文件；get_spec/get_spec_class │ └── <engine_kv_format>.py # 每种格式一个文件 └── detectors/ # 按引擎划分的检测层 ├── base.py # EngineDetector ABC + measure_structure() ├── registry.py # 自动发现 detector 文件；get_detector └── <engine>.py # 每种引擎一个文件（单个 discover()）

仓库中实际已存在的 spec 文件覆盖了全部 15 种格式（如 nb_nl_two_bs_nh_hs.py、nl_x_nb_bs_hs.py、two_x_nl_x_nbbs_nh_hs.py 等），detector 文件则按引擎分为 vllm.py、sglang.py、trtllm.py、atom.py。

每个格式一个 spec 类：几何 + 静态事实

每个EngineKVFormat恰好对应一个KVFormatSpec子类，该子类知道如何对该格式的值进行索引。spec只描述布局——类和文件都以格式成员命名（例如nb_nl_two_bs_nh_hs.py中的NB_NL_TWO_BS_NH_HS_Spec），是几何编码，永远不代表引擎：

get_spec(kv, fmt)返回一个几何实例；
get_spec_class(fmt)返回类，用于读取静态事实（is_mla、is_hnd、is_cross_layer、attention_backends）。

KVFormatSpec是抽象基类（specs/base.py），通过@abstractmethod强制每个 spec 实现num_layers、num_blocks、block_size、num_heads、hidden_dim、head_size、dtype、data_ptrs等访问器——缺失任何一个，第一次get_spec时就会触发TypeError，这正是 golden 测试要兜住的。

静态事实是类属性，声明一次

结构形态（is_cross_layer/is_kv_list/is_layer_list，恰好一个为 True）以及is_mla/is_hnd/is_fused_packed/is_two_major/is_pbs_fused修饰符，都是 spec 上的布尔类属性，默认False，只有适用的才声明。消费方通过get_spec_class(fmt)读取——调用点绝不允许重新列出格式（如fmt in (A, B, ...)），这正是过去同一谓词在torch_ops.py、传输 Helper 和 connector 之间漂移的原因。例如 nl_x_two_nb_bs_nh_hs.py 中就是is_layer_list = True。

设备内核在 csrc/engine_kv_format.h 中维护自己的一份拷贝，test_kv_format_classification.py 将 Python 侧与 C++ 侧钉在一起，并强制结构划分的一致性。

后端标签是诊断性的

一个EngineKVFormat可能由多个（引擎，attention 后端）组合产生，所以每个 spec 在attention_backends（元组）中列出它们，第一项即规范代表。get_attention_backend(fmt)（utils 门面）返回第一项用于日志。这些标签只用于诊断，绝不驱动几何决策，并取代了旧的、手工维护的EngineKVFormat → label字典。

枚举是唯一身份

每个 spec 在类体内声明自己的engine_kv_format，注册表由它派生——没有单独的字符串 id，没有引擎属性。C++ 的EngineKVFormat枚举是"存在哪些格式"的唯一权威。其成员名本身就是布局图例：_连接的 token 序列中X标记列表嵌套边界（TWO_X_NL_X_NBBS_NH_HS→2 x NL x [PBS, NH, HS]）。符号化的shape_desc(fmt)与数值化的concrete_shape(fmt, size)都从该名称渲染而来（specs/base.py），因此它们既不会与枚举漂移，也不会互相漂移。

每个格式一个文件，在一处自动发现

每个 spec 位于自己的specs/<engine_kv_format>.py（以格式命名）。specs/registry.py 通过pkgutil.iter_modules导入文件夹中的每个文件，按声明的engine_kv_format索引进SPECS表。新增格式 = 只需丢一个新文件；发现逻辑就是registry.py中一段可读的循环（没有__init_subclass__，没有散落各处的注册）。

代码库刻意避免继承分类法：格式在 ≥5 个正交轴（引擎、per-/cross-layer、MLA/MHA、NHD/HND、fused/separate PBS）上变化，单一继承脊线无法建模而无孤儿类。只在出现具体需求时才加结构。

检测是唯一引擎感知的层

detect_format先做引擎无关的连续视图恢复，然后按EngineType分发到对应的EngineDetector。每个 detectors/<engine>.py 就是单个discover(kv, hints)，它把引擎的原始布局重塑为规范形式并识别格式，一步返回(format, kv)；detectors/registry.py 用与specs/相同的方式自动发现它们。spec 层永远看不到EngineType。新增引擎 = 只需丢一个新的 detector 文件。

detector 的基类提供measure_list_depth_until_tensor(kv_caches)（detectors/base.py），返回(list_depth, tensor_ndim, first_tensor)——这是 detector 内部做格式分派的关键依据，但它属于kv_format内部实现细节，不对外暴露。

新增一个 KV 格式的完整步骤

按设计文档，新增格式只需以下 4 步，其他 Python 模块都不应需要修改：

添加枚举值：在 csrc/kv_transfer_types.h（所有加速器后端共享的后端无关定义）中添加枚举值，然后在通用 native pybind 模块中注册——csrc/lmcache_native/pybind.cpp 以及 csrc/sycl/pybind_sycl.cpp（SYCL/XPU）。
在引擎的 detector 中添加分支：在detectors/<engine>.py的discover()中，以measure_list_depth_until_tensor得到的(list_depth, tensor_ndim)为键分派，返回(format, kv)；任何基于 hints 的重塑（例如 TRT-LLM 的 4-Dview到 6-D）必须在 shape 检查之前、同一方法内完成。
添加KVFormatSpec子类：新建specs/<engine_kv_format>.py（以格式命名，声明其engine_kv_format与静态事实——结构标志加上适用的修饰符）。registry.py自动发现它，无需改其他文件。ABC 使必需的访问器显式化。
向 golden 表添加一行：在 test_kv_format_specs.py 和 test_kv_format_classification.py 的 golden 表中各加一行，并在 test_kv_format_detection.py 中补一个检测用例。

设计文档还给出了一个明确的警示信号：如果你为了一个新布局而去改kv_layer_groups.py、gpu_context.py或任何KVLayerGroupInfo消费方——这个分支应该放进 spec 里。

Helper 接口面

utils.py中的每个 Helper 都接受DiscoverableKVCache，在布局相关处接受EngineKVFormat。除此之外的任何代码都不得索引原始 shape。

发现（Discovery）

Helper	返回
`normalize_kv_and_discover_format(kv_caches, engine, layout_hints)`	`tuple[EngineKVFormat, DiscoverableKVCache]`—— 唯一的解析器。返回规范（permute 到连续）的 kv_caches 连同检测出的格式；调用方必须使用返回的张量结构进行后续操作。

Format → 引擎映射表

`EngineKVFormat`	引擎	布局	结构
`NB_NL_TWO_BS_NH_HS`	vLLM cross-layer	NHD	裸 6-D 张量`[NB, NL, 2, BS, NH, HS]`
`NB_NL_TWO_NH_BS_HS`	TRT-LLM cross-layer	HND	裸 6-D 张量`[NB, NL, 2, NH, BS, HS]`
`NL_X_TWO_NB_BS_NH_HS`	vLLM flash-attn	NHD	`NL × [2, NB, BS, NH, HS]`
`NL_X_NB_TWO_BS_NH_HS`	vLLM flash-infer	NHD	`NL × [NB, 2, BS, NH, HS]`
`NL_X_TWO_NB_NH_BS_HS`	vLLM flash-attn	HND	`NL × [2, NB, NH, BS, HS]`
`NL_X_NB_TWO_NH_BS_HS`	vLLM flash-infer	HND	`NL × [NB, 2, NH, BS, HS]`
`NL_X_NB_BS_HS`	vLLM MLA	—	`NL × [NB, BS, HS]`
`TWO_X_NL_X_NBBS_NH_HS`	SGLang MHA	NHD	`[K_list, V_list]`，各自`NL × [PBS, NH, HS]`
`TWO_X_NL_X_NB_BS_NH_HS`	SGLang MHA via MP daemon	NHD	`[K_list, V_list]`，各自`NL × [NB, BS, NH, HS]`
`NL_X_NBBS_ONE_HS`	SGLang MLA	—	`NL × [PBS, 1, HS]`
`NL_X_NB_NH_BS_TWO_HS`	vLLM blocks-first fused (CPU)	HND	DEPRECATED（改用`NL_X_NB_NH_BS_CS`）：`NL × [NB, NH, BS, 2, HS]`，从原始`[NB, NH, BS, 2·HS]`拆分
`NL_X_NB_BS_NH_TWO_HS`	vLLM blocks-first fused	NHD	DEPRECATED（改用`NL_X_NB_BS_NH_CS`）：`NL × [NB, BS, NH, 2, HS]`，从原始`[NB, BS, NH, 2·HS]`拆分
`NL_X_NB_NH_BS_CS`	vLLM blocks-first fused (unified KV cache)	HND	`NL × [NB, NH, BS, CS]`，原始注册；CS = content size = 2·HS（K/V 打包）
`NL_X_NB_BS_NH_CS`	vLLM blocks-first fused (unified KV cache)	NHD	`NL × [NB, BS, NH, CS]`，原始注册；CS = content size = 2·HS（K/V 打包）
`NL_X_NB_NH_ONE_BS_HS`	vLLM-RBLN attention	HND	`NL × [2, NB, NH, 1, BS, HS]`；axis 3 恒为 1（RBLN attention 后端的硬性要求）

两个 cross-layer 格式（NB_NL_TWO_*）共享单一 base 指针，内核内部通过shape_desc.nl遍历各层。用get_spec_class(fmt).is_cross_layer做该分派，用.is_hnd检测块内 head-major 布局。

基于 Hints 的重塑（TRT-LLM）

TRT-LLM 交给 LMCache 的是 4-D pool 张量[NB, NL, 2, num_kv_heads * tokens_per_block * head_dim]（HND，K 和 V 在第 2 维交错）。normalize_kv_and_discover_format在连续性检查之前，使用layout_hints["num_kv_heads" | "tokens_per_block" | "head_dim"]将它重塑为规范的 6-D 形式。函数还会把"1 元素列表包一个 6-D 张量"折叠为裸 6-D 张量，使检测落在list_depth == 0。适配器既可以传 4-D 裸张量也可以传[4-D]，函数两者都处理。

标量访问器

以下访问器全部按EngineKVFormat分派。其中可能按层变化的接受可选参数layer_idx: int = 0；传入显式索引即可在异构分组（heterogeneous groups）中做按层查询，无需任何中间 Helper。

Helper	按层？	说明
`get_num_layers(kv, fmt)`	否	总层数。
`get_num_blocks(kv, fmt)`	否	分页块数（组级）。
`get_block_size(kv, fmt)`	否	每块 token 数。
`get_page_buffer_size(kv, fmt)`	否	已废弃（见 utils.py 中的`@lmcache_deprecate`注释：仅旧的非 MP 进程内 connector 使用）。
`get_tokens_per_layer(kv, fmt)`	否	每层 token 容量。
`get_elements_per_layer(kv, fmt)`	否	每层元素数（非 MLA 含 K 和 V）。
`get_num_heads(kv, fmt, layer_idx=0)`	是	每层 KV head 数。
`get_head_size(kv, fmt, layer_idx=0)`	是	每 head 维度。
`get_hidden_dim_size(kv, fmt, layer_idx=0)`	是	隐藏维度（`num_heads * head_size`）。
`get_dtype(kv, fmt, layer_idx=0)`	是	张量 dtype。
`is_mla(fmt)`	—	格式谓词；其余静态事实通过`get_spec_class(fmt)`读取。
`get_device(kv)`	—	与格式无关（下探到任意叶子）。

此外 utils 门面还提供了get_kv_size（K/V 轴大小，2 表示拆分，1 表示 fused）与按层格式探测normalize_and_discover_per_layer_formats（utils.py）——后者为混合模型（如主缓存kv_size=2与 key-only MLA index 缓存kv_size=1并存）逐层报告正确格式。

指针与描述符构造

Helper	返回	说明
`get_group_data_ptrs(kv, fmt, layer_indices)`	`list[int]`	内核期望顺序的指针数组：cross-layer 为`[base]`（忽略`layer_indices`）；SGLang MHA 为`[K_0…K_N, V_0…V_N]`；其余 per-layer 为扁平。与 csrc/cuda/mp_mem_kernels.cu 中的分派一致。指针数组形状是格式的属性——调用方从不问"这个格式有没有 per-layer 指针？"。
`make_page_buffer_shape_desc(kv, fmt, layer_idx, num_layers_in_group, num_blocks, block_size, block_stride_elems)`	`PageBufferShapeDesc`	面向内核的形状结构体。`block_stride_elems`携带每块 dim-0 元素步长；传入`resolve_block_stride_and_log_layout`返回的值，使物理块大小不同的分组（例如压缩的 DeepSeek V4 indexer 组与 dense 层）可以共享同一个 GPU 池。

resolve_block_stride_and_log_layout（utils.py）是KVLayerGroupsManager获取block_stride_elems的唯一入口，同时输出一次性布局审计日志。它的规则是：block 轴格式（_BLOCK_AXIS_FORMATS，当前为NL_X_NB_BS_HS、NL_X_NB_BSV_BSS、NL_X_NB_NH_BS_CS、NL_X_NB_BS_NH_CS）返回stride(0)作为每块步长，大于 tight stride 表示 dim-0 padding（如压缩缓存共享 KV 池）；其他格式返回None并回退 tight stride，若检测到 dim-0 padding 则抛出ValueError（下游内核无法处理）。

连续性

Helper	返回	说明
`attempt_permute_to_contiguous_view(kv)`	`DiscoverableKVCache`	递归、仅元数据操作。已连续则原样返回；无法通过置换恢复（slicing、`as_strided`）则抛`ValueError`。绝不拷贝。遍历整个结构并对每个张量叶子做 permute。由`normalize_kv_and_discover_format`内部调用；对在 discover 流程之外处理张量的调用方（`GPUConnectorInterface.initialize_kvcaches_ptr`、`CudaIPCWrapper.__init__`）仍保持公开。

Consumer 代码中的禁区

原始 shape 索引属于kv_format层（spec 类）内部；consumer 代码必须通过utils.py门面查询，禁止以下任何做法：

用isinstance(kv_cache, (tuple, list))区分布局；
重新列出格式来恢复静态事实（fmt in (A, B)、fmt == NL_X_NB_NH_BS_CS）——改从get_spec_class(fmt)读取，或把缺失的事实补进 spec；
索引原始 shape（tensor.shape[3]、len(shape) == 5）推导维度；
手写列表深度探测（while isinstance(x, list): depth += 1; x = x[0]）——不存在也不应该有公开的 depth Helper，normalize_kv_and_discover_format封装了下探，下游只需要最终得到的EngineKVFormat；
用[tensor]包裹张量来适配 Helper 的列表深度预期——访问器直接接受layer_idx；
手写指针组装（[t.data_ptr() for t in kv_caches]）——用get_group_data_ptrs；
手写设备发现（kv_caches[0][0].device）——用get_device；
手写连续性修复（tensor.contiguous()、.clone()）——用拒绝拷贝的attempt_permute_to_contiguous_view；
编写把kv_caches重写成统一形状再传给 Helper 的"canonicalize"函数——Helper 已经通过接受EngineKVFormat完成规范化，任何确实需要的 reshape/normalize 步骤都住在normalize_kv_and_discover_format内部，调用方从那一次调用就能拿回规范形式。

主要 Consumer

设计文档点名的核心消费方有三个，它们的用法印证了不变量：

kv_layer_groups.py 的KVLayerGroupsManager.__init__：用 5 元组(kv_size, num_heads, head_size, block_size, dtype)对层分组，使用is_mla、get_num_heads、get_head_size、get_block_size、get_dtype（带每层索引）。把block_size纳入分组身份，使压缩组（如物理 slot 更小的 DeepSeek V4 indexer）能与非压缩组在同一个GPUCacheContext下共存。每个组通过make_page_buffer_shape_desc构造PageBufferShapeDesc，传入resolve_block_stride_and_log_layout解析出的block_stride_elems。真实构造函数是唯一入口——没有测试专用捷径、没有缓存的拓扑字段，manager 只暴露kv_layer_groups、num_groups、get_shape_desc。
platform/cuda/cache_context.py 的GPUCacheContext：在 init 时直接构造 manager，把get_shape_desc(group_idx)委托给它，通过get_group_data_ptrs组装每个组的 GPU 指针张量。没有并行的shape_descs_/hidden_dim_sizes_状态。
gpu_connectors.py 的VLLMPagedMemGPUConnectorV3._initialize_kv_cache_pointers：进程内 vLLM 路径调用normalize_kv_and_discover_format（一次完成 HND 支持所需的 permute 与格式检测），并在首次 store/retrieve 时惰性构造metadata.kv_layer_groups_manager。适配器（vllm_v1_adapter.py）不参与格式发现——它只在注册时保存self.kv_caches。

只有normalize_kv_and_discover_format消费layout_hints。内部调用的attempt_permute_to_contiguous_view从 strides 推断置换，不需要 hints。

实现注记：mypy 与递归联合类型

格式分派的原始索引（kv_caches.shape[i]、kv_caches[0][j]）集中在kv_format层：每个specs/<format>.py在文件级设置# mypy: disable-error-code="union-attr,call-overload"（detectors/与contiguity.py只用union-attr，utils.py门面对剩余的结构化 Helper 保留该指令）。engine_kv_format——或者在 spec 中，类身份本身——是索引良定义的证明，但 mypy 无法让这个证明穿过递归 Union，除非逐行加 cast。文件级指令取代了散落的# type: ignore注释；其余所有类型检查保持开启。

总结

LMCache 的 KV Cache 布局设计可以浓缩为一句话：让EngineKVFormat成为唯一身份、让 spec 成为唯一几何来源、让 detector 成为唯一引擎感知层、让utils.py门面成为唯一对外接口。这种分层把"布局解析"这一横切关注点从散落的 consumer 代码中收拢到一个可测试、可扩展的kv_format包内——新增引擎只需加一个 detector 文件，新增格式只需加枚举值、detector 分支、spec 文件与 golden 测试四件事，而下游代码永远只面向EngineKVFormat编程，从根上杜绝了形状猜测带来的漂移。

【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析