From a9871cb1fcbc5cf6170f8bcf989ee6ec0a305d7f Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Fri, 28 Aug 2026 22:50:08 +0000 Subject: [PATCH] [Common] Reduce default cfg's SMEM usage for mhc triton kernel Signed-off-by: Kaining Zhong --- transformer_engine/common/triton/mhc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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):