深入篇 20 / Transformer 逐步拆解 · 2026-11-08

一个 decoder block：残差、RMSNorm 与 SwiGLU

把简洁的数学残差公式还原到 nano-vllm 的双变量 hidden_states / residual 实现。

核心阅读 12–18 分钟练习另需 10–20 分钟前置：15–19：attention 支路

本课只拿下这一件事

把简洁的数学残差公式还原到 nano-vllm 的双变量 hidden_states / residual 实现。

Pre-Norm attention 子层。公式是逻辑视图；内核融合可以延迟加法但要保持相同数据流。 — 公式是逻辑视图；内核融合可以延迟加法但要保持相同数据流。

先用最容易核对的数学版本

x₁ = x + Attention(RMSNorm₁(x))
x₂ = x₁ + MLP(RMSNorm₂(x₁))
RMSNorm(x) = γ ⊙ x / √(mean(x²)+ε)
MLP(x) = [SiLU(xWg) ⊙ (xWu)] Wd

RMSNorm 对最后一个 hidden 维归一化，不减均值；它不是 LayerNorm 的同义词。γ 是可学习的逐通道缩放。ε 避免接近零时的数值问题；本地实现转 float32 算平方均值，再转回原 dtype。SwiGLU 的 SiLU(z)=z·sigmoid(z)，一条投影经非线性控制另一条投影的内容。依据分别是 RMSNorm 论文、GLU Variants 论文和本地 layernorm/activation 源码。

为什么源码没有出现两句直白的 x +=

第一个 decoder 收到 residual=None，先把原始 embedding 留作 residual，再计算 norm。attention 结果和 residual 在 post_attention_layernorm 内相加，得到 x₁ 并保留；归一化后的值送入 MLP。函数返回的 hidden_states 是 MLP 增量，而 residual 是 x₁。下一层的 input_layernorm 再融合相加，完成上一层 x₂。最后一层也需要 final norm 消费剩余残差。

所以不要把层函数返回的 hidden_states 单独当成已经相加完的 x₂。这个分离是为了融合 residual add 和 RMSNorm，减少中间读写；读代码时每走一步都写清“当前 residual 代表哪个数学变量”。

FFN 的参数与形状

本地 gate_up_proj 将两次 H→F 合并为一次 H→2F，再 chunk；down_proj 做 F→H。忽略 bias，三组权重合计3HF个参数。H=1024、F=3072时为9,437,184个参数/层。FFN 的激活大小依赖本轮 token 数；权重大小不随 prompt 长度改变。注意力负责跨位置混合，FFN 负责当前位置的通道变换；经过多层后两者共同构成上下文相关表示。