你有没有想过,为什么大模型生成第一个字需要几秒钟,但后续的字却快得多?为什么同样的模型,处理1000字的上下文比处理100字消耗更多显存?为什么有些70B参数的模型反而比7B模型的KV缓存更小?

答案都指向同一个技术:KV Cache。这个听起来简单的"键值缓存",实际上是大模型推理效率的核心,也是理解大模型部署成本的钥匙。

从一个直观的问题开始

假设你让一个语言模型续写这句话:“人工智能正在改变”。

模型生成"世界"这个词时,需要看前面的所有词——“人工智能”、“正在”、“改变”。当它生成第二个词"的"时,理论上需要重新计算"人工智能"、“正在”、“改变”、“世界"这些词的注意力表示。

但等等——“人工智能"这个词的注意力表示,在生成"世界"时已经算过了,为什么生成"的"的时候还要重新算?

这就像你每次做数学题都要重新背诵乘法口诀——明明上次已经背过了,为什么还要重来?

KV Cache就是这个"记忆”。它把计算过的Key和Value向量存下来,下次用到时直接取用,不需要重新计算。

注意力机制的本质回顾

要理解KV Cache,先要理解Transformer的注意力机制在算什么。

自注意力机制的核心是三个向量:Query(查询)、Key(键)、Value(值)。每个输入token都会产生这三个向量。Query是"我想找什么”,Key是"我是什么",Value是"我的内容"。

graph LR
    subgraph "Token处理流程"
        A[输入Token] --> B[线性变换]
        B --> C[Query Q]
        B --> D[Key K]
        B --> E[Value V]
        C --> F[注意力分数计算]
        D --> F
        F --> G[加权求和]
        E --> G
        G --> H[输出表示]
    end

注意力分数的计算公式是:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

这里有个关键点:每个位置的Query只需要跟它之前所有位置的Key和Value计算。这就是因果掩码(Causal Mask)的作用——当前位置不能"看见"未来的位置。

自回归生成时,第$t$个token需要与第$1$到第$t-1$个token的Key和Value计算注意力。问题是:每生成一个新token,之前所有token的Key和Value都要重新算吗?

KV Cache的工作原理

不需要。这正是KV Cache解决的问题。

KV Cache的思路很简单:把每个token的Key和Value向量缓存起来。生成新token时,只需要:

  1. 计算新token的Query、Key、Value
  2. 把新的Key和Value追加到缓存中
  3. 用新token的Query与缓存中所有的Key计算注意力
  4. 用注意力分数对缓存中所有的Value加权求和
sequenceDiagram
    participant M as 模型
    participant C as KV Cache
    participant O as 输出

    Note over M,C: Step 1: 处理第一个token
    M->>C: 存储K1, V1
    C->>M: 返回K1, V1
    M->>O: 输出token 2

    Note over M,C: Step 2: 处理第二个token
    M->>C: 存储K2, V2
    C->>M: 返回K1, K2, V1, V2
    M->>O: 输出token 3

    Note over M,C: Step n: 处理第n个token
    M->>C: 存储Kn, Vn
    C->>M: 返回K1...Kn, V1...Vn
    M->>O: 输出token n+1

这个简单的设计带来了巨大的效率提升。HuggingFace的基准测试显示,启用KV Cache后,推理速度可以提升5倍以上。

两阶段推理:Prefill与Decode

KV Cache的工作方式将LLM推理分成两个截然不同的阶段:

Prefill阶段:模型一次性处理所有输入token。这阶段是计算密集型的,GPU利用率高,因为可以并行计算所有token的表示。同时,这阶段会建立KV Cache的初始内容。

Decode阶段:模型逐个生成输出token。每生成一个token,就要从KV Cache读取所有之前的Key和Value,计算注意力,然后把新token的Key和Value追加到缓存中。

graph TB
    subgraph "Prefill阶段 - 计算密集"
        P1[输入: 所有prompt tokens] --> P2[并行计算]
        P2 --> P3[建立完整KV Cache]
        P3 --> P4[输出: 第一个生成token]
    end

    subgraph "Decode阶段 - 内存带宽密集"
        D1[输入: 单个新token] --> D2[读取完整KV Cache]
        D2 --> D3[计算注意力]
        D3 --> D4[追加新K,V到缓存]
        D4 --> D5[输出: 下一个token]
    end

    P4 --> D1

这两个阶段的性能特征完全不同。Prefill是计算受限(compute-bound),Decode是内存带宽受限(memory-bound)。在Decode阶段,GPU大部分时间在等待从显存读取KV Cache数据,而不是在计算。

