Add support for bmm for fbgemm config

test/dtypes/test_fbgemm_fp8.py

@@ -34,6 +34,12 @@ def setUp(self):
             weight_dtype=e4m3_dtype,
             output_dtype=torch.bfloat16,
         )
+        self.bmm_config = FbgemmConfig(
+            input_dtype=e4m3_dtype,
+            weight_dtype=e4m3_dtype,
+            output_dtype=torch.bfloat16,
+            transpose_input=True,
+        )
 
     def test_linear(self):
         dtype = torch.bfloat16
@@ -106,6 +112,25 @@ def test_slice_and_copy_(self):
         # making sure param.data is updated
         assert param.data.float8_data[0][0] != orig_value
 
+    def test_bmm(self):
+        class M(torch.nn.Module):
+            def __init__(self, weight):
+                super().__init__()
+                self.weight = weight
+
+            def forward(self, x):
+                return torch.bmm(x, self.weight)
+
+        dtype = torch.bfloat16
+        device = "cuda"
+        input = torch.randn(10, 32, 128, dtype=dtype, device=device)
+        weight = torch.randn(10, 128, 256, dtype=dtype, device=device)
+        m = M(weight).eval()
+        original = m(input)
+        quantize_(m, self.bmm_config, filter_fn=lambda x, fqn: True)
+        quantized = m(input)
+        self.assertTrue(compute_error(original, quantized) > 20)
+
 
 if __name__ == "__main__":
     run_tests()

test/dtypes/test_fbgemm_int4.py

@@ -34,6 +34,13 @@ def setUp(self):
             output_dtype=torch.bfloat16,
             block_size=[1, 128],
         )
+        self.bmm_config = FbgemmConfig(
+            input_dtype=torch.bfloat16,
+            weight_dtype=torch.int4,
+            output_dtype=torch.bfloat16,
+            block_size=[1, 1, 128],
+            transpose_input=True,
+        )
 
     def test_linear(self):
         dtype = torch.bfloat16
@@ -111,6 +118,25 @@ def test_slice_and_copy_(self):
         # making sure param.data is updated
         assert param.data.packed_weight[0][0] != orig_value
 
+    def test_bmm(self):
+        class M(torch.nn.Module):
+            def __init__(self, weight):
+                super().__init__()
+                self.weight = weight
+
+            def forward(self, x):
+                return torch.bmm(x, self.weight)
+
+        dtype = torch.bfloat16
+        device = "cuda"
+        input = torch.randn(10, 32, 128, dtype=dtype, device=device)
+        weight = torch.randn(10, 128, 256, dtype=dtype, device=device)
+        m = M(weight).eval()
+        original = m(input)
+        quantize_(m, self.bmm_config, filter_fn=lambda x, fqn: True)
+        quantized = m(input)
+        self.assertTrue(compute_error(original, quantized) > 18)
+
 
 if __name__ == "__main__":
     run_tests()

torchao/dtypes/fbgemm_fp8_tensor.py

@@ -79,6 +79,7 @@ def from_float(
         cls,
         w: torch.Tensor,
         activation_scale_ub: Optional[float] = None,
+        transpose_input: bool = False,
     ):
         if activation_scale_ub is None:
             activation_scale_ub = 1200.0
@@ -88,6 +89,12 @@ def from_float(
             dtype=torch.float,
             device=w.device,
         )
+        if transpose_input:
+            if w.ndim == 3:
+                w = w.transpose(1, 2)
+            else:
+                w = w.t()
+
         wq, w_scale = torch.ops.triton.quantize_fp8_row(w)
         # wq, w_scale = torch.ops.fbgemm.quantize_fp8_per_row(w)
         dtype = w.dtype
@@ -141,6 +148,39 @@ def _(func, types, args, kwargs):
     return res
 
 
