fix(cpu): correctly handle 1D output shape in dequantize_4bit - #2055
fix(cpu): correctly handle 1D output shape in dequantize_4bit#2055SparshM8 wants to merge 1 commit into
Conversation
caiotheodoro
left a comment
There was a problem hiding this comment.
#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).
Description
Fixes #2047.
The CPU kernel for
dequantize_4bitwas 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
bitsandbytes/backends/cpu/ops.pyto use anout_ptr(an unsqueezed view) for C-library calls while maintaining the originalouttensor's shape.out_ptrpattern toquantize_blockwise,dequantize_blockwise, andgemv_4bitkernels for consistency and safety.tests/test_issue_2047.pythat verifies 1D output shapes for all supported dtypes and quantization types on available devices.Testing
Verified with the new regression test on CPU:
Output:
12 passed in 18.96s