Skip to content

Commit 9001134

Browse files
committed
[Benchmark] Add fp8_attention to tritonbench integration
stack-info: PR: #319, branch: yf225/stack/27
1 parent 0a2a04e commit 9001134

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

benchmarks/run.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,11 @@
7070
"examples.cross_entropy",
7171
"cross_entropy",
7272
),
73+
"fp8_attention": (
74+
"tritonbench.operators.fp8_attention.operator",
75+
"examples.fp8_attention",
76+
"fp8_attention_tritonbench",
77+
),
7378
}
7479

7580

@@ -282,7 +287,7 @@ def helion_method(
282287
if isinstance(attr, Kernel):
283288
attr.reset()
284289

285-
def _inner() -> Callable[..., Any]:
290+
def _inner() -> Callable[..., Any] | object:
286291
# Force autotuning unless HELION_USE_DEFAULT_CONFIG=1 is set
287292
# This ensures we run autotuning even if the kernel has pre-specified configs
288293
if os.environ.get("HELION_USE_DEFAULT_CONFIG", "0") != "1":
@@ -292,7 +297,10 @@ def _inner() -> Callable[..., Any]:
292297
if isinstance(attr, Kernel):
293298
attr.settings.force_autotune = True
294299

295-
return kernel_func(*args)
300+
result = kernel_func(*args)
301+
if callable(result):
302+
return result()
303+
return result
296304

297305
return _inner
298306

0 commit comments

Comments
 (0)