Paged Attention：大模型推理显存优化的核心方案

在本文中，我们将学习 Paged Attention（分页注意力）。这是一种解决 KV Cache 内存浪费问题的技术，能够让大语言模型（LLM）同时为更多用户提供服务。

我们将首先快速回顾 KV Cache，了解它产生的内存问题，通过示例观察 传统方法 如何浪费空间，然后深入探讨 Paged Attention 如何借鉴计算机管理内存的思想来解决这一问题。

KV Cache 快速回顾

当 LLM 生成文本时，它一次预测一个 Token。Token 是文本的一个微小单元，可以是一个单词、单词的一部分，甚至是单个字符。

在生成的每一步中，模型都会为新 Token 计算 Key（键） 和 Value（值），并且需要之前所有 Token 的 Key 和 Value。

如果没有 KV Cache，模型在每一步都需要重新计算之前所有 Token 的 Key 和 Value，这会造成计算资源的极大浪费。KV Cache 通过存储每个 Token 计算后的 Key 和 Value 来解决这个问题，以便在后续步骤中重复使用。

KV Cache 显著提高了文本生成的速度，但也带来了一个权衡：它需要消耗内存来存储所有的 Key 和 Value。

核心问题：KV Cache 中的内存浪费

当用户向 LLM 发送请求时，系统需要为该请求的 KV Cache 预留内存。但挑战在于：系统预先并不知道响应会有多长。

用户可能问一个简单的问题，得到 20 个 Token 的回答；也可能问一个复杂的问题，得到 2,000 个 Token 的回答。系统无法提前预知。

在传统方法中，系统会为每个请求预留一个巨大的连续内存块，足以处理最大可能的响应长度。这意味着即使实际响应只有 50 个 Token，系统也会为（例如）2,048 个 Token 预留内存。

旅馆类比

想象一家旅馆，客人在入住时不知道要住多少晚。有的住 1 晚，有的住 30 晚。

• 传统做法：当客人入住时，旅馆会在同一层楼连续预留 30 间房，以防客人住满 30 晚。即使客人只住了 2 晚，这 30 间房也会一直被占用，不能给其他人住。
• 后果：如果有 10 位客人入住，旅馆就要连续预留 300 间房。如果大多数人只住几晚，就会有数百间房空置浪费。

这正是传统 KV Cache 内存分配中发生的情况。

内存浪费的两种形式

1. 内部浪费（Internal Waste）：内存已预留但从未被使用。如果请求只生成了 50 个 Token，但预留了 2,048 个 Token 的空间，那么剩余 1,998 个 Token 的空间就被浪费了。
2. 外部浪费（External Waste）：由于内存必须分配在连续的大块中，散落在各处的微小空闲内存无法被利用，即使它们加起来足够服务一个新请求。

这种内存浪费是一个严重的问题。由于内存有限，这种浪费意味着系统同时服务的用户数量会大大减少。如果能更有效地利用内存，系统就能同时处理更多的请求。

这就是 Paged Attention 要解决的问题。

什么是 Paged Attention？

Paged Attention 是一种通过将 KV Cache 内存分解为小的、固定大小的块（称为“页”，Pages）来更有效地管理内存的技术。这些页面在内存中不需要是连续的。

“Paged”一词源于操作系统中的概念。操作系统为计算机上运行的所有程序管理内存，它不会给每个程序分配一整块连续内存，而是将内存分成固定大小的页。每个程序根据需要获取页面，这些页面可以散落在内存的任何地方。操作系统使用页表（Page Table）来记录每个页面的实际位置。

Paged Attention 将这一思想应用到了 KV Cache 中。

旅馆类比：分页方法

• 分页做法：旅馆不再为每位客人连续预留 30 间房，而是根据需要逐间提供。客人入住时给 1 间房，需要续住时再给 1 间（只要有空房即可），新房间不需要在旧房间隔壁。旅馆会维护一个清单，记录哪些房间属于哪位客人。
• 结果：如果客人只住 2 晚，就只占用 2 间房。剩余房间可以给其他客人。没有浪费，同样的房间数可以服务更多的客人。

