From 4ccf1321945e093f5c7416b5a8fad284c9f94866 Mon Sep 17 00:00:00 2001 From: Krushna Thakkar <109793947+kru2710shna@users.noreply.github.com> Date: Sun, 5 Jul 2026 18:15:04 +0530 Subject: [PATCH 1/5] Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant Templates int8VectorQuant on T and adds bf16 kernel instantiations plus a cint8_vector_quant_bf16 C ABI entry point, mirroring the existing gemm_4bit_inference_naive fp16/bf16/fp32 pattern. The blockwise absmax reduction now accumulates in float rather than T: required for bf16 to compile cleanly and slightly improves fp16 accuracy (rowStats was already float, so downstream is unaffected). Removes the forced A.to(torch.float16) casts in MatMul8bitLt so bf16 activations quantize natively. Closes #1868. --- bitsandbytes/autograd/_functions.py | 8 +++---- bitsandbytes/backends/cuda/ops.py | 10 ++++----- csrc/kernels.cu | 35 ++++++++++++----------------- csrc/ops.cu | 15 ++++++++++--- csrc/ops.cuh | 3 ++- csrc/pythonInterface.cpp | 9 +++++++- tests/test_ops.py | 9 +++----- 7 files changed, 48 insertions(+), 41 deletions(-) diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 8a069bd10..30b65d966 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -134,10 +134,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(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(A, threshold=state.threshold) CAt = SCAt = None has_grad = False @@ -152,7 +152,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(B) # Handle sparse decomposition if state.threshold > 0.0: @@ -258,7 +258,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(B) B = state.CB CB = state.CB.data.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0)) diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index 09d2e2244..e7c3f4cca 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -56,10 +56,9 @@ def _setup_ctypes(names, argtypes, restype=None): # 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], ) - # 4-bit/8-bit blockwise quantize: (code, A, absmax, out, blocksize, n) _setup_ctypes( [f"cquantize_blockwise_{d}_{q}" for d in ("fp32", "bf16", "fp16") for q in ("nf4", "fp4")] @@ -214,8 +213,8 @@ def _( @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") @@ -237,8 +236,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(), diff --git a/csrc/kernels.cu b/csrc/kernels.cu index 0d313c8d7..6fe5807c9 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -1324,39 +1324,29 @@ template __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; - - // 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; __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((float)__ldcs(&(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; } __syncthreads(); @@ -1364,12 +1354,10 @@ __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); } @@ -1382,7 +1370,12 @@ template __global__ void kInt8VectorQuant( template __global__ void kInt8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols ); - +template __global__ void kInt8VectorQuant( + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols +); +template __global__ void kInt8VectorQuant( + 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 diff --git a/csrc/ops.cu b/csrc/ops.cu index 16eed4e81..1bfffa5f3 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -421,13 +421,14 @@ void dequant_mm_int32_fp16( BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } +template void int8VectorQuant( - half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream + T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ) { if (threshold == 0.0) { - kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); + kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } else { - kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); + kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } @@ -481,6 +482,14 @@ template void gemm_4bit_inference_naive( int ldc, int blocksize, bnb_stream_t stream ); +template void int8VectorQuant( + half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream +); +template void int8VectorQuant( + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, + bnb_stream_t stream +); + template int igemmlt<32, 0>( bnb_blasLt_handle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, bnb_stream_t stream diff --git a/csrc/ops.cuh b/csrc/ops.cuh index c7114bcaa..55dbe183f 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -136,8 +136,9 @@ void cutlass_igemm( void dequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream ); +template void int8VectorQuant( - half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream + T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ); template diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 214c2a2d8..74b9676a1 100644 --- a/csrc/pythonInterface.cpp +++ b/csrc/pythonInterface.cpp @@ -557,7 +557,14 @@ void cdequant_mm_int32_fp16( void cint8_vector_quant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream ) { - int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); + int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); +} + +void cint8_vector_quant_bf16( + __nv_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, + cudaStream_t stream +) { + int8VectorQuant<__nv_bfloat16>(A, out, rowStats, threshold, rows, cols, stream); } void* cget_managed_ptr(size_t bytes) { diff --git a/tests/test_ops.py b/tests/test_ops.py index 69589dcc0..1043b4771 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -36,21 +36,19 @@ def test_int8_linear_matmul_out(self, device): opcheck(torch.ops.bitsandbytes.int8_linear_matmul.out, (A, B, out)) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("threshold", [0.0, 6.0]) @pytest.mark.parametrize("device", get_available_devices()) - def test_int8_vectorwise_quant(self, threshold, device): - A = torch.randn(10, 20, dtype=torch.float16, device=device) + def test_int8_vectorwise_quant(self, dtype, threshold, device): + A = torch.randn(10, 20, dtype=dtype, device=device) A[1][0] = 1000.0 - out_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant(A, threshold=threshold) - assert out_row.shape == (10, 20) assert out_row.dtype == torch.int8 assert out_row.device == A.device assert row_stats.shape == (10,) assert row_stats.dtype == torch.float32 assert row_stats.device == A.device - if threshold > 0.0: assert outlier_cols is not None assert outlier_cols.dim() == 1 @@ -58,7 +56,6 @@ def test_int8_vectorwise_quant(self, threshold, device): assert outlier_cols.device == A.device else: assert outlier_cols is None - opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,)) opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold)) From cdd9e23e0ca1d4e8c9fcbcbd5879cbd29e8cf6c3 Mon Sep 17 00:00:00 2001 From: Krushna Thakkar <109793947+kru2710shna@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:46:31 +0530 Subject: [PATCH 2/5] Address review: restore kernel comments, remove obsolete fp16 cast warning - Restore explanatory comments in kInt8VectorQuant that were dropped during the float-accumulation rewrite (striped-load pattern, sparse-decomp absmax, outlier zeroing). - Remove the 'inputs will be cast to float16' warning in MatMul8bitLt, which no longer applies now that the forced fp16 cast is gone. --- bitsandbytes/autograd/_functions.py | 4 ---- csrc/kernels.cu | 14 ++++++++++---- csrc/ops.cu | 3 +-- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 30b65d966..59749f5a6 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -124,10 +124,6 @@ def forward( input_shape = A.shape - # Cast A to fp16 - if A.dtype != torch.float16 and not _is_compiling(): - logger.warning("MatMul8bitLt: inputs will be cast from %s to float16 during quantization", A.dtype) - if len(A.shape) == 3: A = A.reshape(-1, A.shape[-1]) diff --git a/csrc/kernels.cu b/csrc/kernels.cu index 6fe5807c9..bf71f8466 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -1323,9 +1323,13 @@ __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit1StateBlockwise( template __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) { - + // 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; - __shared__ typename BlockReduceT::TempStorage temp_storage; __shared__ float smem_row_absmax; @@ -1336,7 +1340,8 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ float row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { const float absval = fabsf((float)__ldcs(&(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 < threshold ? absval : row_local_absmax); } else { @@ -1355,8 +1360,9 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ const float scale = __fdividef(127.0f, smem_row_absmax); for (int i = threadIdx.x; i < cols; i += THREADS) { 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] = fabsf(val) < threshold ? __float2int_rn(val * scale) : 0; } else { out[row_id * cols + i] = __float2int_rn(val * scale); diff --git a/csrc/ops.cu b/csrc/ops.cu index 1bfffa5f3..feb7962da 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -486,8 +486,7 @@ template void int8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ); template void int8VectorQuant( - bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, - bnb_stream_t stream + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ); template int igemmlt<32, 0>( From e888a86ee3cbcb7ca8060d5fd41a29e9e9b7d0d9 Mon Sep 17 00:00:00 2001 From: Krushna Thakkar <109793947+kru2710shna@users.noreply.github.com> Date: Wed, 8 Jul 2026 02:28:54 +0530 Subject: [PATCH 3/5] Add bf16 coverage to LLM.int8 tests Extends existing tests to exercise the new bf16 activation-quant path: - test_matmullt: drop the forced fp16 cast so the existing bf16 param actually exercises the bf16 kernel end-to-end. - test_int8_linear_matmul_half: parametrize over fp16/bf16. - test_linear8bitlt_inference and new test_linear8bitlt_forward_dtypes: exercise the Linear8bitLt module forward in both dtypes. Addresses review feedback on #1985. --- tests/test_autograd.py | 2 +- tests/test_functional.py | 12 +++++------- tests/test_linear8bitlt.py | 30 ++++++++++++++++++++++++++++++ tests/test_modules.py | 10 +++++----- 4 files changed, 41 insertions(+), 13 deletions(-) diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 7d273c853..aafb76a4c 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -76,7 +76,7 @@ def test_matmullt( if not transpose[0] and not transpose[1]: B2 = B2.t().contiguous() - state.CB, state.SCB, _ = bnb.functional.int8_vectorwise_quant(B2.to(torch.float16)) + state.CB, state.SCB, _ = bnb.functional.int8_vectorwise_quant(B2) B2 = state.CB if not transpose[0] and transpose[1]: diff --git a/tests/test_functional.py b/tests/test_functional.py index e4cd6a128..2b3704af9 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -397,22 +397,20 @@ def test_int8_linear_matmul(self, device, dim1, dim2, dim3, dim4, dims, ldb): @pytest.mark.parametrize("dim3", [32], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [32], ids=id_formatter("dim4")) @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims")) - def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims): + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) + def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims, dtype): for i in range(k): if dims == 2: - A = torch.normal(0, 0.5, size=(dim1, dim3), device=device).half() + A = torch.normal(0, 0.5, size=(dim1, dim3), device=device).to(dtype) elif dims == 3: - A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device=device).half() - B = torch.randn((dim4, dim3), device=device).half() + A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device=device).to(dtype) + B = torch.randn((dim4, dim3), device=device).to(dtype) torch.nn.init.xavier_uniform_(B) C1 = torch.matmul(A, B.t()) - A = A.view(-1, A.shape[-1]) - CA, statsA, _ = F.int8_vectorwise_quant(A) CB, statsB, _ = F.int8_vectorwise_quant(B) output = F.int8_mm_dequant(F.int8_linear_matmul(CA, CB), statsA, statsB) - torch.testing.assert_close(C1.view(-1, C1.shape[-1]), output, atol=0.025, rtol=0.05) @pytest.mark.parametrize("device", get_available_devices()) diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py index b078e82b7..25bc8bd38 100644 --- a/tests/test_linear8bitlt.py +++ b/tests/test_linear8bitlt.py @@ -13,6 +13,7 @@ from bitsandbytes.nn.modules import Linear8bitLt from tests.helpers import ( TRUE_FALSE, + describe_dtype, get_available_devices, id_formatter, torch_load_from_buffer, @@ -364,3 +365,32 @@ def test_linear8bitlt_device_movement(device): # Accelerator outputs should match both times. torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8) + + +@pytest.mark.parametrize("device", get_available_devices(no_cpu=True)) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) +@pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) +def test_linear8bitlt_forward_dtypes(device, dtype, threshold): + # Exercises the activation-quant path end-to-end through Linear8bitLt for both fp16 and bf16. + linear = torch.nn.Linear(32, 96) + linear_custom = Linear8bitLt( + linear.in_features, + linear.out_features, + linear.bias is not None, + has_fp16_weights=False, + threshold=threshold, + ) + linear_custom.weight = bnb.nn.Int8Params( + linear.weight.data.clone(), + requires_grad=False, + has_fp16_weights=False, + ) + linear_custom.bias = linear.bias + linear_custom = linear_custom.to(device) + + x = torch.randn(4, 32, dtype=dtype, device=device) + out = linear_custom(x) + + assert out.dtype == dtype + assert out.shape == (4, 96) + assert out.isfinite().all() diff --git a/tests/test_modules.py b/tests/test_modules.py index 95f78b6d3..42f38eea7 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -7,7 +7,7 @@ from torch import nn import bitsandbytes as bnb -from tests.helpers import get_available_devices, id_formatter, is_supported_on_hpu +from tests.helpers import describe_dtype, get_available_devices, id_formatter, is_supported_on_hpu @contextlib.contextmanager @@ -61,15 +61,15 @@ def assert_all_approx_close(a, b, atol=1e-8, rtol=1e-5, count=10): @pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("threshold", [0.0, 3.0], ids=id_formatter("threshold")) -def test_linear8bitlt_inference(device, threshold): - l1 = bnb.nn.Linear8bitLt(32, 64, threshold=threshold, has_fp16_weights=False).to(device).half() +def test_linear8bitlt_inference(device, dtype, threshold): + l1 = bnb.nn.Linear8bitLt(32, 64, threshold=threshold, has_fp16_weights=False).to(device).to(dtype) assert l1.weight.device.type == device assert l1.weight.dtype == torch.int8 - l1.eval() for i in range(100): - b1 = torch.randn(16, 8, 32, device=device).half() + b1 = torch.randn(16, 8, 32, device=device, dtype=dtype) o1 = l1(b1) if i == 1: assert l1.state.CB is not None From 786d3f36616e24874510d2031d72a712af3458eb Mon Sep 17 00:00:00 2001 From: Krushna Thakkar <109793947+kru2710shna@users.noreply.github.com> Date: Sat, 11 Jul 2026 03:23:28 +0530 Subject: [PATCH 4/5] Add native bf16 output support to int8_mm_dequant Templates kdequant_mm_int32_fp16 and its launcher on output dtype, adds a cdequant_mm_int32_bf16 C ABI entry point, and dispatches on dtype in the Python backend. Fused bias now works for both fp16 and bf16 when the bias dtype matches the output dtype. Exposes a dtype kwarg on bitsandbytes.functional.int8_mm_dequant so callers can request bf16 output directly instead of relying on a cast afterward. Reverts the earlier test workaround now that int8_mm_dequant genuinely returns bf16 when requested. Addresses review feedback on #1985 (previously deferred as a follow-up). --- bitsandbytes/backends/cuda/ops.py | 21 +++++++++++++-------- bitsandbytes/functional.py | 9 +++++++-- csrc/kernels.cu | 21 ++++++++++++--------- csrc/kernels.cuh | 6 +++--- csrc/ops.cu | 13 +++++++++++-- csrc/ops.cuh | 5 ++++- csrc/pythonInterface.cpp | 9 ++++++++- tests/test_functional.py | 2 +- 8 files changed, 59 insertions(+), 27 deletions(-) diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index e7c3f4cca..eb4c463c1 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -48,9 +48,10 @@ def _setup_ctypes(names, argtypes, restype=None): restype=ct.c_int32, ) +# int8 mm dequant: (A, row_stats, col_stats, out, bias, numRows, numCols, stream) # 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], ) @@ -186,14 +187,18 @@ def _( # 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) + out_dtype = dtype or torch.float16 + if out_dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"dtype must be float16 or bfloat16, got {out_dtype}") + + out = torch.empty_like(A, dtype=out_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 output dtype; otherwise add it afterward. + bias_ptr = bias.data_ptr() if bias is not None and bias.dtype == out_dtype else None + fn = lib.cdequant_mm_int32_bf16 if out_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(), @@ -205,10 +210,10 @@ def _( ) # Add bias separately if not fused in kernel - if bias is not None and bias.dtype != torch.float16: + if bias is not None and bias.dtype != out_dtype: out.add_(bias) - return out.to(dtype or torch.float16) + return out @register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda") diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index bb56e9d50..a238b6208 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -1561,6 +1561,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. @@ -1570,11 +1571,15 @@ 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. Supports `torch.float16` and + `torch.bfloat16`. 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: diff --git a/csrc/kernels.cu b/csrc/kernels.cu index bf71f8466..83490578d 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -1384,10 +1384,10 @@ template __global__ void kInt8VectorQuant( ); #define MM_DEQUANT_CONST 6.200012e-05f // 1.0f/(127.0f*127.0f) -template +template __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 + 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; @@ -1395,7 +1395,7 @@ __global__ void kdequant_mm_int32_fp16( 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]; @@ -1414,7 +1414,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 @@ -1424,9 +1424,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 @@ -1595,10 +1594,14 @@ template __global__ void kgemm_4bit_inference_naive( 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_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 __global__ void kdequant_mm_int32_fp16( + 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); diff --git a/csrc/kernels.cuh b/csrc/kernels.cuh index dc511661b..a43ddd9b1 100644 --- a/csrc/kernels.cuh +++ b/csrc/kernels.cuh @@ -65,10 +65,10 @@ __global__ void kOptimizerStatic8bit1StateBlockwise( const float gnorm_scale, const bool skip_zeros, const int n ); -template +template __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 + 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 diff --git a/csrc/ops.cu b/csrc/ops.cu index feb7962da..32b394d4b 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -407,8 +407,9 @@ int fill_up_to_nearest_multiple(int value, int multiple) { return value + (value % multiple == 0 ? 0 : (multiple - (value % multiple))); } +template void dequant_mm_int32_fp16( - int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream + int* A, float* rowStats, float* colStats, T* out, T* bias, int numRows, int numCols, bnb_stream_t stream ) { const int threads = 512; const int num_per_thread = 4; @@ -416,7 +417,7 @@ void dequant_mm_int32_fp16( const int n = numRows * numCols; const int num_blocks = (n + num_per_block - 1) / num_per_block; - kdequant_mm_int32_fp16 + kdequant_mm_int32_fp16 <<>>(A, rowStats, colStats, out, bias, numRows, numCols, n); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } @@ -489,6 +490,14 @@ template void int8VectorQuant( bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ); +template void dequant_mm_int32_fp16( + int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream +); +template void dequant_mm_int32_fp16( + int* A, float* rowStats, float* colStats, bnb_bfloat16* out, bnb_bfloat16* bias, int numRows, int numCols, + bnb_stream_t stream +); + template int igemmlt<32, 0>( bnb_blasLt_handle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, bnb_stream_t stream diff --git a/csrc/ops.cuh b/csrc/ops.cuh index 55dbe183f..b3452fe51 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -133,9 +133,12 @@ int igemmlt( void cutlass_igemm( bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ); + +template void dequant_mm_int32_fp16( - int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream + int* A, float* rowStats, float* colStats, T* out, T* bias, int numRows, int numCols, bnb_stream_t stream ); + template void int8VectorQuant( T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 74b9676a1..32c1770b8 100644 --- a/csrc/pythonInterface.cpp +++ b/csrc/pythonInterface.cpp @@ -551,7 +551,14 @@ int cigemmlt_8_rowscale( void cdequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, cudaStream_t stream ) { - dequant_mm_int32_fp16(A, rowStats, colStats, out, bias, numRows, numCols, stream); + dequant_mm_int32_fp16(A, rowStats, colStats, out, bias, numRows, numCols, stream); +} + +void cdequant_mm_int32_bf16( + int* A, float* rowStats, float* colStats, __nv_bfloat16* out, __nv_bfloat16* bias, int numRows, int numCols, + cudaStream_t stream +) { + dequant_mm_int32_fp16<__nv_bfloat16>(A, rowStats, colStats, out, bias, numRows, numCols, stream); } void cint8_vector_quant( diff --git a/tests/test_functional.py b/tests/test_functional.py index 2b3704af9..4214f050b 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -410,7 +410,7 @@ def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims, dty A = A.view(-1, A.shape[-1]) CA, statsA, _ = F.int8_vectorwise_quant(A) CB, statsB, _ = F.int8_vectorwise_quant(B) - output = F.int8_mm_dequant(F.int8_linear_matmul(CA, CB), statsA, statsB) + output = F.int8_mm_dequant(F.int8_linear_matmul(CA, CB), statsA, statsB, dtype=dtype) torch.testing.assert_close(C1.view(-1, C1.shape[-1]), output, atol=0.025, rtol=0.05) @pytest.mark.parametrize("device", get_available_devices()) From 42b08e1c48013d94ce67a9501228354d69731cc0 Mon Sep 17 00:00:00 2001 From: kru2710shna Date: Mon, 2 Nov 2026 00:01:40 -0500 Subject: [PATCH 5/5] Fix CUDA<12.4/ROCm builds, restore fp32 support, rename dequant_mm_int32 - kInt8VectorQuant: CUDA < 12.4 only declares the bf16 __ldcs overload for sm_80+. Issue the streaming load through the half overload and reinterpret the bits. - pythonInterface.cpp: use bnb_bfloat16 instead of __nv_bfloat16 so ROCm builds. - Rename kdequant_mm_int32_fp16/dequant_mm_int32_fp16 to kdequant_mm_int32/dequant_mm_int32 now that they are templated on the output type (review feedback). - int8_mm_dequant (CUDA): accept any output dtype again. fp16/bf16 are written natively; other dtypes (fp32) are computed in fp16 and cast, as before. - MatMul8bitLt/MatMul8bitFp: only cast to fp16 when the input is not fp16/bf16, so fp32 inputs keep working and keep their fp32 output dtype. Backward keeps bf16 grads in bf16. - Int8Params: quantize bf16 weights natively instead of via fp16. - Add test_int8_mm_dequant_dtypes covering fp16/bf16/fp32 output and fused/unfused bias. --- bitsandbytes/autograd/_functions.py | 23 ++++++++++++++------- bitsandbytes/backends/cuda/ops.py | 24 ++++++++++------------ bitsandbytes/functional.py | 9 ++++---- bitsandbytes/nn/modules.py | 4 +++- csrc/kernels.cu | 32 +++++++++++++++++++++++++---- csrc/kernels.cuh | 2 +- csrc/ops.cu | 8 ++++---- csrc/ops.cuh | 2 +- csrc/pythonInterface.cpp | 11 +++++----- tests/test_ops.py | 23 +++++++++++++++++++++ 10 files changed, 96 insertions(+), 42 deletions(-) diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 59749f5a6..e28192b08 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/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,16 +129,20 @@ def forward( input_shape = A.shape + # 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: A = A.reshape(-1, A.shape[-1]) # 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, 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, threshold=state.threshold) + CA, SCA, outlier_cols = F.int8_vectorwise_quant(_to_16bit(A), threshold=state.threshold) CAt = SCAt = None has_grad = False @@ -148,7 +157,7 @@ def forward( state.reset_grads() # 2. Quantize B - state.CB, state.SCB, _ = F.int8_vectorwise_quant(B) + state.CB, state.SCB, _ = F.int8_vectorwise_quant(_to_16bit(B)) # Handle sparse decomposition if state.threshold > 0.0: @@ -210,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: @@ -254,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) + 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)) diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index cb8f30125..603a84db5 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -46,7 +46,6 @@ def _setup_ctypes(names, argtypes, restype=None): restype=ct.c_int32, ) -# int8 mm dequant: (A, row_stats, col_stats, out, bias, numRows, numCols, stream) # int8 mm dequant: (A, row_stats, col_stats, out, bias, numRows, numCols, stream) _setup_ctypes( ["cdequant_mm_int32_fp16", "cdequant_mm_int32_bf16"], @@ -58,6 +57,7 @@ def _setup_ctypes(names, argtypes, restype=None): ["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], ) + # 4-bit/8-bit blockwise quantize: (code, A, absmax, out, blocksize, n) _setup_ctypes( [f"cquantize_blockwise_{d}_{q}" for d in ("fp32", "bf16", "fp16") for q in ("nf4", "fp4")] @@ -183,17 +183,15 @@ 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. + # 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 - if out_dtype not in (torch.float16, torch.bfloat16): - raise ValueError(f"dtype must be float16 or bfloat16, got {out_dtype}") + kernel_dtype = torch.bfloat16 if out_dtype == torch.bfloat16 else torch.float16 + out = torch.empty_like(A, dtype=kernel_dtype) - out = torch.empty_like(A, dtype=out_dtype) - - # Fuse bias in the kernel when it already matches the output dtype; otherwise add it afterward. - bias_ptr = bias.data_ptr() if bias is not None and bias.dtype == out_dtype else None - fn = lib.cdequant_mm_int32_bf16 if out_dtype == torch.bfloat16 else lib.cdequant_mm_int32_fp16 + # 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): fn( @@ -201,17 +199,17 @@ def _( 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 != out_dtype: + if bias is not None and not fuse_bias: out.add_(bias) - return out + return out.to(out_dtype) @register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda") diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index ba40d8918..9f7c43979 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -1575,8 +1575,7 @@ 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. Supports `torch.float16` and - `torch.bfloat16`. Defaults to `torch.float16`. + dtype (`torch.dtype`, *optional*): The desired output dtype. Defaults to `torch.float16`. Returns: `torch.Tensor`: The dequantized result with an optional bias, with the requested dtype. @@ -1612,7 +1611,7 @@ def int8_double_quant( 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. @@ -1658,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. diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..783ddf170 100644 --- a/bitsandbytes/nn/modules.py +++ b/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 diff --git a/csrc/kernels.cu b/csrc/kernels.cu index 93f413ccc..9455f388a 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -7,6 +7,8 @@ #include "compat_device.cuh" #include "kernels.cuh" +#include + #define HLF_MAX 65504 #define TH 1024 #define NUM 4 @@ -1328,6 +1330,24 @@ __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 __device__ __forceinline__ float ldcs_as_float(const T* ptr) { + if constexpr (std::is_same_v) { + return __half2float(__ldcs(ptr)); + } else { +#if BNB_HIP + return (float)(*ptr); +#else + const __half_raw bits = __ldcs(reinterpret_cast(ptr)); + __nv_bfloat16_raw raw; + raw.x = bits.x; + return __bfloat162float(__nv_bfloat16(raw)); +#endif + } +} + template __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) { @@ -1347,7 +1367,8 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ // Threads will read the row values in a striped access pattern and find a local absmax. float row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { - const float absval = fabsf((float)__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) { @@ -1360,6 +1381,7 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ // Reduce thread-local absmax across the block. 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; } __syncthreads(); @@ -1368,6 +1390,7 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ const float scale = __fdividef(127.0f, smem_row_absmax); for (int i = threadIdx.x; i < cols; i += THREADS) { 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. @@ -1390,10 +1413,11 @@ template __global__ void kInt8VectorQuant( template __global__ void kInt8VectorQuant( 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 -__global__ void kdequant_mm_int32_fp16( +__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 ) { @@ -1602,11 +1626,11 @@ template __global__ void kgemm_4bit_inference_naive( float* out, int lda, int ldb, int ldc, int blocksize ); -template __global__ void kdequant_mm_int32_fp16( +template __global__ void kdequant_mm_int32( 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_fp16( +template __global__ void kdequant_mm_int32( 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 ); diff --git a/csrc/kernels.cuh b/csrc/kernels.cuh index a43ddd9b1..6ef2fcd67 100644 --- a/csrc/kernels.cuh +++ b/csrc/kernels.cuh @@ -66,7 +66,7 @@ __global__ void kOptimizerStatic8bit1StateBlockwise( ); template -__global__ void kdequant_mm_int32_fp16( +__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 ); diff --git a/csrc/ops.cu b/csrc/ops.cu index 483588eaf..74e3d4936 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -496,7 +496,7 @@ int fill_up_to_nearest_multiple(int value, int multiple) { } template -void dequant_mm_int32_fp16( +void dequant_mm_int32( int* A, float* rowStats, float* colStats, T* out, T* bias, int numRows, int numCols, bnb_stream_t stream ) { const int threads = 512; @@ -505,7 +505,7 @@ void dequant_mm_int32_fp16( const int n = numRows * numCols; const int num_blocks = (n + num_per_block - 1) / num_per_block; - kdequant_mm_int32_fp16 + kdequant_mm_int32 <<>>(A, rowStats, colStats, out, bias, numRows, numCols, n); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } @@ -578,10 +578,10 @@ template void int8VectorQuant( bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ); -template void dequant_mm_int32_fp16( +template void dequant_mm_int32( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream ); -template void dequant_mm_int32_fp16( +template void dequant_mm_int32( int* A, float* rowStats, float* colStats, bnb_bfloat16* out, bnb_bfloat16* bias, int numRows, int numCols, bnb_stream_t stream ); diff --git a/csrc/ops.cuh b/csrc/ops.cuh index b3452fe51..d0e0dc62a 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -135,7 +135,7 @@ void cutlass_igemm( ); template -void dequant_mm_int32_fp16( +void dequant_mm_int32( int* A, float* rowStats, float* colStats, T* out, T* bias, int numRows, int numCols, bnb_stream_t stream ); diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 8c531832e..9b8a17779 100644 --- a/csrc/pythonInterface.cpp +++ b/csrc/pythonInterface.cpp @@ -545,14 +545,14 @@ int cigemmlt_8_rowscale( void cdequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, cudaStream_t stream ) { - dequant_mm_int32_fp16(A, rowStats, colStats, out, bias, numRows, numCols, stream); + dequant_mm_int32(A, rowStats, colStats, out, bias, numRows, numCols, stream); } void cdequant_mm_int32_bf16( - int* A, float* rowStats, float* colStats, __nv_bfloat16* out, __nv_bfloat16* bias, int numRows, int numCols, + int* A, float* rowStats, float* colStats, bnb_bfloat16* out, bnb_bfloat16* bias, int numRows, int numCols, cudaStream_t stream ) { - dequant_mm_int32_fp16<__nv_bfloat16>(A, rowStats, colStats, out, bias, numRows, numCols, stream); + dequant_mm_int32(A, rowStats, colStats, out, bias, numRows, numCols, stream); } void cint8_vector_quant( @@ -562,10 +562,9 @@ void cint8_vector_quant( } void cint8_vector_quant_bf16( - __nv_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, - cudaStream_t stream + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream ) { - int8VectorQuant<__nv_bfloat16>(A, out, rowStats, threshold, rows, cols, stream); + int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); } void* cget_managed_ptr(size_t bytes) { diff --git a/tests/test_ops.py b/tests/test_ops.py index 828e8e56f..aef624f40 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -72,6 +72,29 @@ def test_int8_mm_dequant(self, device): opcheck(torch.ops.bitsandbytes.int8_mm_dequant, (A, row_stats, col_stats)) + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) + @pytest.mark.parametrize( + "bias_dtype", [None, torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("bias_dtype") + ) + def test_int8_mm_dequant_dtypes(self, device, dtype, bias_dtype): + # Covers the native fp16/bf16 kernel outputs, the fused-bias path (bias dtype == output dtype), + # the unfused path, and fp32 output (computed in fp16, then cast). + A = torch.randint(-127 * 127, 127 * 127, (64, 96), dtype=torch.int32, device=device) + row_stats = torch.rand(64, dtype=torch.float32, device=device) * 4 + col_stats = torch.rand(96, dtype=torch.float32, device=device) * 4 + bias = torch.randn(96, dtype=bias_dtype, device=device) if bias_dtype is not None else None + + out = torch.ops.bitsandbytes.int8_mm_dequant(A, row_stats, col_stats, dtype=dtype, bias=bias) + + ref = A.float() * row_stats.unsqueeze(1) * col_stats.unsqueeze(0) / (127.0 * 127.0) + if bias is not None: + ref = ref + bias.float() + + assert out.dtype == dtype + assert out.shape == A.shape + torch.testing.assert_close(out.float(), ref, rtol=1e-2, atol=1e-2) + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("has_bias", TRUE_FALSE)