PaddlePaddle
diff --git a/‎paddle/fluid/inference/tensorrt/plugin/layer_norm_op_plugin.cu
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/inference/tensorrt/plugin/layer_norm_op_plugin.cu
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/inference/tensorrt/plugin/trans_layernorm_op_plugin.cu
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎paddle/fluid/inference/tensorrt/plugin/trans_layernorm_op_plugin.cu
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
diff --git a/‎paddle/phi/kernels/funcs/load_store_util.h
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions b/‎paddle/phi/kernels/funcs/load_store_util.h
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_bias_act_kernel.cu
Copy file name to clipboard
+37-37Lines changed: 37 additions & 37 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_bias_act_kernel.cu
Copy file name to clipboard
+37-37Lines changed: 37 additions & 37 deletions
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_dropout_common.h
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion b/‎paddle/phi/kernels/fusion/gpu/fused_dropout_common.h
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_dropout_helper.h
Copy file name to clipboard
+7-7Lines changed: 7 additions & 7 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_dropout_helper.h
Copy file name to clipboard
+7-7Lines changed: 7 additions & 7 deletions
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_layernorm_kernel.cu
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_layernorm_kernel.cu
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
@@ -102,7 +102,7 @@ int LayerNormPlugin::enqueue(int batch_size,
   int64_t batched_mean_shape = mean_shape_[0] * input_dims.d[0];
   int64_t batched_variance_shape = variance_shape_[0] * input_dims.d[0];
 
-  std::vector<int> input_shape;
+  std::vector<int64_t> input_shape;
   input_shape.push_back(batch_size);
   for (int i = 0; i < input_dims.nbDims; i++) {
     input_shape.push_back(input_dims.d[i]);
@@ -277,7 +277,7 @@ int LayerNormPluginDynamic::enqueue(
   int begin_norm_axis = begin_norm_axis_;
   float eps = eps_;
 
-  std::vector<int> input_shape;
+  std::vector<int64_t> input_shape;
   for (int i = 0; i < input_dims.nbDims; i++) {
     input_shape.push_back(input_dims.d[i]);
   }
 
@@ -299,7 +299,7 @@ int TransLayerNormPluginDynamic::enqueue(
   int begin_norm_axis = begin_norm_axis_;
   float eps = eps_;
 
-  std::vector<int> input_shape;
+  std::vector<int64_t> input_shape;
   for (int i = 0; i < input_dims.nbDims; i++) {
     input_shape.push_back(input_dims.d[i]);
   }
@@ -334,7 +334,7 @@ int TransLayerNormPluginDynamic::enqueue(
 
   // transpose do not change numel
   int trans_result_numel = input_numel;
-  std::vector<int> trans_result_shape{
+  std::vector<int64_t> trans_result_shape{
       input_shape[0], input_shape[2], input_shape[3], input_shape[1]};
 
   const auto input_ddim = common::make_ddim(input_shape);
@@ -481,7 +481,7 @@ int TransLayerNormPluginDynamic::enqueue(
       int sm = getSMVersion();
       // sm >= 60 to support __ldg
       if (sm >= 60) {
-        int hidden = input_shape[1];
+        int64_t hidden = input_shape[1];
         if (hidden % 2 == 0) {
           const size_t rows =
               static_cast<size_t>(input_shape[0] * input_shape[2] *
 
@@ -49,7 +49,7 @@ struct Load {
   explicit Load(const T *src) : src_(src) {}
 
   template <int VecSize>
-  __device__ void load(phi::AlignedVector<T, VecSize> *dst, int idx) {
+  __device__ void load(phi::AlignedVector<T, VecSize> *dst, int64_t idx) {
     phi::Load<T, VecSize>(src_ + idx, dst);
   }
 
@@ -61,7 +61,7 @@ struct Store {
   explicit Store(T *dst) : dst_(dst) {}
 
   template <int VecSize>
-  __device__ void store(phi::AlignedVector<T, VecSize> &src, int idx) {
+  __device__ void store(phi::AlignedVector<T, VecSize> &src, int64_t idx) {
     phi::Store<T, VecSize>(src, dst_ + idx);
   }
 
@@ -74,7 +74,7 @@ struct Store<T, true> {
       : dst_(dst), shift_(shift), smooth_(smooth), cols_(cols) {}
 
   template <int VecSize>
-  __device__ void store(phi::AlignedVector<T, VecSize> &src, int idx) {
+  __device__ void store(phi::AlignedVector<T, VecSize> &src, int64_t idx) {
     using Vec = phi::AlignedVector<T, VecSize>;
     Vec shift_vec;
     Vec smooth_vec;
@@ -100,7 +100,7 @@ struct DequantLoad {
       : src_(src), dequant_scales_(dequant_scales), cols_(cols) {}
 
   template <int VecSize>
-  __device__ void load(phi::AlignedVector<T, VecSize> *dst, int idx) {
+  __device__ void load(phi::AlignedVector<T, VecSize> *dst, int64_t idx) {
     using SrcVec = phi::AlignedVector<int32_t, VecSize>;
     using DstVec = phi::AlignedVector<T, VecSize>;
     using ScaleVec = phi::AlignedVector<float, VecSize>;
@@ -139,7 +139,7 @@ struct QuantStore {
 
   template <int VecSize>
   __device__ void store(phi::AlignedVector<T, VecSize> &src,  // NOLINT
-                        int idx) {                            // NOLINT
+                        int64_t idx) {                        // NOLINT
     using DstVec = phi::AlignedVector<int8_t, VecSize>;
 
     DstVec dst_vec;
@@ -183,7 +183,7 @@ struct QuantStore<T, true> {
 
   template <int VecSize>
   __device__ void store(phi::AlignedVector<T, VecSize> &src,  // NOLINT
-                        int idx) {                            // NOLINT
+                        int64_t idx) {                        // NOLINT
     using DstVec = phi::AlignedVector<int8_t, VecSize>;
     using Vec = phi::AlignedVector<T, VecSize>;
 
 
@@ -28,25 +28,25 @@ template <typename T,
           typename StoreFunc>
 __global__ void ActFFNGlu(const T *bias,
                           Functor act_functor,
-                          const int token_num,
-                          const int hid_dim,
-                          const int elem_num,
+                          const int64_t token_num,
+                          const int64_t hid_dim,
+                          const int64_t elem_num,
                           LoadFunc load_func,
                           StoreFunc store_func) {
   using LoadT = phi::AlignedVector<T, VecSize>;
   LoadT src_vec1;
   LoadT src_vec2;
   LoadT bias_vec1;
   LoadT bias_vec2;
-  const int global_tid = blockIdx.x * blockDim.x + threadIdx.x;
-  for (int i = global_tid * VecSize; i < elem_num;
+  const int64_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int64_t i = global_tid * VecSize; i < elem_num;
        i += gridDim.x * blockDim.x * VecSize) {
-    int bi = i / hid_dim;
-    int idx = i % hid_dim;
+    int64_t bi = i / hid_dim;
+    int64_t idx = i % hid_dim;
 
-    load_func.template load<VecSize>(&src_vec1, bi * hid_dim * 2 + idx);
-    load_func.template load<VecSize>(&src_vec2,
-                                     bi * hid_dim * 2 + idx + hid_dim);
+    int64_t index = bi * hid_dim * 2 + idx;
+    load_func.template load<VecSize>(&src_vec1, index);
+    load_func.template load<VecSize>(&src_vec2, index + hid_dim);
 
     if (bias) {
       phi::Load<T, VecSize>(&bias[idx], &bias_vec1);
@@ -73,13 +73,13 @@ template <typename T,
           typename LoadT = T>
 void LaunchActFFNGlu(const Context &dev_ctx,
                      const T *bias,
-                     const int token_num,
-                     const int hid_dim,
+                     const int64_t token_num,
+                     const int64_t hid_dim,
                      LoadFunc load_func,
                      StoreFunc store_func) {
   constexpr int VecSize = 16;
   constexpr int PackSize = VecSize / sizeof(LoadT);
-  const int elem_cnt = token_num * hid_dim;
+  const int64_t elem_cnt = token_num * hid_dim;
   const int blocksize = 128;
   int grid_size = 1;
   Functor functor;
@@ -110,9 +110,9 @@ template <typename T,
           typename StoreFunc>
 __global__ void BiasAct(const T *bias,
                         Functor act_functor,
-                        const int rows,
-                        const int cols,
-                        const int elem_num,
+                        const int64_t rows,
+                        const int64_t cols,
+                        const int64_t elem_num,
                         LoadFunc load_func,
                         StoreFunc store_func) {
   using LoadT = phi::AlignedVector<T, VecSize>;
@@ -121,16 +121,16 @@ __global__ void BiasAct(const T *bias,
 
 // Zero Initialize BiasVec.
 #pragma unroll
-  for (int unroll_idx = 0; unroll_idx < VecSize; unroll_idx++) {
+  for (int64_t unroll_idx = 0; unroll_idx < VecSize; unroll_idx++) {
     bias_vec[unroll_idx] = 0;
   }
 
-  const int global_tid = blockIdx.x * blockDim.x + threadIdx.x;
-  for (int i = global_tid * VecSize; i < elem_num;
+  const int64_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int64_t i = global_tid * VecSize; i < elem_num;
        i += gridDim.x * blockDim.x * VecSize) {
-    int row_idx = i / cols;
-    int col_idx = i % cols;
-    int linear_idx = row_idx * cols + col_idx;
+    int64_t row_idx = i / cols;
+    int64_t col_idx = i % cols;
+    int64_t linear_idx = row_idx * cols + col_idx;
     load_func.template load<VecSize>(&src_vec, linear_idx);
     if (bias) {
       phi::Load<T, VecSize>(&bias[col_idx], &bias_vec);
@@ -154,13 +154,13 @@ template <typename T,
           typename LoadT = T>
 void LaunchBiasAct(const Context &dev_ctx,
                    const T *bias,
-                   const int token_num,
-                   const int hid_dim,
+                   const int64_t token_num,
+                   const int64_t hid_dim,
                    LoadFunc load_func,
                    StoreFunc store_func) {
   constexpr int VecSize = 16;
   constexpr int PackSize = VecSize / sizeof(LoadT);
-  const int elem_cnt = token_num * hid_dim;
+  const int64_t elem_cnt = token_num * hid_dim;
   const int blocksize = 128;
   int grid_size = 1;
   Functor functor;
@@ -192,8 +192,8 @@ template <typename T,
 void ComputeImpl(const Context &dev_ctx,
                  const T *bias_data,
                  const std::string &act_method,
-                 int rows,
-                 int cols,
+                 int64_t rows,
+                 int64_t cols,
                  LoadFunc load_func,
                  StoreFunc store_func) {
   if (act_method == "geglu") {
@@ -237,8 +237,8 @@ void DispatchComputeImpl(const Context &dev_ctx,
                          const DenseTensor *bias,
                          const DenseTensor *dequant_scales,
                          const std::string &act_method,
-                         int rows,
-                         int cols,
+                         int64_t rows,
+                         int64_t cols,
                          const float quant_scale,
                          const int quant_round_type,
                          const float quant_max_bound,
@@ -286,8 +286,8 @@ void DispatchComputeImpl(const Context &dev_ctx,
                          const DenseTensor *shift,
                          const DenseTensor *smooth,
                          const std::string &act_method,
-                         int rows,
-                         int cols,
+                         int64_t rows,
+                         int64_t cols,
                          const float quant_scale,
                          const int quant_round_type,
                          const float quant_max_bound,
@@ -363,8 +363,8 @@ void DispatchWithDtype(const Context &dev_ctx,
                        const paddle::optional<DenseTensor> &shift,
                        const paddle::optional<DenseTensor> &smooth,
                        const std::string &act_method,
-                       int rows,
-                       int cols,
+                       int64_t rows,
+                       int64_t cols,
                        float quant_scale,
                        int quant_round_type,
                        float quant_max_bound,
@@ -423,8 +423,8 @@ void DispatchWithDtype(const Context &dev_ctx,
                        const paddle::optional<DenseTensor> &shift,
                        const paddle::optional<DenseTensor> &smooth,
                        const std::string &act_method,
-                       int rows,
-                       int cols,
+                       int64_t rows,
+                       int64_t cols,
                        float quant_scale,
                        int quant_round_type,
                        float quant_max_bound,
@@ -446,8 +446,8 @@ void FusedBiasActKernel(const Context &dev_ctx,
                         float quant_max_bound,
                         float quant_min_bound,
                         DenseTensor *out) {
-  int cols = x.dims()[x.dims().size() - 1];
-  int rows = x.numel() / cols;
+  int64_t cols = x.dims()[x.dims().size() - 1];
+  int64_t rows = x.numel() / cols;
   if (x.dtype() == phi::DataType::INT32) {
     if (compute_dtype == "bf16") {
       DispatchWithDtype<phi::dtype::bfloat16, Context>(
 
@@ -70,10 +70,16 @@ inline phi::backends::gpu::GpuLaunchConfig Get1DBlocksAnd2DGrids(
                                       ctx.GetMaxThreadsPerBlock(), 512))));
   const auto blocks_x =
       std::max(static_cast<uint32_t>(1), (tmp_cols + threads - 1) / threads);
-  const auto blocks_y = std::max(static_cast<uint32_t>(1), rows);
+  int blocks_y = std::max(static_cast<uint32_t>(1), rows);
+  int blocks_z = 1;
+  if (blocks_y > 65536) {
+    blocks_z = 1024;
+    blocks_y = (blocks_y + blocks_z - 1) / blocks_z;
+  }
   phi::backends::gpu::GpuLaunchConfig config;
   config.block_per_grid.x = blocks_x;
   config.block_per_grid.y = blocks_y;
+  config.block_per_grid.z = blocks_z;
   config.thread_per_block.x = threads;
   return config;
 }
 
@@ -337,8 +337,8 @@ class FusedDropoutHelper {
   }
 
  protected:
-  int rows_;
-  int cols_;
+  int64_t rows_;
+  int64_t cols_;
   DropoutParam dropout_param_;
   float residual_alpha_;
 };
@@ -351,8 +351,8 @@ class FusedDropoutLayerNormHelper
     : public FusedDropoutHelper<T, MaskType, InType, OutType> {
  public:
   FusedDropoutLayerNormHelper() {}
-  FusedDropoutLayerNormHelper(const int rows,
-                              const int cols,
+  FusedDropoutLayerNormHelper(const int64_t rows,
+                              const int64_t cols,
                               const float epsilon,
                               const float residual_alpha = 1.0) {
     using U = phi::funcs::LayerNormParamType<T>;
@@ -363,8 +363,8 @@ class FusedDropoutLayerNormHelper
   }
 
   FusedDropoutLayerNormHelper(const phi::GPUContext& ctx,
-                              const int rows,
-                              const int cols,
+                              const int64_t rows,
+                              const int64_t cols,
                               const DropoutParam& dropout_param,
                               const float epsilon,
                               const float residual_alpha = 1.0)
@@ -388,7 +388,7 @@ class FusedDropoutLayerNormHelper
     phi::LayerNormDirectCUDAFunctor<InDataType,
                                     phi::funcs::LayerNormParamType<T>>
         layer_norm;
-    std::vector<int> src_shape{this->rows_, this->cols_};
+    std::vector<int64_t> src_shape{this->rows_, this->cols_};
     layer_norm(ctx.stream(),
                reinterpret_cast<const InDataType*>(src),
                src_shape,
 
@@ -1048,8 +1048,8 @@ void FusedLayerNormKernel(const Context& dev_ctx,
   const U* norm_weight_data =
       norm_weight ? norm_weight.get().data<U>() : nullptr;
   const U* norm_bias_data = norm_bias ? norm_bias.get().data<U>() : nullptr;
-  int32_t rows = 1;
-  int32_t cols = 1;
+  int64_t rows = 1;
+  int64_t cols = 1;
   for (int i = 0; i < begin_norm_axis; i++) {
     rows *= x.dims()[i];
   }