题目内容

在混合专家大模型中，为了节省算力，并不是所有的神经网络层都会处理所有的$Token$。

当系统包含$E$个专家时，对于输入的每一个$Token$，路由器会给各个专家打分，并只将该$Token$分配给得分排名前$K$的专家。同时，为了防止某个专家被分配了过多的Token导致显存溢出，系统为每个专家设置了最大处理容量$C$。如果某个专家达到了容量上限，后续被路由到该专家的Token将被强制“掩蔽”并丢弃。

给定$N$个$Token$对$E$个专家的原始打分矩阵$S$（维度$N \times E$），具体流程如下：

1. $Top-K$路由掩码生成：对于打分矩阵的每一行$i$（代表第$i$个Token），找出得分排名前$K$的专家索引。生成一个初始的$N \times E$二进制路由掩码矩阵$M^{(1)}$。

   • 如果专家$j$在前$K$名中$M_{i,j}^{(1)}=1$，否则为$0$。
   • 如果出现得分相同的情况，优先选择索引$j$较小的专家。

2. 动态容量掩码调度：按顺序（从$Token$ $0$到Token $N-1$）依次处理每个$Token$。维护一个长度为$E$的计数器数组$load$，初始全为$0$。遍历初始掩码$M^{(1)}$，生成最终的生效掩码$M^{(2)}$：

   • 如果$M_{i,j}^{(1)}==1$，且此时$load[j]$小于$C$，则该分配生效，$M_{i,j}^{(2)}=1$，$load[j]$加一。
   • 如果$M_{i,j}^{(1)}==1$，且此时$load[j]$已经等于$C$，则触发容量掩码，强制丢弃，$M_{i,j}^{(2)}=0$。
   • 对于同一个$Token$，被选中的$K$个专家的容量判断是相互独立的，均与当前Token处理开始时刻的$load$数组进行比较。

3. 路由惩罚聚合：根据最终的每个专家的实际负载数组$load$，计算当前批次的路由不平衡惩罚值$P$。公式为各专家实际负载的平方和：

   $P=\sum_{j=0}^{E-1}(load[j])^2$

输入描述

第1行：$N\ E\ K\ C$四个整数，以空格分隔，分别为$Token$数量、专家数量、每个$Token$激活的专家数、单个专家最大容量。

接下来$N$行：每行包含$E$个整数，以空格分隔，代表原始打分矩阵$S$。

约束条件：

$1 \le N \le 10^4$

$1 \le E \le 100$

$1 \le K \le E$

$0 \le C \le N$

$0 \le S_{i,j} \le 10^4$

输出描述

第1行：一个整数，代表最终的路由惩罚值$P$。

第2行：$E$个整数，以空格分隔，代表每个专家最终处理的$Token$数量。

样例1

输入：

4 3 2 2
1 5 4
8 1 2
3 6 5
2 7 9

输出：

9
1 2 2

解释：

1. Top-K路由：

   Token0，选专家1, 2

   Token1，选专家0, 2

   Token2，选专家1, 2

   Token3，选专家1, 2

2. 容量掩码调度：

   Token0：专家1, 2容量均充足。load变为[0, 1, 1]

   Token1：专家0, 2容量均充足。load变为[1, 1, 2]

   Token2：专家1充足，专家2超载。load变为[1, 2, 2]

   Token3：专家1, 2均超载。load保持[1, 2, 2]

3. 惩罚聚合：

   最终load数组为[1, 2, 2]

   惩罚值为load数组中数值的平方和为9

样例2

输入：

4 4 2 2
5 5 5 1
2 8 8 9
1 2 9 9
7 7 1 1

输出：

13
2 2 1 2

解释：

1. Top-K路由：

   Token0，前三名得分同为5。根据“优先选择索引较小”规则，因此选专家0, 1

   Token1，最高分专家3；第二高分专家1, 2得分同为8，优先选较小的专家1。因此选专家1, 3

   Token2，选专家2, 3

   Token3，选专家0, 1

2. 容量掩码调度：

   Token0：专家0, 1容量均充足。load变为[1, 1, 0, 0]

   Token1：专家1, 3容量均充足。load变为[1, 2, 0, 1]

   Token2：专家2, 3容量均充足。load变为[1, 2, 1, 2]

   Token3：专家0充足，专家1超载。load变为[2, 2, 1, 2]

3. 惩罚聚合：

   最终load数组为[2, 2, 1, 2]

   惩罚值为load数组中数值的平方和为13