Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant

bitsandbytes/autograd/_functions.py

@@ -98,6 +98,11 @@ def reset_grads(self):
         self.CBt = None
 
 
+def _to_16bit(t: torch.Tensor) -> torch.Tensor:
+    """Pass fp16/bf16 tensors through unchanged; cast anything else to fp16 for the int8 quantization kernels."""
+    return t if t.dtype in (torch.float16, torch.bfloat16) else t.to(torch.float16)
+
+
 class MatMul8bitLt(torch.autograd.Function):
     @staticmethod
     def forward(
@@ -124,8 +129,8 @@ def forward(
 
         input_shape = A.shape
 
-        # Cast A to fp16
-        if A.dtype != torch.float16 and not _is_compiling():
+        # fp16 and bf16 are quantized natively; other dtypes (e.g. fp32) are cast to fp16 for quantization only.
+        if A.dtype not in (torch.float16, torch.bfloat16) and not _is_compiling():
             logger.warning("MatMul8bitLt: inputs will be cast from %s to float16 during quantization", A.dtype)
 
         if len(A.shape) == 3:
@@ -134,10 +139,10 @@ def forward(
         # 1. Quantize A. Note that as a side-effect, outliers are suppressed in CA/CAt.
         if ctx.needs_input_grad[1]:
             # Slower path
-            CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(A.to(torch.float16), threshold=state.threshold)
+            CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(_to_16bit(A), threshold=state.threshold)
         else:
             # Fast path
-            CA, SCA, outlier_cols = F.int8_vectorwise_quant(A.to(torch.float16), threshold=state.threshold)
+            CA, SCA, outlier_cols = F.int8_vectorwise_quant(_to_16bit(A), threshold=state.threshold)
             CAt = SCAt = None
 
         has_grad = False
@@ -152,7 +157,7 @@ def forward(
                 state.reset_grads()
 
                 # 2. Quantize B
-                state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16))
+                state.CB, state.SCB, _ = F.int8_vectorwise_quant(_to_16bit(B))
 
         # Handle sparse decomposition
         if state.threshold > 0.0:
@@ -214,19 +219,19 @@ def backward(ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor
             # compute grad_bias first before changing grad_output dtype
             grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias)
 
-        # Cast grad_output to fp16
         if len(grad_output.shape) == 3:
             grad_output = grad_output.reshape(-1, grad_output.shape[-1]).contiguous()
 
         if req_gradB:
-            Cgrad, _, _, SCgradt, _ = F.int8_double_quant(grad_output.to(torch.float16))
+            grad_output_16 = _to_16bit(grad_output)
+            Cgrad, _, _, SCgradt, _ = F.int8_double_quant(grad_output_16)
 
             grad_B = torch.ops.bitsandbytes.int8_scaled_mm.default(
                 Cgrad.t().contiguous(),
                 CAt.t(),
                 SCgradt,
                 SCAt,
-                dtype=torch.float16,
+                dtype=grad_output_16.dtype,
             )
 
             if state.threshold > 0.0 and subA is not None and subA.numel() > 0:
@@ -258,7 +263,7 @@ def forward(ctx, A, B, out=None, bias=None, state=MatmulLtState):
 
             if (state.is_training and not has_grad) or state.CB is None or state.SCB is None:
                 state.reset_grads()
-                state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16))
+                state.CB, state.SCB, _ = F.int8_vectorwise_quant(_to_16bit(B))
                 B = state.CB
 
         CB = state.CB.data.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))

bitsandbytes/backends/cuda/ops.py

@@ -48,13 +48,13 @@ def _setup_ctypes(names, argtypes, restype=None):
 
 # int8 mm dequant: (A, row_stats, col_stats, out, bias, numRows, numCols, stream)
 _setup_ctypes(
-    ["cdequant_mm_int32_fp16"],
+    ["cdequant_mm_int32_fp16", "cdequant_mm_int32_bf16"],
     [ct.c_void_p] * 5 + [ct.c_int32, ct.c_int32, ct.c_void_p],
 )
 
 # int8 vectorwise quant: (A, out, row_stats, threshold, rows, cols, stream)
 _setup_ctypes(
-    ["cint8_vector_quant"],
+    ["cint8_vector_quant", "cint8_vector_quant_bf16"],
     [ct.c_void_p] * 3 + [ct.c_float, ct.c_int32, ct.c_int32, ct.c_void_p],
 )
 
