#2 · Track M · Qwen-Omni 模型实现 · 书 p.23

M02 · Thinker Attention：QK-norm + GQA + KV cache

Qwen3 的 attention 比 Llama 多一步：对每个 head 的 q、k 先 RMSNorm 再 RoPE；同一份实现要同时支持整段 prefill、分块 prefill 和逐 token decode，且三者逐行一致。

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

题面速览

在系统中的位置：thinker 和 talker 的每一层都有一个这样的 attention。serving 时它几乎从不“整段”运行：长 prompt 被 chunked prefill 切成多块，decode 时每步只有 1 个新 token，历史 K/V 来自 cache（vLLM 里是 paged KV，见 P06/P09；这里用简单的 list 拼接）。
计算步骤（单层，新 token 数为 T）：

q = x @ Wq^T -> [T, H, D]     k = x @ Wk^T -> [T, KVH, D]     v = x @ Wv^T -> [T, KVH, D]
q = RMSNorm_head(q; q_norm)   k = RMSNorm_head(k; k_norm)     # 对最后一维 D 做，每个 head 独立
q, k = RoPE(q, k; cos, sin)                                   # cos/sin 只对应这 T 个新 token
K, V = cat(cache.K, k), cat(cache.V, v)                       # S = past + T
score[h, t, s] = q[t, h] · K[s, h // (H/KVH)] / sqrt(D)，且 s ≤ past + t（因果）
out = softmax_fp32(score) @ V -> [T, H*D] @ Wo^T

为什么三种调用方式必须一致：调度器（P07）决定这一步算多少 token；模型代码只能依赖“新 token 的绝对位置 = len(cache) + i”这一条约定。因果 mask 必须按绝对位置算，而不是块内下标。
GQA：KV head 少于 Q head（这里 4 vs 2），head h 用 kv head h // (H/KVH)，即 repeat_interleave，不是 repeat。
边界：cache 为 None 就是普通 prefill；T=1 的 decode 和 T>1 的 chunk 用同一份代码。

验收要点：整段 prefill 与 HF 一致 · 先 prefill 5 个再逐个 decode 7 次逐行等于整段 · 分块 prefill（4 + 7）等于整段且 cache 长度正确。

直觉与类比

迟到的人读会议纪要

每个人发言都会记进纪要（KV cache）。迟到的人（新 token）不需要大家重开一遍会，只要读纪要。纪要上记的是当时每句话的时间戳——对应“cache 里的 k 已经用当时的绝对位置旋转过了”，以后永远不用重算。新来的人只能看他到场之前（含自己）的发言（因果 mask）。关键是要用会议总时钟判断“之前”，而不是“我今天是第 0 个到的”——否则 decode 的 token 以为自己在位置 0，只能看到第一句话。

GQA：同桌合用一本参考书

每个 Q head 是一个学生，每个 KV head 是一本参考书（H = 4、KVH = 2）：相邻座位的 H/KVH 个学生合用一本——学生 0、1 用书 0，学生 2、3 用书 1。这就是 h // (H // KVH)，等价于把书 repeat_interleave。误用 repeat 就变成“隔一个座位共用一本”（0、2 用书 0）——而在 H == KVH 的测试里你根本发现不了。

QK-norm：比赛前统一称重

在 RoPE 之前对每个 head 的 q、k 各自做 RMSNorm（对 head_dim），把每个 head 拉到同一“量级”，点积不会因为某些 head 激活特别大而爆掉。顺序是 proj → QK-norm → RoPE：RMS 对旋转不变，但逐维的 q_norm 权重不是，顺序反了数值就不同。v 不做 norm。

框架视角 · shape trace

记号：新 token 数 T，past = 调用前的 len(cache)，S = past + T。

#	tensor	shape	发生了什么
1	`x`	`[T, hidden]`	只有新 token
2	`q` / `k, v`	`[T, H, D]` / `[T, KVH, D]`	`x @ proj.T` 后 view；q_proj 的行按 head 分块，最后一维连续 D 个数属于同一个 head
3	QK-norm	同上	对最后一维 D 做 RMSNorm，所有 head 共享 q_norm（k 用 k_norm）；v 不动
4	RoPE	同上	cos/sin `[T, D]` → `[T, 1, D]` 广播到每个 head
5	cache	K, V: `[S, KVH, D]`	先读 past，再追加本次 k、v，再拼接全部块
6	GQA	`[S, H, D]`	每个 KV head 复制给相邻的 H/KVH 个 Q head
7	scores	`[H, T, S]` fp32	`q·K · D^-0.5`；`k_pos > past + t` 的位置置 -inf
8	probs @ V	`[T, H, D]`	softmax 在 key 维（fp32），加权求和后转回 x.dtype
9	输出	`[T, H·D] @ o_proj.T → [T, hidden]`	reshape 与第 2 步的 view 对称

易错点

在追加之后才读 len(cache)：past 多算了 T，mask 整体右移，token 能“偷看”同一块后面的 token。
mask 用块内下标 arange(T) 而不是 past + arange(T)：decode 的 token 只能看见 key 0。
GQA 用 repeat / expand + 错误 reshape 得到 0,1,0,1 映射；H == KVH 时看不出来。
QK-norm 放在 RoPE 之后，或对整个 H·D 做 norm 而不是对每个 head 的 D。
softmax 在 dim=-2 上做（每列和为 1），或在 bf16 里做 softmax。
对 cache 里的旧 k 重新做 RoPE（它们早已用正确的绝对位置旋转过）。

逐步图：toy 例子上的 prefill 与 decode

▶ hidden=2, H=KVH=1, D=2, 全单位阵投影, cos=1/sin=0

1. 输入与 QK-norm

x = [[1, 0],
     [0, 1],
     [3, 3]]                       # T = 3
QK-norm 之后: q = k = [[1.414, 0], [0, 1.414], [1, 1]]   # 每行除以自身 RMS
v 不做 norm:  v = x

scale = D^-0.5 = 1/1.414。

2. 一次性 prefill（past = 0）

mask 是普通下三角：

✓·· token0 只看见 key0
✓✓· token1 看见 key0..1
✓✓✓ token2 看见全部

token0 -> [1.0000, 0.0000]          # 只能看见自己
token1: scores [0, 1.414] -> softmax [0.1956, 0.8044] -> [0.1956, 0.8044]
token2: scores [1, 1, 1.414] -> softmax [0.2846, 0.2846, 0.4307] -> [1.5768, 1.5768]

3. 分两次调用：第一次 x[:2]

cache = KVCache()
a = attention_forward(w, x[:2], cos[:2], sin[:2], cache)  # -> [[1, 0], [0.1956, 0.8044]]
len(cache) == 2; cache.k == [tensor of shape [2, 1, 2]]

past = 0（在 append 之前读的）。

4. 第二次 x[2:]（decode，past = 2）

b = attention_forward(w, x[2:], cos[2:], sin[2:], cache)  # -> [[1.5768, 1.5768]]
len(cache) == 3; cache.k 里有两个块: [2,1,2] 和 [1,1,2]

q_pos = past + 0 = 2，k_pos = [0, 1, 2] 全部可见：✓✓✓

cos/sin 只有 1 行——它属于绝对位置 2 的那个新 token。cat([a, b]) 与一次性 prefill 逐行相同。

5. 两种错误的样子

错误	结果
完全没有 causal mask	`[[1.5665, 1.1690], [1.1690, 1.5665], [1.5768, 1.5768]]`——只有最后一行碰巧正确
decode 时 mask 用块内下标（q_pos = 0 而不是 2）	第三个 token 只能看见 key 0：✓✗✗ 输出 `[1, 0]` 而不是 `[1.5768, 1.5768]`

6. GQA 映射（H=8, KVH=2）

q head h：01234567

正确 h // 4（repeat_interleave）：00001111

错误 h % 2（repeat）：01010101

三级提示

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

见上面的逐步图：toy（hidden=2，H=KVH=1，D=2，全单位阵，q_norm=k_norm=[1,1]，eps=0，cos 全 1、sin 全 0）上整段 prefill 得 [[1,0],[0.1956,0.8044],[1.5768,1.5768]]，分两次（2 + 1）调用得到相同结果。

关键问题

q_proj 的形状是 [H·D, hidden]。x @ q_proj.T 得到 [T, H·D] 后，怎样 view 成 [T, H, D] 才能保证“第 h 个 head”是连续的 D 列？
QK-norm 对哪一维做？q_norm 形状 [D] 怎样广播到 [T, H, D]？v 要不要 norm？
cos/sin 只有 T 行，K 有 S = past + T 行。cache 里的旧 K 需要重新做 RoPE 吗？
第 i 个新 token（块内下标 i）的绝对位置是多少？它能看见哪些 key？
[S, KVH, D] → [S, H, D] 用 repeat 还是 repeat_interleave？H=4, KVH=2 时各给出什么映射？
“softmax in fp32”指哪几步在 fp32 里做？最后在哪里转回 x.dtype？
先把新 k/v 追加进 cache 再 attention，还是先 attention 再追加？哪个能让 token 看见它自己？

细节

“RMSNorm EACH HEAD … over head_dim”：把 q 看成 [T, H, D]，对最后一维 D 独立算 RMS；所有 head 共享 q_norm [D]。eps 取 w.eps。
“rows belong to the NEW tokens”：调用方已按绝对位置算好 cos/sin 并只切出新 token 的行；函数内部不需要、也不应该自己算位置。cache 中保存的是 RoPE 过的 k（以及原样的 v）。
KVCache：k、v 是块的 list，每块 [t, KVH, D]；len(cache) 是所有块 t 之和——所以 past 必须在 append 之前读。
causal 按绝对位置：p = past + i，可见 key 0..p。整段 prefill 时退化为普通下三角；decode（T=1）时能看见全部 S 个 key。
输出：[T, H, D] 先 reshape 成 [T, H·D] 再 @ o_proj.T。

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

步骤拆分：(1) 三个投影 + view；(2) per-head RMSNorm helper（复用 M01：fp32 归一化 → 转回 → 乘 weight）；(3) rotate_half + RoPE helper（cos/sin 在 head 维广播）；(4) cache 处理得到 past 和全量 K/V；(5) GQA 展开；(6) scores + mask + softmax + 加权求和；(7) reshape + o_proj。
cache 顺序：先读 past = len(cache)，再 append，再 torch.cat 全部块。先 append 再读 len → past 多算 T，mask 整体偏移。cache 为 None 时 past = 0。
GQA：k.repeat_interleave(H // KVH, dim=1) → 0,0,1,1 映射。
scores：einsum("thd,shd->hts") 再乘 D ** -0.5，在 fp32 里做（q、K 先 .float()）。
mask：q_pos = arange(past, past+T)[:, None]，k_pos = arange(S)[None, :]，k_pos > q_pos 处填 -inf，[T, S] 自动广播到 [H, T, S]。每个 query 至少能看见 key 0，不会整行 -inf。
softmax 与加权：softmax(scores, dim=-1)，再 einsum("hts,shd->thd")，然后 .to(x.dtype)。softmax 在 dim=-2 上做时 toy 的 token0 不再是 v0。
自测：随机 T=6 输入，(a) 一次调用；(b) 3+3 两次；(c) 6 次 T=1，三者 cat 后 allclose。

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

关键问题答案

(x @ q_proj.T).view(T, H, D)：q_proj 的行按 head 分块（前 D 行是 head 0），view 后最后一维连续 D 个数属于同一个 head。
对最后一维 D；[D] 自动广播到 [T, H, D]；v 不 norm。
不需要：cache 里存的就是当时已用正确绝对位置旋转过的 k，位置不会变。
绝对位置 = past + i；可见 key 下标 0..past+i（含）。
repeat_interleave：H=4, KVH=2 → [0,0,1,1]（正确）；repeat → [0,1,0,1]（错误）。
q·k、缩放、mask、softmax、与 V 的乘法都在 fp32；乘完 V 后转回 x.dtype，再做 o_proj。
先追加再 attention：新 token 的 key 也在 K 里，mask 允许 s == p，token 能看见自己。

形状追踪（H=4, KVH=2, D=8, hidden=32, 已缓存 past=5，新 T=2）：

x             [2, 32]
q             [2, 4, 8]      k, v [2, 2, 8]
QK-norm       same shapes
RoPE          q, k with cos/sin [2, 8] -> [2, 1, 8] broadcast
cache cat     K, V [7, 2, 8]             (S = 5 + 2)
GQA           K, V [7, 4, 8]
scores        [4, 2, 7]       mask [2, 7]: row0 sees s ≤ 5, row1 sees s ≤ 6
probs @ V     [2, 4, 8] -> reshape [2, 32] -> @ o_proj.T -> [2, 32]

逐行核对 toy 的 token2（past=2，T=1）：q = [1,1]，K = [[1.414,0],[0,1.414],[1,1]]，scores = [1.414, 1.414, 2]/1.414 = [1, 1, 1.414]；q_pos = 2 全部可见；softmax = [0.2846, 0.2846, 0.4307]；输出 = 0.2846·[1,0] + 0.2846·[0,1] + 0.4307·[3,3] = [1.5768, 1.5768]。得到 [1, 0] 说明 q_pos 用了 0；第一行不是 [1, 0] 说明 mask 缺失或 softmax 维度错了。

与 serving 的联系：“past = len(cache)，新 token 绝对位置 = past + i”正是 vLLM 里 seq_lens / query_start_loc 携带的信息；paged attention 只是把 torch.cat(cache.k) 换成按 block table 读取，mask 逻辑不变。

面试官追问 & 对照源码

把 list 拼接换成 P06 的 paged KV + P09 的 paged attention，接口要怎么变？
QK-norm 和 RoPE 能不能融合成一个 kernel？vllm-omni 为什么曾经因为编译融合出问题？

对照源码：model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py:588-596 · HF Qwen3OmniMoeThinkerTextAttention.forward