【AI Infra小白入门⑦】Transformer架构工程拆解从Attention到MLA的优化密码本文是「AI Infra 小白入门」系列的第七篇。Transformer是所有大模型的骨架也是AI Infra优化的主战场。本文不从算法角度讲Transformer而是从工程视角拆解——每个组件的计算过程是什么复杂度来自哪里由此引出了哪些优化技术一、写在前面为什么AI Infra要懂Transformer你可能会问“我是搞AI Infra的又不是搞算法的需要懂Transformer吗”答案是太需要了因为AI Infra的所有优化都是围绕Transformer的结构特点展开的为什么FlashAttention这么火因为Attention有O(N²)复杂度为什么KV Cache这么重要因为推理时每步都要读历史KV为什么量化主要针对FFN层因为FFN占了参数的2/3为什么GQA/MLA这么火因为KV Cache是推理显存的大头不懂Transformer结构你就不知道为什么要做这些优化也不知道优化的重点在哪里。先上一张架构全景图二、Self-Attention大模型的心脏2.1 计算过程大白话版Self-Attention的核心思想是让每个token都看看其他token然后根据相关性把它们的信息加权汇总起来。公式大家都见过QXWQ,KXWK,VXWVQ XW_Q,\quad K XW_K,\quad V XW_VQXWQ​,KXWK​,VXWV​Asoftmax(QKTdk),OAVA \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right),\quad O AVAsoftmax(dk​​QKT​),OAV用人话拆解成四步投影输入X分别乘三个矩阵WQ、WK、WV得到Q、K、V三个矩阵打分Q和K做点积QK^T得到注意力分数——分数越高两个token越相关归一化分数除以√dk防止数值太大再过softmax变成0-1之间的权重加权求和用权重对V加权求和得到最终输出2.2 O(N²)复杂度一切优化的根源重点来了Attention的复杂度是O(N²)——N是序列长度。为什么因为QK^T这个操作Q的形状是(N, d)K^T的形状是(d, N)相乘结果是(N, N)的矩阵FLOPs 2 × N × N × d 2N²d序列长度翻倍计算量和显存占用都变成4倍这就是为什么有这么多Attention优化技术FlashAttention不把完整的N×N矩阵存到显存里分块计算减少显存访存PagedAttention用分页的方式管理KV Cache提高显存利用率稀疏/滑窗Attention不是每个token都看所有token只看一部分减少计算量记住这个结论Attention的O(N²)复杂度是大模型推理和训练优化的核心矛盾之一。2.3 训练 vs 推理瓶颈完全不同Attention在训练和推理时的性能瓶颈完全不一样这个一定要搞清楚阶段计算方式瓶颈在哪训练整段序列一次算完并行计算瓶颈N²矩阵的计算量很大推理Prefill输入prompt时和训练类似计算瓶颈同样是N²矩阵推理Decode每步只算1个新token和历史KV的点积访存瓶颈KV Cache的显存读取带宽Decode阶段为什么是访存瓶颈因为decode时每生成一个新token都要读所有历史KV很大和新token的Q做点积计算量很小计算量很小但要读的数据量很大 →算术强度极低 → memory-bound这就是为什么LLM推理优化中KV Cache管理如此重要——怎么让KV更小、怎么让KV读得更快是关键。三、FFN参数量的大头FFNFeed-Forward Network前馈网络是Transformer的另一个核心组件。公式很简单FFN(x)W2 σ(W1xb1)b2\text{FFN}(x) W_2\,\sigma(W_1 x b_1) b_2FFN(x)W2​σ(W1​xb1​)b2​就是两个线性层中间加一个激活函数通常是SiLU/GELU。3.1 FFN占了多少参数通常FFN的中间维度是hidden size的4倍。比如hidden size是4096FFN中间维度就是16384。这样算下来FFN的参数量大约占整个Transformer模型的2/3。工程意义量化优化的重点在FFN——参数多收益大张量并行TP的切分重点也在FFN——参数量大切了之后省显存四、位置编码RoPE是主流Attention本身是位置无关的——它只看内容相关性不管token的顺序。但语言是有顺序的所以需要额外加入位置信息。常见的位置编码方式方法思想代表模型工程影响Sinusoidal用正弦/余弦函数生成固定位置编码原始Transformer无参数长序列外推差RoPE旋转位置编码对Q和K施加旋转来注入位置信息LLaMA系列、Qwen等主流大模型现代LLM标配长序列外推需NTK/YaRN插值RoPE的工程要点RoPE是对q和k施加旋转变换可以理解为复数乘法在实际的kernel中RoPE通常和Attention融合在一起算子融合减少访存长序列场景需要做NTK插值或YaRN插值来提升外推能力五、LayerNorm与残差连接小但重要5.1 Pre-Norm vs Post-Norm结构形式特点Post-NormLN(x sublayer(x))原始Transformer用的训练不稳定Pre-Normx sublayer(LN(x))现代LLM标配梯度更稳定现在的大模型几乎清一色用Pre-Norm。5.2 工程优化点LayerNorm 残差连接看起来计算量不大但因为它们在每层都有每个Transformer Block有两个而且通常是内存受限的算子所以常常被优化。常见优化算子融合Fusion把LayerNorm、残差连接、以及前后的算子融合成一个kernel减少多次的显存读写——虽然计算量没变但访存量少了速度就快了。六、完整的Transformer前向过程把上面的组件拼起来一个Decoder-only的大模型比如GPT、LLaMA的前向传播是这样的token ids输入的token编号 ↓ Embedding查表把编号变成向量 ↓ × L个Decoder BlockL是层数通常几十层: ├── Pre-LN → Multi-Head Attention → 残差 └── Pre-LN → FFN → 残差 ↓ Final LN最后一层归一化 ↓ LM Head和Embedding共享权重把隐藏状态变成词表大小的logits ↓ logits → softmax → 采样 → 下一个token自回归生成自回归生成每步只生成1个新token然后把它加到输入里再生成下一个。历史的KV Cache会被复用不用每次都重新算。七、Attention变种演进MHA → MQA → GQA → MLAAttention的变种有很多演进的主线就是一句话压缩KV Cache让推理更快、更省显存。变种K/V头数KV Cache大小代表特点MHA和Q头数一样1×原始Transformer效果最好但KV Cache最大MQA11/hh是头数早期一些模型速度快但质量略降GQAg组比如8组g/hLLaMA2-70B等主流模型效果和速度的折中现在的主流MLA低秩压缩更小DeepSeek-V2/V3把KV压缩成低秩latent进一步减小怎么算KV Cache大小公式FP16为例KV Cache大小 2 × L × (KV头数) × head_dim × seq_len × batch_size × 2字节解释一下2K和V两个矩阵LTransformer层数KV头数 × head_dimK/V的隐藏维度seq_len序列长度batch_sizebatch大小2字节FP16每个元素占2字节为什么GQA/MLA这么火因为它们把KV Cache降下来了可以支持更长的序列、更大的batch从而提高推理吞吐。这就是我们之前说的不可能三角——牺牲一点精度换取更多显存和更高速度。八、从Infra视角看Transformer优化点全在这最后我们从AI Infra的视角总结一下Transformer各个组件对应的优化方向组件特点对应的Infra优化AttentionO(N²)复杂度、decode时memory-boundFlashAttention、PagedAttention、KV Cache量化、Prefix CacheFFN参数量大约2/3量化、张量并行切分、算子融合LayerNorm/残差访存密集、每层都有算子融合FusionRoPE计算量小和Attention融合整体结构层次化、可并行3D并行DP/TP/PP、流水线调度九、总结一下今天这篇文章我们从工程视角拆解了Transformer架构Self-AttentionO(N²)复杂度是一切优化的根源训练时compute-bounddecode时memory-boundFFN参数量占约2/3是量化和TP切分的重点位置编码RoPE是主流长序列需要插值LayerNorm残差小但每层都有常做算子融合MHA→MQA→GQA→MLA演进主线就是压缩KV Cache降低推理显存下一篇预告我们来聊聊AI Infra的主战场——PyTorch框架。从Infra工程师的角度看PyTorch有哪些必须掌握的特性Tensor、autograd、Module、checkpoint、profiler……一网打尽。本系列导航第1篇一文搞懂什么是AI基础设施五大方向与不可能三角第2篇大厂AI Infra组织架构大揭秘第3篇AI Infra岗位路线图级别、薪资与成长路径第4篇AI Infra学习路线图从入门到精通的完整指南第5篇AI Infra必备编程基础Python、C与Linux速通指南第6篇AI Infra数学基础速成线代、概率与微积分工程向第7篇Transformer架构工程拆解从Attention到MLA的优化密码本文第8篇PyTorch框架Infra实战指南Tensor、autograd与性能调试更多内容持续更新中……