第 8 章 KV Cache：让推理提速 5 倍

本章来源：本文整理自 LLMs-from-scratch 仓库的 ch04/03_kv-cache bonus 材料，原书为 Sebastian Raschka《Build a Large Language Model (From Scratch)》。

这章要解决什么问题

前七章，我们已经把一个 GPT 从零造了出来：会分词、会注意力、会搭块、会预训练、会微调。到这一步，模型能"说话"了——但它说得很慢。

慢在哪？就慢在生成文本这件事上。第 4 章那句"LLM 一次只生成一个词"，背后藏着一个巨大的浪费：每生成一个新词，模型都要把前面所有词重新算一遍。 你让它写 200 个词，它就把同一批历史计算反复做了将近 200 次。

这一章的任务非常聚焦：在不改变输出结果的前提下，把这份重复计算省掉，让推理快起来。

作者在 Mac Mini（M4 芯片，纯 CPU）上实测：同一个 124M 的 GPT，生成 200 个 token，朴素写法每秒只能吐 27 个词，加上 KV Cache 后飙到 144 个词/秒——提速约 5 倍。而两者输出的文字一模一样。

所以这章讲的不是"让模型更聪明"，而是"让模型更省力"。这是所有真正部署 LLM 的系统里，都必须上的一块优化。

先划重点：KV Cache 只在**推理（inference）**时用，训练时用不上。它不加嗅模型能力，只是把已经算过的东西记下来，别再算第二遍。

用大白话讲原理

先看清楚：朴素生成到底浪费在哪

假设模型正在接着写句子，提示词（prompt）是 "Time flies"。第 3 章我们讲过，注意力会让每个词都生成三个向量：

• 查询（Query，简称 Q）：当前这个词"想找什么信息"；
• 键（Key，简称 K）：每个词"能提供什么信息"的索引；
• 值（Value，简称 V）：每个词"实际要提供的内容"。

算注意力时，是用 Q 去和所有 K 打分（看该关注谁），再按分数把 V 加权求和。用公式写就是：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V$$

现在模型生成了下一个词 "fast"，prompt 变成了 "Time flies fast"。按最朴素的写法，模型会把整句重新喂进去，重新算一遍 "Time"、"flies"、"fast" 三个词各自的 Q、K、V。

问题就来了：把新句子和旧句子一对，你会发现 "Time" 和 "flies" 这两个词的 K 和 V，跟上一轮算出来的完全相同。它们只跟自己以及自己前面的词有关，跟后面新加进来的 "fast" 毫无关系。既然如此，每轮都重算它们，纯属白干。

打个比方：你在做一道需要不断"回看前文"的阅读理解题，每读到一个新句子，就把前面读过的所有句子从头抄一遍再重新分析。答案没错，但时间全浪费在抄写上了。

浪费有多大？从 O(n²) 说起

把这份浪费用复杂度算清楚，你就知道为什么它是"大问题"。

假设要生成 $n$ 个 token。第 $t$ 步时，模型手里的历史长度是 $t$。朴素写法下，每一步都要把长度为 $t$ 的历史重新前向一遍，其中注意力那一步还要让当前 query 去和 $t$ 个历史 key 逐个比较。

把所有步骤加起来：

$$\sum_{t=1}^{n} t = \frac{n(n+1)}{2} = O(n^2)$$

这就是二次复杂度——生成长度翻一倍，计算量翻四倍。写短句子还好，一旦生成上千 token，代价就爆炸了。

KV Cache 的核心：把算过的 K 和 V 存起来

洞见只有一句话：K 和 V 只要算过一次就永远不变，那就存起来复用。

• 新 token 只贡献它自己的 Q、K、V；
• 它的 Q 需要跟全部历史 K 比较——所以我们保留所有历史 K；
• 加权求和要用到全部历史 V——所以我们保留所有历史 V。

于是每生成一个新词，我们只算这一个词的 Q、K、V，把新的 K、V 追加进缓存，然后拿出一整份 K、V 来用。每步的计算量从"和 $t$ 个历史重算"变成"只处理 1 个新 token + 查一次缓存"，总复杂度降到：

$$O(n)$$

这就是"线性"和"二次"的差距，也是那 5 倍提速的来源（生成越长，差距越夸张）。

缓存是怎么"随上下文推进"的

缓存不是简单堆一个数组就完事，还要解决一个容易踩坑的细节：因果掩码（causal mask）的偏移。

回想第 3 章：因果掩码负责挡住"未来"，保证每个词只能看到自己及之前的词。在朴素写法里，掩码就是左下三角，位置天经地义。但用了缓存之后，每一步喂进来的只有那一个新 token：

