PaddlePaddle
diff --git a/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions b/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/fusion.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/phi/infermeta/fusion.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/CMakeLists.txt
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/phi/kernels/CMakeLists.txt
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_stack_quant_kernel.cu
Copy file name to clipboard
+121Lines changed: 121 additions & 0 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_stack_quant_kernel.cu
Copy file name to clipboard
+121Lines changed: 121 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/fusion/gpu/fused_stack_transpose_quant.h
Copy file name to clipboard
+121Lines changed: 121 additions & 0 deletions b/‎paddle/phi/kernels/fusion/gpu/fused_stack_transpose_quant.h
Copy file name to clipboard
+121Lines changed: 121 additions & 0 deletions
@@ -3922,6 +3922,108 @@ void FusionSeqExpandConcatFCInferMeta(const std::vector<const MetaTensor*>& x,
   out->share_lod(*x[0]);
 }
 
+std::tuple<int64_t, int64_t, int64_t> FusedStackQuantCommonCheck(
+    const std::vector<const MetaTensor*>& x) {
+  PADDLE_ENFORCE_GT(x.size(),
+                    0UL,
+                    common::errors::InvalidArgument(
+                        "Number of Inputs(x) must be larger than 0, but"
+                        " received value is:%d.",
+                        x.size()));
+  int64_t N = x.size();
+  for (int i = 0; i < N; ++i) {
+    PADDLE_ENFORCE_EQ(
+        x[i]->dtype(),
+        DataType::BFLOAT16,
+        common::errors::InvalidArgument(
+            "input must be bfloat16, but received dtype: %s", x[i]->dtype()));
+  }
+  auto input_dims = x[0]->dims();
+  PADDLE_ENFORCE_EQ(
+      input_dims.size(),
+      2U,
+      common::errors::InvalidArgument(
+          "input must be 2-D, but received dims: %s", input_dims.to_str()));
+  int64_t M = input_dims[0];
+  int64_t K = input_dims[1];
+  for (int i = 1; i < N; ++i) {
+    input_dims = x[i]->dims();
+    PADDLE_ENFORCE_EQ(input_dims.size(),
+                      2U,
+                      common::errors::InvalidArgument(
+                          "input must be 2-D, but received input[%d] dims: %s",
+                          i,
+                          input_dims.to_str()));
+    PADDLE_ENFORCE_EQ(
+        input_dims[0],
+        M,
+        common::errors::InvalidArgument(
+            "input [%d] must be shape %d, %d, but received dims: %s",
+            i,
+            M,
+            K,
+            input_dims.to_str()));
+    PADDLE_ENFORCE_EQ(
+        input_dims[1],
+        K,
+        common::errors::InvalidArgument(
+            "input [%d] must be shape %d, %d, but received dims: %s",
+            i,
+            M,
+            K,
+            input_dims.to_str()));
+  }
+  PADDLE_ENFORCE_LE(N,
+                    65535,
+                    common::errors::InvalidArgument(
+                        "The batch size (N) must be no larger than 65535."));
+  PADDLE_ENFORCE_EQ(M % 128,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The upper dim (M) must be multiple of 128."));
+  PADDLE_ENFORCE_EQ(K % 128,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The lower dim (K) must be multiple of 128."));
+  return {N, M, K};
+}
+
+void FusedStackTransposeQuantInferMeta(const std::vector<const MetaTensor*>& x,
+                                       MetaTensor* out,
+                                       MetaTensor* scale) {
+  int64_t N, M, K;
+  std::tie(N, M, K) = FusedStackQuantCommonCheck(x);
+
+  std::vector<int64_t> out_shape = {N * K, M};
+  std::vector<int64_t> scale_shape = {N * K / 128, M / 128};
+  out->set_dims(common::make_ddim(out_shape));
+  scale->set_dims(common::make_ddim(scale_shape));
+  out->set_dtype(DataType::FLOAT8_E4M3FN);
+  scale->set_dtype(DataType::FLOAT32);
+  out->share_lod(*x.at(0));
+  scale->share_lod(*x.at(0));
+  out->set_layout(x.at(0)->layout());
+  scale->set_layout(x.at(0)->layout());
+}
+
+void FusedStackQuantInferMeta(const std::vector<const MetaTensor*>& x,
+                              MetaTensor* out,
+                              MetaTensor* scale) {
+  int64_t N, M, K;
+  std::tie(N, M, K) = FusedStackQuantCommonCheck(x);
+
+  std::vector<int64_t> out_shape = {N * M, K};
+  std::vector<int64_t> scale_shape = {N * M / 128, K / 128};
+  out->set_dims(common::make_ddim(out_shape));
+  scale->set_dims(common::make_ddim(scale_shape));
+  out->set_dtype(DataType::FLOAT8_E4M3FN);
+  scale->set_dtype(DataType::FLOAT32);
+  out->share_lod(*x.at(0));
+  scale->share_lod(*x.at(0));
+  out->set_layout(x.at(0)->layout());
+  scale->set_layout(x.at(0)->layout());
+}
+
 // Current constraint is appropriate for GemmEpilogueOp but relaxed for FcOp
 void FCInferMeta(const MetaTensor& input,
                  const MetaTensor& w,
 
@@ -904,6 +904,14 @@ void FusionSeqExpandConcatFCInferMeta(const std::vector<const MetaTensor*>& x,
                                       MetaTensor* out,
                                       MetaTensor* fc_out);
 
+void FusedStackTransposeQuantInferMeta(const std::vector<const MetaTensor*>& x,
+                                       MetaTensor* out,
+                                       MetaTensor* scale);
+
+void FusedStackQuantInferMeta(const std::vector<const MetaTensor*>& x,
+                              MetaTensor* out,
+                              MetaTensor* scale);
+
 void FusedBiasDropoutResidualLnInferMeta(
     const MetaTensor& x,
     const MetaTensor& residual,
 
@@ -70,7 +70,9 @@ if(((WITH_GPU) AND (CUDA_VERSION VERSION_LESS 12.0))
     "legacy/gpu/moe_gate_dispatch_kernel.cu"
     "legacy/gpu/moe_gate_dispatch_grad_kernel.cu"
     "legacy/gpu/int_bincount.cu"
-    "legacy/gpu/layer_norm_cuda_kernel.cu")
+    "legacy/gpu/layer_norm_cuda_kernel.cu",
+    "fusion/gpu/fused_stack_transpose_quant_kernel.cu",
+    "fusion/gpu/fused_stack_quant_kernel.cu")
 endif()
 
 if(NOT WITH_DGC)
 
@@ -0,0 +1,121 @@
+// 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/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/fusion/gpu/fused_stack_transpose_quant.h"
+
+namespace phi {
+namespace fusion {
+
+template <typename OutT>
+__global__ void __launch_bounds__(1024)
+    FusedStackQuantGPUKernel(const int64_t* __restrict__ X_ptrs,
+                             OutT* __restrict__ out,
+                             float* __restrict__ scale,
+                             size_t M,
+                             size_t K,
+                             FastDiv K_div_128) {
+  size_t block_y = K_div_128.Div(blockIdx.x);
+  size_t block_x = blockIdx.x - block_y * (K / 128);
+
+  // Load 128x128 elements from X
+  __nv_bfloat16 input[4][4];
+  BlockLoad(X_ptrs, input, K, block_y, block_x);
+
+  // Find the maximum in all elements
+  __nv_bfloat16 amax = BlockReduceMax(input);
+
+  // Compute scale and store back
+  float scale_inv = ComputeScale<__nv_bfloat16, OutT>(amax, 0.0f);
+  float scale_out = __frcp_rn(scale_inv);
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    size_t idx_n = blockIdx.z;
+    size_t idx_m = block_y;
+    size_t idx_k = block_x;
+    size_t idx = (idx_n * (M / 128) + idx_m) * (K / 128) + idx_k;
+    scale[idx] = scale_out;
+  }
+
+  // Scale X and store to out
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_n = blockIdx.z;
+    size_t idx_m = block_y * 128 + threadIdx.y + i * 32;
+    size_t idx_k = block_x * 128 + threadIdx.x * 4;
+    size_t idx = (idx_n * M + idx_m) * K + idx_k;
+
+    using StoreT = VecType<OutT, 4>;
+    StoreT data;
+    for (int j = 0; j < 4; j++) {
+      float input_fp32 = static_cast<float>(input[i][j]);
+      float output_scaled = input_fp32 * scale_inv;
+      data[j] = static_cast<OutT>(output_scaled);
+    }
+    *reinterpret_cast<StoreT*>(out + idx) = data;
+  }
+}
+
+template <typename T, typename Context>
+void FusedStackQuantKernel(const Context& dev_ctx,
+                           const std::vector<const DenseTensor*>& x,
+                           DenseTensor* out,
+                           DenseTensor* scale) {
+  int N = static_cast<int>(x.size());
+
+  // zero sized tensor case
+  if (x[0]->numel() == 0) {
+    dev_ctx.template Alloc<phi::dtype::float8_e4m3fn>(out);
+    dev_ctx.template Alloc<float>(scale);
+    return;
+  }
+
+  int64_t M = x[0]->dims()[0];
+  int64_t K = x[0]->dims()[1];
+
+  DenseTensor x_ptrs_cpu;
+  x_ptrs_cpu.Resize({N});
+
+  int64_t* cpu_data = dev_ctx.template HostAlloc<int64_t>(&x_ptrs_cpu);
+  for (int64_t i = 0; i < N; i++) {
+    cpu_data[i] = reinterpret_cast<int64_t>(x[i]->data<T>());
+  }
+  DenseTensor x_ptrs_gpu;
+  int64_t* x_ptrs_gpu_data = dev_ctx.template Alloc<int64_t>(&x_ptrs_gpu);
+  phi::Copy(dev_ctx, x_ptrs_cpu, dev_ctx.GetPlace(), true, &x_ptrs_gpu);
+
+  dim3 grid((M / 128) * (K / 128), 1, N);
+  dim3 block(32, 32);
+  auto* out_data = dev_ctx.template Alloc<phi::dtype::float8_e4m3fn>(out);
+  auto* scale_data = dev_ctx.template Alloc<float>(scale);
+  FusedStackQuantGPUKernel<phi::dtype::float8_e4m3fn>
+      <<<grid, block, 0, dev_ctx.stream()>>>(x_ptrs_gpu.data<int64_t>(),
+                                             out_data,
+                                             scale_data,
+                                             M,
+                                             K,
+                                             FastDiv(K / 128));
+}
+
+}  // namespace fusion
+}  // namespace phi
+
+PD_REGISTER_KERNEL(fused_stack_quant,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::fusion::FusedStackQuantKernel,
+                   phi::dtype::bfloat16) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT8_E4M3FN);
+  kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
+}
@@ -0,0 +1,121 @@
+// 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.
+#ifndef PADDLE_PHI_KERNELS_FUSION_GPU_FUSED_STACK_TRANSPOSE_QUANT_H_
+#define PADDLE_PHI_KERNELS_FUSION_GPU_FUSED_STACK_TRANSPOSE_QUANT_H_
+
+#include "paddle/phi/kernels/fusion/gpu/quant_utils.h"
+
+template <typename T, int VecSize>
+struct __align__(sizeof(T) * VecSize) VecType {
+  T val[VecSize];
+  __host__ __device__ inline T& operator[](size_t i) { return val[i]; }
+  __host__ __device__ inline const T& operator[](size_t i) const {
+    return val[i];
+  }
+};
+
+struct FastDiv {
+  FastDiv() {}
+  explicit FastDiv(
+      uint64_t d) {  // Single-parameter constructors should be marked explicit.
+    for (shift_val = 0; shift_val < 64; ++shift_val) {
+      uint64_t shift_limit = uint64_t(1) << shift_val;
+      if (shift_limit >= d) break;
+    }
+
+    // quotient = ((uint128_t)n_hi << 64) / d
+    uint64_t quotient = 0;
+    uint64_t n_hi = (uint64_t(1) << shift_val) - d, n_lo = 0;
+    for (int i = 63; i >= 0; --i) {
+      uint64_t d_hi = i == 0 ? 0 : d >> (64 - i);
+      uint64_t d_lo = d << i;
+      if (n_hi == 0 && n_lo == 0) break;
+      if ((d_hi < n_hi) || (d_hi <= n_hi && d_lo <= n_lo)) {
+        quotient |= uint64_t(1) << i;
+        n_hi -= d_hi + (d_lo > n_lo);
+        n_lo -= d_lo;
+      }
+    }
+    multiplier = quotient + 1;
+  }
+
+  __device__ uint64_t Div(uint64_t n) const {
+    uint64_t t = __umul64hi(n, multiplier);
+    return (t + n) >> shift_val;
+  }
+
+  int shift_val;
+  uint64_t multiplier;
+};
+
+__device__ __forceinline__ void BlockLoad(const int64_t* __restrict__ X_ptrs,
+                                          __nv_bfloat16 input[4][4],
+                                          size_t K,
+                                          size_t block_y,
+                                          size_t block_x) {
+  const __nv_bfloat16* X =
+      reinterpret_cast<const __nv_bfloat16*>(X_ptrs[blockIdx.z]);
+
+  for (size_t i = 0; i < 4; i++) {
+    size_t idx_m = block_y * 128 + threadIdx.y + i * 32;
+    size_t idx_k = block_x * 128 + threadIdx.x * 4;
+    size_t idx = idx_m * K + idx_k;
+
+    using LoadT = VecType<__nv_bfloat16, 4>;
+    LoadT data = *reinterpret_cast<const LoadT*>(X + idx);
+    for (int j = 0; j < 4; j++) {
+      input[i][j] = data[j];
+    }
+  }
+}
+
+__device__ __forceinline__ __nv_bfloat16 WarpReduceMax(__nv_bfloat16 x) {
+  for (int offset = 16; offset > 0; offset /= 2) {
+    __nv_bfloat16 t = __shfl_down_sync(0xffffffff, x, offset);
+    x = BF16_MAX(x, t);
+  }
+  return x;
+}
+
+__device__ __forceinline__ __nv_bfloat16
+BlockReduceMax(__nv_bfloat16 input[4][4]) {
+  // [(4), 32, 32, (4)] => [32, 32]
+  __nv_bfloat16 local_max;
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j++) {
+      __nv_bfloat16 t = BF16_ABS(input[i][j]);
+      local_max = (i == 0 && j == 0) ? t : BF16_MAX(local_max, t);
+    }
+  }
+
+  // [32, (32)] => [32]
+  __nv_bfloat16 warp_max = WarpReduceMax(local_max);
+
+  // [(32)] => [1]
+  __shared__ __nv_bfloat16 block_max[32];
+  if (threadIdx.x == 0) {
+    block_max[threadIdx.y] = warp_max;
+  }
+  __syncthreads();
+  if (threadIdx.y == 0) {
+    warp_max = WarpReduceMax(block_max[threadIdx.x]);
+    if (threadIdx.x == 0) {
+      block_max[0] = warp_max;
+    }
+  }
+  __syncthreads();
+
+  return block_max[0];
+}
+#endif  // PADDLE_PHI_KERNELS_FUSION_GPU_FUSED_STACK_TRANSPOSE_QUANT_H_