ops(5)：激活函数与残差连接的 CUDA 实现

作者: 紫气东来
发布时间: 编辑于 2024-05-11 11:28
原文链接: https://zhuanlan.zhihu.com/p/695703671

注意: 由于技术限制无法完整抓取原文，本文档基于该系列其他文章和CUDA编程知识进行扩展整理。

概述

激活函数和残差连接是深度学习模型中的基础组件，都属于 element-wise 操作。本文探讨如何在CUDA中高效实现这些操作，包括常见激活函数（ReLU、GELU、SiLU等）的优化实现以及算子融合技术。

一、常见激活函数的 CUDA 实现

1.1 ReLU (Rectified Linear Unit)

数学定义:

ReLU(x) = max(0, x)

特点:

• 最简单的激活函数
• 计算开销极小
• 存在"神经元死亡"问题

CPU 基准实现:

void relu_forward_cpu(float* out, const float* inp, int N) {
    for (int i = 0; i < N; i++) {
        out[i] = inp[i] > 0.0f ? inp[i] : 0.0f;
    }
}

CUDA 基础实现:

__global__ void relu_forward_kernel(float* out, const float* inp, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        out[idx] = inp[idx] > 0.0f ? inp[idx] : 0.0f;
    }
}

反向传播:

__global__ void relu_backward_kernel(float* dinp, const float* dout, 
                                     const float* inp, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        dinp[idx] = inp[idx] > 0.0f ? dout[idx] : 0.0f;
    }
}

1.2 GELU (Gaussian Error Linear Unit)

数学定义:

GELU(x) = x * Φ(x) = x * P(X ≤ x), where X ~ N(0,1)
其中 Φ(x) 是标准正态分布的累积分布函数

GELU 可以使用高斯误差函数进行精确计算：

GELU(x) = 0.5 * x * (1 + erf(x/√2))

Tanh 近似公式（更常用）:

GELU(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))

CUDA 实现:

__global__ void gelu_forward_kernel(float* out, const float* inp, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float x = inp[idx];
        float cube = 0.044715f * x * x * x;
        float tanh_arg = sqrtf(2.0f / M_PI) * (x + cube);
        out[idx] = 0.5f * x * (1.0f + tanhf(tanh_arg));
    }
}

反向传播:

__global__ void gelu_backward_kernel(float* dinp, const float* dout,
                                     const float* inp, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float x = inp[idx];
        float cube = 0.044715f * x * x * x;
        float tanh_arg = sqrtf(2.0f / M_PI) * (x + cube);
        float tanh_out = tanhf(tanh_arg);
        float coshf_out = coshf(tanh_arg);
        float sech_out = 1.0f / (coshf_out * coshf_out);
        
        float local_grad = 0.5f * (1.0f + tanh_out) + 
                          0.5f * x * sech_out * sqrtf(2.0f / M_PI) * 
                          (1.0f + 3.0f * 0.044715f * x * x);
        dinp[idx] = local_grad * dout[idx];
    }
}

1.3 SiLU / Swish (Sigmoid Linear Unit)

数学定义:

SiLU(x) = x * σ(x) = x / (1 + e^(-x))

CUDA 实现:

__global__ void silu_forward_kernel(float* out, const float* inp, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float x = inp[idx];
        float sigmoid = 1.0f / (1.0f + expf(-x));
        out[idx] = x * sigmoid;
    }
}

二、残差连接的 CUDA 实现

2.1 基本残差连接

数学定义: output = input + residual

CUDA 实现:

__global__ void residual_forward_kernel(float* out, const float* inp,
                                        const float* residual, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        out[idx] = inp[idx] + residual[idx];
    }
}

三、算子融合优化

3.1 融合激活函数与残差连接

融合 GELU + Residual:

__global__ void fused_residual_gelu_kernel(float* out, const float* inp,
                                           const float* residual, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float x = inp[idx] + residual[idx];
        float cube = 0.044715f * x * x * x;
        float tanh_arg = sqrtf(2.0f / M_PI) * (x + cube);
        out[idx] = 0.5f * x * (1.0f + tanhf(tanh_arg));
    }
}

四、性能对比

4.1 不同激活函数性能（A100, 100M元素）

4.2 算子融合性能提升

五、最佳实践

1. 选择合适的激活函数: Transformer 推荐 GELU 或 SiLU
2. 算子融合: 融合相邻 element-wise 操作减少内存访问
3. 向量化访存: 使用 float4 提高带宽利用率
4. 数值稳定性: 验证 fast math 函数的精度

参考资料

1. llm.c - Karpathy's CUDA implementations
2. NVIDIA CUDA C Programming Guide
3. GELU 论文

总结: 激活函数和残差连接通过算子融合、向量化访存等优化技术可显著提升性能。