这就是为什么第一个token生成很慢(Prefill阶段需要处理整个输入),而后续token生成较快(Decode阶段每次只处理一个新token)。但当上下文很长时,Decode阶段也会变慢——因为要读取的KV Cache变大了。

KV Cache到底占用多少内存?

这是部署大模型时必须回答的问题。让我们推导精确的计算公式。

存储结构分析

对于单个注意力层,单个token,需要存储:

  • Key向量:形状为 [num_kv_heads, head_dim]
  • Value向量:形状为 [num_kv_heads, head_dim]

对于整个模型,所有层,序列长度为$T$:

  • Key缓存:[num_layers, batch_size, num_kv_heads, T, head_dim]
  • Value缓存:[num_layers, batch_size, num_kv_heads, T, head_dim]
graph TD
    subgraph "KV Cache存储结构"
        A["KV Cache Tensor<br/>形状: [L, B, H_kv, T, D_h]"]
        A --> B["L: Transformer层数"]
        A --> C["B: Batch Size"]
        A --> D["H_kv: KV头数"]
        A --> E["T: 序列长度"]
        A --> F["D_h: 头维度"]
    end

    subgraph "内存计算"
        G["每个元素: 2 bytes (FP16)"]
        H["总元素数 = L × B × H_kv × T × D_h"]
        I["总内存 = 2 × 总元素数 × 2 bytes"]
    end

    A --> G
    G --> H
    H --> I

内存计算公式

KV Cache的总内存为:

$$M_{kv} = 2 \times L \times B \times T \times H_{kv} \times D_h \times \text{bytes}$$

其中:

  • $2$:Key和Value两个矩阵
  • $L$:Transformer层数
  • $B$:batch size(并发请求数)
  • $T$:序列长度(token数)
  • $H_{kv}$:Key-Value头数
  • $D_h$:每个头的维度
  • $\text{bytes}$:每个元素的字节数(FP16为2字节,FP32为4字节)

对于标准的多头注意力(MHA),$H_{kv}$等于查询头数$H_q$,且$H_q \times D_h = d_{model}$(隐藏层维度)。此时公式可以简化为:

$$M_{kv} = 2 \times L \times B \times T \times d_{model} \times \text{bytes}$$

实际计算示例

让我们用几个热门模型来计算:

模型 层数$L$ 隐藏维度$d_{model}$ KV头数$H_{kv}$ 4K上下文单请求KV Cache
GPT-2 Small 12 768 12 ~75 MB
LLaMA 7B 32 4096 32 ~2.0 GB
LLaMA 2 70B 80 8192 8 ~1.0 GB
LLaMA 3 70B 80 8192 8 ~1.0 GB

等等,70B模型的KV Cache反而比7B模型小?这正是Grouped Query Attention的功劳,后面会详细解释。

对于Llama-3-70B,8K上下文,32个并发请求:

$$M_{kv} = 2 \times 80 \times 32 \times 8192 \times 8 \times 128 \times 2 = 81.9 \text{ GB}$$

这已经超过了一张A100 80GB的显存——仅用于KV Cache,还没算模型权重。

上下文长度的线性代价

KV Cache内存与序列长度成正比。这意味着:

graph LR
    subgraph "KV Cache随上下文长度增长"
        A["2K上下文<br/>~2GB"] --> B["4K上下文<br/>~4GB"]
        B --> C["8K上下文<br/>~8GB"]
        C --> D["32K上下文<br/>~32GB"]
        D --> E["128K上下文<br/>~128GB"]
    end

    style A fill:#90EE90
    style B fill:#FFD700
    style C fill:#FFA500
    style D fill:#FF6347
    style E fill:#FF0000,color:#fff
  • 2K上下文 → 2GB KV Cache
  • 8K上下文 → 8GB KV Cache
  • 32K上下文 → 32GB KV Cache
  • 128K上下文 → 128GB KV Cache

这个线性增长关系没有捷径可绕。每个token的Key和Value必须存储,否则无法正确计算注意力。

更关键的是,当上下文超过某个阈值后,KV Cache会比模型权重占用更多内存。这个交叉点可以通过以下公式计算:

$$T_{crossover} = \frac{M_{weights}}{2 \times L \times d_{model} \times \text{bytes}}$$

对于LLaMA 7B(FP16权重约14GB),这个交叉点大约在32K token。超过这个长度,KV Cache成为主要的内存消费者。

PagedAttention:从操作系统借来的智慧

传统的KV Cache实现有一个致命问题:内存碎片化

连续分配的困境

