Add dequant gemm fusion kernel with cutlass

Author: xiaolil1

CMakeLists.txt

@@ -28,7 +28,7 @@ set(CUDA_FILES csrc/ops.cu csrc/kernels.cu)
 set(HIP_FILES csrc/ops.hip csrc/kernels.hip)
 set(MPS_FILES csrc/mps_ops.mm)
 set(METAL_FILES csrc/mps_kernels.metal)
-set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp csrc/xpu_cutlass.cpp csrc/xpu_cutlass-cute.cpp)
+set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp csrc/xpu_cutlass.cpp csrc/xpu_cutlass-cute.cpp csrc/xpu_cutlass_fusion.cpp)
 # C++ sources are always included
 list(APPEND SRC_FILES ${CPP_FILES})
 

build_xpu.sh

@@ -3,4 +3,4 @@ cmake -DCMAKE_CXX_STANDARD=17 -DSYCL_INTEL_TARGET=ON -DENABLE_INTEL_XMX=ON -DCUT
 #cmake -DCOMPUTE_BACKEND=xpu -DCMAKE_BUILD_TYPE=Debug -DCMAKE_CXX_FLAGS_DEBUG="-O0 -g" -S .
 #cmake --build . --config Release -j
 make
-pip install -e .
+#pip install -e .

csrc/pythonInterface.cpp

@@ -379,6 +379,13 @@ void gemv_4bit_inference_fp16(
     //gemv_4bit_fusion<sycl::half, 16>(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
 }
 
+void gemm_4bit_inference_bf16(
+    int m, int n, int k, sycl::ext::oneapi::bfloat16 * A,  unsigned char* B,  float *absmax, float *datatype, float * out,
+    int lda, int ldb, int ldc, int blocksize, sycl::queue* stream
+) {
+    gemm_4bit_inference_cutlass_dequant<sycl::ext::oneapi::bfloat16, 16>(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
+}
+
 void gemv_4bit_inference_bf16(
     int m, int n, int k,  sycl::ext::oneapi::bfloat16 * A,  sycl::ext::oneapi::bfloat16* B,  float *absmax, float *datatype,
     float * out,  int lda, int ldb, int ldc, int blocksize, sycl::queue* stream

csrc/xpu_cutlass.h

@@ -107,6 +107,11 @@ void gemv_4bit_inference_cutlass_cute(int m, int n, int k, T *A, T *B,
                          float *absmax, float *datatype, float *out, int lda,
                          int ldb, int ldc, int blocksize, sycl::queue *stream);
 
+template <typename T, int BITS>
+void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned char *B,
+                         float *absmax, float *datatype, float *out, int lda,
+                         int ldb, int ldc, int blocksize, sycl::queue *stream);
+
 template <typename T, int BITS>
 void gemv_4bit_inference_cutlass(int m, int n, int k, T *A, T *B,
                          float *absmax, float *datatype, float *out, int lda,