KV Cache 与 PagedAttention：为什么你的 LLM 服务器会变慢

你的 LLM 服务器运行缓慢。

你在四块 A100 GPU 上部署了一个 70B 模型。早上 8 点一切正常。到了午饭时间，延迟翻了一倍。你检查了内存，发现大部分都被“张量缓冲区”（tensor buffers）占用了。这些实际上是旧对话的缓存状态。

这就是 KV cache 问题。它是生产环境 LLM 推理服务中最大的瓶颈。

什么是 KV cache？

每个 Transformer 模型都是逐个生成 token 的。为了生成一个新的 token，模型需要之前所有 token 的 Key 和 Value 张量。每次都重新计算这些内容速度太慢了。因此，推理引擎会将它们存储起来。这种存储方式就是 KV cache。

内存问题：

对于 Llama 3.1 70B 模型，单个 4096-token 的序列大约需要 1.3 GB 内存。

如果你同时有 256 个用户，则需要 336 GB 内存。这已经超过了四块 A100 GPU 的容量。KV cache 的增长速度非常快，以至于它占用的内存往往比模型权重本身还要多。

传统的内存管理之所以失败，是因为：

  • 内部碎片（Internal fragmentation）：你为 4096 个 token 分配了空间，但实际只用了 300 个。你浪费了 93% 的空间。
  • 无法共享：两个使用相同系统提示词（system prompt）的用户会各自存储该提示词的副本。
  • 全有或全无的逐出机制（All-or-nothing eviction）：当内存耗尽时，你必须将整个序列移动到 CPU。这会导致 GPU 停顿。

PagedAttention 如何解决这个问题：

PagedAttention 的工作方式类似于操作系统。它将 KV cache 划分为称为“页”（pages）的小型固定大小块。

这解决了三个主要问题：

  • 按需分配：序列随着其增长才占用相应的页。你不会在未使用的容量上浪费内存。
  • 支持共享前缀：多个用户可以为共同的系统提示词共享相同的物理页。这利用了“写时复制”（copy-on-write）逻辑来节省大量内存。
  • 细粒度逐出：当内存满时，系统会将小的页移动到 CPU，而不是移动庞大的整个序列。

结果：

使用 PagedAttention（vLLM 内部采用的技术）与传统方法相比，吞吐量可以提高 2 到 4 倍。

适用场景：

  • 高并发。
  • 不同长度的序列。
  • 具有相同开头（前缀）的提示词。

不适用场景：

  • 单用户本地推理。
  • 非常小的模型。
  • 每个序列长度完全相同的任务。

KV Cache 与 PagedAttention：它们的作用以及为什么它们如此重要

在大语言模型（LLM）的推理过程中，速度和效率是开发者面临的核心挑战。如果你深入研究过 LLM 的架构，你可能已经听说过 KV Cache 和 PagedAttention。这两个概念对于实现高效、高吞吐量的模型推理至关重要。

本文将深入探讨这两个概念，解释它们的工作原理，以及为什么它们对现代 AI 基础设施如此重要。

什么是 KV Cache？

要理解 KV Cache，首先需要了解 LLM 是如何生成文本的。

LLM（如 GPT-4 或 Llama 3）是**自回归（Autoregressive）**模型。这意味着它们在生成文本时，是一个 token 接一个 token 地生成。在生成第 $n$ 个 token 时，模型需要查看之前生成的所有 $n-1$ 个 token。

注意力机制的计算

在 Transformer 架构中，核心是注意力机制（Attention Mechanism）。对于每一个 token，模型会计算三个向量：

  • Query (Q)：代表当前 token “想要寻找什么”。
  • Key (K)：代表当前 token “能提供什么信息”。
  • Value (V)：代表当前 token “包含的具体信息”。

为了计算当前 token 与之前所有 token 之间的注意力，模型需要执行以下操作：

  1. 计算当前 token 的 $Q$。
  2. 获取之前所有 token 的 $K$ 和 $V$。
  3. 计算 $Q$ 与所有 $K$ 的点积，得到注意力权重。
  4. 使用这些权重对所有的 $V$ 进行加权求和。

