Building an LLM Inference Engine from Scratch Part I / Module 2
Part I · Foundations Module 2：Transformer Block，算子逐个看
从 embedding 和 residual stream 出发，准确推导 RMSNorm、linear、attention、GQA、RoPE、SwiGLU 与 LM head 的数学形状。
学习目标 能从张量 shape 角度解释 decoder layer 中每个算子的输入和输出。 理解 RMSNorm、RoPE、GQA、SwiGLU 的数学定义和工程直觉。 能把公式对应到 xinfer_core::cpu 的参考实现。 能解释为什么 matmul / GEMV 是推理性能的核心工作负载。 完成 Lab 2：实现每个 op 的 CPU 参考函数并写单元测试。 2.1 Embedding 与 residual streamtokenizer 输出的是整数 token id。进入 transformer 前，模型用 embedding table 做一次查表，将每个 id 映射为 hidden vector：
X 0 [ t ] = E [ x t ] , E ∈ R ∣ V ∣ × d model X_0[t] = E[x_t], \qquad
E \in \mathbb{R}^{|V|\times d_{\text{model}}} X 0 [ t ] = E [ x t ] , E ∈ R ∣ V ∣ × d model 这里 X 0 ∈ R S × H X_0\in\mathbb{R}^{S\times H} X 0 ∈ R S × H S S S H = d model H=d_{\text{model}} H = d model config.json 给出 H = 896 H=896 H = 896 ∣ V ∣ = 151936 |V|=151936 ∣ V ∣ = 151936 rope_theta=1000000.0、rms_norm_eps=1e-6、tie_word_embeddings=true。这条 S × H S\times H S × H residual stream 。
Xℓ
[S,H]
RMSNorm
attention norm
Attention
QKV + RoPE + GQA
+
Uℓ
[S,H]
RMSNorm
mlp norm
MLP
gate/up + SwiGLU
+
Xℓ+1
[S,H]
橙色虚线是 residual skip：子层输出只负责提供“要加上的变化量”
图 2-1： 一个 decoder layer 由 attention 子层和 MLP 子层组成，每个子层都带 residual add。Qwen2 的 decoder layer 采用 pre-norm 结构。归一化先作用于子层输入，子层输出再加回 residual stream：
U ℓ = X ℓ + AttentionBlock ℓ ( Norm ( X ℓ ) ) X ℓ + 1 = U ℓ + MLP ℓ ( Norm ( U ℓ ) ) \begin{aligned}
U_\ell &= X_\ell + \operatorname{AttentionBlock}_\ell(\operatorname{Norm}(X_\ell)) \\
X_{\ell+1} &= U_\ell + \operatorname{MLP}_\ell(\operatorname{Norm}(U_\ell))
\end{aligned} U ℓ X ℓ + 1 = X ℓ + AttentionBlock ℓ ( Norm ( X ℓ )) = U ℓ + MLP ℓ ( Norm ( U ℓ )) 一层 Qwen2 decoder 的完整展开
图 2-1 省略了具体投影矩阵。按 xinfer-model 的 run_forward 顺序，一层 Qwen2 decoder 先对 X ℓ X_\ell X ℓ o_proj 回到 hidden size，并与原始 X ℓ X_\ell X ℓ U ℓ U_\ell U ℓ
A = RMSNorm ( X ℓ ; g attn ) Q = A W Q ⊤ + b Q , K = A W K ⊤ + b K , V = A W V ⊤ + b V Q ~ = RoPE ( Q ) , K ~ = RoPE ( K ) C = GQA ( Q ~ , K ~ , V ) O = C W O ⊤ U ℓ = X ℓ + O \begin{aligned}
A &= \operatorname{RMSNorm}(X_\ell; g_{\text{attn}}) \\
Q &= A W_Q^\top + b_Q,\quad
K = A W_K^\top + b_K,\quad
V = A W_V^\top + b_V \\
\tilde Q &= \operatorname{RoPE}(Q),\quad
\tilde K = \operatorname{RoPE}(K) \\
C &= \operatorname{GQA}(\tilde Q,\tilde K,V) \\
O &= C W_O^\top \\
U_\ell &= X_\ell + O
\end{aligned} A Q Q ~ C O U ℓ = RMSNorm ( X ℓ ; g attn ) = A W Q ⊤ + b Q , K = A W K ⊤ + b K , V = A W V ⊤ + b V = RoPE ( Q ) , K ~ = RoPE ( K ) = GQA ( Q ~ , K ~ , V ) = C W O ⊤ = X ℓ + O 随后对 U ℓ U_\ell U ℓ gate_proj 与 up_proj 产生两个 [ S , 4864 ] [S,4864] [ S , 4864 ] down_proj 投回 [ S , 896 ] [S,896] [ S , 896 ]
B = RMSNorm ( U ℓ ; g mlp ) G = B W gate ⊤ , P = B W up ⊤ S = SiLU ( G ) ⊙ P D = S W down ⊤ X ℓ + 1 = U ℓ + D \begin{aligned}
B &= \operatorname{RMSNorm}(U_\ell; g_{\text{mlp}}) \\
G &= B W_{\text{gate}}^\top,\quad
P = B W_{\text{up}}^\top \\
S &= \operatorname{SiLU}(G)\odot P \\
D &= S W_{\text{down}}^\top \\
X_{\ell+1} &= U_\ell + D
\end{aligned} B G S D X ℓ + 1 = RMSNorm ( U ℓ ; g mlp ) = B W gate ⊤ , P = B W up ⊤ = SiLU ( G ) ⊙ P = S W down ⊤ = U ℓ + D
Qwen2 decoder layer：Attention 子层 + MLP 子层的算子位置
Xℓ
RMSNorm
q_proj
k_proj
v_proj
RoPE
RoPE
GQA softmax(V)
o_proj
+
Uℓ
RMSNorm
gate_proj
up_proj
SwiGLU SiLU(gate) ⊙ up
down_proj
+
Xℓ+1
紫色是线性层（matmul/GEMV）；白色是逐元素或归一化算子；橙色圆圈是 residual add。
图 2-2： 每个算子在真实 Qwen2 decoder layer 中的位置。Attention 子层先产生 $U_\ell$，MLP 子层再产生 $X_{\ell+1}$。步骤 符号 典型 shape（Qwen2.5-0.5B） 对应算子 输入 residual X_\ell [ S , 896 ] [S,896] [ S , 896 ] 上一层输出 Attention norm A A A [ S , 896 ] [S,896] [ S , 896 ] RMSNorm Q 投影 Q Q Q [ S , 14 , 64 ] [S,14,64] [ S , 14 , 64 ] linear + bias K/V 投影 K , V K,V K , V [ S , 2 , 64 ] [S,2,64] [ S , 2 , 64 ] linear + bias（GQA） 位置编码 \tilde Q,\tilde K 同 Q/K RoPE 注意力输出 C C C [S,14,64]\Rightarrow[S,896] GQA attention 输出投影 + 残差 U_\ell [ S , 896 ] [S,896] [ S , 896 ] o_proj + add MLP norm B B B [ S , 896 ] [S,896] [ S , 896 ] RMSNorm 门控 MLP S S S [ S , 4864 ] [S,4864] [ S , 4864 ] gate/up + SwiGLU down 投影 + 残差 X_{\ell+1} [ S , 896 ] [S,896] [ S , 896 ] down_proj + add
2.2 RMSNorm：只归一化均方根Qwen2 使用 RMSNorm。对一行 hidden vector x ∈ R H x\in\mathbb{R}^H x ∈ R H
RMSNorm ( x ) i = x i 1 H ∑ j = 1 H x j 2 + ε g i \operatorname{RMSNorm}(x)_i
= \frac{x_i}{\sqrt{\frac{1}{H}\sum_{j=1}^{H}x_j^2+\varepsilon}}\; g_i RMSNorm ( x ) i = H 1 ∑ j = 1 H x j 2 + ε x i g i 其中 g ∈ R H g\in\mathbb{R}^H g ∈ R H
实现对应 xinfer_core::cpu::rms_norm 的外层循环遍历 token row，内层先计算 mean_sq，再写出 row[i] * inv * weight[i]。shaders/rms_norm.hlsl 用 64 个线程组成一个 threadgroup 处理一行，在 groupshared memory 中归约平方和。
2.3 Linear layer / matmul：主要计算负载Transformer block 中的大部分浮点计算来自线性层。给定输入 X ∈ R S × K X\in\mathbb{R}^{S\times K} X ∈ R S × K W ∈ R N × K W\in\mathbb{R}^{N\times K} W ∈ R N × K
Y = X W ⊤ + b , Y ∈ R S × N Y = XW^\top + b,\qquad
Y\in\mathbb{R}^{S\times N} Y = X W ⊤ + b , Y ∈ R S × N prefill 阶段 S S S S = 1 S=1 S = 1 896 → 151936 896\rightarrow151936 896 → 151936 896 × 151936 ≈ 1.36 × 10 8 896\times151936\approx1.36\times10^8 896 × 151936 ≈ 1.36 × 1 0 8
图 2-3： decode 阶段许多线性层是 GEMV；LM head 因词表巨大而特别重。2.4 Self-attention：Q / K / V、缩放点积与因果 maskSelf-attention 先从同一组 hidden states 生成 query、key、value：
Q = X W Q ⊤ , K = X W K ⊤ , V = X W V ⊤ Q=XW_Q^\top,\qquad K=XW_K^\top,\qquad V=XW_V^\top Q = X W Q ⊤ , K = X W K ⊤ , V = X W V ⊤ 对单个 head，query q t q_t q t k j k_j k j
s t , j = q t k j ⊤ d head s_{t,j} = \frac{q_t k_j^\top}{\sqrt{d_{\text{head}}}} s t , j = d head q t k j ⊤ 除以 d head \sqrt{d_{\text{head}}} d head t t t 0.. t 0..t 0.. t
s t , j = { q t k j ⊤ / d head , j ≤ t − ∞ , j > t s_{t,j} =
\begin{cases}
q_tk_j^\top/\sqrt{d_{\text{head}}}, & j\le t \\
-\infty, & j>t
\end{cases} s t , j = { q t k j ⊤ / d head , − ∞ , j ≤ t j > t 然后对 j j j α t , j \alpha_{t,j} α t , j
Attn ( q t , K , V ) = ∑ j ≤ t α t , j v j \operatorname{Attn}(q_t,K,V)=\sum_{j\le t}\alpha_{t,j}v_j Attn ( q t , K , V ) = j ≤ t ∑ α t , j v j
Causal Attention Mask（深色 = 可见，浅色 = 被 mask）
query 0
query 1
query 2
query 3
query 4
key 0
key 1
key 2
key 3
key 4
第 t 行只能看 key 0..t：这就是 decoder-only 模型的自回归约束。
图 2-4： 因果 mask 保证模型不能偷看未来 token。2.5 Multi-head 与 Grouped-Query Attention（GQA）Multi-head attention 将 hidden 维度拆成若干 head。每个 query head 在独立子空间中计算 attention：
H = n heads ⋅ d head H = n_{\text{heads}}\cdot d_{\text{head}} H = n heads ⋅ d head Qwen2.5-0.5B 中 n heads = 14 n_{\text{heads}}=14 n heads = 14 d head = 64 d_{\text{head}}=64 d head = 64 H = 896 H=896 H = 896 n kv = 2 n_{\text{kv}}=2 n kv = 2 Grouped-Query Attention （GQA）。
kv_head ( h ) = ⌊ h n heads / n kv ⌋ \text{kv\_head}(h)=\left\lfloor \frac{h}{n_{\text{heads}}/n_{\text{kv}}} \right\rfloor kv_head ( h ) = ⌊ n heads / n kv h ⌋ GQA 保留较多 query head，同时减少需要存入 KV cache 的 K/V head 数。对该模型，KV cache 的 head 数从 14 降到 2，容量与读取带宽按比例降低为 MHA 的 2 / 14 2/14 2/14
图 2-5： GQA 减少 KV head 数，直接降低 KV cache 显存。2.6 RoPE：把位置编码成旋转纯 attention 对序列顺序没有内建感知。RoPE（Rotary Position Embedding）不把位置向量加到 hidden 上，而是在 Q/K 的 head 维度中按二维平面旋转。对一 对维度 ( a , b ) (a,b) ( a , b ) p p p
[ a ′ b ′ ] = [ cos θ p − sin θ p sin θ p cos θ p ] [ a b ] \begin{bmatrix}
a'\\ b'
\end{bmatrix}
=
\begin{bmatrix}
\cos\theta_p & -\sin\theta_p\\
\sin\theta_p & \cos\theta_p
\end{bmatrix}
\begin{bmatrix}
a\\ b
\end{bmatrix} [ a ′ b ′ ] = [ cos θ p sin θ p − sin θ p cos θ p ] [ a b ] xinfer_core::cpu::rope 和 shaders/rope.hlsl 采用 Qwen2/HuggingFace 的 GPT-NeoX rotate_half 约定：前半维与后半维配对。频率由 rope_theta 控制，Qwen2.5-0.5B 中该值为 1000000.0：
θ p , i = p ⋅ θ base − 2 i / d head \theta_{p,i}=p\cdot \theta_{\text{base}}^{-2i/d_{\text{head}}} θ p , i = p ⋅ θ base − 2 i / d head RoPE 的作用 RoPE 将绝对位置写入 Q/K 的相位。两个位置的相对距离会反映到 q ⊤ k q^\top k q ⊤ k
2.7 MLP / FFN 与 SwiGLUAttention 在 token 之间聚合信息；MLP 对每个 token 的 hidden vector 独立做非线性变换。Qwen2 使用 gated MLP，其逐元素非线性为 SwiGLU：
SiLU ( x ) = x ⋅ σ ( x ) , SwiGLU ( x ) = SiLU ( x W g ⊤ ) ⊙ ( x W u ⊤ ) \operatorname{SiLU}(x)=x\cdot\sigma(x),\qquad
\operatorname{SwiGLU}(x)=\operatorname{SiLU}(xW_g^\top)\odot(xW_u^\top) SiLU ( x ) = x ⋅ σ ( x ) , SwiGLU ( x ) = SiLU ( x W g ⊤ ) ⊙ ( x W u ⊤ ) 然后再投影回 hidden size：
MLP ( x ) = SwiGLU ( x ) W d ⊤ \operatorname{MLP}(x)=\operatorname{SwiGLU}(x)W_d^\top MLP ( x ) = SwiGLU ( x ) W d ⊤ 其中 ⊙ \odot ⊙ up_proj 产生内容分支，gate_proj 经 SiLU 后产生门控分支；两者相乘后再由 down_proj 回到 hidden size。shaders/swiglu.hlsl 对每个元素执行 silu(gate) * up。
2.8 LM head 与 tied embeddings24 个 decoder layer 结束后，模型得到 X L X_L X L
Z = Norm ( X L ) W lm ⊤ , Z ∈ R S × ∣ V ∣ Z = \operatorname{Norm}(X_L)W_{\text{lm}}^\top,\qquad
Z\in\mathbb{R}^{S\times |V|} Z = Norm ( X L ) W lm ⊤ , Z ∈ R S × ∣ V ∣ 自回归生成只使用最后一行 Z [ S − 1 ] Z[S-1] Z [ S − 1 ] tie_word_embeddings=true；xinfer-model 在这种情况下直接使用 embedding 权重作为 LM head，即 W lm = E W_{\text{lm}}=E W lm = E
性能提醒 LM head 的输出维度等于词表大小。Qwen2.5-0.5B 的 ∣ V ∣ = 151936 |V|=151936 ∣ V ∣ = 151936 151936 × 4 ≈ 608 151936\times4\approx608 151936 × 4 ≈ 608 u32 token id。
小结 本模块把 Qwen2.5-0.5B-Instruct 的 decoder layer 拆成可验证的算子链：embedding 查表得到 [ S , 896 ] [S,896] [ S , 896 ] o_proj 与 residual add，再执行 RMSNorm、gate/up projection、SwiGLU、down_proj 与 residual add。GQA 使 14 个 query head 共享 2 个 KV head，降低 KV cache 容量与读取带宽。最终的 LM head 使用 tied embeddings，将 hidden vector 投影到 151936 维 logits。
Lab 2 实现每个 op 的 CPU 参考函数本实验对应 xinfer-core::cpu。目标是实现可读、确定的 FP32 reference op，作为 DirectML/HLSL kernel 的 correctness oracle。
函数 输入 shape 输出 shape 测试建议 rms_norm[seq, hidden][seq, hidden]单位权重时 RMS 接近 1 matmul[m,k]×[k,n][m,n]乘单位矩阵不变 rope[seq, heads, head_dim]in-place 旋转保持 L2 norm attentionQ,K,V[seq_q, heads, head_dim]mask 后不能看未来 swiglugate, upsame 与手算 SiLU 对比
pub fn rms_norm(x: &[f32], weight: &[f32], seq: usize, hidden: usize, eps: f32) -> Vec<f32> {
let mut out = vec![0.0; x.len()];
for t in 0..seq {
let row = &x[t * hidden..(t + 1) * hidden];
let mean_sq = row.iter().map(|v| v * v).sum::<f32>() / hidden as f32;
let inv = 1.0 / (mean_sq + eps).sqrt();
for i in 0..hidden {
out[t * hidden + i] = row[i] * inv * weight[i];
}
}
out
}
思考与练习 基础 写出 RMSNorm 与 LayerNorm 的主要区别。LayerNorm 同时减去均值并除以标准差：x − μ σ 2 + ε γ + β \frac{x-\mu}{\sqrt{\sigma^2+\varepsilon}}\gamma+\beta σ 2 + ε x − μ γ + β x 1 H ∑ x j 2 + ε g \frac{x}{\sqrt{\frac{1}{H}\sum x_j^2+\varepsilon}}g H 1 ∑ x j 2 + ε x g g g g
基础 若 X ∈ R 4 × 8 X\in\mathbb{R}^{4\times 8} X ∈ R 4 × 8 W ∈ R 16 × 8 W\in\mathbb{R}^{16\times 8} W ∈ R 16 × 8 X W ⊤ XW^\top X W ⊤ W ⊤ ∈ R 8 × 16 W^\top\in\mathbb{R}^{8\times16} W ⊤ ∈ R 8 × 16 X W ⊤ ∈ R 4 × 16 XW^\top\in\mathbb{R}^{4\times16} X W ⊤ ∈ R 4 × 16 X X X K K K W W W N N N [seq=4, out=16]。
进阶 证明二维旋转矩阵保持向量长度：a 2 + b 2 = a ′ 2 + b ′ 2 a^2+b^2=a'^2+b'^2 a 2 + b 2 = a ′2 + b ′2 设 a ′ = a cos θ − b sin θ a'=a\cos\theta-b\sin\theta a ′ = a cos θ − b sin θ b ′ = b cos θ + a sin θ b'=b\cos\theta+a\sin\theta b ′ = b cos θ + a sin θ
a ′ 2 + b ′ 2 = ( a cos θ − b sin θ ) 2 + ( b cos θ + a sin θ ) 2 a'^2+b'^2=(a\cos\theta-b\sin\theta)^2+(b\cos\theta+a\sin\theta)^2 a ′2 + b ′2 = ( a cos θ − b sin θ ) 2 + ( b cos θ + a sin θ ) 2
= a 2 cos 2 θ − 2 a b cos θ sin θ + b 2 sin 2 θ + b 2 cos 2 θ + 2 a b cos θ sin θ + a 2 sin 2 θ =a^2\cos^2\theta-2ab\cos\theta\sin\theta+b^2\sin^2\theta + b^2\cos^2\theta+2ab\cos\theta\sin\theta+a^2\sin^2\theta = a 2 cos 2 θ − 2 ab cos θ sin θ + b 2 sin 2 θ + b 2 cos 2 θ + 2 ab cos θ sin θ + a 2 sin 2 θ
交叉项相消，余下 a 2 ( cos 2 θ + sin 2 θ ) + b 2 ( cos 2 θ + sin 2 θ ) = a 2 + b 2 a^2(\cos^2\theta+\sin^2\theta)+b^2(\cos^2\theta+\sin^2\theta)=a^2+b^2 a 2 ( cos 2 θ + sin 2 θ ) + b 2 ( cos 2 θ + sin 2 θ ) = a 2 + b 2
进阶 Qwen2.5-0.5B 的 n heads = 14 n_{\text{heads}}=14 n heads = 14 n kv = 2 n_{\text{kv}}=2 n kv = 2 group size = n heads / n kv = 14 / 2 = 7 =n_{\text{heads}}/n_{\text{kv}}=14/2=7 = n heads / n kv = 14/2 = 7 kv_head ( h ) = ⌊ h / 7 ⌋ \text{kv\_head}(h)=\lfloor h/7\rfloor kv_head ( h ) = ⌊ h /7 ⌋
挑战 实现 GQA attention 的 CPU 版本，支持 seq_q != seq_k 和 q_pos_base。核心循环对每个 (head h h h t q t_q t q k v = h / group kv=h/\text{group} k v = h / group key_count = min(q_pos_base + tq + 1, seq_k)；对 0.. key_count 0..\text{key\_count} 0.. key_count q ⋅ k / d q\cdot k/\sqrt{d} q ⋅ k / d V V V
① seq_q 与 seq_k 分开传入，使其同时支持 prefill（s e q q = s e q k seq_q=seq_k se q q = se q k s e q q = 1 seq_q=1 se q q = 1 s e q k = seq_k= se q k = q_pos_base 表示 query 的绝对起始位置，用于 causal 上限；③ K/V 用 n_kv_heads 索引而 Q 用 n_heads 索引。可对照 xinfer_core::cpu::attention 的参考实现。