PaddlePaddle
diff --git a/‎paddle/phi/infermeta/multiary.cc
Copy file name to clipboard
+103Lines changed: 103 additions & 0 deletions b/‎paddle/phi/infermeta/multiary.cc
Copy file name to clipboard
+103Lines changed: 103 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/multiary.h
Copy file name to clipboard
+20Lines changed: 20 additions & 0 deletions b/‎paddle/phi/infermeta/multiary.h
Copy file name to clipboard
+20Lines changed: 20 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+111Lines changed: 111 additions & 0 deletions b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+111Lines changed: 111 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/CMakeLists.txt
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/kernels/CMakeLists.txt
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
@@ -2273,6 +2273,109 @@ void FakeQuantOrWithDequantMovingAverageAbsMaxInferMeta(
   out->share_lod(x);
 }
 
+void Fp8GemmBlockwiseInferMeta(const MetaTensor& A,
+                               const MetaTensor& A_scale,
+                               const MetaTensor& B,
+                               const MetaTensor& B_scale,
+                               const MetaTensor& input_result,
+                               const MetaTensor& bias,
+                               const MetaTensor& pre_gelu,
+                               const MetaTensor& workspace,
+                               bool transa,
+                               bool transb,
+                               bool grad,
+                               bool accumulate,
+                               bool use_split_accumulator,
+                               int math_sm_count,
+                               bool is_A_1d_scaled,
+                               bool is_B_1d_scaled,
+                               MetaTensor* output,
+                               MetaTensor* pre_gelu_out,
+                               MetaTensor* workspace_out) {
+  PADDLE_ENFORCE_EQ(
+      use_split_accumulator,
+      true,
+      errors::InvalidArgument("Only split accumulator is supported"));
+
+  auto A_dims = A.dims();
+  auto B_dims = B.dims();
+
+  PADDLE_ENFORCE_EQ(
+      transa,
+      true,
+      errors::InvalidArgument("Only transa == true is supported"));
+
+  PADDLE_ENFORCE_EQ(
+      transb,
+      false,
+      errors::InvalidArgument("Only transb == false is supported"));
+
+  PADDLE_ENFORCE_EQ(
+      A_dims.size(),
+      2,
+      errors::InvalidArgument("Input A should have 2 dimensions"));
+
+  PADDLE_ENFORCE_EQ(
+      B_dims.size(),
+      2,
+      errors::InvalidArgument("Input B should have 2 dimensions"));
+
+  const auto IsFp8Dtype = [](const paddle::DataType dtype) {
+    return dtype == phi::DataType::FLOAT8_E4M3FN ||
+           dtype == phi::DataType::FLOAT8_E5M2;
+  };
+
+  PADDLE_ENFORCE_EQ(IsFp8Dtype(A.dtype()),
+                    true,
+                    errors::InvalidArgument("A must be FP8 dtype"));
+
+  PADDLE_ENFORCE_EQ(IsFp8Dtype(B.dtype()),
+                    true,
+                    errors::InvalidArgument("B must be FP8 dtype"));
+
+  PADDLE_ENFORCE_EQ(
+      A_scale.dtype(),
+      phi::DataType::FLOAT32,
+      errors::InvalidArgument(
+          "The dtype of A_scale must be float32, but got %d", A_scale.dtype()));
+
+  PADDLE_ENFORCE_EQ(
+      B_scale.dtype(),
+      phi::DataType::FLOAT32,
+      errors::InvalidArgument(
+          "The dtype of B_scale must be float32, but got %d", B_scale.dtype()));
+
+  PADDLE_ENFORCE_EQ(input_result.dtype() == phi::DataType::FLOAT32 ||
+                        input_result.dtype() == phi::DataType::BFLOAT16,
+                    true,
+                    errors::InvalidArgument(
+                        "out_dtype must be bfloat16 or float32, but got %d",
+                        input_result.dtype()));
+
+  // Validate scaling modes
+  PADDLE_ENFORCE_EQ(is_A_1d_scaled || is_B_1d_scaled,
+                    true,
+                    errors::InvalidArgument("2Dx2D scaling is not supported"));
+
+  // Validate matrix dimension compatibility
+  PADDLE_ENFORCE_EQ(
+      transa ? A_dims[1] : A_dims[0],
+      transb ? B_dims[0] : B_dims[1],
+      errors::InvalidArgument(
+          "Matrix inner dimensions must match for multiplication. "
+          "A inner dim: %d, B inner dim: %d",
+          transa ? A_dims[1] : A_dims[0],
+          transb ? B_dims[0] : B_dims[1]));
+
+  // Set output dimensions and dtype
+  output->set_dims(input_result.dims());
+  output->set_dtype(input_result.dtype());
+  pre_gelu_out->set_dims(pre_gelu.dims());
+  pre_gelu_out->set_dtype(pre_gelu.dtype());
+  workspace_out->set_dims(workspace.dims());
+  workspace_out->set_dtype(workspace.dtype());
+}
+
 void FtrlInferMeta(const MetaTensor& param,
                    const MetaTensor& squared_accumulator,
                    const MetaTensor& linear_accumulator,
 
@@ -478,6 +478,26 @@ void FakeQuantOrWithDequantMovingAverageAbsMaxInferMeta(
     MetaTensor* out_state,
     MetaTensor* out_accum);
 
+void Fp8GemmBlockwiseInferMeta(const MetaTensor& A,
+                               const MetaTensor& A_scale,
+                               const MetaTensor& B,
+                               const MetaTensor& B_scale,
+                               const MetaTensor& input_result,
+                               const MetaTensor& bias,
+                               const MetaTensor& pre_gelu,
+                               const MetaTensor& workspace,
+                               bool transa,
+                               bool transb,
+                               bool grad,
+                               bool accumulate,
+                               bool use_split_accumulator,
+                               int math_sm_count,
+                               bool is_A_1d_scaled,
+                               bool is_B_1d_scaled,
+                               MetaTensor* output,
+                               MetaTensor* pre_gelu_out,
+                               MetaTensor* workspace_out);
+
 void FtrlInferMeta(const MetaTensor& param,
                    const MetaTensor& squared_accumulator,
                    const MetaTensor& linear_accumulator,
 
@@ -2115,6 +2115,117 @@ void FrameInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
+void Fp8QuantBlockwiseInferMeta(const MetaTensor& X,
+                                float epsilon,
+                                bool using_1x128_vec_quant,
+                                bool input_transpose,
+                                bool output_scale_transpose,
+                                bool using_e5m2,
+                                bool using_pow2_scale,
+                                MetaTensor* out,
+                                MetaTensor* scale,
+                                MetaTensor* out_transposed,
+                                MetaTensor* scale_transposed) {
+  auto x_dims = X.dims();
+  PADDLE_ENFORCE_EQ(
+      x_dims.size(),
+      2,
+      common::errors::InvalidArgument("Input must have exactly 2 dimensions."));
+  PADDLE_ENFORCE_EQ(
+      using_e5m2,
+      false,
+      common::errors::InvalidArgument("currently e5m2 is not support."));
+  PADDLE_ENFORCE_EQ(X.dtype(),
+                    phi::DataType::BFLOAT16,
+                    common::errors::InvalidArgument(
+                        "currently only support bfloat16 input."));
+
+  const int64_t rows = x_dims[0];
+  const int64_t cols = x_dims[1];
+  PADDLE_ENFORCE_LE(rows,
+                    65535 * 128,
+                    common::errors::InvalidArgument(
+                        "Currently only supports the first dim of "
+                        "Input(X) <= 65535 * 128, but got %d",
+                        rows));
+  PADDLE_ENFORCE_EQ(
+      rows % 128,
+      0,
+      common::errors::InvalidArgument("The first dim of "
+                                      "Input(X) should be exactly divided "
+                                      "by 128 , but got %d",
+                                      rows));
+  PADDLE_ENFORCE_EQ(cols % 128,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The last dim of Input(X) should be exactly divided "
+                        "by 128 , but got %d",
+                        cols));
+
+  const int64_t row_quantized = (rows + 127) / 128;
+  const int64_t col_quantized = (cols + 127) / 128;
+
+  int64_t output_outer_dim = -1;
+  int64_t output_inner_dim = -1;
+  int64_t scale_outer_dim = -1;
+  int64_t scale_inner_dim = -1;
+  int64_t scale_transposed_outer_dim = -1;
+  int64_t scale_transposed_inner_dim = -1;
+
+  // output_fp8's shape should exactly match input x.
+  output_outer_dim = rows;
+  output_inner_dim = cols;
+
+  if (using_1x128_vec_quant) {
+    // 1x128 w/wo transpose
+    scale_outer_dim = output_scale_transpose ? col_quantized : rows;
+    scale_inner_dim = output_scale_transpose ? rows : col_quantized;
+    scale_transposed_outer_dim = output_scale_transpose ? row_quantized : cols;
+    scale_transposed_inner_dim = output_scale_transpose ? cols : row_quantized;
+  } else {
+    // 128x128 w/wo transpose
+    scale_outer_dim = output_scale_transpose ? col_quantized : row_quantized;
+    scale_inner_dim = output_scale_transpose ? row_quantized : col_quantized;
+    scale_transposed_outer_dim = scale_inner_dim;
+    scale_transposed_inner_dim = scale_outer_dim;
+  }
+
+  PADDLE_ENFORCE_GT(output_outer_dim,
+                    0,
+                    common::errors::InvalidArgument(
+                        "invalid shape encountered in output outer dim."));
+  PADDLE_ENFORCE_GT(output_inner_dim,
+                    0,
+                    common::errors::InvalidArgument(
+                        "invalid shape encountered in output inner dim."));
+  PADDLE_ENFORCE_GT(scale_outer_dim,
+                    0,
+                    common::errors::InvalidArgument(
+                        "invalid shape encountered in scale outer dim."));
+  PADDLE_ENFORCE_GT(scale_inner_dim,
+                    0,
+                    common::errors::InvalidArgument(
+                        "invalid shape encountered in scale inner dim."));
+
+  if (X && out && scale) {
+    out->set_dims(common::make_ddim({output_outer_dim, output_inner_dim}));
+    out->set_dtype(phi::DataType::FLOAT8_E4M3FN);
+    scale->set_dims(common::make_ddim({scale_outer_dim, scale_inner_dim}));
+    scale->set_dtype(phi::DataType::FLOAT32);
+    if (input_transpose) {
+      out_transposed->set_dims(
+          common::make_ddim({output_inner_dim, output_outer_dim}));
+      out_transposed->set_dtype(phi::DataType::FLOAT8_E4M3FN);
+      scale_transposed->set_dims(common::make_ddim(
+          {scale_transposed_outer_dim, scale_transposed_inner_dim}));
+      scale_transposed->set_dtype(phi::DataType::FLOAT32);
+    }
+  } else {
+    PADDLE_THROW(
+        common::errors::InvalidArgument("invalid input or output tensor"));
+  }
+}
+
 void FullBatchSizeLikeInferMeta(const MetaTensor& x,
                                 const std::vector<int>& shape,
                                 const Scalar& val,
 
@@ -356,6 +356,18 @@ void FrameInferMeta(const MetaTensor& x,
                     MetaTensor* out,
                     MetaConfig = MetaConfig());
 
+void Fp8QuantBlockwiseInferMeta(const MetaTensor& X,
+                                float epsilon,
+                                bool using_1x128_vec_quant,
+                                bool input_transpose,
+                                bool output_scale_transpose,
+                                bool using_e5m2,
+                                bool using_pow2_scale,
+                                MetaTensor* out,
+                                MetaTensor* scale,
+                                MetaTensor* out_transposed,
+                                MetaTensor* scale_transposed);
+
 void FullBatchSizeLikeInferMeta(const MetaTensor& x,
                                 const std::vector<int>& shape,
                                 const Scalar& val,
 
@@ -71,6 +71,8 @@ if(((WITH_GPU) AND (CUDA_VERSION VERSION_LESS 12.0))
     "legacy/gpu/moe_gate_dispatch_grad_kernel.cu"
     "legacy/gpu/int_bincount.cu"
     "legacy/gpu/layer_norm_cuda_kernel.cu"
+    "legacy/gpu/fp8_gemm_blockwise_kernel.cu"
+    "legacy/gpu/fp8_quant_blockwise_kernel.cu"
     "fusion/gpu/fused_act_dequant_kernel.cu"
     "fusion/gpu/fused_stack_transpose_quant_kernel.cu"
     "fusion/gpu/fused_stack_quant_kernel.cu"