Migrate the grouped_gemm inputs from qwen3

xuzhao9 · web-flow · commit 0bb2da4f7c9b · 2025-10-07T21:01:23.000Z
Differential Revision: D84006584 Pull Request resolved: #522
diff --git a/tritonbench/data/input_configs/qwen3/grouped_gemm.json b/tritonbench/data/input_configs/qwen3/grouped_gemm.json
@@ -0,0 +1,280 @@
+{
+    "grouped_gemm": [
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(5120, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(2048, 2048)', '(5120, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(7168, 2048)', '(8192, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(16, 768)', '(16, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(6144, 2048)', '(4096, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(5120, 2048)', '(8192, 2048)', '(7168, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(2048, 2048)', '(7168, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(8192, 2048)', '(5120, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(7168, 2048)', '(7168, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(4096, 2048)', '(8192, 2048)', '(1024, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(8192, 768)', '(5120, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(3072, 2048)', '(8192, 2048)', '(1024, 2048)', '(1024, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(8192, 2048)', '(1024, 2048)', '(8192, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(5120, 768)', '(7168, 768)', '(7168, 768)', '(2048, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(5120, 768)', '(8192, 768)', '(7168, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(4096, 2048)', '(1024, 2048)', '(2048, 2048)', '(8192, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(3072, 768)', '(1024, 768)', '(7168, 768)', '(8192, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(4096, 768)', '(1024, 768)', '(2048, 768)', '(8192, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(2048, 2048)', '(7168, 2048)', '(1024, 2048)', '(5120, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(7168, 768)', '(4096, 768)', '(5120, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(3072, 2048)', '(1024, 2048)', '(2048, 2048)', '(4096, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(5120, 768)', '(1024, 768)', '(7168, 768)', '(3072, 768)', '(7168, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(6144, 768)', '(1024, 768)', '(1024, 768)', '(2048, 768)', '(1024, 768)', '(4096, 768)', '(1024, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(2048, 2048)', '(7168, 2048)', '(1024, 2048)', '(4096, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(2048, 768)', '(7168, 768)', '(1024, 768)', '(4096, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(7168, 768)', '(5120, 768)', '(8192, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(5120, 2048)', '(1024, 2048)', '(7168, 2048)', '(3072, 2048)', '(7168, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(5120, 2048)', '(3072, 2048)', '(6144, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(2048, 768)', '(3072, 768)', '(8192, 768)', '(1024, 768)', '(1024, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(3072, 768)', '(2048, 768)', '(6144, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(8192, 2048)', '(1024, 2048)', '(7168, 2048)', '(6144, 2048)', '(7168, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(2048, 768)', '(1024, 768)', '(3072, 768)', '(2048, 768)', '(4096, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(2048, 2048)', '(3072, 2048)', '(8192, 2048)', '(1024, 2048)', '(1024, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(7168, 2048)', '(4096, 2048)', '(5120, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(3072, 2048)', '(2048, 2048)', '(1024, 2048)', '(4096, 2048)', '(3072, 2048)', '(7168, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(3072, 2048)', '(1024, 2048)', '(4096, 2048)', '(3072, 2048)', '(1024, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(2048, 2048)', '(5120, 2048)', '(4096, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(6144, 2048)', '(1024, 2048)', '(1024, 2048)', '(1024, 2048)', '(3072, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(2048, 2048)', '(5120, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(1024, 768)', '(4096, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(3072, 768)', '(2048, 768)', '(6144, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(6144, 2048)', '(1024, 2048)', '(2048, 2048)', '(2048, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(2048, 768)', '(5120, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(2048, 768)', '(1024, 768)', '(3072, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(4096, 2048)', '(2048, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(2048, 2048)', '(1024, 2048)', '(3072, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(5120, 2048)', '(1024, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(3072, 2048)', '(2048, 2048)', '(5120, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(1024, 2048)', '(1024, 2048)', '(4096, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(6144, 768)', '(1024, 768)', '(2048, 768)', '(2048, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(5120, 768)', '(1024, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(4096, 768)', '(2048, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(3072, 768)', '(2048, 768)', '(5120, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(2048, 768)', '(5120, 768)', '(4096, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(3072, 2048)', '(2048, 2048)', '(6144, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(3072, 768)', '(1024, 768)', '(4096, 768)', '(3072, 768)', '(1024, 768)', '(1024, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(768, 768)', '(640, 768)', '(64, 768)', '(1024, 768)', '(320, 768)', '(256, 768)', '(256, 768)', '(128, 768)', '(960, 768)', '(704, 768)', '(640, 768)', '(640, 768)', '(128, 768)', '(576, 768)', '(256, 768)', '(512, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(256, 2048)', '(512, 2048)', '(64, 2048)', '(832, 2048)', '(576, 2048)', '(512, 2048)', '(704, 2048)', '(448, 2048)', '(448, 2048)', '(2560, 2048)', '(1152, 2048)', '(640, 2048)', '(384, 2048)', '(64, 2048)', '(384, 2048)', '(1024, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(320, 2048)', '(384, 2048)', '(320, 2048)', '(128, 2048)', '(64, 2048)', '(832, 2048)', '(128, 2048)', '(832, 2048)', '(320, 2048)', '(448, 2048)', '(384, 2048)', '(1024, 2048)', '(704, 2048)', '(960, 2048)', '(768, 2048)', '(320, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(1024, 768)', '(384, 768)', '(448, 768)', '(768, 768)', '(640, 768)', '(448, 768)', '(640, 768)', '(1152, 768)', '(704, 768)', '(384, 768)', '(64, 768)', '(384, 768)', '(512, 768)', '(128, 768)', '(704, 768)', '(1600, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(128, 2048)', '(128, 2048)', '(256, 2048)', '(448, 2048)', '(704, 2048)', '(1152, 2048)', '(192, 2048)', '(512, 2048)', '(960, 2048)', '(384, 2048)', '(256, 2048)', '(384, 2048)', '(1600, 2048)', '(128, 2048)', '(512, 2048)', '(1600, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(64, 768)', '(64, 768)', '(64, 768)', '(128, 768)', '(192, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(64, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(64, 2048)', '(64, 2048)', '(64, 2048)', '(128, 2048)', '(192, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(256, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(128, 2048)', '(64, 2048)', '(64, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(256, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(320, 768)', '(64, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(64, 2048)', '(64, 2048)', '(64, 2048)', '(320, 2048)', '(64, 2048)', '(64, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(256, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(64, 768)', '(128, 768)', '(64, 768)', '(64, 768)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(2048, 2048)', 'A_list': ['(256, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(64, 2048)', '(320, 2048)', '(64, 2048)']})"
+        },
+        {
+            "count": 1,
+            "inputs": "((), {'B': '(768, 768)', 'A_list': ['(64, 768)', '(64, 768)', '(64, 768)', '(320, 768)', '(64, 768)', '(64, 768)']})"
+        }
+    ]
+}
diff --git a/tritonbench/data/loader.py b/tritonbench/data/loader.py
@@ -1,10 +1,11 @@
 import importlib
+import os
 from pathlib import Path
 from typing import Any
 
 from tritonbench.utils.env_utils import is_fbcode
 
-SUPPORTED_INPUT_OPS = ["highway_self_gating"]
+SUPPORTED_INPUT_OPS = ["highway_self_gating", "grouped_gemm"]
 
 INPUT_CONFIG_DIR = Path(__file__).parent.joinpath("input_configs")
 INTERNAL_INPUT_CONFIG_DIR = (
@@ -19,7 +20,15 @@ def get_input_loader(tritonbench_op: Any, op: str, input: str):
         hasattr(tritonbench_op, "aten_op_name") or op in SUPPORTED_INPUT_OPS
     ), f"Unsupported op: {op}. "
     op_module = ".".join(tritonbench_op.__module__.split(".")[:-1])
-    generator_module = importlib.import_module(f"{op_module}.input_loader")
-    input_iter_getter = generator_module.get_input_iter
-    input_iter = input_iter_getter(tritonbench_op, op, input)
-    return input_iter
+    generator_module = importlib.import_module(op_module)
+    input_loader_cls = generator_module.InputLoader
+    if os.path.exists(input):
+        input_file_path = Path(input)
+    elif INPUT_CONFIG_DIR.joinpath(input).exists():
+        input_file_path = INPUT_CONFIG_DIR.joinpath(input)
+    elif INTERNAL_INPUT_CONFIG_DIR.joinpath(input).exists():
+        input_file_path = INTERNAL_INPUT_CONFIG_DIR.joinpath(input)
+    else:
+        raise RuntimeError(f"Input file {input} does not exist.")
+    input_loader = input_loader_cls(tritonbench_op, op, input_file_path)
+    return input_loader.get_input_iter()
diff --git a/tritonbench/operator_loader/aten/input_loader.py b/tritonbench/operator_loader/aten/input_loader.py
@@ -7,7 +7,6 @@
 import logging
 import math
 from collections import Counter, defaultdict
-from pathlib import Path
 from typing import Any, Callable, Generator
 
 import torch
@@ -164,7 +163,7 @@ def deserialize_args(inps):
     return eval(inps.strip().strip("'").strip('"'), global_vals)
 
 
-class OperatorInputsLoader:
+class OperatorInputLoader:
     def __init__(self, op_name: str, json_file_path: str):
         self.op_name = op_name
         self.operator_db = defaultdict(Counter)
@@ -229,5 +228,5 @@ def merge(self, other):
 def get_input_iter(tritonbench_op: Any, op: str, input: str) -> Generator:
     aten_op_name = tritonbench_op.aten_op_name
     input_file_path = INPUT_CONFIG_DIR.joinpath(input)
-    operator_inputs_loader = OperatorInputsLoader(aten_op_name, input_file_path)
+    operator_inputs_loader = OperatorInputLoader(aten_op_name, input_file_path)
     return operator_inputs_loader.get_input_iter()
diff --git a/tritonbench/operators/grouped_gemm/__init__.py b/tritonbench/operators/grouped_gemm/__init__.py
@@ -1 +1,2 @@
-from .operator import Operator
+from .input_loader import InputLoader  # noqa
+from .operator import Operator  # noqa
diff --git a/tritonbench/operators/grouped_gemm/input_loader.py b/tritonbench/operators/grouped_gemm/input_loader.py
@@ -0,0 +1,31 @@
+"""
+Input loader for Grouped GEMM operator.
+"""
+
+from typing import Callable
+
+from tritonbench.operator_loader.aten.input_loader import OperatorInputLoader
+
+
+class InputLoader(OperatorInputLoader):
+    def __init__(self, tritonbench_op: str, op_name: str, json_file_path: str):
+        super().__init__(op_name, json_file_path)
+        self.op = tritonbench_op
+
+    def get_input_iter(
+        self,
+    ) -> Callable:
+        shapes = [eval(inp)[1] for inp, _cnt in self.operator_db[self.op_name].items()]
+        parsed = []
+        for entry in shapes:
+            B_shape = (
+                eval(entry["B"]) if isinstance(entry["B"], str) else tuple(entry["B"])
+            )
+            A_shapes = [
+                eval(a) if isinstance(a, str) else tuple(a) for a in entry["A_list"]
+            ]
+            parsed.append((A_shapes, B_shape))
+
+        # Set shapes on the operator
+        self.op.external_shapes = parsed
+        return self.op.get_input_iter

Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`		`-from .operator import Operator`
	`1`	`+from .input_loader import InputLoader # noqa`
	`2`	`+from .operator import Operator # noqa`