From f30306c177204928b03fb2037dd18e3ec04074b4 Mon Sep 17 00:00:00 2001 From: SparshM8 <1.88700067e+08+SparshM8@users.noreply.github.com> Date: Sat, 22 Aug 2026 18:46:07 +0000 Subject: [PATCH] fix(cpu): correctly handle 1D output shape in dequantize_4bit --- bitsandbytes/backends/cpu/ops.py | 39 +++++++++++++++++--------------- tests/test_issue_2047.py | 38 +++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 18 deletions(-) create mode 100644 tests/test_issue_2047.py diff --git a/bitsandbytes/backends/cpu/ops.py b/bitsandbytes/backends/cpu/ops.py index 44fb5dceb..c6908473d 100755 --- a/bitsandbytes/backends/cpu/ops.py +++ b/bitsandbytes/backends/cpu/ops.py @@ -43,13 +43,14 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty(A.shape, device=A.device, dtype=torch.uint8) + out_ptr = out if A.dtype == torch.float32: lib.cquantize_blockwise_cpu_fp32( get_ptr(code), get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(n), ) @@ -58,7 +59,7 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor get_ptr(code), get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(n), ) @@ -67,7 +68,7 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor get_ptr(code), get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(n), ) @@ -97,12 +98,13 @@ def _( ) -> torch.Tensor: A = A.contiguous() out = torch.empty_like(A, dtype=dtype) + out_ptr = out if dtype == torch.float32: lib.cdequantize_blockwise_cpu_fp32( get_ptr(code), get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(A.numel()), ) @@ -111,7 +113,7 @@ def _( get_ptr(code), get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(A.numel()), ) @@ -120,7 +122,7 @@ def _( get_ptr(code), get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(A.numel()), ) @@ -137,7 +139,7 @@ def _( return out @register_kernel("bitsandbytes::dequantize_4bit", "cpu") - def _( + def dequantize_4bit_cpu( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, @@ -169,21 +171,22 @@ def _( if absmax.dtype != torch.float32: absmax = absmax.float() - if len(shape) == 1: - shape = (1, shape[0]) - m = prod(shape[:-1]) n = shape[-1] A = A.reshape(m, n // 2) + out = torch.empty(shape, dtype=dtype, device=A.device) + out_ptr = out + if len(shape) == 1: + out_ptr = out.unsqueeze(0) if quant_type == "fp4": if dtype == torch.float32: lib.cdequantize_blockwise_cpu_fp4_fp32( get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(m), ct.c_longlong(n), @@ -192,7 +195,7 @@ def _( lib.cdequantize_blockwise_cpu_fp4_bf16( get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(m), ct.c_longlong(n), @@ -201,7 +204,7 @@ def _( lib.cdequantize_blockwise_cpu_fp4_fp16( get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(m), ct.c_longlong(n), @@ -211,7 +214,7 @@ def _( lib.cdequantize_blockwise_cpu_nf4_fp32( get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(m), ct.c_longlong(n), @@ -220,7 +223,7 @@ def _( lib.cdequantize_blockwise_cpu_nf4_bf16( get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(m), ct.c_longlong(n), @@ -229,7 +232,7 @@ def _( lib.cdequantize_blockwise_cpu_nf4_fp16( get_ptr(A), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_longlong(blocksize), ct.c_longlong(m), ct.c_longlong(n), @@ -296,7 +299,7 @@ def _( get_ptr(A), get_ptr(B), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_int64(blocksize), ct.c_int64(x_strideM), ct.c_int64(out_strideM), @@ -309,7 +312,7 @@ def _( get_ptr(A), get_ptr(B), get_ptr(absmax), - get_ptr(out), + get_ptr(out_ptr), ct.c_int64(blocksize), ct.c_int64(x_strideM), ct.c_int64(out_strideM), diff --git a/tests/test_issue_2047.py b/tests/test_issue_2047.py new file mode 100644 index 000000000..bb10b08ab --- /dev/null +++ b/tests/test_issue_2047.py @@ -0,0 +1,38 @@ + +import pytest +import torch +import bitsandbytes as bnb +from bitsandbytes import functional as F +from tests.helpers import get_available_devices + +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) +@pytest.mark.parametrize("blocksize", [64, 128]) +def test_dequantize_4bit_1d_shape(device, dtype, quant_type, blocksize): + """ + Regression test for issue #2047: + CPU dequantize_4bit returns shape (1, n) for even-length 1-D inputs. + """ + if device == "hpu": + pytest.skip("Skipping on HPU") + + input_size = 256 + shape = (input_size,) + + # Create dummy quantized data + n = input_size + blocks = -(n // -blocksize) + + # 4-bit packed data + A = torch.randint(0, 255, (n // 2,), dtype=torch.uint8, device=device) + absmax = torch.randn((blocks,), dtype=torch.float32, device=device) + + # Call dequantize_4bit through the torch op (which calls our kernel) + out = torch.ops.bitsandbytes.dequantize_4bit.default( + A, absmax, blocksize, quant_type, shape, dtype + ) + + assert out.shape == shape, f"Expected shape {shape}, got {out.shape} on {device}" + assert out.dtype == dtype + assert out.device.type == torch.device(device).type