expand search space for hstu gemm

nmacchioni · facebook-github-bot · commit b151b84011ec · 2024-11-19T23:43:13.000-08:00
Summary: before we were using just a single config

Reviewed By: xuzhao9

Differential Revision: D66213893

fbshipit-source-id: 6d5a149d496b24cb80e9e56a8cb1d78010da5e0b
diff --git a/tritonbench/operators/gemm/operator.py b/tritonbench/operators/gemm/operator.py
@@ -40,6 +40,12 @@
 )
 
 if IS_FBCODE:
+    import hammer.oss.generative_recommenders.ops.triton.triton_addmm as hstu_triton_addmm
+
+    # without this set we can only pick a single config for AMD, Nvidia has 8
+    # with this set AMD will pick from 256 different configs (not the actual full
+    # tuning space, so some perf may be left on the table)
+    hstu_triton_addmm.ENABLE_FULL_TURNING_SPACE = True
     from hammer.ops.triton.triton_matmul import (
         triton_matmul as hstu_triton_matmul_kernel,
     )

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,12 @@`
`40`	`40`	`)`
`41`	`41`
`42`	`42`	`if IS_FBCODE:`
	`43`	`+ import hammer.oss.generative_recommenders.ops.triton.triton_addmm as hstu_triton_addmm`
	`44`	`+`
	`45`	`+ # without this set we can only pick a single config for AMD, Nvidia has 8`
	`46`	`+ # with this set AMD will pick from 256 different configs (not the actual full`
	`47`	`+ # tuning space, so some perf may be left on the table)`
	`48`	`+ hstu_triton_addmm.ENABLE_FULL_TURNING_SPACE = True`
`43`	`49`	`from hammer.ops.triton.triton_matmul import (`
`44`	`50`	`triton_matmul as hstu_triton_matmul_kernel,`
`45`	`51`	`)`