Skip to content

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

Open

Description

@2sumtech

The native CPU path of bitsandbytes::dequantize_4bit reshapes a 1-D shape argument to (1, shape[0]) before allocating the output (bitsandbytes/backends/cpu/ops.py, if len(shape) == 1:), so a quantize/dequantize round trip of a 1-D tensor comes back 2-D on CPU.

import torch, bitsandbytes.functional as F
A = torch.randn(4096)
q, state = F.quantize_4bit(A, blocksize=64, quant_type="nf4")
F.dequantize_4bit(q, state).shape   # torch.Size([1, 4096]) on CPU

Every other implementation returns exactly shape: the CUDA backend allocates torch.empty(shape), the default and MPS backends reshape to shape, and the registered fake kernel returns torch.empty(shape) — so eager CPU also contradicts the op's fake-tensor contract under torch.compile. The smoking gun that this is a bug and not a design choice: the CPU backend is inconsistent with itself — an odd-length 1-D input (n=4095) takes the generic fallback and correctly returns (4095,), while n=4096 returns (1, 4096).

Present since the native CPU 4-bit kernel landed in #1789. Reproduced on macOS arm64 (NEON path), torch 2.13, main @ 95f9af3. Fix is a 2-line deletion — m = prod(shape[:-1]) already evaluates to 1 for 1-D inputs and the kernel writes a flat contiguous buffer; PR incoming.

Activity

  1. kvnloo commented on Sep 28, 2026

    @kvnloo

    The proposed removal of the special 1-D reshape looks consistent with every other backend's contract.

    I would review #2048 with shape parity as the primary invariant:

    dequantize_4bit(quantize_4bit(x)).shape == x.shape

    for:

    • even-length 1-D;
    • odd-length 1-D;
    • 2-D;
    • both NF4 and FP4.

    It is also worth exercising the fake/meta registration because the current bug means eager CPU and the fake op disagree about output rank.

    The CPU optimization can still internally view the data as (1, n) if needed, but that implementation detail should not escape into the public result shape.

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions