   for(; l < nchannel; ++l) {
     welford_online_sum_step(col_vals[l], mean, sigma2, count);
   }
-  /*
-  for (int l = tid; l < nchannel; l += nthread) {
-    welford_online_sum_step(col_vals[l], mean, sigma2, count);
-  }
-  */
-}
-
-/*
-template<>
-__device__ __forceinline__ void _block_welford_online_sum(const int tid,
-                                                          const int nthread,
-                                                          const float* __restrict__ col_vals,
-                                                          const int nchannel,
-                                                          float& mean,
-                                                          float& sigma2,
-                                                          float& count) {
-  // Calculate the shift which makes sure that the `col_vals + shift` is aligned to float4
-  //  float (4 bytes) --> float4 (16 bytes)
-  // The address of the pointer is `addr = static_cast<size_t>(col_vals)`
-  // The shift is (4 - (addr >> 2) & 3) & 3
-  size_t addr = reinterpret_cast<std::size_t>(col_vals);
-  int alignment_shift = (4 - (addr >> 2) & 3) & 3;
-  // 1) Shift the elements to make sure that the pointer is aligned to float4
-  for(int i = tid; i < alignment_shift && i < nchannel; i += nthread) {
-    welford_online_sum_step(col_vals[tid], mean, sigma2, count);
-  }
-  // 2) Use float4 to load the middle part of the input columns.
-  //  alignment (float), middle (divisible by 4, float4), rest elements (float)
-  const float4* col_vals_float4 = reinterpret_cast<const float4*>(col_vals + alignment_shift);
-  int mid_length = (nchannel > alignment_shift) ? (nchannel - alignment_shift) >> 2 : 0;
-  for (int i = tid; i < mid_length; i += nthread) {
-    float4 vec_vals = col_vals_float4[i];
-    welford_online_sum_step(vec_vals.x, mean, sigma2, count);
-    welford_online_sum_step(vec_vals.y, mean, sigma2, count);
-    welford_online_sum_step(vec_vals.z, mean, sigma2, count);
-    welford_online_sum_step(vec_vals.w, mean, sigma2, count);
-  }
-  // 3) Handling the remanining values
-  for(int i = alignment_shift + 4 * mid_length + tid; i < nchannel; i += nthread) {
-    welford_online_sum_step(col_vals[i], mean, sigma2, count);
-  }
-}
-*/
+}
 
 /* Fused CUDA kernel for layer normalization. It computes the LayerNorm when axis=-1.
  * Shape of the input tensors:

                                                          DType& mean,
                                                          DType& sigma2,
                                                          DType& count) {
  int l = 4 * tid;
  for (; l + 3 < nchannel; l += 4 * nthread) {
    for (int i = 0; i < 4; ++i) {
      welford_online_sum_step(col_vals[l + i], mean, sigma2, count);
    }

  DType count = 0;
  DType mean = 0;
  DType sigma2 = 0;
  extern __shared__ char buf[];  // Shared memory

  if (bid < nbatch) {
    int tid = threadIdx.x + threadIdx.y * blockDim.x;
    const DType* col_vals = in_data + bid * nchannel;
    _block_welford_online_sum(tid, nthread, col_vals, nchannel, mean, sigma2, count);

    // Merge the mean/sigma2 within a warp
