注意力机制核心拆解：QKV投影与KV Cache推理优化-北京尧图网络

很多人看注意力机制的示意图看了无数遍Q、K、V三条线画个点积再来个softmax最后加权求和好像就这么回事。但真到了自己写代码、排查推理性能瓶颈、或者调KV Cache显存的时候才发现脑子里那幅图是模糊的——Q和K到底怎么从输入向量里变出来的那个除以根号d_k到底在图里哪一步每个头的输出又怎么拼回d_model维更绕的是训练时明明一张图搞定为什么推理时要搞KV Cache它到底缓存了什么这篇文章我把“一次注意力计算”从头到尾拆开配合可以直接跑的小例子和PyTorch代码把每一步的数据形状变化、矩阵运算含义、数值示例都摆出来再把KV Cache的来龙去脉讲清楚。适合刚接触Transformer但想真正读懂源码的人也适合做推理优化、被显存和延迟问题困扰的朋友。1. 从词向量到QKV一次投影到底做了什么1.1 为什么非要拆成三个空间注意力机制里最反直觉的设计就是同一个输入向量为什么要复制三份然后用三个不同的矩阵去变换直接理解是一个token在序列里承担多个角色。当它作为“当前我要处理的对象”时它是query负责发出需求我要从上下文里找什么信息。当它作为“被其他token查看的对象”时它是key相当于一个索引标签告诉别人我这里有什么内容。而当信息真正被提取时它是value提供实际的语义内容。这三个角色如果共用一份向量就会出现一个问题某个token既想被“高亮”检索又想让别人只取走它的一部分信息这俩诉求在同一个向量空间里无法独立表达。拆成三个空间之后模型就有了自由度——同样一个词它的query可以偏向“我要找动词搭配”它的key可以偏向“我是一个名词短语的起始”它的value可以偏向“我携带了具体实体信息”。这套拆分的自由度正是attention能建模复杂依赖关系的基础。实现上Q、K、V分别通过三个独立的线性映射得到q x W_q b_q k x W_k b_k v x W_v b_v不共享权重是因为三个角色需要学完全不同的特征。共享了模型就退化成只能在同一个空间里折中表达力会大打折扣。1.2 QKV投影的数学定义与参数量估算假设输入hidden state的维度是d_model多头注意力的head数量是h每个head的维度是head_dim那么通常有d_model num_heads * head_dim例如常见的配置d_model768num_heads12head_dim64。在这个配置下单头注意力的Q、K、V投影矩阵形状分别是(d_model, head_dim)也就是768×64。把所有头合起来看一个多头注意力模块的QKV总投影参数量大约是参数量 ≈ 4 * d_model * d_model这里的4来自W_q、W_k、W_v、W_o四份投影。以d_model768为例就是4×768×768≈236万参数。对于一个40层、d_model4096的模型单层attention的参数量就是4×4096×4096≈6710万40层加起来约27亿参数——足以说明attention里的线性投影不是“小配件”本身就是大模型参数的重要构成部分。这里有个很多初学者容易忽略的点虽然每个head的维度只有head_dim但在PyTorch的nn.Linear实现里通常直接用(d_model, d_model)的矩阵做整体投影然后按head_dim切块。比如q nn.Linear(d_model, d_model)出来之后view(batch, seq_len, num_heads, head_dim)再转置。这样做的好处是一次矩阵乘完成所有head的投影计算效率高坏处是如果你去读shape会看到整个head的Q是同一个大矩阵容易误以为没有“分头”。实际上分头只是把结果沿head维度切成h段。1.3 数据形状变化一条信息流的完整走位为了后面讨论KV Cache我们必须对形状变化非常敏感。一次标准的多头注意力计算输入输出形状变化是这样的假设输入x的形状是(batch, seq_len, d_model)。经过Q、K、V三个投影后得到三个同样形状的张量。然后做一个reshape和permute把单头的顺序摆正q的shape: (batch, seq_len, d_model) → view成 (batch, seq_len, num_heads, head_dim) → transpose成 (batch, num_heads, seq_len, head_dim)为什么要调成(batch, num_heads, seq_len, head_dim)因为注意力分数的计算是每个head独立做的形状里把head维放在第二维后面做batch矩阵乘时可以一次把全部头的计算并起来q k.transpose(-2, -1)得到(batch, num_heads, seq_len, seq_len)的分数矩阵。要注意的是这个分数矩阵的维度含义是(batch, num_heads, query_len, key_len)。也就是每一行代表一个目标位置query每一列代表一个源位置key。这决定了后面mask、KV Cache拼接时的方向。2. 注意力分数点积、缩放与归一化2.1 点积相似度怎么衡量“相关”拿到Q和K之后最核心的一步是计算每个query与所有key的相似度。采用的方法是内积点积。内积的直觉很朴素两个向量如果在方向上越一致点积越大。如果方向正交点积接近0。如果方向相反点积为负。把query想象成“我要找的内容概要”把key想象成“我这里的标签”概要跟标签越匹配内积就越大这个位置也就越值得被关注。但这里有个细节Q和K的维度在注意力内部是head_dim不是d_model。为什么用低维空间算因为点积本身的作用是打分不需要把全部语义维度都拉进来head_dim一般64或128已经足够表达相似度了。高维反而容易让分数分布极端。使用低维投影相当于先做了一次降维聚焦这也让模型能在多个head中各自关注不同属性。2.2 缩放因子根号d_k的数学直觉点积之后为什么要除以根号d_k这个操作叫缩放点积注意力。如果q和k的每个分量都是均值为0、方差为1的随机变量那么它们点积的期望方差是d_kVar(q·k) d_k也就是说head_dim越大点积分数的绝对值越大、分布越散。当这些分数进入softmax时会发生什么指数函数对大的输入非常敏感如果有个分数是10另一个是20e^20与e^10的比值大约是e^10≈22026倍softmax输出几乎变成one-hot——模型直接“锁死”在某一个位置梯度传不过去其他位置的信息全部丢失。除以根号d_k之后分数方差回到1左右softmax的输入分布相对平缓训练稳定得多。这就是为什么公式里写的是attention_weight softmax(Q K^T / sqrt(d_k))而不是直接softmax(Q K^T)。实话说这个缩放因子的具体数值不一定是最优的但它是目前最稳妥、可训练性最好的设定PyTorch、各种框架默认都是这么实现的。2.3 手算一个极小例子这部分我们实际算一遍不玩虚的。下面构造一个只有3个token、head_dim2的小场景Q、K、V直接给出来模拟投影以后的输出。token1: q[1.0, 0.0], k[0.8, 0.2], v[1.0, 0.0] token2: q[0.0, 1.0], k[0.1, 0.9], v[0.5, 0.5] token3: q[0.5, 0.5], k[0.5, 0.5], v[0.0, 1.0]第一步算所有query和所有key的点积矩阵S Q K^TS [[0.8, 0.1, 0.5], [0.2, 0.9, 0.5], [0.5, 0.5, 0.5]]第二步除以根号2得到缩放后的分数S / 1.414 ≈ [[0.566, 0.071, 0.354], [0.141, 0.636, 0.354], [0.354, 0.354, 0.354]]第三步按行做softmax。以第一行为例e^0.566 ≈ 1.761 e^0.071 ≈ 1.073 e^0.354 ≈ 1.425 sum 4.259 权重 [0.413, 0.252, 0.335]可以看到token1的query对token1自己的key最关注权重0.413对token2的key关注度最低权重0.252。第二行同理token2最关注自己0.423。第三行三个分数相等softmax之后恰好是均等的1/3。第四步用注意力权重对V加权求和output1 0.413*[1,0] 0.252*[0.5,0.5] 0.335*[0,1] [0.539, 0.461] output2 0.258*[1,0] 0.423*[0.5,0.5] 0.319*[0,1] [0.470, 0.530] output3 (1/3)*[1,0] (1/3)*[0.5,0.5] (1/3)*[0,1] [0.500, 0.500]这个例子最直观的价值是output1明显偏向token1的Voutput2明显偏向token2的Voutput3因为对所有位置一视同仁输出正好就是V的均值。注意力机制本质上干的活就是根据相关性做“加权平均”把V里有用的信息抽出来。3. 从分数到上下文加权求和、输出投影与多头整合3.1 加权求和上下文向量是怎么来的softmax权重确定后最后一步就是attention_weight V。这一步从形状上看是把(seq_len, seq_len)的权重矩阵乘以(seq_len, head_dim)的V矩阵结果是(seq_len, head_dim)。这里的含义要反过来理解对于第i个位置它的输出是V矩阵所有行的加权和权重就是第i行softmax之后的分布。所以每个位置的输出都不是原token本身的向量而是“从整个序列中按相关性提取出来的混合信息”。这也是为什么叫上下文向量——它包含了其他token的信息。我见过有人把这一步误解成“当前位置的V保留其他位置的V清零”不对。注意力的输出不是选择而是软性混合。所有历史位置的V都会对输出有贡献只是贡献比例不同。这个“软”特性让梯度可以顺畅地流向所有相关位置这也是attention能捕捉长距离依赖的核心原因。3.2 输出投影为什么还多一步多头注意力的每个head都会产生一个head_dim维的上下文向量把所有head拼起来得到的是num_heads * head_dim d_model维的向量。然后还需要经过一个输出投影W_o把多头信息融合回统一空间。这一步常常被初学者忽略但它很重要。如果没有W_o模型各个head的输出只是机械地拼接彼此之间没有交叉融合的机会。W_o相当于一个“混合层”它决定最终输出里每个head的贡献如何组合。在代码实现里W_o就是那个(d_model, d_model)的线性层输入是拼接后的向量输出是融合后的向量。注意一个细节QKV投影里W_q、W_k、W_v的bias在某些模型中是不用的比如很多开源decoder模型为了避免位置编码和bias冲突把bias置为False但W_o的bias一般保留。这个小事影响不大只是如果你在复现某个模型时发现参数量对不上可以留意一下是不是bias的设置差异。3.3 多头注意力到底在并行算什么多头注意力并不是“把序列切成多段分别处理”而是同一份输入被投影到h个不同的低维子空间在每个子空间独立地计算注意力。每个head相当于一个不同的“视角”。举个类比一群人同时看同一张照片有人关注颜色有人关注边缘有人关注物体之间的相对位置。每个head就是一个拿着不同“滤镜”的观察者它的W_q、W_k、W_v就是滤镜本身。最终W_o把这些观察结果综合起来。从可解释性工作的经验来看不同head确实会学到不同的模式有的head偏向关注句法依赖有的head偏向关注指代关系有的head则主要关注相邻位置的局部信息。这些行为不是人为设计的而是训练中自组织出来的。所以多头的意义不是“让模型更大更强”而是“让模型有多个独立的信息通道”减少单一空间下的相互干扰。3.4 PyTorch手写一个完整的attention理论说完了上代码。比起直接看公式手写一遍能暴露很多理解上的漏洞。下面是一个多头注意力的紧凑实现没有做mask假设是普通的全局注意力import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.head_dim d_model // num_heads self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.w_o nn.Linear(d_model, d_model) def forward(self, x, maskNone): batch, seq_len, _ x.shape q self.w_q(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k self.w_k(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v self.w_v(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) scores q k.transpose(-2, -1) / math.sqrt(self.head_dim) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights torch.softmax(scores, dim-1) out attn_weights v out out.transpose(1, 2).contiguous().view(batch, seq_len, -1) out self.w_o(out) return out, attn_weights代码里容易写错的是那个reshape链先.view(batch, seq_len, num_heads, head_dim)再.transpose(1, 2)把(batch, seq_len, num_heads, head_dim)变成(batch, num_heads, seq_len, head_dim)。这里如果忘了contiguous()后面的view会报错因为transpose之后张量在内存里不连续。我建议你实际跑一遍这个代码用x torch.randn(2, 5, 768)单层前向一下把每步的shape都打印出来。这个东西比看十遍图都管用。4. 推理时的KV Cache一次注意力的性能关键4.1 解码阶段为什么有大量重复计算训练的时候一个句子的所有token是一次性全部并行送入模型的所以attention计算的矩阵是L×L一次做完。但在推理生成阶段模型是一个token一个token往外蹦的。每生成一个新token都要重新计算一次attention新token需要attend到所有历史token。如果不做任何缓存每生成一个新token时都要把之前所有token的Q、K、V从头到尾重新过一遍投影层和注意力层。但问题来了之前那些token的K和V是由它们的hidden state经过固定的W_k、W_v算出来的而它们的hidden state在之前生成时就已经算过了K、V并不会因为新token的到来而改变。重复计算它们纯粹就是浪费算力。打个比方你要在图书馆里面对一本书做摘要每次读到新一页都把前面所有页的目录重新编一遍这不是傻是什么KV Cache做的就是把已经编好的“目录”K和“卡片内容”V存下来新来一页只编新页的然后拼到目录后面。4.2 KV Cache缓存的是什么、如何写入KV Cache本质上就是两个张量一个缓存K一个缓存V形状是(batch, num_heads, seq_len, head_dim)。在生成开始时它们为空。每生成一个新token执行以下流程# 假设当前token的hidden_state形状: (batch, 1, d_model) # 已经完成QKV投影得到q/k/v形状: (batch, num_heads, 1, head_dim) # 把当前token的K、V写入缓存末尾即seq_len维度上拼接 k_cache torch.cat([k_cache, k], dim2) # (batch, num_heads, old_len 1, head_dim) v_cache torch.cat([v_cache, v], dim2) # 当前token的query只保留一个位置 # scores形状: (batch, num_heads, 1, old_len 1) scores q k_cache.transpose(-2, -1) / math.sqrt(head_dim) # 如果需要causal mask这里不需要因为缓存里本来就是历史token没有未来信息 attn_weights torch.softmax(scores, dim-1) out attn_weights v_cache # (batch, num_heads, 1, head_dim)整个过程中每个step只需要计算“当前token的Q”和“当前token的K、V”历史token的K、V直接从缓存里取。复杂度从重新计算的O(L²)降到O(L)。这里有个很关键的工程细节上面的代码用torch.cat实现图省事可以但生产环境千万别这么干。每生成一个token就c反一次意味着每次都要重新申请内存、拷贝数据生成长序列时总时间翻好几倍。正确做法是预分配一个(batch, num_heads, max_len, head_dim)的连续buffer然后每步原地写入k_cache[:, :, step:step1] k v_cache[:, :, step:step1] v这样每步只拷贝一个token的数据速度快很多显存也稳定。这是我调试推理服务时踩过的坑用torch.cat版本跑512个token的生成延迟比预分配版本慢了接近一倍。4.3 为什么不缓存QKV Cache只存K和V不存Q这不是为了省事而是因为Q本身没有复用的价值。每个新token到来时新token需要计算自己的Q去跟所有历史K做点积这一步无法避免。但历史token的Q呢它们在之前的step里已经算完并且已经用于当时的attention计算、已经产生了当时的输出。等新token到来时历史token的Q不会参与任何计算——attention分数永远是“当前query”对“所有历史key”的打分不存在历史query与未来key的打分。所以历史Q是真正的一次性用品缓存它纯粹浪费显存。等生成结束、要计算下一个step时当前token的K、V会进入缓存而它的Q在下一个step就不再需要了。整个生命周期里每个token的Q只被计算一次、使用一次、然后丢弃。这跟K、V“计算一次、反复使用”的性质完全不同。4.4 KV Cache的显存开销与计算量分析显存是KV Cache最直观的代价。假设一个40层decoder-only模型d_model4096num_heads32head_dim128以fp16精度存储每层每token的KV大小 2K和V× num_heads × head_dim × 2 bytes 2 × 32 × 128 × 2 16384 bytes 40层每token 16384 × 40 655360 bytes ≈ 640 KB也就是说这个配置下每生成一个tokenKV Cache要增加约640KB显存。生成2048个token时KV Cache大约会占用1.25GB4096个token就是2.5GB。如果有更大的batch还要再乘以batch size。这也解释了为什么业界会出现各种“省显存”的注意力变体。把每层的KV都缓存起来对长上下文推理来说是一笔沉重负担。常见的做法Multi-Query Attention所有head共享同一个K和V只保留Q的多头KV占用直接除以head数。Grouped-Query Attention把head分成若干组每组共享K、V是MQA和标准MHA之间的折中。Sliding Window Attention / 稀疏注意力限制每个query只attend附近固定长度的历史token。代价是长距离依赖变弱但显存和计算量都变成线性。计算量方面更直观。如果没有KV Cache第t步生成时attention的矩阵乘规模是t×t总计算量大约是O(N^3)——生成N个token就要做1^22^2...N^2次级别的矩阵运算。有了KV Cache第t步规模是1×t总计算量变成O(N^2)。长序列下这个差距是数量级的。所以KV Cache不是“锦上添花”的优化而是“没有它推理就跑不动”的必需品。5. 实战中的常见问题与排查方法5.1 softmax数值溢出与fp16下的坑标准实现里softmax前通常先减去每行最大值再做指数运算。这是数值稳定性的基本操作PyTorch内置的torch.softmax已经处理了。但当你手写attention、或者做某些自定义mask操作时很容易踩两个坑。第一个坑是把mask位置填成0而不是float(-inf)。如果mask位置是0softmax会把它当成一个普通分数被掩码的token照样参与加权求和。正确的是掩码位置填负无穷这样exp之后变成0注意力权重为0。如果你用float(-inf)已经填了就不需要再担心softmax出NaN因为exp(-inf)0是安全值。第二个坑在fp16混合精度下容易出现。float(-inf)本身没问题但如果mask前把scores做了很激进的截断比如有些实现为了匹配某些算子会用-10000这种大负数代替-inf。在fp16下-10000虽然也能让exp之后趋近于0但在某些kernel实现里极大负值参与计算时精度表现不如-inf干净。我建议优先用-inf除非你的推理框架明确要求用大负数。5.2 位置编码与KV Cache的对齐问题KV Cache缓存的是K、V而位置编码对K的效果不同处理方式也不同。如果你用的是绝对位置编码直接把位置向量加到token embedding上那位置信息已经进入hidden stateK、V天然包含位置信息直接缓存就行。如果用的是相对位置编码或旋转位置编码情况就微妙了。以旋转位置编码为例K在进入缓存之前需要先根据它当前的绝对位置做旋转。历史token的K已经按它当时的位置旋转过缓存后不需要再动当前新token的K要按当前step对应的旋转角度旋转然后再写入缓存。也就是说K的旋转是在“入缓存前”完成的这与Q每step都要单独处理当前角度的逻辑是一样的。一个常见的错误是在推理时沿用训练时的写法把全序列的Q、K一起做RoPE但KV Cache里存的是“已经旋转过的旧K”再对整段做一次RoPE就会把位置信息叠加错。正确的做法是只对当前step的Q和K做旋转再写缓存。5.3 KV Cache的显存规划不当推理服务部署时最容易看到的问题KV Cache按max_len一次性预分配结果显存被占得死死的batch size上不去不预分配又频繁申请内存导致延迟抖动。一个比较实用的经验是先量算模型在最坏情况下的KV显存需求再结合可用的剩余显存决定max_len和max_batch_size。公式就是上面写的那样KV显存 2 × num_layers × batch × num_heads × seq_len × head_dim × 2字节比如一个40层、d_model4096的模型如果想要单batch跑4096上下文KV就要2.5GB。如果你只有16GB显存模型权重已经占了8GB那KV最多只能给2GB左右4096上下文就放不下要么减小max_len要么换MQA/GQA结构。这个账应该在部署前算清楚别等OOM了再调。5.4 常见问题速查表与避坑清单问题现象原因解决方法softmax输出异常注意力权重几乎one-hot训练不收敛忘记除以根号d_k按缩放公式计算mask失效被mask的位置仍然影响输出mask位置填0而不是负无穷填float(-inf)shape报错view或transpose后无法连续忘了contiguous()在view前调用contiguous()推理越跑越慢生成速度随tokens增加明显下降KV Cache用torch.cat反复拼接预分配buffer原地写入显存爆掉长上下文推理OOMKV Cache未做预算按公式估算并限制max_len长序列位置错乱生成内容重复或语义跳跃KV Cache与RoPE处理不对齐只对当前step的Q/K做旋转后再缓存最后再分享一个个人经验。我当时是从零写attention起步的最强烈的感受是这个东西的每一行代码都不难但放在一起很容易在形状和维度上绕晕。我的建议是不要嫌麻烦把q、k、v、scores、attn_weights、out每一步的shape都打印出来核对特别是从(batch, seq_len, d_model)到(batch, num_heads, seq_len, head_dim)那一步很多人就是在这里掉进坑里的。把这个流程走通了再看KV Cache、看GQA、看各种inference优化都会轻松很多。

注意力机制核心拆解：QKV投影与KV Cache推理优化

相关推荐

评价模型分析总和：从指标归一化到加权汇总的完整落地指南

T611镗床主轴箱传动与尾柱设计：转速图规划到出图规范全解析

基于STM32的智能婴儿床：传感器融合与哭声识别实战

更多资讯

从 20 个图标到 5 个：Ice 菜单栏管理新手完整指南

AI小说创作神器

一文读懂Llama-3.2-1B_rai_1.8.0_npu_4K文件结构：从model.onnx到40多个dd_metastate文件全解析

gvim快捷键声明wire和reg

如何免费听遍全网歌曲：LX Music 五音源聚合搜索、批量下载与定制播放完整教程

搭建机器人仿真资产库：V2D Challenge的GLB纹理网格与URDF文件完整使用指南

今日更新

电商智能运营体系从0到1：架构设计、落地步骤与避坑指南

压缩空气储能与热网热惯性联合优化调度：热电联产系统建模与求解

若依框架AI代码生成不稳定？用Agent Skill三层结构+五步法搞定

本周精选

金融AI落地指南：从调研报告到技术选型与工程实践

大模型上下文成本与注意力优化技术深度解析——为何O(n²)平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测

基于GMM的语音识别系统原理与MATLAB实现详解

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？