PaddlePaddle
diff --git a/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+195Lines changed: 195 additions & 0 deletions b/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+195Lines changed: 195 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/fusion.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions b/‎paddle/phi/infermeta/fusion.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/CMakeLists.txt
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎paddle/phi/kernels/CMakeLists.txt
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_act_dequant_kernel.cu
Copy file name to clipboard
+128Lines changed: 128 additions & 0 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_act_dequant_kernel.cu
Copy file name to clipboard
+128Lines changed: 128 additions & 0 deletions
@@ -853,6 +853,51 @@ void FcXPUInferMeta(const MetaTensor& x,
   out_max->set_layout(x.layout());
 }
 
+void FusedActDequantInferMeta(const MetaTensor& x,
+                              const MetaTensor& x_scale,
+                              MetaTensor* out) {
+  auto x_dims = x.dims();
+
+  PADDLE_ENFORCE_EQ(
+      x.dtype(),
+      phi::DataType::FLOAT8_E4M3FN,
+      phi::errors::InvalidArgument(
+          "The data type of X should be FLOAT8_E4M3FN, but received %s.",
+          x.dtype()));
+
+  PADDLE_ENFORCE_EQ(
+      x_scale.dtype(),
+      phi::DataType::FLOAT32,
+      phi::errors::InvalidArgument(
+          "The data type of X_scale should be FLOAT32, but received %s.",
+          x_scale.dtype()));
+
+  PADDLE_ENFORCE_EQ(x_dims.size(),
+                    2,
+                    phi::errors::InvalidArgument(
+                        "The input X should be a 2D tensor, but received %dD.",
+                        x_dims.size()));
+
+  int64_t rows = x_dims[0];
+  int64_t cols = x_dims[1];
+
+  PADDLE_ENFORCE_GT(
+      rows,
+      0,
+      phi::errors::InvalidArgument(
+          "The rows of X should be positive, but received %d.", rows));
+
+  PADDLE_ENFORCE_GT(
+      cols,
+      0,
+      phi::errors::InvalidArgument(
+          "The cols of X should be positive, but received %d.", cols));
+
+  out->set_dims(x_dims);
+  out->set_dtype(phi::DataType::BFLOAT16);
+  out->set_layout(x.layout());
+}
+
 void FusedAttentionInferMeta(const MetaTensor& x,
                              const MetaTensor& ln_scale,
                              const MetaTensor& ln_bias,
@@ -5637,6 +5682,156 @@ static phi::DDim GetBitmaskDims(std::vector<int> out_shape) {
   return common::make_ddim(bitmask_shape);
 }
 
+void FusedSwigluWeightedBwdInferMeta(const MetaTensor& o1,
+                                     const MetaTensor& do2_s,
+                                     const MetaTensor& unzipped_probs,
+                                     MetaTensor* do1,
+                                     MetaTensor* probs_grad,
+                                     MetaTensor* o2_s) {
+  PADDLE_ENFORCE_EQ(
+      o1.dtype(),
+      phi::DataType::BFLOAT16,
+      phi::errors::InvalidArgument("The data type of o1 must be bfloat16. "
+                                   "But received o1 dtype: %s",
+                                   phi::DataTypeToString(o1.dtype())));
+
+  PADDLE_ENFORCE_EQ(
+      do2_s.dtype(),
+      phi::DataType::BFLOAT16,
+      phi::errors::InvalidArgument("The data type of do2_s must be bfloat16. "
+                                   "But received do2_s dtype: %s",
+                                   phi::DataTypeToString(do2_s.dtype())));
+
+  PADDLE_ENFORCE_EQ(unzipped_probs.dtype(),
+                    phi::DataType::FLOAT32,
+                    phi::errors::InvalidArgument(
+                        "The data type of unzipped_probs must be float32. "
+                        "But received unzipped_probs dtype: %s",
+                        phi::DataTypeToString(unzipped_probs.dtype())));
+
+  auto o1_dims = o1.dims();
+  auto do2_s_dims = do2_s.dims();
+  auto probs_dims = unzipped_probs.dims();
+
+  PADDLE_ENFORCE_EQ(
+      o1_dims.size(),
+      do2_s_dims.size(),
+      phi::errors::InvalidArgument(
+          "o1 and do2_s should have the same number of dimensions. "
+          "But received o1 dims: %d, do2_s dims: %d",
+          o1_dims.size(),
+          do2_s_dims.size()));
+
+  PADDLE_ENFORCE_EQ(
+      o1_dims.size(),
+      probs_dims.size(),
+      phi::errors::InvalidArgument(
+          "o1 and unzipped_probs should have the same number of dimensions. "
+          "But received o1 dims: %d, unzipped_probs dims: %d",
+          o1_dims.size(),
+          probs_dims.size()));
+
+  int64_t o1_last_dim = o1_dims[o1_dims.size() - 1];
+  int64_t do2_s_last_dim = do2_s_dims[do2_s_dims.size() - 1];
+
+  PADDLE_ENFORCE_EQ(o1_last_dim,
+                    do2_s_last_dim * 2,
+                    phi::errors::InvalidArgument(
+                        "The last dimension of o1 should be twice the last "
+                        "dimension of do2_s. "
+                        "But received o1 last dim: %d, do2_s last dim: %d",
+                        o1_last_dim,
+                        do2_s_last_dim));
+
+  int64_t o1_batch_size = 1;
+  int64_t do2_s_batch_size = 1;
+  int64_t probs_batch_size = 1;
+
+  for (int i = 0; i < o1_dims.size() - 1; i++) {
+    o1_batch_size *= o1_dims[i];
+    do2_s_batch_size *= do2_s_dims[i];
+    probs_batch_size *= probs_dims[i];
+  }
+
+  PADDLE_ENFORCE_EQ(o1_batch_size,
+                    do2_s_batch_size,
+                    phi::errors::InvalidArgument(
+                        "o1 and do2_s should have the same batch size (product "
+                        "of all dimensions except last). "
+                        "But received o1 batch size: %d, do2_s batch size: %d",
+                        o1_batch_size,
+                        do2_s_batch_size));
+
+  PADDLE_ENFORCE_EQ(o1_batch_size,
+                    probs_batch_size,
+                    phi::errors::InvalidArgument(
+                        "o1 and unzipped_probs should have the same batch size "
+                        "(product of all dimensions except last). "
+                        "But received o1 batch size: %d, probs batch size: %d",
+                        o1_batch_size,
+                        probs_batch_size));
+
+  do1->set_dims(o1_dims);
+  do1->set_dtype(o1.dtype());
+  do1->set_layout(o1.layout());
+
+  probs_grad->set_dims(probs_dims);
+  probs_grad->set_dtype(phi::DataType::FLOAT32);
+  probs_grad->set_layout(unzipped_probs.layout());
+
+  o2_s->set_dims(do2_s_dims);
+  o2_s->set_dtype(do2_s.dtype());
+  o2_s->set_layout(do2_s.layout());
+}
+
+void FusedWeightedSwigluActQuantInferMeta(const MetaTensor& x,
+                                          const MetaTensor& prob,
+                                          bool using_pow2_scaling,
+                                          MetaTensor* out,
+                                          MetaTensor* scale) {
+  PADDLE_ENFORCE_EQ(
+      x.dtype(),
+      DataType::BFLOAT16,
+      common::errors::InvalidArgument(
+          "The dtype of Input(x) must be BFLOAT16, but received %s",
+          x.dtype()));
+  if (prob) {
+    PADDLE_ENFORCE_EQ(
+        prob.dtype(),
+        DataType::FLOAT32,
+        common::errors::InvalidArgument(
+            "The dtype of Input(prob) must be FLOAT32, but received %s",
+            prob.dtype()));
+  }
+  int64_t rows = 1;
+  for (int i = 0; i < x.dims().size() - 1; ++i) {
+    rows *= x.dims()[i];
+  }
+  int64_t cols = x.dims()[x.dims().size() - 1];
+  PADDLE_ENFORCE_EQ(cols % 2,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The last dim of Input(X) should be exactly divided "
+                        "by 2 , but got %d",
+                        cols));
+  if (prob) {
+    PADDLE_ENFORCE_EQ(prob.dims()[0],
+                      rows,
+                      common::errors::InvalidArgument(
+                          "The first dim of Input(x) should be equal to the "
+                          "first dim of Input(prob) but got X.shape[0]: %d, "
+                          "prob.shape[0]: %d",
+                          rows,
+                          prob.dims()[0]));
+  }
+
+  out->set_dims(common::make_ddim({rows, cols / 2}));
+  out->set_dtype(DataType::FLOAT8_E4M3FN);
+
+  scale->set_dims(common::make_ddim({rows, ((cols / 2) + 127) / 128}));
+  scale->set_dtype(DataType::FLOAT32);
+}
+
 void ResnetUnitInferMeta(const MetaTensor& x,
                          const MetaTensor& filter_x,
                          const MetaTensor& scale_x,
 
@@ -296,6 +296,10 @@ void MultiEncoderXPUInferMeta(
     MetaTensor* x_fp16,
     MetaTensor* out_fp16);
 
+void FusedActDequantInferMeta(const MetaTensor& x,
+                              const MetaTensor& x_scale,
+                              MetaTensor* out);
+
 void FusedAttentionInferMeta(const MetaTensor& x,
                              const MetaTensor& ln_scale,
                              const MetaTensor& ln_bias,
@@ -1212,6 +1216,19 @@ void FusionSeqpoolConcatInferMeta(const std::vector<const MetaTensor*>& x,
                                   MetaTensor* out,
                                   MetaConfig config = MetaConfig());
 
+void FusedSwigluWeightedBwdInferMeta(const MetaTensor& o1,
+                                     const MetaTensor& do2_s,
+                                     const MetaTensor& unzipped_probs,
+                                     MetaTensor* do1,
+                                     MetaTensor* probs_grad,
+                                     MetaTensor* o2_s);
+
+void FusedWeightedSwigluActQuantInferMeta(const MetaTensor& x,
+                                          const MetaTensor& prob,
+                                          bool using_pow2_scaling,
+                                          MetaTensor* out,
+                                          MetaTensor* scale);
+
 void ResnetUnitInferMeta(const MetaTensor& x,
                          const MetaTensor& filter_x,
                          const MetaTensor& scale_x,
 
@@ -71,9 +71,12 @@ 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"
+    "fusion/gpu/fused_act_dequant_kernel.cu"
     "fusion/gpu/fused_stack_transpose_quant_kernel.cu"
     "fusion/gpu/fused_stack_quant_kernel.cu"
-    "fusion/gpu/fused_transpose_split_quant_kernel.cu")
+    "fusion/gpu/fused_transpose_split_quant_kernel.cu"
+    "fusion/gpu/fused_swiglu_weighted_bwd_kernel.cu"
+    "fusion/gpu/fused_weighted_swiglu_act_quant_kernel.cu")
 endif()
 
 # Get flag for CUDA arch >= 80
 
@@ -0,0 +1,128 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+
+template <typename T, int N>
+struct alignas(16) VectorType {
+  T data[N];
+};
+
+__global__ void FusedActDequant(const phi::float8_e4m3fn* __restrict__ Xin,
+                                const float* __restrict__ Xscale,
+                                phi::bfloat16* __restrict__ out,
+                                const int rows,
+                                const int cols) {
+  const int this_row_idx = blockIdx.x;
+  if (this_row_idx >= rows) return;
+
+  const int Xscale_stride = (cols + 127) / 128;
+
+  const int vector_size = 16;
+
+  const int num_vectors = cols / vector_size;
+  const int remaining_elements = cols % vector_size;
+
+  const int tid = threadIdx.x;
+
+  for (int vec_idx = tid; vec_idx < num_vectors; vec_idx += blockDim.x) {
+    int x_offset = vec_idx * vector_size;
+    int64_t X_idx = (int64_t)this_row_idx * (int64_t)cols + (int64_t)x_offset;
+
+    const VectorType<__nv_fp8_e4m3, vector_size>* X_vec_ptr =
+        reinterpret_cast<const VectorType<__nv_fp8_e4m3, vector_size>*>(Xin +
+                                                                        X_idx);
+    VectorType<__nv_fp8_e4m3, vector_size> X_vec = X_vec_ptr[0];
+
+    int64_t scale_idx =
+        (int64_t)this_row_idx * (int64_t)Xscale_stride + (x_offset / 128);
+    float this_scale = Xscale[scale_idx];
+
+    VectorType<__nv_bfloat16, vector_size> out_vec;
+
+#pragma unroll
+    for (int i = 0; i < vector_size; ++i) {
+      float X_value = static_cast<float>(X_vec.data[i]);
+      X_value *= this_scale;
+      out_vec.data[i] = __float2bfloat16(X_value);
+    }
+
+    VectorType<__nv_bfloat16, vector_size>* out_vec_ptr =
+        reinterpret_cast<VectorType<__nv_bfloat16, vector_size>*>(out + X_idx);
+    out_vec_ptr[0] = out_vec;
+  }
+
+  if (remaining_elements > 0) {
+    int x_offset = num_vectors * vector_size;
+    int64_t X_idx = (int64_t)this_row_idx * (int64_t)cols + (int64_t)x_offset;
+    int64_t idx = X_idx + tid;
+    if (tid < remaining_elements) {
+      float X_value = static_cast<float>(Xin[idx]);
+      X_value *= Xscale[(int64_t)this_row_idx * (int64_t)Xscale_stride +
+                        (x_offset / 128)];
+      out[idx] = __float2bfloat16(X_value);
+    }
+  }
+}
+
+template <typename T, typename Context>
+void FusedActDequantKernel(const Context& dev_ctx,
+                           const DenseTensor& x,
+                           const DenseTensor& x_scale,
+                           DenseTensor* out) {
+  auto x_dims = x.dims();
+  int rows = x_dims[0];
+  int cols = x_dims[1];
+
+  out->Resize({rows, cols});
+  dev_ctx.template Alloc<phi::dtype::bfloat16>(out);
+
+  auto out_ptr =
+      reinterpret_cast<void*>(out->template data<phi::dtype::bfloat16>());
+  cudaMemsetAsync(
+      out_ptr, 0, sizeof(phi::dtype::bfloat16) * rows * cols, dev_ctx.stream());
+
+  dim3 grid(rows);
+  dim3 block(256);
+
+  FusedActDequant<<<grid, block, 0, dev_ctx.stream()>>>(
+      x.data<phi::dtype::float8_e4m3fn>(),
+      x_scale.data<float>(),
+      out->data<phi::dtype::bfloat16>(),
+      rows,
+      cols);
+
+#ifdef PADDLE_WITH_CUDA_CHECK
+  auto cuda_error = cudaGetLastError();
+  PADDLE_ENFORCE_GPU_SUCCESS(cuda_error);
+#endif
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(fused_act_dequant,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::FusedActDequantKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float8_e4m3fn) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BFLOAT16);
+}