LMCache Hybrid KV Cache Groups：面向多进程 vLLM 连接器的混合 KV 缓存分组设计

LMCache Hybrid KV Cache Groups：面向多进程 vLLM 连接器的混合 KV 缓存分组设计

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

导读

本文深入解析 LMCache 为多进程（multiprocess）vLLM 连接器设计的混合内存分配器（HMA，Hybrid Memory Allocator）方案——docs/design/integration/vllm/hybrid-kv-cache-groups.md。当 vLLM 以多个异构 KV cache group（如滑动窗口组与全注意力组并存、MLA/Indexer 压缩缓存、Mamba 循环状态）运行时，LMCache 必须按物理布局（kv_size / num_heads / head_size / block_size / dtype）执行传输，同时严格隔离各引擎组的块 ID 地址空间。读完本文，你将掌握EngineGroupInfo与KVLayerGroupInfo两个核心抽象的分工、注册/存储/检索的完整数据流、跨层 KV 共享与槽位压缩的处理策略，以及如何用仓库源码（group_view.py、kv_layer_groups.py、kv_cache_groups.py）验证这些设计。

背景：为什么需要"混合"KV 缓存分组

引擎视角与传输视角的错位

vLLM 的混合 KV 缓存特性（hybrid KV cache）允许不同层使用不同的缓存行为——例如google/gemma-4-E4B-it中滑动窗口组与全注意力组并存，DeepSeek-V4-Flash 则混合了 256/64/8/4 四种tokens_per_block。vLLM 通过KVCacheConfig.kv_cache_groups对层进行分组，每个引擎组是一个独立的分页块地址空间，块 ID 仅在组内有意义。

而 LMCache 的 GPU 传输内核需要按物理传输身份（tensor 的秩、隐藏维度、块大小、dtype）来分组，才能让一组层共享一次内核启动。两者并不对齐：

• vLLM 按"缓存行为"分组；
• LMCache 必须按"物理布局 + 块 ID 空间"分组；
• 同一个引擎组内，各层可能因隐藏维度不同而需要不同的拷贝内核（如 5-D 的 key+value 缓存[NB, 2, BS, NH, HS]与 3-D 的 key-only 缓存[NB, BS, HS]同处一个UniformTypeKVCacheSpecs组）。

因此文档将三个概念严格分离（hybrid-kv-cache-groups.md的 Summary 部分）：

目标与非目标

• 保持 ZMQ API 引擎中立：所有 vLLM 字段读取被限定在lmcache.integration.vllm包内；
• 注册即定义协议可见的组顺序：store/retrieve 的块 ID 按该顺序索引；
• 复用同一个分组原语：group_layers_by_identity在 vLLM 侧与服务器侧共用，保证组顺序一致；
• 明确不在范围内：滑动窗口 load-plan 裁剪、非 GPU 传输路径上的 HMA（该路径直接拒绝多组传输）、移除layout_hints（仍用于 tensor 布局检测）。

两个核心类型：EngineGroupInfo与KVLayerGroupInfo

EngineGroupInfo：引擎中立、可跨进程编码

EngineGroupInfo定义在 group_view.py，是一个frozen=True的msgspec.Struct，字段如下：

关键语义（源码 docstring 与测试共同印证）：

• 多个 info 可共享同一engine_group_id：一个引擎组按物理传输身份被拆成多个 LMCache 组时；
• 列表顺序即协议可见组顺序：REGISTER_KV_CACHE负载中逐字携带list[EngineGroupInfo]，由消息队列完成编码/解码；
• 空列表 = 单一非混合组：引擎不上报组元数据时的默认行为。

group_view.py 提供了若干Sequence[EngineGroupInfo]上的辅助函数：

• num_engine_groups：返回引擎组数（max(engine_group_id) + 1，空列表返回 1）；
• num_engine_group_infos：返回 LMCache 组数（len(groups)，空列表返回 1）；
• expand_engine_block_ids：将引擎侧按引擎组索引的块 ID 展开为按 LMCache 组索引的list[list[int]]；
• get_engine_group_indices：为每个已注册 tensor 映射引擎组索引，未归属任何组的 tensor 标记为EXCLUDED_ENGINE_GROUP；
• slice_block_ids_per_group：按 token 范围切分每个引擎组的块 ID（见下文"以 token 为统一记账单位"）。

KVLayerGroupInfo：服务器端运行时内核分发单元

在 kv_layer_groups.py 中，KernelGroupInfo（其别名即KVLayerGroupInfo）是运行时的传输内核分发单元，包含：

