能推理:贪心生成、对拍与 norm_topk 发现

1. 贪心生成

generate.cc 的 GreedyGenerate 是最朴素的循环:每步跑一次整图前向,取最后位置 logits 的 argmax,追加,再来。没有采样、没有 chat 模板——目标是先把”能推理”这件事锚死,花哨的后面再说。

2. 对拍:全项目最重要的方法论

自己写的前向怎么知道对不对?逐 token 对拍 llama.cpp。 这是贯穿 M3 的纪律:

1. 用 llama-tokenize 把 prompt 变成 token id;
2. 用 nuthatch 跑生成,记下每步的 token id;
3. 用 llama-cli 贪心跑同样输入,逐 token 比对。

只要有一层的形状、permute、RoPE 参数错了,logits 就会偏,argmax 就会分岔。对拍把”差不多对”变成”逐 id 精确一致”——这是数值正确性唯一可信的判据。

3. norm_topk:对拍抓到的关键 bug

第一次生成,nuthatch 给出:

The capital of France is called Paris.

而 llama.cpp 是:

The capital of France is Paris.

首 token 就分岔了(nuthatch 选了 “ called”,llama.cpp 选了 “ Paris”)。顺着对拍往回查,问题出在 MoE 的路由权重:很多 MoE 实现会把 top-k 的门控权重归一到和为 1;但 OLMoE 的配置是 norm_topk_prob=false——不归一。

把 norm_topk 设成 false 后,首 token 精确一致(id 7785,即 ` Paris`)。整句输出也对上了:

prompt ids:    510 5347 273 6181 310         ← 与 llama-tokenize 一致
generated ids: 7785 15 187 187 510 14731
text:          The capital of France is Paris.

教训:对拍是最好的调试器。 一个布尔配置的差别,靠肉眼读代码很难发现,但对拍让它在第一个 token 就无所遁形。这次经历也验证了”每加一层就锚一次正确性”的价值——如果等到十几层都写完再对拍,定位会难得多。

4. 从这里长出的东西

GreedyGenerate 后来分裂出几个变体,每个都由一条 parity 单测守护:

• GreedyGenerateCached(第 8 章):带 KV cache,== 不带 cache 逐 token 一致;
• GreedyGenerateCachedTrace(第 9 章):额外导出真实路由 trace,生成结果与不导时一致(trace 是旁路);
• StreamingGenerate(第 10 章):物理流式,== 常驻路径逐 token 一致。

“每个新变体都用 == 老路径的单测锚住”——这套纪律让引擎在不断长出新能力(缓存、流式)的同时,数值正确性始终有据可查。

下一章:自包含分词器:BPE + PCRE2。