AI 陪读顶会论文（三）：KV Cache 显存占用计算与 PagedAttention 原理拆解

在当下的大语言模型（LLM）推理系统（Inference Serving）中，显存（VRAM）是极其昂贵的核心算力瓶颈。
在部署像 LLaMA-70B、Qwen-72B 或 DeepSeek 这种大模型时，很多同学以为显存主要被模型本身的权重参数（Weights）占满了。
但在高并发、长文本（如 32K/128K 上下文）的实际生产服务中，KV Cache（键值缓存）所吞噬的动态显存往往会超过模型权重本身的数倍，成为限制并发吞吐量（Throughput）的最致命元凶！

UC 伯克利团队发表在 SOSP 2023 顶会上的著名论文《Efficient Memory Management for Large Language Model Serving with PagedAttention》（vLLM 的基石），正是为了解决这一痛点。

今天我们借助大模型辅助，把 KV Cache 的显存数学账本与 PagedAttention 的虚拟内存分页哲学彻底推演一遍。

为什么自回归解码（Autoregressive Decoding）必须使用 KV Cache？

在 Transformer 解码阶段，大模型是逐个 Token（One-by-one）生成的。
在生成第 $t$ 个 Token 时，自注意力机制需要计算当前 Token 与前 $t-1$ 个历史 Token 的相关性：
$$\text{Attention} = \text{softmax}\left(\frac{q_t K_{1:t}^T}{\sqrt{d}}\right)V_{1:t}$$

如果没有缓存：

• 生成第 1 个 Token，计算 1 次；
• 生成第 1000 个 Token，需要将前 999 个历史 Token 的 $K$ 和 $V$ 矩阵全部重新算一遍！
• 时间复杂度直接退化为 $O(N^2)$。

KV Cache 的做法：
将前面所有历史生成的 Key 向量和 Value 向量保存在显存中。每次生成新 Token 时，只需要计算当前 Token 对应的单个 $q_t, k_t, v_t$，并将新算出的 $k_t, v_t$ 追加到显存缓存中，直接将每一步的计算开销降至 $O(1)$！

核心账本：单次请求 KV Cache 显存占用的精确数学公式

大厂面试必考算力账本：“一个 70B 的大模型，在 4096 上下文长度、Batch Size = 16 时，KV Cache 到底占用多少显存？”

计算公式如下：
$$\text{KV Cache 显存 (Bytes)} = 2 \times 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{seq_len} \times \text{batch_size} \times \text{dtype_bytes}$$

各个参数的物理含义：

1. 第一个 $2$：分别保存 Key 矩阵和 Value 矩阵；
2. $n_{\text{layers}}$：模型 Transformer 的总层数（如 80 层）；
3. $n_{\text{heads}} \times d_{\text{head}} = d_{\text{model}}$：注意力头数乘以头维度，即隐藏层维度（如 8192）；
4. $\text{seq_len}$：当前上下文总序列长度（如 4096）；
5. $\text{batch_size}$：并发请求批次大小（如 16）；
6. $\text{dtype_bytes}$：采用 FP16/BF16 存储时，每个浮点数占 2 字节。

实战测算（LLaMA-70B 真实参数）：

• $n_{\text{layers}} = 80$, $d_{\text{model}} = 8192$, $\text{seq_len} = 4096$, $\text{batch_size} = 16$, $\text{dtype} = 2$ 字节；
• 若采用经典多头注意力（MHA）：
  $$\text{显存} = 2 \times 80 \times 8192 \times 4096 \times 16 \times 2 \approx 171,798,691,840 \text{ 字节} \approx \mathbf{160 \text{ GB}!}$$
  160GB 仅仅是 KV Cache 的动态显存开销！两张顶级的 A100 (80GB) 甚至连一个 Batch 的缓存都装不下！

传统 KV Cache 显存分配的严重浪费：内部与外部碎片

在传统的 PyTorch 推理框架（如 HuggingFace Transformers）中，为了在 GPU 上进行高效的张量运算，系统必须为每个请求在显存中开辟一段连续的物理内存空间（Contiguous Memory）。

这引发了巨大的浪费：

1. 预分配浪费（Over-allocation）：因为不知道用户这次对话会生成多少个字，系统不得不按最大输出长度（如 2048）提前开辟显存。如果用户只回答了 50 个字，剩下的显存全被白白占着无法被其他请求使用；
2. 外部内存碎片（Memory Fragmentation）：请求在不同时刻开始与结束，导致连续显存空间被切得稀碎，系统空有 30% 的空闲显存，却因为找不出一块足够大的连续空间而无法接收新请求。

PagedAttention 的破局哲学：操作系统虚拟内存分页在 GPU 上的完美重现

vLLM 团队提出了 PagedAttention 算法，其核心思想直接借鉴了操作系统经典的虚拟内存分页管理（Virtual Memory Paging）：

graph TD
    A[逻辑连续的 KV Cache 序列 (Token 1 到 100)] --> B[逻辑块 Block 0: Token 1~16]
    A --> C[逻辑块 Block 1: Token 17~32]
    A --> D[逻辑块 Block 2: Token 33~48]

    B -->|页表 Block Table 映射| E[物理显存非连续物理块 Physical Block #42]
    C -->|页表 Block Table 映射| F[物理显存非连续物理块 Physical Block #105]
    D -->|页表 Block Table 映射| G[物理显存非连续物理块 Physical Block #12]

三大工程创新：

1. 非连续物理存储（Non-contiguous Storage）：
   将 KV Cache 划分为固定大小的“块（Block，如每个 Block 容纳 16 个 Token）”。物理上这些 Block 可以散落在显存的任意离散角落；
2. 动态按需分配（On-demand Allocation）：
   只有当模型真正生成了第 17 个 Token 时，才去显存中申请第 2 个 Block，内部内存碎片被彻底压缩至最后一个 Block 的微小空间（$\le 4%$）；
3. 极速零拷贝分支共享（Copy-on-Write for Parallel Sampling）：
   在进行 Beam Search 或多分支生成时，不同分支可以共享相同的 Prompt 物理块，只有发生分叉时才执行写时复制（CoW），显存节约达 55% 以上！

实习生的学术总结

从计算机操作系统（OS）的虚拟内存分页，到大模型推理引擎的 PagedAttention，计算机底层经典的基础设计思想（分页、按需分配、CoW、逻辑地址映射）在 AI 时代再次展现了其强大的生命力。
搞懂了 KV Cache 与分页内存，在面对大模型算力成本优化与系统架构设计时，你将拥有极为扎实的理论底座。