diff --git a/bitsandbytes/_ops.py b/bitsandbytes/_ops.py index 43efd8609..4a72ad729 100644 --- a/bitsandbytes/_ops.py +++ b/bitsandbytes/_ops.py @@ -27,8 +27,8 @@ def _( out = torch.empty(shapeC, device=A.device, dtype=A.dtype) - outlier_cols = torch.library.get_ctx().new_dynamic_size() - subA = A.new_empty(outlier_cols, dtype=torch.int64) + num_outlier_cols = 0 if outlier_cols is None else outlier_cols.shape[0] + subA = A.new_empty((A.shape[0], num_outlier_cols)) return out, subA diff --git a/bitsandbytes/backends/default/ops.py b/bitsandbytes/backends/default/ops.py index 521802922..844a04695 100644 --- a/bitsandbytes/backends/default/ops.py +++ b/bitsandbytes/backends/default/ops.py @@ -73,22 +73,25 @@ def _( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: subB = None - if outlier_cols is not None and outlier_cols.numel(): - # Extract the inputs with outliers in original precision + if outlier_cols is not None: + # Extract the inputs with outliers in original precision. Keep this + # two-dimensional even when there are no outliers so the output + # metadata is data-independent for FakeTensor and torch.compile. subA = A[:, outlier_cols].contiguous() - # Dequantize the corresponding weight columns - subB = ( - torch.ops.bitsandbytes.int8_vectorwise_dequant.default(CB[:, outlier_cols].contiguous(), SCB) - .to(A.dtype) - .t() - ) + if outlier_cols.numel(): + # Dequantize the corresponding weight columns + subB = ( + torch.ops.bitsandbytes.int8_vectorwise_dequant.default(CB[:, outlier_cols].contiguous(), SCB) + .to(A.dtype) + .t() + ) - # TODO: if state.has_fp16_weights: subB = B[:, outlier_cols].t() + # TODO: if state.has_fp16_weights: subB = B[:, outlier_cols].t() else: # Needed for torch.compile when there are no outliers. - subA = torch.empty(0, device=A.device, dtype=A.dtype) + subA = A.new_empty((A.shape[0], 0)) # Int8 Matmul + Dequant + Bias output = torch.ops.bitsandbytes.int8_scaled_mm.default(CA, CB, SCA, SCB, bias=bias, dtype=A.dtype) diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..4f3bc063f 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -92,6 +92,40 @@ def test_int8_scaled_mm(self, device, dtype, has_bias): opcheck(torch.ops.bitsandbytes.int8_scaled_mm, (A, B, row_stats, col_stats, bias, dtype)) + @pytest.mark.parametrize("has_outliers", TRUE_FALSE) + @pytest.mark.parametrize("device", get_available_devices()) + def test_int8_mixed_scaled_mm(self, device, has_outliers): + A = torch.full((10, 20), 0.5, dtype=torch.float16, device=device) + threshold = 6.0 + if has_outliers: + A[1, 0] = 10.0 + + CA, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant(A, threshold) + B = torch.randn(30, 20, dtype=torch.float16, device=device) + CB, col_stats, _ = torch.ops.bitsandbytes.int8_vectorwise_quant(B) + + out, subA = torch.ops.bitsandbytes.int8_mixed_scaled_mm(A, CA, CB, row_stats, col_stats, outlier_cols) + + assert out.shape == (10, 30) + assert out.dtype == A.dtype + assert out.device == A.device + assert subA.shape == (10, int(has_outliers)) + assert subA.dtype == A.dtype + assert subA.device == A.device + + opcheck( + torch.ops.bitsandbytes.int8_mixed_scaled_mm.default, + (A, CA, CB, row_stats, col_stats, outlier_cols), + ) + + if not has_outliers: + _, subA_without_outlier_cols = torch.ops.bitsandbytes.int8_mixed_scaled_mm(A, CA, CB, row_stats, col_stats) + assert subA_without_outlier_cols.shape == (10, 0) + opcheck( + torch.ops.bitsandbytes.int8_mixed_scaled_mm.default, + (A, CA, CB, row_stats, col_stats), + ) + class TestInt8BlockwiseQuantOps: @pytest.mark.parametrize("device", get_available_devices())