nano-vllm-omni · 提示 ← M07P02 →
#8 · Track M · Qwen-Omni 模型实现 · 书 p.44

M08 · 串起来：tiny thinker 的 KV-cache greedy 解码

把 M01/M02/M03/M07 拼成一个只做文本的 thinker：prefill 一次、之后每步只算 1 个 token，生成的序列要和 HF 每步全量重算的 greedy 完全一致。

⏱ 60 min 难度 中 实现 mine/m08_tiny_thinker_e2e.py 验收 ./check.sh m08 提示 A-8 · 参考解 C-8 solutions/m08_tiny_thinker_e2e.py

题面速览

h = embed_tokens[ids]                                          # [T, hidden]
for layer:  h = h + Attn(rms(h, ln1), cos, sin, cache)          # M02（cache 每层一个）
            h = h + MoE(rms(h, ln2))                            # M03
logits = rms(h, norm) @ lm_head^T                               # [T, vocab]

验收要点：prefill logits 与 HF 一致 · KV-cache greedy 12 个 token 与 HF 全量重算 greedy 完全一致（两组随机权重）。

直觉与类比

新盖的房子要编门牌号

KV cache 是一条街：已经盖好的房子（缓存的 token）门牌 0, 1, …, L-1，每栋的 K 在入住时就按自己的门牌“旋转”好了。新房子的门牌必须是 L——街上已有房子的数目。如果编成 L+1，它和每个邻居的距离都多算 1：不报错、看起来差不多，但邻里关系悄悄变了——往往几步之后才“走错路”。

接龙游戏

prefill = 先把题目整段读完，顺手说出第一个接龙词。之后每轮只听最新的一个词（只喂 out[-1:]），前面的内容都记在笔记本（cache）上。说 max_new 个词只需要 max_new 轮：第一轮的词来自读题，最后一个词说出口之后不必再听一轮。每轮开始时确认：笔记本上记了几个词（len(cache)），这一轮的新词就排第几（start）。

框架视角 · shape trace

阶段 张量 形状（prefill / decode） 发生了什么
位置 positions [3, 5] / [3, 1] 三行都是 start … start+T-1；此时 interleaved MRoPE 与 1D RoPE 相同（实测 allclose），但仍走 M01 的接口
cos, sin [5, 16] / [1, 16] 只覆盖新 token；所有层共享
embed h = embed[ids] [5, 64] / [1, 64] 查表
每层 h + attn(rms(h, ln1), cos, sin, cache) [T, 64] pre-norm 残差；cache 就地追加 T 个：0→5 / 5→6
h + moe(rms(h, ln2)) [T, 64] top-2 of 6 experts
输出 rms(h, norm) @ lm_head.T [5, 97] / [1, 97] lm_head [97, 64] 来自 state["thinker.lm_head.weight"]，不在 thinker.model. 下
采样 int(logits[-1].argmax()) 标量 只用最后一行

易错点

逐步图：greedy(prompt=[10, 20, 30, 40], max_new=4)

▶ 每次 forward 的参数
1. 第 1 次 forward（prefill）

ids = [10, 20, 30, 40]，start = 0，positions = [[0,1,2,3]]×3，返回 logits [4, V]。

forward 后每层 cache 长度 = 4；out[0] = argmax(logits[-1])。

2. 第 2 次 forward

ids = [out[0]]，start = 4，positions = [[4]]×3，logits [1, V]；cache 4 → 5；得 out[1]。

规律：start = len(prompt) + len(out) − 1，并且每次 forward 前 start == len(cache)。

3. 第 3、4 次 forward

start = 5 → cache 6 → out[2]；start = 6 → cache 7 → out[3]。

max_new=4 只需要 4 次 forward：prefill 本身就产生了第一个 token，最后一个 token 不需要再 forward。

4. 真实数字（vocab=97, hidden=64, 2 层, H=4, KVH=2, D=16, 6 expert top-2）
forward(tensor(prompt), 0, caches) -> logits [4, 97]，argmax = 34，caches 长度 [4, 4]
forward(tensor([34]), 4, caches)   -> logits [1, 97]，argmax = 75，caches 长度 [5, 5]
greedy(prompt, 6)                  -> [34, 75, 75, 75, 17, 1]
HF 每步全量重算的 greedy            -> [34, 75, 75, 75, 17, 1]   ✓
5. 两种 bug 的样子
decode 时 start 多加 1           -> [34, 75, 75, 75, 75, 75]   # 第 5 个 token 开始不同
每步把完整序列从 start=0 喂进同一组 caches -> [34, 77, 2]，3 步后 cache 长度已是 15

位置错误往往不会立刻改变 argmax，而是在若干步后才“分叉”——所以要比较整段序列。

三级提示

L1I/O 示例与关键问题I/O 示例 + 关键问题 + 细节（轻度剧透）

见逐步图的表格与真实数字。另一个例子：纯文本 MRoPE 等于普通 1D RoPE：

positions = arange(4, 6).expand(3, 2) -> [[4, 5], [4, 5], [4, 5]]
cos, sin = mrope_cos_sin(rope_inv_freq(16, 10000.0), positions, [4, 2, 2], interleaved=True)
cos.shape -> [2, 16]，与“只用 t 轴位置 [4, 5] 的普通 RoPE”的 cos 完全相同

