enable multi-batch

Author: xiaolil1

bitsandbytes/backends/xpu/ops.py

@@ -74,10 +74,14 @@ def _gemv_4bit_impl(
     blocksize: int,
     out: torch.Tensor,
 ) -> None:
+    #import pdb
     m = ct.c_int32(A.shape[-2])#ct.c_int32(1)
     n = ct.c_int32(shapeB[0])
     k = ct.c_int32(shapeB[1])
-    #import pdb
+    l = 1
+    #pdb.set_trace()
+    if A.dim() == 3:
+      l = A.shape[0]
     lda = m
     ldb = ct.c_int32((A.shape[-1] + 1) // 2)
     ldc = m
@@ -106,6 +110,7 @@ def _gemv_4bit_impl(
             m,
             n,
             k,
+            l,
             get_ptr(A),
             get_ptr(B),
             get_ptr(absmax),

csrc/pythonInterface.cpp

@@ -381,10 +381,10 @@ void gemv_4bit_inference_fp16(
 
 #if 1
 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 m, int n, int k, int l, 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);
+    gemm_4bit_inference_cutlass_dequant<sycl::ext::oneapi::bfloat16, 16>(m, n, k, l, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
 }
 #endif
 
@@ -826,10 +826,10 @@ void cgemv_4bit_inference_fp16(
 
 #if 1
 void cgemv_4bit_inference_bf16(
-    int m, int n, int k, sycl::ext::oneapi::bfloat16 * A,  unsigned char* B,   float *absmax, float *datatype,
+    int m, int n, int k, int l, 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_bf16(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
+    gemm_4bit_inference_bf16(m, n, k, l, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);
 }
 #else
 void cgemv_4bit_inference_bf16(

csrc/xpu_cutlass.h

@@ -108,7 +108,7 @@ void gemv_4bit_inference_cutlass_cute(int m, int n, int k, T *A, T *B,
                          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,
+void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, int l, T *A, unsigned char *B,
                          float *absmax, float *datatype, float *out, int lda,
                          int ldb, int ldc, int blocksize, sycl::queue *stream);
 

csrc/xpu_cutlass_fusion.cpp

@@ -163,7 +163,7 @@ class kgemm_4bit_inference_cutlass_dequant {
   };
 
   struct Params {
-    int m, n, k;
+    int m, n, k, l;
     T* A;
     uint8_t* B;
     float* out;
@@ -278,7 +278,7 @@ class kgemm_4bit_inference_cutlass_dequant {
     int M = params.m;
     int N = params.n;
     int K = params.k;
-    int L = 1;
+    int L = params.l;
 
     //Total Threads number
     static constexpr auto Num_SGs = ATOM_N * ATOM_M * ATOM_K; //32 //2
@@ -578,7 +578,7 @@ printf("\n");
 };
 
 template <typename T, int BITS>
-void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned char *B,
+void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, int l, T *A, unsigned char *B,
                          float *absmax_, float *datatype, float *out, int lda,
                          int ldb, int ldc, int blocksize, sycl::queue *stream) {
   ////std::cout<<"this is gemm_4bit_inference_cutlass_dequant ......................!!!!!!\n";
@@ -599,7 +599,7 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
 
   //static constexpr int smem_size= 512; // (16 * 32) for quant_map
   static constexpr int smem_size= 256; // (16 * 16) for quant_map
-  int l = 1;
+  //int l = 1;
 
   auto problem_size = ProblemShape{m, n, k, l};
 
@@ -610,6 +610,7 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
   params.m = m;
   params.n = n;
   params.k = k;
+  params.l = l;
   params.A = A;
   params.B = B;
   params.out = out;
@@ -701,7 +702,7 @@ void gemm_4bit_inference_cutlass_dequant(int m, int n, int k, T *A, unsigned cha
 }
 
 template void gemm_4bit_inference_cutlass_dequant<sycl::ext::oneapi::bfloat16, 16>(
-    int m, int n, int k, sycl::ext::oneapi::bfloat16 *A, unsigned char *B,
+    int m, int n, int k, int l, 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);