• layer_indices：该组层索引（内核迭代顺序）；
• shape_desc：PageBufferShapeDesc，一次构造时打上全部八个字段（kv_size, nl, nb, bs, nh, hs, element_size, block_stride_elems）；
• dtype：torch dtype，用于内核模板实例化（element_size只是字节宽度，无法区分同为 2 字节的 bfloat16 与 float16）；
• engine_kv_format：每组的 Engine KV format（混合格式模型按组分发）；
• tokens_per_block/slots_per_block：逻辑 token 数 / 物理槽位数；
• engine_group_idx、sw_size_tokens、extra_object_group_tag、recurrent_state。

注意：KVLayerGroupInfo由真实 tensor 派生（kv_layer_groups.py的模块注释），绝不是 API 契约的一部分——协议可见的只有EngineGroupInfo列表。

数据流总览

文档给出了完整的注册 → 存储/检索调用链：

vLLM KVCacheConfig + registered kv_caches | integration.vllm.kv_cache_groups.create_engine_group_infos_from_vllm v list[EngineGroupInfo] --REGISTER_KV_CACHE (msgspec)--> server msgspec-decode | KVLayerGroupsManager validates against real tensors v KVLayerGroupInfo list --STORE/RETRIEVE block_ids per info--> transfer kernels

实际调用链在 lmcache_mp_connector.py 的register_kv_caches中可完整验证：

1. 读取layout_hints，先调用apply_kv_cache_group_edits对 KV tensor 做"组编辑"（见 Mamba/线性注意力章节）；
2. 调用create_engine_group_infos_from_vllm(kv_cache_config, kv_caches, layout_hints=..., dcp_size=...)构建engine_group_infos；
3. 将其与编辑后的kv_caches一起交给worker_adapter.register_kv_caches。

在 ZMQ 协议层，engine.py 中REGISTER_KV_CACHE的 payload 为(instance_id, KVCache, model_name, world_size, EngineType, LayoutHints, list[EngineGroupInfo])，注释明确说明engine_group_infos是"由消息队列 msgspec 编码的引擎中立 KV cache 组元数据"。STORE/RETRIEVE的 payload 则是(key, instance_id, gpu_block_ids, event_ipc_handle)，其中gpu_block_ids: list[list[int]]按 LMCache KV 组索引。

注册流程：create_engine_group_infos_from_vllm

这是唯一读取 vLLMKVCacheConfig的函数（kv_cache_groups.py），其实现与文档的四个步骤一一对应：

步骤 1：按布局（而非引擎组）检测每层的 Engine KV format。调用normalize_and_discover_per_layer_formats。检测粒度是布局：同一个引擎组可以包含注册 tensor 形状不同的层——例如 5-D 的 key+value 缓存（[NB, 2, BS, NH, HS]，kv_size=2）与 3-D 的 key-only 缓存（[NB, BS, HS]，kv_size=1）同处一个UniformTypeKVCacheSpecs组，两种布局分别被检测并上报。（"5-D"/"3-D" 指一层注册 KV tensor 的维度数。）

步骤 2：把每个已注册层映射到引擎组索引。以layer_to_idx = {name: idx ...}建立层名到注册索引的映射。所有引擎组layer_names都未提及的层（跨层 KV 共享层）被标记为EXCLUDED_ENGINE_GROUP并丢弃。

步骤 3：group_layers_by_identity按传输身份拆分。分组原语定义在 kv_layer_groups.py：

class KernelGroupIdentity(NamedTuple): kv_size: int num_heads: int head_size: int block_size: int engine_group_idx: int dtype: torch.dtype engine_kv_format: "lmcache_native.EngineKVFormat"

• engine_group_idx使不同引擎组中形状完全相同的层也不会合并（块 ID 只在组内有效）；
• engine_kv_format使共享一个引擎组的不同布局（如步骤 1 的 5-D 与 3-D）保持分离；
• 函数按每个组的首层索引排序输出，保证组顺序跨运行确定性（kv_layer_groups.py的返回注释）。

步骤 4：每个身份产出一个EngineGroupInfo。同时合并该组的滑动窗口大小（_merge_layer_sw_sizes）与循环状态标记（_merge_layer_recurrent），并把tokens_per_block从group_tokens_per_block中取出（注意get_tokens_per_block在 DCP 解码上下文并行时会按block_size * dcp_size放大注意力组的逻辑块覆盖 token 数，而 Mamba 循环状态是复制的、不放大，见 kv_cache_groups.py）。

引擎组 vs LMCache 组的数量

group_view.py 中的两个函数区分了两种计数：num_engine_groups（引擎块 ID 列表数，即每次传输请求携带的块 ID 列表数）与num_engine_group_infos（协议可见的 LMCache 组数）。两者可以不同——一个引擎组按传输身份拆成多个 info 时，num_engine_group_infos > num_engine_groups。

Store / Retrieve：块 ID 的展开与按组索引