• 它的 query 只有 1 行，但它的 key 有 $t+1$ 列（历史 + 自己）；
• 这个新 token 的"第几个位置"，得靠一个游标（position pointer）来记。

所以缓存的机制里有三样东西要协同：

1. cache_k / cache_v：一路把新的 K、V 拼接（concatenate）进来；
2. 位置指针 ptr_current_pos：记住"当前算到第几个 token 了"，位置嵌入（positional embedding）也用它来给新 token 分配正确的位置编号；
3. 掩码切片：取掩码时从 ptr_current_pos 处开始切，保证因果性不错乱。

还有一件事必须做：用完要清空缓存。两次独立的生成任务之间，缓存绝不能串味，否则第二次生成会"偷看"到上一次的内容。为此要提供 reset_cache。

原理讲完了，下面看作者怎么用五步把它实现出来——代码完全忠于原仓库。

8.1 第一步：注册缓存缓冲区

第 3 章的 MultiHeadAttention 构造函数里，我们加两个缓冲区（buffer）。这里用 register_buffer 而不是普通属性，是为了让它们跟着模型一起正确地放到 GPU/CPU 上，并且在保存模型时不写入权重文件（persistent=False）：

# 在 MultiHeadAttention.__init__ 中新增
self.register_buffer("cache_k", None, persistent=False)  # 缓存的键 K
self.register_buffer("cache_v", None, persistent=False)  # 缓存的值 V
self.ptr_current_pos = 0                                 # 当前 token 位置游标

初始值都是 None，表示"还没有缓存"，等到第一次真正前向时再初始化。

8.2 第二步：forward 里加 use_cache 分支

这是整章最核心的改动。forward 多接一个 use_cache 开关。算出这一批新 token 的 keys_new、values_new 之后，分两条路走：

def forward(self, x, use_cache=False):
    b, num_tokens, d_in = x.shape

    keys_new = self.W_key(x)      # 只对"这次进来的 token"算 K
    values_new = self.W_value(x)  # 只对"这次进来的 token"算 V
    queries = self.W_query(x)

    # 拆成多头: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
    keys_new = keys_new.view(b, num_tokens, self.num_heads, self.head_dim)
    values_new = values_new.view(b, num_tokens, self.num_heads, self.head_dim)
    queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

    # NEW：缓存分支
    if use_cache:
        if self.cache_k is None:                      # 第一次：直接用新算的填满缓存
            self.cache_k, self.cache_v = keys_new, values_new
        else:                                         # 之后：把新 K/V 拼接到历史缓存后
            self.cache_k = torch.cat([self.cache_k, keys_new], dim=1)
            self.cache_v = torch.cat([self.cache_v, values_new], dim=1)
        keys, values = self.cache_k, self.cache_v     # 用"全部历史 K/V"
    else:
        keys, values = keys_new, values_new           # 朴素路径：只用本次的 K/V

    # ...（转置、打分、softmax 等与原版一致）...

注意 dim=1 这个细节：拼接是沿着 token 维度接的，所以 cache_k 的长度会随着生成过程一步步从 2、3、4…… 线性增长。它记的不是"某一层"，而是"这个注意力头见过的所有 token"。

紧接着是掩码的处理，也就是前面说的"游标偏移"。缓存路径下，query 只有新 token 的那么几行（记作 num_tokens_Q），而 key 却有整段长度（记作 num_tokens_K）；于是掩码从 self.mask 里按 [ptr_current_pos : ptr_current_pos + num_tokens_Q, :num_tokens_K] 切片取出，再把 ptr_current_pos 往前推进 num_tokens_Q。朴素路径则直接从 0 取下三角即可，位置天然对齐。

为什么用了缓存后掩码要"切片"？因为缓存路径里 query 只有新 token 的那么几行，而 key 却有整段长度。掩码必须从"新 token 真实所处的位置"开始算，否则因果性就错了。这就是那个游标 ptr_current_pos 存在的全部意义。

8.3 第三步：清空缓存 reset_cache

两次独立生成之间必须把缓存清掉。为此给 MultiHeadAttention 加一个重置方法：

def reset_cache(self):
    self.cache_k, self.cache_v = None, None   # 丢掉所有历史 K/V
    self.ptr_current_pos = 0                  # 游标归零

这一步千万别漏。 如果不清，第二次生成会直接接着上次的缓存往下算，输出结果就错了。缓存是"一次生成会话内"的东西，不能跨会话复用。

8.4 第四步：把 use_cache 一路传下去

改动要贯穿整个模型。前面 MultiHeadAttention 变了，同样地：

TransformerBlock 的 forward 也要多接一个 use_cache 开关，并在调用注意力子块时把它转发进去（self.att(x, use_cache=use_cache)）——其余部分（LayerNorm、dropout、捷径连接、前馈子块）都不变。

GPTModel 有两处改动。第一，把原来一句话的 nn.Sequential 换成显式的 nn.ModuleList 循环，这样才能逐块把 use_cache 传进去；第二，加一个模型级的位置游标 self.current_pos。缓存路径下，新 token 的位置编号不能再从 0 数起，而是用游标接着上次往后排：pos_ids = torch.arange(self.current_pos, self.current_pos + seq_len)，算完再把 self.current_pos += seq_len；朴素路径照旧从 0 数到 seq_len。拿到位置编号后照常取位置嵌入、加上 token 嵌入，再逐块调用 blk(x, use_cache=use_cache)。

这里的 self.current_pos 解决的是**位置嵌入（positional embedding）**问题：新 token 在整个序列里到底是第几个。

最后，还要加一个模型级的批量重置 reset_kv_cache：遍历所有 transformer 块、逐层调用注意力的 reset_cache，再把模型级游标 current_pos 归零，这样就能一次性清掉全部缓存。

8.5 第五步：在生成函数里用起来

万事俱备。生成函数 generate_text_simple_cached 和原版的区别，全在"每次喂给模型多少东西"：

def generate_text_simple_cached(model, idx, max_new_tokens,
                                context_size=None, use_cache=True):
    model.eval()
    ctx_len = context_size or model.pos_emb.num_embeddings

    with torch.no_grad():
        if use_cache:
            # 用完整 prompt 初始化缓存
            model.reset_kv_cache()
            logits = model(idx[:, -ctx_len:], use_cache=True)

            for step in range(max_new_tokens):
                # a) 贪心采样：选概率最高的 token
                next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)
                # b) 拼到运行中的序列末尾
                idx = torch.cat([idx, next_idx], dim=1)
                # c) 只把"新 token"喂给模型（这才是省算力的关键）
                if step + 1 < max_new_tokens:
                    logits = model(next_idx, use_cache=True)
        else:
            # 朴素路径：每一步都把整段历史重新喂进去
            for _ in range(max_new_tokens):
                logits = model(idx[:, -ctx_len:], use_cache=False)
                next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)
                idx = torch.cat([idx, next_idx], dim=1)

    return idx

把两个分支对照着看，差别一目了然：

• 缓存路径：第一轮喂完整 prompt 建立缓存，之后每轮只喂一个新 token（model(next_idx, ...)），因为历史 K/V 已经在缓存里了；
• 朴素路径：每轮都喂 idx[:, -ctx_len:]——整段历史重来一遍，因为它没有存下任何可以复用的 K/V。

代码里 if step + 1 < max_new_tokens 是个贴心的小优化：最后一个 token 生成完就不必再前向了，省掉一次多余的调用。

8.6 实测：27 → 144 tokens/sec

跑起来看效果。书里用 124M 的 GPT、给一句 4 个 token 的提示 "Hello, I am"、让它生成 200 个新 token，分别跑一遍朴素版（gpt_ch04.py）和缓存版（gpt_with_kv_cache.py）。在 Mac Mini（M4 芯片、CPU）上的结果：

实现	Tokens/sec
gpt_ch04.py（朴素）	27
gpt_with_kv_cache.py（KV Cache）	144

约 5 倍提速，而这还只是个 124M 的小模型、200 token 的短序列。模型越大、序列越长，KV Cache 的收益越明显。

有两个点必须说清楚：

第一，两边输出的文字完全一样。 这一点比速度更重要——它说明缓存没有改变语义，只是省了重复劳动。作者特意强调：缓存的索引很容易写错，一旦搞错，得到的文本就会偏离，所以"输出一致"是判断实现是否正确的最硬标准。

第二，输出仍然是乱码。 因为这里用的模型还没训练，权重是随机的。看，它吐出来的是这种东西：

Hello, I am Featureiman Byeswickattribute argue logger Normandy Compton ...

这很正常，第 5 章才训练模型。这里故意用未训练的模型，是为了让代码更干净、聚焦在缓存本身。等模型训好了，随时可以把 KV Cache 套上去生成通顺的文本。

8.7 代价与优化：天下没有免费的午餐

KV Cache 用"记住历史"换来了速度，但这份记忆是要占地方的。

代价：显存随上下文线性增长。 每来一个新 token，就往缓存里追加一份 K 和 V。序列越长，缓存越大。对大模型、长上下文来说，这份开销可能大到、甚至贵到占满 GPU 显存。这就引出了缓存优化。

