Focus on attention and gemm

Author: yf225

benchmarks/run.py

@@ -28,53 +28,53 @@
 # Maps tritonbench op names to Helion kernel examples
 KERNEL_MAPPINGS: dict[str, tuple[str, str, str]] = {
     # <tritonbench_op_name>: (<tritonbench_module_path>, <helion_kernel_module_path>, <helion_kernel_function_name>)
-    "vector_add": ("tritonbench.operators.vector_add.operator", "examples.add", "add"),
-    "embedding": (
-        "tritonbench.operators.embedding.operator",
-        "examples.embedding",
-        "embedding_tritonbench",
-    ),
-    "vector_exp": (
-        "tritonbench.operators.vector_exp.operator",
-        "examples.exp",
-        "exp_tritonbench",
-    ),
-    "rms_norm": (
-        "tritonbench.operators.rms_norm.operator",
-        "examples.rms_norm",
-        "rms_norm_tritonbench",
-    ),
-    "sum": ("tritonbench.operators.sum.operator", "examples.sum", "sum_tritonbench"),
-    "softmax": (
-        "tritonbench.operators.softmax.operator",
-        "examples.softmax",
-        "softmax",
-    ),
-    "jagged_mean": (
-        "tritonbench.operators.jagged_mean.operator",
-        "examples.jagged_mean",
-        "jagged_mean_tritonbench",
-    ),
-    "fp8_gemm": (
-        "tritonbench.operators.fp8_gemm.fp8_gemm",
-        "examples.fp8_gemm",
-        "fp8_gemm_tritonbench",
-    ),
+    # "vector_add": ("tritonbench.operators.vector_add.operator", "examples.add", "add"),
+    # "embedding": (
+    #     "tritonbench.operators.embedding.operator",
+    #     "examples.embedding",
+    #     "embedding_tritonbench",
+    # ),
+    # "vector_exp": (
+    #     "tritonbench.operators.vector_exp.operator",
+    #     "examples.exp",
+    #     "exp_tritonbench",
+    # ),
+    # "rms_norm": (
+    #     "tritonbench.operators.rms_norm.operator",
+    #     "examples.rms_norm",
+    #     "rms_norm_tritonbench",
+    # ),
+    # "sum": ("tritonbench.operators.sum.operator", "examples.sum", "sum_tritonbench"),
+    # "softmax": (
+    #     "tritonbench.operators.softmax.operator",
+    #     "examples.softmax",
+    #     "softmax",
+    # ),
+    # "cross_entropy": (
+    #     "tritonbench.operators.cross_entropy.operator",
+    #     "examples.cross_entropy",
+    #     "cross_entropy",
+    # ),
+    # "jagged_mean": (
+    #     "tritonbench.operators.jagged_mean.operator",
+    #     "examples.jagged_mean",
+    #     "jagged_mean_tritonbench",
+    # ),
     "flash_attention": (
         "tritonbench.operators.flash_attention.operator",
         "examples.attention",
         "attention",
     ),
-    "cross_entropy": (
-        "tritonbench.operators.cross_entropy.operator",
-        "examples.cross_entropy",
-        "cross_entropy",
-    ),
     "fp8_attention": (
         "tritonbench.operators.fp8_attention.operator",
         "examples.fp8_attention",
         "fp8_attention_tritonbench",
     ),
+    "fp8_gemm": (
+        "tritonbench.operators.fp8_gemm.fp8_gemm",
+        "examples.fp8_gemm",
+        "fp8_gemm_tritonbench",
+    ),
 }
 
 

helion/autotuner/base_search.py

@@ -137,8 +137,8 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
         except CompilationError:
             self.log.debug("Benchmarking failed: Triton CompilationError")
         except Exception as e:
-            if not _expected_errors_regexp.search(str(e)) and not "exceeds triton maximum tensor numel" in str(e):
-                raise exc.TritonError(f"{type(e).__qualname__}: {e}", config) from e
+            # if not _expected_errors_regexp.search(str(e)) and not "exceeds triton maximum tensor numel" in str(e):
+            #     raise exc.TritonError(f"{type(e).__qualname__}: {e}", config) from e
             self.log.debug(f"Benchmarking failed: {type(e).__name__}: {e}")
         return inf