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以防止溢出。