LLM高效推理：KV缓存与分页注意力机制深度解析
<div id="container" data-v-1d7a5742="" data-element="root" contentScore="2356">随着大型语言模型（LLM）规模和复杂性的持续增长，高效推理的重要性日益凸显。KV（键值）缓存与分页注意力是两种优化LLM推理的关键技术。本文将深入剖析这些概念，阐述其重要性，并探讨它们在仅解码器（decoder-only）模型中的工作原理。常规推理机制
首先，我们通过一个简单的例子来理解Transformer模型中典型的推理过程。假设我们需要生成短语：
“The quick brown fox jumped”
以下是常规推理的简化实现：
import numpy as np# 简化的嵌入表示，仅用于演示embeddings = { 'The': np.array(), 'quick': np.array(), 'brown': np.array(), 'fox': np.array(), 'jumped': np.array()}# 权重矩阵（简化）W_Q = W_K = W_V = np.array([, , , ])def compute_attention(self, input_words): # 将单词转换为嵌入向量 E = np.array( for word in input_words]) # 计算所有token的K和V矩阵 K = E @ W_K# 形状: (seq_len, 2) V = E @ W_V# 形状: (seq_len, 2) # 计算最后一个token的Q矩阵 Q = E[-1] @ W_Q# 形状: (1, 2) # 计算缩放的点积注意力得分 scale = np.sqrt(2)# 缩放因子，为key/query维度（此处为2）的平方根 scores = (Q @ K.T) / scale# 形状: (1, seq_len) # 应用Softmax函数，获得注意力权重 attention_weights = self.softmax(scores)# 形状: (1, seq_len) # 将注意力权重应用于V矩阵 output = attention_weights @ V# 形状: (1, 2) return output
[*]1.
[*]2.
[*]3.
[*]4.
[*]5.
[*]6.
[*]7.
[*]8.
[*]9.
[*]10.
[*]11.
[*]12.
[*]13.
[*]14.
[*]15.
[*]16.
[*]17.
[*]18.
[*]19.
[*]20.
[*]21.
[*]22.
[*]23.
[*]24.
[*]25.
[*]26.
[*]27.
[*]28.
[*]29.
[*]30.
[*]31.
[*]32.
[*]33.
[*]34.
[*]35.
[*]36.
[*]37.
[*]38.
以下是逐步生成的过程：
# 步骤1: 生成 "brown"input_words_step1 = ['The', 'quick']output_step1 = compute_attention(input_words_step1)# 步骤2: 生成 "fox"input_words_step2 = ['The', 'quick', 'brown']output_step2 = compute_attention(input_words_step2)# 步骤3: 生成 "jumped"input_words_step3 = ['The', 'quick', 'brown', 'fox'] output_step3 = compute_attention(input_words_step3)
[*]1.
[*]2.
[*]3.
[*]4.
[*]5.
[*]6.
[*]7.
[*]8.
[*]9.
冗余计算：观察上述代码可以发现对于每个新生成的token：
[*]需要为所有先前的token重新计算K和V矩阵。
[*]矩阵的大小随着token数量的增加而增大。
[*]存在大量不必要的重复计算。
KV缓存机制
当使用Transformer模型生成文本时，通过缓存键（K）和值（V）矩阵，可以显著优化推理过程。下图展示了KV缓存的工作原理：
在上图中：
[*]q_new表示最新token的查询向量。
[*]K_prev和V_prev是从先前计算中缓存得到的键和值矩阵。
[*]k_new和v_new仅为当前新token计算。
[*]蓝色箭头表示如何利用缓存值和新值计算注意力。
以下是KV缓存的实现示例：
def compute_attention_with_cache(self, input_words): """使用KV缓存计算注意力""" # 获取新token（序列中的最后一个单词） new_word = input_words[-1] e_new = embeddings # 计算新token的K和V矩阵 K_new = e_new @ W_K# 形状: (2,) V_new = e_new @ W_V# 形状: (2,) # 更新缓存的K和V矩阵 if self.cached_K is None: self.cached_K = K_new.reshape(1, -1)# 形状: (1, 2) self.cached_V = V_new.reshape(1, -1)# 形状: (1, 2) else: self.cached_K = np.vstack()# 形状: (seq_len, 2) self.cached_V = np.vstack()# 形状: (seq_len, 2) # 计算最后一个token的Q矩阵 Q = e_new @ W_Q# 形状: (2,) # 使用缓存的K矩阵计算缩放的点积注意力得分 scale = np.sqrt(2)# 缩放因子，为key/query维度（此处为2）的平方根 scores = (Q @ self.cached_K.T) / scale# 形状: (1, seq_len) # 应用Softmax函数，获得注意力权重 attention_weights = self.softmax(scores)# 形状: (1, seq_len) # 使用缓存的V矩阵计算注意力输出 output = attention_weights @ self.cached_V# 形状: (1, 2) return output
[*]1.
[*]2.
[*]3.
[*]4.
[*]5.
[*]6.
[*]7.
[*]8.
[*]9.
[*]10.
[*]11.
[*]12.
[*]13.
[*]14.
[*]15.
[*]16.
[*]17.
[*]18.
[*]19.
[*]20.
[*]21.
[*]22.
[*]23.
[*]24.
[*]25.
[*]26.
[*]27.
[*]28.
[*]29.
[*]30.
[*]31.
[*]32.
以下是逐步生成的过程：
# 步骤1: 生成 "brown"input_words_step1 = ['The', 'quick']output_step1 = compute_attention_with_cache(input_words_step1)# 步骤2: 生成 "fox"input_words_step2 = ['The', 'quick', 'brown']output_step2 = compute_attention_with_cache(input_words_step2)# 步骤 3: 生成 "jumped"input_words_step3 = ['The', 'quick', 'brown', 'fox'] output_step3 = compute_attention_with_cache(input_words_step3)
[*]1.
[*]2.
[*]3.
[*]4.
[*]5.
[*]6.
[*]7.
[*]8.
[*]9.
比较有无KV缓存的推理计算
内存需求与挑战
我们来看一个使用典型模型参数的实际例子：
[*]序列长度: 4096
[*]层数: 32
[*]注意力头数: 32
[*]头维度: 128
[*]精度: FP16 (2 bytes)
每个token所需的内存：
KV_cache_per_token = 2×num_layers×(num_heads×head_dim)×precision = 2 × 32 × (32 × 128) × 2 bytes = 2 × 32 × 4096 × 2 bytes = 524,288 bytes ≈ 0.5 MB
[*]1.
[*]2.
[*]3.
[*]4.
[*]5.
KV缓存的低效性
尽管KV缓存显著提高了计算效率，但它也带来了内存管理方面的挑战。以下是三种主要的内存低效类型：
内部碎片
[*]由因未知输出长度而导致的过度分配引起。
[*]示例：在图像中，2040个槽位从未被使用。
[*]影响：可能浪费高达60-80％的已分配内存。
[*]解决方案：更精确的输出长度估计或动态分配策略。
预留浪费
[*]为将来的token生成而预留的内存。
[*]在图像中显示为“3 slots future used (reserved)”。
[*]维持生成连续性的必要措施。
[*]可以通过更好地预测所需的未来槽位来优化。
外部碎片
[*]由处理具有不同序列长度的多个请求导致。
[*]在不同请求之间创建内存间隙。
[*]解决方案包括内存碎片整理和智能请求批处理。
如上图所示，通常仅有20-40％的KV缓存被用于存储实际的token状态。
分页注意力：解决内存低效的方案
为了应对这些内存挑战，可以采用分页注意力机制。
分页注意力是一种用于有效处理Transformer模型中长序列的技术，它通过将注意力计算分解为更小、更易于管理的“页”或“块”来实现。这种方法降低了内存消耗和计算复杂度，从而能够处理原本因过大而无法放入内存的序列。
def compute_attention_with_paging(self, input_words): """使用分页KV缓存计算注意力""" # 获取新token（序列中的最后一个单词） new_word = input_words[-1] e_new = embeddings # 计算新token的K和V矩阵 K_new = e_new @ W_K# 形状: (2,) V_new = e_new @ W_V# 形状: (2,) # 确定当前页的索引 total_tokens = sum(len(K_page) for K_page in self.cached_K_pages) + 1 current_page_idx = (total_tokens - 1) // PAGE_SIZE # 如果需要，初始化新页 if len(self.cached_K_pages)
页:
[1]