深度解构 LLM 推理优化：从 KV Cache 冗余到 PagedAttention 内存分页映射

缓解AI焦虑
0 阅读3分钟

在 LLM 生产级部署中，推理性能的瓶颈已从算力受限（Compute-bound）转向显存带宽与容量受限（Memory-bound） 。对于 Llama-3、Mixtral 等长上下文模型，KV Cache 的管理效率直接决定了系统的吞吐量（Throughput）与延迟（Latency）。本文将从底层原理出发，剖析 KV Cache 的物理开销与 PagedAttention 的工程实现。

一、 KV Cache：自回归生成的“状态机”

image-20260307221716073

Transformer 推理的 Prefill 阶段与 Decoding 阶段存在计算不对称性。在 Decoding 阶段，模型呈现典型的自回归特性，KV Cache 实际上充当了模型的“时间状态机”。

1.1 显存占用的精确量化

以 Llama-3-8B (BF16/FP16) 为例，由于其采用了 GQA (Grouped-Query Attention) ，Key/Value 矩阵的头数（8）少于 Query 头数（32）。

每个 Token 占用的物理显存计算如下：

Memory∗pertoken=2×Layer×Head∗kv×Dim∗head×Bytes∗precisionMemory *{per_token} = 2 \times Layer \times Head*{kv} \times Dim *{head} \times Bytes*{precision}

MemoryL3−8B=2×32×8×128×2=131,072 Bytes≈128 KBMemory_{L3-8B} = 2 \times 32 \times 8 \times 128 \times 2 = 131,072 \text{ Bytes} \approx 128 \text{ KB}

1.2 传统管理的失效：内存碎片（Fragmentation）

在原生 PyTorch 或 FasterTransformer 中，KV Cache 通常采用预分配连续显存。

  • Internal Fragmentation: 为了支持最大长度 LL，系统必须为每个 Request 预留 LL 的空间，但实际生成长度往往远小于 LL，导致空间空置。
  • Reservation Gap: 在请求处理过程中，由于无法预测输出长度，系统必须始终锁定剩余的最大可能空间，导致显存无法被其他 Batch 共享。

二、 PagedAttention：借鉴 OS 内核的非连续内存映射

image-20260307221747550

vLLM 提出的 PagedAttention 核心突破在于将虚拟内存的思想引入 GPU 显存管理，实现了逻辑 Token 序列与物理显存块（Physical Blocks）的解耦。

2.1 核心数据结构：页表（Page Table）

PagedAttention 将 KV Cache 划分为固定大小的 Block（通常 Block Size=16）。

  • Logical Blocks: 模型感知的 KV 序列，逻辑上是连续的。
  • Physical Blocks: 显存中离散分布的存储块。
  • Block Table: 维护逻辑块到物理块的映射关系。

2.2 动态按需分配（Dynamic Allocation）

在生成过程中，当一个 Block 被填满后，内核才会从全局 Free Block Pool 中申请下一个物理块。这种机制彻底消除了预留空间的浪费，将显存利用率从 ~60% 榨取至 96% 以上。

2.3 工程细节：Copy-on-Write (CoW) 与并行采样

PagedAttention 天然支持高效的并行采样（Parallel Sampling）和束搜索（Beam Search）。多个请求可以共享同一个 Prompt 的物理块，仅在需要写入不同的新 Token 时触发 CoW 机制，极大地降低了多样化输出时的内存冗余。

三、 性能实测与部署权衡

3.1 吞吐量（Throughput）与有效并发

在相同的硬件环境（如单机 8×H100）下，PagedAttention 允许显著提高 Effective Batch Size。

  • Continuous Batching: 配合分页机制，vLLM 可以在同一 Batch 中处理不同阶段（Prefill/Decoding）的请求，消除了传统 Batch 必须等待最长请求结束的“木桶效应”。

3.2 真实硬件约束下的极限

针对 RTX 4090 (24GB) 部署 Llama-3-8B 的技术参数：

  • Static Memory: 模型权重占 ~15GB。
  • Dynamic KV Cache: 剩余 ~9GB。
  • Capacity: 在 FP16 精度下可承载约 7.2 万 Token；若开启 FP8 KV Cache Quantization，则可翻倍至 14.4 万 Token。
  • 注： 所谓“百万上下文”必须依赖多卡张量并行（TP）或算力受限的 CPU Offloading，单卡 24GB 无法突破物理定律。

四、 技术演进与未来路径

image-20260307221857807

  1. KV Cache 量化: 引入 INT8/FP8 甚至更低精度的 KV 量化，配合 PagedAttention 进一步压缩单元空间。
  2. 块大小（Block Size）调优: Block Size 过大会增加内部碎片，过小会增加 GPU Kernel 调度开销，通常 16 是性能与碎片的平衡点。
  3. 计算下沉: 结合 FlashAttention-3 等算子级优化，减少数据在 HBM 与 SRAM 之间的搬运延迟。

总结： KV Cache 解决了 O(n^2) 的计算复杂度问题，而 PagedAttention 解决了 O(n) 的空间管理失控问题。对于架构师而言，理解分页映射不仅是为了部署 vLLM，更是为了在自定义推理链中实现更高效的算力调度。

avatar
文章
阅读
粉丝