Add amax as default per-tensor scaling factor for fp8_gemm benchmark (#339)

Summary:
Pull Request resolved: #339

Add `amax` (absolute maximum) as the default scaling factor for per-tensor scaling with fp8 workloads, as is used in practice. Also add a command-line argument that allows the user to define a scaling factor for per-tensor scaling.

Reviewed By: NikhilAPatel

Differential Revision: D80577628

Author: jananisriram

tritonbench/operators/fp8_gemm/fp8_gemm.py

@@ -41,6 +41,8 @@ def parse_args(args):
     parser.add_argument("--m", type=int)
     parser.add_argument("--k", type=int)
     parser.add_argument("--n", type=int)
+    parser.add_argument("--per-tensor-scale-a", type=float, default=None)
+    parser.add_argument("--per-tensor-scale-b", type=float, default=None)
     return parser.parse_args(args)
 
 
@@ -54,7 +56,20 @@ def __init__(
         super().__init__(tb_args, extra_args)
         self.extra_args = parse_args(extra_args)
 
+    def _get_dtype(self):
+        if self.extra_args.scaling_rowwise:
+            return torch.bfloat16
+        else:
+            return torch.float16
+
     def get_input_iter(self):
+        def _get_scale_per_tensor(x: torch.Tensor, custom_scale: float = None) -> torch.Tensor:
+            # For tensor-wise scaling, kernel requires a float32 scale tensor
+            if custom_scale:
+                return torch.tensor(custom_scale, dtype=torch.float32, device=x.device)
+            scale = torch.finfo(torch.float8_e4m3fn).max / x.abs().max()
+            return scale.to(torch.float32)
+
         def args(m, n, k):
             a = torch.randn(m, k, device=self.device).to(torch.float16)
             b = torch.randn(k, n, device=self.device).to(torch.float16).T.contiguous().T
@@ -64,8 +79,8 @@ def args(m, n, k):
                 scale_a = torch.ones((M, 1), dtype=torch.float32, device=a.device)
                 scale_b = torch.ones((1, N), dtype=torch.float32, device=b.device)
             else:
-                scale_a = torch.tensor(1.0, device=a.device)
-                scale_b = torch.tensor(1.0, device=a.device)
+                scale_a = _get_scale_per_tensor(a, custom_scale=self.extra_args.per_tensor_scale_a)
+                scale_b = _get_scale_per_tensor(b, custom_scale=self.extra_args.per_tensor_scale_b)
 
             # Kernels expect dtype=float8_e4m3fn
             a = a.to(torch.float8_e4m3fn)
@@ -103,16 +118,10 @@ def get_x_val(self, example_inputs) -> float:
         _, n = b.size()
         return (m, n, k)
 
-    def _get_out_dtype(self):
-        if self.extra_args.scaling_rowwise:
-            return torch.bfloat16
-        else:
-            return torch.float16
-
     @register_benchmark(baseline=True)
     def torch_fp8_gemm(self, a, b, scale_a, scale_b):
         return lambda: torch._scaled_mm(
-            a, b, scale_a, scale_b, use_fast_accum=True, out_dtype=self._get_out_dtype()
+            a, b, scale_a, scale_b, use_fast_accum=True, out_dtype=self._get_dtype()
         )
 
     @register_benchmark()
@@ -129,7 +138,7 @@ def pt2_fp8_gemm(self, a, b, scale_a, scale_b) -> Callable:
                 scale_a,
                 scale_b,
                 use_fast_accum=True,
-                out_dtype=self._get_out_dtype(),
+                out_dtype=self._get_dtype(),
             )
             compiled = torch.compile(f, dynamic=False)
             compiled(a, b)