support fp_zp quant

Author: zhewang1-intc

include/common/core/common_types.hpp

@@ -27,7 +27,11 @@ enum class grf_mode : uint8_t { normal = 0, double_grf = 1 };
 
 enum class mem_layout : uint8_t { row_major = 0, col_major = 1 };
 
-enum class quant_mode : uint8_t { S4_ASYM = 0, S4_FULLRANGE_NO_ZP = 1 };
+enum class quant_mode : uint8_t {
+  S4_ASYM = 0,
+  S4_FULLRANGE_NO_ZP = 1,
+  INT4_ASYM_FP_ZERO = 2
+};
 
 struct quant_info {
   quant_mode quant_mode;

include/common/core/memory.hpp

@@ -355,13 +355,9 @@ __XETLA_API xetla_vector<T, N> xetla_load_global(
       __ESIMD_NS::cache_hint_L1<gpu::xetla::detail::get_cache_hint(L1H)>,
       __ESIMD_NS::cache_hint_L2<gpu::xetla::detail::get_cache_hint(L2H)>,
       __ESIMD_NS::alignment<alignment>};
-  if constexpr (sizeof(T) * N < sizeof(uint32_t) || N == 1) {
-    xetla_vector<T, N> ret;
-#pragma unroll
-    for (uint32_t i = 0; i < N; i++) {
-      ret[i] = ptr[i + byte_offset / sizeof(T)];
-    }
-    return ret;
+  if constexpr (sizeof(T) * N < sizeof(uint32_t)) {
+    xetla_vector<uint32_t, N> offsets(byte_offset, sizeof(T));
+    return __ESIMD_NS::gather<T, N, uint32_t>(ptr, offsets);
   } else {
     return __ESIMD_NS::block_load<T, N>(ptr, byte_offset, props);
   }
@@ -505,11 +501,9 @@ __XETLA_API void xetla_store_global(
       __ESIMD_NS::cache_hint_L2<gpu::xetla::detail::get_cache_hint(L2H)>,
       __ESIMD_NS::alignment<alignment>};
 
-  if constexpr (sizeof(T) * N < sizeof(uint32_t) || N == 1) {
-#pragma unroll
-    for (uint32_t i = 0; i < N; i++) {
-      ptr[i + byte_offset / sizeof(T)] = vals[i];
-    }
+  if constexpr (sizeof(T) * N < sizeof(uint32_t)) {
+    xetla_vector<uint32_t, N> offsets(byte_offset, sizeof(T));
+    return __ESIMD_NS::scatter<T, N, uint32_t>(ptr, offsets, vals);
   } else {
     __ESIMD_NS::block_store<T, N>(ptr, byte_offset, vals, props);
   }

include/experimental/group/gemm/impl/int4_dequantize_xe.hpp

@@ -101,12 +101,6 @@ class gemm_t<
       std::is_same<remove_const_t<dtype_b>, remove_const_t<int4x2>>::value ||
           std::is_same<remove_const_t<dtype_b>, remove_const_t<int4x8>>::value,
       "this is for 4bit matB ");
-  static_assert(
-      std::is_same<remove_const_t<dtype_zero_pt>, remove_const_t<int4x2>>::
-              value ||
-          std::is_same<remove_const_t<dtype_zero_pt>, remove_const_t<int4x8>>::
-              value,
-      "this is for 4bit zero_pt ");
 
   /******** set memory attribute **********/
   static constexpr mem_space mem_space_a = mem_desc_a_t::space;
@@ -284,12 +278,20 @@ class gemm_t<
       arch_tag>;
 
   // compress int4 along N dimensions
-  using zero_pt_tile_desc_t = subgroup::tile_desc_t<
-      (tile_size_x_b + pack_ratio - 1) / pack_ratio,
-      tile_size_y_zero_pt,
-      (block_size_x_b + pack_ratio - 1) / pack_ratio,
-      block_size_y_zero_pt,
-      reg_layout::tiled>;
+  using zero_pt_tile_desc_t = std::conditional_t<
+      quant_info_.quant_mode != quant_mode::INT4_ASYM_FP_ZERO,
+      subgroup::tile_desc_t<
+          (tile_size_x_b + pack_ratio - 1) / pack_ratio,
+          tile_size_y_zero_pt,
+          (block_size_x_b + pack_ratio - 1) / pack_ratio,
+          block_size_y_zero_pt,
+          reg_layout::tiled>,
+      subgroup::tile_desc_t<
+          tile_size_x_b,
+          tile_size_y_zero_pt,
+          block_size_x_b,
+          block_size_y_zero_pt,
+          reg_layout::tiled>>;
 
   using zero_pt_t = subgroup::tile_t<dtype_zero_pt, zero_pt_tile_desc_t>;
   using zero_pt_payload_t = subgroup::mem_payload_t<

include/experimental/kernel/gemm/impl/int4_dequantize_kslicing_xe.hpp

@@ -566,8 +566,7 @@ class gemm_universal_t<
     // check for int4x2
     implementable &=
         ((args.matB_ld % pack_ratio == 0) && (args.matrix_n % pack_ratio == 0));
-    if constexpr (
-        gemm_t::compute_policy::quant_mode != quant_mode::S4_FULLRANGE_NO_ZP) {
+    if constexpr (gemm_t::compute_policy::quant_mode == quant_mode::S4_ASYM) {
       implementable &= (args.zero_pt_ld % pack_ratio == 0);
     }
 
@@ -618,7 +617,10 @@ class gemm_universal_t<
     int start_x_scale = start_n;
     int start_y_scale = start_k / dequant_s;
 
-    int start_x_zero_pt = start_n / pack_ratio;
+    int start_x_zero_pt = gemm_t::compute_policy::quant_mode ==
+            quant_mode::INT4_ASYM_FP_ZERO
+        ? start_n
+        : start_n / pack_ratio;
     int start_y_zero_pt = start_k / dequant_s;
 
     // set up arguments
@@ -672,7 +674,8 @@ class gemm_universal_t<
           inner_loop_start,
           inner_loop_count,
           mem_desc_scale);
-    } else {
+    } else if constexpr (
+        gemm_t::compute_policy::quant_mode == quant_mode::S4_ASYM) {
       mem_desc_zero_pt_t mem_desc_zero_pt(
           args.zero_pt_base,
           {(args.matrix_n + pack_ratio - 1) / pack_ratio,
@@ -686,6 +689,24 @@ class gemm_universal_t<
           inner_loop_count,
           mem_desc_scale,
           mem_desc_zero_pt);
+    } else if constexpr (
+        gemm_t::compute_policy::quant_mode ==
+        quant_mode::INT4_ASYM_FP_ZERO) {
+      mem_desc_zero_pt_t mem_desc_zero_pt(
+          args.zero_pt_base,
+          {args.matrix_n,
+           ((args.matrix_k + dequant_s - 1) / dequant_s),
+           args.zero_pt_ld},
+          {start_x_zero_pt, start_y_zero_pt});
+      gemm_args = gemm_args_t(
+          mem_desc_a,
+          mem_desc_b,
+          inner_loop_start,
+          inner_loop_count,
+          mem_desc_scale,
+          mem_desc_zero_pt);
+    } else {
+      assert(0);
     }
     matAcc_t matAcc;
     matAcc.init(0);

include/subgroup/tile/impl/tile_op_functor.hpp

@@ -149,15 +149,26 @@ struct dequant_int4_weight_t {
               cvt_blk_i8.xetla_select<step, 1>(jj * block_size_y_b + ii) =
                   cvt_blk_i8.xetla_select<step, 1>(jj * block_size_y_b + ii) -
                   zero_pt_i8;
-            } else if constexpr (quant_mode == quant_mode::S4_FULLRANGE_NO_ZP) {
+            } else if constexpr (
+                quant_mode == quant_mode::S4_FULLRANGE_NO_ZP ||
+                quant_mode == quant_mode::INT4_ASYM_FP_ZERO) {
               cvt_blk_i8.xetla_select<step, 1>(jj * block_size_y_b + ii) =
                   cvt_blk_i8.xetla_select<step, 1>(jj * block_size_y_b + ii) -
                   int8_t(8);
             }
             dst_blk.xetla_select<step, 1>(jj * block_size_y_b + ii) =
                 cvt_blk_i8.xetla_select<step, 1>(jj * block_size_y_b + ii) *
                 scale.reg[scale_idx];
-
+            if constexpr (quant_mode == quant_mode::INT4_ASYM_FP_ZERO) {
+              uint32_t zero_pt_idx =
+                  offset_y_in_tile / dequant_s * zero_pt_t::block_size_x +
+                  offset_x_in_tile;
+              native_type_t<typename zero_pt_t::dtype> zero_pt_pack =
+                  zero_pt.reg[zero_pt_idx];
+              dst_blk.xetla_select<step, 1>(jj * block_size_y_b + ii) =
+                  dst_blk.xetla_select<step, 1>(jj * block_size_y_b + ii) +
+                  zero_pt_pack;
+            }
             // sycl::ext::oneapi::experimental::printf(
             //     "scale[%d] %f \n",
             //     scale_idx,

tests/integration/gemv/int4/main.cpp

@@ -40,7 +40,8 @@ class test_col_major_1 {
   static constexpr size_t sg_k = 512 / sg_m;
   static constexpr size_t dequant_s = 128;
   // static constexpr quant_mode quant_mode = quant_mode::S4_ASYM;
-  static constexpr quant_mode quant_mode = quant_mode::S4_FULLRANGE_NO_ZP;
+  // static constexpr quant_mode quant_mode = quant_mode::S4_FULLRANGE_NO_ZP;
+  static constexpr quant_mode quant_mode = quant_mode::INT4_ASYM_FP_ZERO;
 
   static constexpr size_t local_kslicing = 1;
   static constexpr size_t global_kslicing = 1;
@@ -131,13 +132,19 @@ std::vector<fp16> convert_int4(
     data_type_zero_pt zero_pt) {
   std::vector<fp16> dequant_fp16(sizeof(data_type_b) * 2);
 
-  int8_t zero_pt_i8 = zero_pt & 0xf;
+  int8_t zero_pt_i8;
+  if constexpr (quant_mode != quant_mode::INT4_ASYM_FP_ZERO)
+    zero_pt_i8 = zero_pt & 0xf;
   for (uint32_t i = 0; i < dequant_fp16.size(); i++) {
     int8_t dequant_8bit = data_b & 0xf;
     if constexpr (quant_mode == quant_mode::S4_FULLRANGE_NO_ZP) {
       dequant_fp16[i] = scale * (dequant_8bit - 8);
-    } else {
+    } else if constexpr (quant_mode == quant_mode::S4_ASYM) {
       dequant_fp16[i] = scale * (dequant_8bit - zero_pt_i8);
+    } else if constexpr (quant_mode == quant_mode::INT4_ASYM_FP_ZERO) {
+      dequant_fp16[i] = scale * (dequant_8bit - 8) + zero_pt;
+    } else {
+      assert(0);
     }
     data_b = data_b >> 4;
   }
@@ -169,15 +176,17 @@ std::vector<data_type_acc_in> dequantize_weight(
     for (uint32_t j = 0; j < width; j += step) {
       int start_b_in = i * width + j;
       int start_scale_in = start_b_in / step;
-      int start_zero_pt_in =
-          (j / step) * (matrix_n / pack_radio) + i / pack_radio;
+      int start_zero_pt_in = quant_mode == quant_mode::INT4_ASYM_FP_ZERO
+          ? (j / step) * matrix_n + i
+          : (j / step) * (matrix_n / pack_radio) + i / pack_radio;
       int start_out =
           layout_b == mem_layout::row_major ? 0 : i * matrix_k + j * pack_radio;
+      data_type_zero_pt zp_value = zero_pt[start_zero_pt_in];
+      if constexpr (quant_mode != quant_mode::INT4_ASYM_FP_ZERO)
+        zp_value = zp_value >> (4 * (i % pack_radio));
       for (uint32_t jj = 0; jj < step; jj++) {
         std::vector<fp16> dequant_fp16 = convert_int4<quant_mode>(
-            b[start_b_in + jj],
-            scale[start_scale_in],
-            zero_pt[start_zero_pt_in] >> (4 * (i % pack_radio)));
+            b[start_b_in + jj], scale[start_scale_in], zp_value);
         for (uint32_t jjj = 0; jjj < dequant_fp16.size(); jjj++) {
           b_out[start_out + pack_radio * jj + jjj] = dequant_fp16[jjj];
         }
@@ -215,7 +224,10 @@ void dequantize_gemv_run(int iter) {
   using data_type_a = typename Test::data_type_a;
   using data_type_b = typename Test::data_type_b;
   using data_type_c = typename Test::data_type_c;
-  using data_type_zero_pt = data_type_b;
+  using data_type_zero_pt = std::conditional_t<
+      Test::quant_mode == quant_mode::INT4_ASYM_FP_ZERO,
+      data_type_c,
+      data_type_b>;
   using data_type_scale = fp16;
   using data_type_acc_in = fp16;
   using data_type_acc = float;
@@ -225,7 +237,7 @@ void dequantize_gemv_run(int iter) {
   constexpr mem_layout layout_b = Test::layout_b;
 
   constexpr size_t size_a = matrix_m * matrix_k;
-  constexpr size_t size_b = matrix_k * matrix_n / (2 * sizeof(data_type_b));
+  constexpr size_t size_b = matrix_k * matrix_n / 2;
 
   constexpr size_t size_scale_k = matrix_k / dequant_s;
   constexpr size_t size_scale_n = matrix_n;
@@ -234,7 +246,9 @@ void dequantize_gemv_run(int iter) {
   constexpr size_t size_zero_pt_k = matrix_k / dequant_s;
   constexpr size_t size_zero_pt_n = matrix_n;
   constexpr size_t size_zero_pt =
-      size_zero_pt_k * size_zero_pt_n / (2 * sizeof(data_type_b));
+      Test::quant_mode != quant_mode::INT4_ASYM_FP_ZERO
+      ? size_zero_pt_k * size_zero_pt_n / 2
+      : size_zero_pt_k * size_zero_pt_n;
 
   constexpr size_t size_c = matrix_m * matrix_n;
   constexpr size_t size_bias = matrix_n;
@@ -405,16 +419,18 @@ void dequantize_gemv_run(int iter) {
     scale_h[i] = INFINITY;
   }
   for (unsigned i = 0; i < size_zero_pt + UNDEFINED_DATA_SIZE; ++i) {
-    if constexpr (std::is_same_v<int4x2, data_type_b>) {
+    if constexpr (std::is_same_v<int4x2, data_type_zero_pt>) {
       zero_pt_h[i] = random_uint8();
 #ifdef UT_DEBUG
       zero_pt_h[i] = 0x12 << i;
 #endif
-    } else if constexpr (std::is_same_v<int4x8, data_type_b>) {
+    } else if constexpr (std::is_same_v<int4x8, data_type_zero_pt>) {
       zero_pt_h[i] = random_uint32();
 #ifdef UT_DEBUG
       zero_pt_h[i] = 0x33333333;
 #endif
+    } else if constexpr (std::is_same_v<fp16, data_type_zero_pt>) {
+      zero_pt_h[i] = random_float();
     }
   }
 
@@ -491,7 +507,9 @@ void dequantize_gemv_run(int iter) {
             Acc_d,
             Cnt_d,
             epilogue_args);
-  } else if constexpr (compute_policy::quant_mode == quant_mode::S4_ASYM) {
+  } else if constexpr (
+      compute_policy::quant_mode == quant_mode::S4_ASYM ||
+      compute_policy::quant_mode == quant_mode::INT4_ASYM_FP_ZERO) {
     gemm_arg =
         typename gemm_op_t::template arguments_t<compute_policy::quant_mode>(
             matrix_m,