From 7c2a8feef4d9ef4024cd52b5c2d57668ed079a7c Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Thu, 15 Oct 2026 03:42:36 +0000 Subject: [PATCH] Fix XPU quantize_4bit returning padded tensors that violate the op contract The Triton quantize_4bit path pads the input up to a multiple of blocksize and then sizes its outputs from the padded element count. The padded sizes leak into the returned tensors, so on XPU both `packed` and `absmax` come back larger than the shapes declared by register_fake in bitsandbytes/_ops.py whenever numel() is not a multiple of blocksize. For a 20x20 fp16 tensor with blocksize=64 the fake kernel, the CPU backend and the CUDA backend all return packed (200, 1) / absmax (7,), but XPU returns packed (224, 1) / absmax (8,). This breaks meta tracing and torch.compile, makes 4-bit checkpoints non-portable between XPU and CPU/CUDA, and causes downstream shape validation (for example in diffusers) to reject XPU-quantized weights. The padding is unnecessary: the kernels in backends/triton/kernels_4bit.py already mask both their loads and their stores. Drop it, keep absmax over-allocated to an even block count since the kernel processes blocks in pairs, and trim it before returning. Size the packed byte buffer from the real element count, rounding up only as far as needed to keep the quant_storage view legal, then slice to the declared length - this also fixes odd element counts combined with a non-uint8 quant_storage, which the padding used to mask. Quantized bytes are bit-identical to before for inputs divisible by blocksize, and identical on the valid prefix otherwise. --- bitsandbytes/backends/triton/ops.py | 30 ++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/bitsandbytes/backends/triton/ops.py b/bitsandbytes/backends/triton/ops.py index 2dfe7e758..2940371d4 100644 --- a/bitsandbytes/backends/triton/ops.py +++ b/bitsandbytes/backends/triton/ops.py @@ -73,31 +73,35 @@ def quantize_4bit( if A.dtype not in (torch.bfloat16, torch.float16, torch.float32): raise ValueError(f"Blockwise 4bit quantization only supports 16/32-bit floats, but got {A.dtype}") + # The kernel addresses the input as a flat buffer, so it must be contiguous. + A = A.contiguous() + n = A.numel() - # Pad to next multiple of blocksize so the kernel always processes full blocks - 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 kernel processes blocks in pairs, so absmax is over-allocated and trimmed below. + blocks = -(n // -blocksize) + paired_blocks = -(n // -(blocksize * 2)) + + # Packed element count promised by the fake kernel. + packed_elems = (n + 1) // (quant_storage.itemsize * 2) - blocks = -(n // -(blocksize * 2)) + # One byte per pair of elements, rounded up so the view() to quant_storage is valid. + nbytes = n - n // 2 + nbytes += -nbytes % quant_storage.itemsize - absmax = torch.empty((blocks * 2,), device=A.device, dtype=A.dtype) - # Use n - n//2 instead of (n+1)//2 to avoid integer overflow for large n - out = torch.empty((n - n // 2, 1), device=A.device, dtype=torch.uint8) + absmax = torch.empty((paired_blocks * 2,), device=A.device, dtype=A.dtype) + out = torch.empty((nbytes, 1), device=A.device, dtype=torch.uint8) with torch_accelerator_module.device(A.device): kernels_4bit.quantize_4bit_blockwise_triton( - A, blocksize, quant_type, blocks, absmax, num_elements=n, quantized_out=out + A, blocksize, quant_type, paired_blocks, absmax, num_elements=n, quantized_out=out ) packed = out if quant_storage != torch.uint8: - packed = out.squeeze().view(quant_storage).unsqueeze(1) + packed = out.view(-1).view(quant_storage)[:packed_elems].unsqueeze(1) - return packed, absmax.float() + return packed, absmax[:blocks].float() def dequantize_4bit(