背景
TopKRouter是混合专家模型（Mixture of Experts, MoE）中的门控路由器。它接收一个批次中所有 Token 的隐藏状态，为每个 Token 选出最合适的K 个专家，并计算对应的权重。路由包含三个关键步骤：
- 用线性层产生每个专家的 logits；
- 对 logits 做 softmax 得到概率分布；
- 选取概率最大的 K 个专家，并重归一化这些权重。
三个 TODO 分别对应第 2、3 步以及重归一化。
TODO 1：对全量 Logits 进行 Softmax 获取所有专家的概率分布
routing_probs=F.softmax(router_logits.float(),dim=-1)作用：将门控层输出的原始分数router_logits（形状[total_tokens, num_experts]）转换为概率分布，使每个 Token 在所有专家上的权重和为 1。
细节与必要性：
.float()：强制将router_logits转换为float32再进行 softmax。在混合精度训练（如bfloat16或float16）中，logits 可能数值范围大，直接 softmax 容易溢出或损失精度。提升到 FP32 能保证数值稳定性。F.softmax(..., dim=-1)：在最后一个维度（专家维度）上计算 softmax。结果routing_probs的每一行是概率向量，表示该 Token 与各专家的匹配程度。- 这一步是路由的基础，为后续 Top-K 选择提供标准的概率分布。
TODO 2：从概率分布中截取 Top-K 最大的概率及其索引
routing_weights,selected_experts=torch.topk(routing_probs,self.top_k,dim=-1)作用：对于每个 Token，从routing_probs中选出概率最大的self.top_k个专家，得到它们的权重和专家索引。
细节：
torch.topk(routing_probs, self.top_k, dim=-1)返回两个张量：routing_weights：形状[total_tokens, top_k]，每个 Token 选出的 K 个最高的概率值。selected_experts：形状[total_tokens, top_k]，这些概率对应的专家索引（long类型）。
- 这一步体现了 MoE 的“稀疏性”：每个 Token 只激活 K 个专家，其余专家不参与计算。大大节省了计算量。
- 选取后的
routing_weights不再是一个完整的概率分布（因为舍弃了其余专家），需要下一步的处理。
TODO 3：对截取后的 routing_weights 进行重归一化（Re-normalize）
routing_weights=routing_weights/routing_weights.sum(dim=-1,keepdim=True)作用：让选出的 K 个专家的权重按比例放大，使其总和为 1，变成一个合法的概率分布（仅在这 K 个专家上）。
为何需要重归一化：
- Top-K 截取后，被选中的 K 个概率之和通常小于 1（除非这 K 个恰好就是概率最大的全部且剩余概率为 0）。
- 后续在
SparseMoEBlock中，每个选中的专家输出将乘以routing_weights再累加到最终结果。如果直接使用截取后的概率，累加的总贡献会变小（相当于多了一个衰减因子），这不符合“这 K 个专家完全负责该 Token”的预期，会改变梯度的尺度。 - 重归一化保证了“被激活的这 K 个专家共同解释该 Token 的全部输出”，使得输出幅值稳定。
- 计算方式：
routing_weights / routing_weights.sum(dim=-1, keepdim=True)，对每个 Token 在top_k维度上除以各自的和。
最后，代码通过routing_weights = routing_weights.to(hidden_states.dtype)将权重转回输入的精度（如bfloat16），以便与后续计算保持一致。
额外说明：SparseMoEBlock 如何使用这些结果
虽然这不是 TODO，但理解上下文有助于理解路由的目的：
在SparseMoEBlock.forward中，对每个专家使用torch.where(selected_experts == expert_idx)找到所有选中该专家的 Token，分别计算专家输出，然后乘以对应的重归一化权重routing_weights[token_idx, kth_expert]，累加到最终输出中。这体现了 MoE 的“稀疏加权汇聚”：每个 Token 只被自己选中的 K 个专家处理，最后按权重合并。路由器返回的routing_weights和selected_experts正是驱动这一过程的核心数据结构。
总结
- TODO 1用 FP32 执行 softmax，得到每个 Token 对所有专家的完整概率分布。
- TODO 2通过
torch.topk从中选出 K 个最大的概率和对应的专家索引，实现稀疏选择。 - TODO 3对这 K 个权重进行重归一化，使它们和为 1，保证激活专家的输出幅值稳定。
这三步共同构成了 MoE 路由的标准流程，平衡了计算效率和模型容量。
