Sigmoid Linear Unit

原始题目：LeetGPU - Sigmoid Linear Unit

题目描述

编写一个 GPU 程序，对一维输入向量实现 SiLU（Sigmoid Linear Unit，也称 Swish）激活函数的前向传播。SiLU 定义为：

$$\sigma(x) = \frac{1}{1 + e^{-x}}$$

$$\text{SiLU}(x) = x \cdot \sigma(x)$$

实现要求

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

示例

示例 1

Input:  input = [0.5, 1.0, -0.5]  (N=3)
Output: [0.3112295, 0.731059, -0.1887705]

示例 2

Input:  input = [-1.0, -2.0, -3.0, -4.0, -5.0]  (N=5)
Output: [-0.26894143, -0.23840584, -0.14227763, -0.07194484, -0.03346425]

约束条件

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

解题思路

SiLU 在 LLaMA 等现代大模型中广泛使用。每个元素需要计算一次 sigmoid 和一次乘法。sigmoid 中的指数运算 $e^{-x}$ 是性能瓶颈，可以使用 __expf() 内置函数或查表法加速。注意：对于大的负输入，$e^{-x}$ 可能溢出到 $+\infty$，需要用条件判断处理。

代码实现

CUDA

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

__global__ void silu_kernel(const float* input, float* output, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        float x = input[i];
        output[i] = x / (1.0f + expf(-x));  // SiLU(x) = x * sigmoid(x)
    }
}

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

Triton

import triton
import triton.language as tl

@triton.jit
def silu_kernel(input_ptr, output_ptr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N
    x = tl.load(input_ptr + idx, mask=mask)
    tl.store(output_ptr + idx, x * tl.sigmoid(x), mask=mask)