Fix mixed INT8 FakeTensor output metadata

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
 

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)

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())