早期系统(如HuggingFace Transformers)为每个请求预分配一块连续内存,大小按最大序列长度计算。这导致三种内存浪费:

  1. 内部碎片:请求的实际长度往往远小于最大长度,预分配的空间大量闲置
  2. 外部碎片:不同请求的预分配大小不同,buddy allocator无法高效利用
  3. 预留浪费:生成过程中预留的"未来token空间"在当前时刻无法被其他请求使用

vLLM论文中的测量数据显示,传统系统只有20-38%的KV Cache内存实际存储了有用的token状态,其余60-80%都是浪费。

分页内存管理

2023年,UC Berkeley的vLLM团队提出了PagedAttention,核心思想来自操作系统的虚拟内存管理:

  • 把KV Cache切分成固定大小的块(blocks)
  • 每个块可以存储若干个token的Key和Value
  • 块在物理内存中不需要连续存储
  • 通过**块表(block table)**映射逻辑块到物理块
graph TB
    subgraph "逻辑视图(每个请求看到的)"
        L1[逻辑块0] 
        L2[逻辑块1]
        L3[逻辑块2]
    end

    subgraph "物理存储(实际GPU内存)"
        P1[物理块7]
        P2[物理块1]
        P3[物理块3]
        P4[物理块9]
    end

    subgraph "块表"
        BT["块表映射:<br/>逻辑0 → 物理7<br/>逻辑1 → 物理1<br/>逻辑2 → 物理3"]
    end

    L1 --> BT
    L2 --> BT
    L3 --> BT
    BT --> P1
    BT --> P2
    BT --> P3

这个设计带来了几个重要优势:

近乎零的内存浪费:只有最后一个块可能有少量未填满的空间,整体碎片率低于4%。相比传统系统的60-80%浪费,这是质的飞跃。

按需分配:生成过程中,只有当需要存储新token时才分配新块。不需要预先为"可能生成的token"预留空间。

内存共享:这是PagedAttention最精妙的设计。当多个请求共享相同的prompt前缀(比如系统提示词),它们的KV Cache可以共享相同的物理块。

Copy-on-Write机制

当共享的物理块需要被修改时(比如不同请求生成了不同的token),PagedAttention采用写时复制(Copy-on-Write)机制:

stateDiagram-v2
    [*] --> 检测共享块
    检测共享块 --> 引用计数>1?: 检查引用计数
    引用计数>1? --> 分配新物理块: 是
    引用计数>1? --> 直接写入: 否
    分配新物理块 --> 复制内容
    复制内容 --> 更新块表
    更新块表 --> 减少原块引用计数
    减少原块引用计数 --> 写入完成
    直接写入 --> 写入完成
    写入完成 --> [*]

这个机制使得beam search、parallel sampling等解码算法的内存效率大幅提升。vLLM的测试显示,beam search可以节省37-66%的内存。

性能影响

PagedAttention对性能的影响是双面的:

开销:非连续内存访问增加了attention kernel的复杂度。vLLM的attention kernel比FasterTransformer慢20-26%。

收益:由于内存利用率大幅提升,可以支持更大的batch size,整体吞吐量提升2-4倍。

这是一笔划算的交易:单次计算慢一点,但能同时处理更多请求。

GQA:用架构设计减少缓存

如果说PagedAttention是在系统层面优化内存管理,那么Grouped Query Attention(GQA)则从模型架构层面减少KV Cache的大小。

从MHA到MQA到GQA

标准的多头注意力(MHA):每个查询头都有自己对应的Key头和Value头。如果模型有32个查询头,就有32个Key头和32个Value头。

Multi-Query Attention(MQA):所有查询头共享同一组Key和Value。32个查询头,只有1个Key头和1个Value头。KV Cache减少32倍,但模型质量下降明显。

Grouped Query Attention(GQA):折中方案。将查询头分成若干组,每组共享一组Key和Value。比如64个查询头分成8组,每组8个查询头共享一组KV。KV Cache减少8倍,模型质量下降很小。

graph LR
    subgraph "MHA (32头) - 每头独立KV"
        M1[Q1→K1,V1]
        M2[Q2→K2,V2]
        M3["..."]
        M4[Q32→K32,V32]
        M5["KV Cache: 32x"]
        style M5 fill:#FF6347
    end

    subgraph "MQA (32头) - 全共享"
        Q1[Q1]
        Q2[Q2]
        Q3["..."]
        Q4[Q32]
        K1[共享的K,V]
        Q1 --> K1
        Q2 --> K1
        Q3 --> K1
        Q4 --> K1
        M6["KV Cache: 1x"]
        style M6 fill:#90EE90
    end

    subgraph "GQA (32头, 8组) - 分组共享"
        G1["组1: Q1-Q8 → K1,V1"]
        G2["组2: Q9-Q16 → K2,V2"]
        G3["..."]
        G4["组8: Q25-Q32 → K8,V8"]
        G5["KV Cache: 8x"]
        style G5 fill:#FFD700
    end

