From 900113447ad3348452b1ae74fa640c945c3e4890 Mon Sep 17 00:00:00 2001 From: Will Feng Date: Mon, 14 Jul 2025 12:55:15 -0700 Subject: [PATCH] [Benchmark] Add fp8_attention to tritonbench integration stack-info: PR: https://github.com/pytorch-labs/helion/pull/319, branch: yf225/stack/27 --- benchmarks/run.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/benchmarks/run.py b/benchmarks/run.py index 3f9a2a7f..3a87130d 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -70,6 +70,11 @@ "examples.cross_entropy", "cross_entropy", ), + "fp8_attention": ( + "tritonbench.operators.fp8_attention.operator", + "examples.fp8_attention", + "fp8_attention_tritonbench", + ), } @@ -282,7 +287,7 @@ def helion_method( if isinstance(attr, Kernel): attr.reset() - def _inner() -> Callable[..., Any]: + def _inner() -> Callable[..., Any] | object: # Force autotuning unless HELION_USE_DEFAULT_CONFIG=1 is set # This ensures we run autotuning even if the kernel has pre-specified configs if os.environ.get("HELION_USE_DEFAULT_CONFIG", "0") != "1": @@ -292,7 +297,10 @@ def _inner() -> Callable[..., Any]: if isinstance(attr, Kernel): attr.settings.force_autotune = True - return kernel_func(*args) + result = kernel_func(*args) + if callable(result): + return result() + return result return _inner