算法面试实战：手撕 MoE (混合专家模型) 路由优化算法
随着大模型（如 GPT-4、DeepSeek）纷纷转向 MoE（Mixture of Experts）架构，关于 MoE 的底层机制成为了今年算法岗面试的"重灾区"。
今天，我们来看一道非常经典的华为算法面试真题：要求手写实现一个针对设备限制（Device-Limited）的 MoE 路由优化算法。
🌟 核心面试概念扫盲：什么是 MoE 的路由优化？
在写代码之前，我们要先理解为什么要做这件事。
在 MoE 模型训练时，Token 会根据概率被发送到 Top-K 个不同的专家（Experts）进行计算。在分布式集群中，这些专家分布在多个 NPU（或 GPU）卡上。
如果 Token 满天飞，频繁跨卡传输，会产生巨大的通信成本 （网络带宽往往是大模型训练的瓶颈）。
为了降低通信开销，我们引入了 Device-Limited Routing（设备限制路由）算法 ：它的核心思想是限制 Token 只能路由到特定的 PPP 个 NPU 上，从而将网络通信范围锁定在局部。
补充面试考点 (Cheat Sheet)：
面试官在出这道题时，大概率还会追问以下两个 MoE 核心概念：
- 辅助损失 (Auxiliary Loss)： MoE 模型非常容易出现"赢者通吃"的现象------少数几个专家越来越强，所有的 Token 都涌向它们，导致其他专家被"饿死"。为了解决这个问题，我们需要在损失函数中加入辅助损失（公式：Laux=α⋅∑i=1EfiPiL_{aux} = \alpha \cdot \sum_{i=1}^{E} f_i P_iLaux=α⋅∑i=1EfiPi），强迫模型尽可能均匀地将 Token 分配给所有专家。
- 容量约束 (Capacity Constraints)： 显存和算力是有限的。我们会设定每个专家最多能接收多少 Token（容量阈值）。如果排队的 Token 超过了这个限制，多出来的 Token 就会被直接丢弃（Drop），不参与本次训练。
💻 算法题面与解题思路
题目描述
- 把 nnn 个专家平均分配 在 mmm 个 NPU 上，每个 NPU 上的专家组成一个"专家组"。
- 每个专家对应一个被路由到的概率。我们需要找出每个 NPU 组里的最大概率作为该组的代表。
- 从所有的 NPU 组中，挑选出代表概率最大的 ppp 个组（这就限制了跨卡通信范围）。
- 最后，把这 ppp 个 NPU 上的所有专家放在一起，选出概率最大的 kkk 个专家，作为最终的路由目标。
面试做题技巧
Tips： 在面试中遇到这种长描述的代码题，**千万不要直接闷头写代码！**一定要先在白板（或编辑器）上写下解题的 1、2、3、4 步，和面试官对齐思路后，再开始填充代码。
解题四步走：
- 分组： 确定每个 NPU 上包含哪些编号的专家。
- 选组： 计算每组的最大概率，排序选出 Top-ppp 个组。
- 聚合： 把这 ppp 个组里的所有专家及其概率收集到一个大列表里。
- 决胜： 对大列表按概率全局排序，选出最终的 Top-kkk 专家。
🚀 完整可运行的 Python 代码实现
下面是直接可以复制运行的优质代码实现，逻辑清晰，适合在面试中手撕：
python
import numpy as np
def moe_routing_optimization(n, m, p, k, expert_probs):
"""
n: 专家总数
m: NPU数量 (组数)
p: 限制路由目标的 NPU 数量
k: 最终选择的专家数量
expert_probs: 长度为 n 的数组，表示每个专家的概率
"""
# === Step 1: 将 n 个专家平均分配到 m 个 NPU 上 ===
experts_per_npu = n // m
groups = []
for i in range(m):
start = i * experts_per_npu
# 处理最后一个 NPU 可能无法完全均分的情况
end = (i + 1) * experts_per_npu if i != m - 1 else n
groups.append((i, start, end)) # 记录 (NPU编号, 起始专家编号, 结束专家编号)
print(f'每个组的专家编号范围: {groups}\n')
# === Step 2: 计算每组的最大概率，选择概率最大的 p 个组 ===
group_max_probs = []
for npu_id, start, end in groups:
# 切片获取该组内所有专家的概率
group_prob = expert_probs[start:end]
max_group_prob = np.max(group_prob)
group_max_probs.append((npu_id, max_group_prob))
# 按组的最大概率降序排序，选择前 p 个组
group_max_probs.sort(key=lambda x: x[1], reverse=True)
selected_groups = group_max_probs[:p]
# === Step 3: 收集选中的 p 个 NPU 对应的所有专家 ===
selected_experts = []
for npu_id, _ in selected_groups:
# 找到对应 NPU 组的起始和结束索引
start, end = next((s, e) for idx, s, e in groups if idx == npu_id)
npu_experts_probs = expert_probs[start:end]
expert_indices = np.arange(start, end)
# 将局部排序的结果加入候选列表
sorted_indices = expert_indices[np.argsort(npu_experts_probs)[::-1]]
sorted_probs = npu_experts_probs[np.argsort(npu_experts_probs)[::-1]]
selected_experts.extend(list(zip(sorted_indices, sorted_probs)))
# === Step 4: 从所有候选专家中，全局选出概率最大的 k 个 ===
selected_experts.sort(key=lambda x: x[1], reverse=True)
final_experts = selected_experts[:k]
return final_experts
# ====================
# 测试代码
# ====================
if __name__ == "__main__":
n = 8 # 专家总数
m = 2 # NPU 数量
p = 1 # 限制选择的组数量
k = 2 # 最终选择的专家数量
# 随机生成每个专家的概率
expert_probs = np.random.rand(n)
print(f'随机生成每个专家的概率:\n{expert_probs}\n')
# 执行路由算法
final_routing = moe_routing_optimization(n, m, p, k, expert_probs)
print("最终路由的专家编号及概率:")
for expert_id, prob in final_routing:
print(f"专家编号: {expert_id}, 概率: {prob:.4f}")总结
这道题完美融合了数组切片操作、自定义排序规则、以及大模型底层的系统设计理念 。在准备面试时，建议大家不仅仅是背诵代码，更要能流利地说出引入 p（NPU 组数限制）是为了减少通信开销的工程考量，这样才能在一众候选人中脱颖而出。
