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 是关键。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void moe_gate(const float* logits, float* weights, int* indices, int M, int E, int k) {
    int row=blockIdx.x; if(row>=M)return;
    float vals[256]; int idxs[256];
    for(int e=0;e<E;e++){vals[e]=logits[row*E+e];idxs[e]=e;}
    // Select top-k (partial sort)
    for(int i=0;i<k;i++){
        int mi=i;
        for(int j=i+1;j<E;j++) if(vals[j]>vals[mi])mi=j;
        float tv=vals[i];vals[i]=vals[mi];vals[mi]=tv;
        int ti=idxs[i];idxs[i]=idxs[mi];idxs[mi]=ti;
    }
    // Softmax on top-k
    float mx=vals[0],sm=0.0f;
    for(int i=0;i<k;i++)mx=fmaxf(mx,vals[i]);
    for(int i=0;i<k;i++){vals[i]=expf(vals[i]-mx);sm+=vals[i];}
    for(int i=0;i<k;i++){weights[row*k+i]=vals[i]/sm;indices[row*k+i]=idxs[i];}
}
extern "C" void solve(const float* logits, float* weights, int* indices, int M, int E, int k) {
    moe_gate<<<M,1>>>(logits,weights,indices,M,E,k);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def moe_gate(logits_ptr,weights_ptr,indices_ptr, M,E,k, BLOCK:tl.constexpr):
    row=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=row<M
    e_range=tl.arange(0,E)
    vals=tl.load(logits_ptr+row[:,None]*E+e_range[None,:],mask=mask[:,None])
    topk_vals,topk_idx=tl.sort(vals,descending=True)
    topk_vals=topk_vals[:,:k]; topk_idx=topk_idx[:,:k]
    w=tl.softmax(topk_vals)
    k_range=tl.arange(0,k)
    tl.store(weights_ptr+row[:,None]*k+k_range[None,:],w,mask=mask[:,None])
    tl.store(indices_ptr+row[:,None]*k+k_range[None,:],topk_idx,mask=mask[:,None])