GQA的内存节省

GQA的KV Cache节省比例可以直接计算:

$$\text{节省比例} = \frac{H_q}{H_{kv}}$$

其中$H_q$是查询头数,$H_{kv}$是Key-Value头数。

LLaMA 2 70B和LLaMA 3 70B使用64个查询头,8个KV头。这意味着:

  • 传统MHA:KV Cache大小 = $2 \times 80 \times 8192 \times 64 \times 128 \times 2$
  • GQA配置:KV Cache大小 = $2 \times 80 \times 8192 \times 8 \times 128 \times 2$
  • 节省:$64/8 = 8$倍

这就是为什么70B模型的KV Cache反而比7B模型小的原因。

主流模型的GQA配置

模型 查询头数 KV头数 KV Cache节省
LLaMA 7B 32 32 1x (无节省)
LLaMA 2 70B 64 8 8x
LLaMA 3 70B 64 8 8x
Mistral 7B 32 8 4x
Mixtral 8x7B 32 8 4x
DeepSeek-V2 128 2 64x

GQA已经成为大模型的标准配置。它在不显著损害模型质量的前提下,大幅降低了推理的内存需求。

KV Cache量化:进一步压缩

即使使用了GQA,长上下文场景下KV Cache仍然是内存大户。量化是另一个重要的优化方向。

精度与内存的权衡

数据类型 每元素字节数 相对FP16节省
FP32 4 -100% (更大)
FP16/BF16 2 基准
FP8 (E4M3/E5M2) 1 50%
INT8 1 50%
INT4 0.5 75%
graph LR
    subgraph "量化后的KV Cache大小对比 (以LLaMA 7B 4K上下文为例)"
        A["FP32: ~8GB"] --> B["FP16: ~4GB"]
        B --> C["FP8: ~2GB"]
        C --> D["INT4: ~1GB"]
    end

    style A fill:#FF0000,color:#fff
    style B fill:#FFA500
    style C fill:#FFD700
    style D fill:#90EE90

FP8量化是目前最受欢迎的方案。它将KV Cache从FP16的2字节压缩到1字节,内存减半,同时对模型质量的影响很小。

INT8和INT4量化需要额外的缩放因子(scale factor),但可以进一步压缩。实测表明,INT8 KV Cache几乎无损,INT4在大多数场景下也可以接受。

主流框架支持

vLLM、TensorRT-LLM、llama.cpp等主流推理框架都已经支持KV Cache量化。以vLLM为例:

# FP8 KV Cache量化
--kv-cache-dtype fp8

# INT8 KV Cache量化  
--kv-cache-dtype int8

量化的一个额外好处是减少内存带宽压力。在Decode阶段,瓶颈往往是内存带宽而不是计算。更小的KV Cache意味着更少的数据传输,即使计算量不变,整体速度也会提升。

Prefill与Decode:两个世界的性能特征

理解KV Cache的另一个关键角度是Prefill和Decode阶段的根本差异。

性能瓶颈分析

可以用算术强度(Arithmetic Intensity)来量化:

$$I = \frac{\text{FLOPs}}{\text{Bytes transferred}}$$

对于A100 GPU,峰值算力312 TFLOPS(FP16),峰值带宽2 TB/s。计算-带宽比约为156 FLOPs/Byte。

如果操作的计算强度低于156,就是内存受限;高于156,就是计算受限。

graph LR
    subgraph "A100 GPU性能边界"
        A["计算-带宽比: 156 FLOPs/Byte"]
        B["计算强度 < 156: 内存受限"]
        C["计算强度 > 156: 计算受限"]
    end

    D["Prefill阶段"] --> E["计算强度高<br/>计算受限"]
    F["Decode阶段"] --> G["计算强度低(~0.5)<br/>严重内存受限"]

    style E fill:#90EE90
    style G fill:#FF6347

Decode阶段的attention计算强度约为:

$$I_{decode} \approx \frac{2 \times d_{model}}{4 \times d_{model}} = 0.5 \text{ FLOPs/Byte}$$

这远远低于156的阈值,证实了Decode阶段是严重的内存带宽受限。

分离式推理

理解了两个阶段的差异后,一个自然的想法是:能不能把它们分开处理?

DistServe和Mooncake等系统提出了分离式推理架构:

graph TB
    subgraph "分离式推理架构"
        Client[客户端请求] --> Router[请求路由]
        Router --> Prefill[Prefill节点<br/>高算力GPU]
        Router --> Decode[Decode节点<br/>大显存带宽]

        Prefill -->|传输KV Cache| Decode
        Decode -->|返回结果| Client
    end

    subgraph "Prefill节点特点"
        P1[计算密集型]
        P2[需要高FLOPs]
        P3[处理完后可复用]
    end

    subgraph "Decode节点特点"
        D1[内存带宽密集]
        D2[需要大显存]
        D3[保持KV Cache]
    end
  • Prefill节点:专注于处理prompt,需要高算力
  • Decode节点:专注于生成token,需要大显存带宽

两个节点之间通过高速网络传输KV Cache。这种架构可以:

  1. 针对不同阶段选择不同的硬件
  2. Prefill节点处理完后可以立即开始处理下一个请求
  3. Decode节点保持KV Cache,支持多轮对话

Mooncake的测试显示,分离式架构可以提升525%的吞吐量。

实践指南:容量规划

部署大模型时,如何估算显存需求?

总内存公式

$$M_{total} = M_{weights} + M_{kv} + M_{activation} + M_{overhead}$$
  • $M_{weights}$:模型权重,FP16下约等于参数量×2字节
  • $M_{kv}$:KV Cache,使用前面的公式计算
  • $M_{activation}$:激活内存,约等于权重的10-20%
  • $M_{overhead}$:框架开销,约1-2GB

快速估算示例

以LLaMA 3 8B为例,假设8K上下文,batch size = 1:

模型权重:8B参数 × 2字节 = 16 GB
KV Cache:2 × 32 × 1 × 8192 × 32 × 128 × 2 = 4.3 GB (使用GQA后更小)
激活内存:约 1.6 GB
框架开销:约 1 GB
总计:约 23 GB

一张24GB的RTX 4090刚好够用。

如果batch size = 8:

KV Cache:4.3 GB × 8 = 34.4 GB
总计:16 + 34.4 + 1.6 + 1 = 53 GB

需要A100 80GB才能容纳。

graph TB
    subgraph "显存分配决策树"
        A[选择模型] --> B{上下文长度}
        B -->|短 2K| C[KV Cache小<br/>权重主导]
        B -->|中 8K| D[KV Cache中等<br/>需平衡]
        B -->|长 32K+| E[KV Cache大<br/>KV主导]

        C --> F{Batch Size}
        D --> F
        E --> F

        F -->|单请求| G[可用消费级GPU]
        F -->|多并发| H[需要数据中心GPU]

        G --> I[计算总内存需求]
        H --> I
    end

关键决策点

  1. 上下文长度:每翻倍,KV Cache翻倍
  2. Batch size:每翻倍,KV Cache翻倍
  3. 是否使用GQA:根据模型架构决定
  4. 量化策略:FP8可以减半KV Cache

KV Cache优化的演进历程

从2017年Transformer诞生至今,KV Cache的优化经历了几个重要阶段:

timeline
    title KV Cache优化技术演进
    2017 : Transformer诞生
         : 标准MHA架构
         : KV Cache概念出现
    2019 : MQA提出
          : 所有头共享KV
          : 质量损失明显
    2020 : GQA提出
          : 分组共享KV
          : 质量损失小
    2023 : PagedAttention
          : vLLM发布
          : 解决内存碎片
          : 吞吐量提升2-4x
    2024 : KV Cache量化
          : FP8成为主流
          : 进一步减少内存
    2025 : 分离式推理
          : Prefill/Decode分离
          : 大幅提升吞吐量

总结

KV Cache是大模型推理的核心技术,它把重复的注意力计算转化为内存查找,实现了空间换时间的优化。理解KV Cache的本质、内存消耗和优化技术,是部署大模型的必备知识。

几个关键认知:

KV Cache的内存增长是线性的,与上下文长度和batch size成正比。长上下文场景下,KV Cache可能超过模型权重成为主要的内存消费者。

PagedAttention解决了内存碎片化问题,将KV Cache利用率从20-38%提升到96%以上。这是系统层面的优化。

GQA从架构层面减少KV Cache,现代大模型普遍采用8个KV头配合64个查询头的配置,实现8倍的内存节省。

Prefill和Decode有根本不同的性能特征。Prefill是计算受限,Decode是内存带宽受限。分离式推理是面向未来的架构方向。

当你下次思考"这个模型能跑在什么显卡上"时,记得把KV Cache算进去——它往往比模型权重更需要你的关注。