[Metax] optimize wint4 quantization implementation

backends/metax_gpu/kernels/impl/metax_weight_quantize_kernel_impl.h

@@ -25,26 +25,26 @@
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
-
+template <typename DataType>
 void show_2d_cpu_tensor(const DenseTensor& tensor,
                         const int64_t row_num = 3,
                         const int64_t col_num = 3) {
   const int64_t rows = tensor.dims()[0];
   const int64_t cols = tensor.dims()[1];
   printf("\nTensor shape = [%d, %d]\n", rows, cols);
 
-  const int8_t* cpu_ptr = tensor.data<int8_t>();
+  const DataType* cpu_ptr = tensor.data<DataType>();
 
   for (int r = 0; r < row_num; r++) {
     for (int c = 0; c < col_num; c++) {
-      int8_t val = *(cpu_ptr + r * cols + c);
-      printf("%d ", val);
+      DataType val = *(cpu_ptr + r * cols + c);
+      printf("%#x ", val);
     }
     printf("\n");
   }
   printf("\n\n");
 }
-
+template <typename DataType>
 void show_2d_gpu_tensor(const CustomContext& dev_ctx,
                         const DenseTensor& tensor,
                         const int64_t row_num = 3,
@@ -58,18 +58,39 @@ void show_2d_gpu_tensor(const CustomContext& dev_ctx,
   const int64_t cols = cpu_tensor.dims()[1];
   printf("\nTensor shape = [%d, %d]\n", rows, cols);
 
-  const int8_t* cpu_ptr = cpu_tensor.data<int8_t>();
+  const DataType* cpu_ptr = cpu_tensor.data<DataType>();
 
   for (int r = 0; r < row_num; r++) {
     for (int c = 0; c < col_num; c++) {
-      int8_t val = *(cpu_ptr + r * cols + c);
-      printf("%d ", val);
+      DataType val = *(cpu_ptr + r * cols + c);
+      printf("%#x ", val);
     }
     printf("\n");
   }
   printf("\n\n");
 }
 
+template <typename DataType>
+void show_1d_gpu_tensor(const CustomContext& dev_ctx,
+                        const DenseTensor& tensor,
+                        const int64_t num = 3) {
+  phi::CPUPlace cpu_place;
+
+  DenseTensor cpu_tensor;
+  phi::Copy(dev_ctx, tensor, cpu_place, true, &cpu_tensor);
+
+  const int64_t nums = cpu_tensor.numel();
+  printf("\nTensor shape = [%d]\n", nums);
+
+  const DataType* cpu_ptr = cpu_tensor.data<DataType>();
+
+  for (int n = 0; n < num; n++) {
+    DataType val = *(cpu_ptr + n);
+    printf("%#x ", val);
+  }
+  printf("\n\n");
+}
+
 void cpu_2d_tensor_transpose(const DenseTensor& input_data,
                              DenseTensor* transposed_data) {
   const int64_t input_data_rows = input_data.dims()[0];

backends/metax_gpu/kernels/metax_kernel/weight_quantize_kernel_register.cu

@@ -116,7 +116,7 @@ void WeightQuantizeKernel(const Context& dev_ctx,
     dev_ctx.template Alloc<T>(scale);
     weight_quant_gpu<T, Context>(dev_ctx,
                                  x.data<T>(),
-                                 out->data<int8_t>(),
+                                 quanted_x.data<int8_t>(),
                                  scale->data<T>(),
                                  weight_shape,
                                  arch,
@@ -141,7 +141,13 @@ void WeightQuantizeKernel(const Context& dev_ctx,
     //                             arch,
     //                             algo);
 #endif
-    MetaxQuantizedWeightLayoutTrans<Context>(dev_ctx, algo, weight_shape, out);
+    quanted_x.Resize({m / 2, n});
+
+    std::vector<int> axis = {1, 0};
+    funcs::Transpose<Context, int8_t, 2> trans;
+    trans(dev_ctx, quanted_x, out, axis);
+
+    out->Resize({n / 2, m});
   } else if (algo == "w4a8") {
     weight_permute_gpu_w4a8<Context>(dev_ctx,
                                      x.data<int8_t>(),