RMSNorm kernel

1 Background

kernel表述：

RMSNorm — PyTorch 2.13 documentation

其中，$x$, $y$ 对应输入输出行向量，$\gamma$对应权重向量。

而实际上输入的$x$的shape应该为..., hidden_dim，其中...表示batch_like shape。该计算仅在最后一个维度上进行。在本问题中定义$x$的shape为rows hidden_dim

2 Analytics & Implementation

Abstraction

可以拆解成并行规约与逐元素乘法

1. 并行规约$RMS(x)$

   对每行采用并行规约算法计算出平方和，然后在合适的时机取均值加上$\epsilon$开根号获取$RMS$值。

2. 逐元素乘法$x_{i} * \gamma_i * \frac{1}{RMS(x)}$

   第二步进行并行逐元素乘法即可。

Arithmentic Intensity

假设数据类型为float32，我们以行为粒度进行计算。定义d = hidden_dim。

For every row, Load line of x, and $\gamma$, which is 2 * d elements.

Write one line of y, which is 1 * d elements.

For one row It has approximately 2 * d(rms) + 2 * d(per-element-multiply) FP operations.

对于并行规约，可以把每行当做一个BLOCK，进行基于warp_shuffle的规约。

由于还不知道输入的形状分布如何，认为是随机分布的，因此暂时还没有一个fine-grained的实现。

template <typename T>
/**
 * @brief fused kernel of RMSNorm
 * 每个BLOCK接管一行，行内通过warp_reduce进行求和
 * 然后进行后续计算
 * @tparam T float
 * @param output \sum {x_i ^ 2} upcast to float
 * @param input input x
 * @param weight weight
 * @param eps epsilon of RMSNorm
 * @param rows rows
 * @param hidden_dim hidden dimension
 * @param n element count of input , in this case its hidden_dim
 * @return __global__ 
 */
__global__ void fused_RMSNorm_kernel(T *output, const T *input, const T *weight, float eps, size_t rows, size_t hidden_dim) {
    extern __shared__ float smem[];
    __shared__ float multiplier;
    size_t tid = threadIdx.x;
    int rowStartIdx = blockIdx.x * hidden_dim;

    float sum = 0;
    for (size_t i = tid; i < hidden_dim; i += blockDim.x) 
    {
        int idx = rowStartIdx + i;
        sum += (float)input[idx] * (float)input[idx];
    }

    float warp_sum = warp_reduce(sum);

    if (tid % 32 == 0) {
        smem[tid / 32] = warp_sum;
    }
    __syncthreads();

    if (tid < 32) {
        float block_sum = (tid < (blockDim.x + 31) / 32) ? smem[tid] : float(0);

        block_sum = warp_reduce(block_sum);
        if(tid == 0) {
            multiplier = rsqrt(block_sum / hidden_dim + eps);
        }
    }

    __syncthreads();

    #pragma unroll
    for (int i = tid; i < hidden_dim; i += blockDim.x) {
        int idx = rowStartIdx + i;
        output[idx] = (T)((float)input[idx] * multiplier * (float)weight[i]);
    }
}

Performance

1. 耗时分布

   -------(512,1592)------------
   [float]CPU耗时：3728us
   [float]GPU耗时：4051us
   -------(512,1592)------------
   [half]CPU耗时：58268us
   [half]GPU耗时：2789us
   -------(512,2048)------------
   [float]CPU耗时：4960us
   [float]GPU耗时：2586us
   -------(512,2048)------------
   [half]CPU耗时：77099us
   [half]GPU耗时：1523us
   

   可以发现半精度确实由于数据加载时间缩短导致耗时降低。

   

可以发现在数据规模较小的时候，由于kernel launch等各种overhead，CPU的实现会比GPU的各种实现更加快。

在大规模的时候，GPU的并行带来的性能才逐渐显现出来，且使用half会比float更加efficiency。

不过读者需要注意到该naiive实现中，RMS的计算会upcast到float32以防止溢出。

3 Optimization

3.1 Vectorize

后面的点乘为逐元素运算，因此可以选择向量化进行加速。

向量化加速的原理其一为缩短指令条数，其二为增大了运算部件的利用率。类似于SIMD的用法。

template<typename T>
/**
 * @brief vectorized kernel
 * 每个BLOCK接管一行，行内通过warp_reduce进行求和
 * 然后进行后续计算
 * @tparam T float/half
 * @param output \sum {x_i ^ 2} upcast to float
 * @param input input x
 * @param weight weight
 * @param eps epsilon of RMSNorm
 * @param rows rows
 * @param hidden_dim hidden dimension
 * @param n element count of input , in this case its hidden_dim
 * @return __global__ 
 */
__global__ void vec_fused_RMSNorm_kernel(T *output, const T *input, const T *weight, float eps, size_t rows, size_t hidden_dim) {
    extern __shared__ float smem_vec[];
    __shared__ float multiplier;
    size_t tid = threadIdx.x;
    int rowStartIdx = blockIdx.x * hidden_dim;

    float sum = 0;
    for (size_t i = tid; i < hidden_dim; i += blockDim.x) 
    {
        int idx = rowStartIdx + i;
        sum += (float)input[idx] * (float)input[idx];
    }

    float warp_sum = warp_reduce(sum);

    if (tid % 32 == 0) { /* each warp, thread 0 takes sum of warp. */
        smem_vec[tid / 32] = warp_sum;
    }
    __syncthreads();

    if (tid < 32/* 32 is warpsize */) {
        float block_sum_float = (tid < (blockDim.x + 31) / 32) ? smem_vec[tid] : 0;

        block_sum_float = warp_reduce(block_sum_float);
        if(tid == 0) {
            multiplier = rsqrt(block_sum_float / hidden_dim + eps); // Assign row rms multiplier.
        }
    }

    __syncthreads();

    if constexpr (std::is_same_v<T, float>) {
        const float4 *vec_input = reinterpret_cast<const float4*>(input);
        const float4 *vec_weight = reinterpret_cast<const float4*>(weight);
        float4 *vec_output = reinterpret_cast<float4*>(output);
        int rowStartFloat4 = rowStartIdx / 4;
        #pragma unroll
        for (int i = tid; i < hidden_dim / 4; i += blockDim.x / 4) {
            int idx = rowStartFloat4 + i;
            vec_output[idx] = mul(mul(vec_input[idx], float4{multiplier,multiplier,multiplier,multiplier}), vec_weight[i]);
        }
    } else if constexpr (std::is_same_v<T, half>) {
        const half2 *vec_input = reinterpret_cast<const half2*>(input);
        const half2 *vec_weight = reinterpret_cast<const half2*>(weight);
        half2 *vec_output = reinterpret_cast<half2*>(output);
        int rowStartHalf2 = rowStartIdx / 2;
        #pragma unroll
        for (int i = tid; i < hidden_dim / 2; i += blockDim.x / 2) {
            int idx = rowStartHalf2 + i;
            vec_output[idx] = mul(mul(vec_input[idx], half2{(half)multiplier,(half)multiplier}), vec_weight[i]);
        }
    }
}

性能

-------(512,1592)------------
[float]CPU耗时：4115us
[float]GPU耗时：2211us
-------(512,1592)------------
[half]CPU耗时：54678us
[half]GPU耗时：1754us
-------(512,2048)------------
[float]CPU耗时：4696us
[float]GPU耗时：1818us
-------(512,2048)------------
[half]CPU耗时：69593us
[half]GPU耗时：2192us

从数据可以发现确实降低了相当的执行耗时。

从图上分析可知依旧是小batch时float CPU领先，General Case时GPU是一个trade-off选择。