推理主流程

本文档从入口 main 出发，跟踪一次前向的骨架（step / layers_forward / layer_forward），以及采样与投机解码的主循环。注意力和 MoE 的细节分别在 06 和 07，MTP 投机在 08。

1. 前向的三层骨架 

step / step_all               ← 一次前向（embed → layers → norm → lm_head）
  └─ layers_forward           ← 循环所有层
       └─ layer_forward       ← 单层（attention + moe/dense_mlp）
            ├─ attention       ← MLA（见 06）
            └─ moe / dense_mlp ← MoE 路由 / 稠密（见 07）

1.1 layer_forward / layers_forward

layer_forward（1349 行）算一层：先做 attention，再根据 l->sparse 分派 moe 或 dense_mlp（1360 行），中间可能触发投机预取（PILOT/SPEC，见 08）。

layers_forward（1363 行）分配 [S,D] 的 scratch，循环所有 n_layers 层调用 layer_forward——当 S>=8（prefill）时打印 [prefill] layer .../... 进度（1369 行）。这是 prefill / decode / score / teacher-forcing 共用的核心。

1.2 step / step_all

1.3 KV 分配

kv_alloc（1376 行）分配 KV 缓冲——注意它为 NR = n_layers+1 个层分配（1386 行），多出的那一行给 MTP 层。kv_bind（1393 行） 把 model 的各个 KV 别名指向某个 KVState（用于多 slot 切换，见 10）。

2. 采样：温度 + nucleus

colibrì 实现了真正的温度 + 核采样（nucleus / top-p），默认值针对 int4 现实调紧——官方的 1.0/0.95 会把量化噪声从尾部采出来。

2.1 默认值

参数 全局 默认 说明
温度 g_temp（545 行） -1（auto） run_text 里 auto → 0.7（1745 行）
nucleus g_nuc（547 行） 0.95 → main 里覆盖为 0.90（2406 行） 词表 top-p
topk g_topk（548 行） 0 专家 top-k 覆盖（见 07）

2.2 采样内核

2.3 停止词

g_stop[9]/g_nstop（1559 行）、is_stop（1560 行）、stops_arm（1561 行）——把 config 的 stop_ids + EOS 装进活动停止集。

3. spec_decode：生成主循环

spec_decode（1576 行）是所有真实生成的主循环，即使 g_draft=0 也走它（此时退化成逐 token）。它把”起草 + 验证”合到一次批量前向里：

spec_decode 循环:
  1. 起草：has_mtp ? mtp_draft : ngram_draft  →  得到 G 个候选 draft
  2. 验证：step_all 一次批量前向，算出 [真 token + G 个 draft] 的所有 logits
  3. 接受循环：逐个比对 draft 与主模型输出
       - 贪心（g_temp<=0）：argmax 精确匹配即接受
       - 采样（g_temp>0）：拒绝采样，接受概率 = p(draft[k])（见 08）
  4. mtp_acc += 接受数；mtp_absorb 把接受的 token 吸收进 MTP 层 KV
  5. pick_tok 对第一个不匹配位置采样，作为新的真 token

一个自适应保护：如果 mtp_prop>=24 && mtp_acc*10 < mtp_prop（接受率 <10%），就把 g_draft=0 关掉投机（1589 行）——冷缓存下投机反而更慢，这个 guard 会自动放弃。详见 08。

4. 运行模式与入口

main（2384 行）是 env-var 驱动的：位置参数只有 argv[1]=缓存 cap、argv[2]=专家 bits、argv[3]=稠密 bits，其余全靠环境变量。必需的是 SNAP=<dir>（2387 行）。模式按顺序互斥分派：

环境变量 入口 行 用途
SCORE run_score 2480 对每个候选续写算 log-likelihood（quality benchmark）
SERVE run_serve 2483 持久服务模式（见 10）
PROMPT run_text 2486 真实文本生成
REPLAY run_replay 2503 回放 oracle 序列，测 decode tok/s
TF forward_all 2509 teacher-forcing 验证
（默认） generate 2494 对 ref_glm.json oracle 验证

4.1 run_text —— 真实文本生成

run_text（1740 行）是 PROMPT 模式的主体：加载 tokenizer → 解析 EOS + stops_arm → auto 温度 0.7（1745 行） → 编码 prompt → kv_alloc → step 做 prefill（1755 行） → spec_decode 配 emit_stream 流式输出（1757 行） → 打印吞吐/专家命中率/投机与 MTP 接受率统计 → profile_print（1773 行） → usage_save（1781 行）。

4.2 输出回调与剖析

  • emit_store（1629 行）：把 token 累积到数组（验证模式用）。
  • emit_stream（1632 行）：detokenize 并流式输出到 stdout，每 16 个 token 打一行心跳（RSS、命中率、tok/s、tok/fw）。
  • profile_print（1707 行）：把耗时分解成 专家磁盘读 / 专家矩阵乘 / 注意力 / lm_head / 其他——这正是 README 里”66% disk vs 57% matmul”那种 profile 的来源。

5. prefill vs decode 的两种成本结构

理解 colibrì 的性能，关键是区分两个阶段：

prefill（处理 prompt，S 个 token 一次批量）:
  - 每个专家只需读一次磁盘，然后应用到所有路由到它的位置（batch-union）
  - 计算受矩阵乘限制（S 大）
  - 一次性成本

decode（逐 token 生成，S=1 或 S=1+draft）:
  - 每层选 8 个专家，几乎每个都要从磁盘读（冷缓存下 ~11 GB/token）
  - 计算受磁盘 I/O 限制
  - 每个输出 token 都付这个成本
  - MTP 投机在缓存暖后能把有效成本大致减半

这个区分解释了为什么 colibrì 把这么多工程投在专家缓存（07）和投机解码（08）上——它们都是在攻击 decode 阶段的磁盘瓶颈。

6. 小结

组件 出处 角色
单层前向 layer_forward:1349 attention + moe/dense 分派
全层前向 layers_forward:1363 prefill/decode/score 共用核心
单/全 logit step:1399, step_all:1414 decode 用 step，投机验证用 step_all
采样 dist_build:1525, pick_tok:1551 温度 + nucleus
生成主循环 spec_decode:1576 起草 + 验证，自适应关闭
模式分派 main:2384 env-var 驱动
This site uses Just the Docs, a documentation theme for Jekyll.