@@ -183,37 +183,39 @@ def _(
     if col_stats.dtype != torch.float32:
         raise ValueError(f"col_stats must be float32, got {col_stats.dtype}")
 
-    # Note: cuda kernel only currently supports fp16 output.
-    # We'll later cast to desired dtype if needed.
-    out = torch.empty_like(A, dtype=torch.float16)
+    # The kernel writes fp16 or bf16 directly. Any other requested dtype (e.g. fp32) is computed
+    # in fp16 and cast afterwards, matching the previous behavior.
+    out_dtype = dtype or torch.float16
+    kernel_dtype = torch.bfloat16 if out_dtype == torch.bfloat16 else torch.float16
+    out = torch.empty_like(A, dtype=kernel_dtype)
 
-    # Note: fused bias in the kernel is only supported for fp16
-    # TODO(matthewdouglas): Consider supporting bf16 fused bias
-    bias_ptr = bias.data_ptr() if bias is not None and bias.dtype == torch.float16 else None
+    # Fuse bias in the kernel when it already matches the kernel dtype; otherwise add it afterward.
+    fuse_bias = bias is not None and bias.dtype == kernel_dtype
+    fn = lib.cdequant_mm_int32_bf16 if kernel_dtype == torch.bfloat16 else lib.cdequant_mm_int32_fp16
 
     with _cuda_device_of(A):
-        lib.cdequant_mm_int32_fp16(
+        fn(
             A.data_ptr(),
             row_stats.data_ptr(),
             col_stats.data_ptr(),
             out.data_ptr(),
-            bias_ptr,
+            bias.data_ptr() if fuse_bias else None,
             A.numel() // A.shape[-1],
             A.shape[-1],
             _get_raw_stream(A.device.index),
         )
 
     # Add bias separately if not fused in kernel
-    if bias is not None and bias.dtype != torch.float16:
+    if bias is not None and not fuse_bias:
         out.add_(bias)
 
-    return out.to(dtype or torch.float16)
+    return out.to(out_dtype)
 
 
 @register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda")
 def _(A: torch.Tensor, threshold=0.0):
-    if A.dtype != torch.float16:
-        raise ValueError(f"A must be float16, got {A.dtype}")
+    if A.dtype not in (torch.float16, torch.bfloat16):
+        raise ValueError(f"A must be float16 or bfloat16, got {A.dtype}")
     if threshold < 0.0:
         raise ValueError("threshold must be non-negative")
 
@@ -235,8 +237,9 @@ def _(A: torch.Tensor, threshold=0.0):
             # Needed for torch.compile support.
             outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64)
 
+    fn = lib.cint8_vector_quant_bf16 if A.dtype == torch.bfloat16 else lib.cint8_vector_quant
     with _cuda_device_of(A):
-        lib.cint8_vector_quant(
+        fn(
             A.data_ptr(),
             out_row.data_ptr(),
             row_stats.data_ptr(),

bitsandbytes/functional.py

@@ -1565,6 +1565,7 @@ def int8_mm_dequant(
     col_stats: torch.Tensor,
     out: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
+    dtype: Optional[torch.dtype] = None,
 ):
     """Performs dequantization on the result of a quantized int8 matrix multiplication.
 
@@ -1574,11 +1575,14 @@ def int8_mm_dequant(
         col_stats (`torch.Tensor`): The column-wise quantization statistics for the rhs operand of the matrix multiplication.
         out (`torch.Tensor`, *optional*): A pre-allocated tensor to store the output of the operation.
         bias (`torch.Tensor`, *optional*): An optional bias vector to add to the result.
+        dtype (`torch.dtype`, *optional*): The desired output dtype. Defaults to `torch.float16`.
 
     Returns:
-        `torch.Tensor`: The dequantized result with an optional bias, with dtype `torch.float16`.
+        `torch.Tensor`: The dequantized result with an optional bias, with the requested dtype.
     """
-    result = torch.ops.bitsandbytes.int8_mm_dequant.default(A, row_stats, col_stats, dtype=torch.float16, bias=bias)
+    result = torch.ops.bitsandbytes.int8_mm_dequant.default(
+        A, row_stats, col_stats, dtype=dtype or torch.float16, bias=bias
+    )
 
     # TODO(matthewdouglas): Deprecate out kwarg
     if out is not None:
@@ -1607,7 +1611,7 @@ def int8_double_quant(
     </Tip>
 
     Args:
-        A (`torch.Tensor` with dtype `torch.float16`): The input matrix.
+        A (`torch.Tensor` with dtype `torch.float16` or `torch.bfloat16`): The input matrix.
         col_stats (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the column-wise quantization scales.
         row_stats (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the row-wise quantization scales.
         out_col (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the column-wise quantized data.
@@ -1653,12 +1657,12 @@ def int8_vectorwise_dequant(A: torch.Tensor, stats: torch.Tensor):
 
 
 def int8_vectorwise_quant(A: torch.Tensor, threshold=0.0):
-    """Quantizes a tensor with dtype `torch.float16` to `torch.int8` in accordance to the `LLM.int8()` algorithm.
+    """Quantizes a tensor with dtype `torch.float16` or `torch.bfloat16` to `torch.int8` in accordance to the `LLM.int8()` algorithm.
 
     For more information, see the [LLM.int8() paper](https://arxiv.org/abs/2208.07339).
 
     Args:
-        A (`torch.Tensor` with dtype `torch.float16`): The input tensor.
+        A (`torch.Tensor` with dtype `torch.float16` or `torch.bfloat16`): The input tensor.
         threshold (`float`, *optional*):
             An optional threshold for sparse decomposition of outlier features.
 

bitsandbytes/nn/modules.py

@@ -739,7 +739,9 @@ def _quantize(self, device):
             return super().to(device)
 
         # We quantize the weight and store in 8bit row-major
-        B = self.data.contiguous().to(device=device, dtype=torch.float16)
+        # bf16 weights are quantized natively; everything else goes through fp16.
+        dtype = torch.bfloat16 if self.data.dtype == torch.bfloat16 else torch.float16
+        B = self.data.contiguous().to(device=device, dtype=dtype)
         CB, SCB, _ = bnb.functional.int8_vectorwise_quant(B)
         self.data = CB
         self.CB = CB

csrc/kernels.cu

@@ -7,6 +7,8 @@
 #include "compat_device.cuh"
 #include "kernels.cuh"
 
+#include <type_traits>
+
 #define HLF_MAX 65504
 #define TH 1024
 #define NUM 4
@@ -1328,41 +1330,56 @@ __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit1StateBlockwise(
 // Outputs:
 //  rowStats [rows]
 //  out [rows, cols]
+// Streaming (evict-first) load of a 16-bit float, widened to fp32.
+// CUDA < 12.4 only declares the __nv_bfloat16 overload of __ldcs for sm_80+. The half overload is available on
+// every architecture, so for bf16 we issue the streaming load through it and reinterpret the 16 raw bits.
+template <typename T> __device__ __forceinline__ float ldcs_as_float(const T* ptr) {
+    if constexpr (std::is_same_v<T, half>) {
+        return __half2float(__ldcs(ptr));
+    } else {
+#if BNB_HIP
+        return (float)(*ptr);
+#else
+        const __half_raw bits = __ldcs(reinterpret_cast<const half*>(ptr));
+        __nv_bfloat16_raw raw;
+        raw.x = bits.x;
+        return __bfloat162float(__nv_bfloat16(raw));
+#endif
+    }
+}
+
 template <typename T, int THREADS, int SPARSE_DECOMP>
 __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__
     void kInt8VectorQuant(T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols) {
-
-    using BlockReduceT = bnb_cub::BlockReduce<T, THREADS>;
-
     // One block per row.
     // Threads load column values in a striped arrangement.
     // e.g. t0 reads row[0], row[0+nthreads], ..
     // and  t1 reads row[1], row[1+nthreads], ..
     // Each thread will determine its local absmax.
     // We then do a blockwise reduction to determine the row's absmax.
-
+    using BlockReduceT = bnb_cub::BlockReduce<float, THREADS>;
     __shared__ typename BlockReduceT::TempStorage temp_storage;
-    __shared__ T smem_row_absmax;
+    __shared__ float smem_row_absmax;
 
     const int row_id = blockIdx.x;
     const T* row_data = A + (row_id * cols);
 
     // Threads will read the row values in a striped access pattern and find a local absmax.
-    T row_local_absmax = -FLT_MIN;
+    float row_local_absmax = -FLT_MIN;
     for (int i = threadIdx.x; i < cols; i += THREADS) {
-        const T absval = fabsf(__ldcs(&(row_data[i])));
+        const float absval = fabsf(ldcs_as_float(&row_data[i]));
 
         // For sparse decomposition, values outside of the threshold are not to be
         // included when calculating the row's absmax.
         if constexpr (SPARSE_DECOMP) {
-            row_local_absmax = fmaxf(row_local_absmax, absval < T(threshold) ? absval : row_local_absmax);
+            row_local_absmax = fmaxf(row_local_absmax, absval < threshold ? absval : row_local_absmax);
         } else {
             row_local_absmax = fmaxf(row_local_absmax, absval);
         }
     }
 
     // Reduce thread-local absmax across the block.
-    const T row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, BNB_MAX_OP, cols);
+    const float row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, BNB_MAX_OP, cols);
     if (threadIdx.x == 0) {
         // Save our block's absmax to shared memory for the quantization step.
         rowStats[row_id] = smem_row_absmax = row_absmax;
@@ -1372,12 +1389,12 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__
     // Quantize row-wise.
     const float scale = __fdividef(127.0f, smem_row_absmax);
     for (int i = threadIdx.x; i < cols; i += THREADS) {
-        float val = row_data[i];
+        float val = (float)row_data[i];
 
         if constexpr (SPARSE_DECOMP) {
             // For sparse decomposition, we do not want to quantize the outliers.
             // Instead they're zeroed out.
-            out[row_id * cols + i] = fabs(val) < threshold ? __float2int_rn(val * scale) : 0;
+            out[row_id * cols + i] = fabsf(val) < threshold ? __float2int_rn(val * scale) : 0;
         } else {
             out[row_id * cols + i] = __float2int_rn(val * scale);
         }
@@ -1390,21 +1407,27 @@ template __global__ void kInt8VectorQuant<half, 1024, 0>(
 template __global__ void kInt8VectorQuant<half, 1024, 1>(
     half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols
 );
+template __global__ void kInt8VectorQuant<bnb_bfloat16, 1024, 0>(
+    bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols
+);
+template __global__ void kInt8VectorQuant<bnb_bfloat16, 1024, 1>(
+    bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols
+);
 
 #define MM_DEQUANT_CONST 6.200012e-05f // 1.0f/(127.0f*127.0f)
 
-template <int ITEMS_PER_THREAD, int THREADS>
-__global__ void kdequant_mm_int32_fp16(
-    int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out,
-    half* __restrict__ const bias, const int numRows, const int numCols, const int n
+template <typename T, int ITEMS_PER_THREAD, int THREADS>
+__global__ void kdequant_mm_int32(
+    int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, T* out,
+    T* __restrict__ const bias, const int numRows, const int numCols, const int n
 ) {
     const int n_out = numRows * numCols;
 
     int block_offset = blockIdx.x * THREADS * ITEMS_PER_THREAD;
     int thread_offset = threadIdx.x * ITEMS_PER_THREAD;
 
     int local_values[ITEMS_PER_THREAD];
-    half local_output[ITEMS_PER_THREAD];
+    T local_output[ITEMS_PER_THREAD];
 
     float local_rowStats[ITEMS_PER_THREAD];
     float local_colStats[ITEMS_PER_THREAD];
@@ -1423,7 +1446,7 @@ __global__ void kdequant_mm_int32_fp16(
 
         local_colStats[j] = col_idx >= numCols ? 0.0f : __ldg(&colStats[col_idx]);
         local_rowStats[j] = row_idx >= numRows ? 0.0f : __ldg(&rowStats[row_idx]);
-        local_biasValue[j] = ((bias == nullptr) || col_idx >= numCols) ? 0.0f : __half2float(bias[col_idx]);
+        local_biasValue[j] = ((bias == nullptr) || col_idx >= numCols) ? 0.0f : (float)(bias[col_idx]);
     }
 
     // Each block loads THREADS * ITEMS_PER_THREAD values from A
@@ -1433,9 +1456,8 @@ __global__ void kdequant_mm_int32_fp16(
 
 #pragma unroll ITEMS_PER_THREAD
     for (int j = 0; j < ITEMS_PER_THREAD; ++j) {
-        local_output[j] = __float2half(
-            fmaf(local_values[j] * local_rowStats[j] * local_colStats[j], MM_DEQUANT_CONST, local_biasValue[j])
-        );
+        local_output[j] =
+            (T)(fmaf(local_values[j] * local_rowStats[j] * local_colStats[j], MM_DEQUANT_CONST, local_biasValue[j]));
     }
 
 #pragma unroll ITEMS_PER_THREAD
@@ -1604,10 +1626,14 @@ template __global__ void kgemm_4bit_inference_naive<float, 128, 32>(
     float* out, int lda, int ldb, int ldc, int blocksize
 );
 
-template __global__ void kdequant_mm_int32_fp16<4, 512>(
+template __global__ void kdequant_mm_int32<half, 4, 512>(
     int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out,
     half* __restrict__ const bias, const int numRows, const int numCols, const int n
 );
+template __global__ void kdequant_mm_int32<bnb_bfloat16, 4, 512>(
+    int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats,
+    bnb_bfloat16* out, bnb_bfloat16* __restrict__ const bias, const int numRows, const int numCols, const int n
+);
 
 template __device__ unsigned char dQuantize<0>(float* smem_code, const float rand, float x);
 template __device__ unsigned char dQuantize<1>(float* smem_code, const float rand, float x);

csrc/kernels.cuh

@@ -65,10 +65,10 @@ __global__ void kOptimizerStatic8bit1StateBlockwise(
     const float gnorm_scale, const bool skip_zeros, const int n
 );
 
-template <int ITEMS_PER_THREAD, int THREADS>
-__global__ void kdequant_mm_int32_fp16(
-    int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out,
-    half* __restrict__ const bias, const int numRows, const int numCols, const int n
+template <typename T, int ITEMS_PER_THREAD, int THREADS>
+__global__ void kdequant_mm_int32(
+    int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, T* out,
+    T* __restrict__ const bias, const int numRows, const int numCols, const int n
 );
 
 template <typename T, int THREADS, int SPARSE_DECOMP>