Router 打分：每个 token 先看所有专家，再选 Top-K
🎯 先猜一猜（下注解锁）
Router 有多种等价写法。本关 notebook 明确要求先得到完整概率表，再从概率里取 Top-K。
先判断：为什么 notebook 要先对全量 router_logits 做 softmax？
先补的知识
- MoE 的目标是“大容量、低激活”：模型有很多专家，但每个 token 只去少数几个专家那里计算。
- Router 是一个很小的线性层，把每个 token 的 hidden state 映射成 num_experts 个 logit 分数。
- 本 notebook 的实现路径是：先对所有专家做 softmax 得到完整概率表，再从概率里取 Top-K，最后对 Top-K 权重重归一化。
- Router 不负责真正执行专家网络，它只输出两张清单：选了哪些专家 selected_experts，以及每个专家拿多少权重 routing_weights。
图解原理
把 Router 想成分诊台。每个 token 先拿自己的特征名片给所有专家打分，softmax 把分数变成全局概率，再只保留最适合的 K 个专家。Notebook 的 TODO 1/2 本质就是这两步。初学时先盯住一行：一行代表一个 token，这一行里所有列代表它对所有专家的偏好。
TopKRouter.forward 的输入仍是三维句子张量，但输出是按 token 排好的两张清单。先把这个合同读清楚，后面代码就不会绕。
先算每个 token 对所有专家的概率，再保留 Top-K。真正执行专家网络是后面的 SparseMoEBlock，不要在 Router 里混进去。
这里的关键不是画很多线，而是抓住最后一维的含义：展平后每一行是一枚 token，最后一维是全部专家。
如果有 4 个专家，一行 logits 就是这个 token 对 4 个专家的偏好分数。softmax 和 topk 都只在这一行内部做。
1.2 E1
0.1 E2
2.0 E3
-0.5
dim=-1 就是专家维
不要把不同 token 混在一起归一化。Router 要让“同一个 token 的所有专家概率”加起来等于 1。
notebook 里 TODO 后面有 zero 占位变量。写完 TODO 后，要返回你真正算出的 routing_weights 和 selected_experts，不要让占位值覆盖结果。
MoE 的专家像并行工作台，Router 只决定每个 token 送去哪几个工作台，以及每个工作台占多少权重。调度得好，模型容量变大但每个 token 只激活少数专家；调度得差，热门专家拥堵，冷门专家闲置。
下一关的负载均衡损失，就是专门约束这个调度器不要总偏心同几个专家。
语法热身：先用一行玩具分数练 softmax 和 topk
# 1 个 token, 4 个专家
toy_logits = torch.tensor([[1.2, 0.1, 2.0, -0.5]])
# dim=-1：只在“专家这一行”里归一化，得到每个专家的概率
toy_probs = F.softmax(toy_logits.float(), dim=-1)
# topk 同时返回值(values)和下标(indices)
toy_weights, toy_ids = torch.topk(
toy_probs, k=2, dim=-1
)
print(toy_weights)
print(toy_ids)举一反三：把玩具张量换成 notebook 变量
玩具例子只有 1 个 token。notebook 只是把很多 token 排成 [tokens, experts]，同样沿最后一维做 softmax/topk。
迁移口诀：先把输入展平成 [tokens, hidden]，gate 后得到 [tokens, experts]；之后所有 softmax/topk 都沿 dim=-1，因为最后一维就是专家维。
巩固一下
hidden_states shape 是 [2,4,16]，num_experts=8，top_k=2。展平并 topk 后 selected_experts 的 shape 是？
回到 notebook 的作业
这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。