Skip to content

fix(cpu): correctly handle 1D output shape in dequantize_4bit #2055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SparshM8 wants to merge 1 commit into
base: main
Choose a base branch
from
+59 −18
Open

fix(cpu): correctly handle 1D output shape in dequantize_4bit #2055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 21 additions & 18 deletions bitsandbytes/backends/cpu/ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
)
Expand All @@ -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),
)
Expand All @@ -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),
)
Expand Down Expand Up @@ -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()),
)
Expand All @@ -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()),
)
Expand All @@ -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()),
)
Expand All @@ -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,
Expand Down Expand Up @@ -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),
Expand All @@ -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),
Expand All @@ -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),
Expand All @@ -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),
Expand All @@ -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),
Expand All @@ -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),
Expand Down Expand Up @@ -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),
Expand All @@ -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),
Expand Down
38 changes: 38 additions & 0 deletions tests/test_issue_2047.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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