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/4] 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/4] 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/4] 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/4] 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())