vLLM 如何管理混合架构模型的 KV Cache：从 Kimi Linear / Kimi K3 说起

本文以 Moonshot 的 Kimi Linear（2025-10）与 Kimi K3（vLLM PR #50000，2026-07 合入 main）为例，梳理 vLLM 在面对"线性注意力 + 全注意力"混合架构时，KV cache 是如何组织、分配与复用的。内容基于 vLLM main 分支源码整理。

一、问题的起点：什么是"混合架构"

传统 LLM 每一层都是同一种注意力。而 Kimi Linear / Kimi K3 这类模型逐层交替使用两种机制：

KDA 层（Kimi Delta Attention）：本质是 GatedDeltaNet 线性注意力，属于类 Mamba 的 RNN 结构。它不逐 token 保存 KV，而是维护一个固定大小的"循环状态矩阵"。
MLA 层（Multi-head Latent Attention）：标准全注意力，需要逐 token 缓存压缩后的 latent KV。

在 vLLM 里，模型通过配置决定每层的类型：

# kimi_linear.py: KimiDecoderLayer.__init__
if config.is_kda_layer(layer_idx):
    self.self_attn = KimiGatedDeltaNetAttention(...)   # 线性注意力
else:
    self.self_attn = KimiMLAAttention(...)             # 全注意力

这两种层对 KV cache 的需求根本不同：

	KDA（线性注意力）	MLA（全注意力）
缓存什么	固定大小的 `conv_state` + `recurrent_state`	逐 token 的 latent KV
大小随序列长度	不变	线性增长
一个请求用几个块	通常 1~2 个（状态块）	很多（每 block_size 个 token 一块）

这个差异是后面一切设计的根源。

二、第一层抽象：KVCacheSpec 与"为什么要分组"

2.1 每种层有自己的 KVCacheSpec

vLLM 用 KVCacheSpec 描述"一层的 KV cache 长什么样"。混合模型里：

KDA 层 → MambaSpec（记录 conv/recurrent 两个状态张量的形状）
MLA 层 → MLAAttentionSpec（记录分页 latent KV 的形状）

KDA 状态形状的计算（与序列长度无关，只和头数/维度有关）：

# 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 的值相等：

# 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。

想象同一个请求，序列长 10000 token：

在某个 MLA 层，它可能用了 40 个块，block table 是 [b0, b1, ..., b39]。
在某个 KDA 层，它只需要 1 个状态块，block table 是 [b0]。

这两种块的"块↔token 映射关系"完全不同，不可能共用一张 block table。所以必须分成不同的 group，每个 group 独立维护自己的 block 分配和 block table。KV cache manager 把"同组的所有层当作一层来管理"——只维护 2~3 张 block table，而不是为几十层各管一张。

分组还让不同类型能应用差异化策略：sliding window 层可以驱逐窗口外的块，full attention 不能；KDA 用"状态快照"做 prefix cache，MLA 用"逐 token KV"。这些都靠分组隔离。

关键点：分组是必需的，它和"要不要 padding"是两个正交的问题。

三、第二层抽象：分组之后，如何塞进同一块显存

分好组后，问题变成：KDA 组的块很小、MLA 组的块很大，怎么放进同一块显存池、还共用一套全局 block_id？vLLM 有两条路径。

3.1 路径 A：共享 Tensor + 统一 page + padding（默认，一般混合模型走这条）

这是历史实现。核心是让来自不同组的层轮流复用同一个 Tensor 的同一偏移：

# 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：

# 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 逻辑统一。

3.2 路径 B：packed 布局（DSV4 默认，其他需 `--enable-cross-layers`）

这条路径就是很自然的另一种想法：“一个逻辑块在所有层的总大小是固定的，按这个总大小算 num_blocks，所有层共享 block_id，各组保留自己真实的 page，不 padding。”

vLLM 的实现完全对应这个思路：

# _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）：

# gpu_model_runner.py: _allocate_kv_cache_tensors
if kv_cache_tensor.block_stride > 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。 这正是"不强行 page 对齐"的正解。

3.3 那为什么 packed 不是所有混合模型的默认？

唯一的硬约束是：kernel 能不能读这种布局。

packed 下，一个层的 KV 不是"连续的 num_blocks 个 page"，而是"每隔 block_stride 字节才有自己的一个 page"（中间夹着别层的数据）。要读它，attention/mamba kernel 必须支持带 block_stride 的跳跃式寻址。vLLM 用一个字段标记 backend 是否支持：

