[TRITON]: Benchmarking scripts updates

aiter/ops/triton/configs/MI350X-EXTEND_ATTENTION.json

@@ -0,0 +1,10 @@
+{
+  "default": {
+      "BLOCK_M": 64,
+      "BLOCK_N": 64,
+      "waves_per_eu": 1,
+      "num_warps": 4,
+      "matrix_instr_nonkdim": 16,
+      "kpack": 2
+  }
+}

aiter/ops/triton/configs/moe/MI350X-MOE_ROUTING_SIGMOID_TOPK1.json

@@ -0,0 +1,70 @@
+{
+  "N16": {
+    "small" :{
+      "BLOCK_M": 16,
+      "BLOCK_K": 256,
+      "num_warps": 4,
+      "num_stages": 2,
+      "waves_per_eu": 3,
+      "kpack": 1
+    },
+    "medium" :{
+      "BLOCK_M": 16,
+      "BLOCK_K": 256,
+      "num_warps": 4,
+      "num_stages": 2,
+      "waves_per_eu": 3,
+      "kpack": 1
+    },
+    "large" :{
+      "BLOCK_M": 16,
+      "BLOCK_K": 256,
+      "num_warps": 4,
+      "num_stages": 2,
+      "waves_per_eu": 3,
+      "kpack": 2
+    },
+    "xlarge" :{
+      "BLOCK_M": 32,
+      "BLOCK_K": 128,
+      "num_warps": 8,
+      "num_stages": 2,
+      "waves_per_eu": 2,
+      "kpack": 2
+    }
+  },
+  "N128": {
+    "small" :{
+      "BLOCK_M": 16,
+      "BLOCK_K": 256,
+      "num_warps": 8,
+      "num_stages": 1,
+      "waves_per_eu": 0,
+      "kpack": 1
+    },
+    "medium" :{
+      "BLOCK_M": 16,
+      "BLOCK_K": 256,
+      "num_warps": 8,
+      "num_stages": 1,
+      "waves_per_eu": 0,
+      "kpack": 2
+    },
+    "large" :{
+      "BLOCK_M": 16,
+      "BLOCK_K": 256,
+      "num_warps": 8,
+      "num_stages": 1,
+      "waves_per_eu": 2,
+      "kpack": 2
+    },
+    "xlarge" :{
+      "BLOCK_M": 32,
+      "BLOCK_K": 128,
+      "num_warps": 8,
+      "num_stages": 2,
+      "waves_per_eu": 2,
+      "kpack": 2
+    }
+  }
+}

aiter/ops/triton/moe_routing_sigmoid_top1_fused.py

@@ -11,67 +11,6 @@
 from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
 
 
-def get_config_heuristic(M, K, N):
-    """
-    Return the best Triton configuration based on input dimensions.
-
-    Args:
-        M: Batch dimension
-        K: Hidden dimension
-        N: Number of experts (16 or 128)
-        TOPK: Top-k value (default: 1)
-
-    Returns:
-        triton.Config: Configuration for the Triton kernel
-    """
-    # Determine M bucket (small: <2048, medium: 2048-4095, large: 4096-8191, very_large: 8192+)
-    m_bucket = (
-        "very_large"
-        if M >= 8192
-        else "large" if M >= 4096 else "medium" if M >= 2048 else "small"
-    )
-
-    # Create parameter configuration using nested dictionaries
-    configs = {
-        # Format: {N: {m_bucket: (BLOCK_M, BLOCK_K, num_warps, num_stages, waves_per_eu, kpack)}}
-        16: {
-            "small": (16, 256, 4, 2, 3, 1),
-            "medium": (16, 256, 4, 2, 3, 1),
-            "large": (16, 256, 4, 2, 3, 2),
-            "very_large": (32, 256, 4, 2, 0, 1),
-        },
-        128: {
-            "small": (16, 256, 8, 1, 0, 1),
-            "medium": (16, 256, 8, 1, 0, 2),
-            "large": (16, 256, 8, 1, 2, 2),
-            "very_large": (32, 128, 8, 2, 2, 2),
-        },
-        256: {
-            "small": (16, 64, 8, 1, 0, 1),
-            "medium": (16, 64, 8, 1, 0, 2),
-            "large": (16, 64, 8, 1, 2, 2),
-            "very_large": (16, 64, 8, 2, 2, 2),
-        },
-    }
-
-    # Get configuration parameters
-    BLOCK_M, BLOCK_K, num_warps, num_stages, waves_per_eu, kpack = configs[N][m_bucket]
-
-    # Return Triton configuration
-    return triton.Config(
-        {
-            "BLOCK_M": BLOCK_M,
-            "BLOCK_K": BLOCK_K,
-            "matrix_instr_nonkdim": 16,  # Always 16
-            "waves_per_eu": waves_per_eu,
-            "kpack": kpack,
-        },
-        num_warps=num_warps,
-        num_stages=num_stages,
-        num_ctas=1,
-    )
-
-
 @triton.jit
 def _routing_sigmoid_top1_kernel(
     X_ptr,

op_tests/op_benchmarks/triton/bench_batched_gemm_afp4wfp4.py

@@ -18,28 +18,39 @@
 from aiter.ops.triton.batched_gemm_afp4wfp4 import (
     batched_gemm_afp4wfp4 as batched_gemm_afp4wfp4,
 )
+import aiter.ops.triton.utils.arch_info as arch_info
 
 
 def model_benchmark_shapes(args):
     config_file = args.model_configs
     configs = get_model_configs(config_path=config_file, models=args.model)
-    M_list = [args.M] if args.model == "all" else [2**i for i in range(0, 15)]
+    M_list = [4096] if args.model == "all" else [2**i for i in range(0, 15)]
     shapes = []
     for M in M_list:
-        for _, config in configs.items():
+        for model_name, config in configs.items():
             N = config["intermediate_size"]
             K = config["hidden_size"]
 
             shapes.append(
-                (M, N, K, 16)
+                (model_name, M, N, K, 16)
             )  # rearrange batch to last dim so M is graph x-axis
 
     return shapes
 
 
-def bench_gemm_fn(batch, M, N, K, metric):
+def bench_gemm_fn(
+    batch: int, M: int, N: int, K: int, metric: str, layout: str, model_name=None
+):
     c_dtype = torch.bfloat16
-    x, w, x_scale, w_scale = generate_batched_gemm_afp4wfp4_inputs(batch, M, N, K)
+    x, w, x_scale, w_scale, y = generate_batched_gemm_afp4wfp4_inputs(
+        batch,
+        M,
+        N,
+        K,
+        c_dtype,
+        layout=layout,
+        output=True,
+    )
     # print(f"M: {M}, N: {N}, K: {K}, x.shape: {x.shape}, x.stride(): {x.stride()}, w.shape: {w.shape}, w.stride(): {w.stride()}")
     # flops
     flops = 2.0 * M * N * K * batch
@@ -51,12 +62,9 @@ def bench_gemm_fn(batch, M, N, K, metric):
     )
     mem_write = (M * N) * 2  # TODO: Fix for c_dtype != bf16
     mem = mem_read + mem_write
