MoE Top-K Gating

原始题目：LeetGPU - MoE Top-K Gating

题目描述

编写一个 GPU 程序，实现混合专家模型（MoE）的 Top-K 门控。给定形状为 $[M, E]$ 的 logit 矩阵（$M$ 为 token 数，$E$ 为专家数），对每行找出 $k$ 个最大值，提取其索引，并应用 softmax 得到混合权重。

对于每一行 $i$：

$$\begin{aligned} indices_i, vals_i &= \text{TopK}(logits_i, k) \\ vals_i &= logits_i[indices_i] \\ weights_i &= \text{Softmax}(vals_i) \end{aligned}$$

选出的专家必须按 logit 值降序排列，topk_weights 必须按位置对应 topk_indices 的顺序。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 topk_weights 和 topk_indices 数组中。

示例

Input:  logits = [[1,2,3,4],[4,3,2,1]], M=2, E=4, k=2
Output: topk_weights = [[0.7311, 0.2689],[0.7311, 0.2689]]
        topk_indices  = [[3, 2],[0, 1]]
        (Row 0: Top-2 = [4,3] at idx [3,2], softmax = [0.7311,0.2689])
        (Row 1: Top-2 = [4,3] at idx [0,1], softmax = [0.7311,0.2689])

约束条件

• $1 \le M \le 10{,}000$, $1 \le E \le 256$, $1 \le k \le E$。
• 所有张量存储在 GPU 上，logits 为 32 位浮点数，indices 为 32 位整数。
• 性能测试在 $M = 1{,}024,\ k = 2$ 的规模下进行。

解题思路

MoE 门控的瓶颈在 Top-K 选择——每行要从 $E$ 个专家中选出 $k$ 个最大值。当 $E$ 较小（$\le 256$）时，每行可以用一个 warp 在寄存器或共享内存中维护一个大小为 $k$ 的最小堆。Softmax 只在选出的 $k$ 个值上计算（而非全部 $E$ 个），开销较小。对于大 $M$，充分利用所有 SM 并行处理不同 token 是关键。欢迎在 GitHub Discussions 分享你的解法。