   CHECK_EQ(gamma_grad.CheckContiguous(), true);
   CHECK_EQ(beta_grad.CheckContiguous(), true);
   dim3 part_grad_block_dim, part_grad_grid_dim, gb_block_dim, gb_grid_dim;
-  int row_repeat, npart;
+  int npart;
   GetGammaBetaGradKernelParams(nbatch, nchannel, &part_grad_block_dim, &part_grad_grid_dim,
-                               &gb_block_dim, &gb_grid_dim, &row_repeat, &npart);
-  if(gamma_grad_req != kNullOp || beta_grad_req != kNullOp) {
+                               &gb_block_dim, &gb_grid_dim, &npart);
+  if (gamma_grad_req != kNullOp || beta_grad_req != kNullOp) {
     MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, DType, {
-      Tensor<gpu, 1, DType> workspace =
-        ctx.requested[0].get_space_typed<gpu, 1, DType>(Shape1(2 * npart * nchannel), s);
-      DType* part_gamma_grad_ptr = workspace.dptr_;
-      DType* part_beta_grad_ptr = workspace.dptr_ + npart * nchannel;
+      typedef std::conditional<std::is_same<DType, mshadow::half::half_t>::value,
+                               float, DType>::type AType;
+      Tensor<gpu, 1, AType> workspace =
+        ctx.requested[0].get_space_typed<gpu, 1, AType>(Shape1(2 * npart * nchannel), s);
+      AType* part_gamma_grad_ptr = workspace.dptr_;
+      AType* part_beta_grad_ptr = workspace.dptr_ + npart * nchannel;
       const int nshared_K1 = 2 * (part_grad_block_dim.x + 1)
-                               * part_grad_block_dim.y * sizeof(DType);
-      const int nshared_K2 = 2 * gb_block_dim.x * gb_block_dim.y * sizeof(DType);
+                               * part_grad_block_dim.y * sizeof(AType);
+      const int nshared_K2 = 2 * gb_block_dim.x * gb_block_dim.y * sizeof(AType);
       DType* gamma_grad_ptr = (gamma_grad_req != kNullOp) ? gamma_grad.dptr<DType>() : nullptr;
       DType* beta_grad_ptr = (beta_grad_req != kNullOp) ? beta_grad.dptr<DType>() : nullptr;
       LayerNormFusedBackwardKernel_PartGammaBeta
         <<<part_grad_grid_dim, part_grad_block_dim, nshared_K1, stream>>>
-        (nbatch, nchannel, row_repeat, in_data.dptr<DType>(), out_grad.dptr<DType>(),
+        (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(),
          mean_data.dptr<DType>(), std_data.dptr<DType>(), part_gamma_grad_ptr, part_beta_grad_ptr);
       MSHADOW_CUDA_POST_KERNEL_CHECK(LayerNormFusedBackwardKernel_PartGammaBeta);
-      if(gamma_grad_req == kAddTo && beta_grad_req != kAddTo) {
+      if (gamma_grad_req == kAddTo && beta_grad_req != kAddTo) {
         LayerNormFusedBackwardKernel_GammaBeta<true, false>
           <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
           (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,
            gamma_grad_ptr, beta_grad_ptr);
-      } else if(gamma_grad_req != kAddTo && beta_grad_req == kAddTo) {
+      } else if (gamma_grad_req != kAddTo && beta_grad_req == kAddTo) {
         LayerNormFusedBackwardKernel_GammaBeta<false, true>
           <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
           (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,
             gamma_grad_ptr, beta_grad_ptr);
-      } else if(gamma_grad_req == kAddTo && beta_grad_req == kAddTo) {
+      } else if (gamma_grad_req == kAddTo && beta_grad_req == kAddTo) {
         LayerNormFusedBackwardKernel_GammaBeta<true, true>
           <<<gb_grid_dim, gb_block_dim, nshared_K2, stream>>>
           (nbatch, nchannel, npart, part_gamma_grad_ptr, part_beta_grad_ptr,

  }
  const dim3 data_block_dim(32, nthread_y);
  const int LOAD_UNROLL = 4;
  if(data_grad_req != kNullOp) {
    MSHADOW_REAL_TYPE_SWITCH(in_data.type_flag_, DType, {
      int nshared = data_block_dim.y > 1 ? data_block_dim.y * data_block_dim.x * sizeof(DType) : 0;
      CheckLaunchParam(data_grid_dim, data_block_dim);
      if(data_grad_req == kAddTo) {
        LayerNormFusedBackwardKernel_Data<LOAD_UNROLL, true>
          <<<data_grid_dim, data_block_dim, nshared, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma.dptr<DType>(), data_grad.dptr<DType>());
      } else {
        LayerNormFusedBackwardKernel_Data<LOAD_UNROLL, false>
          <<<data_grid_dim, data_block_dim, nshared, stream>>>
          (nbatch, nchannel, in_data.dptr<DType>(), out_grad.dptr<DType>(), mean_data.dptr<DType>(),
           std_data.dptr<DType>(), gamma.dptr<DType>(), data_grad.dptr<DType>());

}

template <typename DType>
__device__ __forceinline__ DType warp_shfl_xor(DType value, int laneMask,
                                               int width = 32, unsigned int mask = 0xffffffff) {
#if CUDA_VERSION >= 9000
  return __shfl_xor_sync(mask, value, laneMask, width);
