       }
       break;
     case 7:  // d_x != nullptr, d_scale != nullptr, d_bias != nullptr
-      switch (block_dim) {
-        FIXED_BLOCK_DIM_FIXED_BLOCK_NUM_CASE(
-            feature_size, kMaxBlockNum,
-            LayerNormBackwardGradientAll<
-                T, U, kBlockDim, true><<<block_num, kBlockDim, 0, stream>>>(
-                x, d_y, d_scale, d_bias, d_x, mean, var, scale, epsilon,
-                batch_size, feature_size, col_offset));
-      }
-      switch (GetDesiredBlockDim(feature_size)) {
-        FIXED_BLOCK_DIM_CASE(
-            LayerNormBackwardPostProcessToCalculateDX<
-                T, U, kBlockDim><<<batch_size, kBlockDim, 0, stream>>>(
-                x, d_x, mean, var, epsilon, feature_size));
-      }
+    {
+      constexpr int VPT = 4;
+      constexpr int BDIMX2 = 32;
+      constexpr int BDIMY2 = 4;
+      dim3 threads2(BDIMX2, BDIMY2, 1);
+      constexpr int part_size = BDIMY2 * VPT;
+      const dim3 blocks2((feature_size + BDIMX2 - 1) / BDIMX2, part_size, 1);
+
+      auto part_grad_gamma_ptr =
+          memory::Alloc(dev_ctx, part_size * feature_size * sizeof(U));
+      auto part_grad_beta_ptr =
+          memory::Alloc(dev_ctx, part_size * feature_size * sizeof(U));
+      U *part_grad_gamma = reinterpret_cast<U *>(part_grad_gamma_ptr->ptr());
+      U *part_grad_beta = reinterpret_cast<U *>(part_grad_beta_ptr->ptr());
+
+      LayerNormBackwardPartGradGammaBeta<T, U, BDIMX2, BDIMY2,
+                                         VPT><<<blocks2, threads2, 0, stream>>>(
+          d_y, x, batch_size, feature_size, mean, var, epsilon, part_grad_gamma,
+          part_grad_beta);  // compute part_grad_gamma, beta
+
+      constexpr int BDIMX3 = 32;
+      constexpr int BDIMY3 = 8;
+      dim3 threads3(BDIMX3, BDIMY3, 1);
+      const dim3 blocks3((feature_size + BDIMX2 - 1) / BDIMX2, 1, 1);
+      LayerNormBackwardSumGradGammaBeta<
+          T, U, BDIMX3, BDIMY3><<<blocks3, threads3, 0, stream>>>(
+          part_grad_gamma, part_grad_beta, part_size, batch_size, feature_size,
+          d_scale, d_bias);
+
+      constexpr int BDIMX1 = 32;
+      constexpr int BDIMY1 = 4;
+      dim3 threads1(BDIMX1, BDIMY1, 1);
+      const dim3 blocks1(1, batch_size, 1);
+      LayerNormBackwardComputeGradInput<
+          T, U, BDIMX1, BDIMY1><<<blocks1, threads1, 0, stream>>>(
+          d_y, x, batch_size, feature_size, mean, var, epsilon, scale, d_x);
       break;
+    }
     default:
       break;
   }

static void LayerNormBackward(const T *x, const T *d_y, const U *scale,
                              const U *mean, const U *var, T *d_x, U *d_scale,
                              U *d_bias, float epsilon, int batch_size,
                              int feature_size, cudaStream_t stream) {
  const int kMaxBlockDim = 512;
  const int kMaxBlockNum = 128;
  int gradient_flag = ((d_x != nullptr ? 1 : 0) << 2) |

    int batch_size = static_cast<int>(matrix_dim[0]);
    int feature_size = static_cast<int>(matrix_dim[1]);

    LayerNormBackward<T, U>(x_data, d_y_data, scale_data, mean_data, var_data,
                            d_x_data, d_scale_data, d_bias_data, epsilon,
                            batch_size, feature_size, ctx);
  }
};

