PaddlePaddle
diff --git a/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+6-7Lines changed: 6 additions & 7 deletions b/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+6-7Lines changed: 6 additions & 7 deletions
diff --git a/‎paddle/phi/infermeta/multiary.cc
Copy file name to clipboard
+13-7Lines changed: 13 additions & 7 deletions b/‎paddle/phi/infermeta/multiary.cc
Copy file name to clipboard
+13-7Lines changed: 13 additions & 7 deletions
diff --git a/‎paddle/phi/infermeta/multiary.h
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions b/‎paddle/phi/infermeta/multiary.h
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+6-11Lines changed: 6 additions & 11 deletions b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+6-11Lines changed: 6 additions & 11 deletions
diff --git a/‎paddle/phi/kernels/cpu/quant_for_compress_kernel.cc
Copy file name to clipboard
+15-9Lines changed: 15 additions & 9 deletions b/‎paddle/phi/kernels/cpu/quant_for_compress_kernel.cc
Copy file name to clipboard
+15-9Lines changed: 15 additions & 9 deletions
diff --git a/‎paddle/phi/kernels/gpu/llm_int8_mat_mul_kernel.cu renamed to ‎paddle/phi/kernels/gpu/llm_int8_matmul_kernel.cu
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions b/‎paddle/phi/kernels/gpu/llm_int8_mat_mul_kernel.cu renamed to ‎paddle/phi/kernels/gpu/llm_int8_matmul_kernel.cu
Copy file name to clipboard
+5-5Lines changed: 5 additions & 5 deletions
diff --git a/‎paddle/phi/kernels/gpu/weight_only_mat_mul_kernel.cu
Copy file name to clipboard
-73Lines changed: 0 additions & 73 deletions b/‎paddle/phi/kernels/gpu/weight_only_mat_mul_kernel.cu
Copy file name to clipboard
-73Lines changed: 0 additions & 73 deletions
@@ -1367,14 +1367,14 @@
   data_transform :
     skip_transform : out_size, size_tensor, scale_tensor
 
-- op : llm_int8_mat_mul
+- op : llm_int8_matmul
   args : (Tensor x, Tensor weight, Tensor weight_scale, float threshold=6.0)
   output : Tensor(out)
   infer_meta :
-    func : LLMInt8MatMulInferMeta
+    func : LLMInt8MatmulInferMeta
     param : [x, weight]
   kernel :
-    func : llm_int8_mat_mul
+    func : llm_int8_matmul
     data_type : x
 
 - op : log
@@ -2602,14 +2602,13 @@
   intermediate: warprnntgrad
   backward : warprnnt_grad
 
-- op : weight_only_mat_mul
+- op : weight_only_matmul
   args : (Tensor x, Tensor weight, Tensor weight_scale)
   output : Tensor(out)
   infer_meta :
-    func : WeightOnlyMatMulInferMeta
-    param : [x, weight]
+    func : WeightOnlyMatmulInferMeta
   kernel :
-    func : weight_only_mat_mul
+    func : weight_only_matmul
     data_type : x
 
 - op : weighted_sample_neighbors
 
@@ -3572,7 +3572,7 @@ void WeightedSampleNeighborsInferMeta(const MetaTensor& row,
   out_count->set_dtype(DataType::INT32);
 }
 
-void LLMInt8MatMulInferMeta(const MetaTensor& x,
+void LLMInt8MatmulInferMeta(const MetaTensor& x,
                             const MetaTensor& weight,
                             MetaTensor* out) {
   auto x_dims = x.dims();
@@ -3595,25 +3595,31 @@ void LLMInt8MatMulInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
-void WeightOnlyMatMulInferMeta(const MetaTensor& x,
+void WeightOnlyMatmulInferMeta(const MetaTensor& x,
                                const MetaTensor& weight,
+                               const MetaTensor& weight_scale,
                                MetaTensor* out) {
   auto x_dims = x.dims();
   auto w_dims = weight.dims();
+  auto n = weight_scale.dims()[0];
   PADDLE_ENFORCE_EQ(
       w_dims.size(),
       2UL,
       errors::InvalidArgument("The input(weight) must be a 2D Tensor."));
+  PADDLE_ENFORCE_EQ(
+      weight_scale.dims().size(),
+      1UL,
+      errors::InvalidArgument("The input(weight_scale) must be a 1D Tensor."));
   PADDLE_ENFORCE_EQ(
       x_dims[x_dims.size() - 1],
-      w_dims[0],
+      w_dims[1],
       errors::InvalidArgument(
-          "Input(X) dim[-1] and Input(Weight) dim[0] should be euqal."
-          "But received Input(X) dim[-1](%s) != Input(Weight) dim[0](%s)",
+          "Input(X) dim[-1] and Input(Weight) dim[1] should be euqal."
+          "But received Input(X) dim[-1](%s) != Input(Weight) dim[1](%s)",
           x_dims[x_dims.size() - 1],
-          w_dims[0]));
+          w_dims[1]));
   auto out_dims = x_dims;
-  out_dims[out_dims.size() - 1] = w_dims[1];
+  out_dims[out_dims.size() - 1] = n;
   out->set_dims(out_dims);
   out->set_dtype(x.dtype());
 }
 
@@ -690,12 +690,13 @@ void FusedMultiHeadAttentionVariableInferMeta(const MetaTensor& query,
                                               bool causal,
                                               MetaTensor* out);
 
-void LLMInt8MatMulInferMeta(const MetaTensor& x,
+void LLMInt8MatmulInferMeta(const MetaTensor& x,
                             const MetaTensor& weight,
                             MetaTensor* out);
 
-void WeightOnlyMatMulInferMeta(const MetaTensor& x,
+void WeightOnlyMatmulInferMeta(const MetaTensor& x,
                                const MetaTensor& weight,
+                               const MetaTensor& weight_scale,
                                MetaTensor* out);
 
 void FusedRopeInferMeta(const MetaTensor& q,
 
@@ -5086,22 +5086,17 @@ void QuantForCompressInferMeta(const MetaTensor& x,
           x_dims[0]));
   std::vector<int64_t> dim_scale({x_dims[1]});
   std::vector<int64_t> dim_out;
-  if (layout == "weight_only") {
-    dim_out = std::vector<int64_t>({x_dims[0], x_dims[1]});
-  } else if (layout == "llm.int8") {
+  if (bits == 8) {
     dim_out = std::vector<int64_t>({x_dims[1], x_dims[0]});
+  } else if (bits == 4) {
+    dim_out = std::vector<int64_t>({x_dims[1] / 2, x_dims[0]});
   } else {
-    phi::errors::InvalidArgument(
-        "The layout must be weight_only or llm.int8, but got %s", layout);
+    phi::errors::InvalidArgument("The bit must be 8 or 4, but got %d", bits);
   }
   out->set_dims(phi::make_ddim(dim_out));
 
-  // TODO(lizhenyun) support weight_only int4
-  if (bits == 8) {
-    out->set_dtype(DataType::INT8);
-  } else {
-    phi::errors::Fatal("The bits only support 8, but got[%d]", bits);
-  }
+  out->set_dtype(DataType::INT8);
+
   scale->set_dims(phi::make_ddim(dim_scale));
   scale->set_dtype(DataType::FLOAT32);
 }
 
@@ -22,7 +22,7 @@
 
 namespace phi {
 
-template <typename DeviceContext, typename T, typename D>
+template <typename DeviceContext, typename T, typename D, int bits>
 void quant_compute(const DeviceContext& dev_ctx,
                    const DenseTensor& x,
                    DenseTensor* out,
@@ -59,15 +59,15 @@ void quant_compute(const DeviceContext& dev_ctx,
 
   per_channel_scale(scale_data, x_data, m, n);
 
-  per_channel_quant(x_int_data, x_data, scale_data, m, n);
+  per_channel_quant<T, bits>(x_int_data, x_data, scale_data, m, n);
   if (layout == "weight_only") {
-    permute_B_rows_for_mixed_gemm(
+    permute_B_rows_for_mixed_gemm<bits>(
         int_processed_data, x_int_data, std::vector<size_t>{m, n}, (int64_t)80);
-    row_major_to_column_major(
+    subbyte_transpose_impl<bits>(
         int_processed_2_data, int_processed_data, std::vector<size_t>{m, n});
-    interleave_column_major_tensor(
+    interleave_column_major_tensor<bits>(
         out_data, int_processed_2_data, std::vector<size_t>{m, n});
-    add_bias_and_interleave_int8s_inplace(out_data, num);
+    add_bias_and_interleave_inplace<bits>(out_data, num);
   } else if (layout == "llm.int8") {
     std::vector<int> axis = {1, 0};
     funcs::Transpose<DeviceContext, int8_t, 2> trans;
@@ -88,9 +88,16 @@ void QuantForCompressKernel(const Context& dev_ctx,
   if (bits == 8) {
     dev_ctx.template Alloc<int8_t>(out);
     dev_ctx.template Alloc<float>(scale);
-    quant_compute<Context, T, int8_t>(dev_ctx, x, out, scale, layout);
+    quant_compute<Context, T, int8_t, 8>(dev_ctx, x, out, scale, layout);
+  } else if (bits == 4 && layout == "weight_only") {
+    dev_ctx.template Alloc<int8_t>(out);
+    dev_ctx.template Alloc<float>(scale);
+    quant_compute<Context, T, int8_t, 4>(dev_ctx, x, out, scale, layout);
   } else {
-    phi::errors::Unimplemented("The bits only support 8, but got[%d]", bits);
+    phi::errors::Unimplemented(
+        "The bits only support 8 or weight_only 4, but got[%s] [%d]",
+        layout,
+        bits);
   }
   // VLOG(0) << "x: " << x.dtype() << x;
   // VLOG(0) << "out: " << out->dtype() << *out;
@@ -102,5 +109,4 @@ PD_REGISTER_KERNEL(quant_for_compress,
                    CPU,
                    ALL_LAYOUT,
                    phi::QuantForCompressKernel,
-                   float,
                    phi::dtype::float16) {}
@@ -12,12 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/llm_int8_mat_mul_kernel.h"
+#include "paddle/phi/kernels/llm_int8_matmul_kernel.h"
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/core/kernel_registry.h"
 #ifndef PADDLE_WITH_HIP
-#include "paddle/phi/kernels/impl/llm_int8_mat_mul_kernel_impl.h"
+#include "paddle/phi/kernels/impl/llm_int8_matmul_kernel_impl.h"
 #endif
 
 namespace phi {
@@ -56,7 +56,7 @@ void llm_int8_compute(const Context& dev_ctx,
 }
 
 template <typename T, typename Context>
-void LLMInt8MatMulKernel(const Context& dev_ctx,
+void LLMInt8MatmulKernel(const Context& dev_ctx,
                          const DenseTensor& x,
                          const DenseTensor& weight,
                          const DenseTensor& weight_scale,
@@ -68,8 +68,8 @@ void LLMInt8MatMulKernel(const Context& dev_ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(llm_int8_mat_mul,
+PD_REGISTER_KERNEL(llm_int8_matmul,
                    GPU,
                    ALL_LAYOUT,
-                   phi::LLMInt8MatMulKernel,
+                   phi::LLMInt8MatmulKernel,
                    phi::dtype::float16) {}