Refactor fp8_gemm benchmark to simplify addition of new scaling modes (#500)

Summary:

Refactor the `fp8_gemm` benchmark in TritonBench to accept scaling modes as an argument. This diff enables us to extend the `fp8_gemm` benchmark to new scaling modes without adding new benchmarking arguments.

Differential Revision: D83617233

Author: jananisriram

tritonbench/operators/fp8_gemm/fp8_gemm.py

@@ -1,6 +1,11 @@
 import argparse
+
+from enum import IntEnum
+
 import logging
 
+from torch._inductor.kernel.mm import ScalingMode
+
 from typing import Any, Callable, List, Optional
 
 import torch
@@ -46,7 +51,7 @@
 def parse_args(args):
     parser = argparse.ArgumentParser(description="TritonBench fp8_gemm")
     parser.add_argument("--llama", action="store_true")
-    parser.add_argument("--scaling_rowwise", action="store_true")
+    parser.add_argument("--scaling-mode", type=str, default="tensor")
     parser.add_argument("--m", type=int)
     parser.add_argument("--k", type=int)
     parser.add_argument("--n", type=int)
@@ -55,6 +60,15 @@ def parse_args(args):
     return parser.parse_args(args)
 
 
+def get_scaling_mode_int(scaling_mode: str) -> int:
+    if scaling_mode == "tensor":
+        return ScalingMode.TENSOR
+    elif scaling_mode == "row":
+        return ScalingMode.ROW
+    else:
+        raise ValueError(f"Invalid scaling mode: {scaling_mode}")
+
+
 class Operator(BenchmarkOperator):
     DEFAULT_METRICS = ["tflops", "gbps", "latency"]
     DEFAULT_PRECISION = "fp8"
@@ -65,11 +79,12 @@ def __init__(
         super().__init__(tb_args, extra_args)
         self.extra_args = parse_args(extra_args)
 
+        self.scaling_mode_int = get_scaling_mode_int(self.extra_args.scaling_mode).value
+
     def _get_dtype(self):
-        if self.extra_args.scaling_rowwise:
-            return torch.bfloat16
-        else:
+        if self.scaling_mode_int == ScalingMode.TENSOR:
             return torch.float16
+        return torch.bfloat16
 
     def get_input_iter(self):
         def _get_scale_per_tensor(
@@ -102,10 +117,10 @@ def args(m, n, k):
             a = torch.randn(m, k, device=self.device).to(self._get_dtype())
             b = torch.randn(n, k, device=self.device).to(self._get_dtype())
 
-            if self.extra_args.scaling_rowwise:
+            if self.scaling_mode_int == ScalingMode.ROW:
                 scale_a = _get_scale_per_row(a)
                 scale_b = _get_scale_per_row(b)
-            else:
+            else:  # self.scaling_mode_int == ScalingMode.TENSOR
                 scale_a = _get_scale_per_tensor(
                     a, custom_scale=self.extra_args.per_tensor_scale_a
                 )
@@ -191,7 +206,7 @@ def blackwell_persistent_tma_fp8_gemm(self, a, b, scale_a, scale_b):
                 scale_a,
                 scale_b,
                 self._get_dtype(),
-                self.extra_args.scaling_rowwise,
+                self.scaling_mode_int,
             )
 
         @register_benchmark(enabled=True)

tritonbench/operators/fp8_gemm/persistent.py

@@ -1,4 +1,9 @@
+from enum import IntEnum
+
 from functools import lru_cache
+
+from torch._inductor.kernel.mm import ScalingMode
+
 from typing import Optional
 
 import torch
@@ -23,7 +28,6 @@
     except (ImportError, IOError, AttributeError):
         pass
 
-
 def _matmul_launch_metadata(grid, kernel, args):
     ret = {}
     M, N, K = args["M"], args["N"], args["K"]
@@ -410,9 +414,7 @@ def matmul_tma_persistent(a, b, c, desc_a, desc_b, desc_c):
 #       - 1 warp = 32 threads, so each thread block requires 128 / 32 = 4 warps
 
 
-def blackwell_persistent_tma(
-    a, b, scale_a_ptr, scale_b_ptr, acc_dtype, scaling_rowwise
-):
+def blackwell_persistent_tma(a, b, scale_a_ptr, scale_b_ptr, acc_dtype, scaling_mode):
     configs = matmul_configs_blackwell()
 
     # Check constraints.
@@ -471,7 +473,7 @@ def alloc_fn(size: int, align: int, stream: Optional[int]):
         NUM_SMS=NUM_SMS,  #
         num_stages=configs[shape_dtype]["num_stages"],  #
         num_warps=configs[shape_dtype]["num_warps"],  #
-        SCALING_ROWWISE=scaling_rowwise,
+        SCALING_MODE=scaling_mode,  #
         WARP_SPECIALIZE=configs[shape_dtype]["WARP_SPECIALIZE"],  #
         EPILOGUE_SUBTILE=configs[shape_dtype]["EPILOGUE_SUBTILE"],  #
     )
@@ -504,7 +506,7 @@ def blackwell_persistent_tma_kernel(
     GROUP_SIZE_M: tl.constexpr,  #
     ACC_TYPE: tl.constexpr,
     NUM_SMS: tl.constexpr,
-    SCALING_ROWWISE: tl.constexpr,  #
+    SCALING_MODE: tl.constexpr,  #
     WARP_SPECIALIZE: tl.constexpr,
     EPILOGUE_SUBTILE: tl.constexpr,
 ):  #
@@ -538,7 +540,7 @@ def blackwell_persistent_tma_kernel(
     tile_id_c = start_pid - NUM_SMS
     num_pid_in_group = GROUP_SIZE_M * num_pid_n
 
-    if SCALING_ROWWISE:
+    if SCALING_MODE == ScalingMode.ROW:
         # For row-wise scaling, we'll use the pointers as-is
         scale_a = scale_a_ptr
         scale_b = scale_b_ptr
@@ -563,7 +565,7 @@ def blackwell_persistent_tma_kernel(
             b_block = b_desc.load([offs_bn, offs_k])
             accumulator = tl.dot(a_block, b_block.T, accumulator, out_dtype=tl.float32)
 
-        if SCALING_ROWWISE:
+        if SCALING_MODE == ScalingMode.ROW:
             offs_scale_m = offs_am + tl.arange(0, BLOCK_SIZE_M)
             offs_scale_n = offs_bn + tl.arange(0, BLOCK_SIZE_N)