引擎侧重索引：expand_engine_block_ids

vLLM 按引擎组上报块 ID。worker 适配器（vllm_multi_process_adapter.py）调用：

return expand_engine_block_ids(self.engine_group_infos, op.block_ids)

每个 info复用其源引擎组的块 ID（group_view.py的 docstring：element i is the engine group id that LMCache group i reads block IDs from），使STORE/RETRIEVE收到按 info 顺序索引的list[list[int]]。服务器循环因此退化为平凡逻辑：对 infoi，直接用gpu_block_ids[i]。同时该函数对旧版（非混合）连接器发出的扁平Sequence[int]做了向后兼容归一化。

测试用例 精确复现了文档示例：三个 info（组 0 拆成两个）下，块 ID{group 0: [10,11], group 1: [20,21]}被展开为[[10,11], [10,11], [20,21]]。

以 token 为统一记账单位：slice_block_ids_per_group

调度器侧连接器在token（每个组共享的唯一单位）中做所有记账（命中计数、store/retrieve 范围），并按token_range / tokens_per_block_g切分每组的块 ID（slice_block_ids_per_group）。文档示例：同样 256 个 token，tokens_per_block=64的组得到 4 个块 ID，而tokens_per_block=256的组只得到 1 个。若 token 范围未对齐某个组的块边界，该函数直接抛出ValueError（group_view.py）。

服务器侧则计算blocks_per_chunk = lmcache_tokens_per_chunk // tokens_per_block（每组），见 lmcache_driven_transfer.py 附近对blocks_per_chunk的逐组计算，以及KVLayerGroupsManager.calculate_num_blocks用num_tokens * slots_per_block // tokens_per_block推导物理槽位（kv_layer_groups.py）。

每组的块大小与槽位压缩

不存在单一的"引擎块大小"

文档强调：每个组有两个按组计量的量，其余全部由它们推导：

• tokens_per_block：一个分页块（一个块 ID）覆盖的逻辑 token 数。来自引擎 KV cache spec 的block_size，初始化时读取并随EngineGroupInfo携带。混合模型可以自由混用：google/gemma-4-E4B-it的滑动窗口组为 32、全注意力组为 16；DeepSeek-V4-Flash 为 256/64/8/4。
• slots_per_block：一个分页块的物理槽位数，注册时从真实 tensor 检测（batch-size 维度，shape_desc.bs），仅按内核组可用。

压缩：不存储compress_ratio

当tokens_per_block > slots_per_block时该组被视为"压缩"（每个物理槽打包tokens_per_block // slots_per_block个逻辑 token）：普通注意力每槽一个 token，而 DeepSeek-V4-Flash 的 MLA / indexer 缓存分别打包 4 和 128 个。仓库不存储任何compress_ratio——需要比率的地方都从这两个地面真值量内联计算（见 kv_layer_groups.py 的tokens_per_block回退逻辑，以及calculate_slots/calculate_num_blocks中基于几何的推导）。

LMCache chunk 大小必须是每个组tokens_per_block的整数倍，该约束在连接器初始化和注册时双重校验（_validate_block_chunk_size_config，kv_layer_groups.py）。校验规则：

• tokens_per_block % slots_per_block == 0；
• lmcache_tokens_per_chunk % tokens_per_block == 0；
• 子块滑动窗口（小于 chunk 大小）也必须是tokens_per_block的整数倍。

跨层 KV 共享与失败语义

EXCLUDED_ENGINE_GROUP：共享层不建组

某些模型（如google/gemma-4-E4B-it）让一层复用另一层的 KV 缓存（kv_caches[layer] = kv_caches[target]）。vLLM 只在kv_cache_groups中列出缓存所有者层；共享层不在任何组的layer_names中。由于该层的 KV 物理上存放在目标所有者的块里，存储/检索所有者即已覆盖它。

因此注册流程把这些未列出的层标记为EXCLUDED_ENGINE_GROUP（kv_layer_groups.py中= -1的哨兵值，kv_layer_groups.py），group_layers_by_identity直接跳过它们（kv_layer_groups.py）。文档解释了原因：若把它们放进组，会重复工作；若其块大小与默认落入的组不同，还会破坏按组计数的块 ID 数量。

Store 全有或全无（fail-closed）

文档明确存储语义：

• 若块 ID 未完全覆盖每个组的每个 chunk（例如调用方 bug），或某次拷贝失败，整个 store 被跳过、不提交任何内容——后续 retrieve 自然未命中，引擎重算即可；
• 非 GPU 传输路径直接拒绝多组传输。

源码侧，lmcache_driven_transfer.py 在块 ID 数量与num_chunks * blocks_per_chunk不符时记录日志并跳过存储（"skipping the store"），与文档描述一致。

示例推演：两个引擎组、三种 info

