[{"content":" 本文以 Moonshot 的 Kimi Linear（2025-10）与 Kimi K3（vLLM PR #50000，2026-07 合入 main）为例， 梳理 vLLM 在面对\u0026quot;线性注意力 + 全注意力\u0026quot;混合架构时，KV cache 是如何组织、分配与复用的。 内容基于 vLLM main 分支源码整理。\n一、问题的起点：什么是\u0026quot;混合架构\u0026quot; 传统 LLM 每一层都是同一种注意力。而 Kimi Linear / Kimi K3 这类模型逐层交替使用两种机制：\nKDA 层（Kimi Delta Attention）：本质是 GatedDeltaNet 线性注意力，属于类 Mamba 的 RNN 结构。它不逐 token 保存 KV，而是维护一个固定大小的\u0026quot;循环状态矩阵\u0026quot;。 MLA 层（Multi-head Latent Attention）：标准全注意力，需要逐 token 缓存压缩后的 latent KV。 在 vLLM 里，模型通过配置决定每层的类型：\n# kimi_linear.py: KimiDecoderLayer.__init__ if config.is_kda_layer(layer_idx): self.self_attn = KimiGatedDeltaNetAttention(...) # 线性注意力 else: self.self_attn = KimiMLAAttention(...) # 全注意力 这两种层对 KV cache 的需求根本不同：\nKDA（线性注意力） MLA（全注意力） 缓存什么 固定大小的 conv_state + recurrent_state 逐 token 的 latent KV 大小随序列长度 不变 线性增长 一个请求用几个块 通常 1~2 个（状态块） 很多（每 block_size 个 token 一块） 这个差异是后面一切设计的根源。\n二、第一层抽象：KVCacheSpec 与\u0026quot;为什么要分组\u0026quot; 2.1 每种层有自己的 KVCacheSpec vLLM 用 KVCacheSpec 描述\u0026quot;一层的 KV cache 长什么样\u0026quot;。混合模型里：\nKDA 层 → MambaSpec（记录 conv/recurrent 两个状态张量的形状） MLA 层 → MLAAttentionSpec（记录分页 latent KV 的形状） KDA 状态形状的计算（与序列长度无关，只和头数/维度有关）：\n# mamba_utils.py: kda_state_shape conv_state_shape = (conv_dim // tp, conv_kernel_size - 1) recurrent_state_shape = (num_heads // tp, head_dim, head_dim) # d×d 记忆矩阵 2.2 分组（KV cache group）的本质 vLLM 把同类型的层归为一个 KV cache group，依据是 KVCacheSpec 的值相等：\n# kv_cache_utils.py same_type_layers = defaultdict(list) for layer_name, layer_spec in kv_cache_spec.items(): same_type_layers[layer_spec].append(layer_name) 为什么必须分组？ 一句话：不同类型的层需要不同的 block table。\n想象同一个请求，序列长 10000 token：\n在某个 MLA 层，它可能用了 40 个块，block table 是 [b0, b1, ..., b39]。 在某个 KDA 层，它只需要 1 个状态块，block table 是 [b0]。 这两种块的\u0026quot;块↔token 映射关系\u0026quot;完全不同，不可能共用一张 block table。所以必须分成不同的 group，每个 group 独立维护自己的 block 分配和 block table。KV cache manager 把\u0026quot;同组的所有层当作一层来管理\u0026quot;——只维护 2~3 张 block table，而不是为几十层各管一张。\n分组还让不同类型能应用差异化策略：sliding window 层可以驱逐窗口外的块，full attention 不能；KDA 用\u0026quot;状态快照\u0026quot;做 prefix cache，MLA 用\u0026quot;逐 token KV\u0026quot;。这些都靠分组隔离。\n关键点：分组是必需的，它和\u0026quot;要不要 padding\u0026quot;是两个正交的问题。\n三、第二层抽象：分组之后，如何塞进同一块显存 分好组后，问题变成：KDA 组的块很小、MLA 组的块很大，怎么放进同一块显存池、还共用一套全局 block_id？vLLM 有两条路径。\n3.1 路径 A：共享 Tensor + 统一 page + padding（默认，一般混合模型走这条） 这是历史实现。核心是让来自不同组的层轮流复用同一个 Tensor 的同一偏移：\n# kv_cache_utils.py: general case kv_cache_tensors.append( KVCacheTensor(size=page_size * num_blocks, shared_by=shared_by) ) 因为一个块的物理偏移是 block_id * page_size，而不同组共用同一个 block_id 和 num_blocks，所以 page_size 必须是全局唯一常量。于是要把小的 KDA page padding 撑大到 MLA page：\n# unify_kv_cache_spec_page_size elif isinstance(layer_spec, MambaSpec): # Mamba page 与 block_size 无关，只能填充 new_spec = replace(layer_spec, page_size_padded=max_page_size) 代价：KDA 层浪费了 padding 部分的显存。但因为混合模型里 KDA 层通常不多，浪费总量有限。收益：单一 block 空间、零碎片、prefix cache 逻辑统一。\n3.2 路径 B：packed 布局（DSV4 默认，其他需 --enable-cross-layers） 这条路径就是很自然的另一种想法：\u0026ldquo;一个逻辑块在所有层的总大小是固定的，按这个总大小算 num_blocks，所有层共享 block_id，各组保留自己真实的 page，不 padding。\u0026rdquo;\nvLLM 的实现完全对应这个思路：\n# _get_packed_kv_cache_layout for group in kv_cache_groups: byte_offset = 0 for layer_name in group.layer_names: page_size = ... # 各层各自的 page，不强行相等 byte_offset += page_size # 累加 block_stride = max(block_stride, byte_offset) # _get_kv_cache_config_packed num_blocks = available_memory // block_stride # 按总大小算 num_blocks 物理上只分配一块 backing，所有层通过 (offset, block_stride) 用 as_strided 切出自己那段，共享同一个 num_blocks（即共享 block_id）：\n# gpu_model_runner.py: _allocate_kv_cache_tensors if kv_cache_tensor.block_stride \u0026gt; 0: if packed_backing is None: packed_backing = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, ...) tensor = packed_backing # 所有 packed 层 alias 同一块 backing packed 免除了 page 尺寸 padding。 这正是\u0026quot;不强行 page 对齐\u0026quot;的正解。\n3.3 那为什么 packed 不是所有混合模型的默认？ 唯一的硬约束是：kernel 能不能读这种布局。\npacked 下，一个层的 KV 不是\u0026quot;连续的 num_blocks 个 page\u0026quot;，而是\u0026quot;每隔 block_stride 字节才有自己的一个 page\u0026quot;（中间夹着别层的数据）。要读它，attention/mamba kernel 必须支持带 block_stride 的跳跃式寻址。vLLM 用一个字段标记 backend 是否支持：\n# AttentionSpec.indexes_kv_by_block_stride padding 逻辑里明确写了这个前提：\nelif isinstance(layer_spec, AttentionSpec) and layer_spec.indexes_kv_by_block_stride: new_spec = replace(layer_spec, page_size_padded=max_page_size) else: raise NotImplementedError( \u0026#34;Padding is only supported for attention layers whose backend \u0026#34; \u0026#34;indexes KV pages by the block stride ...\u0026#34;) 所有相关 backend 都支持 stride 寻址 → 可 packed → 无 page padding。DSV4 的 FlashMLA/FlashInfer 支持，所以走 packed。 有 backend（尤其 Mamba/KDA 的 causal_conv1d、chunk_kda kernel）假设 page 连续 → 只能退回路径 A 的 padding。 结论：padding 不是\u0026quot;设计者想浪费\u0026quot;，而是\u0026quot;部分 kernel 不支持跳跃寻址\u0026quot;的兜底。packed 已实现，是首选，只是不能假设所有 kernel 都支持。\n3.4 澄清一个常见误解：DSV4 的 padding 是\u0026quot;层数\u0026quot;不是\u0026quot;page 尺寸\u0026quot; DSV4 是纯 MLA 但存在多种 page 尺寸（C4 indexer / C4 attn / C128），它全部是 UniformTypeKVCacheSpecs，跟 Mamba 无关。它走 packed，page_size 不需要对齐。但它仍有另一种 padding——layer-tuple 层数 padding：\n# _get_kv_cache_groups_uniform_groups num_layer_tuples = _approximate_gcd(num_layer_tuples_per_group, ...) num_layer_tuples_per_group = [round_up(x, num_layer_tuples) for x in ...] 这是因为不同 page 尺寸的层数不成整数比，为了让\u0026quot;每个逻辑块里各类型取相同层数\u0026quot;而 round_up。这与 KDA 撑到 MLA 大小是完全不同维度的浪费。\n四、KV cache 在内存里到底长什么样 理解\u0026quot;物理\u0026quot;和\u0026quot;逻辑\u0026quot;分离很重要。\n4.1 物理：一维 int8 缓冲 底层就是一块 int8 字节缓冲（size = page_size × num_blocks），且被多个层共享（shared_by）。\n4.2 逻辑：各 backend 的 reshape 视图 MLA 层：(num_blocks, block_size, kv_lora_rank + qk_rope_head_dim)。因为只存压缩 latent、无 K/V 分离，单 token 一个大 latent 向量——这就是 MLA page 大（~6K 量级）的来源。 KDA 层：先保持字节页视图 (num_blocks, 1, 1, page_size_bytes)，再由 bind_kv_cache 拆成 conv_state / recurrent_state 两个子视图。 # gpu_model_runner.py: _reshape_kv_cache_tensors (Mamba 分支) kv_caches[layer_name] = raw_tensor[: num_blocks * page_size_bytes] \\ .view(num_blocks, 1, 1, page_size_bytes) 五、KDA 的计算原理：为什么能并行、为什么 prefix 能省算 5.1 状态即\u0026quot;历史的压缩\u0026quot; KDA 每层每请求保存两样固定大小的东西：\nconv_state：短卷积滑动窗口（最近 kernel-1 个 token）。 recurrent_state：d×d 的记忆矩阵，编码整个历史。 递归本质（简化版，忽略门控）：\nS_t = S_{t-1} + β_t · k_t v_tᵀ # 状态更新 o_t = q_t · S_t # 输出 S_t 永远是 d×d，与 t 无关——这是\u0026quot;状态大小固定\u0026quot;。但 S_t 的内容是前 t 个 token 的函数——这是 prefix 能复用的前提。\n5.2 为什么串行依赖却能并行（chunk scan） 关键：线性递归可以展开成求和，求和可以并行。\n把 S_t 展开：S_t = S_0 + Σ_{j≤t} k_j v_jᵀ，于是\no_t = q_t·S_0 + Σ_{j≤t} (q_t·k_j) v_j └──── 这就是 attention 形式，可并行 ────┘ Chunk-wise parallel scan 把两种视角结合：\n序列切 chunk: [chunk0] [chunk1] [chunk2] S_0 ─▶ S_C ─────▶ S_2C ─────▶ S_3C （chunk 间：串行传 d×d 状态，步数 = N/C） 每个 chunk 内 token t 的输出 = 两部分之和： ① inter-chunk: q_t · S_{chunk起始} ← 一个 GEMM 覆盖整个 chunk（并行） ② intra-chunk: Σ_{j在chunk内且j≤t} (q_t·k_j)v_j ← chunk 内 C×C causal attn（并行） chunk 内：矩阵乘吃满 GPU。 chunk 间：只传 d×d 状态，串行深度仅 N/C。 对应代码：prefill 用 chunk_kda_with_fused_gate，decode 用 fused_recurrent_kda（逐步）。带门控时衰减因子可累乘（A_log / dt_bias 在 fused_kda_gate 预计算），不破坏可展开性。\n5.3 prefix cache 省的是\u0026quot;扫描 token 数\u0026quot;，不是存储 这是最容易误解的点。KDA 的 prefix cache 缓存的是处理完某前缀后的 recurrent_state（那个 d×d 矩阵）。\n省算发生在两处配合：\nScheduler 层面：命中前缀被标为 num_computed_tokens，不进入本次 forward 的 q/k/v。 状态层面：从缓存载入对应 recurrent_state 作为 initial_state： # kimi_gdn_linear_attn.py: _forward zero_idx = state_indices[~has_initial_state] recurrent_state[zero_idx] = 0 # 未命中 → 清零从头 initial_state = recurrent_state[state_indices] # 命中 → 用缓存状态接着算 举例：序列 1000 token，前 512 命中：\n无 cache：chunk scan 处理 1000 token，initial_state=0。 有 cache：chunk scan 只处理 488 token（512~1000），initial_state=缓存里处理完 512 token 的 S。 省掉的就是前 512 token 重新扫一遍的 FLOPs。\u0026ldquo;状态大小固定\u0026quot;不妨碍省算，因为省的是计算量而非存储。\n六、Kimi K3 的点睛之笔：用更细的 prefix_match_unit 救回命中率 6.1 矛盾 为了和 KDA 对齐 + 省管理开销，K3 的 MLA block_size 很大（page ~6K 量级）。但标准 prefix cache 命中只能落在 block_size 整数倍边界——block 越大，命中粒度越粗，长共享前缀经常只命中一小截。\n共享前缀 400 token, block_size=256： 块边界: 0 ────── 256 ────── 512 命中只能到 256，256~400 这 144 token 明明相同却复用不了 6.2 解法：把\u0026quot;命中粒度\u0026quot;从\u0026quot;物理块大小\u0026quot;解耦出来 K3 引入三个独立的 block size：\n名称 取值 作用 各组 block_size MLA 大、KDA 小 物理块大小（显存布局） scheduler_block_size lcm(各组 block_size) 调度对齐 hash_block_size = prefix_match_unit gcd(各组) 或用户指定更小值 命中粒度 # kv_cache_utils.py: get_kv_cache_configs hash_block_size = prefix_match_unit or math.gcd(*group_block_sizes) scheduler_block_size = math.lcm(*group_block_sizes) 物理块仍然大（省显存/管理），但命中判定用细粒度。\n6.3 三个机制协同 ① hash 永远按细粒度、链式计算\n# get_request_block_hasher：按 hash_block_size 切，每个 hash 折进前一个 链式使得\u0026quot;第 k 个细粒度 hash = 前 k×hash_block_size 个 token 整个前缀的指纹\u0026rdquo;。\n② 粗粒度组借用细粒度 hash（无需重算） 因为链式，一个大块的 hash = 它最后一个细粒度子块的 hash：\n# BlockHashListWithBlockSize._get_value_at return self.block_hashes[(idx + 1) * self.scale_factor - 1] 一套 hash，两种粒度共用。\n③ 命中查找两阶段\n# FullAttentionManager.find_longest_cache_hit # Phase 1: 按大 block_size 找从头连续命中的完整大块 # Phase 2 (fine_grained): 在首个未满大块内部，按细粒度边界高→低探测 hit_length = (fine_idx + 1) * alignment_tokens # 命中落在细粒度边界 于是 400-token 例子里，命中能从 256 延伸到 384，多复用 128 token。\n④ 写入侧配套：请求在细粒度边界注册 partial 尾块（_cache_partial_tail_block），供后续请求命中。\n6.4 整体闭环 prefix_match_unit（细粒度，如 64） │ ┌───────────┼───────────┐ ▼ ▼ ▼ hash 计算 命中查找 缓存写入 按 64 链式 Phase1 大块 请求在 384(细边界) 一套通用 Phase2 块内探测 注册 partial 尾块 │ │ │ └─ MLA 组按 scale 抽粗粒度 hash ─┘ 物理块仍 6K（省显存不变）→ 命中粒度 64（命中率大增） 精髓：物理块大小与命中粒度是两个诉求，用两个独立的量各自满足，不必为命中率牺牲块大小。\n七、mamba_cache_mode：all vs align KDA 的 prefix cache 需要在某些位置保存中间 recurrent_state 快照，两种模式的区别在于保存哪些位置：\nall：在每个 block 边界都存快照。命中点密、复用机会多，但存储/写入开销大，需模型显式支持（supports_mamba_prefix_caching）。 align：只在scheduler step 恰好停在 block 边界时存快照。命中点稀疏、存储省，靠块内运行状态 + causal_conv1d 对齐拷贝补足；是通用 fallback，要求 chunked prefill。 # MambaModelConfig.verify_and_update_config cache_config.mamba_cache_mode = \u0026#34;all\u0026#34; if supports_mamba_prefix_caching else \u0026#34;align\u0026#34; 八、总结 混合架构的 KV cache = 不同类型层用不同 KVCacheSpec（KDA→MambaSpec，MLA→MLAAttentionSpec），分成不同 group 各自管理 block table。分组是必需的。 分组后放进同一显存有两条路： 共享 Tensor + 统一 page + padding（默认，KDA 被撑大），或 packed（各组 page 累加、共享 block_id、无 page padding）。packed 就是\u0026quot;不强行对齐\u0026quot;的正解，已实现（DSV4 默认），但受限于 kernel 是否支持 block_stride 寻址。 DSV4 走 packed 与 Mamba 无关（它是纯 MLA 多 page 尺寸），且它残留的是层数 padding 而非 page 尺寸 padding。 KDA 靠固定大小的 recurrent_state 概括历史；chunk scan 通过\u0026quot;线性递归展开成求和\u0026quot;实现并行；prefix cache 省的是重扫前缀的计算，靠载入缓存状态作 initial_state。 Kimi K3 用独立的 prefix_match_unit 把命中粒度从大物理块解耦，通过\u0026quot;细粒度链式 hash + 两阶段命中 + partial 尾块\u0026quot;救回大 MLA page 的命中率。 本文由技术讨论整理而成，结论待实机（DeepSeek V4 / Kimi K3）验证。\n","permalink":"https://adeline-ly-blog.pages.dev/posts/vllm-hybrid-kvcache/","summary":"以 Kimi Linear 与 Kimi K3 为例，梳理 vLLM 面对线性注意力 + 全注意力混合架构时，KV cache 如何组织、分配与复用。","title":"vLLM 如何管理混合架构模型的 KV Cache：从 Kimi Linear / Kimi K3 说起"}]