很多人第一次接触 Transformer 时注意力公式只有一行Attention(Q,K,V) softmax(QK^T / √d_k) V。但一旦真正动手写推理引擎、调显存占用、优化生成速度就会发现这一行公式背后全是工程细节QKV 投影出来的张量形状怎么变化、多头怎么拆、因果 mask 放在哪一步、分数为什么要缩放以及最关键的——生成下一个 token 时之前算过的 K 和 V 到底能不能复用、怎么复用。这篇文章的目标就是把“一次注意力计算”从头到尾拆开。我会从输入张量进入投影矩阵开始一步一步走完自注意力里的所有张量变换和数值操作然后把推理场景下的 KV Cache 单独拎出来讲透最后给出一份带缓存的 Attention 实现伪代码以及我自己踩过的坑。如果你正在自己实现 Transformer 推理、读推理框架源码或者只是想把注意力这块的基础打牢这篇应该能帮你把“公式”和“代码”之间的缝隙填上。1. 从输入到 QKV一次投影里发生了什么1.1 输入如何变成三份张量先明确一个前提我这里说的是一次“标准多头自注意力”输入是上一层的输出序列记作x形状是[batch, seq_len, hidden_dim]。hidden_dim是整个模型的隐藏维度通常也叫d_model常见的取值是 512、768、1024、4096 这类。注意力需要三份不同的“视角”Query查询、Key键、Value值。Transformer 的做法是用三个独立的线性投影矩阵分别把x映射过去Q x W_q K x W_k V x W_v其中W_q、W_k、W_v的形状都是[hidden_dim, hidden_dim]在部分实现里是[hidden_dim, num_heads * head_dim]效果一样。乘完以后Q、K、V的形状仍然是[batch, seq_len, hidden_dim]。也就是说投影这一步不改变序列长度也不改变 feature 数量只是把每个位置上的向量从“当前语义空间”映射到“用于匹配/检索的语义空间”。这里有一个很多初学者忽略的点Q、K、V 并不是三个“新的序列”而是同一个序列在不同线性变换下的表示。Q代表“我在找什么”K代表“我有什么”V代表“我能提供什么内容”。可以类比图书馆检索Q是你输入的关键词K是每本书的索引标签V是书本身的内容。注意力就是用关键词去匹配索引找到内容再按匹配程度取加权内容。1.2 多头视角reshape、转置和维度含义投影完的Q/K/V不会直接去做点积而是先拆成多个头。隐藏维度hidden_dim会被均分成num_heads份每份大小是head_dim hidden_dim / num_heads。比如常见配置中hidden_dim512num_heads8那么head_dim64。拆分的过程在张量运算里长这样先把最后一维hidden_dimreshape 成[num_heads, head_dim]所以形状从[batch, seq_len, hidden_dim]变成[batch, seq_len, num_heads, head_dim]。然后要把“头”这一维搬到序列长度前面变成[batch, num_heads, seq_len, head_dim]。搬维度在 PyTorch 里就是transpose(1,2)在 NumPy 里就是swapaxes。为什么要搬维度因为点积注意力需要在“每个头内部”独立计算[seq_len, seq_len]的注意力矩阵。如果不把num_heads放到 batch 和 seq_len 之间后续批量矩阵乘法的语义会非常别扭。搬完之后每个头就是一个小 batch相互之间完全独立。实际实现里这一步的 reshape 和 transpose 成本不高但特别容易错。我记得有一次自己写 attention忘了在 reshape 之后补一个contiguous()结果后面 matmul 直接报错。原因是transpose返回的 tensor 在内存里不是连续的必须contiguous()把数据重新排一遍才能继续。1.3 投影权重的初始化与为什么 head_dim 要小W_q、W_k、W_v并不是随便初始化的一般会从正态分布采样标准差通常是0.02大模型常用或按xavier_uniform_初始化。初始化不好会影响训练稳定性和收敛速度不过推理阶段权重已经固定不用关心这一点。关键问题是为什么head_dim这么小很多模型只有 64 或 128而不是直接把整个hidden_dim拿去做点积原因有二。第一每个头只在一个较小的子空间里做匹配多个头并行能捕捉不同维度的依赖关系类似多个“专家”从不同角度做检索第二点积的数值大小和维度成正比如果head_dim是 4096点积分数很容易到几百甚至几千softmax 会饱和到几乎 one-hot梯度非常小训练也容易炸。小维度配合√d_k缩放能让点积分数落在合理区间。这一点在下一章详细说。2. 分数、掩码、归一化点积注意力中间每一步2.1 点积分数怎么算为什么除以 √d_k现在Q和K的形状都是[batch, num_heads, seq_len, head_dim]。注意力分数是每个 query 位置和所有 key 位置的点积scores Q K.transpose(-2, -1)K.transpose(-2, -1)把[seq_len, head_dim]变成[head_dim, seq_len]所以最终scores的形状是[batch, num_heads, seq_len_q, seq_len_k]。在普通自注意力里seq_len_q seq_len_k L所以就是[batch, num_heads, L, L]。这个L×L矩阵就是后面所有内存和计算瓶颈的来源。然后要除一个缩放因子1/√d_k。为什么假设Q和K里的每个元素都是均值 0、方差 1 的独立随机变量那么两个d_k维向量的点积方差会变成d_k。也就是说维度越大点积分数的波动范围越大。如果直接过 softmax大数值会让概率分布非常尖锐梯度趋近于零。除以√d_k可以把点积方差重新拉回 1 左右让 softmax 的输入落在比较温和的区域。这个细节在训练和推理精度上都很重要。实际工程中scores都是 FP16 或 BF16 存储如果点积数值过大很容易溢出或者丢失精度。缩放是必须的。2.2 因果 mask 写在哪个位置在自回归模型GPT 这类 decoder-only 模型里注意力必须满足因果性位置i只能看到0..i的历史信息不能看到未来。这个约束通过 mask 实现。标准做法是把未来的 key 位置在scores里对应的值设为-inf或一个非常大的负数这样经过 softmax 后这些位置的权重会变成 0。具体来说对于形状[batch, num_heads, L, L]的scores行代表 query 位置列代表 key 位置那么矩阵右上三角不包括对角线就全部置为-infmask torch.triu(torch.ones(L, L), diagonal1).bool() scores scores.masked_fill(mask, -inf)对角线保留表示当前位置能看到自己。这一步必须在 softmax 之前做不能先 softmax 再 mask因为 softmax 的分母会包含未来位置的 exp 值mask 就失去意义了。有个容易搞错的地方训练阶段因为并行给整个序列算注意力所以需要完整的上三角 mask。推理阶段如果用了 KV Cache每次只生成一个新 queryscores的形状是[batch, num_heads, 1, cache_len]此时所有cache_len个 key 都是历史位置天然满足因果性不需要再额外 mask。但有些实现为了统一还是会生成一个全 False 的 mask注意别弄反。2.3 softmax 后的加权求和数值稳定和输出形状softmax 作用于scores的最后一维也就是对每个 query 位置在“所有 key 位置”上做归一化probs softmax(scores, dim-1)probs的形状依旧是[batch, num_heads, L, L]。然后加权求和output probs V这里V的形状是[batch, num_heads, L, head_dim]probs V得到[batch, num_heads, L, head_dim]。这一步的含义是每个输出位置都是所有 value 向量的加权平均权重就是注意力概率。数值稳定性是容易被忽视的点。理论上 softmax 直接算会有溢出风险因为exp(大数)可能会溢出。标准 softmax 实现会先减去每行的最大值m scores.max(dim-1, keepdimTrue).values probs exp(scores - m) / sum(exp(scores - m))PyTorch 的softmax内部已经做了这个处理所以你不太需要手动写。但如果自己在 CUDA kernel 里实现注意力或者移植到别的框架一定要记得减最大值否则遇到异常大分数时结果会变 NaN。2.4 一次完整计算中的中间张量到底有多大我们用一个具体配置来感受一下batch1seq_len2048num_heads32head_dim128。单层单头的scores形状是[1, 32, 2048, 2048]如果是 FP16每个元素占 2 字节那么一个scores张量需要32 * 2048 * 2048 * 2 268,435,456字节约 256MB。这还只是一层的一个中间结果。如果模型有 32 层同时把每层的中间结果都存下来的话就是 8GB 级别。所以你会看到 Flash Attention 这类优化把重点放在避免物化完整的L×L矩阵上。它的核心思想是分块计算用 online softmax 去流式地把最终结果算出来而不是先算完整 scores 再算 probs 再算输出。这样不仅省显存还能减少 HBM 读写性能自然就上来了。我在实际写推理代码时也习惯用torch.nn.functional.scaled_dot_product_attention因为 PyTorch 自带的这个函数在特定硬件上会自动走 Flash Attention / Memory-Efficient 路径比自己手写Q K^T快很多。3. 多头结果如何汇合输出投影与残差结构3.1 合并多头并做输出投影经过probs V后每个头输出一个[batch, num_heads, L, head_dim]的张量。接下来要把这些头合并回hidden_dim维度。操作是transpose再reshapeattn output.transpose(1, 2) # [batch, L, num_heads, head_dim] attn attn.contiguous().view(batch, L, hidden_dim)合并后的形状又变回[batch, seq_len, hidden_dim]和输入x的形状一致。最后再乘一个输出投影矩阵W_o形状是[hidden_dim, hidden_dim]attn_out attn W_o为什么要做这步输出投影有两个作用。一是把多头的信息混合起来让不同头学到的不同子空间特征可以互相融合二是把维度保持在hidden_dim方便后面接残差和 LayerNorm。很多简化图里只画到多头合并实际模型里W_o是存在的别漏了。3.2 残差和归一化在整体 Block 中的位置注意力层的完整结构是输入先过 LayerNorm再过 Attention加上残差最后再过一个 MLPFeed-Forward再加一次残差。这叫做 Pre-LN 结构现代 GPT 模型基本都是这种。在 Pre-LN 里Attention 部分的代码逻辑是x_norm layernorm(x) attn_out attention(x_norm) x x attn_outattention(x_norm)内部已经包含了 QKV 投影、多头注意力和输出投影。残差连接的意义在于即使注意力内部数值变化很大叠加到原来的x上也不会破坏已经学好的信息。同时残差让梯度可以跨层直接回传训练深层模型时更稳定。工程上有一个实用的经验检查一个注意力实现是否正确可以先看它的残差连接和输出投影形状是否保持[batch, seq_len, hidden_dim]。如果形状变了说明结构写错了。3.3 为什么形状不变是工程的关键Transformer Block 的输入输出形状完全一致意味着它可以像搭积木一样随意堆叠。这个特性在推理引擎里非常重要因为所有层都可以共用一套内存管理逻辑不需要为不同层分配不同形状的缓冲区。我自己做推理优化时经常利用这个性质把多层 Attention 的临时 buffer 复用同一块显存省下不少峰值内存。4. KV Cache让一次注意力在推理时被复用4.1 自回归解码为什么不重新计算历史 K/V推理时模型是一个 token 一个 token 地生成。假设当前要生成第t个 token输入包括之前生成的t-1个 token 和当前新 token。如果没有 KV Cache最粗暴的做法是把长度为t的完整序列重新送入模型从第 1 层开始重新计算所有位置的所有 Q、K、V然后算完整的t×t注意力。这样每生成一个新 token都要重新计算一遍历史所有位置的 K 和 V计算量随序列长度平方增长完全无法接受。但仔细想想第1..t-1个位置的 K 和 V是依赖位置本身的输入和之前的所有层计算得到的。当前新 token 进来后并不能改变历史 token 在某一层经过 LayerNorm 后的向量所以历史 token 对应的 K、V 在数学上是固定的。既然固定就能缓存下来不用重复计算。这就是 KV Cache 的最核心依据。Q 不需要缓存是因为每次新生成的 token 只有一个它的 Q 只用于和所有历史 K 做点积随后不会再被别的 token 使用。换句话说Q 是一次性消费的K 和 V 是会被后续所有新 token 反复消费的。4.2 具体缓存什么、形状与内存估算KV Cache 缓存的是每个 Transformer 层、每个注意力头的 K 和 V 张量。具体形状是key_cache: [batch, num_heads, cache_len, head_dim] value_cache: [batch, num_heads, cache_len, head_dim]cache_len是已经缓存的 token 数量。生成新 token 时先计算出新的k_new和v_new形状是[batch, num_heads, 1, head_dim]然后追加到缓存末尾key_cache torch.cat([key_cache, k_new], dim2) value_cache torch.cat([value_cache, v_new], dim2)这里重点注意缓存的是投影后的 K 和 V不是原始 hidden state。每一层都有自己的 KV Cache不同层之间的 cache 互相独立。层数一多显存压力就上来了。我们来算一个实际例子。假设模型有 32 层num_heads32head_dim128batch1生成序列长度 2048使用 FP16每个元素 2 字节。单个 K cache 的大小是1 * 32 * 2048 * 128 * 2 bytes 16,777,216 bytes ≈ 16MBV cache 同样 16MB。一层就是 32MB。32 层就是 1GB。如果 batch 是 16那就是 16GB。很多长上下文的服务集群显存大部分都被 KV Cache 吃掉了模型权重反而不是最占显存的部分。4.3 工程上的缓存管理优化预分配、GQA 与 PagedAttention因为每次cat都会重新分配内存、复制数据实际推理框架一般会预分配一块最大长度的连续显存比如max_seq_len4096然后维护一个“当前长度”的游标。新 token 的 K/V 直接写入预分配内存的第current_len位置再更新游标。这能避免频繁的内存申请和数据拷贝明显提升长序列生成速度。另一个方向是减少 KV Cache 的量。目前主流大模型大量使用 GQAGrouped Query Attention。原始 MHAMulti-Head Attention每个 query 头配一个独立的 key/value 头KV Cache 很大。GQA 让多个 query 头共享一组 key/value 头比如 32 个 query 头只配 4 个 key/value 头KV Cache 直接缩小到原来的 1/8效果损失很小。MQAMulti-Query Attention更极端所有 query 头共享一组 K/V 头推理速度提升非常明显短文本任务下质量损失也可控。做工程的时候如果模型是 MHA 结构但又想压缩 KV Cache可以考虑做 GQA 的权重合并转换不过需要重新训练或者微调不能直接改。服务端还有 PagedAttention 的思路把连续的逻辑 KV Cache 分散到不连续的内存块中管理按需分配减少碎片和浪费。这属于框架层面的事但如果自研推理服务很值得借鉴。4.4 KV Cache 与训练阶段的注意力区别训练阶段没有 KV Cache因为训练时是一次性把整个序列都喂进去每个位置都同时作为 query 又作为 key/value不存在增量生成。训练阶段需要计算完整的L×L注意力矩阵并用因果 mask 遮住未来。推理阶段有了 KV Cache每个 token 只算一个新的 query 和完整的历史 key 的点积矩阵是1×L然后 softmax 后加权。两者数学上完全等价但计算模式和内存模式差别很大。理解这一点后你就知道为什么很多论文里会说“训练时是计算密集型推理时是访存密集型”。推理时每个 token 的主要操作是读取所有缓存的 K 和 V计算量不大但数据量很大所以瓶颈经常在 HBM 带宽而不是算力。5. 带 KV Cache 的 Attention 实现一份可直接参考的伪代码5.1 完整数据流从新 token 输入到更新缓存先梳理一遍带 KV Cache 的单步推理数据流以 decoder-only 模型的一层 Attention 为例当前 token 经过 embedding加上位置信息得到x形状[batch, 1, hidden_dim]。x过 LayerNorm 得到x_norm。用W_q、W_k、W_v分别投影得到q_new、k_new、v_new形状都是[batch, 1, hidden_dim]。拆成多头并换轴q_new变成[batch, num_heads, 1, head_dim]k_new、v_new同理。把k_new、v_new追加到对应层的key_cache、value_cache里。用q_new和完整key_cache做点积得到scores形状[batch, num_heads, 1, cache_len]。缩放1/√d_k如果已在 QKV 投影后缩放可以省略。单步推理不需要 causal mask直接 softmax。probs与value_cache做加权求和输出[batch, num_heads, 1, head_dim]。换轴合并头乘W_o得到[batch, 1, hidden_dim]加上残差后进入下一层。注意第 6 步这里key_cache已经包含了当前新 token 的 K所以当前 token 能看到自己符合自回归规则。5.2 PyTorch 风格的注意力层代码下面是一段我经常用来做实验的伪代码实现了带 KV Cache 的单层多头注意力。它省略了 LayerNorm、MLP 和 embedding专注展示 Attention 本身。import torch import torch.nn as nn import torch.nn.functional as F class CachedAttention(nn.Module): def __init__(self, hidden_dim, num_heads, max_seq_len, dtypetorch.float16): super().__init__() self.hidden_dim hidden_dim self.num_heads num_heads self.head_dim hidden_dim // num_heads self.max_seq_len max_seq_len self.w_q nn.Linear(hidden_dim, hidden_dim, biasFalse) self.w_k nn.Linear(hidden_dim, hidden_dim, biasFalse) self.w_v nn.Linear(hidden_dim, hidden_dim, biasFalse) self.w_o nn.Linear(hidden_dim, hidden_dim, biasFalse) # 预分配缓存: [batch, num_heads, max_seq_len, head_dim] # 实际使用时会先知道 batch因此在 forward 里初始化。 self.key_cache None self.value_cache None self.current_len 0 def _reshape_to_heads(self, x): batch, seq_len, _ x.shape x x.view(batch, seq_len, self.num_heads, self.head_dim) x x.transpose(1, 2) # [batch, num_heads, seq_len, head_dim] return x def forward(self, x, use_cacheTrue): batch, seq_len, _ x.shape q self._reshape_to_heads(self.w_q(x)) k self._reshape_to_heads(self.w_k(x)) v self._reshape_to_heads(self.w_v(x)) if use_cache: # 如果缓存未初始化先分配 if self.key_cache is None: self.key_cache torch.zeros( batch, self.num_heads, self.max_seq_len, self.head_dim, dtypex.dtype, devicex.device ) self.value_cache torch.zeros_like(self.key_cache) # 新 token 的 k, v 写入缓存 self.key_cache[:, :, self.current_len:self.current_len seq_len, :] k self.value_cache[:, :, self.current_len:self.current_len seq_len, :] v self.current_len seq_len # 只取当前 query q q[:, :, -1:, :] k self.key_cache[:, :, :self.current_len, :] v self.value_cache[:, :, :self.current_len, :] else: k k v v scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) # 训练或非cache模式下需要causal mask单步cache模式不需要。 if not use_cache or seq_len 1: L scores.size(-1) causal_mask torch.triu(torch.ones(L, L, dtypetorch.bool, devicescores.device), diagonal1) scores scores.masked_fill(causal_mask, float(-inf)) probs F.softmax(scores, dim-1) out torch.matmul(probs, v) # 合并多头 out out.transpose(1, 2).contiguous().view(batch, -1, self.hidden_dim) out self.w_o(out) return out这段代码为了演示清晰牺牲了一点效率。实际工程里会把预分配和缓存管理放到推理引擎层面因为不同层的 cache 要统一调度。但核心逻辑就是这样的把新 K/V 写进预分配位置然后只对当前 Q 做注意力。5.3 多 batch、padding 与 mask 的兼容处理上面的代码处理单 batch 很舒服但真实场景往往一个 batch 里有多个请求每个请求序列长度不同。此时 KV Cache 就不能是一个简单的连续块而要做成“每个序列独立管理”。常见做法是每个序列维护自己的current_len。在 batch 维度上做 padding 到统一长度padding 位置不参与计算mask 里置为-inf。对 KV Cache 来说padding 的位置仍然占用显存所以很多框架会动态移动 cache 以压缩内存或者用 PagedAttention 分页管理。我自己在还没用成熟框架、手写服务时踩过的坑是把多序列的 KV Cache 直接cat成一个 batch但不同序列长度不同结果短的序列会被长的序列“污染”。后来老老实实做了 per-sequence cache才解决。如果只是做实验直接 pad 到 max_len 就够用如果要上生产还是用成熟的推理框架更靠谱。6. 实践中的坑与排查口诀6.1 转置和维度导致的经典错误我在写注意力时最常犯的错误是在reshape和transpose后忘了contiguous()。原因很简单view要求张量在内存中是连续的而transpose会改变内存步长导致不连续。解法是x x.transpose(1, 2).contiguous().view(batch, seq_len, self.num_heads, self.head_dim)另一个经典错误是Q K^T时把转置写反。可以用一个小技巧验证先构造一个简单的随机输入让seq_len1如果 shapes 不报错再把scores的形状打印出来确认是[batch, num_heads, seq_len_q, seq_len_k]。如果seq_len_q和seq_len_k对调了说明转置维度错了。6.2 mask 位置放错造成的信息泄漏有一次我写训练代码把 mask 放到了 softmax 之后结果损失下降很慢而且模型明显的“作弊”式收敛——验证时效果极差。因为先 softmax 再 mask未来位置的权重并没有从分母里去掉相当于泄露了未来信息。正确顺序是scores - masked_fill(-inf) - softmax。排查时可以取一个很小的序列如L3肉眼检查 softmax 之后的下三角矩阵是否严格为 0。如果对角线以下的非零元素出现在右上方就是 mask 的顺序或位置错了。6.3 KV Cache 显存爆炸我第一次给长文本开发生成服务时预期只考虑模型权重显存结果跑了 2 万个 token 的 prompt 之后直接 OOM。原因就是没有预估 KV Cache。建议你写一个小函数估算kv_bytes 2 * num_layers * batch * max_seq_len * num_heads * head_dim * bytes_per_element其中2代表 K 和 V。然后对比显卡容量。如果超了优先考虑三种手段缩短单请求长度、降低 batch、改用 GQA 结构的模型。实在不行做 KV Cache 量化把 K/V 压缩成 int8精度损失相对可控。6.4 性能优化方向与经验如果你发现推理速度越来越慢尤其是序列变长后速度下降明显原因大概率不是算力不够而是显存带宽不够。每次生成一个 token都要把完整 KV Cache 从 HBM 读到计算单元里。序列越长读的数据量越大。优化方向按性价比排序用 FlashAttention 等融合算子减少中间张量读写。用 GQA/MQA 减少 KV Cache 大小。预分配 KV Cache 内存避免运行时cat。对长上下文可以考虑把 KV Cache 放在更高带宽的存储层级或者做量化。如果你是在 PyTorch 里做实验直接用F.scaled_dot_product_attention它能自动选择核函数比自己手写的matmul softmax快很多。但它的参数里有attn_mask单步推理时别传因果 mask否则可能多一次不必要的内存操作。6.5 精度问题半精度下的注意力分数FP16 的表示范围有限如果head_dim比较大或者模型权重方差较大缩放前的QK^T可能超过 FP16 能表达的最大值。此时 scores 会出现 infsoftmax 后变 NaN。排查方法很简单在scores后面加一行打印最大值看看是否在合理范围通常应该在 10 以内。如果出现上百甚至上千的值检查是否忘了除以√d_k或者在混合精度训练时需要把某些计算保留在 FP32。BF16 比 FP16 的指数范围大很多所以现在新模型多采用 BF16 来降低这种风险。这个内容后续还可以往两个方向扩展一是把 Flash Attention 的 online softmax 原理单独拆开讲二是深入 KV Cache 的量化与分页管理。如果你正准备优化自己的推理服务我建议先从复现上面那段 CachedAttention 开始把维度和 mask 写清楚再引入更高级的优化这样每一步都能验证不会栽在莫名其妙的形状错误上。