#3 · Track M · Qwen-Omni 模型实现 · 书 p.27

M03 · Sparse MoE：路由、按 expert 分组计算、shared expert

Qwen3-Omni-30B-A3B 的 “A3B” 就是 MoE：每个 token 只激活少数 expert。thinker 是纯路由 MoE，talker 还多一个带 sigmoid 门控的 shared expert。

⏱ 60 min 难度中实现 mine/m03_sparse_moe.py 验收 ./check.sh m03 提示 A-3 · 参考解 C-3 solutions/m03_sparse_moe.py 标记完成

题面速览

在系统中的位置：thinker 和 talker 的 FFN 都是 MoE。serving 时同一个 expert 的所有 token 应该一次矩阵乘完成（vLLM 的 FusedMoE），而不是逐 token 循环。
路由（必须和 HF 一致）：

logits = x @ router^T                        # [N, E]
probs  = softmax(logits, dim=-1, fp32)       # 在全部 E 个 expert 上 softmax
w, ids = topk(probs, k)                      # 先 softmax 再 top-k
if norm_topk_prob: w = w / w.sum(-1)         # Qwen3 默认重新归一化
out[n] = Σ_j w[n,j] · Expert_{ids[n,j]}(x[n]),   Expert(h) = (silu(h Wg^T) ⊙ (h Wu^T)) Wd^T

“serving 的写法”：① 把 N×k 个 (token, slot) 对按 expert id 排序（argsort(stable=True)）；② bincount 得到每个 expert 的 token 数，跳过 0；③ 每个 expert 对自己那批 token 做一次 SwiGLU；④ 乘权重后 index_add_ scatter 回原 token。这就是 FusedMoE 里 “moe_align_block_size + grouped GEMM” 的 CPU 版。
talker 的 shared expert：out += sigmoid(x @ g^T) · SharedMLP(x)，每个 token 都过，门控是标量。
边界：N=1（decode）时只有 k 个 expert 有活；单独算一个 token 与在 batch 里算结果一致（批不变）。

验收要点：路由与参考完全一致 · thinker MoE 与 HF 一致 · talker MoE（含 shared）与 HF 一致 · norm_topk_prob=False 变体 · 单 token 与 batch 一致。

直觉与类比

医院分诊台

router 是分诊台：给每位病人（token）在所有科室（expert）上打一个可能性，先在全部科室上归一（softmax），再挑分最高的 k 个挂号（top-k）。norm_topk_prob = 只在挂了号的这几个科室之间重新分配 100% 话语权。 serving 的写法：每个科室把今天分到它的病人集中起来一起看（一次矩阵乘），没人挂号的科室休息。最后每位病人拿到 k 份诊断书，按话语权加权累加进病历——第二位医生的诊断书不能覆盖第一位的（index_add_，不是赋值）。

shared expert = 家庭医生

talker 的每位病人除了专科还都会见一次家庭医生。家庭医生意见占多大分量，由只看病人本身的小开关决定：sigmoid(x @ gate.T)，每个 token 一个 0~1 的标量。

框架视角 · shape trace

记号：N 个 token，H = hidden，E 个 expert，I = expert 中间维，k = top_k，n_e = 分给 expert e 的 token 数。

#	tensor	shape	发生了什么
1	logits	`[N, E]`	`x @ router.T`
2	probs	`[N, E]` fp32	在全部 E 个 expert 上 softmax
3	weights, ids	`[N, k]`	top-k；可选 renorm（fp32），weights 转回 x.dtype
4	展平	`[N·k]` ×3	每个 (token, slot) 对：expert id、权重、token id（三者一一对应，row-major）
5	分组	`order[N·k]`, `counts[E]`	按 expert id 稳定排序；每个 expert 的对数（和 = N·k）
6	expert e	`x[toks]: [n_e, H] → [n_e, I] → [n_e, H]`	gather 这批 token，做一次 SwiGLU；n_e = 0 跳过
7	scatter	`out[N, H]`	乘权重后按 token id 累加
8	shared（talker）	`[N, 1] · [N, H]`	`sigmoid(x @ gate.T)` 乘全量 SwiGLU，加到 out

易错点

先 top-k logits 再 softmax：norm_topk_prob=True 时恰好等价，False 时完全不同（权重和应 < 1）。
token id 用 arange(N).repeat(k)（0,1,2,0,1,2）而不是逐个重复（0,0,1,1,2,2），和 ids.reshape(-1) 对不上。
scatter 用赋值 out[toks] = y：同一 token 被第二个 expert 覆盖。
gate/up 投影写成 h @ w_gate[e]（少了转置），[n_e, H] @ [I, H] 对不上。
路由权重留在 fp32，结果 dtype 被提升；或对空 expert 做 matmul。
shared expert 的门控当成向量，或忘了 w.shared is None（thinker）时跳过。

逐步图：按 expert 分组的 bookkeeping

▶ N = 3, k = 2, E = 5, ids = [[2, 0], [1, 2], [2, 3]]

1. 展平成 N·k = 6 个 (token, slot) 对

flat expert id:  [2, 0, 1, 2, 2, 3]       # ids.reshape(-1)
flat token id:   [0, 0, 1, 1, 2, 2]       # arange(N).repeat_interleave(k)

两者第 p 个元素描述同一个对。用 repeat(k) 会得到 [0,1,2,0,1,2]——错位。

2. 稳定排序 + 计数

stable argsort by expert:  [1, 2, 0, 3, 4, 5]
per-expert counts:         [1, 1, 3, 1, 0]       # expert 4 收到 0 个 token

counts 的前缀和就是每段在 order 里的起点。

3. 每个 expert 一次矩阵乘

expert	pairs（order 的一段）	tokens	计算
0	[1]	[0]	1 行
1	[2]	[1]	1 行
2	[0, 3, 4]	[0, 1, 2]	一次 matmul 处理 3 行
3	[5]	[2]	1 行
4	—	—	完全跳过

4. index_add_ 累加回各 token

token0 收到 expert 2 和 expert 0 的贡献，token1 收到 expert 1 和 2，token2 收到 expert 2 和 3。

out.index_add_(0, toks, y * flat_w[sel, None])——每个 token 的两份贡献必须相加。

5. 标量版端到端核对

H = 1、I = 1、E = 3、top_k = 2、norm_topk_prob=True，router = [[1],[0],[-1]]，w_gate=[1,2,0.5]、w_up=[1,1,2]、w_down=[1,-1,3]，x = [[2.0], [-1.0]]：

token0: logits [2, 0, -2] -> ids [0, 1], w [0.8808, 0.1192]
        out = 0.8808·3.5232 + 0.1192·(-7.8561) = 2.1667
token1: logits [-1, 0, 1] -> ids [2, 1], w [0.7311, 0.2689]
        out = 0.7311·1.1326 + 0.2689·(-0.2384) = 0.7639

若 token0 得到 -0.9365（只剩 expert 1 的贡献），说明你用了赋值而不是累加。

三级提示

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

route：E = 4，H = 4，router = eye(4)（logits 就等于 x），top_k = 2。

x = [[2.0, 1.0, 0.0, -1.0],
     [0.0, 0.5, 3.0, -0.5]]
softmax over ALL 4 experts:
     [[0.6439, 0.2369, 0.0871, 0.0321],
      [0.0428, 0.0706, 0.8605, 0.0260]]
top-2 ids            = [[0, 1], [2, 1]]
norm_topk_prob=True  -> weights [[0.7311, 0.2689], [0.9241, 0.0759]]
norm_topk_prob=False -> weights [[0.6439, 0.2369], [0.8605, 0.0706]]

weights 是 x.dtype（x 为 bf16 时也是 bf16），ids 是 int64。

moe_forward 标量版（见逐步图最后一步）：每个 expert 输出

	e=0	e=1	e=2
x=2	3.5232	-7.8561	8.7727
x=-1	0.2689	-0.2384	1.1326

moe_forward(x, w) = [[2.1667], [0.7639]]。再加 shared expert（gate_proj = up_proj = down_proj = [[1]]，gate = [[0]] → sigmoid = 0.5）：SharedExpert(x) = silu(x)·x = [3.5232, 0.2689]，结果 [[3.9283], [0.8984]]。

关键问题

softmax 是在全部 E 个 expert 上做再取 top-k，还是先 top-k 再在 k 个里 softmax？norm_topk_prob=True/False 时两者是否相同？
top-k 之后 weights 的 dtype 应该是什么？在哪一步转换？
展平后共有多少个 (token, slot) 对？怎样构造“每个对属于哪个 token”的向量？
排序后怎样知道 expert e 那一段的起止位置？
同一 token 的两个 slot 落到不同 expert，怎样加回同一行？out[all_toks] += y_all 会怎样？
w_gate 形状 [E, I, H]，对 h[n_e, H] 做 gate 投影该乘 w_gate[e] 还是它的转置？
shared expert 的门控是标量还是向量？sigmoid(x @ gate.T) 形状是多少？

细节

概率先在全部 E 个上归一，再选 k 个；不 renorm 时 k 个权重和 < 1（上例 0.6439+0.2369 = 0.8808）。
路由概率与 renorm 在 fp32，最后 weights 转回 x.dtype；ids 是 int64。
SwiGLU：Wg/Wu 是 [I, H]，Wd 是 [H, I]：h @ Wg.T → [n, I]，逐元素乘后 @ Wd.T → [n, H]。
“the serving way”：不是对每个 token 循环 k 次，也不是每个 expert 在全部 N 个 token 上算再 mask，而是每个 expert 只在路由给它的那批 token 上算一次。
scatter-add 必须是“加”，且要用支持重复下标的累加 API。
shared expert：每个 token 都过一个普通 SwiGLU，乘 sigmoid(x @ gate.T)（[N, 1]）后加到路由结果；w.shared is None 跳过。

L2思路与常见坑思路拆分与常见坑（中度剧透）

route 三行：logits → fp32 softmax（dim=-1）→ torch.topk(probs, k, dim=-1)；可选 renorm（keepdim 的 sum）；返回前转回 x.dtype。
展平：flat_e = ids.reshape(-1)、flat_w = weights.reshape(-1)、flat_tok = arange(N).repeat_interleave(k)。
分组：order = argsort(flat_e, stable=True)；counts = bincount(flat_e, minlength=E)；游标 start 累加得到每段 order[start:start+n]。也可以 (flat_e == e).nonzero()——更直观但每个 expert 扫一遍全数组。
每个 expert：toks = flat_tok[sel]，h = x[toks]，SwiGLU 得 y[n_e, H]，乘 flat_w[sel][:, None]。
scatter-add：out = zeros_like(x)；out.index_add_(0, toks, y_weighted)。out[toks] += y 在 toks 有重复时只保留最后一次写入。
跳过空 expert：if n == 0: continue。
shared：torch.sigmoid(x @ shared.gate.T)（[N, 1]）沿 H 广播。
批不变性自测：moe_forward(x)[i] vs moe_forward(x[i:i+1])[0]；不一致通常是 flat_tok 构造错或 scatter 丢了累加。

L3详解：答案与形状追踪详解：答案、逐步追踪、骨架提示（重度剧透）

关键问题答案

先全量 softmax 再 top-k。norm_topk_prob=True 时与“对 top-k logits 做 softmax”等价（分母被消掉）；False 时不等价——全量 softmax 的 top-k 权重和 < 1。
x.dtype；在 renorm 之后、return 之前转换。
N·k 个；arange(N).repeat_interleave(k) → 0,0,1,1,2,2（k=2）。
bincount 得到每个 expert 的计数，按 expert 顺序累加（前缀和）就是各段起点。
index_add_（或 scatter_add_）。out[idx] += y 是读-加-写，idx 有重复时后写覆盖先写；逐 expert 循环里 toks 恰好不重复所以能用，一次性全部写回就会丢贡献。
h @ w_gate[e].T：[n_e, H] @ [H, I] → [n_e, I]；down：[n_e, I] @ w_down[e].T → [n_e, H]。
每 token 一个标量：x @ gate.T 是 [N, 1]，sigmoid 后与 [N, H] 广播相乘。

形状追踪（N=5, H=16, E=8, I=32, k=2）：

logits      x[5,16] @ router.T[16,8]        -> [5, 8]
probs       softmax fp32                     -> [5, 8]
topk        weights [5, 2], ids [5, 2]
flatten     flat_e / flat_w / flat_tok       -> [10] each
order       argsort(flat_e, stable)          -> [10]
counts      bincount(minlength=8)            -> [8]     (sum = 10)
expert e    h = x[toks]                      -> [n_e, 16]
            gate/up                          -> [n_e, 32]
            down                             -> [n_e, 16]
scatter     out.index_add_(0, toks, ·)       -> out [5, 16]
shared      sigmoid(x @ gate.T) [5,1] * SwiGLU_s(x) [5,16]

与 serving 的联系：argsort + bincount 对应 FusedMoE 的 moe_align_block_size（把 token 按 expert 对齐并 pad 到 block）；循环里的每个 expert matmul 对应 grouped GEMM 的一个 group；Expert Parallel 时 pair_token / pair_expert 就是 all-to-all dispatch 的路由表。

面试官追问 & 对照源码

Expert Parallel 时 token 要跨卡 all-to-all，路由结果怎么用？
为什么 vllm-omni 的 deploy 里给 Qwen3-Omni 指定 moe_backend: triton？

对照源码：model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py:356-393 · deploy/qwen3_omni_moe.yaml:28-31 · HF Qwen3OmniMoeThinkerTextSparseMoeBlock · HF Qwen3OmniMoeTalkerTextSparseMoeBlock