Fix XPU quantize_4bit returning padded tensors that violate the op contract

[jiqing-feng]

Problem

On XPU, torch.ops.bitsandbytes.quantize_4bit returns tensors whose shapes do not match the ones declared by its register_fake implementation (and produced by the CPU/CUDA backends) whenever the input element count is not a multiple of blocksize.

import torch
import bitsandbytes  # noqa: F401

A = torch.randn(20, 20, dtype=torch.float16, device="xpu")
packed, absmax = torch.ops.bitsandbytes.quantize_4bit(A, 64, "nf4", torch.uint8)
print("real:", tuple(packed.shape), tuple(absmax.shape))

with torch._subclasses.FakeTensorMode():
    A = torch.randn(20, 20, dtype=torch.float16, device="xpu")
    packed, absmax = torch.ops.bitsandbytes.quantize_4bit(A, 64, "nf4", torch.uint8)
    print("fake:", tuple(packed.shape), tuple(absmax.shape))

Before:

real: (224, 1) (8,)
fake: (200, 1) (7,)

After:

real: (200, 1) (7,)
fake: (200, 1) (7,)

Root cause

bitsandbytes/backends/triton/ops.py::quantize_4bit pads the input up to the next multiple of blocksize and then derives the output sizes from the padded element count, so the padding leaks into the returned shapes:

remainder = n % blocksize
if remainder != 0:
    padding = blocksize - remainder
    A = torch.nn.functional.pad(A.view(-1), (0, padding), value=0.0)
    n = A.numel()

The padding is unnecessary: the Triton kernel already receives num_elements=n and masks out-of-range lanes itself.

Fix

Drop the padding and size the outputs from the real element count. absmax is still allocated to an even block count because the kernel walks blocks in pairs, and is trimmed to the real block count before returning. The packed buffer is rounded up to a whole number of quant_storage elements so the final view(quant_storage) stays valid, then sliced to the size the fake kernel promises.

This also fixes a latent crash the padding used to hide: with an odd element count and a non-uint8 quant_storage, the reinterpreting view raised RuntimeError: self.size(-1) must be divisible by 2 to view Byte as Half.

triton_ops.quantize_4bit is only registered for "xpu"; CPU, CUDA, MPS and HPU have their own implementations and are unaffected.

Verification

Hardware: Intel Arc Pro B60, torch 2.14.0+xpu.

Shape contract against register_fake, across shapes, blocksizes, quant types, dtypes and quant_storage variants: 47 failures before, all passing after.

No numerical change. Quantized bytes were compared bit-for-bit between stock and patched builds over 144 combinations (8 shapes x 3 blocksizes x 2 quant types x 3 dtypes) with a fixed seed: the 54 cases divisible by blocksize are bit-identical, and the 90 non-divisible cases are bit-identical over the non-padding prefix.

Performance is unchanged for divisible shapes and improves for non-divisible ones, since the pad allocation and copy are gone (200 iterations x 5 repeats, fp4):

shape	blocksize	before	after
(1024, 1000)	64	0.1134 ms	0.1134 ms
(1024, 1000)	256	0.1087 ms	0.1088 ms
(20, 20)	64	0.1257 ms	0.1066 ms

Test suite:

BNB_TEST_DEVICE=xpu pytest tests/test_ops.py tests/test_functional.py tests/test_linear4bit.py \
    tests/test_modules.py tests/test_autograd.py tests/test_parametrize.py tests/test_generation.py

3136 passed, 42 skipped, 10 deselected, 30 xfailed

[jiqing-feng]

Hi @matthewdouglas . Would you please review this PR? Thanks!