跳到主要内容

数据截至 (上游 commit b6c0bfe04c82)

04 · KV cache 体系

这一章讲什么: 自回归推理的显存大头与性能命脉——KV cache——在 v5 里怎么组织。核心是一个二级抽象:Cache 是容器,CacheLayerMixin 的子类们是每层的存取策略。读完你能选对 cache 类,也知道滑窗、静态、量化、线性注意力各自的代价。

1. 它要解决的小问题​

自回归生成的第 t 步,attention 需要的是所有历史 token 的 K/V。不重算就得存下来——这就是 KV cache,显存占用随 batch × 层数 × 头数 × 序列长度 线性增长,长序列下比模型权重还大。

难的不是「存起来」,而是不同架构、不同后端要的存法根本不一样:

  • 普通层要无限增长的 K/V;滑窗层只留最后 W 个;线性注意力(Mamba 类)存的是固定大小的循环状态,根本没有 per-token KV。
  • 动态图(eager)喜欢「拼上去」;torch.compile/CUDA graph 要求张量地址固定,只能「预分配 + 原位写」。
  • 显存不够时想把部分层甩到 CPU,还得用另一个 CUDA stream 预取下一层,不能让主流程等。

一个抽象要同时装下这些,就不能是一个「past_key_values 元组」了事。

2. 思路:容器归容器,策略归策略​

v5 的拆法(cache_utils.py):

Cache(容器,一个模型一个)         cache_utils.py:1269
  ├── layers: [Layer, Layer, ...]   一层一个
  ├── update(k, v, layer_idx)      → 转交 layers[layer_idx].update()
  ├── reorder_cache(beam_idx)      → beam search 重排,逐层转发
  └── offload / prefetch           → CPU offloading 的流管理

CacheLayerMixin(策略,一层一个)    cache_utils.py:27
  ├── DynamicLayer           全量,torch.cat 往上拼
  ├── DynamicSlidingWindowLayer  只留 window-1 个
  ├── StaticLayer            预分配定长,index_copy_ 原位写
  ├── QuantizedLayer         动态存,但存量化后的
  └── LinearAttentionLayer   无 KV,存 conv/循环状态

模型代码完全不感知这些差异。 attention 里只有一行(models/llama/modeling_llama.py:262):

key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

返回值是「本次 attention 该看的完整 K/V」——对 DynamicLayer 是拼完的全量,对滑窗层是窗口内,对 StaticLayer 是整个定长 buffer(靠 mask 遮住还没写的部分)。

3. 图示:prefill 与 decode 时 cache 里发生什么​

prefill(prompt 长度 L)                    decode(每次 1 个 token)
─────────────────────                     ─────────────────────────
模型 forward 整段                          模型 forward 只喂新 token
   │                                         │
   ▼ 每层:                                   ▼ 每层:
Layer.update(K[L], V[L])                   Layer.update(K[1], V[1])
   │  Dynamic: cat → 长度变 L                  │  Dynamic: cat → 长度 +1
   │  Static:  原位写 [0:L]                    │  Static:  原位写 [t:t+1]
   ▼                                         ▼
返回完整 K/V 给 attention                   返回完整 K/V(含历史)

循环侧的配合在 prepare_inputs_for_generation 的默认实现(generation/utils.py:519):有 cache 时只切出最后 next_sequence_length 个 token 喂模型(:552),并给 input_ids.clone(memory_format=torch.contiguous_format)——注释(:551)指向 issue #32227,保持 stride 一致是为编译内核准备的。

4. 原理演示​

Dynamic 与 Static 的本质差别,12 行演完:

# 示意,非源码
class ToyDynamicLayer:          # 对应 DynamicLayer
    def update(self, k, v):
        self.keys = torch.cat([self.keys, k], dim=-2)   # 新张量,地址每次都变
        return self.keys

class ToyStaticLayer:           # 对应 StaticLayer
    def __init__(self, max_len):
        self.keys = torch.zeros(B, H, max_len, D)       # 预分配,地址终身不变
        self.t = 0
    def update(self, k, v):
        pos = torch.arange(k.shape[-2]) + self.t
        self.keys.index_copy_(2, pos, k)                # 原位写
        self.t += k.shape[-2]
        return self.keys                                 # 返回整个 buffer,靠 mask 遮空位

重点看:Static 返回的是含大量「未写区域」的定长 buffer,正确性由 attention mask 保证——这就是为什么它必须与 mask 生成(masking_utils.py 的 create_causal_mask 会询问每层的 get_mask_sizes)配套。

5. 真实实现:三个 Layer 的关键三行​

5.1 DynamicLayer:拼​

DynamicLayer.update(cache_utils.py:127)的全部核心(:145-146):

self.keys = torch.cat([self.keys, key_states], dim=-2)
self.values = torch.cat([self.values, value_states], dim=-2)

简单、永远正确,但每步都新分配并拷贝全部历史——长序列下这是 eager 模式的主要浪费。

5.2 StaticLayer:原位写​

StaticLayer.update(cache_utils.py:455)用累计长度算写入位置,然后 index_copy_ 原位写(:480-481):

self.keys.index_copy_(2, cache_position, key_states)
self.values.index_copy_(2, cache_position, value_states)

注释(:475)点破原因:"has to be performed in-place, as we have a static address that we need to keep"——CUDA graph / torch.compile fullgraph 录制后张量地址不能变。MPS 不支持 index_copy_,留了切片赋值 fallback(:482-486)。

