     DType* out_col_val = out_data + bid * nchannel;
 
     if (gamma != NULL && beta != NULL) {
-      for (int i = tid; i < nchannel; i += nthread) {
-        out_col_val[i] = gamma[i] * invstd_eps * (col_vals[i] - mean) + beta[i];
+      int l = 4 * tid;
+      for (; l + 3 < nchannel; l += 4 * nthread) {
+        for (int i = 0; i < 4; ++i) {
+          out_col_val[l + i] = gamma[l + i] * invstd_eps * (col_vals[l + i] - mean) + beta[l + i];
+        }
+      }
+      for(; l < nchannel; l++) {
+        out_col_val[l] = gamma[l] * invstd_eps * (col_vals[l] - mean) + beta[l];
       }
     } else if (gamma == NULL && beta != NULL) {
       for (int i = tid; i < nchannel; i += nthread) {

  DType count = 0;
  DType mean = 0;
  DType sigma2 = 0;
  // const int N_ACCUM = 4;  // TODO(sxjscience) Profile
  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;
    // Each thread takes charge of 4 consecutive numbers
    // To minimize branch divergence, we split the for-loop into two parts.
    _block_welford_online_sum(tid, nthread, col_vals, nchannel, mean, sigma2, count);

    // Merge the mean/sigma2 within a warp

  for(; l < nchannel; ++l) {
    welford_online_sum_step(col_vals[l], mean, sigma2, count);
  }
}

/* Fused CUDA kernel for layer normalization. It computes the LayerNorm when axis=-1.
 * Shape of the input tensors:
