From 14a3fb9bef19c49693c4f8b5775756ff2f26e599 Mon Sep 17 00:00:00 2001 From: Chulian Zhang Date: Thu, 3 Sep 2026 17:51:15 +0000 Subject: [PATCH] [https://nvbugs/6193837][fix] Include FINALIZE-fusion workspace for SM>=90 in the MoE autotuner calcMaxWorkspaceSize() only accounted for the FINALIZE-fusion epilogue workspace when sm_ == 90. On SM>=90 the autotuner can still select FINALIZE-fusion tactics, but their workspace requirement was excluded from the computed max, allowing the launcher to size the workspace buffer too small for the tactic actually picked at runtime. Widen the guard to sm_ >= 90 so Blackwell and newer architectures also reserve the FINALIZE-fusion workspace in the max-size calculation. Hopper (sm_ == 90) behavior is unchanged. Signed-off-by: Chulian Zhang Signed-off-by: farazkh80 <58580514+farazkh80@users.noreply.github.com> --- .../cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h index f65e4cd1911c..565ca7fe48f3 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h @@ -988,7 +988,7 @@ size_t MoeGemmRunner::calcMaxWorkspace } while (0) CALC_SIZE_FUSION(TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE); - if (sm_ == 90) + if (sm_ >= 90) { CALC_SIZE_FUSION(TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE); }