Fix fp8 matmul reference for sm90

Currently we need `use_fast_accum` on `torch._scaled_mm` to match
the behavior of `ct.mma` and `ct.scaled_mma` with fp8 on sm90.

Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

samples/BatchMatMul.py

@@ -103,7 +103,8 @@ def torch_batch_matmul_fp8(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
         A_row = A[i].contiguous()
         B_col = B[i].transpose(-2, -1).contiguous().transpose(-2, -1)
         C[i] = torch._scaled_mm(
-            A_row, B_col, scale_a=inv_sa, scale_b=inv_sb, out_dtype=torch.float32
+            A_row, B_col, scale_a=inv_sa, scale_b=inv_sb, out_dtype=torch.float32,
+            use_fast_accum=True
         )
     return C
 

samples/templates/BatchMatMul.py

@@ -64,7 +64,8 @@ def torch_batch_matmul_fp8(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
         A_row = A[i].contiguous()
         B_col = B[i].transpose(-2, -1).contiguous().transpose(-2, -1)
         C[i] = torch._scaled_mm(
-            A_row, B_col, scale_a=inv_sa, scale_b=inv_sb, out_dtype=torch.float32
+            A_row, B_col, scale_a=inv_sa, scale_b=inv_sb, out_dtype=torch.float32,
+            use_fast_accum=True
         )
     return C
 

test/bench_matmul.py

@@ -181,7 +181,8 @@ def torch_batch_matmul(bs, A, B, C):
             A_row = A[i].contiguous()
             B_col = B[i].transpose(-2, -1).contiguous().transpose(-2, -1)
             C[i] = torch._scaled_mm(
-                A_row, B_col, scale_a=inv_sa, scale_b=inv_sb, out_dtype=torch.float32
+                A_row, B_col, scale_a=inv_sa, scale_b=inv_sb, out_dtype=torch.float32,
+                use_fast_accum=True
             )
 
 

test/test_mma.py

@@ -125,7 +125,7 @@ def test_mma_fp8(tile_size, case):
     C = torch.ones((m, n), dtype=case.acc_dtype, device="cuda")
     scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
     try:
-        ref = torch._scaled_mm(A, B.T, scale, scale, out_dtype=C.dtype) + C
+        ref = torch._scaled_mm(A, B.T, scale, scale, out_dtype=C.dtype, use_fast_accum=True) + C
     except (RuntimeError, ValueError) as e:
         assert 'Multiplication of two Float8_e5m2 matrices is not supported' in str(e)
         ref = None
@@ -279,7 +279,8 @@ def test_matmul_fp8(tile_size, dtype):
     C = torch.zeros((m, n), dtype=dtype, device="cuda")
     scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
     try:
-        ref = torch._scaled_mm(A, B.T, scale, scale, out_dtype=torch.float16).to(dtype)
+        ref = torch._scaled_mm(A, B.T, scale, scale,
+                               out_dtype=torch.float16, use_fast_accum=True).to(dtype)
     except (RuntimeError, ValueError) as e:
         assert 'Multiplication of two Float8_e5m2 matrices is not supported' in str(e)
         ref = None