add 3Dx3D test

Author: danielvegamyhre

torchao/prototype/grouped_mm/test_grouped_mm.py

@@ -8,8 +8,6 @@
 @pytest.mark.parametrize("use_fast_accum", [True, False])
 @pytest.mark.parametrize("strided", [True, False])
 def test_grouped_gemm_2d_3d(use_fast_accum, strided):
-    # unit test ensuring parity between torchao and pytorch core grouped_gemm
-    # https://github.com/pytorch/pytorch/blob/87bfd66c3c7061db6d36d8daa62f08f507f90e39/test/test_matmul_cuda.py#L1204
     device = "cuda"
     s_int = int(strided)
     m, n, k, n_groups = 16, 32, 16, 4
@@ -25,6 +23,24 @@ def test_grouped_gemm_2d_3d(use_fast_accum, strided):
     )
     assert isinstance(result, torch.Tensor)
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.parametrize("use_fast_accum", [True, False])
+@pytest.mark.parametrize("strided", [True, False])
+def test_grouped_gemm_3d_3d(use_fast_accum, strided):
+    device = "cuda"
+    s_int = int(strided)
+    m, n, k, n_groups = 16, 32, 16, 4
+    a = torch.randn(n_groups * (1 + s_int), m, k * (1 + s_int), device=device)[::(1 + s_int), :, :k]
+    b = torch.randn(n_groups * (1 + s_int), n, k * (1 + s_int), device=device)[::(1 + s_int), :, :k]
+    result = grouped_mm(
+        a, b,
+        float8_recipe=Float8LinearRecipeName.ROWWISE, 
+        out_dtype=torch.bfloat16, 
+        use_fast_accum=use_fast_accum
+    )
+    assert isinstance(result, torch.Tensor)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 def test_tensorwise_scaling_not_supported():
     device = "cuda"