SwiGLU MLP Block

原始题目：LeetGPU - SwiGLU MLP Block

题目描述

实现 SwiGLU MLP 块——LLaMA、Mistral、Gemma 及大多数现代大语言模型中使用的前馈网络。给定输入矩阵 $x$（$[M, d_{model}]$）和三个权重矩阵 $W_{gate}$（$[d_{model}, d_{ffn}]$）、$W_{up}$（$[d_{model}, d_{ffn}]$）和 $W_{down}$（$[d_{ffn}, d_{model}]$），计算：

$$output = \left(\text{SiLU}(x \times W_{gate}) \odot (x \times W_{up})\right) \times W_{down}$$

其中 $\text{SiLU}(z) = z \cdot \sigma(z)$（sigmoid 门控），$\odot$ 表示逐元素乘法。所有张量为 float32。

实现要求

• 实现 solve 函数，签名保持不变。
• 不允许使用外部库。
• 将结果写入 output。

示例

Input:  M=2, d_model=2, d_ffn=4
        x = [[1,0],[0,1]], W_gate=W_up = [[1,0,0,0],[0,1,0,0]], W_down = [[1,0],[0,1],[0,0],[0,0]]
        gate = x @ W_gate = [[1,0,0,0],[0,1,0,0]]
        SiLU(gate) ≈ [[0.7311,0,0,0],[0,0.7311,0,0]]
        hidden = SiLU(gate) ⊙ up ≈ [[0.7311,0,0,0],[0,0.7311,0,0]]
Output: hidden @ W_down ≈ [[0.7311,0],[0,0.7311]]

约束条件

• $1 \le M \le 65{,}536$，$1 \le d_{model} \le 8{,}192$，$1 \le d_{ffn} \le 32{,}768$。
• 性能测试在 $M = 512,\ d_{model} = 4{,}096,\ d_{ffn} = 14{,}336$ 下进行。

解题思路

SwiGLU MLP 包含两次矩阵乘法（gate 和 up projection）和一次下游投影。gate 和 up 的投影可以合并为一个矩阵乘法（拼在一起），然后 split+SiLU+multiply，减少 kernel launch 次数。SiLU 是逐元素操作，可以融合到上一个或下一个 kernel 中。$W_{down}$ 投影是最昂贵的部分（$d_{ffn}$ 通常远大于 $d_{model}$）。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void swiglu_mlp(const float* x, const float* Wg, const float* Wu, const float* Wd,
    float* out, int M, int D, int F) {
    int m=blockIdx.x, d=threadIdx.x+blockIdx.y*blockDim.x;
    if(m<M && d<D) {
        // gate = SiLU(x @ Wg); up = x @ Wu; hidden = gate * up; out = hidden @ Wd
        float sum_g=0,sum_u=0;
        for(int f=0;f<F;f++){float xf=x[m*D+f];sum_g+=xf*Wg[f*D+d];sum_u+=xf*Wu[f*D+d];}
        float gate = sum_g/(1.0f+expf(-sum_g)); // SiLU
        out[m*D+d] = gate * sum_u;
    }
}
extern "C" void solve(const float* x, const float* Wg, const float* Wu, const float* Wd,
    float* out, int M, int D, int F) {
    dim3 t(16,4), b(M,(D+15)/16);
    swiglu_mlp<<<b,t>>>(x,Wg,Wu,Wd,out,M,D,F); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def swiglu_mlp(x_ptr,Wg_ptr,Wu_ptr,Wd_ptr,out_ptr, M,D,F, BLOCK:tl.constexpr):
    m=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=m<M
    f_range=tl.arange(0,F)
    xf=tl.load(x_ptr+m[:,None]*D+f_range[None,:],mask=mask[:,None])
    gate=tl.sum(xf*tl.load(Wg_ptr+f_range[:,None]*D+tl.arange(0,D)[None,:]),axis=1)
    up=tl.sum(xf*tl.load(Wu_ptr+f_range[:,None]*D+tl.arange(0,D)[None,:]),axis=1)
    hidden=tl.sigmoid(gate)*gate*up  # SiLU(x) * up
    tl.store(out_ptr+m[:,None]*D+tl.arange(0,D)[None,:],hidden@tl.load(Wd_ptr).T,mask=mask[:,None])