Skip to content

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

Open
SparshM8 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
SparshM8:fix/cpu-dequantize-1d-shape
Open

fix(cpu): correctly handle 1D output shape in dequantize_4bit#2055
SparshM8 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
SparshM8:fix/cpu-dequantize-1d-shape

Conversation

@SparshM8

Copy link
Copy Markdown

Description

Fixes #2047.

The CPU kernel for dequantize_4bit was unconditionally unsqueezing 1D shapes to (1, n) to satisfy C-library pointer requirements, but it failed to preserve the original 1D shape in the returned tensor. This resulted in an unexpected (1, N) output for a requested (N,) shape, which could cause dimension mismatch errors in downstream code.

Changes

  • Modified bitsandbytes/backends/cpu/ops.py to use an out_ptr (an unsqueezed view) for C-library calls while maintaining the original out tensor's shape.
  • Applied the same out_ptr pattern to quantize_blockwise, dequantize_blockwise, and gemv_4bit kernels for consistency and safety.
  • Added a new regression test tests/test_issue_2047.py that verifies 1D output shapes for all supported dtypes and quantization types on available devices.

Testing

Verified with the new regression test on CPU:

PYTHONPATH=. pytest tests/test_issue_2047.py

Output: 12 passed in 18.96s

@caiotheodoro caiotheodoro left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#2048 fixes the same root cause in 4 lines: shape gets reassigned to (1, shape[0]) and that mutated value is what torch.empty(shape, ...) uses for out, but m/n don't need the reassignment (prod(()) and prod((1,)) are both 1). Deleting the reassignment fixes it without touching anything else.

This PR's out_ptr change also touches quantize_blockwise and dequantize_blockwise (lines 43 and 97), which don't have the shape bug at all — out_ptr = out there is a no-op alias, pure diff noise on unrelated functions. Worth closing this in favor of #2048, or at minimum scoping it down to just dequantize_4bit and dropping the rename of the registered kernel from _ to dequantize_4bit_cpu (every other kernel in this file keeps the anonymous _ convention).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CPU dequantize_4bit returns shape (1, n) for even-length 1-D inputs; all other backends return (n,)

2 participants