常见坑

• 内存碎片与反复分配：前面用 torch.cat 一次次拼接张量，每次拼接都要重新分配一块更大的内存、把旧数据拷过去，在长序列下这会成为性能瓶颈。
• 内存无上限增长：如果不加处理，缓存的体积对超长序列来说不现实。

优化 1：预分配内存

与其一遍遍拼接，不如一开始就按预期的最大序列长度，预分配好足够大的张量，之后往里"填格子"，避免反复分配和拷贝。做法是用 torch.zeros((batch_size, num_heads, max_seq_len, head_dim)) 按最大序列长度开好 cache_k / cache_v，推理时只需往对应切片里写数据，不再反复 cat。

优化 2：滑窗截断缓存

为了不让显存被撑爆，可以让缓存只保留最近的 window_size 个 token，更早的就被"滑出去"丢掉。实现上就是每次取 K/V 时沿序列维取最后 window_size 个，例如 cache_k = cache_k[:, :, -window_size:, :]，V 同理。

作者把这些优化实现在 gpt_with_kv_cache_optimized.py 里。同样是 200 token 生成、窗口大小设为上下文长度（以保证结果一致），Mac Mini（M4，CPU）上的对比如下：

实现	Tokens/sec
gpt_ch04.py（朴素）	27
gpt_with_kv_cache.py（KV Cache）	144
gpt_with_kv_cache_optimized.py（优化版）	166

一个反直觉的提醒：这些速度优势在 CUDA 设备上反而可能消失。因为这个模型太小了，设备间数据搬移和通信的开销，会盖过 KV Cache 带来的好处。换句话说，KV Cache 的性价比是"模型越大、序列越长、越值"，小模型上不一定划算。

一句话记住本章：K 和 V 算一次就不会变，那就存下来别再算——这就是 KV Cache。它把生成从 $O(n^2)$ 降到 $O(n)$，代价是显存随上下文线性增长。

关键概念

• KV Cache（键值缓存）：推理时把已经算过的 Key/Value 存下来复用，下一步只计算新 token 的 K/V；只在推理用，训练时用不上。
• 查询/键/值（Query / Key / Value）：注意力的三个向量。Q 是"我想找什么"，K 是"我能被谁找到"，V 是"我实际提供的内容"。
• 因果掩码（causal mask）：遮住"未来" token 的掩码，保证每个词只看自己和之前的词。用缓存后需要靠位置游标对齐切片。
• 位置游标（ptr_current_pos / current_pos）：记录当前算到第几个 token，供掩码切片和位置嵌入使用，随生成推进而增长。
• register_buffer：把张量注册为模型的缓冲区，随模型一起迁移设备，但不作为可训练权重、不写入权重文件。
• reset_cache / reset_kv_cache：清空缓存并把游标归零；两次独立生成之间必须调用。
• 二次复杂度 O(n²)：朴素生成的总计算量。每一步都重算历史，$n$ 步加起来是 $\sum_{t=1}^n t$。
• 线性复杂度 O(n)：KV Cache 后的总计算量。每个 K/V 只算一次，之后复用。
• 预分配内存（pre-allocation）：按最大序列长度预先开好缓存张量，避免反复 torch.cat 带来的分配与拷贝开销。
• 滑窗截断（sliding window）：缓存只保留最近 window_size 个 token，控制显存占用。

动手练习

1. 跑一遍对比：分别运行 gpt_ch04.py 和 gpt_with_kv_cache.py，量一下自己机器上的 tokens/sec 差距，再确认两边输出的文本完全一致——后者才是验证缓存正确性的关键。

2. 故意不清缓存：试着把 generate_text_simple_cached 里的 model.reset_kv_cache() 注释掉，连续调用两次生成函数，观察第二次的输出是否跑偏，亲手体会"为什么必须清缓存"。

3. 亲手改坏掩码：把 forward 里缓存路径的掩码切片 self.ptr_current_pos:self.ptr_current_pos + num_tokens_Q 改成从 0 开始，看看结果会不会偏离——理解那个游标到底在保护什么。

4. 上优化版：阅读并运行 gpt_with_kv_cache_optimized.py，对照 8.7 节的两种优化（预分配 + 滑窗），看看 166 tokens/sec 是怎么来的。

5. 动手实现滑窗：给 cache_k / cache_v 加上 8.7 节那段截断逻辑，把窗口大小设成 128，观察显存占用和生成结果的变化。想清楚：截断缓存后，模型"忘掉"了哪些信息？