diff --git a/transformer_engine/common/triton/mhc.py b/transformer_engine/common/triton/mhc.py
index ddad878e1a..cc8f4fae91 100644
--- a/transformer_engine/common/triton/mhc.py
+++ b/transformer_engine/common/triton/mhc.py
@@ -51,7 +51,7 @@ def projection_prune_fwd(configs, named_args, **kwargs):
         block_k = align_to(K, 32)
         step_k = [256]
         warps = [2, 8]
-        stages = [3, 4]
+        stages = [2, 3, 4]
 
         pruned_configs = []
         for bm, sk, w, s in itertools.product(block_m, step_k, warps, stages):