你有没有遇到过这种情况同一个大模型第一次提问时响应还挺快但只要把对话变长或者一次性喂进一大段上下文生成速度就像被什么东西拖住了显存占用也肉眼可见地往上爬。如果你部署过推理服务、或者本地跑过开源模型大概率已经听过两个名词多头注意力机制Multi-Head Attention和 KV cache。这篇文章就把这两个东西一次性讲清楚从它们到底解决什么问题到每一层计算里到底发生了什么再到显存占用怎么算、工程上常见的坑有哪些全按我实际调试模型时的理解来写。这个话题适合谁如果你是刚开始接触大模型原理、正在做推理性能优化、或者单纯想搞清楚为什么长对话推理那么慢这篇都值得读完。我不会堆公式重要的数学部分会配合直观解释和实际数字一起讲尽量做到看完之后你能自己推导出一个模型的 KV cache 大概吃多少显存。1. 先从一次模型推理说起慢在哪里、费在何处我在本地跑一个 7B 量级模型的时候做过一个简单测试让模型写一段 200 字的短文输入只有一句话模型生成得很快显卡功耗也平稳。但同样是这个模型当我给它的对话历史增加到 4000 token 之后每个新词的生成时间明显变长显存占用也涨了近 2GB。当时第一反应是模型难道越长越笨其实不是模型变笨了是注意力机制的计算特性和缓存策略导致的。要理解这个问题得先知道自回归模型是怎么生成内容的。以 ChatGPT 这类模型为例它不是一个字一个字想好再输出而是每预测一个词都要把当前已有的整个序列重新过一遍模型。也就是说生成第 100 个词时模型要看前 99 个词生成第 101 个词时又要重新看前 100 个词。如果每次都把之前算过的内容全部重算一遍计算量就会随着序列长度平方级增长。这显然不划算。于是 KV cache 这个工程优化就出现了既然前面那些词的 Key 和 Value 已经算过了为什么不把它们缓存下来下次直接复用这个想法朴素但极其有效。后面我会详细展开它到底缓存了什么、省了多少计算、又带来了什么新的问题。先把注意力机制本身讲清楚因为 KV cache 里的 K 和 V正是从注意力机制里来的。2. 注意力机制的基础Q、K、V 从哪里来要到哪里去2.1 从查找资料理解 Query、Key、Value说起注意力机制很多人的第一反应是三个字母Q、K、V。但 QKV 到底代表什么很多人是背下来的不是理解下来的。我常用一个类比想象你在图书馆查资料。假设你想了解多头注意力的数学推导你心里会形成一个查询这就是 Query。你会去目录里找关键词目录里每个条目对应一个键这就是 Key。如果你找到了匹配的条目你会把对应的那本书抽出来读书里的内容就是 Value。注意力机制做的事就是根据 Query 和 Key 的匹配程度决定从 Value 里取多少信息。在 Transformer 里每个 token 会生成三个向量一个 Query一个 Key一个 Value。生成方式是对输入向量分别做三次线性变换。所谓线性变换本质上就是拿输入向量乘三个不同的权重矩阵。这三个权重矩阵是模型在训练过程中学出来的参数。如果一句话让你记住Q 代表我要找什么K 代表我有什么可被找到的标签V 代表我真正提供的信息内容。打分决定了我该关注谁加权求和决定了我最终从每个人那里拿到多少信息。2.2 缩放点积注意力打分与加权求和注意力机制的核心计算可以拆成四步第一步对序列中每两个 token 计算 Q 和 K 的点积。点积的结果就是注意力分数它表示当前这个 token 应该对另一个 token 投入多少关注。点积越大说明两个向量在方向上的相关性越强。第二步对这些分数进行缩放也就是除以根号下头的维度大小。这一步不是玄学后面单独解释。第三步沿着序列维度做 Softmax把分数变成一组和为 1 的概率分布。这样可以保证所有历史 token 的注意力权重非负、且总量为 1。第四步用这些权重对不同位置的 Value 做加权求和得到最终的注意力输出。用一个具体的数字感受一下假设输入序列有 5 个 token每个 token 的维度是 512。那么注意力分数矩阵就是 5×5 的矩阵Softmax 之后依然 5×5最终用它对 5×512 的 Value 矩阵做加权得到 5×512 的输出。这个过程中的计算量是序列长度的平方关系。序列越长注意力矩阵越大计算和显存开销都会明显上涨。这也是长上下文推理迟迟难以做快的重要原因。2.3 为什么点积要缩放一个看似微小但影响巨大的细节很多人初看注意力公式时会对那个除以根号 d_k感到迷惑。为什么不直接除以 d_k为什么偏偏是根号原因在于方差控制。Q 和 K 都是多维向量如果每个维度是均值为 0、方差为 1 的随机变量那么两个向量的点积其方差会随着维度数 d_k 线性增长。也就是说维度越大点积结果的绝对值越大。等你把它送进 Softmax 时容易进入梯度极小的饱和区导致训练时候梯度消失模型学不动。除以根号 d_k 之后点积的方差被拉回 1 左右Softmax 的输入分布就不会因为维度过大而过于极端。我曾经自己跑过一个对比实验同样的模型结构去掉缩放项训练时损失收敛明显变慢而且容易不稳定。加上缩放之后训练曲线立刻正常。这个细节在调试自定义 Transformer 时非常值得注意。3. 多头注意力一个头不够那就多来几个3.1 多头到底在做一件什么事单头注意力的能力是有限的。它就像一个学生读文章时只能从一种角度理解内容。但一篇文章里的关系是多样的有的词之间是语法关系有的词之间是语义关联有的词之间有指代关系。单头注意力很难同时捕捉这么多层次的信息。多头注意力的思路很简单不搞一个头而是把模型维度切分成多份让每个头在低维子空间里独立做注意力。每个头可以学习到不同的关系模式最后把所有头的输出拼起来再做一次线性投影。一个好的说法是多头注意力不是让一个模型学得更用力而是让多个小专家从不同角度观察同一个序列再把意见汇总。3.2 完整的计算流程与矩阵维度假设模型的隐藏维度 d_model 是 512头数 h 是 8那么每个头的维度 d_k 就是 512 / 8 64。具体流程如下第一步输入向量分别经过三组线性变换生成 Q、K、V各自的维度都是 512。第二步把 Q、K、V 按头数切分。512 维向量切成 8 份每份 64 维。于是第 i 个头拿到的 Q_i、K_i、V_i 维度都是 64。第三步每个头独立计算注意力输出得到维度为 64 的输出向量。第四步把 8 个头的输出拼接起来恢复成 512 维。第五步经过输出投影矩阵 Wo形状是 512×512得到最终的多头注意力输出。如果写成伪代码可以这样看def multi_head_attention(x, W_q, W_k, W_v, W_o, num_heads): # x: [seq_len, d_model] Q x W_q # [seq_len, d_model] K x W_k V x W_v d_model x.shape[-1] d_head d_model // num_heads # 拆头: [seq_len, num_heads, d_head] Q Q.reshape(seq_len, num_heads, d_head).transpose(0, 1) K K.reshape(seq_len, num_heads, d_head).transpose(0, 1) V V.reshape(seq_len, num_heads, d_head).transpose(0, 1) # 每头独立做注意力 scores Q K.transpose(-2, -1) / sqrt(d_head) weights softmax(scores, dim-1) out weights V # [num_heads, seq_len, d_head] # 合并头再做输出投影 out out.transpose(0, 1).reshape(seq_len, d_model) out out W_o return out这里有一个容易搞混的点在推理时我们经常说一个 32 层模型有 32 个头其实指的是每层有 32 个注意力头而不是整个模型只有 32 个头。层数、头数、头的维度是三个独立的超参数。3.3 多头带来的实际价值不同头关注不同的内容我在训练一个自研小模型时曾经把注意力头的权重可视化出来。比较明显的是一个规律有的头高度关注相邻词之间的关系表现出很强的局部性有的头会跨越很长距离去关注句子开头的主题词表现出全局依赖还有的头会关注句法结构里的主语和谓语。这些模式不是人为设计的是模型在训练中自然形成的。这也是多头相对单头的核心优势并行的多个子空间让模型有容量去同时建模不同类型的依赖关系。在机器翻译、文本分类、代码生成等任务上多头基本上都比单头有效。不过头数也不是越多越好。头数太多每个头的维度被压缩得很小单个头的表达能力反而下降。而且头与头之间可能出现冗余。在工程实践中8 头到 32 头是比较常见的范围具体头数需要结合模型规模和任务类型来调。4. KV Cache推理加速背后的缓存魔法4.1 生成式模型的两场考试预填充与逐字生成理解了注意力机制就可以解释 KV cache 了。在自回归生成中模型把一个输入序列从头到尾计算一遍称为 Pre-fill 阶段也就是预填充阶段。这个阶段把用户输入一次性处理完生成第一个输出 token。之后进入 Decode 阶段也就是逐字生成阶段每生成一个 token都要把新 token 拼到已有序列后面再做一次完整的前向计算。如果没有 KV cache每次生成新 token 时模型都要对历史所有 token 重新计算一遍 Key 和 Value。这就像考试时每做一道新题都要把前面所有题目重新读一遍而不是只看一眼自己的草稿纸。这种做法既浪费计算又消耗时间。KV cache 做的事就是在预填充阶段算完所有历史 token 的 K 和 V 后把它们缓存住。后续每生成一个新 token只需要计算当前 token 的 Q、K、V用它的 Q 去和缓存的所有 K 做注意力再用注意力权重对缓存的所有 V 加权求和。历史 token 的 K、V 完全不需要重算。4.2 KV Cache 到底缓存了什么、放在哪里KV cache 缓存的具体内容是每一层 Transformer 中注意力模块的 Key 和 Value 矩阵。模型每一层都有自己的 K 和 V所以缓存也是分层存的。形状方面假设 batch size 为 1当前序列长度为 S模型有 L 层每层有 H 个头每个头的维度为 D那么每层的 K 缓存形状是 [S, H, D]V 缓存同样为 [S, H, D]。全部层加起来就是 2 × L × S × H × D。很多人会忽略的一个点是 KV cache 的存储位置。它一般不在显存中单独划一块高速缓存区而是分配在推理框架管理的内存池里。比如当前主流的推理引擎在启动时会预先分配一块显存专门用于存放不断增长的 KV cache而不是每次动态申请。这样做的原因是动态申请显存开销大而且容易产生碎片。我在本地跑推理时观测到当对话长度逐渐变长时显存占用会稳步上升而上升的这部分几乎全部来自 KV cache 的增长。模型参数占用的显存是固定的差异完全在 KV 缓存上。4.3 显存占用怎么算带数字的实战计算这个计算是理解 KV cache 开销的必修课。公式如下KV cache 显存 2 × 层数 × 头数 × 头维度 × 序列长度 × batch size × 字节数公式里的 2 是因为 K 和 V 各占一份。字节数取决于精度FP16 是 2 字节FP32 是 4 字节INT8 是 1 字节。用实际例子来感受一下。假设模型是 7B 参数32 层每层 32 个头每个头维度 128。那么隐藏维度就是 32 × 128 4096。如果序列长度是 2048batch size 是 1使用 FP16每个 token 在每层需要的 K 缓存大小是 32 × 128 × 2 字节 8192 字节V 缓存同样 8192 字节。一层加起来是 16384 字节。32 层就是 524288 字节约 0.5 MB。再乘以序列长度 2048得到约 1 GB。所以一个 7B 模型在序列长度 2048 时单条请求的 KV cache 大约占 1 GB 显存。如果你用 batch size 8 同时处理 8 个请求那就是 8 GB。这时候你会意识到模型权重可能只占 14 GB 左右但一批并发请求的 KV cache 就能轻松超过权重大小。这个计算还可以反推如果你想在 24 GB 显存的卡上跑 7B 模型模型权重 FP16 约 14 GB剩下约 10 GB 给 KV cache大概能支撑 batch size 4、序列长度 2048 的并发。超过这个规模就得用量化或者其他优化手段。5. KV Cache 在实际工程中的常见问题与优化思路5.1 显存不够怎么办浅谈 GQA、PagedAttention 与量化KV cache 最大的工程问题就是显存。为了解决它业界出现了几个有代表性的优化方向。第一个方向是减少 KV 缓存本身的数量。典型做法是分组查询注意力GQA和多查询注意力MQA。在标准多头注意力中每个头都有自己的 K 和 V。GQA 的做法是让多个 Query 头共享一组 K、V 头。比如 32 个 Query 头共享 8 组 K、V 头。这样 KV cache 直接缩小到原来的四分之一。很多开放权重模型的 32B、70B 版本都在用这种结构。第二个方向是解决显存碎片和动态分配问题。PagedAttention 的思路借鉴了操作系统的虚拟内存分页机制把 KV cache 分成固定大小的块按需分配避免碎片化。这在长上下文场景里提升明显。第三个方向是 KV cache 量化。把 FP16 的 K、V 值压缩成 INT8 甚至 INT4能大幅减少显存占用。代价是有一定精度损失需要根据任务实测效果决定是否可用。我在实际部署中用过 INT8 的 KV cache 量化对于通用对话任务效果差别很小生成速度没有明显变化但显存占用降低了不少。对于需要精确推理的任务比如数学推理或代码生成建议先做验证再启用。5.2 长上下文场景下KV cache 的另一个隐患长上下文场景里KV cache 不只是显存变大这么简单。还有一个容易被忽视的问题随着序列变长模型需要从更多的历史信息里检索相关内容注意力分布会被稀释。而 KV cache 存储的 Key 与 Value 是原始精度的一旦量化过度长距离依赖信息会丢失模型的表现会有明显下降。另外某些推理框架为了快速访问 KV cache会把 K、V 矩阵搬运到 GPU 上做计算而 KV cache 本身也在显存里。即使显存够用带宽也可能成为瓶颈。序列越长每次生成新 token 需要读取的 KV cache 数据量越大这块读取时间往往超过了计算本身。我调过一个长文档问答项目输入文档 8000 token生成第一个 token 之前预填充阶段耗时约 2 秒但之后每个 token 的生成延迟也明显比短文本高。后来用 profiling 工具一看时间大部分花在了从 KV cache 读取数据上。这种场景下单纯增加计算卡解决不了问题瓶颈在内存带宽。5.3 给实际部署和实验的一些建议根据我个人踩坑的经验如果你要在自己的项目里用到 KV cache下面这几条建议很有参考价值。第一先估算后优化。在你选推理框架之前用前面那个公式算一下自己业务的 KV cache 峰值。别等跑起来显存爆了再查那时定位问题会非常痛苦。第二别盲目上量化。KV cache 量化确实省显存但它和模型权重量化叠加时效果需要实测。我建议先用原始精度跑通业务再做量化对比用业务指标说话。第三留意框架的预分配策略。有些推理框架默认预留大量显存给 KV cache即使当前请求的序列很短。如果你同时要做模型训练或多路推理服务记得调整内存池的上限否则会出现明明显存还很多但新请求被拒绝的现象。第四如果序列长度差异很大考虑动态批处理。把不同长度的请求打包在一起跑能显著提升 GPU 利用率但会增加调度复杂度。6. 最后说点实践里的真实感受多头注意力和 KV cache 这两个机制一个是模型能力的地基一个是推理效率的骨架二者原理上各自独立工程上却紧密耦合。模型的层数、头数、头维度决定了 KV cache 的体量而 KV cache 又反过来限制了你能处理多长的序列、能支撑多大的并发。我在不少项目里看到过类似的误区有人一味增大模型规模却不关注注意力头的配置是否合理有人把 KV cache 当作黑盒只管开大缓存池结果显存余量被吃得干干净净。其实只要把这两个机制的核心逻辑想明白大半性能问题都能提前规避。最后送一个实操小技巧调试模型推理性能时先把 KV cache 的显存占用单独打点打印出来同时记录每轮生成的 token 延迟。把这两条曲线叠在一起看你往往能找到性能拐点。之前我就是在一条增长曲线里发现某个推理框架在序列超过 4096 后会有异常的带宽抖动换掉之后速度立刻恢复正常。这些散落在实现细节里的坑不亲手踩一遍光看文档是体会不到的。