M05 · Talker MTP：一帧内预测 15 个残差 codebook
talker 每一步只采样 codebook 0；剩下 15 个残差 codebook 由一个小 transformer（code predictor）在同一步内自回归补全——每个 talker token 背后是“1 次大模型 + 15 次小模型”。
mine/m05_talker_mtp.py
验收 ./check.sh m05
提示 A-5 · 参考解 C-5 solutions/m05_talker_mtp.py
题面速览
- 在系统中的位置：talker（stage 1）输出 RVQ 音频 codec，每帧 G=16 个 codebook。talker LLM 采样 CB0 后，code predictor 以
[talker 最后一层 hidden, CB0 的 embedding]为起点逐个预测 CB1..CB15；下一步 talker 的输入 embedding 是这一帧 16 个 codec embedding 之和 + 本步对应的 thinker 文本向量。
seq = [h_talker, E0(cb0)] # 长度 2
for j in 0..G-2: # 预测 codebook j+1
hj = CP(seq)[:, -1] # 小 transformer，整段重算
code = argmax(lm_head[j](hj))
seq.append(codec_embedding[j](code)) # 注意 embedding 表也按 j 取
next_input = E0(cb0) + Σ_j codec_embedding[j](code_j) + text_step
- 为什么 vllm-omni 选 re-prefill（不用 KV cache）：序列最长 17，O(T²) 重算几乎免费；没有 KV 状态 → 形状固定、可整体包成一张 CUDA graph（P17），也不用管 per-request cache。batch 很大时又提供了可选的增量 KV 版本——两者必须给出完全相同的 code。
- 要实现：
predict_residual_codes（re-prefill，greedy，对拍 HFcode_predictor.generate(do_sample=False)）；predict_residual_codes_kv（transformers.DynamicCache增量版）；next_talker_input（对拍 HF talker 的prepare_inputs_for_generation）。 - 只能用：
cp.config.num_code_groups、cp.model(inputs_embeds=..., [past_key_values=, use_cache=]).last_hidden_state、cp.lm_head[j]、cp.model.codec_embedding[j]（j = 0..G-2）。真实服务用 top-k 50 / top-p 0.8 采样，本题为可对拍用 greedy。
验收要点：re-prefill 与 HF generate（KV cache）完全一致 · KV 版与 re-prefill 一致 · batch 各行互不影响 · 下一步 talker 输入与 HF 的求和方式一致。
直觉与类比
CB0 像是先报“大概 1 米 7”。CB1 不是重量一遍，而是回答“还差多少？”——“再加 3 厘米”；CB2 再修正“再减 2 毫米”。每一级只修正上一级留下的误差，所以必须知道前面几级报了什么。talker 大模型 = 用卷尺量出“1 米 7”的人；code predictor = 拿游标卡尺的助手。每一级用的是不同的刻度尺：第 j 步有自己的 lm_head[j]（读数）和 codec_embedding[j]（把读数记到纸上）。
re-prefill：每想下一句都把已写好的整首诗从头念一遍——诗只有十几行时几乎不花时间，而且每次动作一模一样（形状固定），很好“彩排”（CUDA graph）。KV cache：把读过的每一行要点记在笔记本上，下次只读新的一行。两种做法读出的下一句必须完全相同——因果注意力里前面每一行的“要点”不会因为后面加了新行而改变。
框架视角 · shape trace
G=5，B=2，H=64，vocab=12：
| 步 | 张量 | 形状 | 发生了什么 |
|---|---|---|---|
| init | seq | [2, 2, 64] |
两个 [B, H] 叠出长度维：位置 0 = talker 隐状态，位置 1 = CB0 embedding |
| j=0 | cp.model(inputs_embeds=seq).last_hidden_state |
[2, 2, 64] |
取最后位置 → [2, 64] |
lm_head[?](h) → argmax |
[2, 12] → [2] |
得 CB1 = [8, 7] | |
codec_embedding[?](code) |
[2, 64] → [2, 1, 64] |
补出长度维，拼到 seq 末尾 | |
| j=1 | seq | [2, 3, 64] |
CB2 = [0, 1] |
| j=2 | seq | [2, 4, 64] |
CB3 = [0, 4] |
| j=3 | seq | [2, 5, 64] |
CB4 = [6, 5]，最后一个码 |
| out | 4 个 [2] |
[2, 4] |
按列拼：[[8,0,0,6],[7,1,4,5]] |
| next | next_talker_input | [2, 64] |
G 个 codec embedding 之和 + text_step |
表里的 ? 填什么
都是 j：预测 CB(j+1) 用 lm_head[j]，嵌入它用 codec_embedding[j]。两张表都只覆盖残差 codebook（长度 G-1），下标 0 都指 CB1。
易错点
- 所有步都用
lm_head[0]：第一列对、后面开始错。G=4 例子会从[[4,5,6],[7,4,9]]变成[[4,5,4],[7,4,2]]——只有第一列稳定正确，这是典型症状。 - KV 版每步仍喂整段序列：cache 重复存旧位置、位置编号被推远，结果悄悄偏离（同例得
[[4,5,6],[7,4,2]]，有的行对有的行错）。每步之后 cache 长度应为 2, 3, 4, … - 循环写成
range(G)（越界访问lm_head[G-1]）或range(G-2)（少一个 code）。 - 起始序列用
cat而不是stack(dim=1)：得到[2B, H]或[B, 2H]。 - embedding 输出
[B, H]追加前忘了补长度维。 next_talker_input原地+=改坏调用方的cb0_embed；或把整个[B, G-1]丢进同一张表。- 忘了
torch.no_grad()：结果不变，但更慢更占内存。
逐步图：re-prefill vs KV（G=4, B=2, H=16, vocab=10）
seq0 = stack([hid, cb0], dim=1) # [2, 2, 16]
re-prefill 与 KV 版第一次都喂这 2 个位置。
re-prefill: cp.model(seq0) [2, 2, 16] -> [:, -1] [2, 16]
KV: cp.model(seq0, cache) cache 长度 0 -> 2
lm_head[0](.) [2, 10] -> argmax = [4, 7]
codec_embedding[0](code) [2, 16] -> [:, None] [2, 1, 16]
re-prefill: seq1 = cat([seq0, e0], 1) [2, 3, 16] 整段再喂
KV: 只喂 e0 [2, 1, 16] cache 2 -> 3
lm_head[1] -> CB2 = [5, 4]
re-prefill: [2, 4, 16] KV: 只喂 e1, cache 3 -> 4
lm_head[2] -> CB3 = [6, 9]
最后一个 code 算出后不必再跑 transformer；但 next_talker_input 需要它的 embedding（codec_embedding[G-2]）。
stack([c0, c1, c2], dim=1) = [[4,5,6],[7,4,9]] # [B, G-1]
next = cb0_embed
+ codec_embedding[0](codes[:, 0]) # 查 4 / 7
+ codec_embedding[1](codes[:, 1]) # 查 5 / 4
+ codec_embedding[2](codes[:, 2]) # 查 6 / 9
+ text_step # [2, 16]
一共 G = 4 个 codec embedding 相加（1 个 CB0 + G-1 个残差）再加文本向量。
三级提示
L1I/O 示例与关键问题I/O 示例 + 关键问题 + 细节（轻度剧透）
例 1：G=4，B=2，H=16，vocab=10（随机 tiny code predictor，1 层）：
predict_residual_codes(cp, hid, cb0) -> tensor([[4, 5, 6], [7, 4, 9]]) # [B, G-1]
predict_residual_codes_kv(cp, hid, cb0) -> tensor([[4, 5, 6], [7, 4, 9]]) # 必须逐元素相同
第 0 列是 CB1（lm_head[0]），第 1 列 CB2，第 2 列 CB3。只取第 1 行单独跑得到 [[7, 4, 9]]——batch 各行互不影响。KV 版 3 步后 cache 长度依次 2 → 3 → 4。
关键问题
- 预测第 j+1 个 codebook 用 lm_head 的第几个、codec_embedding 的第几个？为什么下标相同而不是差 1？
- 起始序列两个
[B, H]变成[B, 2, H]用 stack 还是 cat？沿哪维？ - transformer 输出
[B, L, H]，取哪个位置算 logits？为什么不是位置 0？ - 最后一个 code 算出来后，re-prefill 版还要再跑一次 transformer 吗？next_talker_input 需要它的 embedding 吗？
- KV 版第一次喂几个位置？第二次几个？第二次仍喂完整序列会怎样？
- KV 版新位置的 position id 由谁递增？需要手动传 position_ids 吗？
- next_talker_input 一共几个 embedding 相加？CB0 那一项为什么直接用
cb0_embed？
细节
- G = num_code_groups 是每帧 codebook 总数（含 CB0）；lm_head 与 codec_embedding 都是长度 G-1 的 ModuleList，下标 0 对应 CB1。
- re-prefill：每步把完整序列交给
cp.model(inputs_embeds=...)，不传 past_key_values。 - greedy：
logits.argmax(-1)得[B]long，可直接索引 nn.Embedding。 codec_embedding[j]是 code predictor 自己的表，与 talker 主模型里 CB0 的表不是同一张。- DynamicCache：创建空对象，作为
past_key_values=传入并use_cache=True，模型就地追加每层 K/V；它知道已缓存多少位置，RoPE 位置自动接上。
L2思路与常见坑思路拆分与常见坑（中度剧透）
- 三个函数共享“单步”模式：最后位置 hidden →
lm_head[j]→ argmax →codec_embedding[j]嵌入。两个 predict 的区别只在于每步喂什么。 - 序列用
[B, L, H]维护，每步torch.cat(..., dim=1)追加[B, 1, H]；code 收集到 list，最后torch.stack(codes, dim=1)。 - 循环
for j in range(G - 1)。 - KV 版用变量 x 表示“这一步要喂的新位置”：第一次是
[B, 2, H]，之后是刚得到的 code 的 embedding[B, 1, H]；cache 在循环外创建一次。 - 用
@torch.no_grad()包起来。 - next_talker_input：第 j 列用
codec_embedding[j]查；用a = a + b不要a += b。
L3详解：答案详解：答案、逐步追踪、骨架提示（重度剧透）
关键问题答案
- 预测 CB(j+1) 用
lm_head[j]；嵌入用codec_embedding[j]。两个表都只覆盖残差 codebook，下标 0 都指 CB1。 torch.stack([last_talker_hidden, cb0_embed], dim=1)→[B, 2, H]。last_hidden_state[:, -1]：因果 LM 用最后一个位置预测下一个。- 不必再跑 transformer；但 next_talker_input 需要它（
codec_embedding[G-2]）。 - 第一次 2 个位置，之后每次 1 个。重复喂整段会让 cache 出现重复 K/V，位置也错。
- 不需要：DynamicCache 记录已缓存长度，HF 据此推出新位置的 cache_position。
- G 个 codec embedding（CB0 + G-1 残差）加 text_step。CB0 的 embedding 来自 talker 主模型的表，code predictor 表里没有 CB0。
residual_codes[:, j] 是 [B]，查表得 [B, H]；不要把整个 [B, G-1] 丢进同一张表（会得到 [B, G-1, H] 且所有列用了同一张表）。
为什么两版一定一致：因果注意力下位置 p 的输出只依赖位置 0..p 的输入。re-prefill 每步重算旧位置，但旧位置的 K/V 与上一步完全相同；KV 版只是把它们存下来。只要 RoPE 位置一致（DynamicCache 自动保证），最后位置的 hidden 几乎相同，argmax 相同。
自检：循环里打印 cache.get_seq_length()，G=4 时应依次为 2、3、4；两版输出 torch.equal。
面试官追问 & 对照源码
- 为什么 seed 化的逐行采样会阻止整张 CUDA graph replay？vllm-omni 怎么处理？
- MRv2 的 “eager frames” 为什么能在采样出 CB0 的同一步就完成整帧？
对照源码：model_executor/models/qwen3_omni/qwen3_omni.py:1044-1069 · model_executor/models/common/qwen3_code_predictor.py:706-722 · :1357-1440 · worker/gpu_model_runner.py:243-272 · HF Qwen3OmniMoeTalkerCodePredictorModelForConditionalGeneration · HF Qwen3OmniMoeTalkerForConditionalGeneration.prepare_inputs_for_generation