From 60db6380bc4b854ec98b7bfe5f92ea2b69d1b8ef Mon Sep 17 00:00:00 2001 From: 2sumtech <2sumtech@gmail.com> Date: Wed, 26 Aug 2026 17:46:21 -0700 Subject: [PATCH] Fix default int8 kernels for inputs with more than 2 dimensions The default backend implementations of int8_mm_dequant, int8_vectorwise_quant and int8_vectorwise_dequant assumed a 2-D input. int8_mm_dequant flattened the leading dimensions and returned the flattened result, and the two vectorwise kernels broadcast against the wrong dimension and raised a RuntimeError. The registered fake kernels and the CUDA implementations all treat the leading dimensions as rows and preserve the shape of A, so this was a shape contract violation that torch.library.opcheck catches on any non-CUDA device. Flatten the leading dimensions for the computation and restore the original shape on the way out. Output for 2-D inputs is bit for bit unchanged. --- bitsandbytes/_ops.py | 4 +++- bitsandbytes/backends/default/ops.py | 23 +++++++++++++---------- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/bitsandbytes/_ops.py b/bitsandbytes/_ops.py index 43efd8609..c1fe15244 100644 --- a/bitsandbytes/_ops.py +++ b/bitsandbytes/_ops.py @@ -118,7 +118,9 @@ def _(A: torch.Tensor, stats: torch.Tensor) -> torch.Tensor: @register_kernel("bitsandbytes::int8_vectorwise_dequant", "default") def _(A: torch.Tensor, stats: torch.Tensor): # To dequantize we divide by 127, or multiply by the reciprocal. - return A * stats.view(-1, 1) * 7.874015718698502e-3 + # A may have more than 2 dimensions, so flatten the leading dimensions into rows first. + out = A.reshape(-1, A.shape[-1]) * stats.view(-1, 1) * 7.874015718698502e-3 + return out.reshape(A.shape) torch.library.define( diff --git a/bitsandbytes/backends/default/ops.py b/bitsandbytes/backends/default/ops.py index 521802922..fa302a82a 100644 --- a/bitsandbytes/backends/default/ops.py +++ b/bitsandbytes/backends/default/ops.py @@ -50,7 +50,7 @@ def _( if col_stats.dtype != torch.float32: raise ValueError(f"col_stats must be float32, got {col_stats.dtype}") - A_calc = A.view(-1, A.shape[-1]) + A_calc = A.reshape(-1, A.shape[-1]) row_stats = row_stats.reshape(-1).unsqueeze(-1) col_stats = col_stats.reshape(-1).unsqueeze(0) @@ -58,7 +58,8 @@ def _( if bias is not None: out += bias - return out.to(dtype or torch.float16) + # The op contract preserves the shape of A, which may have more than 2 dimensions. + return out.to(dtype or torch.float16).reshape(A.shape) @register_kernel("bitsandbytes::int8_mixed_scaled_mm", "default") @@ -141,30 +142,32 @@ def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: Optional[tor @register_kernel("bitsandbytes::int8_vectorwise_quant", "default") def _(A: torch.Tensor, threshold=0.0): - rows = A.numel() // A.shape[-1] + # Rows are the leading dimensions flattened together; A may have more than 2 dimensions. + A_calc = A.reshape(-1, A.shape[-1]) + rows = A_calc.shape[0] outlier_cols = None outlier_restore = None if threshold > 0.0: - outliers = A.abs() >= threshold + outliers = A_calc.abs() >= threshold if outliers.any(): # Determine which columns contain outliers, and zero out the # outliers ahead of quantization. We need to keep a backup of these # outliers to restore them after quantization. outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1) - outlier_restore = A[outliers].clone() - A[outliers] = 0 + outlier_restore = A_calc[outliers].clone() + A_calc[outliers] = 0 else: # Needed for torch.compile support. outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64) # Get absmax for each row. - row_stats = torch.max(A.abs(), dim=1).values.float() + row_stats = torch.max(A_calc.abs(), dim=-1).values.float() # Quantize row-wise to int8. - out_row = torch.round(A * (127.0 / row_stats.unsqueeze(-1))).to(torch.int8) + out_row = torch.round(A_calc * (127.0 / row_stats.unsqueeze(-1))).to(torch.int8) # Zero out values from outlier columns across all rows. if rows > 1 and outlier_cols is not None: @@ -172,9 +175,9 @@ def _(A: torch.Tensor, threshold=0.0): # Restore outliers. if outlier_restore is not None: - A[outliers] = outlier_restore + A_calc[outliers] = outlier_restore - return out_row, row_stats, outlier_cols + return out_row.reshape(A.shape), row_stats, outlier_cols @register_kernel("bitsandbytes::quantize_blockwise", "default")