冗余计算的问题

如果没有 KV Cache，在生成每一个新 token 时，模型都必须重新计算之前所有 token 的 $K$ 和 $V$ 向量。这会导致大量的重复计算，随着序列变长，计算量会呈平方级增长，导致推理速度极慢。

KV Cache 的解决方案： KV Cache 通过在内存中存储之前生成的 token 的 $K$ 和 $V$ 向量，避免了重复计算。在生成下一个 token 时，我们只需要计算当前 token 的 $Q, K, V$，然后直接从缓存中读取之前的 $K$ 和 $V$ 即可。

这极大地提高了推理速度，但也带来了一个新的挑战：显存（VRAM）占用。

问题所在：内存碎片

随着序列长度的增加，KV Cache 所需的显存会迅速膨胀。在传统的推理实现中，通常会为每个请求预先分配一块连续的显存空间，其大小基于模型可能达到的最大上下文长度（Max Context Length）。

这导致了两个严重的问题：

  1. 内部碎片（Internal Fragmentation）：如果一个请求只生成了 10 个 token，但我们为它预分配了 2048 个 token 的空间，那么剩下的 2038 个 token 的空间就被浪费了，无法给其他请求使用。
  2. 外部碎片（External Fragmentation）：由于显存被预分配成不连续的大块，即使总剩余显存足够，也可能因为找不到足够大的连续空间而无法处理新请求。

这种低效的内存管理限制了系统的吞吐量（Throughput），即单位时间内系统能处理的请求数量。

什么是 PagedAttention？

为了解决 KV Cache 带来的内存管理难题，vLLM 团队提出了 PagedAttention。

PagedAttention 的灵感来源于操作系统中的**虚拟内存（Virtual Memory）和分页（Paging）**机制。

工作原理

在操作系统中，物理内存被划分为固定大小的“页”（Pages），而进程看到的内存是连续的“虚拟地址空间”。通过页表（Page Table），操作系统可以将不连续的物理页映射到连续的虚拟地址上。

PagedAttention 将这一思想应用到了 KV Cache 的管理中：

  1. 分块存储（Block-based Storage）：不再为每个请求分配一整块连续的显存，而是将 KV Cache 划分为固定大小的块（Blocks）。
  2. 非连续分配：这些块可以存储在显存中任何不连续的位置。
  3. 块表（Block Table）：系统维护一个块表，记录每个请求的 KV 缓存分布在哪些物理块中。当模型需要访问某个 token 的 KV 值时，通过块表找到对应的物理块。

PagedAttention 的优势

  • 几乎零碎片：由于块的大小是固定的且可以动态分配，内部碎片被限制在最后一个块中（非常小），几乎消除了外部碎片。
  • 极高的显存利用率：由于不再需要预分配最大长度的空间，系统可以更高效地利用显存，从而在同一块 GPU 上运行更多的并发请求。
  • 大幅提升吞吐量：更高的显存利用率意味着可以同时处理更多的请求，从而显著提升了系统的整体吞吐量。
  • 支持高效的共享内存：PagedAttention 使得多个请求可以轻松共享相同的 KV 块（例如在并行采样或 Beam Search 时），进一步节省了空间。

总结

  • KV Cache 是通过存储已计算的键值对来避免重复计算、加速推理的关键技术。
  • 传统的 KV Cache 管理 存在严重的内存碎片问题，限制了显存利用率和吞吐量。
  • PagedAttention 通过借鉴操作系统的分页机制，实现了非连续的、动态的 KV Cache 管理，彻底解决了内存碎片问题，是实现高吞吐量 LLM 推理的核心技术之一。

理解这两者对于理解现代大模型推理引擎（如 vLLM, TensorRT-LLM）的工作原理至关重要。

Source: https://dev.to/tech_nuggets/kv-cache-and-pagedattention-what-they-do-and-why-they-matter-jce

Optional learning community: https://t.me/GyaanSetuAi