物理流式:让”流式”名副其实
命中率结论已经拿到,但引擎其实还全常驻——专家仍占着 ~4GB 内存。研究不需要物理执行,但要真省内存,得走完最后一步。这是全项目最难的一块。
1. 选择性加载:非专家常驻,专家留盘
专家 3D 张量约占 OLMoE 九成参数。真省内存必须不把它们载入。streaming_model.cc 的 StreamingModel::Load:
gguf_init用no_alloc=true只拿元数据 + 形状;- 非专家张量(注意力/norm/router/embd/lm_head)在常驻 ctx 建同形张量、从文件定位读其数据;
- 专家张量(名含
_exps.weight)跳过,交给ExpertReader按需流式。
常驻内存从 ~4GB 降到 ~250MB。判断”是不是专家张量”就一句 strstr(name, "_exps.weight")——router 是 ffn_gate_inp 不含,正好排除。
2. 两段式前向:难在哪
难点在于——专家选择依赖运行时激活。你必须先算完这一层的路由、知道选了哪 8 个专家,才能去磁盘装载、再计算。所以一层没法一趟图算完,得拆成两段(streaming_forward.cc):
逐 token(T=1):
段A 图: 注意力(带 KV cache)+ 残差 ─► ffn_norm ─► 路由 ─► 选中的全局专家 id
└─ 读回 host:h_mid、weights、selected
host: caches[l].Ensure(selected) ← miss 真从盘 pread 装进槽
全局 id ─► slot id(重映射)
段B 图: 在槽张量上 mul_mat_id(cur, slot_id)算专家 FFN ─► + 残差 ─► h_next
h 在两段之间走 host 中转——因为逐 token 处理,活值只有 [n_embd, 1],极小,这个中转很便宜。段 B 里 mul_mat_id 的专家张量是 expert_slot_cache 的有界槽张量 [n_embd, n_ff, C](C = 每层容量),selected 被重映射成 0..C 的槽号。约束:C ≥ 每步选中的不同专家数(逐 token 下 = n_expert_used = 8)。
3. 有界槽缓存
ExpertSlotCache::Ensure(ids) 是物理流式的心脏:命中即用;miss 则占槽(空槽优先,否则淘汰 tick 最小的)+ 从 ExpertReader 读入该专家的 gate/up/down 三类权重;返回 专家id → 槽号 供重映射。这里 P20 用 LRU;它同时也是第 9 章那些策略在物理世界里的一个实现。
4. 正确性锚点:streaming == resident
流式路径的输出必须与常驻逐 token 一致——因为槽里装的是和常驻专家完全相同的量化字节,同样的 mul_mat_id,结果理应一致。单测 StreamingMatchesResident 断言 StreamingGenerate == GreedyGenerateCached(norm_topk 两种、含淘汰路径)。
真机结果:
| 路径 | 输出 | 峰值内存 |
|---|---|---|
| 流式(每层缓存 16 专家) | “…is Paris.” | 1.2 GB |
| 常驻 | 同上,逐 token 一致 | 3.0 GB |
同一模型、逐 token 完全一致的输出,内存降到 ~40%。 1.2GB ≈ 250MB 非专家 + 16 层×16 槽×~3MB 专家。至此,nuthatch 是真正的流式 MoE 引擎,不是模拟器。
5. 提速 2×
流式每 token 有 ~34 趟小图,起初慢。瓶颈是两件蠢事:每趟 ggml_init 都新 malloc 一块 64MB compute buffer、还以 4 线程 compute——对 T=1 的极小图纯属浪费。改成复用一块 thread_local compute buffer + 极小图单线程(只有末端 lm_head 大乘保留多线程),省掉每 token 578 次线程池 spin,提速 2×(2.73s → 1.36s,12 token 含加载)。输出不变、parity 仍成立。
异步预取(把 miss 的磁盘读与 compute 重叠)还没做,是后续。
下一章:工程判断与适用边界。