File tree Expand file tree Collapse file tree 1 file changed +10
-2
lines changed Expand file tree Collapse file tree 1 file changed +10
-2
lines changed Original file line number Diff line number Diff line change 70
70
"examples.cross_entropy" ,
71
71
"cross_entropy" ,
72
72
),
73
+ "fp8_attention" : (
74
+ "tritonbench.operators.fp8_attention.operator" ,
75
+ "examples.fp8_attention" ,
76
+ "fp8_attention_tritonbench" ,
77
+ ),
73
78
}
74
79
75
80
@@ -282,7 +287,7 @@ def helion_method(
282
287
if isinstance (attr , Kernel ):
283
288
attr .reset ()
284
289
285
- def _inner () -> Callable [..., Any ]:
290
+ def _inner () -> Callable [..., Any ] | object :
286
291
# Force autotuning unless HELION_USE_DEFAULT_CONFIG=1 is set
287
292
# This ensures we run autotuning even if the kernel has pre-specified configs
288
293
if os .environ .get ("HELION_USE_DEFAULT_CONFIG" , "0" ) != "1" :
@@ -292,7 +297,10 @@ def _inner() -> Callable[..., Any]:
292
297
if isinstance (attr , Kernel ):
293
298
attr .settings .force_autotune = True
294
299
295
- return kernel_func (* args )
300
+ result = kernel_func (* args )
301
+ if callable (result ):
302
+ return result ()
303
+ return result
296
304
297
305
return _inner
298
306
You can’t perform that action at this time.
0 commit comments