add draft test for simple gemm kernel

Author: xiaolil1

bitsandbytes/backends/xpu/ops.py

@@ -74,7 +74,9 @@ def _gemv_4bit_impl(
     blocksize: int,
     out: torch.Tensor,
 ) -> None:
-    m = ct.c_int32(1)
+    import pdb
+    pdb.set_trace()
+    m = ct.c_int32(*A.shape[:-1])
     n = ct.c_int32(shapeB[0])
     k = ct.c_int32(shapeB[1])
 
@@ -180,7 +182,7 @@ def _(
         blocksize: int,
     ) -> torch.Tensor:
         shape = (*A.shape[:-1], shapeB[0])
-        out = torch.empty(shape, device=A.device, dtype=A.dtype)
+        out = torch.empty(shape, device=A.device, dtype=torch.float32)
         _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
         return out
 

bitsandbytes/functional.py

@@ -1630,39 +1630,39 @@ def gemv_4bit(
         raise ValueError("state cannot be None. gemv_4bit() requires the state from quantize_4bit()")
 
     absmax = state.absmax
-    if state.nested:
-        absmax = dequantize_blockwise(absmax, state.state2) + state.offset
-
-    if getattr(state, "ipex", False) and state.quant_type == "nf4":
-        # compute_dtype: 1 indicates fp16, 2 indicates bf16
-        compute_dtype = 2 if A.dtype == torch.bfloat16 else 1
-        out = torch.ops.torch_ipex.woq_linear(
-            A,
-            B,
-            "nf4",
-            state.shape,
-            state.new_scales,
-            state.new_zeros,
-            None,
-            None,
-            state.blocksize,
-            compute_dtype,
-            1,
-            state.compensation,
-        )
-        return out
-
-    if out is not None:
-        torch.ops.bitsandbytes.gemv_4bit.out(
-            A,
-            B,
-            state.shape,
-            absmax,
-            state.code,
-            state.blocksize,
-            out=out,
-        )
-        return out
+    #if state.nested:
+    #    absmax = dequantize_blockwise(absmax, state.state2) + state.offset
+
+    #if getattr(state, "ipex", False) and state.quant_type == "nf4":
+    #    # compute_dtype: 1 indicates fp16, 2 indicates bf16
+    #    compute_dtype = 2 if A.dtype == torch.bfloat16 else 1
+    #    out = torch.ops.torch_ipex.woq_linear(
+    #        A,
+    #        B,
+    #        "nf4",
+    #        state.shape,
+    #        state.new_scales,
+    #        state.new_zeros,
+    #        None,
+    #        None,
+    #        state.blocksize,
+    #        compute_dtype,
+    #        1,
+    #        state.compensation,
+    #    )
+    #    return out
+
+    #if out is not None:
+    #    torch.ops.bitsandbytes.gemv_4bit.out(
+    #        A,
+    #        B,
+    #        state.shape,
+    #        absmax,
+    #        state.code,
+    #        state.blocksize,
+    #        out=out,
+    #    )
+    #    return out
 
     return torch.ops.bitsandbytes.gemv_4bit.default(
         A,

build_xpu.sh

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

csrc/pythonInterface.cpp

@@ -375,12 +375,12 @@ void gemv_4bit_inference_fp16(
     int m, int n, int k, sycl::half * A,  unsigned char* B,  float *absmax, float *datatype, sycl::half * out,
     int lda, int ldb, int ldc, int blocksize, sycl::queue* stream
 ) {
-    gemv_4bit_inference<sycl::half, 16>(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
+    //gemv_4bit_inference<sycl::half, 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,  unsigned char* B,  float *absmax, float *datatype,
-    sycl::ext::oneapi::bfloat16 * out,  int lda, int ldb, int ldc, int blocksize, sycl::queue* stream
+    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
 ) {
     gemv_4bit_inference<sycl::ext::oneapi::bfloat16, 16>(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
 }
@@ -389,7 +389,7 @@ void gemv_4bit_inference_fp32(
     int m, int n, int k, float * A,  unsigned char* B,  float *absmax, float *datatype, float * out, int lda,
     int ldb, int ldc, int blocksize, sycl::queue* stream
 ) {
-    gemv_4bit_inference<float, 32>(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
+    //gemv_4bit_inference<float, 32>(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
 }
 
 #endif
@@ -811,8 +811,8 @@ void cgemv_4bit_inference_fp16(
 }
 
 void cgemv_4bit_inference_bf16(
-    int m, int n, int k, sycl::ext::oneapi::bfloat16 * A,  unsigned char* B,  float *absmax, float *datatype,
-    sycl::ext::oneapi::bfloat16 * out,  int lda, int ldb, int ldc, int blocksize, sycl::queue* stream
+    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
 ) {
     gemv_4bit_inference_bf16(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
 }

csrc/xpu_cutlass.cpp

@@ -6,8 +6,8 @@
 #include <sycl/sycl.hpp>
 
 template <typename T, int BITS>
-void gemv_4bit_inference(int m, int n, int k, T *A, unsigned char *B,
-                         float *absmax, float *datatype, T *out, int lda,
+void gemv_4bit_inference(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) {
   std::cout<<"this is gemv_4bit_inference path !!!\n";
 
@@ -24,22 +24,6 @@ void gemv_4bit_inference(int m, int n, int k, T *A, unsigned char *B,
       typename TiledMMAHelper<MMA_Atom<XE_8x16x16_F32BF16BF16F32_TT>, Layout<TileShape>,
                                     Layout<Shape<_8, _4, _1>, Stride<_4, _1, _0>>>::TiledMMA;
 
-  // GEMM Mainloop - iteration over blocks in K dimension
-  using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
-          cutlass::gemm::MainloopIntelPVC<2>, //use PipelineStages = 2
-          TileShape,
-          bfloat16_t, // data_type of input: A
-          cutlass::gemm::TagToStrideA_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
-          bfloat16_t, // data_type of input: B
-          cutlass::gemm::TagToStrideB_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
-          TiledMma,
-          XE_2D_U16x32x32_LD_N, // 2D block copy operations used for A 
-          void, void, cute::identity,
-          XE_2D_U16x32x32_LD_V, // 2D block copy operations used for B
-          void, void, cute::identity
-  >;
-
-  //TODO: Remove the Epilogue as BNB may not need it.
   // Create the Epilogue
   using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<float /*data_type of GEMM output*/, ElementComputeEpilogue, ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
 
@@ -58,20 +42,35 @@ void gemv_4bit_inference(int m, int n, int k, T *A, unsigned char *B,
           XE_2D_U32x8x16_ST_N, // The copy atom used to store matrix D
           void, void>;
 
+  // GEMM Mainloop - iteration over blocks in K dimension
+  using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
+          cutlass::gemm::MainloopIntelPVC<2>, //use PipelineStages = 2
+          TileShape,
+          bfloat16_t, // data_type of input: A
+          cutlass::gemm::TagToStrideA_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
+          bfloat16_t, // data_type of input: B
+          cutlass::gemm::TagToStrideB_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
+          TiledMma,
+          XE_2D_U16x32x32_LD_N, // 2D block copy operations used for A
+          void, void, cute::identity,
+          XE_2D_U16x32x32_LD_V, // 2D block copy operations used for B
+          void, void, cute::identity
+  >;
+
   // Define the whole kernel (mainloop and epilogue)
   using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
           Shape<int, int, int, int>, // Defer global problem shape definition to runtime
           CollectiveMainloop,
           CollectiveEpilogue
   >;
 
-  auto problem_size = typename GemmKernel::ProblemShape{m, n, k, ldb};
+  auto problem_size = typename GemmKernel::ProblemShape{m, n, k, 1};
   auto [M, N, K, L]  = cute::append<4>(problem_size, 1);
 
   typename GemmKernel::StrideA stride_A = cutlass::make_cute_packed_stride(typename GemmKernel::StrideA{}, cute::make_shape(M, K, L));
   typename GemmKernel::StrideB stride_B = cutlass::make_cute_packed_stride(typename GemmKernel::StrideB{}, cute::make_shape(N, K, L));
   typename GemmKernel::StrideC stride_C = cutlass::make_cute_packed_stride(typename GemmKernel::StrideC{}, cute::make_shape(M, N, L));
-  typename GemmKernel::StrideD stride_D = cutlass::make_cute_packed_stride(typename GemmKernel::StrideD{}, cute::make_shape(M, N, L));
+  //typename GemmKernel::StrideD stride_D = cutlass::make_cute_packed_stride(typename GemmKernel::StrideD{}, cute::make_shape(M, N, L));
 
   //
   // Data members
@@ -80,12 +79,12 @@ void gemv_4bit_inference(int m, int n, int k, T *A, unsigned char *B,
   cutlass::DeviceAllocation<typename Gemm::ElementA> block_A;
   cutlass::DeviceAllocation<typename Gemm::ElementB> block_B;
   cutlass::DeviceAllocation<typename Gemm::ElementC> block_C;
-  cutlass::DeviceAllocation<typename CollectiveEpilogue::ElementOutput> block_D;
+  //cutlass::DeviceAllocation<typename CollectiveEpilogue::ElementOutput> block_D;
 
   block_A.reset(M * K * L);
   block_B.reset(K * N * L);
   block_C.reset(M * N * L);
-  block_D.reset(M * N * L);
+  //block_D.reset(M * N * L);
 
   //TODO: check whether need the fake data?
   uint64_t seed = 0;
@@ -95,40 +94,49 @@ void gemv_4bit_inference(int m, int n, int k, T *A, unsigned char *B,
 
   cutlass::KernelHardwareInfo hw_info;
   hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
-
+std::cout<<"log1 ...\n";
   float alpha=1.0; //Maybe will not be used by BNB, just keep it temporarily.
   float beta=0.f; //Maybe will not be used by BNB, just keep it temporarily.
   typename GemmKernel::Arguments arguments{
     cutlass::gemm::GemmUniversalMode::kGemm,
     problem_size,
+#if 1    
+    {(bfloat16_t*)A, stride_A, (bfloat16_t*)B, stride_B},
+    {{alpha, beta}, nullptr, stride_C, out, stride_C}, hw_info};
+#else
     {block_A.get(), stride_A, block_B.get(), stride_B},
-    {{alpha, beta}, block_C.get(), stride_C, block_D.get(), stride_D}, hw_info};
+    {{alpha, beta}, nullptr, stride_C, block_C.get(), stride_C}, hw_info};
+#endif
 
+std::cout<<"log2 ...\n";
   Gemm gemm_op;
 
   size_t workspace_size = Gemm::get_workspace_size(arguments);
   cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
 
   CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get()));
-
+std::cout<<"log3 ...\n";
   // Run the GEMM
   CUTLASS_CHECK(gemm_op.run());
 
+std::cout<<"log4 ...\n";
   syclcompat::wait();
+
+  //out = (T*)out;
 }
 
-template void gemv_4bit_inference<sycl::half, 16>(
-    int m, int n, int k, sycl::half *A, unsigned char *B, float *absmax,
-    float *datatype, sycl::half *out, int lda, int ldb, int ldc, int blocksize,
-    sycl::queue *stream);
+//template void gemv_4bit_inference<sycl::half, 16>(
+//    int m, int n, int k, sycl::half *A, unsigned char *B, float *absmax,
+//    float *datatype, sycl::half *out, int lda, int ldb, int ldc, int blocksize,
+//    sycl::queue *stream);
 template void gemv_4bit_inference<sycl::ext::oneapi::bfloat16, 16>(
-    int m, int n, int k, sycl::ext::oneapi::bfloat16 *A, unsigned char *B,
-    float *absmax, float *datatype, sycl::ext::oneapi::bfloat16 *out, int lda,
+    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);
-template void gemv_4bit_inference<float, 32>(int m, int n, int k, float *A,
-                                             unsigned char *B, float *absmax,
-                                             float *datatype, float *out,
-                                             int lda, int ldb, int ldc,
-                                             int blocksize,
-                                             sycl::queue *stream);
-
+//template void gemv_4bit_inference<float, 32>(int m, int n, int k, float *A,
+//                                             unsigned char *B, float *absmax,
+//                                             float *datatype, float *out,
+//                                             int lda, int ldb, int ldc,
+//                                             int blocksize,
+//                                             sycl::queue *stream);
+//

csrc/xpu_cutlass.h

@@ -71,6 +71,6 @@ class kgemv_4bit_inference_cutlass {
 #endif
 
 template <typename T, int BITS>
-void gemv_4bit_inference(int m, int n, int k, T *A, unsigned char *B,
-                         float *absmax, float *datatype, T *out, int lda,
+void gemv_4bit_inference(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);