文档示例：vLLM 暴露两个引擎组——组 0：层 [0,2,4]，组 1：[1,3]。若层 0–3 形状相同、层 4 形状不同，注册产出：

info 0: engine group 0, layers [0, 2] info 1: engine group 1, layers [1, 3] info 2: engine group 0, layers [4]

块 ID{group 0: [10,11], group 1: [20,21]}按 info 顺序发送为[[10,11], [20,21], [10,11]]（info 0 与 info 2 共享组 0 的块 ID）。这一行为由 test_kv_cache_groups.py 与 test_kv_layer_groups_manager.py 中的EngineGroupInfo(0, (0, 2)) / EngineGroupInfo(1, (1, 3))用例直接验证。

不变量（Invariants）

• list[EngineGroupInfo]的顺序是协议可见的组顺序，调用方为每个 info 发送一个块 ID 列表；
• vLLM 特有的访问全部留在lmcache.integration.vllm；info 只携带中立元数据；
• 服务器端用同一个group_layers_by_identity复现分组；真实 tensor 是 shape/dtype/stride 的唯一事实来源——因此服务器在 KVLayerGroupsManager 构造时校验每个 info 的layer_indices与内核组覆盖的层完全一致（不一致抛ValueError，kv_layer_groups.py），并校验 info 数量与内核组数量一一对应（kv_layer_groups.py）。

Mamba / 线性注意力混合与组编辑

注册前的视图重排

Mamba/线性注意力混合模型通过注册时的 tensor 视图重排（re-view）得到支持，入口是 kv_cache_group_edits.py，配套设计文档见 kv_cache_group_edits.md。两类编辑：

1. Mamba 状态页（如 Qwen3.5 GDN）：层注册[conv_state, ssm_state]两个形状/dtype 不同的 tensor、在一个填充页中连续排布，原始配对会扰乱格式发现；编辑将其重解释为一个 bf16 tensor(num_blocks, 2, block_size, 1, head_size)；
2. 子分页全注意力：vLLM 为统一混合组页大小而膨胀注意力逻辑块大小（如 Qwen3.5-0.8B 为 544），而注意力后端按自己的内核块大小（混合模型上 FlashAttention 为 32）重新分页；编辑将 tensor 视图为(num_kernel_pages / k, 2, logical_block_size, 1, head_size)，避免_derive_compression_metadata把仅1/k的 KV 当作压缩组误传。

关键契约：编辑后的组是"字节不透明"的——store/retrieve 经同一双射的"块 ID → 字节"映射可正确往返，但不支持内容感知处理（serde 压缩、blending、头重分片、布局转换），也不支持跨选择不同内核块大小的注意力后端的缓存共享。编辑均为同一存储上的纯view()，绝不拷贝。

声明的压缩不走编辑

规格声明槽位压缩的组（MLAAttentionSpec.compress_ratio > 1，即 DeepSeek-V4 槽位打包；或TQFullAttentionSpec.tq_slot_size > 0）必须原样进入 kv_layer_groups.py 的压缩路径——编辑规则通过matches的字段检查主动排除它们。DeepSeek-V3.2 的fp8_ds_mla缓存打包的是每槽字节而非槽数，其compress_ratio == 1，同样无需编辑。

启动校验validate_kv_cache_groups（连接器初始化与注册时各执行一次）会聚合拒绝传输路径无法正确服务的规格：CrossAttentionSpec与mamba_cache_mode != "align"的 Mamba（无可复用快照）。

相关配置项与运行入口

HMA 相关的配置分散在服务器端上下文配置中，可通过 multiprocess 服务器的 CLI 参数开启：

• --separate-object-groups：开启后每个滑动窗口大小拆成一个对象组（否则所有内核组合并进单个全注意力对象组），配置项定义于 config.py，消费于 engine_context.py 与 kv_layer_groups.py 的_detect_object_groups；
• 对象组（ObjectGroupInfo）用于区分需要不同前缀匹配逻辑的滑动窗口 / Mamba KV（与全注意力 KV 分开存储）。

此外，benchmark 场景可用--kvcache-shape-spec直接声明多组几何（kv_layer_groups.py 的parse_kvcache_shape_spec），例如异构模型：(2,1024,16,8,128):float16:30;(1,1024,16,4,64):bfloat16:2，供lmcache bench server回显解析结果（format_kvcache_shape_spec为其逆操作）。

代码地图速查

测试覆盖了msgspec往返（IPC 路径无损编解码）、旧 payload 默认字段兼容（sw_size_tokens缺失解码为-1）、块 ID 展开、token 范围切分与对齐校验，以及KVLayerGroupsManager对 info/内核组一一对应的强校验——是理解这套混合分组设计最直接的验证入口。

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