# AttentionSpec.indexes_kv_by_block_stride

padding 逻辑里明确写了这个前提：

elif 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(
        "Padding is only supported for attention layers whose backend "
        "indexes KV pages by the block stride ...")

所有相关 backend 都支持 stride 寻址 → 可 packed → 无 page padding。DSV4 的 FlashMLA/FlashInfer 支持，所以走 packed。
有 backend（尤其 Mamba/KDA 的 causal_conv1d、chunk_kda kernel）假设 page 连续 → 只能退回路径 A 的 padding。

结论：padding 不是"设计者想浪费"，而是"部分 kernel 不支持跳跃寻址"的兜底。packed 已实现，是首选，只是不能假设所有 kernel 都支持。

3.4 澄清一个常见误解：DSV4 的 padding 是"层数"不是"page 尺寸"

DSV4 是纯 MLA 但存在多种 page 尺寸（C4 indexer / C4 attn / C128），它全部是 UniformTypeKVCacheSpecs，跟 Mamba 无关。它走 packed，page_size 不需要对齐。但它仍有另一种 padding——layer-tuple 层数 padding：

# _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 尺寸的层数不成整数比，为了让"每个逻辑块里各类型取相同层数"而 round_up。这与 KDA 撑到 MLA 大小是完全不同维度的浪费。

四、KV cache 在内存里到底长什么样

理解"物理"和"逻辑"分离很重要。

4.1 物理：一维 int8 缓冲

底层就是一块 int8 字节缓冲（size = page_size × num_blocks），且被多个层共享（shared_by）。

4.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 状态即"历史的压缩"

KDA 每层每请求保存两样固定大小的东西：

conv_state：短卷积滑动窗口（最近 kernel-1 个 token）。
recurrent_state：d×d 的记忆矩阵，编码整个历史。

递归本质（简化版，忽略门控）：

S_t = S_{t-1} + β_t · k_t v_tᵀ      # 状态更新
o_t = q_t · S_t                     # 输出

S_t 永远是 d×d，与 t 无关——这是"状态大小固定"。但 S_t 的内容是前 t 个 token 的函数——这是 prefix 能复用的前提。

5.2 为什么串行依赖却能并行（chunk scan）

关键：线性递归可以展开成求和，求和可以并行。

把 S_t 展开：S_t = S_0 + Σ_{j≤t} k_j v_jᵀ，于是

o_t = q_t·S_0 + Σ_{j≤t} (q_t·k_j) v_j
                └──── 这就是 attention 形式，可并行 ────┘

Chunk-wise parallel scan 把两种视角结合：

序列切 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 预计算），不破坏可展开性。

5.3 prefix cache 省的是"扫描 token 数"，不是存储

这是最容易误解的点。KDA 的 prefix cache 缓存的是处理完某前缀后的 recurrent_state（那个 d×d 矩阵）。

省算发生在两处配合：

Scheduler 层面：命中前缀被标为 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 命中：

无 cache：chunk scan 处理 1000 token，initial_state=0。
有 cache：chunk scan 只处理 488 token（512~1000），initial_state=缓存里处理完 512 token 的 S。

省掉的就是前 512 token 重新扫一遍的 FLOPs。“状态大小固定"不妨碍省算，因为省的是计算量而非存储。

六、Kimi K3 的点睛之笔：用更细的 `prefix_match_unit` 救回命中率

6.1 矛盾

为了和 KDA 对齐 + 省管理开销，K3 的 MLA block_size 很大（page ~6K 量级）。但标准 prefix cache 命中只能落在 block_size 整数倍边界——block 越大，命中粒度越粗，长共享前缀经常只命中一小截。

共享前缀 400 token, block_size=256：
块边界: 0 ────── 256 ────── 512
命中只能到 256，256~400 这 144 token 明明相同却复用不了

6.2 解法：把"命中粒度"从"物理块大小"解耦出来

K3 引入三个独立的 block size：

