1. 从“慢吞吞”到“飞起来”为什么我们需要KV Cache和GQA如果你玩过最近几年流行的大语言模型不管是写代码、写文章还是聊天肯定都遇到过这种情况模型在生成第一个词的时候好像要“思考”很久但一旦开始后面的词就“哗啦啦”地出来了。这背后的“功臣”就是我们今天要手撕的KV Cache。而另一个技术GQA则是为了让这个“功臣”变得更轻便、更高效而生的。想象一下你正在读一本很厚的书每次读新的一页你都需要把前面所有页的内容都重新翻看一遍才能理解当前页在讲什么。这效率得多低啊Transformer模型在生成文本比如让AI续写故事时最初就是这么干的。它每生成一个新的词token都需要把之前生成的所有词连同你给的提示词全部重新计算一遍注意力。这个过程计算量巨大尤其是当序列越来越长时速度会呈平方级下降这就是所谓的O(n²)复杂度。KV Cache解决的正是这个“重复劳动”的问题。它的核心思想特别像我们人类的短期记忆把之前已经计算过的、不会改变的信息存起来下次直接用。在Transformer的解码过程中对于给定的历史序列其对应的Key和Value张量在后续生成步骤中是固定不变的。那么何不把它们缓存起来呢这样模型在生成下一个词时只需要计算新词的Query然后去和缓存里所有的Key、Value做注意力计算就行了。这相当于把每次的“全文重读”变成了“查阅笔记”推理速度尤其是在长文本生成时会有质的飞跃。而GQA则是在KV Cache的基础上进一步做“内存优化”。传统的多头注意力机制中每个注意力头都有一组独立的Key、Value。在自回归生成时这些KV对都要被缓存内存开销很大。GQA的想法很巧妙让多个Query头共享同一组Key和Value头。这就好比一个小组讨论不是每个人都要准备一套完整的资料而是几个人共享一份核心资料然后各自从自己的角度Query去提问和理解。这样需要缓存的KV数量就减少了在几乎不影响模型效果的前提下显著降低了内存占用和带宽压力让推理更快、更省资源。简单来说KV Cache让你不用重复计算GQA让你少存点东西。两者结合是让大模型在实际部署中能够流畅运行的关键技术。接下来我们就抛开理论直接上手用代码把它们实现出来看看它们到底是怎么工作的。2. 手撕KV Cache告别重复计算的“内存小本本”2.1 KV Cache到底在缓存什么我们先把视角拉回到最基础的Scaled Dot-Product Attention计算。注意力机制的核心公式是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。在自回归生成比如GPT那种逐词生成的场景下假设我们已经生成了前t-1个token现在要生成第t个token。在这个过程中第t步的输入其实只有一个新tokenseq_len1。但是为了计算这个新token应该关注之前哪些部分我们需要它作为Query去和之前所有的token作为Key进行匹配。这里的“之前所有的token”既包括最初的提示词也包括模型自己刚刚生成的第1到第t-1个token。关键点来了从第1步到第t-1步每一步都已经计算过它们自己的Key和Value了而且这些Key和Value在后续生成中永远不会改变。那么如果我们把这些历史步骤的Key和Value都保存下来在第t步时就只需要计算当前新token的Key和Value然后和历史缓存拼接起来一起参与注意力计算。这就是KV Cache的全部秘密。我刚开始理解的时候总喜欢把它想象成一个不断变长的“记忆磁带”。每生成一个词就把这个词对应的Key和Value“录制”到磁带末尾。下次生成时只需要播放整盘磁带然后加上新录制的那一小段就行了。这个“磁带”就是我们的Cache。2.2 维度变换与代码拼接torch.cat的魔法光说不练假把式我们直接看代码。首先我们要明确一下在代码中这些张量长什么样。在标准的实现里经过线性投影和切头后Q、K、V的形状通常是[batch_size, num_heads, seq_len, head_dim]。batch_size一次处理多少条数据。num_heads注意力头的数量让模型可以从不同角度理解信息。seq_len当前这一步序列的实际长度。在自回归生成中除了第一步通常都是1。head_dim每个注意力头对应的向量维度通常是d_model / num_heads。现在假设我们正在进行第t步生成。我们手头有past_key缓存的历史Key形状为[batch_size, num_heads, past_seq_len, head_dim]。其中past_seq_len t-1。current_k当前新token计算出的Key形状为[batch_size, num_heads, 1, head_dim]。我们的任务就是把它们拼起来形成完整的Key序列。注意拼接发生在seq_len这个维度上也就是张量的第2维因为batch是第0维num_heads是第1维。在PyTorch中我们使用torch.cat函数# 将当前步的Key与缓存的Key拼接 if past_key is not None: k torch.cat([past_key, k], dim2) # dim2 对应 seq_len 维度 # Value的拼接完全同理 if past_value is not None: v torch.cat([past_value, v], dim2)拼接完成后新的k和v的形状就变成了[batch_size, num_heads, past_seq_len 1, head_dim]。你看只有序列长度这个维度发生了变化其他维度都保持不变。这个操作虽然简单但却是KV Cache高效性的基石——它用一次简单的内存拼接替代了海量的重复计算。2.3 完整流程模拟一步步拆解生成过程让我们模拟一个完整的文本生成流程把KV Cache的作用看得更清楚。假设我们的提示词是“人工智能”模型要生成“是未来”。第一步处理提示词 “人工智能”输入input_ids对应“人工智能”假设长度为4。past_key,past_value初始为None。模型计算“人工智能”这四个字对应的 K1, V1, K2, V2, K3, V3, K4, V4。输出模型预测的下一个词比如“是”。缓存更新将计算出的所有Key和ValueK1~V4缓存起来作为下一步的past_key和past_value。第二步生成第一个词 “是”输入上一步生成的词“是”长度为1。past_key,past_value形状为[batch, heads, 4, dim]即“人工智能”的缓存。模型计算“是”这个字对应的 K5, V5。拼接操作k cat([past_key, K5], dim2)形状变为[batch, heads, 5, dim]。Value同理。注意力计算用“是”的Query去和拼接后的5个Key/Value做注意力从而决定输出。输出模型预测的再下一个词比如“未”。缓存更新将拼接后的完整K和VK1~K5, V1~V5缓存起来。第三步及之后重复第二步的过程。每一步序列长度seq_len都增加1缓存也越来越长。但无论多长模型每一步都只需要计算当前一个token的K和V计算复杂度从O(n²)降到了O(n)。这就是为什么模型生成会越往后“感觉”越快的原因因为主要开销变成了读取和拼接缓存而非重复计算。2.4 代码实现与踩坑指南下面是一个支持KV Cache的多头注意力模块的完整代码。我对比了原始MHA只在forward函数中增加了对past_key和past_value的处理以及返回当前的K/V供下一步使用。import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttentionWithKVCache(nn.Module): def __init__(self, d_model512, num_heads8): super().__init__() assert d_model % num_heads 0, d_model must be divisible by num_heads self.d_model d_model self.num_heads num_heads self.head_dim d_model // num_heads # 线性投影层 self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, q, k, v, past_keyNone, past_valueNone, maskNone): batch_size, seq_len, _ q.shape # 1. 线性投影 q self.q_proj(q) k self.k_proj(k) v self.v_proj(v) # 2. 重塑为多头 [batch, num_heads, seq_len, head_dim] q q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 3. KV Cache 核心在seq_len维度拼接历史缓存 if past_key is not None: k torch.cat([past_key, k], dim2) # dim2 是 seq_len 维度 if past_value is not None: v torch.cat([past_value, v], dim2) # 保存当前的K/V作为下一步的“历史缓存” present (k, v) # 4. 缩放点积注意力 attn_scores torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) if mask is not None: # 通常用于掩盖未来信息解码器自注意力 attn_scores attn_scores.masked_fill(mask.unsqueeze(1), -1e9) attn_weights F.softmax(attn_scores, dim-1) attn_output torch.matmul(attn_weights, v) # 5. 合并多头输出 attn_output attn_output.transpose(1, 2).contiguous() attn_output attn_output.view(batch_size, seq_len, self.d_model) # 最终投影 output self.out_proj(attn_output) return output, present # 返回输出和当前的K/V缓存几个容易踩的坑拼接维度一定要确认是在dim2即seq_len维度拼接。我一开始在dim1num_heads维度折腾了半天结果完全不对。掩码处理在自回归生成中mask通常是一个下三角矩阵包括对角线用来防止当前位置关注到未来的词。在拼接了KV Cache后mask的形状也需要相应调整以匹配新的序列长度。上面的代码中mask.unsqueeze(1)是为了广播到多头维度。缓存初始化第一次调用时past_key和past_value是None直接使用当前计算的k和v。从第二次开始才需要进行拼接。内存增长这是KV Cache最大的代价。缓存会随着生成序列的长度线性增长。对于非常长的对话或文档生成需要有相应的缓存管理或丢弃策略。提示在实际的Transformer解码器如GPT中q、k、v的输入通常都是同一个张量自注意力并且会有一个更外层的循环来管理每一步的输入token和传递past_key_values。我们这个模块是其中最核心的注意力计算部分。3. 深入GQA用“小组讨论”优化注意力3.1 MHA、MQA与GQA一场关于共享的权衡在弄明白GQA之前我们得先看看它的“前辈们”MHA和MQA。MHA标准的多头注意力。num_heads个头每个头都有自己独立的Q、K、V投影矩阵。好处是模型容量大能捕捉丰富的信息。坏处是在自回归生成中每个头的K和V都要缓存内存和带宽开销大。MQA多查询注意力。这是另一个极端num_heads个Query头但只共享一组Key和Value头。这极大地减少了KV缓存的大小但实验发现这通常会导致模型性能的明显下降因为Key/Value的表达能力被严重削弱了。那么有没有折中的方案呢GQA应运而生。它的全称是Grouped-Query Attention即分组查询注意力。它的思想非常直观我们不搞“一人一套”也不搞“全班共用一套”而是分小组共享。具体来说我们把所有的Query头分成G个组。每个组内的所有Query头共享同一组Key头和Value头。这样Key/Value头的数量num_kv_heads就等于分组数G它小于Query头的数量num_heads但大于1。类型Query头数量Key/Value头数量共享关系特点MHAHH1对1性能好但KV缓存大MQAH1所有Q共享1个K/VKV缓存极小但性能损失大GQAHG (1 G H)每G/H个Q共享1个K/V平衡性能和内存实践中的优选你可以把它想象成公司开会。MHA是每个人Query头都自带一份完整的报告K/VMQA是只有老板有一份报告所有人都对着这份报告提问而GQA是分部门开会每个部门组内部共享一份核心报告部门成员从不同角度Query进行讨论。这样既保证了讨论的深度每组有自己的焦点又避免了材料KV缓存的过度冗余。3.2 GQA的实现核心repeat_interleave的妙用GQA的实现相比MHA主要改动在于K和V的投影层以及后续的“扩展”操作。我们直接上代码结合注释来看import torch import torch.nn as nn import torch.nn.functional as F class GQA(nn.Module): def __init__(self, embed_dim, num_heads, num_groups4, head_dimNone, biasTrue): super().__init__() self.embed_dim embed_dim self.num_heads num_heads # Query头的总数 H self.num_groups num_groups # 分组数 G也是KV头的数量 self.head_dim head_dim if head_dim else embed_dim // num_heads # 关键检查Query头数必须能被分组数整除 assert num_heads % num_groups 0, num_heads must be divisible by num_groups # 投影层Q的输出维度是 H * head_dim K/V的输出维度是 G * head_dim self.q_proj nn.Linear(embed_dim, num_heads * self.head_dim, biasbias) self.k_proj nn.Linear(embed_dim, num_groups * self.head_dim, biasbias) self.v_proj nn.Linear(embed_dim, num_groups * self.head_dim, biasbias) self.out_proj nn.Linear(num_heads * self.head_dim, embed_dim, biasbias) def forward(self, q, k, v, maskNone, past_keyNone, past_valueNone): batch_size, seq_len, _ q.shape # 1. 线性投影 q self.q_proj(q) # [B, L, H*D] k self.k_proj(k) # [B, L, G*D] v self.v_proj(v) # [B, L, G*D] # 2. 重塑并转置为多头格式 q q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # [B, H, L, D] # 注意K和V的头数是 num_groups (G) k k.view(batch_size, seq_len, self.num_groups, self.head_dim).transpose(1, 2) # [B, G, L, D] v v.view(batch_size, seq_len, self.num_groups, self.head_dim).transpose(1, 2) # [B, G, L, D] # 3. 支持KV Cache和之前逻辑完全一样 if past_key is not None: k torch.cat([past_key, k], dim2) if past_value is not None: v torch.cat([past_value, v], dim2) present (k, v) # 4. GQA核心将K/V沿“头”维度扩展以匹配Q的数量 # 例如H8, G4则每个KV头需要被复制 8//4 2 次 expansion_factor self.num_heads // self.num_groups k k.repeat_interleave(expansion_factor, dim1) # [B, G, L, D] - [B, H, L, D] v v.repeat_interleave(expansion_factor, dim1) # [B, G, L, D] - [B, H, L, D] # 5. 标准的缩放点积注意力计算此时Q, K, V的head维度都是H scale self.head_dim ** -0.5 attn torch.matmul(q, k.transpose(-2, -1)) * scale # [B, H, L, L] if mask is not None: attn attn.masked_fill(mask.unsqueeze(1), float(-inf)) attn F.softmax(attn, dim-1) x torch.matmul(attn, v) # [B, H, L, D] # 6. 合并多头并输出投影 x x.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) # [B, L, H*D] x self.out_proj(x) return x, present代码中最精妙的部分是第4步的repeat_interleave。这里有一个常见的疑问为什么用repeat_interleave而不是repeatrepeat是在整个张量的某个维度上进行整体复制。比如k.repeat(1, 2, 1, 1)会把[B, G, L, D]在头维度复制成[B, 2*G, L, D]但这不符合“分组内共享”的语义。它会把第一个KV头复制成两个连续的、完全相同的头但这可能不对应分组关系。repeat_interleave是“交错重复”它是在元素级别进行复制。当expansion_factor2时它会把第一个KV头复制2份紧接着把第二个KV头复制2份……这样复制后的第1、2个头对应原第1个KV头即第一组第3、4个头对应原第2个KV头即第二组完美实现了“组内共享组间不同”的GQA思想。3.3 效果对比与参数选择为了直观感受GQA带来的收益我们可以做一个简单的对比。假设模型d_model1024,num_heads16序列长度seq_len1024head_dim64。MHA需要缓存的KV总参数量 batch * num_heads * seq_len * head_dim * 2(K和V)。对于单batch就是1 * 16 * 1024 * 64 * 2 2,097,152个标量。GQA (G4)num_kv_heads 4。缓存参数量 1 * 4 * 1024 * 64 * 2 524,288。MQAnum_kv_heads 1。缓存参数量 1 * 1 * 1024 * 64 * 2 131,072。可以看到GQA (G4) 将KV缓存大小降到了MHA的1/4同时保留了4组独立的Key/Value表达在性能和效率之间取得了很好的平衡。这也是为什么像Llama 2 70B这样的大模型会采用GQA8个KV头的原因。那么num_groups该怎么选呢这没有绝对答案但有一些经验法则num_groups越小越接近MQA内存收益越大但性能风险越高。num_groups越大越接近MHA性能保留越好但内存节省越少。常见的实践是取num_heads的1/2、1/4 或 1/8。例如32个Query头配8个或4个KV头。你可以把它看作一个可以调节的“性能-效率”旋钮。在资源紧张时调小追求极致性能时调大。4. 实战优化将KV Cache与GQA投入生产4.1 与现有模型架构集成理解了原理和基础实现后我们如何将它们用到真实的模型里呢通常我们会在Transformer的解码器块中集成这些技术。下面是一个简化的解码器层示例展示了如何将GQA模块嵌入其中并管理KV Cache的生命周期。class TransformerDecoderLayerWithGQA(nn.Module): def __init__(self, d_model, num_heads, num_groups, ffn_dim, dropout0.1): super().__init__() # 使用我们实现的GQA作为自注意力层 self.self_attn GQA(embed_dimd_model, num_headsnum_heads, num_groupsnum_groups) self.cross_attn None # 编码器-解码器注意力此处省略 self.feed_forward nn.Sequential( nn.Linear(d_model, ffn_dim), nn.ReLU(), nn.Linear(ffn_dim, d_model), ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, encoder_outputNone, self_attn_maskNone, past_key_valueNone, use_cacheFalse): # past_key_value: 来自上一时间步的 (self_attn_k, self_attn_v) 缓存 self_attn_past past_key_value if past_key_value is not None else (None, None) # 1. 带KV Cache的自注意力 attn_output, present_key_value self.self_attn( qx, kx, vx, maskself_attn_mask, past_keyself_attn_past[0], past_valueself_attn_past[1] ) x x self.dropout(attn_output) x self.norm1(x) # 2. 前馈网络 ff_output self.feed_forward(x) x x self.dropout(ff_output) x self.norm2(x) # 根据是否使用缓存决定返回什么 if use_cache: return x, present_key_value else: return x在实际的生成循环中你需要维护一个past_key_values列表其中每个元素对应一个解码器层返回的present_key_value元组。每一步生成时将这个列表传入下一时间步。4.2 性能实测与 profiling 技巧理论归理论优化效果到底如何还得跑起来看。我们可以写一个简单的测试脚本来对比不同配置下的内存和速度。import time import torch from memory_profiler import memory_usage def benchmark_generation(model, input_ids, max_length100, use_cacheTrue): 基准测试生成过程 model.eval() past_key_values None generated input_ids start_time time.time() with torch.no_grad(): for _ in range(max_length - len(input_ids)): outputs model(input_idsgenerated[:, -1:], # 只输入最后一个token past_key_valuespast_key_values, use_cacheuse_cache) next_token_logits outputs.logits[:, -1, :] next_token torch.argmax(next_token_logits, dim-1, keepdimTrue) generated torch.cat([generated, next_token], dim-1) if use_cache: past_key_values outputs.past_key_values end_time time.time() avg_time_per_token (end_time - start_time) / (max_length - len(input_ids)) return generated, avg_time_per_token # 假设我们有一个实现了GQA的模型类 MyGPTWithGQA model_mha MyGPTWithGQA(num_heads16, num_groups16) # 等价于MHA model_gqa MyGPTWithGQA(num_heads16, num_groups4) # GQA model_mqa MyGPTWithGQA(num_heads16, num_groups1) # 等价于MQA # 测试 prompt torch.tensor([[1, 2, 3, 4]]) # 模拟输入 print(Benchmarking MHA (num_groups16)...) _, t_mha benchmark_generation(model_mha, prompt, max_length50) print(f Avg time per token: {t_mha:.4f}s) print(Benchmarking GQA (num_groups4)...) _, t_gqa benchmark_generation(model_gqa, prompt, max_length50) print(f Avg time per token: {t_gqa:.4f}s) print(f Speedup: {t_mha/t_gqa:.2f}x) print(Benchmarking MQA (num_groups1)...) _, t_mqa benchmark_generation(model_mqa, prompt, max_length50) print(f Avg time per token: {t_mqa:.4f}s)除了时间更要关注内存。你可以使用torch.cuda.memory_allocated()来监控GPU内存变化特别是在长序列生成时观察KV Cache的内存增长曲线。你会发现使用GQA后内存增长的速度会显著慢于MHA这对于在资源有限的设备上部署大模型至关重要。4.3 高级优化技巧与避坑指南在实际项目中应用这些技术时还有一些细节需要注意缓存格式选择为了极致优化工业级框架如vLLM、TGI不会直接存储巨大的[batch, heads, seq_len, dim]张量。它们会使用更紧凑的格式比如把K和V的head_dim维度放到最后或者使用分页缓存来管理碎片这能进一步提升内存利用率和推理速度。连续生成与中断续接KV Cache使得中断后继续生成成为可能。你需要将之前的完整缓存序列化保存续接时反序列化加载然后从断点开始继续生成。这要求你的缓存管理代码足够健壮。批处理中的变长序列在实际服务中不同用户的输入长度可能不同。这要求KV Cache能处理“锯齿状”的批次。常用的技巧是使用注意力掩码attention mask来屏蔽填充部分并确保缓存拼接时维度的正确对齐。GQA的微调如果你从一个标准的MHA模型比如原始的LLaMA检查点切换到GQA架构直接加载权重可能会出问题因为K/V投影层的形状变了。通常的作法是对K/V投影层的权重进行“复制”或“平均”来初始化新的、更小的投影层。例如将原来16个头的K投影权重根据分组规则平均成4组用这4组的权重来初始化GQA的K投影层。量化与KV Cache为了进一步降低部署成本KV Cache本身也可以进行量化如INT8甚至INT4。这能大幅减少内存占用但需要仔细处理量化带来的精度损失确保不影响生成质量。从我自己的经验来看KV Cache和GQA是推理优化中“性价比”最高的两项技术。它们不需要改变模型架构的核心思想只是对计算和存储模式进行了巧妙的重新组织就能带来数倍的性能提升。尤其是在结合了量化、算子融合等底层优化后效果更加显著。刚开始实现时多在维度变换和缓存管理上花点时间调试理解每一步张量的形状变化后面就会一马平川了。