5.3 DynamicSlidingWindowLayer:留窗口​

DynamicSlidingWindowLayer.update(cache_utils.py:229)先拼出全量,再只留窗口(:249-254):

full_key_states = torch.cat([self.keys, key_states], dim=-2)
...
self.keys = full_key_states[:, :, -self.sliding_window + 1 :, :]

存 -window + 1 而不是 -window:本步的 query 对应的 K/V 也要参与本轮 attention,缓存里只需保留「对下一个 token 有用」的部分,即窗口减去当前这一个。

6. 容器层:Cache 在做什么​

Cache(cache_utils.py:1269)自身不存张量,它的 update(:1356)做三件事:

  1. 懒建层:若构造时没给 layers 而给了 layer_class_to_replicate,遇到新 layer_idx 就 append 一个新层(:1375-1376)。
  2. offloading 编排:开了 offloading 就先 wait_stream 等预取流,再 prefetch(layer_idx + 1) 提前拉下一层(:1378-1381),更新完把本层 offload 到 CPU(:1385-1386)。预取走独立 torch.Stream,与主计算流重叠。
  3. 转发:reorder_cache/batch_repeat_interleave 逐层转发——beam search 每轮换 beam 归属时靠它重排历史。

6.1 按 config 自动选层类型​

DynamicCache(config=...) 构造时(:1789-1793)调 get_layer_types_and_kwargs(config)(:1701)读出每层的 layer_types,再查 DYNAMIC_LAYER_TYPE_MAPPING(:1223)逐个实例化。映射表本身就是一份「现代架构层类型」小百科(:1224-1238):

layer_types 值Layer 类用于
full_attentionDynamicLayer标准注意力层
sliding_attention / chunked_attentionDynamicSlidingWindowLayerMistral 系滑窗 / Llama4 chunked
conv / linear_attentionLinearAttentionLayerMamba 类、线性注意力
hybrid / hybrid_slidingLinearAttentionAndFullAttentionLayer 等同一层里线性状态 + 全量注意力(Jamba 类)
deepseek_sparse_attention / qwen_sparse_attentionDynamicIndexedLayer带 indexer 的稀疏注意力

这就是「hybrid 架构免费获得 cache 支持」的机制:模型只要在 config 里声明每层的类型,容器自动拼出正确的层序列。

6.2 两个兄弟变体​

  • StaticCache(:1829):全层静态,配合 compile_config 的 CUDA graph 路径。
  • QuantizedCache(:1884)/QuantizedLayer(:703):存量化后的 K/V,attention 前反量化——用精度换显存。
  • EncoderDecoderCache(:1947):encoder-decoder 模型里 self-attention 用动态层、cross-attention 缓存 encoder 侧 K/V 的合体容器。

7. 关键细节与坑​

  • Dynamic 的 cat 是 O(长度²) 的内存搬运。 每一步都拷贝全部历史。交互/调试无所谓;吞吐场景请换 Static + compile,或直接上 vllm 的 paged attention。
  • Static 的代价是 max_cache_len 定死。 超长直接报错而不是变慢;batch 内不同长度的序列共享同一个 buffer 宽度,短序列浪费显存。
  • beam search 会物理重排 cache。 每轮按新 beam 归属 reorder_cache(Cache.reorder_cache,:1604)——cache 越大这轮拷贝越贵,beam 数别开太大。
  • offloading 默认只甩非滑窗层。 滑窗层本来就小,甩了反而慢,所以基类 Cache.__init__ 里 offload_only_non_sliding 默认 True(:1294);但 DynamicCache 把它翻成 False(:1785),静态侧的 StaticCache 又是 True(:1874)——三个默认值不一样,用前看一眼签名。
  • assisted decoding 强制 dynamic。 每轮收下的 token 数不固定,定长 buffer 处理不了(第 3 章 §6)。
  • 线性注意力层没有 per-token KV。 它们的 update_conv_state/update_recurrent_state(LinearAttentionCacheLayerMixin,:891)走另一条接口;把这类模型当普通 transformer 调试 cache 形状会困惑很久。

8. 代码地图​

主题文件路径符号名
容器抽象src/transformers/cache_utils.pyCache(update、reorder_cache、prefetch、offload)
层策略基类src/transformers/cache_utils.pyCacheLayerMixin
动态层src/transformers/cache_utils.pyDynamicLayer.update、DynamicCache
滑窗层src/transformers/cache_utils.pyDynamicSlidingWindowLayer
静态层src/transformers/cache_utils.pyStaticLayer.update(index_copy_、cumulative_length)、StaticCache
量化层src/transformers/cache_utils.pyQuantizedLayer、QuantizedCache
线性注意力层src/transformers/cache_utils.pyLinearAttentionCacheLayerMixin、LinearAttentionLayer、LinearAttentionAndFullAttentionLayer
层类型分派src/transformers/cache_utils.pyDYNAMIC_LAYER_TYPE_MAPPING、get_layer_types_and_kwargs
encoder-decoder cachesrc/transformers/cache_utils.pyEncoderDecoderCache
模型侧挂点src/transformers/models/llama/modeling_llama.pyLlamaModel.forward(:383-384 建 cache)、LlamaAttention.forward(:262 写 cache)
解码侧配合src/transformers/generation/utils.pyprepare_inputs_for_generation
mask 配套src/transformers/masking_utils.pycreate_causal_mask、create_masks_for_generate