refine code

Author: xiaolil1

csrc/xpu_cutlass_fusion.cpp

@@ -241,11 +241,11 @@ class kgemm_4bit_inference_cutlass_dequant {
       //if(syclcompat::global_id::x() == 2 && syclcompat::global_id::y() ==0 && syclcompat::global_id::z() ==0 )
         //printf("syclcompat::global_id::x() = %d, syclcompat::global_id::y() = %d, syclcompat::global_id::z() = %d, thread_idx = %d, i = %d, in[i].ptr_ = %x, in[i].idx_=%x, value_bit = %x, value = %d, quant_map[value] = %f, out[i] = %f\n",syclcompat::global_id::x(), syclcompat::global_id::y(), syclcompat::global_id::z(), thread_idx, i, in[i].ptr_, in[i].idx_, value, static_cast<int>(value), quant_map[value], static_cast<float>(out[i]));
       }
-      int scale_number = decltype(size(tCrS_input))::value;
-      for(int i=0; i<scale_number; i++){
-        auto s_value = tCrS_input(i);
-        if(cute::thread0()) printf("scale_number = %d, tCrS_input[%d] = %f\n",scale_number, i, s_value);
-      }
+    }
+    int scale_number = decltype(size(tCrS_input))::value;
+    for(int i=0; i<scale_number; i++){
+      auto s_value = tCrS_input(i);
+      if(cute::thread0()) printf("scale_number = %d, tCrS_input[%d] = %f\n",scale_number, i, s_value);
     }
 #else    
     static constexpr auto N = decltype(size<1>(in))::value;
@@ -297,6 +297,7 @@ class kgemm_4bit_inference_cutlass_dequant {
     int N = params.n;
     int K = params.k;
     int L = 1;
+    static constexpr int BLK_K = 64;
 
     T* A = params.A;
     uint8_t* B = params.B;
@@ -383,11 +384,11 @@ class kgemm_4bit_inference_cutlass_dequant {
 
 	  Tensor dequant_frag = make_tensor<ElementB>(mma_B.layout());
 
-    static constexpr auto scale_traits_size = decltype(size(typename GmemTiledCopyScale::BlockShape{}))::value / SubgroupSize;
-    static constexpr auto scale_traits_num = SG_QNT_WIDTH / size<1>(typename GmemTiledCopyScale::BlockShape{});
+    static constexpr auto scale_traits_size = 16 / SubgroupSize;
+    static constexpr auto scale_traits_num = 64 / 16;
     using FragScaleLayout = Layout<Shape<Int<scale_traits_size>, Int<scale_traits_num>, _1>>;
     Tensor fragment_scale = make_tensor<ElementScale>(FragScaleLayout{});
-    if(cute::thread0()) printf("scale_traits_size = %d, scale_traits_num = %d, SG_QNT_WIDTH = %d\n", scale_traits_size, scale_traits_num, SG_QNT_WIDTH);
+    if(cute::thread0()) printf("scale_traits_size = %d, scale_traits_num = %d\n", scale_traits_size, scale_traits_num);
 
     static_assert(std::is_same_v<typename decltype(dequant_frag)::value_type, ElementQuant>);
     static_assert(std::is_same_v<typename decltype(mma_A)::value_type, ElementMMA>);
@@ -414,15 +415,15 @@ class kgemm_4bit_inference_cutlass_dequant {
 
 // Run mainloop
     auto [m_idx, n_idx, k_idx, l_idx] = blk_coord_mnkl;
-    const int n_coord_s = n_idx * BLK_N + (get_sub_group_id() % ATOM_N) * SG_N;
+    const int n_coord_s = n_idx * 64 + (get_sub_group_id() % 2) * 32;
     const int l_coord_s = l_idx;
 
     if(cute::thread0()) printf("m_idx = %d, n_idx = %d, k_idx = %d, l_idx = %d, n_coord_s = %d, l_coord_s = %d\n",m_idx, n_idx, k_idx, l_idx, n_coord_s, l_coord_s);
 
     auto copy_iter_s = [&](){
         return make_tensor(make_inttuple_iter(make_coord(n_coord_s, 0, l_coord_s)),
                           make_layout(make_shape(Int<scale_traits_size>{}, Int<scale_traits_num>{}, _1{}, k_tile_count),
-                                      make_stride(E<0>{} * _16{}, E<0>{} * size<1>(typename GmemTiledCopyScale::BlockShape{}), _0{}, E<1>{} * _1{})));
+                                      make_stride(E<0>{} * _16{}, E<0>{} * 16, _0{}, E<1>{} * _1{})));
 
     }();
 #if 1