KV Cache与PagedAttention深度解析：vLLM内存管理革命

一、KV Cache基础与内存瓶颈

1.1 自回归生成的内存代价

大语言模型的自回归生成（Autoregressive Generation）有一个根本性的效率问题：每生成一个新token，都需要重新计算之前所有token的Key（键）和Value（值）向量。对于一个长度为N的序列，生成第N+1个token时，Transformer的Attention计算需要所有历史token的KV对。

这个问题的直观理解是：假设你在读一本书，每读新的一页，都要从头重新读一遍之前的所有页——这显然是荒谬的。KV Cache（键值缓存）就是解决这个问题：把每页的"理解结果"（KV向量）保存下来，后面直接复用。

1.2 KV Cache的内存占用公式

KV Cache的内存占用可以用一个精确的公式描述。对于一个包含L层、隐藏维度d、注意力头数h的模型，处理一个长度为S的序列时：

单请求KV Cache内存占用（字节）：

每个token的KV：
  K: [h, d/h, S] → 每层每个头需要 d/h × S 个参数存储K
  V: [h, d/h, S] → 每层每个头需要 d/h × S 个参数存储V
  
单层KV = 2 × h × (d/h) × S = 2 × d × S 个参数
L层总KV = L × 2 × d × S 个参数

以FP16（2字节）存储：
Memory = L × 2 × d × S × 2 bytes

具体示例（LLaMA-2 70B）：
  L = 80层, d = 8192, S = 4096（序列长度）
  Memory = 80 × 2 × 8192 × 4096 × 2 bytes
         = 80 × 2 × 8192 × 4096 × 2
         = 10,995,116,800 bytes ≈ 10.2 GB

一个请求就占10GB！而且这是静态的——如果S=8192，内存直接翻倍到20GB。

1.3 生产环境的内存爆炸问题

上面的计算还只是单个请求。在生产环境中，推理服务需要同时处理多个并发请求（例如batch_size=32），而且每个请求的序列长度动态变化。这带来两个关键问题：

二、Transformer中的KV Cache机制

2.1 Attention计算回顾

要理解KV Cache，必须先理解Transformer的Attention计算。标准的Scaled Dot-Product Attention计算公式为：

Attention(Q, K, V) = softmax(QK^T / √d_k) × V

其中：
  Q (Query): 当前token的查询向量，形状 [1, h, d/h]
  K (Key): 所有历史token的键向量，形状 [S, h, d/h]
  V (Value): 所有历史token的值向量，形状 [S, h, d/h]
  d_k = d/h: 每个头的维度

对于自回归生成：
  生成第t个token时，需要 K[1:t] 和 V[1:t]
  即：所有历史token的KV都必须参与当前计算

2.2 KV Cache的工作原理

KV Cache的核心思想是：在计算完某个token的K和V之后，将其保存下来，后续生成新token时直接读取，而不是重新计算。具体流程如下：

传统方式（无KV Cache）：
  生成token t时：
    for i in 1 to t:
        K_i = W_K(x_i)  # 重新计算所有历史K
        V_i = W_V(x_i)  # 重新计算所有历史V
    Attention(Q_t, K[1:t], V[1:t])
  
  时间复杂度：O(t²)，每个新token都要重新计算所有历史

KV Cache方式：
  生成token t时：
    K_t = W_K(x_t)  # 只计算当前token的K
    V_t = W_V(x_t)  # 只计算当前token的V
    将 K_t, V_t 追加到缓存
    Attention(Q_t, K[1:t], V[1:t])  # 直接从缓存读取历史KV
  
  时间复杂度：O(t)，只需计算当前token的KV

2.3 多请求并发的内存管理挑战

在生产推理服务中，往往同时处理多个请求（batch inference）。每个请求都有自己的KV Cache，而且长度不同、增长动态。传统的内存管理方式是为每个请求预分配一块连续的最大长度内存（例如4096 token），这导致：

可以看出，传统方式下平均内存利用率不到50%，大量内存被浪费。更严重的是，当请求数增加时，即使GPU显存还有剩余，也可能因为"内存碎片"和"连续内存要求"而无法分配新的KV Cache。

三、内存碎片与GPTQ的教训

3.1 内存碎片的形成机制

内存碎片（Memory Fragmentation）是动态内存分配中的经典问题，在KV Cache场景下尤为严重。碎片分为两种：

内存碎片类型：

1. 内部碎片（Internal Fragmentation）：
   分配的内存块 > 实际使用量
   例：预分配4096 token，实际使用128 token → 浪费3968 token的空间
   
