fix group_qkv

Author: sunjiweiswift

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

@@ -598,12 +598,8 @@ class gemm_universal_t<
     int start_n = group_swizzle.template get_tile_idx<2>(item) * wg_tile_n;
     int start_k = 0;
     uint32_t wg_tile_k = args.matrix_k;
-    uint32_t boundary_n = (start_n + wg_tile_n) > args.matrix_n
-        ? args.matrix_n
-        : (start_n + wg_tile_n);
-    uint32_t boundary_m = (start_m + wg_tile_m) > args.matrix_m
-        ? args.matrix_m
-        : (start_m + wg_tile_m);
+    uint32_t boundary_n = std::min(start_n + wg_tile_n, args.matrix_n);
+    uint32_t boundary_m = std::min(start_m + wg_tile_m, args.matrix_m);
     uint32_t boundary_k = wg_tile_k;
     if constexpr (num_global_kslicing > 1) {
       wg_tile_k = (wg_tile_k + num_global_kslicing - 1) / num_global_kslicing;

include/group/epilogue/impl/default_xe.hpp

@@ -95,6 +95,7 @@ class epilogue_t<
     using mat_tile_desc = typename matAcc_t::tile_desc;
     using matC_t = subgroup::tile_t<dtype_c, mat_tile_desc>;
 
+    // static constexpr msg_type msg_type_c = msg_type::unaligned_2d;
     static constexpr msg_type msg_type_c =
         subgroup::msg_type_v<mat_tile_desc, mem_desc_c_t>;
     using matC_payload_t = subgroup::
@@ -192,9 +193,9 @@ class epilogue_t<
     using mat_tile_desc = typename matAcc_t::tile_desc;
     using matC_t = subgroup::tile_t<dtype_c, mat_tile_desc>;
 
-    // static constexpr msg_type msg_type_c = msg_type::block_2d;
-    static constexpr msg_type msg_type_c =
-        subgroup::msg_type_v<mat_tile_desc, mem_desc_c_t>;
+    static constexpr msg_type msg_type_c = msg_type::block_2d;
+    // static constexpr msg_type msg_type_c =
+    //     subgroup::msg_type_v<mat_tile_desc, mem_desc_c_t>;
 
     using matC_payload_t = subgroup::
         mem_payload_t<mem_desc_c_t, mat_tile_desc, msg_type_c, arch_tag>;

include/kernel/gemm/impl/kslicing_xe.hpp

@@ -426,12 +426,8 @@ class gemm_universal_t<
     int start_n = group_swizzle.template get_tile_idx<2>(item) * wg_tile_n;
     int start_k = 0;
     uint32_t wg_tile_k = args.matrix_k;
-    uint32_t boundary_n = (start_n + wg_tile_n) > args.matrix_n
-        ? args.matrix_n
-        : (start_n + wg_tile_n);
-    uint32_t boundary_m = (start_m + wg_tile_m) > args.matrix_m
-        ? args.matrix_m
-        : (start_m + wg_tile_m);
+    uint32_t boundary_n = std::min(start_n + wg_tile_n,args.matrix_n);
+    uint32_t boundary_m = std::min(start_m + wg_tile_m, args.matrix_m);
     uint32_t boundary_k = wg_tile_k;
     if constexpr (num_global_kslicing > 1) {
       wg_tile_k = (wg_tile_k + num_global_kslicing - 1) / num_global_kslicing;

include/subgroup/tile/impl/payload_xe.hpp

@@ -435,9 +435,15 @@ struct mem_payload_t<
   uint64_t base_offset;
   dtype* base_ptr;
   uint32_t pitch_in_bytes;
+  uint32_t height_in_elems;
+  uint32_t width_in_elems;
+  uint32_t payload_bytes;
 
   inline mem_payload_t(mem_desc_t& mem_tdesc) {
     pitch_in_bytes = mem_tdesc.shape.stride * sizeof(dtype);
+    width_in_elems = mem_tdesc.shape.x;
+    height_in_elems = mem_tdesc.shape.y;
+    payload_bytes = width_in_elems * height_in_elems * sizeof(dtype);
     uint32_t offset_x = mem_tdesc.coord.x;
     uint32_t offset_y = mem_tdesc.coord.y;
     base_offset = mem_transpose
@@ -448,14 +454,17 @@ struct mem_payload_t<
 
   inline mem_payload_t(
       dtype* p,
-      [[maybe_unused]] int surface_width,
-      [[maybe_unused]] int surface_height,
+      int surface_width,
+      int surface_height,
       int surface_pitch,
       int surface_offset_x,
       int surface_offset_y) {
     pitch_in_bytes = surface_pitch * sizeof(dtype);
     uint32_t offset_x = surface_offset_x;
     uint32_t offset_y = surface_offset_y;
+    width_in_elems = surface_width;
+    height_in_elems = surface_height;
+    payload_bytes = width_in_elems * height_in_elems * sizeof(dtype);
     base_offset = mem_transpose
         ? offset_x * pitch_in_bytes + offset_y * sizeof(dtype)
         : offset_y * pitch_in_bytes + offset_x * sizeof(dtype);
@@ -466,6 +475,9 @@ struct mem_payload_t<
     pitch_in_bytes = mem_tdesc.shape.stride * sizeof(dtype);
     uint32_t offset_x = mem_tdesc.coord.x;
     uint32_t offset_y = mem_tdesc.coord.y;
+    width_in_elems = mem_tdesc.shape.x;
+    height_in_elems = mem_tdesc.shape.y;
+    payload_bytes = width_in_elems * height_in_elems * sizeof(dtype);
     base_offset = mem_transpose
         ? offset_x * pitch_in_bytes + offset_y * sizeof(dtype)
         : offset_y * pitch_in_bytes + offset_x * sizeof(dtype);
@@ -474,14 +486,17 @@ struct mem_payload_t<
 
   __XETLA_API void init(
       dtype* p,
-      [[maybe_unused]] int surface_width,
-      [[maybe_unused]] int surface_height,
+      int surface_width,
+      int surface_height,
       int surface_pitch,
       int surface_offset_x,
       int surface_offset_y) {
     pitch_in_bytes = surface_pitch * sizeof(dtype);
     uint32_t offset_x = surface_offset_x;
     uint32_t offset_y = surface_offset_y;
+    width_in_elems = surface_width;
+    height_in_elems = surface_height;
+    payload_bytes = width_in_elems * height_in_elems * sizeof(dtype);
     base_offset = mem_transpose
         ? offset_x * pitch_in_bytes + offset_y * sizeof(dtype)
         : offset_y * pitch_in_bytes + offset_x * sizeof(dtype);
@@ -492,13 +507,19 @@ struct mem_payload_t<
     this->base_offset = rhs.base_offset;
     this->base_ptr = rhs.base_ptr;
     this->pitch_in_bytes = rhs.pitch_in_bytes;
+    this->width_in_elems = rhs.width_in_elems;
+    this->height_in_elems = rhs.height_in_elems;
+    this->payload_bytes = rhs.payload_bytes;
   }
 
   inline mem_payload_t() = default;
   inline this_payload_t& operator=(const this_payload_t& rhs) {
     this->base_offset = rhs.base_offset;
     this->base_ptr = rhs.base_ptr;
     this->pitch_in_bytes = rhs.pitch_in_bytes;
+    this->width_in_elems = rhs.width_in_elems;
+    this->height_in_elems = rhs.height_in_elems;
+    this->payload_bytes = rhs.payload_bytes;
     return *this;
   }
 

include/subgroup/tile/impl/store_xe.hpp

@@ -283,34 +283,38 @@ tile_store(tile_t& tile, payload_t& payload) {
   static constexpr uint32_t store_len = tile_t::tile_elems;
   static constexpr gpu_arch arch_tag = payload_t::arch_tag;
 
-  using load_store_attr = load_store_attr_t<msg_type::block_1d, arch_tag>;
-  static constexpr uint32_t max_store_vec_len =
-      load_store_attr::max_store_vec_len;
-  static constexpr uint32_t max_store_vec_elems =
-      max_store_vec_len / sizeof(dtype);
+  if (payload.base_offset <= payload.payload_bytes) {
+    using load_store_attr = load_store_attr_t<msg_type::block_1d, arch_tag>;
+    static constexpr uint32_t max_store_vec_len =
+        load_store_attr::max_store_vec_len;
+    static constexpr uint32_t max_store_vec_elems =
+        max_store_vec_len / sizeof(dtype);
+    static constexpr uint32_t store_iter_steps =
+        store_len / max_store_vec_elems;
 
-  static constexpr uint32_t store_iter_steps = store_len / max_store_vec_elems;
-  if constexpr (store_len >= max_store_vec_elems) {
+    if constexpr (store_len >= max_store_vec_elems) {
 #pragma unroll
-    for (uint32_t i = 0; i < store_iter_steps; i++) {
-      uint32_t offset = i * max_store_vec_elems;
-      auto reg_sub = tile.reg.xetla_select<max_store_vec_elems, 1>(offset);
-      uint32_t address_offset = offset * sizeof(dtype);
+      for (uint32_t i = 0; i < store_iter_steps; i++) {
+        uint32_t offset = i * max_store_vec_elems;
+        auto reg_sub = tile.reg.xetla_select<max_store_vec_elems, 1>(offset);
+        uint32_t address_offset = offset * sizeof(dtype);
 
-      xetla_store_global<dtype, max_store_vec_elems, L1, L2>(
-          payload.base_ptr,
-          payload.base_offset + address_offset,
-          reg_sub.xetla_format<dtype>());
+        xetla_store_global<dtype, max_store_vec_elems, L1, L2>(
+            payload.base_ptr,
+            payload.base_offset + address_offset,
+            reg_sub.xetla_format<dtype>());
+      }
     }
+    constexpr uint32_t tail_len =
+        store_len % max_store_vec_elems * sizeof(dtype);
+    uint32_t tail_offset = store_iter_steps * max_store_vec_len;
+    detail::process_1d_tail<
+        tail_len,
+        (max_store_vec_len >> 1),
+        detail::process_flag::store,
+        L1,
+        L2>(tile, payload, tail_offset);
   }
-  constexpr uint32_t tail_len = store_len % max_store_vec_elems * sizeof(dtype);
-  uint32_t tail_offset = store_iter_steps * max_store_vec_len;
-  detail::process_1d_tail<
-      tail_len,
-      (max_store_vec_len >> 1),
-      detail::process_flag::store,
-      L1,
-      L2>(tile, payload, tail_offset);
 }
 
 /// @brief Is the func storing data from register file to unaligned global
@@ -348,7 +352,6 @@ tile_store(
   constexpr uint32_t num_channel_y = payload_t::num_channel_y;
   constexpr uint32_t store_elems = num_channel_y * payload_t::num_channel_x;
   constexpr uint32_t scale_factor = payload_t::scale_factor;
-
 #pragma unroll
   for (uint32_t i = 0; i < tile_desc::tile_size_y / tile_desc::block_size_y;
        i++) {

tests/integration/gemm/fp16/common.hpp

@@ -52,7 +52,7 @@ class TestBase {
 class Test0 : public TestBase {
  public:
   static constexpr size_t mat_m = 1;
-  static constexpr size_t mat_n = 1280;
+  static constexpr size_t mat_n = 64;
   static constexpr size_t mat_k = 8192;
   static constexpr size_t wg_m = 8;
   static constexpr size_t wg_n = 32;

tests/integration/gemv/int4/main.cpp

@@ -32,12 +32,12 @@ class test_col_major_1 {
   // Extract the parameters required by different test cases
   static constexpr size_t mat_m = 1;
   static constexpr size_t mat_n = 4096;
-  static constexpr size_t mat_k = 11008;
+  static constexpr size_t mat_k = 4096;
   static constexpr size_t wg_m = 1;
   static constexpr size_t wg_n = 1;
   static constexpr size_t sg_m = 1;
   static constexpr size_t sg_n = 1;
-  static constexpr size_t sg_k = 256 / 1;
+  static constexpr size_t sg_k = 1024 / 1;
   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;
@@ -47,7 +47,7 @@ class test_col_major_1 {
   static constexpr mem_layout layout_a = mem_layout::row_major;
   static constexpr mem_layout layout_b = mem_layout::col_major;
   static constexpr mma_engine mma_eng = mma_engine::fpu;
-  static constexpr gpu_arch arch = gpu_arch::XeHpc;
+  static constexpr gpu_arch arch = gpu_arch::XeLpg;
   using data_type_a = fp16;
   using data_type_b = int4x8;
   using data_type_c = fp16;