KV Cache 原理与优化：从产生到管理，一个老工程师的实战笔记

大家好，我是老张。今天咱们聊聊 KV Cache。

说实话，我刚接触大模型推理时，也被这玩意儿搞晕过。明明模型参数没变，怎么推理时显存就爆了？后来才明白，罪魁祸首往往就是 KV Cache。今天我就把这块掰开了揉碎了讲清楚。

KV Cache 是怎么产生的？

先看一个最简单的场景。你让模型生成一句话，比如“今天天气真好”。

模型是自回归的，一个字一个字往外蹦。生成“今天”时，它要算一遍注意力。生成“天气”时，它又要算一遍注意力。但注意了，算“天气”时，“今天”的 Key 和 Value 其实已经算过了。

那能不能复用呢？

当然能。这就是 KV Cache 的由来。我们把每个 token 算出来的 Key 和 Value 存起来，后面生成新 token 时直接拿来用，不用重新算。

核心思想：用空间换时间。多占点显存，但省掉大量重复计算。

我刚开始做推理优化时，就犯过一个低级错误。我把所有层的 KV Cache 都存了，但忘了存位置编码。结果生成出来的句子语义完全不对。嗯，这个坑你们别踩。

显存占用计算：一个公式搞定

KV Cache 到底占多少显存？我给大家一个公式：

显存占用 = 2 × 层数 × 头数 × 头维度 × 序列长度 × 精度字节数

为什么是 2？因为 Key 和 Value 各一份。

举个例子。一个 7B 模型，32 层，32 个头，头维度 128，序列长度 2048，用 FP16（2 字节）。

算一下：

2 × 32 × 32 × 128 × 2048 × 2 = 1,073,741,824 字节 ≈ 1GB

你看，光 KV Cache 就吃掉 1GB 显存。如果 batch size 是 8，那就是 8GB。这还没算模型参数和中间激活。

模型规模 层数 序列长度 KV Cache 占用（FP16）
7B 32 2048 ~1GB
13B 40 4096 ~5GB
70B 80 8192 ~40GB

注意：序列长度翻倍，显存占用也翻倍。长文本场景下，KV Cache 是显存杀手。

PageAttention 原理：vLLM 的杀手锏

传统做法有个大问题。每个请求的 KV Cache 是连续分配的。但请求长度不一样，有的长有的短。这就导致显存碎片化严重。

我见过一个极端案例。一个 80GB 的 A100，跑着跑着显存还剩 20GB，但就是分配不出一个 1GB 的连续块。为什么？碎片太多。

vLLM 的 PageAttention 就是来解决这个问题的。

它的思路很简单：把 KV Cache 切成固定大小的“页”。每个页大小是固定的，比如 16 个 token 的 KV Cache。然后像操作系统的虚拟内存一样，用页表来管理。

这样做的好处：

  • 消除碎片：不要求连续内存，小页可以散落在各处
  • 按需分配：生成多少 token，分配多少页，不浪费
  • 共享方便：多个请求可以共享同一个页（比如前缀相同的情况）

说白了，就是让显存管理从“大块连续”变成“小块灵活”。

传统 KV Cache 分配（连续） 请求A 碎片 请求B 空闲（但碎片化） PageAttention 分配（分页） 页1 页2 页3 页4 空闲 页5 页6 空闲 空闲 页表（逻辑页 → 物理页） 逻辑页0 → 物理页1 逻辑页1 → 物理页2 逻辑页2 → 物理页4 逻辑页3 → 物理页6

我的经验：PageAttention 在长序列场景下效果特别明显。我曾经把一个 32K 序列的推理，显存占用从 40GB 降到了 28GB。碎片少了，吞吐自然上去了。

vLLM 的 PagedKV Cache 实现

vLLM 的实现细节，我挑几个关键点说说。

第一，页大小怎么定？

页大小一般是 16 或 32 个 token。太小了，页表开销大。太大了，又回到连续分配的老路。我个人习惯用 16，平衡性比较好。

第二，页表怎么管理？

每个请求有一个逻辑页表。生成新 token 时，先查页表，看当前页有没有空间。有就直接写。没有就分配一个新物理页。

第三，共享怎么实现？

多个请求如果前缀相同（比如系统提示词），可以共享前面的页。vLLM 用引用计数来管理。页被引用一次，计数加一。释放时减一。减到零才真正回收。

# 伪代码示意
class PagedKVBlock:
    def __init__(self, page_size=16):
        self.page_size = page_size
        self.blocks = {}  # 物理块号 -> 数据
        self.ref_counts = {}  # 物理块号 -> 引用计数

    def allocate(self):
        block_id = find_free_block()
        self.blocks[block_id] = torch.zeros(...)
        self.ref_counts[block_id] = 1
        return block_id

    def share(self, block_id):
        self.ref_counts[block_id] += 1

    def release(self, block_id):
        self.ref_counts[block_id] -= 1
        if self.ref_counts[block_id] == 0:
            del self.blocks[block_id]

注意：共享页虽然省显存，但要注意写时复制。如果一个请求要修改共享页的内容，必须先复制一份。否则会影响到其他请求。

TensorRT-LLM 中的 KV Cache 管理策略

TensorRT-LLM 的做法和 vLLM 不太一样。它更注重确定性。

策略一：预分配 + 池化

TensorRT-LLM 会在初始化时，一次性分配好所有可能的 KV Cache 空间。然后做成一个池子。推理时从池子里取，用完归还。

这样做的好处是：没有运行时分配开销，性能更稳定。坏处是：灵活性差一点，最大序列长度必须提前定好。

策略二：支持多种缓存模式

TensorRT-LLM 提供了几种模式：

  • 连续模式：传统方式，适合短序列
  • 分页模式：类似 vLLM，适合长序列
  • 混合模式：前几层用连续，后几层用分页。我试过，在某些场景下能省 10% 显存

策略三：缓存压缩

TensorRT-LLM 支持对 KV Cache 做量化。比如从 FP16 降到 INT8。显存直接减半。代价是精度略有损失。

我曾经在一个线上服务里，把 KV Cache 从 FP16 量化到 INT8。显存占用从 32GB 降到了 16GB。精度损失呢？BLEU 分数掉了 0.3。对于大多数场景，这个 trade-off 是值得的。

缓存模式 适用场景 显存效率 性能
连续模式 短序列（<1024） 低（有碎片） 高
分页模式 长序列（>2048） 高 中
混合模式 中等长度 中 高

我的建议：如果你用 TensorRT-LLM，优先试试分页模式。如果序列长度变化不大，连续模式也够用。量化是最后的手段，毕竟精度损失不可逆。

好了，KV Cache 这块就讲到这里。记住一句话：KV Cache 是显存大户，但也是优化空间最大的地方。把这块吃透了，你的推理服务性能能上一个台阶。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →