#1 · Track M · Qwen-Omni 模型实现 · 书 p.19

M01 · RMSNorm 与多模态 RoPE（interleaved / chunked）

thinker 的每一层都从 RMSNorm 开始，每个 attention 都要对 q/k 做三维 (t,h,w) 旋转位置编码；Qwen3 与 Qwen2.5 的频率布局不同——错一位不会报错，只会悄悄变笨。

⏱ 45 min 难度中实现 mine/m01_rmsnorm_mrope.py 验收 ./check.sh m01 提示 A-1 · 参考解 C-1 solutions/m01_rmsnorm_mrope.py 标记完成

题面速览

在系统中的位置：thinker（stage 0）每个 decoder layer：h = h + Attn(RMSNorm(h))、h = h + MoE(RMSNorm(h))，Attention 里对 q、k 施加旋转位置编码。talker 用同样的组件；M02、M08 会复用这里的函数。
RMSNorm 数值约定：y = w · x / sqrt(mean(x²) + eps)。HF/vLLM 的细节：先把 x 转 fp32 做归一化，转回输入 dtype 之后再乘 weight。bf16 下顺序不同就不 bitwise 一致。
RoPE：head_dim = D，频率 j（0 ≤ j < D/2）的角速度 inv_freq[j] = θ^(-2j/D)，位置 p 的角度为 p·inv_freq[j]；cos/sin 形状 [T, D]（两半各复制一次）；旋转 x·cos + rotate_half(x)·sin，rotate_half([a,b]) = [-b, a]。
MRoPE 三维位置：每个 token 有 (t, h, w)。文本三者相同；图像 patch 有各自网格坐标；音频沿时间线性。mrope_section（真实 Qwen3-Omni 为 [24, 20, 20]，和为 64 = 128/2）给出每个轴分到几个频率。

模型	布局（D/2 = 8，section = [4,2,2] 为例）
Qwen2-VL / Qwen2.5-Omni：chunked	TTTTHHWW
Qwen3-VL / Qwen3-Omni：interleaved	THWTHWTT

interleaved 规则：j%3==1 且 j<3·section[1] → H；j%3==2 且 j<3·section[2] → W；其余 → T。动机：每个轴都同时拿到高频和低频分量，而 chunked 让 W 轴只拿到最低频。

要实现：rms_norm、rope_inv_freq、mrope_cos_sin(inv_freq, positions[3,T], section, interleaved)、rotate_half、apply_rotary(x[T,heads,D], cos, sin)。本套题 q/k 用 [T, heads, D]（vLLM 风格，token 在前），HF 用 [B, heads, T, D]，测试会帮你转。

验收要点：RMSNorm 在 fp32 / bf16 下都与 HF 一致 · inv_freq 一致 · interleaved cos/sin 与 HF Qwen3-Omni 一致（含图像网格位置）· 纯文本时两种布局都退化为 1D RoPE · apply_rotary 后 q/k 一致 · chunked 与 HF Qwen2.5-Omni 一致。

直觉与类比

音量归一化 + 均衡器

一个歌单里每首歌音量不同。RMSNorm 先把每首歌（每个 token 的 hidden 向量）的整体响度拉到同一水平（除以 RMS），再用均衡器（weight，每个维度一个旋钮）调每个频段。 “先转回 x.dtype 再乘 weight”就像：先把音频导出成 16-bit 再过均衡器，还是在 32-bit 工程里过完均衡器再导出——耳朵听不出差别，但逐字节比较不一样，parity 测试就是那个逐字节比较的人。

一排转速不同的钟表指针

把 head_dim = D 的向量拆成 D/2 对，每一对是一根指针。第 j 根指针转速 inv_freq[j] = θ^(-2j/D)：j = 0 最快（秒针），j 越大越慢（时针、日历）。token 在位置 p 时，第 j 根指针转到角度 p · inv_freq[j]。两个 token 做点积只留下角度差 → attention 天然感知相对位置。快指针分辨“隔壁 vs 隔两个”，慢指针分辨“隔 10 个 vs 隔 1000 个”。旋转不改变长度，所以 RoPE 不会破坏 QK-norm 刚归一化好的范数。

发牌：一人一摞 vs 轮流发

把 64 张牌按“快 → 慢”排好，发给 T、H、W 三位玩家。chunked：前 24 张给 T、接着 20 张给 H、最后 20 张给 W——W 只拿到最慢的指针，相邻两列 patch 角度几乎没差别。interleaved：轮流发 T H W T H W …，H、W 的份额发完后剩下的全给 T，每个轴都同时拿到快指针和慢指针。

框架视角 · shape trace

记号：T 个 token，D = head_dim，half = D/2，heads 为 q 或 k 的 head 数。

#	tensor	shape / dtype	发生了什么
1	`x → rms_norm`	`[..., H]` 任意 dtype	转 fp32，最后一维 `mean(x²)`（keepdim），乘 `rsqrt(·+eps)`，转回 x.dtype，再乘 weight
2	`inv_freq`	`[half]` fp32	`θ^(-2j/D)`，j = 0 … half-1
3	`positions`	`[3, T]` int	三行依次是 t、h、w；文本 token 三行相同
4	`freqs`	`[3, T, half]` fp32	一次广播：每个轴、每个 token、每根指针的角度
5	`f`	`[T, half]`	对每一列 j 从三个轴里选一个拷过来（只选列，不做新运算）；interleaved / chunked 只差在这里
6	`emb`	`[T, D]`	`cat([f, f], -1)`：第 j 列与第 j+half 列角度相同
7	`cos, sin`	`[T, D]` fp32	`emb.cos()`、`emb.sin()`
8	`apply_rotary(x)`	`x[T, heads, D]`	cos/sin 插入 head 维变 `[T, 1, D]`、转成 x.dtype；`x·cos + rotate_half(x)·sin`

易错点

bf16 下 RMSNorm 在 fp32 里先乘 weight 再转回：数值“差不多”，但不 bitwise 一致。
mean 忘了 keepdim=True，或对全部维度求 mean。
interleaved 时直接改了 freqs[0]（没 clone），或切片终点写成 section 而不是 3 · section。
chunked 时 H 段取了 freqs[1] 的第 0..s_h 列，而不是同一列区间 s_t..s_t+s_h。
把每个频率相邻复制两次（repeat_interleave，GPT-J 风格），与“前后半配对”的 rotate_half 不匹配。
rotate_half 写成 [b, -a]：旋转方向反了。
cos 直接与 [T, heads, D] 相乘：当 T == heads 时不会报错，但 T 维被错误地对齐到了 heads 维。

逐步图：频率分给哪个轴？

▶ 用 acos(cos) “看见”每个频率跟哪个轴转

1. 设定一个可以肉眼读的输入

取 inv_freq = ones(10)（D = 20，half = 10），单个 token，positions = [[1],[2],[3]]（t=1, h=2, w=3）。

这样每个频率的角度 = 它所用轴的位置，acos(cos[0]) 直接告诉你第 j 个频率跟哪个轴转：1 = T 轴，2 = H 轴，3 = W 轴。

这是 L2 推荐的自测技巧：用 ones 的 inv_freq + 互不相同的小整数位置。

2. section = [4,3,3]，interleaved=True

j = 0…9 的来源轴：

1231231231 | 后半段（j = 10…19）是前半段的复制

H 占 j = 1, 4, 7（slice(1, 3·3=9, 3)，共 3 个 = s_h）
W 占 j = 2, 5, 8（slice(2, 9, 3)，共 3 个 = s_w）
其余（0, 3, 6, 9）是 T，共 4 个 = s_t

3. section = [4,3,3]，interleaved=False（chunked）

1111222333 | 后半段同前半段

三段连续：freqs[0][:, 0:4]、freqs[1][:, 4:7]、freqs[2][:, 7:10]——注意 H 段取的是 同一列区间 4..7，不是 freqs[1][:, 0:3]。

4. section = [6,2,2]，interleaved=True

1231231111

H、W 只占前 3·section 个“槽位”：slice(1, 6, 3) → 1, 4；slice(2, 6, 3) → 2, 5。之后（j ≥ 6，高亮）全是 T。

这就是为什么切片终点必须是 3 * section[axis] 而不是 half。

5. 纯文本：三轴位置相同

文本 token 的 positions 三行相同 → freqs[0] == freqs[1] == freqs[2] → 不管每列从哪个轴拷，结果都一样。

所以 interleaved 和 chunked 都 bitwise 退化成普通 1D RoPE（测试会检查这一点）。

6. cat([f, f]) 与 rotate_half 的配对

emb = cat([f, f], -1)：第 j 列和第 j + D/2 列角度相同。

rotate_half([x0, x1 | x2, x3]) = [-x2, -x3 | x0, x1]，于是 x·cos + rotate_half(x)·sin 正好把 (x0, x2)、(x1, x3) 各自做一次二维旋转：

out[0] = x0·cos θ0 − x2·sin θ0
out[2] = x2·cos θ0 + x0·sin θ0

GPT-J 风格（相邻两维配对、repeat_interleave）与这里的 neox 风格不兼容。

7. apply_rotary 的广播

x 是 [T, heads, D]，cos/sin 是 [T, D]。必须 cos[:, None, :] → [T, 1, D]。

直接相乘时 PyTorch 从右对齐：[T, D] 会被当成 [heads, D]——T == heads 时不报错但全错。

bf16 输入时 cos/sin 先 .to(x.dtype) 再乘（与 HF 一致）。

三级提示

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

rms_norm：x = [[3, 4], [1, 1]]（fp32），weight = [1, 2]，eps = 0。

row0: mean(x²) = (9+16)/2 = 12.5, rsqrt = 1/3.5355
      x̂ = [0.8485, 1.1314] -> * weight -> [0.8485, 2.2627]
row1: mean(x²) = 1 -> x̂ = [1, 1] -> * weight -> [1.0000, 2.0000]

对最后一维求 mean（每行独立），输出形状与 x 相同。

rope_inv_freq：rope_inv_freq(8, 10000.0) → [1.0000, 0.1000, 0.0100, 0.0010]（长度 D/2 = 4）；rope_inv_freq(6, 100.0) → [1.0000, 0.2154, 0.0464]。

mrope_cos_sin 的轴分配（inv_freq = ones(10)，positions = t1/h2/w3，acos(cos[0])）：

section / 布局	acos(cos[0])（长度 20）
`[4,3,3]`, interleaved=True	1 2 3 1 2 3 1 2 3 1 \| 1 2 3 1 2 3 1 2 3 1
`[4,3,3]`, interleaved=False	1 1 1 1 2 2 2 3 3 3 \| 1 1 1 1 2 2 2 3 3 3
`[6,2,2]`, interleaved=True	1 2 3 1 2 3 1 1 1 1 \|（后半同前半）

rotate_half / apply_rotary：rotate_half([1, 2, 3, 4]) = [-3, -4, 1, 2]。取 D = 4，inv_freq = [1, 0.01]，纯文本位置 p = 2（section=[1,1,0]，interleaved）：

cos = [cos2, cos0.02, cos2, cos0.02] = [-0.4161, 0.9998, -0.4161, 0.9998]
sin = [ 0.9093, 0.0200, 0.9093, 0.0200]
apply_rotary(x=[[[1,0,0,0]]]) = [[[-0.4161, 0, 0.9093, 0]]]       # 维度 0 与维度 2 组成一对
apply_rotary(x=[[[1,2,3,4]]]) = [[[-3.1440, 1.9196, -0.3391, 4.0392]]]

关键问题

RMSNorm 里“先转回 x.dtype 再乘 weight”和“fp32 里乘完 weight 再转回”在 fp32 输入下有区别吗？bf16 下呢？为什么？
torch.arange(0, head_dim, 2) 得到什么？除以 head_dim 后指数范围是多少？inv_freq 长度是多少？
positions [3, T]、inv_freq [D/2]，怎样一次广播得到 freqs[3, T, D/2]？
interleaved 下 j = 7、section = [4,3,3] 时属于哪个轴？section = [6,2,2] 呢？
rotate_half 中维度 j 与 j+1 配对，还是与 j + D/2 配对？这和 cat([f, f]) 有什么关系？
纯文本（三轴位置相同）时 interleaved 与 chunked 的 cos/sin 是什么关系？
apply_rotary 里 x 是 [T, heads, D]、cos 是 [T, D]，怎样 reshape 才能正确广播？

细节

“normalize in fp32, cast back, THEN multiply by weight”：x * rsqrt(mean(x², -1) + eps) 全程 fp32；结果先 .to(x.dtype) 再乘 weight。bf16 下两种顺序有 1 ulp 量级差异（例：正确顺序 1.9140625，fp32 先乘得 1.90625）。
inv_freq 显式 dtype=torch.float32 最稳；theta ** tensor 是标量的逐元素幂。
mrope_section = [s_t, s_h, s_w]，和为 D/2，是“每个轴分到多少个频率”（维度数是 2 倍）。
cos/sin 形状 [T, head_dim]：先在 [T, D/2] 上完成轴选择，再 cat([f, f], -1)，最后 .cos() / .sin()。
Qwen3 vs Qwen2.5：inv_freq、rotate_half、apply 公式完全相同，唯一区别是“频率 j 用哪个轴的位置”——这就是 interleaved 这个 bool 的全部含义。

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

拆分：五个函数层层组合 rope_inv_freq → mrope_cos_sin → apply_rotary（内部调用 rotate_half），rms_norm 独立。按这个顺序写，每写完一个用 L1 的数字自测。
rms_norm：x32 = x.float() → x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) → .to(x.dtype) 后乘 weight。常见错：忘 keepdim、eps 放在 sqrt 外、对全部维度求 mean。
mrope_cos_sin 的核心是“选择”而非“计算”：先一次算出三条轴各自的 freqs[a]（[T, D/2]），再构造 [T, D/2] 的 f，每一列 j 从某个轴的 freqs 拷过来。
interleaved：以 T 轴为底（拷贝一份，避免改到原 tensor），再用步长 3 的切片覆盖 H 的列（起点 1）和 W 的列（起点 2），终点由对应 section 决定。也可以先构造长度 D/2 的 axis_of_j 向量再 gather。
chunked：维护起点游标，依次从 freqs[0]、freqs[1]、freqs[2] 切出长度 s_t、s_h、s_w 的列段再拼起来。切的是同一个列区间。
rotate_half：chunk(2, dim=-1) 切成 a、b，返回 cat([-b, a], -1)。错误：[b, -a]（sin 符号整体翻转），或按奇偶位切（GPT-J 风格）。
apply_rotary 广播：cos[:, None, :]。直接相乘会把 cos 的 T 维对齐到 x 的 heads 维——T == heads 时不报错但结果错，这是最阴险的 bug。
自测技巧：inv_freq = ones + 互不相同的小整数位置（1/2/3）跑 mrope_cos_sin，取 acos 肉眼看每列来自哪个轴；三轴相同的位置验证两种布局都退化为普通 RoPE。

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

关键问题答案

fp32 输入下 .to(x.dtype) 是 no-op，两种顺序相同；bf16 下“先 round 到 bf16 再乘 bf16 weight”和“fp32 乘完再 round”舍入点不同，可能差 1 ulp。
arange(0, D, 2) = [0, 2, …, D-2]，除以 D 后指数在 [0, 1)；inv_freq 长度 D/2，单调递减，从 1 到接近 1/θ。
positions.float()[:, :, None] * inv_freq[None, None, :] → [3, T, D/2]。
j = 7：7 % 3 == 1；section=[4,3,3] 时 3·3 = 9 > 7 → H；section=[6,2,2] 时 3·2 = 6 ≤ 7 → T。
维度 j 与 j + D/2 配对；cat([f, f]) 让第 j 列和第 j+D/2 列角度相同，rotate_half 正好把这两列组成一个二维旋转。
三轴相同 → 三条 freqs 完全一样 → 无论怎么选列结果都相同，两种布局 bitwise 相同。
cos[:, None, :]（[T, 1, D]）。

mrope_cos_sin 形状追踪（T 个 token，D = head_dim，half = D/2）：

positions [3, T] --float--> [3, T, 1] * inv_freq [1, 1, half] = freqs [3, T, half]
选择列                                                   -> f   [T, half]
cat([f, f], -1)                                          -> emb [T, D]
emb.cos(), emb.sin()                                     -> [T, D], [T, D]

骨架检查：slice(1, 9, 3) 选中 1, 4, 7（3 个 = s_h）；slice(2, 9, 3) 选中 2, 5, 8（3 个 = s_w）。

rotate_half / apply_rotary 验算：x = [1,2,3,4]，cos = [c0,c1,c0,c1]，sin = [s0,s1,s0,s1]，rotate_half(x) = [-3,-4,1,2]： out[0] = 1·c0 − 3·s0 = −0.4161 − 2.7279 = −3.1440，out[2] = 3·c0 + 1·s0 = −1.2484 + 0.9093 = −0.3391——正是二维旋转矩阵作用在 (x0, x2) 上。

最终检查清单：(1) bf16 下 rms_norm 先 cast 再乘 weight；(2) inv_freq 显式 fp32；(3) interleaved 用 clone() 再覆盖，切片终点 3*section[axis]；(4) chunked 的列区间是累积偏移；(5) cat([f, f]) 而不是 repeat_interleave；(6) apply_rotary 插入 head 维并对齐 dtype。

面试官追问 & 对照源码

为什么 vllm-omni 在 CUDA 上要给 thinker 强制 custom_ops: ["+rotary_embedding"]？（编译后的 MRoPE 与 QK-norm 融合改变了音频输出）
decode 阶段只有一个新 token，cos/sin 怎么高效地增量计算？

对照源码：deploy/qwen3_omni_moe.yaml:84-91 · model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py:1694-1785 · HF Qwen3OmniMoeThinkerTextRotaryEmbedding.apply_interleaved_mrope · HF Qwen3OmniMoeThinkerTextRMSNorm