OLMoE 前向:注意力、MoE 与 ggml 建图
要拿真实推理轨迹,得先有能跑的引擎。M3 用 ggml 建图,逐块搭出 OLMoE-1B-7B(16 层、n_embd 2048、16 头、64 专家选 8、无共享专家)。
1. 为什么用 ggml,以及怎么用
护城河在算法不在引擎。 矩阵乘、量化点积这些 ggml 已经做到极致,重写没有意义。所以 nuthatch 复用 ggml 的算子,自己只负责把 OLMoE 的计算图搭对。
ggml 的模式是”先建图、后计算”:用 ggml_mul_mat / ggml_rms_norm / ggml_rope_ext 等把张量运算连成一张 ggml_cgraph,再 ggml_graph_compute_with_ctx 一次算完。权重张量直接来自 GGUF(量化的 mul_mat/get_rows ggml 原生支持)。
2. 模型加载
olmoe_model.cc 用 gguf_init_from_file(no_alloc=false,把权重也载进 ctx 常驻)读出所有张量,建 name → tensor 映射,并把超参解析成 OlmoeConfig。超参解析抽成了公共函数 ParseOlmoeConfig——常驻加载(本章)和显存受限加载(第 10 章)共用它。
3. 注意力块
attention.cc 的 BuildAttention 镜像 llama.cpp 的 OLMoE 路径:
RMSNorm(×attn_norm) ─► Q/K/V 投影
─► QK-norm(OLMoE 对整段 n_embd 做 RMSNorm,不是 per-head)
─► 拆头 [head_dim, n_head, T]
─► NEOX RoPE(全 head_dim 旋转,θ=10000)
─► permute 到 [head_dim, T, n_head]
─► k^T·q 缩放 ─► 因果 mask(diag_mask_inf) ─► softmax
─► ·v ─► 合并回 [n_embd, T] ─► 输出投影
两个 OLMoE 特有点:QK-norm(对整段 n_embd 归一,不是逐头)、无 GQA(n_head_kv == n_head)。这些细节错一个,logits 就偏——靠第 6 章的对拍抓。
4. MoE FFN 块
moe.cc 的 BuildMoe 是护城河后来接入的地方,先看常驻版:
ffn_norm ─► router(mul_mat) ─► softmax ─► argsort_top_k(选 8) ← selected [n_used, T]
─► get_rows 取选中专家的权重
─►(可选)norm_topk 归一 —— OLMoE 用 false(第 6 章的关键发现)
─► up = mul_mat_id(up_exps, x, selected) ← 只对选中的专家算
gate = mul_mat_id(gate_exps, x, selected)
ff = SiLU(gate) · up
experts = mul_mat_id(down_exps, ff, selected)
─► × 各专家权重 ─► 对 n_used 维求和
关键算子是 ggml_mul_mat_id:它接受融合的 3D 专家张量 + 一个 selected 索引张量,只对每个 token 选中的专家做矩阵乘。这正是 MoE 的稀疏性所在,也是后面流式的关键——第 10 章会把这里的”融合张量”换成”有界槽张量”,把 selected 重映射成槽 id。
selected(每 token 每层选中的全局专家 id)后来加了个可选出参暴露出来——这就是真实推理路由 trace 的来源(第 9 章)。
5. 整图前向
forward.cc 的 BuildForward 把上面两块串成完整前向:
h = get_rows(token_embd, token_ids); // [n_embd, T]
for (l in 16 layers) {
h = h + BuildAttention(cfg, aw, h, pos); // 注意力 + 残差
h = h + BuildMoe(cfg, mw, h, norm_topk); // MoE + 残差
}
h = rms_norm(h) × output_norm;
return mul_mat(output.weight, h); // logits [n_vocab, T]
结构简单,但每个算子的形状、permute 顺序、RoPE 参数都必须和 llama.cpp 一致。怎么保证?下一章的对拍。