 
 template <typename T, typename MT>
 __global__ void MomentumLarsKernel(
-    const T* __restrict__ param, const T* __restrict__ grad,
-    const MT* __restrict__ velocity, T* param_out, MT* velocity_out,
-    const MT* __restrict__ master_param, MT* __restrict__ master_param_out,
-    const MT* __restrict__ learning_rate, MT* __restrict__ p_buffer,
-    MT* __restrict__ g_buffer, const MT mu, const MT lars_coeff,
-    const MT lars_weight_decay, const MT epsilon, const MT rescale_grad,
-    const int repeat_times, const int thresh, const int64_t numel,
-    const bool is_amp) {
+    const T* param, const T* __restrict__ grad, const MT* velocity,
+    T* param_out, MT* velocity_out, const MT* master_param,
+    MT* master_param_out, const MT* __restrict__ learning_rate,
+    MT* __restrict__ p_buffer, MT* __restrict__ g_buffer, const MT mu,
+    const MT lars_coeff, const MT lars_weight_decay, const MT epsilon,
+    const MT rescale_grad, const int repeat_times, const int thresh,
+    const int64_t numel, const bool is_amp) {
   int tid = threadIdx.x + blockIdx.x * blockDim.x;
   int grid_stride = gridDim.x * LARS_BLOCK_SIZE;
 #if CUDA_VERSION >= 11000


template <typename T, typename MT, int VecSize, bool IsAmp = false>
__device__ inline void VectorizeLarsUpdate(
    const T* __restrict__ grad, const MT* __restrict__ param,
    const MT* __restrict__ velocity, T* __restrict__ param_out,
    MT* __restrict__ velocity_out, const MT mu, MT local_lr,
    const MT lars_weight_decay, const MT rescale_grad, const int tid,
    const int grid_stride, const int numel,
    MT* __restrict__ master_param_out = nullptr) {
  using VecType = paddle::platform::AlignedVector<T, VecSize>;
  using VecMType = paddle::platform::AlignedVector<MT, VecSize>;
  int main = numel >> (VecSize >> 1);
  int tail_offset = main * VecSize;

  const VecType* __restrict__ grad_vec = reinterpret_cast<const VecType*>(grad);
  const VecMType* __restrict__ param_vec =
      reinterpret_cast<const VecMType*>(param);
  const VecMType* __restrict__ velocity_vec =
      reinterpret_cast<const VecMType*>(velocity);
  VecType* param_out_vec = reinterpret_cast<VecType*>(param_out);
  VecMType* velocity_out_vec = reinterpret_cast<VecMType*>(velocity_out);


template <typename T, typename MT>
__global__ void L2NormKernel(
#endif
    const T* p_data, const T* __restrict__ g_data, MT* __restrict__ p_buffer,
    MT* __restrict__ g_buffer, const int64_t numel, const int repeat_times,
    const MT rescale_grad, const int thresh = 0, MT* __restrict__ p_n = nullptr,
    MT* __restrict__ g_n = nullptr) {
  __shared__ MT s_buffer[2];
  int tid = threadIdx.x + blockDim.x * blockIdx.x;
  int grid_stride = LARS_BLOCK_SIZE * gridDim.x;
  const MT rescale_pow = rescale_grad * rescale_grad;

  MT p_tmp = static_cast<MT>(0);
  MT g_tmp = static_cast<MT>(0);
  while (tid < numel) {
    MT tmp0 = static_cast<MT>(p_data[tid]);
    MT tmp1 = static_cast<MT>(g_data[tid]);
    p_tmp += (tmp0 * tmp0);
    g_tmp += (tmp1 * tmp1);
    tid += grid_stride;
  }
  p_tmp = math::blockReduceSum<MT>(p_tmp, FINAL_MASK);
  g_tmp = math::blockReduceSum<MT>(g_tmp, FINAL_MASK);

  if (threadIdx.x == 0) {
    p_buffer[blockIdx.x] = p_tmp;
    g_buffer[blockIdx.x] = g_tmp;
  }
#if CUDA_VERSION >= 11000
  cg->sync();  // Grid sync for writring partial result to gloabl memory
