save code

Author: xiaolil1

csrc/xpu_cutlass_fusion.cpp

@@ -232,7 +232,7 @@ inline float dDequantizeNF4(unsigned char val) {
                      ? BlockIdxX() : BlockIdxY();
     const int l_coord = BlockIdxZ();
 
-#if 0 
+#if 1 
     //float* quant_map;
    //static constexpr std::array<float, 16> quant_map{};
    // {
@@ -277,7 +277,7 @@ inline float dDequantizeNF4(unsigned char val) {
     Tensor mma_A = make_tensor<ElementMMA>(make_fragment_layout(params.tiled_copy_a, tCgA(_,_,_,0).shape()));
     Tensor mma_B = make_tensor<ElementMMA>(make_fragment_layout(params.tiled_copy_b, tCgB(_,_,_,0).shape()));
 
-#if 0 //SLM: 0, register: 1
+#if 1 //SLM: 0, register: 1
   #if 1 //fragement register
 	  Tensor dequant_frag = make_tensor<ElementB>(mma_B.layout());
   #else //common register
@@ -324,7 +324,7 @@ inline float dDequantizeNF4(unsigned char val) {
 	  const int k_start_idx = crd2idx((*k_tile_iter), make_shape(params.k));
     int prefetch_k = k_start_idx;
 
-#if 1 //SLM
+#if 0 //SLM
   #if 1
   auto dequant = [&] (int k_tile) {
     constexpr int N = decltype(cute::size<1>(mma_B))::value;
@@ -386,11 +386,10 @@ printf("src_compress_size = %d, dst_compress_size = %d, src_vec_size = %d, dst_v
         constexpr int N = decltype(cute::size<1>(mma_B))::value;
         constexpr int K = decltype(cute::size(mma_B))::value / N;
 
- 
         using src_compress_type = uint64_t;
         using dst_compress_type = uint64_t;
         constexpr int src_compress_size = cute::sizeof_bits_v<src_compress_type> / cute::sizeof_bits_v<ElementB>; //16
-        constexpr int dst_compress_size = cute::sizeof_bits_v<dst_compress_type> / cute::sizeof_bits_v<ElementMMA>; //16
+        constexpr int dst_compress_size = cute::sizeof_bits_v<dst_compress_type> / cute::sizeof_bits_v<ElementMMA>; //4
         constexpr int src_vec_size = (K / src_compress_size) >= 16 ? 16 : K / src_compress_size; //4, 16 -> max vec_size of sycl::vec
         constexpr int dst_vec_size = (K / dst_compress_size) >= 16 ? 16 : K / dst_compress_size; //16, 16 -> max vec_size of sycl::vec
         constexpr int src_loop_num = K / src_vec_size / src_compress_size;
@@ -399,11 +398,11 @@ printf("src_compress_size = %d, dst_compress_size = %d, src_vec_size = %d, dst_v
         //if(cute::thread0()) printf("params.group_size = %d, k_reload_factor = %d, k_tile_count = %d, N = %d, K = %d, src_compress_size = %d, src_vec_size = %d, dst_compress_size = %d, dst_vec_size = %d\n",params.group_size, k_reload_factor, k_tile_count, N, K, src_compress_size, src_vec_size, dst_compress_size, dst_vec_size);
 
         src_compress_type src[src_vec_size];
-        ElementMMA dst[dst_compress_size * dst_vec_size];
+        ElementMMA dst[dst_loop_num * dst_compress_size * dst_vec_size];
 
         #pragma unroll
         for (int n = 0; n < N; n++) {
-          float scale_value = fragment_scale(n);
+          //float scale_value = fragment_scale(n);
           #pragma unroll
           for (int l = 0; l < src_loop_num; l++) {
             //src_compress_type src[src_vec_size];
@@ -412,18 +411,19 @@ printf("src_compress_size = %d, dst_compress_size = %d, src_vec_size = %d, dst_v
             #pragma unroll
             for (int v = 0; v < src_vec_size; v++) {
               src_compress_type src_value = src[v];
-              int dst_idx = v * src_compress_size;
+              int dst_base_idx = l * src_vec_size * src_compress_size + v * src_compress_size;
               #pragma unroll
               for (int c = 0; c < src_compress_size; c++) {
                   uint8_t bit_value = (src_value >> (4 * (((c + 1) & 1) + (c >> 1) * 2))) & 0xF;
-                  dst[dst_idx + c] = static_cast<ElementMMA>(quant_map[bit_value] * scale_value);
+                  float scale_value = fragment_scale(n * (BLK_K / GROUP_SIZE) + (dst_base_idx + c) / GROUP_SIZE);
+                  dst[dst_base_idx + c] = static_cast<ElementMMA>(quant_map[bit_value] * scale_value);
               }
             }
           }
 
           #pragma unroll
           for (int l = 0; l < dst_loop_num; l++) {
-            reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(cute::raw_pointer_cast(mma_B.data()))[n*dst_loop_num + l] = reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(dst)[0];
+            reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(cute::raw_pointer_cast(mma_B.data()))[n*dst_loop_num + l] = reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(dst)[l];
           }
         }
       };
@@ -454,9 +454,9 @@ printf("src_compress_size = %d, dst_compress_size = %d, src_vec_size = %d, dst_v
     }
 
     for (int k_tile = k_start_idx, k_s = 0; k_tile < k_tile_count; k_tile++, k_s++, prefetch_k++) {
-#if 0 //SLM: 0, register: 1     
+#if 1 //SLM: 0, register: 1     
       copy(params.tiled_copy_b, tBgB(_,_,_,k_tile), frag_copy_B);
-      copy(params.tiled_copy_scale, tSgS(_, _, _, (k_start_idx + k_s) / k_reload_factor), frag_copy_Scale);
+      copy(params.tiled_copy_scale, tSgS(_, _, _, (k_start_idx + k_s) * BLK_K/params.group_size), frag_copy_Scale);
       copy(params.tiled_copy_a, tAgA(_,_,_,k_tile), frag_copy_A);
       dequant();
 #else
@@ -501,8 +501,11 @@ void gemm_4bit_cutlass(int m, int n, int k, int l, T *A, unsigned char *B,
 
   using GemmKernel = gemm_4bit_cutlass_kernel<T, BITS>;
 
-  //static constexpr int smem_size= (16) * sizeof(float);
+#if 1
+  static constexpr int smem_size= (16) * sizeof(float);
+#else  
   static constexpr int smem_size = BLK_N * BLK_K * sizeof(ElementMMA) * 2 * 2; //aligned with 128B and will be reused for dequant src and dst.
+#endif  
   size_t max_slm_size = q.get_device().get_info<sycl::info::device::local_mem_size>();
   assert(smem_size <= max_slm_size);