Gaussian Error Gated Linear Unit

原始题目：LeetGPU - Gaussian Error Gated Linear Unit

题目描述

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

GEGLU 的定义如下：

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

$$\text{GELU}(x_2) = \frac{1}{2} x_2\left(1 + \text{erf}\left(\frac{x_2}{\sqrt{2}}\right)\right)$$

GEGLU 输出为：

$$\text{GEGLU}(x_1, x_2) = x_1 \cdot \text{GELU}(x_2)$$

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

实现要求

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

示例

示例 1

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

示例 2

Input:  [2.0, -1.0, 1.0, 0.5]  (N=4)
Output: [1.6826895, -0.3457312]

约束条件

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

解题思路

GEGLU 是 GLU 家族中使用 GELU 作为门控函数的变体。erff() 在 GPU 上是软件实现，延迟较高但精度为 IEEE 标准。题目明确规定使用 erf 定义，不可用 tanh 近似替换。输入的前后两半在内存中连续，访问模式友好。

代码实现

CUDA

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

__device__ const float SQRT_2 = 1.41421356237309504880f;

// 解法一：使用 erff() 的精确 GELU（题目要求的实现）
// 注意：erff() 是软件实现，延迟较高但精度为 IEEE 标准
__global__ void geglu_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 gelu = 0.5f * x2 * (1.0f + erff(x2 / SQRT_2));
        output[i] = x1 * gelu;
    }
}

// 解法二：Grid-Stride Loop
__global__ void geglu_grid_stride(const float* input, float* output, int halfN) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < halfN; i += stride) {
        float x1 = input[i];
        float x2 = input[i + halfN];
        float gelu = 0.5f * x2 * (1.0f + erff(x2 / SQRT_2));
        output[i] = x1 * gelu;
    }
}

// 注：实际工程中 GELU 常用 tanh 近似以获得更好的 GPU 性能：
//   GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
// 但本题明确要求使用 erf 定义，不可替换。

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

Triton

import triton
import triton.language as tl
import torch


@triton.jit
def geglu_kernel(
    input_ptr, output_ptr,
    halfN: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    idx = pid * 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)

    # GELU via erf (exact, slow): 0.5 * x * (1 + erf(x / sqrt(2)))
    SQRT2 = 1.4142135623730951
    gelu = 0.5 * x2 * (1.0 + tl.math.erf(x2 / SQRT2))
    tl.store(output_ptr + idx, x1 * gelu, mask=mask)