名称	取值	作用
各组 `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)

物理块仍然大（省显存/管理），但命中判定用细粒度。

6.3 三个机制协同

① hash 永远按细粒度、链式计算

# get_request_block_hasher：按 hash_block_size 切，每个 hash 折进前一个

链式使得"第 k 个细粒度 hash = 前 k×hash_block_size 个 token 整个前缀的指纹”。

② 粗粒度组借用细粒度 hash（无需重算） 因为链式，一个大块的 hash = 它最后一个细粒度子块的 hash：

# BlockHashListWithBlockSize._get_value_at
return self.block_hashes[(idx + 1) * self.scale_factor - 1]

一套 hash，两种粒度共用。

③ 命中查找两阶段

# 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。

④ 写入侧配套：请求在细粒度边界注册 partial 尾块（_cache_partial_tail_block），供后续请求命中。

6.4 整体闭环

        prefix_match_unit（细粒度，如 64）
                 │
     ┌───────────┼───────────┐
     ▼           ▼           ▼
  hash 计算    命中查找     缓存写入
  按 64 链式   Phase1 大块   请求在 384(细边界)
  一套通用     Phase2 块内探测  注册 partial 尾块
     │           │           │
     └─ MLA 组按 scale 抽粗粒度 hash ─┘
  物理块仍 6K（省显存不变）→ 命中粒度 64（命中率大增）

精髓：物理块大小与命中粒度是两个诉求，用两个独立的量各自满足，不必为命中率牺牲块大小。

七、mamba_cache_mode：all vs align

KDA 的 prefix cache 需要在某些位置保存中间 recurrent_state 快照，两种模式的区别在于保存哪些位置：

all：在每个 block 边界都存快照。命中点密、复用机会多，但存储/写入开销大，需模型显式支持（supports_mamba_prefix_caching）。
align：只在scheduler step 恰好停在 block 边界时存快照。命中点稀疏、存储省，靠块内运行状态 + causal_conv1d 对齐拷贝补足；是通用 fallback，要求 chunked prefill。

# MambaModelConfig.verify_and_update_config
cache_config.mamba_cache_mode = "all" if supports_mamba_prefix_caching else "align"

八、总结

混合架构的 KV cache = 不同类型层用不同 KVCacheSpec（KDA→MambaSpec，MLA→MLAAttentionSpec），分成不同 group 各自管理 block table。分组是必需的。
分组后放进同一显存有两条路： 共享 Tensor + 统一 page + padding（默认，KDA 被撑大），或 packed（各组 page 累加、共享 block_id、无 page padding）。packed 就是"不强行对齐"的正解，已实现（DSV4 默认），但受限于 kernel 是否支持 block_stride 寻址。
DSV4 走 packed 与 Mamba 无关（它是纯 MLA 多 page 尺寸），且它残留的是层数 padding 而非 page 尺寸 padding。
KDA 靠固定大小的 recurrent_state 概括历史；chunk scan 通过"线性递归展开成求和"实现并行；prefix cache 省的是重扫前缀的计算，靠载入缓存状态作 initial_state。
Kimi K3 用独立的 prefix_match_unit 把命中粒度从大物理块解耦，通过"细粒度链式 hash + 两阶段命中 + partial 尾块"救回大 MLA page 的命中率。

本文由技术讨论整理而成，结论待实机（DeepSeek V4 / Kimi K3）验证。

一、问题的起点：什么是"混合架构"#

二、第一层抽象：KVCacheSpec 与"为什么要分组"#

2.1 每种层有自己的 KVCacheSpec#

2.2 分组（KV cache group）的本质#

三、第二层抽象：分组之后，如何塞进同一块显存#

3.1 路径 A：共享 Tensor + 统一 page + padding（默认，一般混合模型走这条）#

3.2 路径 B：packed 布局（DSV4 默认，其他需 --enable-cross-layers）#

3.3 那为什么 packed 不是所有混合模型的默认？#

3.4 澄清一个常见误解：DSV4 的 padding 是"层数"不是"page 尺寸"#

四、KV cache 在内存里到底长什么样#

4.1 物理：一维 int8 缓冲#

4.2 逻辑：各 backend 的 reshape 视图#

五、KDA 的计算原理：为什么能并行、为什么 prefix 能省算#

5.1 状态即"历史的压缩"#

5.2 为什么串行依赖却能并行（chunk scan）#

5.3 prefix cache 省的是"扫描 token 数"，不是存储#

六、Kimi K3 的点睛之笔：用更细的 prefix_match_unit 救回命中率#

6.1 矛盾#

6.2 解法：把"命中粒度"从"物理块大小"解耦出来#

6.3 三个机制协同#

6.4 整体闭环#

七、mamba_cache_mode：all vs align#

八、总结#