+@implements(torch.bmm)
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor = (
+        args[0],
+        args[1],
+    )
+    if not input_tensor.is_floating_point():
+        raise NotImplementedError(
+            f"{func} is not implemented for non floating point input"
+        )
+
+    orig_act_size = input_tensor.size()
+
+    # not used
+    num_tokens = torch.empty([input_tensor.size(0)], device=input_tensor.device)
+    xq, x_scale = torch.ops.fbgemm.quantize_fp8_per_row(
+        input_tensor, num_tokens, weight_tensor.activation_scale_ub
+    )
+
+    a_data = xq
+    b_data = weight_tensor.float8_data
+    orig_out_features = b_data.shape[-2]
+
+    res = torch.ops.fbgemm.f8f8bf16_rowwise_batched(
+        a_data,
+        b_data,
+        x_scale,
+        weight_tensor.scale,
+    )
+    res = res.reshape(*orig_act_size[:-1], orig_out_features)
+    return res
+
+
 @implements([aten.detach.default, aten.alias.default])
 def _(func, types, args, kwargs):
     return return_and_correct_aliasing(

torchao/dtypes/fbgemm_int4_tensor.py

@@ -82,13 +82,20 @@ def from_float(
         cls,
         w: torch.Tensor,
         block_size: List[int],
+        transpose_input: bool = False,
     ):
         assert len(block_size) == w.ndim, (
             f"Expecting the length of block_size to be equal to the dimension of the weight, got {block_size=} and {w.ndim=}"
         )
         if int4_row_quantize_zp is None:
             raise ImportError("Requires fbgemm-gpu-genai >= 1.2.0")
 
+        if transpose_input:
+            if w.ndim == 3:
+                w = w.transpose(1, 2)
+            else:
+                w = w.t()
+
         group_size = block_size[-1]
         original_shape = w.shape
 
@@ -146,6 +153,30 @@ def _(func, types, args, kwargs):
     return res
 
 
+@implements(torch.bmm)
+def _(func, types, args, kwargs):
+    input_tensor, weight_tensor = (
+        args[0],
+        args[1],
+    )
+    if not input_tensor.is_floating_point():
+        raise NotImplementedError(
+            f"{func} is not implemented for non floating point input"
+        )
+
+    orig_act_size = input_tensor.size()
+    orig_out_features = weight_tensor.shape[-2]
+
+    res = torch.ops.fbgemm.bf16i4bf16_rowwise_batched(
+        input_tensor,
+        weight_tensor.packed_weight.contiguous(),
+        weight_tensor.scale,
+        weight_tensor.zero_point,
+    )
+    res = res.reshape(*orig_act_size[:-1], orig_out_features)
+    return res
+
+
 @implements([aten.detach.default, aten.alias.default])
 def _(func, types, args, kwargs):
     return return_and_correct_aliasing(

torchao/quantization/quant_api.py

@@ -1991,6 +1991,7 @@ class FbgemmConfig(AOBaseConfig):
     output_dtype: torch.dtype
     block_size: Optional[List[int]] = None
     activation_scale_ub: Optional[float] = None
+    transpose_input: bool = False
 
 
 @register_quantize_module_handler(FbgemmConfig)
@@ -2018,6 +2019,7 @@ def _(module: torch.nn.Module, config: FbgemmConfig) -> torch.nn.Module:
         weight = to_fbgemm_int4(
             module.weight,
             config.block_size,
+            config.transpose_input,
         )
         module.weight = torch.nn.Parameter(weight, requires_grad=False)
         module.extra_repr = types.MethodType(_linear_extra_repr, module)
@@ -2029,6 +2031,7 @@ def _(module: torch.nn.Module, config: FbgemmConfig) -> torch.nn.Module:
         weight = to_fbgemm_fp8(
             module.weight,
             config.activation_scale_ub,
+            config.transpose_input,
         )
         module.weight = torch.nn.Parameter(weight, requires_grad=False)
         module.extra_repr = types.MethodType(_linear_extra_repr, module)