二、KV Cache 与多头注意力变体
二、 KV Cache 与多头注意力变体¶
1. 为什么自回归生成会引入 KV Cache?没有 KV Cache 会怎样?¶
要理解 KV Cache,我们先回到自回归生成的数学本质。Transformer 的自注意力机制在生成第 t 个 token 时,需要计算该 token 与前面所有历史 token 的注意力:
$$ Attention(Q_{t},K_{1:t},V_{1:t})=softmax\left(\frac{Q_{t}K_{1:t}^{T}}{\sqrt{d_{k}}}\right)V_{1:t} $$
其中 $ Q_t $ 是第 $ t $ 个 token 的 Query, $ K_{1:t} $ 和 $ V_{1:t} $ 是所有历史 token(1 到 $ t $)的 Key 和 Value。注意,这里的 Key 和 Value 对于历史 token 来说,在生成后续 token 时是不变的——比如第 3 个 token 的 Key、Value 在计算第 3、4、5 个 token 时都相同。因此,我们可以把已经算出来的 Key 和 Value 缓存起来,避免每次生成新 token 时都重新计算整个历史序列。这就是 KV Cache 的由来。
没有 KV Cache 会怎样?¶
如果没有缓存,每生成一个新 token,模型必须把整个序列(包括 prompt 和之前生成的所有 token)再完整前向传播一次。这会带来两个灾难性的后果:
-
计算量爆炸:在生成第 t 个 token 时,需要计算所有 t 个 token 的注意力,计算复杂度是 $ O(t^{2}) $。而使用 KV Cache 后,每次只需计算当前 token 的注意力,复杂度降为 $ O(t) $。实际中,一个 1000 步的生成,没有缓存时总的计算量将是缓存时的几百倍。
-
显存不可承受:重新计算整个序列意味着要保存所有中间层的激活值,显存消耗同样巨大。而 KV Cache 只保存紧凑的 Key 和 Value 张量,显存开销可控。
用一个直观的比喻:没有KV Cache就像你每写一个字,都要把整本书从头到尾重读一遍,而有了KV Cache,你只需记住之前写了什么,写新字时只需看一眼前面的要点。
公式化对比:¶
假设模型层数为 L,隐藏维度为 d,生成 N 个 token。
• 无缓存时,总计算量 $ \propto L \cdot N^{2} \cdot d $(近似),显存峰值 $ \propto L \cdot N \cdot d $(中间激活)。
- 有缓存时,总计算量 $ \propto L \cdot N \cdot d $(每个 token 的计算),显存额外占用 $ 2 \times L \times N \times d \times \text{sizeof}(\text{dtype}) $(KV Cache)。
因此,KV Cache 是自回归生成得以在有限资源下运行的关键发明。
2. KV Cache 在推理时是如何被读写和更新的?¶
KV Cache 的操作分为两个阶段:Prefill 和 Decode。
Prefill 阶段(预填充)¶
当用户输入一个 prompt 时,模型一次性处理整个 prompt 序列。对于 prompt 中的每一个 token,模型都会通过线性投影计算出它对应的 Key 和 Value,并写入缓存。这个过程是并行的:一个 prompt 如果有 $ S_{prompt} $ 个 token,那么一次前向传播就会产生 $ S_{prompt} $ 组 K、V,全部存入预先分配好的缓冲区。此时,缓存从空变为包含整个 prompt 的 K、V。
Decode 阶段(解码生成)¶
接下来,模型进入逐 token 生成循环:
• 读取:对于当前要生成的 token $ t $,模型只计算它自己的 Query $ Q_t $(以及 Key、Value,用于后续缓存)。注意力计算需要读取所有历史 token 的 Key 和 Value,即从缓存中取出 $ K_{1:t-1} $ 和 $ V_{1:t-1} $,与 $ Q_t $ 计算注意力分数,然后聚合 $ V_{1:t-1} $。这是 Decode 阶段最重的操作,因为每次都要遍历整个缓存。
更新:计算完成后,当前 token 的 Key 和 Value 会被追加写入缓存末尾。于是缓存长度从 $ t-1 $ 增长到 $ t $。
如此循环,直到生成结束符或达到最大长度。
工程实现细节:¶
- 缓存通常预先分配一块连续显存,形状为 [batch_size, num_heads, max_seq_len, head_dim]。有一个 seq_len 指针指示当前有效长度。
在 Prefill 阶段,会一次性写入整个 prompt 的内容到 [:prompt_len] 位置;Decode 阶段每次在 [prompt_len + step] 位置写入新 token。
为了优化显存和速度,现代推理引擎(如 vLLM)采用 PagedAttention,将 KV Cache 分页存储,像操作系统管理虚拟内存那样减少碎片,并支持动态内存分配。
读与带宽分析:在 Decode 的每一步,需要从显存读取整个 KV Cache(大小 $ \propto L \times S \times d $)以及模型权重,这构成了显存带宽瓶颈。这就是为什么长文本推理速度取决于显存带宽而不是算力。
3. 为什么 KV Cache 是长文本推理的主要瓶颈?¶
当上卜又长度从几千扩展到几万甚至几十万时,KV Cache 迅速成为压倒一切的性能瓶颈,主要体现在三个方面:显存容量、显存带宽、计算方式。
显存容量压力¶
以 LLaMA-7B 为例,FP16 精度下,每个 token 在每层所需的 KV Cache 大小为 $ 2 \times 32 $ heads × 128 dim × 2 bytes = 16, 384 字节(即 16 KB)。32 层合计每 token 约 0.5 MB。当上下文长度为 32k 时,KV Cache 总大小 = 0.5 MB × 32, 768 ≈ 16 GB,已经超过了 14 GB 的模型权重!如果扩展到 128k,KV Cache 将达到 64 GB,单张 80 GB 的 A100 几乎无法容纳。所以长文本推理首先面临的是显存放不下的问题。
显存带宽瓶颈¶
如前一个问题所述,Decode 阶段每生成一个 token,就要把整个 KV Cache 从显存读一遍。假如 KV Cache 有 16 GB,而 A100 的显存带宽为 2 TB/s,那么仅读取 KV Cache 就需要 8 毫秒。这 8 毫秒内 GPU 的计算单元几乎完全空闲,因为实际计算量极小。当上下文更长,该读取时间线性增加,成为延迟的直接来源。
计算密度极低¶
注意力计算是“一个查询向量与海量键向量求点积”,每个 token 的计算量与 KV Cache 大小成正比,但其中大部分操作是带宽受限的。硬件利用率(如 Tensor Core 利用率)通常只有百分之几,浪费严重。
这三者叠加,使得KV Cache成了长文本推理的“阿喀琉斯之踵”。所有优化手段(GQA、MQA、KV Cache量化、FlashAttention、PagedAttention)本质上都是在与KV Cache作斗争:要么压缩它的大小,要么更高效地读写它。
4. MHA、MQA、GQA 三者在 KV 头上的核心区别是什么?¶
这三种注意力机制的根本区别在于 Key 和 Value 的头数,以及它们与 Query 头的对应关系。
MHA(Multi-Head Attention,多头注意力)¶
查询(Q)、键(K)、值(V)都有H个头,每个头独立。这是标准Transformer的做法。KV Cache 需要存储H个K头和H个V头,体积最大。
MQA(Multi-Query Attention,多查询注意力)¶
查询 Q 仍然有 H 个头,但键 K 和值 V 只有 1 个头。所有 Q 头共享同一套 K 和 V。KV Cache 大小骤降至原来的 $ 1/H $。
GQA(Grouped-Query Attention,分组查询注意力)¶
Q 有 H 个头,K 和 V 有 G 个头( $ 1 < G < H $,且通常 G 远小于 H)。将 H 个 Q 头分成 G 组,每组内的 Q 头共享一个 K 头和一个 V 头。KV Cache 大小变为原来的 $ G/H $。
核心区别表(假设 H=32,G=8):
| 机制 | Q 头数 | K 头数 | V 头数 | KV Cache 相对大小 |
| MHA | 32 | 32 | 32 | 1 (基准) |
| MQA | 32 | 1 | 1 | 1/32 ≈ 0.031 |
| GQA | 32 | 8 | 8 | 8/32 = 0.25 |
直观理解:¶
· MHA 就像每个专家(Q 头)都有自己的专属资料库(K、V 头),表达力最强,但资料库太多,占地方。
· MQA 是所有专家共用一个资料库,最省地方,但专家们容易打架,表达力受限。
GQA 是把专家分成几个小组,每个组共享一个资料库,兼顾了表达力和存储成本。
5. MQA 为什么能大幅减少 KV Cache 大小?它对模型效果有什么影响?¶
大幅减少的原理:¶
KV Cache 的总大小公式为 $ 2 \times L \times B \times S \times H \times d_h \times \text{sizeof} $,其中 $ H $ 是 $ K $ 和 $ V $ 的头数。MQA 将 $ H $ 降为 1,所以 KV Cache 直接变为原来的 $ 1/H $。例如 LLaMA-7B 的 $ H = 32 $,换成 MQA 后 KV Cache 缩小 32 倍!这简直是颠覆性的压缩。
对模型效果的影响:¶
正向影响(速度与吞吐):Decode 阶段每次读取的 KV Cache 体积极小,带宽瓶颈大幅缓解,生成速度成倍提升。同时,显存压力骤降,可以支持更长的上下文或更大的 batch size。
负向影响(模型能力):所有Q头共享一套K、V,意味着不同的注意力模式被迫使用同一个“信息过滤器”。这会限制模型的表达能力和对不同子空间的捕捉。实践中,MQA往往导致困惑度(PPL)略微上升,在需要细粒度理解的任务(如事实回忆、多跳推理)上性能有所下降。不过,对于许多生成任务,这种轻微的精度损失是可接受的,且可以通过增加训练数据或模型尺寸来弥补。
因此,MQA 是极致追求推理效率的选择,常用于对延迟和吞吐有极端要求的部署场景。
6. GQA 是如何在 MHA 和 MQA 之间折中的?分组数一般取多少?¶
GQA 的设计思想是:保留多查询的缓存优势,同时通过分组恢复部分表达能力。
折中原理:¶
它将 H 个 Q 头均匀分成 G 组,每组内所有 Q 头共享同一个 K 头和 V 头。这样,K、V 的头数降为 G,KV Cache 缩小为原来的 G/H。调整 G 可以在表达能力和缓存大小之间进行滑动调节:
当 $ G=H $时,就是MHA,表达力最强,缓存最大。
当 G = 1 时,就是 MQA,缓存最小,表达力最弱。
当 G 取中间值,比如 4 或 8,则兼顾两者。
分组数一般取多少?¶
实践中,G 通常取 4 到 8。例如:
· LLaMA-2 70B 使用了 H=64,G=8,即 64 个 Q 头分成 8 组,每组 8 个 Q 头共享一个 K、V 头。
· LLaMA-3 8B/70B 延续了 GQA,同样 G=8 左右。
· 较小的模型(如 7B/13B)也经常使用 G=4 或 G=8。
研究表明,当 G=8 时,模型在多数基准上的表现与 MHA 几乎持平,而 KV Cache 缩小到原来的 1/8,这已经是一个非常划算的买卖。如果硬件条件苛刻,也可以进一步降低 G,以轻微性能损失换取更大吞吐。
7. 从 MHA 转换为 GQA,通常怎么做?是否需要重新训练?¶
从 MHA 转为 GQA 主要有两条路径:权重转换(近似转换)和重新训练/微调。
1. 权重转换(离线转换,无需完整重训)¶
既然 MHA 已经训练好了 H 个独立的 K、V 头,我们可以将它们合并成 G 个头。最常见的方法是平均池化:将同一组内的多个 K(或 V)头对应的权重矩阵按元素求平均,得到一个新的 K(或 V)头。具体步骤如下:
· 确定分组方案,例如将 32 个头分成 8 组,每组 4 个头。
对每一组的四个 K 头的权重(形状 [head_dim, d_model])求平均,得到一个 K 头。
同样处理 V 头。
Q 头保持原样不动(因为 GQA 不影响 Q 头数量)。
替换模型中的注意力参数,保存新模型。
优缺点:¶
优点:零训练成本,转换极快,几分钟即可完成。
- 缺点:性能会有一定损失,因为原始头并不是为了共享而训练的,直接平均可能导致信息丢失。有时需要少量微调(例如几百步)来恢复大部分性能。
2. 从头预训练或大量微调¶
如果希望获得最佳效果,可以在预训练阶段就采用 GQA 架构,或者对转换后的模型进行充分的微调(比如在大量语料上继续训练一定步数)。这样模型能够适应共享的 K、V 头,性能损失可降至最低。LLaMA-2 就是在预训练时就使用了 GQA。
实际选择:¶
若仅仅是为了快速部署,追求推理加速,且对轻微的性能下降可以接受,权重转换加极少量微调(甚至零微调)是工业界常用的“捷径”。
若模型是开源基座且资源允许,最好从零开始用 GQA 训练,或使用官方提供的 GQA 版本模型。
8. 为什么现在主流开源模型(如 LLaMA 2/3)都用 GQA?¶
LLaMA-2、LLaMA-3、Mistral、Qwen等主流模型纷纷采用GQA,背后是工业界对“大规模部署可行性”的共识。具体原因:
1. 长文本推理的刚需¶
随着应用场景对长上下文的需求激增(如文档摘要、长对话),KV Cache 的显存和带宽瓶颈变得不可容忍。GQA 可以将 KV Cache 压缩 4~8 倍,直接让 32k 上下文在单卡上运行成为可能,而不需要复杂的模型并行。
2. 近乎无损的精度平衡¶
实验反复证明,当 G=8 或 4 时,模型在语言理解、生成、推理等基准上的表现与 MHA 几乎没有显著差异。也就是说,GQA 用极小的精度代价换来了巨大的效率提升,这是其他压缩技术(如权重量化)难以做到的。
3. 缓解 Decode 带宽瓶颈¶
如前分析,Decode 阶段是带宽受限的。GQA 减少了 K、V 头的数量,也就减少了每步需要读出的数据量,直接提高了生成速度(tokens/s),降低了每 token 的延迟。这对于在线服务至关重要。
4. 硬件适配与生态趋势¶
现代 GPU 的张量核心对矩阵乘法的形状有一定偏好。GQA 使 K、V 头的维度更规整,更容易被高性能算子(如 FlashAttention、PagedAttention)优化。同时,推理框架(VLLM、TensorRT-LLM)都对 GQA 提供了深度支持,形成了强大的生态合力。
5. 架构灵活性与未来兼容¶
GQA 中的组数 G 是一个可调的超参数。硬件富裕时可以设大一些(接近 MHA),硬件紧张时设小一些(接近 MQA)。这种灵活性让同一个模型架构可以适配从云端到边缘的多种部署环境。
总之,GQA是在当前硬件条件和应用需求下的最优解,它让大模型既能“思考”得更深,也能“跑”得更远。
9. 解释 MLA(多头潜在注意力)如何通过低秩压缩减少 KV Cache。¶
MLA(Multi-head Latent Attention)是 DeepSeek-V2 / V3 等最新模型中采用的注意力机制,目的是以比 GQA 更极致的方式压缩 KV Cache。它的核心思想来自 LoRA 的低秩假设:Key 和 Value 的高维表示可以被压缩到一个低维的“潜空间”中,而不会丢失太多信息。
在传统的 MHA 中,我们为每个 token 生成完整的 Key 和 Value 向量。这些向量维度为 $ d = H \times d_h $(通常等于 $ d_{\text{model}} $),每个 token 的 KV Cache 大小就是 2d 个浮点数。MLA 通过引入两个投影矩阵,将 KV 的生成过程分解为三步:
-
将输入 token 的隐藏状态 x 通过一个下投影矩阵 $ W_{\text{down}} \in \mathbb{R}^{d \times d_l} $ 投影到一个低维潜向量 $ c = xW_{\text{down}} $,其中 $ d_l \ll d $。
-
将潜向量 $ c $ 通过两个上投影矩阵 $ W_{K,up} $ 和 $ W_{V,up} $ 恢复到高维的 Key 和 Value: $ K = cW_{K,up} $, $ V = cW_{V,up} $。这两个上投影矩阵的形状是 $ d_l \times d $。
在推理时,MLA 不需要缓存完整的高维 K 和 V,只需要缓存那个低维的潜向量 c。对于每一个 token,缓存的显存从 2d 个元素骤降为 $ d_{l} $ 个元素。当需要计算注意力时,再从潜向量 c 通过上投影矩阵实时恢复出高维 K、V。这种“缓存压缩,计算解压”的策略与 LoRA 在微调时的思路(只训练低秩矩阵)一脉相承,但用在了相反的方向。
KV Cache 的压缩比为 $ \frac{\omega}{d_l} $。例如,DeepSeek-V3 的某些配置中, $ d = 5120 $,而 $ d_l = 512 $ 甚至更小,压缩比可达到 10 倍以上。这比 GQA(靠减少头数)的压缩更彻底,因为 GQA 最多只能压缩头数倍(比如 8 倍),而 MLA 直接在维度上进行压缩,不受头数限制。
与 GQA 的对比:GQA 是从“宽度”上缩减 KV Cache,通过让多个查询头共享一组 K、V 来降低头数;而 MLA 是从“高度”上缩减,直接压缩每个 token 的 K、V 表征维度。MLA 的潜在风险是:如果潜向量维度 $ d_{l} $ 过小,恢复出的 K、V 可能会丢失信息,影响模型精度。因此,MLA 往往需要与更深的模型或更多的训练数据配合,以弥补压缩带来的损失。但 DeepSeek-V2 证明了,在合理的 $ d_{l} $ 下,MLA 可以达到与 MHA 相当甚至更好的性能,同时 KV Cache 缩减一个数量级。
此外,MLA 有一个额外的工程优势:由于缓存的是潜向量 $ c $,而上投影矩阵 $ W_{K,up} $ 和 $ W_{V,up} $ 是层内共享的权重,计算 Key 和 Value 恢复的过程可以与 Attention 计算融合,不会增加过多的延迟。在 Decode 阶段,从缓存读取的潜向量体积小,带宽压力骤降,极长上下文的推理变得可行。
10. MLA 中“潜向量”的维度如何选取?过小会有什么问题?¶
选取潜向量维度 $ d_{l} $ 是 MLA 设计中最关键的参数,本质是在信息保留与压缩率之间寻找平衡。通常的做法是在目标压缩比附近,通过一系列小规模消融实验来确定。
选取原则与经验范围:¶
对于常见的7B~13B模型, $ d_{l} $一般取512到1024,对应的压缩比约为4~10倍(假设原模型d=4096或5120)。
• 对于更大的模型(如 70B),由于模型本身冗余度更高,可以承受更激进的压缩, $ d_{l} $ 甚至可取 256 或 384,压缩比可达 15~20 倍。
最终取值通常是在验证集上测试多个 $ d_{l} $,绘制性能(困惑度或下游任务得分)随 $ d_{l} $变化的曲线,选取性能开始明显下降之前的临界点,再稍微保守一点。
过小的潜向量维度会带来哪些问题?¶
1. 信息丢失,生成质量下降¶
潜向量 c 需要编码生成 Key 和 Value 所需的全部信息。当 $ d_l $ 过小时,这个瓶颈会强制丢弃一些对生成质量至关重要的细节。表现为:模型困惑度升高,生成的文本出现不连贯、事实错误、无法遵循复杂指令等情况,尤其在需要长程记忆和精细推理的任务上更明显。
2. 注意力模式坍缩¶
Key 和 Value 向量的表达能力减弱,可能导致多头注意力的多样性降低。即便物理上仍有多个头,但恢复出的 K、V 趋同,实际效果可能退化为类似 MQA 的效果,但比 MQA 更严重,因为 MQA 只是共享 K、V 头,而没有压缩每个头的维度。
3. 训练收敛变慢或需要更多数据¶
由于表示空间的急剧收缩,模型在训练初期可能难以学习到有效的注意力模式,需要更长的训练时间或更多的训练数据来弥补。如果从 MHA 转换到 MLA 时没有额外的训练,直接降低 $ d_{l} $ 通常会带来显著的性能下降。
4. 对位置编码的兼容性变差¶
一些模型(如 RoPE)对 Key 向量施加位置编码。在 MLA 中,位置编码通常施加在恢复后的 Key 上(即先恢复高维 K,再施加 RoPE),这没问题。但如果直接对潜向量施加位置编码(早期有些尝试),会因为潜空间的低维特性而破坏位置信息的有效编码,导致长文本泛化变差。
因此, $ d_{l} $ 的选取并非越小越好,需要在硬件约束和任务要求之间做权衡。实践中,很多团队会先在目标硬件上定一个能接受的最大延迟和显存,反推可接受的 KV Cache 大小,再算出 $ d_{l} $ 的最大值,然后尽量往这个上限靠,以保留更多模型能力。
11. 在推理时,MLA 实际缓存的是什么?和 GQA 的缓存内容有何不同?¶
MLA 实际缓存的内容:¶
MLA 缓存的是每个 token 在每一层的潜向量 $ c $。这个潜向量是一个低维的密集向量,长度 $ d_l $,通常远小于模型的隐藏维度 $ d $。以 DeepSeek-V3 为例, $ d = 5120 $,而缓存的潜向量维度可能只有 512。对于每个 token 和每一层,我们只需存储这一个向量。在 Decode 阶段生成新 token 时,我们从潜向量 $ c $ 通过上投影矩阵实时计算出高维的 K 和 V。由于上投影矩阵是固定的模型权重,这个计算量虽然存在,但相较于读取大体积 KV Cache 的带宽开销,可以忽略不计。
GQA 实际缓存的内容:¶
GQA 缓存的是低头的 Key 和 Value,即减少了头数的 K 和 V 张量。假设原 MHA 有 32 个 Key/Value 头,GQA 将头数降为 8,那么它缓存的就是这 8 个头的完整 K 和 V,每个头的维度仍然是 $ d_h $。所以缓存的总元素数是 $ 2 \times G \times d_h \times L $(对每个 token 每层)。注意,这里的 Key 和 Value 维度和 MHA 相同,只是头数变少了。
两者缓存内容的根本区别:¶
维度压缩方式:MLA 是从特征维度进行压缩,直接降低每个 token 缓存的向量长度(从 d 降到 $ d_{l} $)。GQA 是从头数维度进行压缩,降低存储的头数,但不改变每个头的维度。
缓存内容的性质:MLA 缓存的是一个抽象的“潜在表示”,需要配合上投影矩阵才能得到可用的 K 和 V;GQA 缓存的仍然是完整的 K 和 V,只是头的数量减少,可以直接参与注意力计算,无需任何恢复过程。
· 计算开销:MLA 在每次计算注意力时需要额外的矩阵乘法(潜向量 × 上投影矩阵)来恢复 K 和 V,这增加了一部分计算量。GQA 则完全没有这一开销。
压缩灵活性:MLA 的压缩比可以通过调整 $ d_{l} $ 连续变化,而 GQA 的压缩比受限于头数整除关系(通常为整数比)。
· 适用场景:GQA 更适合在不大幅修改模型结构的前提下,通过微调或简单权重转换获得加速;MLA 通常需要从预训练阶段就开始设计,因为潜空间的学习需要大量数据支撑,但对长文本的压缩效果更极致。
12. Paged Attention 的核心思想是什么?它解决了什么问题?¶
PagedAttention 是 vLLM 推理框架中的核心技术,灵感来源于操作系统中的虚拟内存和分页机制。它将 KV Cache 的管理从传统的连续静态分配,转变为按页(Page)动态分配。
传统 KV Cache 管理的问题:¶
静态预分配:推理前需要预估最大生成长度,并为每个请求预留一块连续的显存。这导致了严重的内部碎片(一个请求只用了部分预分配空间,剩余浪费)和外部碎片(显存中散布大小不一的空闲块,无法被新请求利用)。
内存低效:无法在请求之间灵活共享相同的 prompt 前缀的 KV Cache,每个请求都要独立存储一份,即使内容相同。
· 扩展性差:随着并发请求增多,碎片问题加剧,显存利用率极低,能支持的吞吐量受限。
PagedAttention 的解决方案:¶
将每个请求的 KV Cache 划分为固定大小的块(页),例如每页 16 或 32 个 token。这些页在显存中不必连续,通过一个页表来记录每页的物理地址。当请求的序列长度超过已分配的页时,系统按需分配新的物理页,并将页表更新。这样,显存分配变得极其灵活:
· 零内部碎片:按需分配,没有预分配浪费。
极低外部碎片:所有页大小相同,显存管理类似于伙伴系统,几乎不存在无法使用的碎片。
高效共享:多个请求的相同 prompt 前缀可以指向同一物理页,实现零拷贝共享。例如,在多轮对话中,系统 prompt 相同的所有请求可共享这部分 KV Cache,显著节省显存。
内存超额分配:类似于虚拟内存,可以允许逻辑上请求超过物理显存的总和,实际通过换负(如果有CPU后备)或限制调度来管理。
解决的问题:¶
· 显存利用率从传统方法的 30%~40% 提升到 90% 以上。
· 支持了极大规模的并发请求和动态 batching,吞吐量成倍提升。
为前缀缓存等高级特性提供了实现基础。
正是 PagedAttention 的这些优势,使得 vLLM 能够在相同的硬件上服务更多的用户,延迟更低,成为当前大模型推理服务的主流方案。
13. Paged Attention 中的“页”大小对推理性能有什么影响?¶
页大小(page size)是 PagedAttention 中一个细微但重要的超参数,直接影响显存碎片、吞吐量、延迟以及页表开销。
常见页大小:通常为 16、32 或 64 个 token。VLLM 默认是 16。
页大小过小(如8或更小)的影响:¶
· 页表开销大:页数增多,管理页表所需的内存和时间增加,GPU 上遍历页表的开销也可能显现。
可能导致显存碎片:虽然页是固定大小的,但过小的页可能导致分配和释放频繁,类似物理内存的“抖动”,且块太小可能无法高效利用GPU的连续数据读取特性,影响带宽利用率。
调度复杂度上升:调度器需要管理更多的页,可能降低调度效率。
页大小过大(如128或更大)的影响:¶
· 内部碎片严重:如果一个请求只多生成了几个 token 就结束,但分配了一整页 128 个 token 的空间,浪费的显存就很大。这又退回到了类似静态预分配的问题。
共享粒度粗糙:前缀缓存要求共享的单位是页。如果页太大,短前缀(如少于128 token)就无法被高效共享,因为即使只有前缀相同,也必须整个页对齐。
灵活性降低:按需分配时,过大的页使得请求的增长粒度太大,不能精细匹配实际的序列长度增长,导致显存利用率下降。
最优页大小的选取:需要在上述因素间权衡。通常根据典型的工作负载来调整:
对于平均生成长度较短的对话应用,较小的页(16 或 32)比较合适,因为可以精细控制显存。
对于需要超长生成的任务(如翻译整篇文章),较大的页(64 甚至 128)可能减少页表开销和调度次数。
· VLLM 团队通过大量实验选择了 16 作为默认值,这在多数场景下表现最佳。
工程上,可以通过性能 Profiling 来确定最优页大小:统计不同页大小下的 GPU 显存占用、吞吐、延迟和页表开销,绘制帕累托前沿进行选择。
14. vLLM 为什么能比传统静态分配显存的方案支持更大吞吐?¶
VLLM 的高吞吐源自多个层面的协同优化,核心是显存利用率和计算效率的双重提升。
1. Paged Attention 带来的高显存利用率¶
如前所述,传统静态分配产生大量碎片,显存的实际可利用率常常不足50%。VLLM的PagedAttention几乎消除了碎片,显存利用率可达90%以上。这意味着同样的硬件可以容纳更多的并发请求,从而直接提升吞吐。
2. 连续批处理(Continuous Batching)¶
传统推理框架在处理请求时,一旦一个 batch 中的某个请求生成结束,其占用的资源不会立即被回收,直到整个 batch 完成。这导致 GPU 在后期部分空闲。VLLM 动态地将新请求插入到正在进行中的 batch 中,当一个请求完成后,立即释放其占用的物理页,并立刻分配给队列中的下一个请求。这种“随到随服务”的模式避免了 GPU 的计算气泡,硬件始终被填满。
3. 高效的前缀共享¶
在多轮对话或批量处理中,大量请求共享相同的 system prompt 或初始上下文。VLLM 通过
PagedAttention 的页表,可以让这些请求在物理上共享同一份 KV Cache 页,无需为每个请求单独复制。这不仅节省了显存,也节省了 Prefill 阶段的计算量(共享部分只需计算一次),使得更多的请求可以同时被服务。
4. 精细的 GPU 内存管理与调度¶
VLLM 使用自定义的 CUDA 内存分配器,避免了 PyTorch 原生分配器的开销和碎片。同时,其调度器能够根据请求长度、优先级等进行智能排队和抢占,最大化整体吞吐。
定量对比:在相同硬件(如 A100)上,使用传统静态分配 + 静态批处理的方案,显存往往被大量未使用的预分配空间占据,实际支持的并发数可能只有 vLLM 的 1/3 到 1/2。vLLM 能轻松将吞吐提升 2-5 倍,具体倍数取决于工作负载。
15. 前缀缓存(Prefix Caching)是什么?在多轮对话中如何工作?¶
前缀缓存指的是在推理时,将多个请求共享的相同前缀部分的KV Cache只计算并存储一次,后续请求直接复用这部分缓存,而无需重新计算。
在多轮对话中的应用:¶
对话系统通常有一个固定的 system prompt(例如“你是一个有帮助的助手”),它在每次对话中都完全相同。
第一轮对话时,模型计算了整个 prompt(包括 system prompt 和第一轮用户问题)的 KV Cache。VLLM 会将 system prompt 对应的 KV Cache 页标记为共享,并保留在显存中。
第二轮对话开始时,新的请求包含了完全相同的 system prompt。调度器检测到此前缀已经缓存,便直接将这些物理页映射到新请求的页表中,而不是重新执行 Prefill 计算。
这样,第二轮及之后的对话只需要计算新增的对话轮次部分的 KV Cache,Prefill 时间大幅缩短,且显存占用不随轮次线性增长。
实现机制:VLLM 使用哈希表来记录已缓存的前缀。当一个新序列的 token 进入时,它逐步计算哈希值并尝试匹配已缓存的页。一旦找到不匹配的 token,就停止复用,后续的 token 生成新的页。
前缀缓存对于多轮对话、批量推理(如大量用户共享相同的上下文)或多次采样(同一个 prompt 生成多个回答)极为有效,能将首 token 延迟降低数倍,并节省显存。
16. 多轮对话时,不同轮次间如何复用和淘汰 KV Cache?¶
在多轮对话中,KV Cache 的管理是一个动态的过程,需要同时处理复用和淘汰,以在有限显存下保持高质量对话。
复用机制:¶
当开始新的一轮对话时,前一轮所有 token 的 KV Cache 已经存储在显存中。只要历史轮次没有被截断,当前轮次可以直接利用这些缓存,无需重新计算。这是最基本的前缀复用。
如果用户在多轮中重复提出相同的问题或命令,前缀缓存可以自动识别并复用已缓存的页。
淘汰机制(当显存不足时):
- LRU(最近最少使用):vLLM 会维护每个缓存页的访问时间。当需要为新 token 分配页面而显存不足时,它选择最长时间未被使用的页面进行淘汰(释放或写回 CPU)。在多轮对话中,较早的对话轮次如果长时间未被重新引用,就可能被逐出。
基于生命周期的淘汰:每个请求有最大长度限制,一旦请求结束,其所有页被回收。在对话系统中,如果某个对话会话结束或用户长时间无交互,相关缓存可以被主动清除。
重计算回退:如果某些历史KV Cache被淘汰,但后续又需要用到(例如用户要求回顾之前的内容),模型将不得不重新计算那些被淘汰的token的KV Cache。这会导致额外的Prefill开销,但保证了无限对话在有限显存下的可行性。
树状缓存结构:在一些更复杂的系统中(如 SGLang),KV Cache 以树状结构组织,不同轮次的回复分支可以共享同一个“主干”(system prompt 和共享的对话历史),而各自的分支(不同轮次的回答)分别存储。淘汰时可以修剪掉不活跃的分支,保留共享的主干,实现更精细的显存控制。
最终目标是在对话质量和显存消耗间取得平衡,确保在有限硬件上能支持尽可能长的多轮对话,同时不因频繁重计算而影响用户体验。
17. 如何对 KV Cache 进行量化?KV8、KV4 会带来什么影响?¶
KV Cache 量化是将缓存的 Key 和 Value 张量从高精度(FP16/BF16)压缩到低精度(如 INT8、INT4)的技术,直接减小 KV Cache 的显存占用和带宽压力。
量化方法:¶
对称量化:计算每个 Key/Value 张量的最大绝对值,映射到 INT8 范围 [-127,127]。这是最基本的方法,实现简单。
逐 token 量化或逐通道量化:由于不同 token 的 KV 分布差异很大,通常按 token 或按通道(head_dim 维度)分别计算量化系数,以减小精度损失。
校准:通常需要一小批校准数据,统计 Key 和 Value 的值分布,确定量化参数(scale, zero-point)。校准可以在推理前离线完成,也可以动态统计(在线量化),后者对分布变化更鲁棒但稍微增加计算。
• 解量化:在注意力计算前,需要将量化后的 K、V 反量化回浮点进行矩阵乘法。因此,KV 量化的真正收益在于显存容量和读取带宽的节省,而非计算量的减少。
KV8 (INT8 量化) 的影响:¶
- 显存和带宽:KV Cache 大小直接减半,带宽需求减半,这在 Decode 阶段非常显著,通常能带来 20-40% 的延迟降低和更高的吞吐。
• 精度损失:对于大多数模型,INT8 量化几乎不会造成可察觉的生成质量下降,因为 Key 和 Value 的数值范围相对集中。
KV4 (INT4 量化) 的影响:¶
更极端的压缩:KV Cache 大小变为原来的 1/4,对长文本推理的收益巨大。
• 精度损失可能显著:INT4 量化精度较低,可能会影响注意力分数的准确性,导致模型在需要精细理解或长程记忆的任务上性能下降,如困惑度升高、多跳推理失败等。但很多研究通过更精细的量化策略(如分组量化、非对称量化、混合精度)将损失控制在可接受范围。
适用场景:¶
若目标是最大化吞吐和长文本长度,KV4 是有吸引力的选项。
若更注重保持原始模型质量,通常选择KV8或FP8。
目前 vLLM 等框架已经支持 KV Cache 的 FP8 和 INT8 量化,INT4 也有实验性实现。
18. 有没有可能对 KV Cache 进行稀疏化或按重要性剪枝?¶
是的,这是一个活跃的研究方向,旨在进一步降低KV Cache 的大小。其核心思想是:并非所有 token 都对未来的生成同等重要,我们可以在推理时动态决定保留哪些 token 的 KV Cache,丢弃或合并不重要的。
主要技术路线:¶
1. 基于注意力分数的剪枝¶
• StreamingLLM / H2O: 观察发现,Decoder 的注意力通常高度集中在最近的 token 和一些关键的“注意力汇聚点”(attention sinks,如起始 token)。这些方法只保留最近的几个 token 和少量高分的历史 token,丢弃中间注意力低的 token 的 KV Cache。例如 H2O 保留最近 512 个 token 和历史上注意力分数最高的 256 个 token。
优点:实现简单,可显著减少缓存。
· 缺点:丢弃的 token 无法恢复,当后续生成需要那些被丢弃的信息时,会引入误差。
2. 令牌合并(Token Merging)¶
将相似的或注意力模式相近的连续 token 进行合并,使用一个代表 token 的 KV 代替多个,类似于图像中的 token 合并。这在 DiT 等扩散模型中也得到应用,对于 LLM 可以减少冗余。
3. 动态稀疏注意力(如 Quest)¶
在解码时,通过近似算法(如 locality-sensitive hashing)快速找到与当前 Query 最相关的 Key,只加载那些 Key 对应的完整 KV 块,忽略其余。这更像一种稀疏加载,而非永久丢弃,更适合长上下文。
4. 训练感知的剪枝¶
在训练阶段就引入对 KV Cache 大小的正则化,例如鼓励 Key/Value 向量的稀疏性,或使用可学习的 mask 动态决定哪些层、哪些头可以丢弃 token。推理时模型自己就能决定保留哪些。
挑战与权衡:¶
精度-效率权衡:激进的剪枝必然丢失信息,可能损害生成质量,尤其在需要精确事实回忆时。
实现复杂度:高效的稀疏读写需要专门的 GPU kernel 支持,否则内存访问的不连续反而可能降低速度。
· 动态性:剪枝策略需要随着生成过程动态调整,且不能带来过高的决策开销。
尽管有这些挑战,稀疏化和剪枝是对 KV Cache 压缩的重要补充。特别是与量化结合时,可以取得“1+1>2”的效果,让百亿模型在消费级显卡上处理超长文本成为可能。
19. StreamingLLM是如何让模型支持无限长输入的?它保留了哪些token的KV Cache?¶
在标准的自回归生成中,我们通常将整个对话历史的所有 token 的 KV Cache 都保存在显存里。当序列长度超过预训练时的上下文窗口时,模型会出现严重的性能退化甚至无法工作。StreamingLLM 发现,这种退化并非因为模型没有能力处理更长的序列,而是因为我们在移除旧 token 时,不经意间丢掉了一些至关重要的“锚点”。
研究人员在分析注意力机制时观察到,初始的几个 token(特别是序列最开始的 1-4 个 token)往往承担着“注意力汇聚点”的角色。无论后续生成多少个 token,注意力权重都会稳定地分配相当大的一部分给这些起始 token。它们像是整段文本的“根基”,一旦被移除,整个注意力的分布就会崩塌,导致模型输出混乱。
✈ 基于这个发现,StreamingLLM 的做法非常直接:始终保留序列最开始的几个 token(通常取 4 个)作为“注意力锚点”,再加上一个固定大小的滑动窗口(比如最近的 512 个 token)。所有位于这两者之间的旧 token,其 KV Cache 会被直接丢弃。
∞ 这样做之所以能支持无限长输入,是因为 KV Cache 的大小被严格限制为了窗口大小 + 锚点数量,不再随序列总长度增长。同时,由于保留了那关键的几个初始 token,注意力机制的稳定性得以维持,模型能够继续生成连贯的文本,即使输入的长度远远超出了训练时的上下文窗口。
用工程直觉来理解:这几颗锚点就像是你在读一本很长的书时,始终记得的“书名”和“第一章的概要”。就算你忘记了中间的第200页具体讲了什么,只要还记得书名和开篇主题,就能继续读下去而不至于完全摸不着头脑。StreamingLLM就是强迫模型用这种“只看头和尾”的方式来处理超长序列。
20. 🔔 H2O 等淘汰策略如何判断哪些 token 的 KV Cache 可以丢弃?¶
H2O(Heavy-Hitter Oracle)这类方法比 StreamingLLM 更进一步,它不预设哪些 token 重要,而是让模型自己在生成过程中“告诉”我们哪些 token 是关键信息。
H2O 的核心思想来自一个观察:在注意力机制中,少数 token 攫取了绝大部分的注意力权重,而大量的 token 获得的关注度极低。这些“重击手”(Heavy Hitters)是模型进行推理和生成时真正依赖的信息节点。
具体做法如下:在生成每个新 token 时,我们都会得到一个当前 token 对所有历史 token 的注意力权重分布。H2O 会累积每个历史 token 被关注的总得分。例如,可以计算每个历史 token 的 Key 向量与所有 Query 的点积之和,作为其“重要性分数”。生成过程中,系统维护一个固定大小的 KV Cache 预算(比如 512 个位置)。当缓存满了需要淘汰旧 token 时,它就毫不留情地把那些累积重要性分数最低的 token 踢出去,只保留分数最高的那一批。
这种方法比固定窗口更聪明,因为它能根据具体内容动态调整记忆。比如,在阅读一本小说时,故事的主人公名字和关键情节会在注意力中反复出现,H2O自然会把它们保留下来;而一些环境描写或过渡语句,因为很少被后续内容引用,就会被优先遗忘。
21. 滚动窗口缓存(Sliding Window)是怎么限制 KV Cache 大小的?会丢失什么信息?¶
滚动窗口缓存是一种最简单的KV Cache管理策略。它的规则就一条:只保留最近生成的N个token的KV Cache,任何超出这个窗口的旧token全部丢弃。
这个 N,比如 512 或 4096,就是窗口大小。它的优势极其明显:实现简单,显存占用完全可预测且恒定,无论输入有多长,都不会出现显存爆炸。Llama 2 等模型在预训练时如果使用了这种注意力机制,推理时就可以无缝启用。
但它丢失的信息也很明确:所有超出窗口的历史上下文。这意味着模型完全无法直接引用窗口之外的任何信息。如果一个关键事实在对话的第200轮被提到,而窗口只有100,那么当用户在第201轮提及时,模型会完全不记得。它就像一个只能记住最近5分钟对话的人,无法处理需要长程记忆的任务,比如整本书的总结、复杂的多跳推理等。
与 StreamingLLM 相比,它连“锚点”也没有,所以当窗口滑动时,可能会突然丢失一些重要的全局上下文,导致注意力分布发生突变。
22. 为什么在长文本推理中,K 和 V 的压缩手段有时比注意力计算优化更紧迫?¶
这是一个很好的工程经济学问题。我们用一个具体的硬件数字来理解。以 NVIDIA A100 为例,它的显存带宽是 2 TB/s。在生成阶段(Decode),每生成一个 token,我们需要将整个 KV Cache 从显存读出。
假设我们正在处理一段 32k tokens 的上下文,使用一个 7B 模型(MHA)。此时的 KV Cache 大约是 16 GB。生成一个 token,仅读取 KV Cache 就需要耗费 16 GB / 2 TB/s = 8 毫秒。而实际的计算时间,因为 Decode 阶段只处理一个新 token,矩阵乘法极小,可能只需要 0.5 毫秒。所以,超过 90% 的时间都花在了等待数据上,而不是计算上。
如果我们能把 KV Cache 压缩 4 倍,变成 4 GB,那么读取时间直接降到 2 毫秒,总延迟会大幅下降。而注意力计算本身,即使我们把它优化到极致,也省不了这 8 毫秒。这就是为什么压缩 KV Cache(减少数据搬运量)在 Decode 阶段是比优化计算更紧迫、更有效的手段。带宽才是瓶颈,算力还在闲置。
23. 请分析在 batch 推理时,不同请求序列长度不一对 KV Cache 管理带来的挑战。¶
在服务端批量处理多个用户请求时,一个很棘手的问题是:每个请求的序列长度都不同,而且差异可能极大。这给KV Cache管理带来了三个层面的挑战。
第一,显存碎片化。如果我们为每个请求都按照 batch 中的最大长度来预分配 KV Cache,那对于那些短请求,大量显存就会被白白浪费(内部碎片)。但如果采用动态分配,不同请求不断追加新的 KV Cache 块,又会产生大量大小不一的外部碎片,就像磁盘用久了会变得支离破碎一样。这正是 vLLM 的 PagedAttention 要解决的核心问题——通过统一大小的页来避免碎片。
第二,计算气泡。一个 batch 中的请求必须等到所有请求都完成当前步的计算,才能一起进入下一步。如果一个请求很快就生成完了(比如回答一个简单的是非题),而另一个请求还在生成长篇大论,那么已经完成的请求所占用的计算资源就会空闲,等待整个 batch 结束。这就是所谓的“计算气泡”,它会拉低 GPU 的利用率。
第三,调度复杂度。系统需要不断地做决策:哪些请求应该被组合到同一个 batch 里?何时应该把完成请求的资源回收,并立即分配给队列里的新请求?这种动态调度需要极低的开销,否则调度本身就会成为新的瓶颈。连续批处理(Continuous Batching)就是为了解决这个问题而生的,它允许请求随时加入和离开batch,最大化 GPU 的“工作时间”。
24. 推理时是否可以动态调整 GQA 的分组数?为什么?¶
答案是:不能,分组数 G 是模型架构的一部分,一旦训练好就固定了。
GQA 的 Key 和 Value 投影矩阵的权重形状直接依赖于分组数 G。在预训练或微调时,模型已经学会了如何在这个特定的 G 下分配注意力模式。如果你在推理时试图改变 G(比如从 G=8 变成 G=4),那么原来的多个 K、V 头就需要被重新合并,而这种合并是模型从未训练过的。这会直接破坏模型的效果,导致输出质量急剧下降。
不过,有一种“技巧”可以在推理时实现类似的效果,那就是将标准MHA模型通过权重转换的方式变成GQA模型。你可以把一个训练好的MHA模型的多个K、V头按组求平均,强行合并成更少的头,从而得到一个GQA模型。这样得到的模型通常需要少量的微调来恢复性能,但一旦转换完成,它的G也就固定了。整个过程改变了模型架构,而不是在推理时动态调整。
25. 你认为未来 KV Cache 优化会朝哪个方向发展?¶
展望未来,我认为KV Cache 优化会沿着两条主线发展:极致的压缩和智能的记忆管理。
在压缩方面,我们会看到越来越多的模型在训练阶段就内置了KV Cache的压缩机制。MLA(多头潜在注意力)就是一个很好的开端,它通过低秩分解将KV Cache压缩到原始大小的几分之一。未来的方向可能是更通用的自适应压缩:模型能够学会根据当前上下文的信息密度,动态决定将KV Cache压缩到什么程度。比如,对于信息量低的闲聊,可以极度压缩;对于需要精确记忆的推理,则保留更多细节。
在记忆管理方面,单纯依赖滑动窗口或简单的累积分数淘汰已经不够。我们会走向更类人的记忆系统:模型会有一个短期记忆(类似窗口缓存),一个长期记忆(类似结构化的知识库或外部记忆),以及一个控制器来决定何时将短期记忆中的信息“巩固”到长期记忆中。这将结合检索增强生成(RAG)和记忆网络的思想,让模型既能记住最近的对话细节,又能随时调用遥远的、关键的信息。
最终,KV Cache 可能会从一个被动的缓存,演变成一个主动的、智能的“记忆管理器”,它能够总结、合并、遗忘和检索信息,真正让大模型在长序列推理中做到效率和能力的统一。