2. 外部碎片（External Fragmentation）：
   空闲内存总量足够，但分散为多个小块，无法满足大块分配请求
   
   GPU显存布局示例：
   | 已用1 | 空闲16MB | 已用2 | 空闲8MB | 已用3 | 空闲16MB |
   
   此时若需要分配32MB的连续空间，虽然总空闲40MB，但无法分配！

3.2 传统推理框架的内存管理

以FasterTransformer、TensorRT-LLM为代表的传统推理框架，采用静态预分配策略：在请求开始时，为整个序列（最大长度）预分配连续的KV Cache内存。这种方式的弊端在生产环境中被放大：

3.3 GPTQ的量化与KV Cache矛盾

GPTQ（Generative Pre-trained Transformer Quantization）是一种流行的权重量化方法，将模型权重压缩到4-bit或8-bit。但在KV Cache场景下，GPTQ引入了一个新的矛盾：

部分的解决方案是：KV Cache使用FP16存储（保证精度），但对Attention计算本身进行量化（如INT8 Attention）。这需要在PagedAttention的实现中特殊处理量化张量，是vLLM等框架的进阶优化方向。

四、PagedAttention：操作系统思想的引入

4.1 核心思想：将操作系统虚拟内存应用到Attention

PagedAttention的革命性在于：它把操作系统管理内存的核心思想——分页（Paging）和页表（Page Table）——应用到了Attention的KV Cache管理中。这个类比非常精妙：

操作系统虚拟内存 vs PagedAttention 类比：

操作系统虚拟内存：
  虚拟地址空间 → 逻辑上的连续地址
  物理内存 → 实际的硬件内存（可能不连续）
  页表（Page Table） → 映射虚拟页到物理页
  页面大小（Page Size） → 固定大小的内存块（如4KB）

PagedAttention：
  逻辑token序列 → 逻辑上的连续token序列
  KV Cache物理存储 → GPU显存中的实际KV存储（可能不连续）
  Block Table（块表） → 映射逻辑token位置到物理KV块
  Block Size（块大小） → 固定大小的KV块（如16个token）

4.2 分页机制详解

PagedAttention将每个请求的KV Cache划分为固定大小的"块"（Block），每个块存储固定数量的token的KV（例如16个token）。块在物理显存中不需要连续，通过一张"块表"（Block Table）来记录逻辑token位置到物理块的映射。

PagedAttention分块示例：

逻辑token序列（长度=100）：
  逻辑位置: [0, 1, 2, ..., 99]
  
分块（Block Size=16）：
  Block 0: token [0-15]   → 物理块 P7
  Block 1: token [16-31]  → 物理块 P2
  Block 2: token [32-47]  → 物理块 P9
  Block 3: token [48-63]  → 物理块 P1
  Block 4: token [64-79]  → 物理块 P5
  Block 5: token [80-95]  → 物理块 P8
  Block 6: token [96-99]  → 物理块 P3（部分填充）

块表（Block Table）：
  逻辑块索引 → 物理块索引
  [0, 1, 2, 3, 4, 5, 6] → [7, 2, 9, 1, 5, 8, 3]

注意：物理块 P7, P2, P9, P1, P5, P8, P3 在显存中不需要连续！
它们可以分散在显存的任何位置，由块表维护映射关系。

4.3 非连续内存访问的Attention计算

传统Attention要求KV在内存中连续，但PagedAttention打破了这个限制。在Attention计算时，通过块表"拼接"不同物理块中的KV，对上层提供一个逻辑上连续的KV视图。这需要在GPU内核（Kernel）层面进行修改：

PagedAttention内核伪代码：

def paged_attention(
    Q,                  # Query [batch, num_heads, 1, head_dim]
    block_table,        # 块表 [batch, max_blocks_per_seq]
    kv_cache,           # KV Cache物理存储 [num_blocks, 2, num_heads, block_size, head_dim]
    context_len         # 当前序列的实际长度
):
    output = zeros_like(Q)
    num_blocks = ceil(context_len / block_size)
    
    for i in range(num_blocks):
        # 通过块表找到第i个逻辑块的物理块索引
        physical_block_idx = block_table[batch_idx][i]
        
        # 从物理块中读取KV
        K_block = kv_cache[physical_block_idx][0]  # K
        V_block = kv_cache[physical_block_idx][1]  # V
        
        # 计算这一块的Attention分数
        attn_scores = Q @ K_block.transpose()
        attn_weights = softmax(attn_scores / sqrt(head_dim))
        output += attn_weights @ V_block
    
    return output

这个设计的精妙之处在于：物理块的访问通过块表间接寻址，完全解耦了逻辑连续性和物理连续性。这使得KV Cache的内存管理变得和操作系统管理内存一样灵活。

五、vLLM架构设计：Page Table与Block管理

5.1 vLLM的整体架构

