Swish-Gated Linear Unit

原始题目：LeetGPU - Swish-Gated Linear Unit

题目描述

编写一个 GPU 程序，对一维输入向量实现 SwiGLU（Swish-Gated Linear Unit）激活函数的前向传播。输入张量形状为 $[N]$，元素类型为 float32。

SWiGLU 的定义如下：

将输入 $x$ 从中间分成两半 $x_1$ 和 $x_2$。对前半部分计算 SiLU：

$$\text{SiLU}(x_1) = x_1 \cdot \sigma(x_1), \quad \sigma(x) = \frac{1}{1 + e^{-x}}$$

SwiGLU 输出为：

$$\text{SwiGLU}(x_1, x_2) = \text{SiLU}(x_1) \cdot x_2$$

输出张量的长度为 $N/2$。

实现要求

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

示例

示例 1

Input:  [1.0, 2.0, 3.0, 4.0]  (N=4)
Output: [2.1931758, 7.0463767]

示例 2

Input:  [0.5, 1.0]  (N=2)
Output: [0.31122968]

约束条件

• $1 \le N \le 100{,}000$，$N$ 为偶数。
• $-100.0 \le$ 输入值 $\le 100.0$。
• 性能测试在 $N = 100{,}000$ 的规模下进行。

解题思路

SwiGLU 是 LLaMA、Mistral、Gemma 等现代大模型使用的 FFN 激活函数。核心计算是分半后的 SiLU 门控 + 逐元素乘法。输入的两半在内存中连续，每个线程只需从相邻位置读取两部分，访存模式友好。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>

__global__ void swiglu_kernel(const float* input, float* output, int halfN) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < halfN) {
        float x1 = input[i];
        float x2 = input[i + halfN];
        float silu = x1 / (1.0f + expf(-x1));  // SiLU(x1)
        output[i] = silu * x2;                   // gate * x2
    }
}

extern "C" void solve(const float* input, float* output, int N) {
    int halfN = N / 2;
    int threadsPerBlock = 256;
    int blocksPerGrid = (halfN + threadsPerBlock - 1) / threadsPerBlock;
    swiglu_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, halfN);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl

@triton.jit
def swiglu_kernel(input_ptr, output_ptr, halfN: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < halfN
    x1 = tl.load(input_ptr + idx, mask=mask)
    x2 = tl.load(input_ptr + idx + halfN, mask=mask)
    tl.store(output_ptr + idx, tl.sigmoid(x1) * x1 * x2, mask=mask)