diff --git a/bitsandbytes/backends/cpu/ops.py b/bitsandbytes/backends/cpu/ops.py index 4efaa98ea..44fb5dceb 100755 --- a/bitsandbytes/backends/cpu/ops.py +++ b/bitsandbytes/backends/cpu/ops.py @@ -37,6 +37,7 @@ def _(A: torch.Tensor, B: torch.Tensor): @register_kernel("bitsandbytes::quantize_blockwise", "cpu") def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: + A = A.contiguous() n = A.numel() blocks = -(n // -blocksize) @@ -94,6 +95,7 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype ) -> torch.Tensor: + A = A.contiguous() out = torch.empty_like(A, dtype=dtype) if dtype == torch.float32: lib.cdequantize_blockwise_cpu_fp32( diff --git a/tests/test_ops.py b/tests/test_ops.py index 69589dcc0..36d4b4e12 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -382,9 +382,6 @@ class TestNonContiguousInputs: @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("blocksize", [64, 128, 256]) def test_quantize_blockwise_non_contiguous(self, device, dtype, blocksize): - if device == "cpu": - pytest.skip("Non-contiguous fix targets CUDA backend only") - code = bitsandbytes.functional.create_dynamic_map().to(device) # Create non-contiguous tensor via slicing @@ -404,9 +401,6 @@ def test_quantize_blockwise_non_contiguous(self, device, dtype, blocksize): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("blocksize", [64, 128, 256]) def test_dequantize_blockwise_non_contiguous(self, device, dtype, blocksize): - if device == "cpu": - pytest.skip("Non-contiguous fix targets CUDA backend only") - code = bitsandbytes.functional.create_dynamic_map().to(device, dtype=torch.float32) # Quantize a contiguous tensor, then create non-contiguous uint8 via transpose