diff --git a/bitsandbytes/backends/cpu/ops.py b/bitsandbytes/backends/cpu/ops.py index 44fb5dceb..7774477e3 100755 --- a/bitsandbytes/backends/cpu/ops.py +++ b/bitsandbytes/backends/cpu/ops.py @@ -169,9 +169,10 @@ def _( if absmax.dtype != torch.float32: absmax = absmax.float() - if len(shape) == 1: - shape = (1, shape[0]) - + # The kernel views the input as a 2-D (m, n) matrix, with m == 1 for 1-D + # inputs. The output buffer keeps the caller-provided shape: the + # CUDA/default/MPS backends and the registered fake kernel all return + # exactly `shape`. m = prod(shape[:-1]) n = shape[-1] diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..c5cf1e007 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -185,6 +185,31 @@ def test_quantize_4bit_not_divisible_by_blocksize(self, device, dtype, quant_typ # Verify output is finite (no NaN/Inf) assert torch.isfinite(out).all(), "Dequantized output contains NaN or Inf" + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) + @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) + @pytest.mark.parametrize("blocksize", [64, 128]) + def test_dequantize_4bit_1d_shape(self, device, dtype, quant_type, blocksize): + """1-D inputs must round-trip to 1-D outputs. + + The output shape contract is the `shape` argument itself: the registered + fake kernel returns torch.empty(shape), and the CUDA/default/MPS backends + allocate the output that way. Regression test for the CPU native path + returning (1, n) for even-length 1-D shapes. + """ + if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): + pytest.skip("This configuration is not supported on HPU.") + + shape = (blocksize * 2,) + A = torch.randn(shape, dtype=dtype, device=device) + + packed, absmax = torch.ops.bitsandbytes.quantize_4bit(A, blocksize, quant_type, torch.uint8) + out = torch.ops.bitsandbytes.dequantize_4bit(packed, absmax, blocksize, quant_type, shape, dtype) + + assert out.shape == shape + assert out.dtype == dtype + assert out.device == A.device + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype"))