num_stages=0 becomes num_stages=2

Differential Revision: D65955484
facebookresearch · Nov 14, 2024 · 663c64b · 663c64b
1 parent ebb0b05
commit 663c64b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/generative_recommenders/ops/triton/triton_addmm.py b/generative_recommenders/ops/triton/triton_addmm.py
@@ -40,7 +40,7 @@ def get_mm_configs() -> List[triton.Config]:
             waves_per_eu_range = [0]
             kpack_range = [1, 2]
             num_warps_range = [4, 8]
-            num_stage_range = [0]
+            num_stage_range = [2]
         else:
             block_m_range = [256]
             block_n_range = [256]
@@ -50,7 +50,7 @@ def get_mm_configs() -> List[triton.Config]:
             waves_per_eu_range = [0]
             kpack_range = [2]
             num_warps_range = [8]
-            num_stage_range = [0]
+            num_stage_range = [2]
 
         return [
             triton.Config(