题目内容

在混合专家大模型中，为了降低计算开销，并非所有 Token 都会经过所有专家。系统中有 $E$ 个专家，路由器会对每个 Token 在所有专家上的得分进行评估，并只把该 Token 路由给得分最高的 $K$ 个专家。同时，每个专家有一个处理容量上限 $C$，表示该专家最多能处理的 Token 数量；一旦某专家的负载达到 $C$，后续原本要分配给它的 Token 会被强制掩蔽并丢弃。

现有一个由 $N$ 个 Token 组成的批次，给定一个 $N$ 行 $E$ 列的得分矩阵 $S$，其中 $S_{i,j}$ 表示 Token $i$ 对专家 $j$ 的分数。需要按照 Top-K 路由、容量掩码调度、惩罚聚合三个步骤进行模拟。

第一步：Top-K 路由掩码。对于每个 Token $i$，从矩阵第 $i$ 行的 $E$ 个专家中找出得分最高的 $K$ 个专家。如果出现同分，则专家索引较小的优先。生成的初始二进制掩码矩阵记为 $M^{(1)}$：如果专家 $j$ 位于 Token $i$ 的前 $K$ 名中，则 $M^{(1)}_{i,j}=1$，否则 $M^{(1)}_{i,j}=0$。

第二步：动态容量掩码调度。按照 Token 顺序从 0 到 $N-1$ 依次处理。维护数组 $load$，长度为 $E$，初始值均为 0。对当前 Token，按 Top-K 顺序检查其候选专家。对于候选专家 $j$，如果 $load[j] < C$，则本次路由生效，最终掩码 $M^{(2)}_{i,j}=1$，并将 $load[j]$ 增加 1；否则该 Token 在专家 $j$ 上被掩蔽，$M^{(2)}_{i,j}=0$，负载保持不变。同一个 Token 内的多个候选专家进行容量判断时互不影响。

第三步：路由惩罚聚合。最终每个专家的实际负载保存在 $load$ 中。本批次的路由不平衡惩罚值 $P$ 定义为所有专家负载的平方和：