【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，然后根据"相关性"把它们的信息加权汇总起来。

公式大家都见过：

Q=XWQ,K=XWK,V=XWVQ = XW_Q,\quad K = XW_K,\quad V = XW_VQ=XWQ,K=XWK,V=XWV

A=softmax(QKTdk),O=AVA = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right),\quad O = AVA=softmax(dk QKT),O=AV

用人话拆解成四步：

投影：输入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：参数量的"大头"

FFN（Feed-Forward Network，前馈网络）是Transformer的另一个核心组件。

公式很简单：

FFN(x)=W2 σ(W1x+b1)+b2\text{FFN}(x) = W_2\,\sigma(W_1 x + b_1) + b_2FFN(x)=W2σ(W1x+b1)+b2

就是两个线性层，中间加一个激活函数（通常是SiLU/GELU）。

3.1 FFN占了多少参数？

通常FFN的中间维度是hidden size的4倍。比如hidden size是4096，FFN中间维度就是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-Norm	`LN(x + sublayer(x))`	原始Transformer用的，训练不稳定
Pre-Norm	`x + sublayer(LN(x))`	现代LLM标配，梯度更稳定

现在的大模型几乎清一色用Pre-Norm。

5.2 工程优化点

LayerNorm + 残差连接，看起来计算量不大，但因为它们在每层都有（每个Transformer Block有两个），而且通常是"内存受限"的算子，所以常常被优化。

常见优化：算子融合（Fusion）

把LayerNorm、残差连接、以及前后的算子融合成一个kernel，减少多次的显存读写------虽然计算量没变，但访存量少了，速度就快了。

六、完整的Transformer前向过程

把上面的组件拼起来，一个Decoder-only的大模型（比如GPT、LLaMA）的前向传播是这样的：

text 复制代码

token ids（输入的token编号）
  ↓
Embedding（查表，把编号变成向量）
  ↓
× L个Decoder Block（L是层数，通常几十层）:
    ├── 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 → MLA

Attention的变种有很多，演进的主线就是一句话：压缩KV Cache，让推理更快、更省显存。

变种	K/V头数	KV Cache大小	代表	特点
MHA	和Q头数一样	1×	原始Transformer	效果最好，但KV Cache最大
MQA	1	1/h（h是头数）	早期一些模型	速度快，但质量略降
GQA	g组（比如8组）	g/h	LLaMA2-70B等主流模型	效果和速度的折中，现在的主流
MLA	低秩压缩	更小	DeepSeek-V2/V3	把KV压缩成低秩latent，进一步减小

怎么算KV Cache大小？

公式（FP16为例）：

KV Cache大小 = 2 × L × (KV头数) × head_dim × seq_len × batch_size × 2字节

解释一下：

2：K和V两个矩阵
L：Transformer层数
KV头数 × head_dim：K/V的隐藏维度
seq_len：序列长度
batch_size：batch大小
2字节：FP16每个元素占2字节

💡 为什么GQA/MLA这么火？

因为它们把KV Cache降下来了，可以支持更长的序列、更大的batch，从而提高推理吞吐。

这就是我们之前说的"不可能三角"------牺牲一点精度，换取更多显存和更高速度。

八、从Infra视角看Transformer：优化点全在这

最后，我们从AI Infra的视角，总结一下Transformer各个组件对应的优化方向：

组件	特点	对应的Infra优化
Attention	O(N²)复杂度、decode时memory-bound	FlashAttention、PagedAttention、KV Cache量化、Prefix Cache
FFN	参数量大（约2/3）	量化、张量并行切分、算子融合
LayerNorm/残差	访存密集、每层都有	算子融合（Fusion）
RoPE	计算量小	和Attention融合
整体结构	层次化、可并行	3D并行（DP/TP/PP）、流水线调度

九、总结一下

今天这篇文章，我们从工程视角拆解了Transformer架构：

Self-Attention：O(N²)复杂度是一切优化的根源；训练时compute-bound，decode时memory-bound
FFN：参数量占约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与性能调试
更多内容持续更新中......