float8 inference: fix bmm semantics

Summary:

Fixes the `Float8Tensor` `torch.bmm` override to match the semantics of the
high precision op. Specifically, input 1 is of shape (B, M, K) and input
2 is of shape (B, K, N).

Previously, the shape expectation from `torch.bmm`, which is confusing.

This is important for quantizing LLaMa 4 MoE variants, which use
`torch.bmm` in the HF implementation.

Test Plan:

```
pytest test/quantization/quantize_/workflows/float8/test_float8_tensor.py -s -x -k bmm
```

Reviewers:

Subscribers:

Tasks:

Tags:
ghstack-source-id: 9e16572
ghstack-comment-id: 3493356198
Pull-Request: #3296

Author: vkuzo

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

@@ -444,25 +444,27 @@ def test_bmm(self):
         # only support per row quantization
         config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
 
-        class M(torch.nn.Module):
+        class Model(torch.nn.Module):
             def __init__(self, weight):
                 super().__init__()
                 self.weight = weight
 
             def forward(self, x):
-                return torch.bmm(x, self.weight)
+                return torch.bmm(x, self.weight.transpose(-2, -1))
 
         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()
+
+        B, M, K, N = 10, 32, 128, 256
+
+        input = torch.randn(B, M, K, dtype=dtype, device=device)
+        weight = torch.randn(B, N, K, dtype=dtype, device=device)
+        m = Model(weight).eval()
         original = m(input)
-        # we need to transpose the weight first for bmm
-        m.weight = torch.nn.Parameter(m.weight.transpose(1, 2).contiguous())
         quantize_(m, config, filter_fn=lambda x, fqn: True)
         quantized = m(input)
-        self.assertTrue(compute_error(original, quantized) > 20)
+        sqnr = compute_error(original, quantized)
+        self.assertTrue(sqnr > 20)
 
     @common_utils.parametrize("granularity", [PerTensor(), PerRow()])
     @common_utils.parametrize(
@@ -551,6 +553,10 @@ def test_cat(self, granularity, sizes):
         self.assertEqual(cat_qweight2.qdata, ref_data)
         self.assertEqual(cat_qweight2.scale, ref_scale)
 
+    # TODO(future PR): add this back
+    @unittest.skip(
+        "This requires rowwise scaling for weight in layout BKN across axis 1 to work"
+    )
     @unittest.skipIf(not is_sm_at_least_90(), "Nedd sm90+")
     @unittest.skipIf(not _is_fbgemm_gpu_genai_available(), "Need fbgemm_gpu_genai")
     def test_moe_weight_reshape_ops(self):

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

@@ -422,24 +422,25 @@ def _(func, types, args, kwargs):
         a_scale = input_tensor.scale
 
         b_data = weight_tensor.qdata
-        b_scale = weight_tensor.scale.squeeze(-1)
-        assert b_data.is_contiguous(), "weight for bmm must be contiguous"
+        b_scale = weight_tensor.scale
 
         assert (
-            all(x == 1 for x in weight_tensor.block_size[:-1])
-            and weight_tensor.block_size[-1] == weight_tensor.shape[-1]
+            weight_tensor.block_size[0] == 1
+            and weight_tensor.block_size[1] == weight_tensor.shape[1]
+            and weight_tensor.block_size[2] == 1
         ), "bmm only works for per row weight quantization"
         assert (
             all(x == 1 for x in input_tensor.block_size[:-1])
             and input_tensor.block_size[-1] == input_tensor.shape[-1]
         ), "bmm only works for per row activation quantization"
 
-        orig_out_features = b_data.shape[-2]
+        orig_out_features = b_data.shape[-1]
 
         res = torch.ops.fbgemm.f8f8bf16_rowwise_batched(
             a_data,
-            b_data,
+            b_data.transpose(-2, -1),
             a_scale,
+            b_scale.transpose(-2, -1),
             b_scale,
         )
         res = res.reshape(*orig_act_size[:-1], orig_out_features)