最近和一场模拟面试的经历撞了个正着候选人很流畅地写出了多头注意力模块手推公式也没问题但当我问起“为什么源码里要把 Q、K、V 一起经过一个大线性层再拆开”“permute 之后张量连续吗”“训练和推理的 mask 到底差在哪”时对方明显卡住了。这不是个例而是很多刚开始接触 Transformer 源码的同学共同的痛点嘴上会讲 attention手上却不知道每一行程式码在干什么。这篇记录就是从那场模拟面试延伸出来的——把 Attention 计算、维度变换、KV Cache、分布式训练这几个高频追问点掰开揉碎讲清楚顺便给出我这些年带人和面试时沉淀的一条源码阅读路线。1. 为什么学了Transformer的人一被追问源码就露馅1.1 一次模拟面试的真实开场这场模拟面试完全没有用真实项目就考察一个最基础的多头注意力模块。候选人 A 同学在共享屏幕上敲出了类似下面这段代码速度很快class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.n_heads n_heads self.d_model d_model self.head_dim d_model // n_heads self.qkv nn.Linear(d_model, 3 * d_model) self.proj nn.Linear(d_model, d_model) def forward(self, x, maskNone): B, T, D x.shape qkv self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] attn q k.transpose(-2, -1) * (self.head_dim ** -0.5) if mask is not None: attn attn.masked_fill(mask 0, float(-inf)) attn attn.softmax(dim-1) out attn v out out.transpose(1, 2).reshape(B, T, D) return self.proj(out)这段代码本身没什么大问题甚至已经在用* (head_dim ** -0.5)替代torch.div来避免额外传参。但面试不是让你默写代码面试官真正关心的是“中间每一步的形状变化是什么”“这样写为什么能工作”“换成view会怎么样”“mask 里填-inf之后 softmax 会不会意外把 NaN 算出来”。当我顺着代码逐行追问之后A 同学承认自己平时只调用封装好的模块很少真的打开源码看 shape 是怎么一路变过来的。这是很多准备模型岗实习的同学的通病背了结构图但没掌握源码的计算图。1.2 面试官真正想考察的底层能力模型后来我把这场模拟面试中反复追问的角度拆解了一下发现面试官考察的其实不是“记性”而是三层能力。第一层是张量形状变换能力。你能不能在心里跟着reshape - permute - transpose - view走完一次 forward能不能说出来什么时候需要.contiguous()什么时候不需要这决定了你调试模型时能不能快速定位维度不匹配的问题。第二层是数值稳定性与计算图意识。为什么缩放因子要除以根号d_k为什么 softmax 前要 mask 成-inf为什么训练设model.train()的时候 dropout 会改输出分布这些细节直接关系到你训练一个大模型时会不会出现 loss 震荡、推理时会不会每步结果都不一样。第三层是工程实现与加速视角。同样的 attention 计算在单卡、多卡、推理、训练这四种场景下源码组织方式完全不同。KV Cache 在哪里生效张量并行时哪些矩阵需要切哪些不需要FlashAttention 到底省的是哪部分显存面试官不会问得太虚但如果你能主动讲到这些和只会背公式的候选人马上就拉开了差距。接下来我按那天追问的顺序把每一条都展开成“源码级”的解析。2. Attention核心计算的源码级拆解Q、K、V到底经历了什么2.1 从公式到实现缩放点积注意力的每个算子的作用先看原始公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。这里最容易被追问的就是sqrt(d_k)到底解决什么问题。假设q和k的每个分量都是均值 0、方差 1 的随机变量它们的点积期望是 0但方差是d_k。如果不做缩放当d_k比较大时点积结果就可能非常大比如落在几百到上千的区间。之后接 softmax这些大的 logits 会直接把概率分布推向极端最大的 logit 对应的概率接近 1其余接近 0。你可以想象一张灰度图几乎只剩纯黑和纯白所有中间细节都没了。这个状态对应的梯度会非常小因为 softmax 在饱和区的导数趋近于 0模型几乎学不动。所以源码里通常写成attn q k.transpose(-2, -1) * (self.head_dim ** -0.5)。这里的head_dim ** -0.5就是1/sqrt(head_dim)不是随手乘的而是在保持点积方差为 1 的尺度。理解这一点后面试官再问你“为什么不用d_model而用head_dim”你就能答上来因为每个 head 内部点积的维度是head_dim d_model // n_heads缩放只作用于这个维度。接着是 softmax。源码里往往写成attn attn.softmax(dim-1)默认不是对最后一维做归一化吗对它的作用是把注意力权重变成总和为 1 的分布并且对负无穷值会自然得到 0 的概率。这里需要注意如果 mask 填充的是-infsoftmax 计算时exp(-inf)是 0不会报错但如果你不小心写成-1e9这样很小的负数exp也可能得到 0不过数值上可能存在上溢风险因此正规实现都用-inf。最后out attn v就是把加权后的 value 聚合起来。很多初学者会忽略这一步的 shapeattn的形状是[B, heads, T, T]v的形状是[B, heads, T, head_dim]矩阵乘出来是[B, heads, T, head_dim]。2.2 reshape、transpose、view的三重陷阱与维度变化细节源码里最让新手头疼的其实是reshape和permute的组合。我们拆开看。输入x的形状是[B, T, D]其中D d_model。经过self.qkv(x)得到[B, T, 3*D]。然后执行qkv.reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)为什么先 reshape 成 5 维再 permute因为qkv最后那个3*D维里顺序是“先放 Q 的全部分量再放 K 的全部分量最后放 V 的全部分量”。如果直接 reshape 成[B, T, 3, n_heads, head_dim]那么这个三维中的第 3 维索引 2实际上表示 Q/K/V 的编号。随后permute(2, 0, 3, 1, 4)把这一维挪到最前面得到形状[3, B, n_heads, T, head_dim]。上面代码里q, k, v qkv[0], qkv[1], qkv[2]其实就在第一维上切“第 0 份、第 1 份、第 2 份”。这里有个很隐蔽的点permute之后张量在内存里不是连续的因为原来的存储顺序是B, T, (Q/K/V), heads, head_dim现在逻辑顺序变成了Q/K/V, B, heads, T, head_dim但底层的物理排列没变。若你在permute之后调用view而不是reshape大概率会报错。正确做法是写q qkv[0].contiguous()之后再用view或者干脆继续用reshape让它内部帮你处理内存拷贝。这部分太适合被连环追问了。面试官通常会这么问“reshape和view的区别”“permute之后为什么需要contiguous()”“如果不加contiguous()后面的矩阵乘能正常跑吗”第一个问题view要求原张量满足内存连续reshape不要求不连续时它会自动复制到连续内存再变换。第二个问题permute只改变逻辑视图不改变物理布局只要后续算子不是对最后一维做简单索引很多情况下也能跑但会多出非连续访存的开销。第三个问题矩阵乘法q k.transpose(...)本身会自动处理步长信息不一定报错但如果你再派生出其他操作就可能因为张量不连续触发行列访问顺序混乱性能骤降、内存占用却激增。2.3 Mask的三种形态padding mask、causal mask、组合mask在源码中mask 处理是 attention 正确性的关键。不要把 mask 理解成一张简单的布尔面具它要跟实际形状匹配。第一种是padding mask用于忽略句子中补齐部分的 token。形状通常是[B, T]或者[B, 1, 1, T]表示哪些位置是有效 token。在注意力分数矩阵里对 key 维度做 mask也就是把[B, heads, T_q, T_k]中对应 key 位置的分数填成-inf。第二种是causal mask用于自回归任务保证第i个 token 只能看到0~i的位置不能看到未来。双向 attention 的 Transformer 如果用在生成任务上不加 causal mask 就等于作弊。源码中常用torch.triu生成一个上三角矩阵把右上角填成-infcausal_mask torch.triu(torch.ones(T, T, dtypetorch.bool), diagonal1)然后在注意力分数上执行attn attn.masked_fill(causal_mask, float(-inf))。第三种是组合 maskpadding mask 和 causal mask 同时存在时要先做逻辑或或者逐个填-inf。实际实现里可能会把两个 mask 都 expand 成相同的[B, 1, T, T]形状再融合成一个 mask。常见错误是把 padding mask 直接 broadcast 到[B, T, T]时维度对不上结果 mask 被作用于 query 维度而不是 key 维度模型看到了一堆错误位置。这里还有一个面试官爱问的小点masked_fill是不是 in-place 的严格说masked_fill在 PyTorch 文档里“不是原位操作”它会返回一个新的张量。很多人会把attn attn.masked_fill(...)里的重新赋值省略写成attn.masked_fill(...)导致 mask 根本没生效。这属于写法上的隐性 bug但面试时能主动提出来说明你真的抠过源码。3. 那些“一读就懂、一问就忘”的Transformer源码细节3.1 多头注意力为什么要按head拆而不是直接算一整块矩阵原始论文给出的 motivation 是让模型在不同的表示子空间里关注不同信息。但用代码的角度看多头的本质是什么是把d_model维的向量切分成n_heads份每一份维度是head_dim然后每一份单独做 attention最后拼回去。如果我们不做多头直接让Q, K, V的维度都是d_model注意力计算就是一个d_model × d_model的矩阵乘相当于用一个高维参数矩阵去拟合所有的注意力模式。这个做法的缺点是所有信息混在同一组矩阵里很难让模型学到“有的头关注位置关系、有的头关注语义关系、有的头关注局部句法”这类分离表示。多头相当于给模型提供了多个相关性度量通道每个头只在自己的低维子空间里建模避免单个参数矩阵把所有模式混成一团。从参数数量上看多头并不会增加多少参数。一个多头注意力模块里qkv线性层的输出还是3*d_modelproj线性层的输出还是d_model。只不过在内部把矩阵分块了。很多同学以为多头会带来参数量爆炸其实没有。面试现场我还追问了 A 同学一句“如果d_model不能整除n_heads代码会发生什么”答案是直接报错因为d_model // n_heads会向下取整剩下的维数要么被丢掉要么导致维度对不上。所以实际实现中d_model必须能被n_heads整除。这在设计模型时就要提前想好。3.2 LayerNorm与残差连接的顺序pre-norm vs post-norm的源码差异Transformer 论文里的结构是x Sublayer(LayerNorm(x))也就是先归一化进子层再加残差。但今天大量开源实现的 Transformer Block 是另一种顺序x x attn(LayerNorm(x)) x x ffn(LayerNorm(x))这就是 pre-norm 结构。区别在哪呢我们先看 post-norm论文原始版本的逻辑先过attention或ffn再归一化再加到输入上。残差分支上直接输入的是未经归一化的表征优化时梯度要穿过一个非线性层再加回来路径比较曲折容易造成训练初期不稳定。pre-norm 把归一化放在子层之前主分支上经过残差加和的是已经归一化后的信息相当于把“尺度归一”这件事提前做掉让每一层输入分布更稳定。另一个好处是即便网络很深梯度也能通过残差短接线直接回传和归一化模块相对独立。因此现在主流大模型几乎全是 pre-norm。面试时还可以补充一个细微差别pre-norm 在输出层的均值和方差统计可能与 post-norm 不同所以有些实现最终还会在最外面接一个 final LayerNorm。这也是为什么你会在源码里看到 Transformer Block 的输出过一个ln_f或final_norm。3.3 推理加速背后的KV Cache源码中缓存的是什么每次生成一个 token 时如果都把之前所有 token 的 Q/K/V 重新算一遍计算量会随序列长度线性增长非常浪费。KV Cache 的思路是在自回归解码中每个新 token 的 Q 需要计算但它对应的 K、V 会参与后续所有 token 的注意力计算所以只要把已经生成的 token 的 K、V 保存下来下一轮直接拼接新 token 的 K、V再用新 query 和拼接后的 cache 做 attention。具体到源码大致是if past_key_value is not None: k torch.cat([past_key_value[0], k], dim-2) v torch.cat([past_key_value[1], v], dim-2)这里在倒数第二维序列维度上拼接新的past_key_value再传给后续层。面试官可能会追问“为什么只缓存 K 和 V不缓存 Q”关键点在于 Q 只用于当前 token 的位置之后的 token 不会再去看“当前 token 的 query”但后续 token 需要和“当前 token 的 key”做匹配也要用“当前 token 的 value”做加权所以 K 和 V 是必须长期保留的。Q 每步都不同用完就没有保留意义了。如果继续追问“为什么 cache 增大了显存反而生成变快”可以从计算量角度答没有 cache 时第 N 步要重新计算前 N 个 token 的 K/V有 cache 后每步只计算新的一个 token 的 K/V操作量从 O(N^2) 降到 O(N)省下的时间远大于拼接操作的开销。当然代价是显存这也是长文本生成时显存容易爆炸的原因之一。4. 从单卡到多卡Transformer源码在分布式场景下的隐藏逻辑4.1 数据并行时梯度同步的瓶颈为什么集中在LayerNorm和Embedding面试官把一个单卡 Transformer 放到多卡场景问的第一句往往是“数据并行怎么做”最直接的答案是把一个小批量切成多份每张卡一个副本前反向独立计算算完梯度后用 all-reduce 做平均同步再更新各自副本上的参数。那为什么说瓶颈集中在 Embedding 和 LayerNorm先看 Embedding 表。假设字典大小是 100k隐藏维度是 4096Embedding 就有 4 亿个参数占了模型参数很大一块。数据并行下每张卡都要维护一份完整的 Embeddingall-reduce 梯度的通信量和参数数量成正比因此 Embedding 的通信开销非常大。LayerNorm 的参数虽然少但它也参与反向传播会生成梯度且在每张卡上都会计算和同步中间包括均值和方差统计同样逃不掉通信。面试中如果答到这里还能主动提一句“所以很多人会做 Embedding 的分片比如让每张卡只管部分词表再通过 all-gather 拼起来”马上会加分。部分并行训练框架里甚至会单独为 Embedding 做词表并行就是为了绕过这个通信瓶颈。4.2 张量并行中Attention的列切分与行切分如何避免大量通信张量并行解决的是“单张卡放不下模型”的问题。思路是把一个矩阵乘法切到多张卡上算然后在合适的位置做通信。对 Transformer Block 里的 MLP 来说第一层线性层一般做“列切分”即把权重按列切成多块每个并行 rank 各算输出的一部分第二层线性层做“行切分”把权重按行切成多块每个 rank 在各自分块上做局部计算最后通过 all-reduce 把各部分结果累加起来。这样设计的好处是只在两层之间做一次 all-reduce中间不需要全通信避免每层都同步。Attention 部分更特殊。多头的天然并行性让它很适合按 head 维度切分把n_heads个 head 平均分给多个 rank每个 rank 算自己负责的那些 head算完输出后把各 rank 的[B, local_heads, T, head_dim]拼回[B, n_heads, T, head_dim]再通过一个通信操作完成拼接。这套方案在 Megatron-LM 之类的并行范式里被称为多头并行。但要注意拼完之后还有一个proj线性层它会对d_model维做线性变换。为了减少通信通常会把这个proj的输入分布到各 rank再做一次 all-reduce 或者把拼接和 proj 融合起来。所以面试时不要只记住“按 head 切分”要能说清楚通信点到底在哪几次矩阵运算之间。4.3 序列并行与FlashAttention的融合技巧面试中的加分回答序列并行是分布式场景里比较进阶的题目。它把序列长度T切到多张卡上每张卡只处理一个序列片段。这样做能同时降低 per-card 的激活显存和计算量但有一个天然矛盾注意力分数矩阵的 softmax 归一化需要跨整个序列的统计量如果序列被切开每个片段算出的 softmax sum 并不全局。FlashAttention 的思路恰好能帮忙。它通过 online softmax把归一化统计量拆成可累加的部分和在扫描分块时动态更新。序列并行场景下每个 rank 可以独立扫描自己的序列块再把局部的统计量通过通信合并。面试时如果能说出“online softmax 维护当前最大值和当前指数和合并分块时用新最大值重新缩放旧的指数和”这个细节几乎等于给了满分。这里还有一层融合注意力计算里的 dropout、mask、矩阵乘都可以在切分块内完成尽力避免把完整的T × T注意力矩阵物化出来。FlashAttention 本身就是不物化大矩阵、利用分块减少 HBM 访问的典型这和序列并行天然契合。回答时可以把两者的共同点概括为“把大的矩阵运算拆成小块用少量额外计算换大量显存和访存”。5. 高频连环追问拆解像面试官一样追问直到源码细节5.1 第一追问链为什么除以根号d、为什么软max、为什么多头这条链我几乎每次模拟面试都会问一遍每个问题都可以往下钻两层。先问“为什么除以根号 d”不缩放会怎样计算方差的公式是什么如果 d 是 512缩放和不缩放的注意力分数量级差多少要能答出“点积方差正比于维度”这一点。再问“为什么用 softmax”能不能换成 sigmoidsoftmax 保证概率和为 1可以处理变长序列sigmoid 不能。当然也存在一些变体用 ReLU attention但那是另一套设计。面试里答“softmax 是为了归一化”还不够要补充“它让较大的 logits 在指数运算下受到抑制梯度不会全部消失”。最后问“为什么多头”除了子空间表征之外还要提并行计算的方便性。多个 head 可以在同一个 batch 内矩阵计算把[B, heads, T, head_dim]的四维张量合并成三维矩阵乘底层 GEMM 效率更高。也就是说多头不只是建模能力的需求也是工程计算效率的需求。这个角度很多人没提提出来会让人眼前一亮。5.2 第二追问链训练时用dropout为什么推理时通常不开启这条链围绕 dropout 展开。源码里nn.Dropout(p0.1)的行为受self.training控制。训练时每个元素有概率p被置 0同时剩余元素除以1-p目的是保持期望不变。推理时如果还开着 dropout模型输出会随机波动生成结果不稳定而且等于在预测时引入了无意义的噪声。所以推理前要调用model.eval()源码里模块内部的dropout.training标志也会被设成 False。面试官通常还会追问一句“如果训练时忘记改成 eval模型会怎样”A 同学当时答“结果可能每次不一样”但更关键的是“性能会下降因为 dropout 会破坏已经习得的表征尤其在生成任务里可能出现明显的随机重复或突变”。能抓住这两层的候选人证明他不是只会调 API。5.3 第三追问链显存不够怎么办从重计算到混合精度的模型策略当模型大到单卡放不下面试官大概率会问“训练时怎么省显存”。大多数人都能想到混合精度但能说出具体手段的人不多。第一个方向是重计算也叫 activation checkpointing / gradient checkpointing。训练时我们不保存每一层的所有激活只保存每层的输入或选定的 checkpoint反向传播用到中间激活时重新算一遍。这样做的效果是前向激活显存从 O(层数 × 每层激活大小) 降到 O(层数 × checkpoint 数量)典型节省比例能到 60% 以上代价是约 30%~50% 的额外计算时间。第二个方向是混合精度。用 FP16 或 BF16 存权重、梯度和激活能大幅减少显存但 FP16 容易造成梯度下溢所以需要 loss scalingBF16 的动态范围更大但精度更低。现在很多框架默认就是 BF16 重计算 序列并行三件套。第三个方向是优化器状态分片比如 ZeRO 的思想把优化器状态平均分到多卡每张卡只维护自己分到的部分。这样在 64 卡上优化器状态每卡只占原来的 1/64。提到这个面试官如果深问就能引出“参数、梯度、优化器状态三层分片”的顺序。这条链回答完毕基本就覆盖了大模型训练中最常被问的显存优化全家桶。6. 三个月吃透Transformer源码的实践路线6.1 用日志和断点把张量维度“看”在手里我建议所有准备模型源码面试的同学不要直接对着源码发呆而是写一个“维度探测器”。比如在.forward里加几个日志打印def log_tensor(name, t): print(f{name}: shape{tuple(t.shape)}, dtype{t.dtype}, fmean{t.mean():.4f}, std{t.std():.4f})然后在每个子层输入输出调用log_tensor。这样你就建立了“抽象维度到实际数值”的感知。比如在 attention 内部你能亲眼看到q在进入 k.transpose(-2, -1)之前的 shape 是[B, heads, T, head_dim]也能感觉到为什么需要转置。善用 PyTorch 的register_forward_hook也能省事但自己 print 更直观。调试时随便跑一个小序列序列长度 8batch 2隐藏维度 8打印出来的维度变化会非常清晰。把每一层的[B,T,D]变化记录在一张纸上基本上不需要死记源码了。6.2 必须手写一遍的模块清单和验证方法我给 A 同学的建议是在三个月内按顺序手写以下几个模块并做一致性验证顺序手写模块验证方法与常见错误1单头自注意力和公式直接计算对比检查缩放因子是否遗漏2多头注意力检查reshape/permute维度对比库实现输出是否一致固定随机种子3带 causal mask 的 attention让第 i 个位置只看前 i 个位置观察梯度是否为下三角4Transformer Blockpre-norm检查残差连接位置注意LayerNorm是放在子层前还是子层后5KV Cache 推理对比开启 cache 和不开 cache 生成的 logits 是否等价6简单的分布式训练脚本至少跑通数据并行观察梯度同步前后的一致性每个模块验证时用同一个输入张量固定所有随机种子对比手写实现和标准实现输出的最大误差max(|out_hand - out_lib|) 1e-5。如果误差偏大先检查自己是否漏了head_dim ** -0.5再检查维度 permute 是否多转了一次。这套排查流程本身就是很宝贵的面试素材。6.3 最后的避坑建议别背源码要理解设计意图吃透源码不等于把每行代码背下来。面试官问你“为什么这里用 reshape 而不用 view”本质上不是考察 API 用法而是考察你是否知道 reshape 和 view 背后的内存连续性问题。同理“为什么推理时要关 dropout”也不是考你 API 开关而是考你对训练和推理分布差异的感知。我自己的体会是读 Transformer 源码最顺的路径是先手推一遍 attention 的数学公式再对照一段最小的多头注意力代码最后再往工程方向扩展mask、KV Cache、分布式、 FlashAttention。每一步都问自己三个问题——输入形状是什么、输出形状是什么、如果我把某个算子换成替代实现会怎样。带着这三个问题去读三个月后你不仅能应对模拟面试还能在真实模型调试时快速定位问题。最后再分享一个实用小技巧很多同学把源码读完了还是记不住维度那就自己画一张“张量流向图”。不需要软件纸上画也行标清楚每层开始时的[B,T,D]以及进了 attention 后变成[B,heads,T,head_dim]出来又变回[B,T,D]。这张图比任何笔记都管用也是你面试时可以随手在草稿纸上复现的重要武器。