vLLM（Virtual LLM）是一个专门为高吞吐量LLM推理设计的框架，其核心就是PagedAttention。vLLM的架构分为几个关键组件：

vLLM架构组件：

┌─────────────────────────────────────────────────────┐
│                    vLLM 推理引擎                      │
├─────────────────────────────────────────────────────┤
│  请求调度器（Scheduler）                            │
│  - 管理待处理请求队列                                │
│  - 决定哪些请求可以被加入当前batch                    │
│  - 分配/释放KV Cache块                              │
├─────────────────────────────────────────────────────┤
│  Block Manager（块管理器）                           │
│  - 维护GPU显存中的空闲块池（Free Block Pool）         │
│  - 分配物理块给新请求                                │
│  - 回收已完成请求的块                                │
│  - 支持Copy-on-Write（写时复制）共享                 │
├─────────────────────────────────────────────────────┤
│  Model Executor（模型执行器）                        │
│  - 执行Transformer前向计算                          │
│  - 调用PagedAttention内核                            │
│  - 管理模型权重和中间激活                            │
├─────────────────────────────────────────────────────┤
│  KV Cache（GPU显存）                                │
│  - 物理块存储 [num_blocks, 2, num_heads,            │
│    block_size, head_dim]                            │
│  - 每个块存储固定数量token的KV                       │
└─────────────────────────────────────────────────────┘

5.2 Block Manager的块分配策略

Block Manager是vLLM内存管理的核心，它维护一个空闲块池，按需分配给请求。块分配遵循以下策略：

5.3 块表的维护与更新

每个请求都维护自己的块表（Block Table），记录逻辑块到物理块的映射。当请求生成新token时，块表需要动态扩展：

块表更新流程（生成第t个token时）：

当前序列长度 = t
需要的块数 = ceil((t + 1) / block_size)

if 当前块表大小 < 需要的块数:
    # 需要分配新的物理块
    new_physical_block = block_manager.allocate_block()
    块表.append(new_physical_block)
    
    # 将新块的引用计数+1
    physical_block.ref_count += 1

# 将新token的KV写入最后一个块
last_block = 块表[-1]
offset_in_block = t % block_size
kv_cache[last_block][offset_in_block] = (K_t, V_t)

注意：块表存储在CPU内存中（因为需要频繁修改），
而KV Cache本身在GPU显存中。
每次Attention计算时，需要将块表复制到GPU。

六、连续批处理与内存碎片消除

6.1 传统静态批处理的局限

传统的推理批处理（Batching）是静态的：一个batch中的所有请求必须同步进行，当一个请求完成后，整个batch需要等待，或者该位置被填充一个"假请求"。这导致GPU利用率低下：

静态批处理的问题：

Batch size = 4，包含4个请求：
  请求A: 生成了128 tokens → 完成
  请求B: 生成了256 tokens → 进行中
  请求C: 生成了64 tokens  → 进行中
  请求D: 生成了512 tokens → 进行中

此时batch中仍有3个活跃请求，但slot A（请求A的位置）被浪费了。
即使有新的请求E等待处理，也无法加入这个batch，
因为静态batch要求batch_size固定！

结果：GPU计算资源浪费约25%（1/4的slot空闲）

6.2 连续批处理（Continuous Batching）

连续批处理（也称为动态批处理）解决了上述问题：每次迭代（生成一个token），都重新决定哪些请求应该包含在当前batch中。完成的请求立即被移除，新的请求可以立即加入。

连续批处理流程：

while 有待处理或活跃请求:
    # 1. 调度器决定当前batch包含哪些请求
    current_batch = scheduler.schedule(
        waiting_requests,      # 等待中的请求
        running_requests,      # 正在运行的请求
        max_batch_size=32      # 最大batch大小
    )
    
    # 2. 为current_batch中的每个请求分配KV Cache块（如果需要）
    for req in current_batch:
        if req.need_new_block():
            block = block_manager.allocate_block()
            req.block_table.append(block)
    
    # 3. 执行一次模型前向计算（生成1个token）
    outputs = model_executor.execute(current_batch)
    
    # 4. 处理输出，更新请求状态
    for req, output in zip(current_batch, outputs):
        req.append_token(output)
        if req.is_finished():
            scheduler.mark_finished(req)
            block_manager.free_blocks(req.block_table)
    
    # 5. 下一个迭代可以加入新请求！

6.3 PagedAttention + 连续批处理的协同效应

PagedAttention和连续批处理是天然的搭档：PagedAttention解决了KV Cache内存不连续的问题，使得不同长度的请求可以灵活共享GPU显存；连续批处理则充分利用了这种灵活性，实现接近100%的GPU利用率。

七、与FlashAttention的协同优化

7.1 FlashAttention回顾

