Add benchmarking on shapes from CSV files to fp8_gemm

Summary: Add ability to benchmark fp8_gemm kernels on shapes from CSV files. Add CLI argument to consume file path for CSV file.

Differential Revision: D80381352

Author: jananisriram

tritonbench/operators/fp8_gemm/fp8_gemm.py

@@ -1,5 +1,7 @@
 import argparse
 import logging
+import csv
+import os
 
 from typing import Any, Callable, List, Optional
 
@@ -33,16 +35,29 @@
     HAS_TMA = False
     logger.warning(f"Failed to import TMA: {e}")
 
-
 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("--m", type=int)
     parser.add_argument("--k", type=int)
     parser.add_argument("--n", type=int)
+    parser.add_argument("--filepath", type=str, default=None)
     return parser.parse_args(args)
 
+def read_fp8_shapes(filepath):
+    fp8_shapes = []
+    try:
+        with open(filepath, 'r', newline='') as csvfile:
+            filtered_lines = (line for line in csvfile if line.strip() and not line.lstrip().startswith('#'))
+            reader = csv.reader(filtered_lines)
+            for row in reader:
+                fp8_shapes.append(tuple(map(int, row)))
+    except Exception as e:
+        logger.error(f"Failed to read fp8 shapes from {filepath}: {e}")
+        raise e
+    return fp8_shapes
+
 
 class Operator(BenchmarkOperator):
     DEFAULT_METRICS = ["tflops", "gbps", "latency"]
@@ -70,6 +85,10 @@ def args(m, n, k):
                 yield args(m, n, k)
         elif self.extra_args.m:
             yield args(self.extra_args.m, self.extra_args.n, self.extra_args.k)
+        elif self.extra_args.filepath:
+            fp8_shapes = read_fp8_shapes(self.extra_args.filepath)
+            for m, n, k in fp8_shapes:
+                yield args(m, n, k)
         else:
             for i in range(10, 15):
                 for j in range(0, 4):
@@ -114,8 +133,8 @@ def pt2_fp8_gemm(self, a, b) -> Callable:
                 scale_b = torch.ones((1, N), dtype=torch.float32, device=b.device)
                 out_dtype = torch.bfloat16
             else:
-                scale_a = torch.tensor(1.0, device=a.device)
-                scale_b = torch.tensor(1.0, device=a.device)
+                scale_a = torch.tensor(1.0, dtype=torch.float32, device=a.device)
+                scale_b = torch.tensor(1.0, dtype=torch.float32, device=a.device)
                 out_dtype = torch.float16
             f = lambda a, b: torch._scaled_mm(
                 a, b, scale_a, scale_b, use_fast_accum=True, out_dtype=out_dtype