fix scale issue

Author: xiaolil1

bitsandbytes/backends/xpu/ops.py

@@ -74,20 +74,15 @@ def _gemv_4bit_impl(
     blocksize: int,
     out: torch.Tensor,
 ) -> None:
-    import pdb
-    #pdb.set_trace()
-    m = ct.c_int32(*A.shape[:-1]) #A.shape[1])
+    m = ct.c_int32(1)
     n = ct.c_int32(shapeB[0])
     k = ct.c_int32(shapeB[1])
-
+    import pdb
     lda = m
     ldb = ct.c_int32((A.shape[-1] + 1) // 2)
     ldc = m
-
-    #absmax = absmax * 10
     #pdb.set_trace()
-    #print("A before kernel: ", A)
-    #print("B before kernel: ", B)
+    absmax = absmax.view(shapeB[0],int(shapeB[1]/blocksize)).transpose(0,1).contiguous()
     stream = _get_tensor_stream(A)
     if A.dtype == torch.float16:
         lib.cgemv_4bit_inference_fp16(
@@ -112,7 +107,7 @@ def _gemv_4bit_impl(
             k,
             get_ptr(A),
             get_ptr(B),
-            get_ptr(absmax.bfloat16()),
+            get_ptr(absmax),
             get_ptr(code),
             get_ptr(out),
             lda,
@@ -186,10 +181,8 @@ def _(
         blocksize: int,
     ) -> torch.Tensor:
         shape = (*A.shape[:-1], shapeB[0])
-        #import pdb
-        #pdb.set_trace()
-        out = torch.zeros(shape, device=A.device, dtype=torch.float32)
-        _gemv_4bit_impl(A, B, shapeB, absmax.bfloat16(), code, blocksize, out=out)
+        out = torch.empty(shape, device=A.device, dtype=A.dtype).float()
+        _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
         return out
 
     @register_kernel("bitsandbytes::gemv_4bit.out", "xpu")

csrc/pythonInterface.cpp

@@ -381,7 +381,7 @@ 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,   sycl::ext::oneapi::bfloat16 *absmax, float *datatype, float * out,
+    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);
@@ -826,7 +826,7 @@ 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,   sycl::ext::oneapi::bfloat16 *absmax, float *datatype,
+    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_bf16(m, n, k, A, B, absmax,  datatype, out, lda, ldb, ldc, blocksize, stream);

csrc/xpu_cutlass.h

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

csrc/xpu_cutlass_fusion.cpp

@@ -289,6 +289,9 @@ struct XE_2D_LD_Unpack {
 
     constexpr auto inst_size_bits = detail::size_of_inst_bits<CopyOp, dtype>;
 
+//if(cute::thread0()){
+//  print("copy base_addr: "); print(base_addr); print("\n");
+//}
     CopyOp::copy(base_addr + l * traits.stride_l,
                  (traits.width * sizeof_bits_v<dtype>) / sizeof_bits_v<int8_t>, traits.height,
                  (traits.pitch * sizeof_bits_v<dtype>) / sizeof_bits_v<int8_t>,
@@ -314,7 +317,9 @@ struct XE_2D_LD_Unpack {
     int y = is_need_reversed ? n : m;
 
     constexpr auto inst_size_bits = detail::size_of_inst_bits<CopyOp, dtype>;
-
+//if(cute::thread0()){
+//  print("prefetch base_addr: "); print(base_addr); print("\n");
+//}
     CopyOp::PREFETCH::copy(base_addr + l * atom.stride_l,
                            (atom.width * sizeof_bits_v<dtype>) / sizeof_bits_v<int8_t>, atom.height,
                            (atom.pitch * sizeof_bits_v<dtype>) / sizeof_bits_v<int8_t>,

include/cute/atom/copy_traits_xe.hpp

@@ -583,7 +583,13 @@ struct ThrMMA : TiledMMA
     auto thr_tensor = make_tensor(static_cast<BTensor&&>(btensor).data(), this->thrfrg_B(btensor.layout()));
 
     auto thr_vnk = make_coord(get<0>(thr_vmnk_), make_coord(get<2>(thr_vmnk_), get<3>(thr_vmnk_)));
-    //if(cute::thread0()) printf("partition_B: get<0>(thr_vmnk_) = %d, get<2>(thr_vmnk_) = %d, get<3>(thr_vmnk_) = %d\n", static_cast<int>(get<0>(thr_vmnk_)),static_cast<int>(get<2>(thr_vmnk_)),static_cast<int>(get<3>(thr_vmnk_)));
+#if 0    
+    if(int(ThreadIdxX()) == 16 && BlockIdxY()==0){
+      printf("partition_B: get<0>(thr_vmnk_) = %d, get<2>(thr_vmnk_) = %d, get<3>(thr_vmnk_) = %d\n", static_cast<int>(get<0>(thr_vmnk_)),static_cast<int>(get<2>(thr_vmnk_)),static_cast<int>(get<3>(thr_vmnk_)));
+      print("  thr_tensor   : "); print(thr_tensor);   print("\n");
+      print("  thr_tensor_return   : "); print(thr_tensor(thr_vnk, make_coord(_, repeat<rank<1,1>(thr_tensor)>(_))));   print("\n");
+    }
+#endif    
     return thr_tensor(thr_vnk, make_coord(_, repeat<rank<1,1>(thr_tensor)>(_)));
   }
 

include/cute/atom/mma_atom.hpp

@@ -30,6 +30,6 @@
 #gdb -args python -m pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
 #pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
 pytest -vs tests/test_xpu.py::TestXPU::test_gemv_4bit
-#python tests/test_xpu_db.py
+##python tests/test_xpu_db.py
 #gdb -args python tests/test_xpu_db.py
 #pytest tests/test_functional.py::TestQuantize4BitFunctional::test_gemv_4bit[dim=256-uint8-bf16-fc1-nf4-DQ_True-xpu]

run_case.sh

@@ -118,11 +118,11 @@ def test_gemm_4bit(self, device, dim, dtype, storage_type, quant_storage, double
           print("qB.t() = ",qB.t())
           C3 = torch.matmul(A, B.t())
           #pdb.set_trace()
-          C2 = F.gemv_4bit(A, qB.t(), state=state)
+          C2 = F.gemv_4bit(A, qB.t(), state=state).bfloat16()
           #pdb.set_trace()
           print("C3.sum() = ", C3.sum())
           print("C2.sum() = ", C2.sum())
-          diff = C2.bfloat16()-C3
+          diff = C2-C3
           print("diff/C2 = ", diff.sum()/C3.sum())
           print(C3)
           print(C2)
@@ -139,7 +139,7 @@ def test_gemm_4bit(self, device, dim, dtype, storage_type, quant_storage, double
           #print("B[0] = ",B[0])
           C3 = torch.matmul(A, B.t())
           #pdb.set_trace()
-          C2 = F.gemv_4bit(A, qB.t(), state=state)
+          C2 = F.gemv_4bit(A, qB.t(), state=state).bfloat16()
           pdb.set_trace()
           #print("C3.sum() = ", C3.sum())
           #print("C2.sum() = ", C2.sum())