FlashAttention是另一个革命性的Attention优化，它通过优化GPU内存访问模式（减少HBM（High Bandwidth Memory，高带宽内存）与SRAM之间的数据传输），实现了2-4倍的Attention计算加速。FlashAttention的核心思想是：

FlashAttention优化原理：

传统Attention:
  Q, K, V 存储在HBM（显存，带宽~1TB/s）
  每次计算需要多次读写HBM → 内存带宽瓶颈
  
FlashAttention:
  将Q, K, V分块（tiling），每块放入SRAM（片上缓存，带宽~19TB/s）
  在SRAM内完成Attention计算
  只将最终结果写回HBM
  
  关键技巧：Online Softmax（在线Softmax计算）
  - 不需要一次性计算所有Attention分数
  - 分块计算，维护运行中的max和sum
  - 避免存储完整的[N, N] Attention矩阵

7.2 PagedAttention与FlashAttention的冲突

理论上，PagedAttention和FlashAttention似乎是冲突的：FlashAttention要求Q/K/V在内存中连续（以便分块加载到SRAM），而PagedAttention故意将KV分散在非连续的块中。这个矛盾需要特殊的工程处理：

vLLM的解决方案是：实现专门的PagedFlashAttention内核，这个内核知道块表的存在，能够正确地从不同物理块中加载KV，同时仍然利用FlashAttention的内存访问优化。

PagedFlashAttention内核伪代码：

def paged_flash_attention(Q, block_table, kv_cache, context_len):
    # Q: [1, num_heads, 1, head_dim]，在SRAM中
    # 不需要将整个KV加载到SRAM，而是分块处理
    
    output = zeros_like(Q)
    m = -inf  # 运行中的max（用于数值稳定性）
    l = 0     # 运行中的sum（Softmax分母）
    
    num_blocks = ceil(context_len / block_size)
    for i in range(num_blocks):
        physical_block = block_table[i]
        K_block = kv_cache[physical_block, 0, :, :, :]  # 从HBM加载一个块
        V_block = kv_cache[physical_block, 1, :, :, :]
        
        # 在SRAM中计算这一块的Attention
        scores = Q @ K_block.transpose() / sqrt(head_dim)
        m_new = max(m, max(scores))
        alpha = exp(m - m_new)  # 修正之前的softmax值
        l = l * alpha + sum(exp(scores - m_new))
        
        output = output * alpha + (exp(scores - m_new) @ V_block)
        m = m_new
    
    return output / l  # 最终的Attention输出

7.3 实际性能收益

结合PagedAttention（内存管理优化）和FlashAttention（计算优化），vLLM实现了全方位的推理加速：

在生产环境中，这些优化叠加的效果是：相同的GPU（例如A100 80GB），vLLM可以处理的并发请求数是传统框架（如FasterTransformer）的3-5倍，且延迟更低。

八、深挖点：Prefix Caching与跨请求KV共享

8.1 问题的提出：重复的前缀

在实际生产中，许多请求共享相同的前缀（Prefix）。典型的场景包括：

• System Prompt（系统提示）：每个请求都以相同的系统提示开头（例如"你是一个有用的助手..."）
• 工具调用描述：Function Calling场景中，工具的描述schema被包含在每个请求中
• 多轮对话历史：同一用户的多条消息共享相同的历史对话

传统方式下，每个请求都独立存储这些共享前缀的KV Cache，造成大量冗余。以一个2KB的System Prompt为例，如果有100个并发请求，就需要200KB的重复KV存储——对于大模型，这轻松达到数GB的浪费。

8.2 Prefix Caching（前缀缓存）机制

Prefix Caching的核心思想：将共享前缀的KV Cache存储一份，多个请求通过"引用"共享这份KV。这类似于操作系统的写时复制（Copy-on-Write，COW）机制：

Prefix Caching实现原理：

1. 哈希前缀内容：
   prefix_hash = hash(system_prompt + tools_description)
   
2. 在全局前缀缓存（Global Prefix Cache）中查找：
   if prefix_hash in global_cache:
       # 缓存命中！重用已有的KV块
       shared_blocks = global_cache[prefix_hash]
       req.block_table.extend(shared_blocks)
       # 增加引用计数（防止被释放）
       for block in shared_blocks:
           block.ref_count += 1
   else:
       # 缓存未命中，正常计算并缓存
       compute_kv_for_prefix(req)
       global_cache[prefix_hash] = req.prefix_blocks
   
3. 当请求完成时：
   for block in req.block_table:
       block.ref_count -= 1
       if block.ref_count == 0:
           block_manager.free_block(block)

8.3 引用计数与块的生命周期

支持Prefix Caching需要一个健壮的引用计数（Reference Counting）机制。每个物理块维护一个引用计数，记录有多少请求正在使用这个块：