fix scale issue

Author: xiaolil1

bitsandbytes/backends/xpu/ops.py

@@ -74,20 +74,16 @@ 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()
+    #pdb.set_trace()
     stream = _get_tensor_stream(A)
     if A.dtype == torch.float16:
         lib.cgemv_4bit_inference_fp16(
@@ -112,7 +108,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,11 +182,9 @@ 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)
-        return 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.bfloat16()
 
     @register_kernel("bitsandbytes::gemv_4bit.out", "xpu")
     def _(

bitsandbytes/functional.py

@@ -937,6 +937,12 @@ def quantize_4bit(
         quant_storage,
     )
 
+    #import pdb
+    #pdb.set_trace()
+    #print("_absmax = ", _absmax)
+    _absmax = _absmax.view(input_shape[0],int(input_shape[1]/blocksize)).transpose(0,1).contiguous()
+    #pdb.set_trace()
+
     code = get_4bit_type(quant_type, device=A.device)
 
     if compress_statistics:
@@ -969,7 +975,6 @@ def quantize_4bit(
     # TODO(matthewdouglas): Deprecate absmax kwarg
     if absmax is not None:
         state.absmax = absmax.copy_(state.absmax)
-
     return out, state
 
 

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())
@@ -294,6 +294,7 @@ def test_gemv_4bit(self, device, dim, dtype, storage_type, quant_storage, double
                 A = torch.randn(1, dim, dtype=dtype, device=device)
                 B = torch.randn(dim * 3, dim, dtype=dtype, device=device) / math.sqrt(dim)
 
+            #pdb.set_trace()
             qB, state = F.quantize_4bit(
                 B,
                 quant_type=storage_type,
@@ -303,10 +304,10 @@ def test_gemv_4bit(self, device, dim, dtype, storage_type, quant_storage, double
             #pdb.set_trace()
             C3 = torch.matmul(A, B.t())
             #pdb.set_trace()
-            C2 = F.gemv_4bit(A, qB.t(), state=state).bfloat16()
+            C2 = F.gemv_4bit(A, qB.t(), state=state)
             #print("C2[0] = ", C2[0])
             A.requires_grad = True
-            C1 = bnb.matmul_4bit(A, qB.t(), state)#.bfloat16()
+            C1 = F.gemv_4bit(A, qB.t(), state=state) #bnb.matmul_4bit(A, qB.t(), state)
             #pdb.set_trace()
 
             err1 = (C1 - C2).abs().float()