Paged Attention 的工作原理

在 Paged Attention 中，KV Cache 内存被划分为固定大小的块。每个块可以存储固定数量 Token 的 Key 和 Value。例如，如果块大小为 4，则每个块可以存储 4 个 Token 的 KV 数据。

让我们通过一个例子来理解。假设模型正在生成响应：“I love teaching AI and Machine Learning at Outcome School”。

传统方法（无 Paged Attention）

系统预先为 16 个 Token 预留一整块连续内存：

内存状态: [__ __ __ __ __ __ __ __ __ __ __ __ __ __ __ __]
         (为该请求预留了 16 个槽位)

随着 Token 的生成，槽位被逐一填满：

内存状态: [I  love  teaching  AI  and  Machine  Learning  at  Outcome  School  __  __  __  __  __  __]
         (已使用: 10 个 Token                                          浪费: 6 个槽位)

剩余的 6 个槽位被预留但从未被使用，造成了内存浪费。

Paged Attention 方法

系统仅在需要时分配 4 个 Token 大小的微小块：

1. 第一步：开始生成。系统从内存中任意可用位置分配 Block 1。

   Block 1: [I  love  teaching  AI]

2. 第二步：Block 1 已满。系统分配 Block 2，它不需要紧邻 Block 1。

   Block 1: [I  love  teaching  AI]
   Block 2: [and  Machine  Learning  at]

3. 第三步：Block 2 已满。系统分配 Block 3。

   Block 1: [I  love  teaching  AI]
   Block 2: [and  Machine  Learning  at]
   Block 3: [Outcome  School  __  __]

在这种情况下，Block 3 仅浪费了 2 个槽位，而传统方法浪费了 6 个。更重要的是，系统没有预先预留任何可能用不到的内存。

块表（Block Table）

如果块散落在内存各处，模型如何找到它们？

系统使用块表（Block Table）记录哪些块属于哪个请求，就像操作系统的页表一样：

请求 1 的块表:
  Block 0 → 内存位置 5
  Block 1 → 内存位置 12
  Block 2 → 内存位置 3

块按顺序编号（0, 1, 2），每个编号指向其在内存中的实际位置。当模型需要读取之前 Token 的 Key 和 Value 时，它会通过块表找到正确的物理位置。

为什么 Paged Attention 如此高效

Paged Attention 解决了前文提到的两种内存浪费：

1. 内部浪费几乎被消除：内存是按需分配的。唯一的浪费只可能出现在最后一个块中（未完全填满）。在我们的例子中，最后一个块只浪费了 2/4 的空间，而传统方法浪费了 6/16。
2. 外部浪费被消除：因为块不需要连续，内存中任何位置的空闲块都可以被利用，不存在无法使用的“碎片间隙”。

结论：Paged Attention 极大地提高了内存利用率。这意味着在相同的内存容量下，系统可以同时处理显著增多的请求。

跨请求的内存共享

Paged Attention 还实现了一个重要的优化：内存共享。

假设两个用户问了同一个问题。对于这两个请求，输入 Token（问题部分）是相同的。在 Paged Attention 下，这两个请求可以共享输入 Token 的相同内存块，因为这些 Token 的 Key 和 Value 是完全一致的。

系统不需要存储两份副本，而是只存一份，让两个请求的块表都指向相同的物理块。这进一步节省了内存。

这种机制在以下场景特别有用：

• 并行采样（Parallel Sampling）：针对同一个问题生成多个不同的回答。由于问题相同，输入块可以在所有响应生成过程中共享。
• 束搜索（Beam Search）：模型同时探索多个可能的生成路径。共享的输入部分只需存储一次。

通过这种方式，Paged Attention 不仅消除了浪费，还通过共享机制进一步降低了总内存需求。

• 原文链接： x.com/amitiitbhu/status/...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～