[TRITON] Add Attention support to the bench_models benchmarking script (#2274)

* Add Attention support to bench_models.py

* Add MHA layout CLI arg

* Add support for batched_gemm_a16wfp4

* Refactor TP logic and _get_handler

* Remove unified attention from this branch

Author: lucas-santos-amd

op_tests/op_benchmarks/triton/bench_batched_gemm_a16wfp4.py

@@ -2,6 +2,13 @@
 import torch
 import triton
 import math
+import aiter.ops.triton.utils._triton.arch_info as arch_info
+from aiter.ops.triton.gemm.batched.batched_gemm_afp4wfp4_pre_quant import (
+    batched_gemm_afp4wfp4_pre_quant,
+)
+from op_tests.triton_tests.gemm.batched.test_batched_gemm_a16wfp4 import (
+    generate_batched_gemm_a16wfp4_inputs,
+)
 from op_tests.op_benchmarks.triton.utils.argparse import (
     get_parser,
     add_argparse_ff,
@@ -14,10 +21,6 @@
     print_vgpr,
     get_caller_name_no_ext,
 )
-from aiter.ops.triton.gemm.batched.batched_gemm_afp4wfp4_pre_quant import (
-    batched_gemm_afp4wfp4_pre_quant,
-)
-import aiter.ops.triton.utils._triton.arch_info as arch_info
 
 
 def bench_gemm_fn(
@@ -29,7 +32,7 @@ def bench_gemm_fn(
     layout: str,
 ):
     c_dtype = torch.bfloat16
-    x, w, x_scale, w_scale, y = generate_batched_gemm_afp4wfp4_pre_quant_inputs(
+    x, w, x_scale, w_scale, y = generate_batched_gemm_a16wfp4_inputs(
         batch, M, N, K, c_dtype, layout=layout, output=True
     )
     # flops
@@ -145,7 +148,7 @@ def run_benchmark(args, defaults):
         run_shape_benchmark(args)
 
 
-def parse_args():
+def parse_args(args: list[str] | None = None):
     parser = get_parser("Batched MXFP4 x MXFP4 GEMM, Pre Quant")
     parser = add_argparse_ff(parser)
     parser.add_argument(
@@ -154,22 +157,22 @@ def parse_args():
         required=False,
         help="Batch size to be used when using --model flag.",
     )
-    return get_ff_args(parser)
+    return get_ff_args(parser, args=args)
 
 
-def main():
+def main(args: list[str] | None = None):
     if not (arch_info.is_fp4_avail()):
         print("MXFP4 is not available on this architecture")
         sys.exit()
 
-    args, defaults = parse_args()
+    args, defaults = parse_args(args=args)
     if args.print_vgpr:
         print("Retrieving VGPR usage for Triton kernels...")
         fun = lambda: run_benchmark(args, defaults)  # noqa: E731
         print_vgpr(fun, get_caller_name_no_ext())
-        return 0
+        return
     run_benchmark(args, defaults)
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    main()

op_tests/op_benchmarks/triton/bench_mla_decode.py

@@ -62,7 +62,6 @@ def model_benchmark_configs(args):
 
 def benchmark(args):
     dtype = str_to_torch_dtype[args.dtype]
-    torch.set_default_dtype(dtype)
 
     configs = []