关键问题

  1. PackedLayer 只有打包后的 qkv 和 w13，M02 要分开的 q/k/v、M03 要分开的 w_gate/w_up，分别按什么尺寸切？
  2. forward(ids, start, caches) 中 positions 形状是什么？三条轴填什么？
  3. cos/sin 每次 forward 都要重算吗？对应“新 token”的位置还是“cache 中全部 token”的位置？
  4. M02 自己根据 len(cache) 决定因果 mask；start 和 len(cache) 应满足什么关系？不满足会怎样？
  5. lm_head 来自哪个 key？为什么不经过 load_thinker_text？logits 是 h @ lm_head 还是 h @ lm_head.T？
  6. greedy 一共调用多少次 forward？第 k 次 decode 喂哪个 token、start 是多少？
  7. caches 应该在 __init__、greedy 开头，还是每次 forward 里创建？

细节

  • Dims 给出切分与 RoPE 所需全部超参；eps 同时用于所有 RMSNorm（包括 q_norm/k_norm）。
  • pre-norm：归一化只作用在子层输入上，残差加的是未归一化的 h；最后一层之后再 rms(h, norm) 才乘 lm_head。
  • One KVCache per layer：caches 长度 num_layers，第 l 层只读写 caches[l]；M02 就地追加 K/V。
  • “Must equal HF full-recompute greedy”：数值有极小浮点差异，但 argmax 序列必须完全一致。
  • forward 返回 [T, V] 全部行（测试会比较 prefill 的全部行）；greedy 只用 logits[-1]。
L2思路与常见坑思路拆分与常见坑（中度剧透）
  • __init__：(1) load_thinker_text(state, dims.num_layers)；(2) 取 state["thinker.lm_head.weight"]；(3) 每个 PackedLayer 转成 AttnWeights + MoeWeights（连同 ln1/ln2）；(4) 预先算 inv_freq = rope_inv_freq(head_dim, rope_theta)。
  • 权重还原：q, k, v = L.qkv.split([H*D, KVH*D, KVH*D], dim=0)；I = L.w13.shape[1] // 2，w_gate = L.w13[:, :I]、w_up = L.w13[:, I:]、w_down = L.w2。
  • forward：positions → cos/sin（一次，所有层共享）→ h = embed[ids] → 逐层残差 → 最终 norm → @ lm_head.T。
  • 与 M02 的契约：M02 把新 token 放在 len(cache) .. len(cache)+T-1 做 mask，RoPE 位置由你的 cos/sin 决定——二者必须一致：start == len(caches[0])，可加 assert。
  • greedy：caches 在开头新建；prefill 得第一个 token；while len(out) < max_new 只喂 out[-1:]，start = len(prompt) + len(out) - 1。
  • 调试顺序：先看 prefill logits 是否 allclose（检查权重、norm、MoE、RoPE）；再看 greedy 序列（检查 cache 与 start）。前者错多半是前置题的问题。
  • @torch.no_grad()；torch.tensor(prompt) 是 long；token 用 int(...) 放进 list。
L3详解：答案详解：答案、逐步追踪、骨架提示（重度剧透）
关键问题答案
  1. qkv 按行 split 为 [H·D, KVH·D, KVH·D]；w13 在 dim 1 对半：前 I 行 gate、后 I 行 up。
  2. [3, T] long，三行都是 arange(start, start+T)。
  3. 每次 forward 重算，只对应新 token；cache 中旧 token 的 K 写入前已用自己的位置旋转过。
  4. 必须 start == len(cache)；否则 RoPE 相对距离和 mask 认为的位置不一致，输出悄悄偏离，常在若干步后分叉。
  5. state["thinker.lm_head.weight"]，形状 [V, hidden]；不在 thinker.model. 下；logits = h @ lm_head.T。
  6. 共 max_new 次：1 次 prefill + (max_new − 1) 次 decode。第 k 次 decode（k 从 1 开始）喂 out[k-1]，start = len(prompt) + k − 1。
  7. greedy 开头、每层一个；放 __init__ 会继承上一次历史，放 forward 里等于没有 cache。

形状追踪（hidden=64, H=4, KVH=2, D=16, V=97, prompt 长 4）：

prefill: ids [4] -> positions [3, 4] -> cos/sin [4, 16]
         h = embed[ids]                      [4, 64]
         rms(h, ln1) -> attention_forward    [4, 64]    cache: 0 -> 4
         rms(h, ln2) -> moe_forward          [4, 64]
         (第 2 层同理)
         rms(h, norm) @ lm_head.T            [4, 97]
decode:  ids [1], start=4 -> positions [3, 1] -> cos/sin [1, 16]
         ...                                 [1, 97]    cache: 4 -> 5

自检：循环里 assert start == len(caches[0])；再写一个“每步全量重算”的慢版本（每次新建 caches，从 start=0 喂 prompt+out），两者 token 序列必须完全相同。

面试官追问 & 对照源码

对照源码：model_executor/models/qwen3_omni/qwen3_omni.py:235-267 · HF Qwen3OmniMoeThinkerTextModel