-    out = torch.empty(
-        x.shape[0], x.shape[1], w.shape[2], device=x.device, dtype=c_dtype
-    )
 
     ms = triton.testing.do_bench(
-        lambda: batched_gemm_afp4wfp4(x, w, x_scale, w_scale, c_dtype, out),
+        lambda: batched_gemm_afp4wfp4(x, w, x_scale, w_scale, c_dtype, y),
         warmup=25,
         rep=100,
     )
@@ -78,13 +86,13 @@ def run_model_benchmark(args):
     benchmark = get_model_benchmark_object(
         plot_name="Batched GEMM MXFP4 x MXFP4 Benchmark",
         args=args,
-        x_names=["M", "hidden_dim", "intermediate_dim", "batch"],
+        x_names=["model_name", "M", "hidden_dim", "intermediate_dim", "batch"],
         model_benchmark_shapes_fn=model_benchmark_shapes,
     )
 
     @triton.testing.perf_report([benchmark])
     def bench_batched_gemm_afp4wfp4(
-        M, hidden_dim, intermediate_dim, batch, metric, layer, **kwargs
+        M, hidden_dim, intermediate_dim, batch, metric, layer, model_name=None, **kwargs
     ):
         if layer == "fc1":
             if args.no_glu:
@@ -99,9 +107,9 @@ def bench_batched_gemm_afp4wfp4(
             K = math.ceil(K / args.tp)
         # print(f"Layer: {layer}, B: {batch}, M: {M}, N: {N}, K: {K}, hidden_dim: {hidden_dim}, intermediate_dim: {intermediate_dim}")
 
-        return bench_gemm_fn(batch, M, N, K, metric)
+        return bench_gemm_fn(batch, M, N, K, metric, layout=args.layout)
 
-    bench_batched_gemm_afp4wfp4.run(save_path=".", print_data=True)
+    bench_batched_gemm_afp4wfp4.run(save_path="." if args.o else None, print_data=True)
 
 
 def run_shape_benchmark(args):
@@ -112,10 +120,10 @@ def run_shape_benchmark(args):
     )
 
     @triton.testing.perf_report([benchmark])
-    def bench_batched_gemm_afp4wfp4(M, N, K, batch, metric, provider):
-        return bench_gemm_fn(batch, M, N, K, metric)
+    def bench_batched_gemm_afp4wfp4(M, N, K, batch, metric, provider, model_name=None):
+        return bench_gemm_fn(batch, M, N, K, metric, layout=args.layout)
 
-    bench_batched_gemm_afp4wfp4.run(save_path=".", print_data=True)
+    bench_batched_gemm_afp4wfp4.run(save_path="." if args.o else None, print_data=True)
 
 
 def run_benchmark(args, defaults):
@@ -124,9 +132,7 @@ def run_benchmark(args, defaults):
     ), "User can specify --shape or --model MODEL -M VAL exclusively"
 
     if args.model:
-        unsupported_args = [
-            "layout",
-        ]
+        unsupported_args = []
         for arg in unsupported_args:
             if getattr(args, arg, None) != getattr(defaults, arg, None):
                 raise Exception(
@@ -154,6 +160,10 @@ def parse_args():
 
 
 def main():
+    if not (arch_info.is_fp4_avail()):
+        print("MXFP4 is not available on this architecture")
+        sys.exit()
+
     args, defaults = parse_args()
     run_benchmark(args, defaults)