From bbfaff47edf1c9d9a3b4c881fdb13d56f91ca33b Mon Sep 17 00:00:00 2001 From: Po-Han Huang Date: Sun, 2 Aug 2026 22:10:00 -0700 Subject: [PATCH] fix(moe): pad BF16 TRTLLM-Gen intermediates to 128 KiB Signed-off-by: Po-Han Huang --- csrc/trtllm_fused_moe_kernel_launcher.cu | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 02485a73e90..f9fb0dc91c5 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -852,10 +852,19 @@ class Bf16MoeLauncher : public FusedMoeLauncher { void prepare_moe(int64_t& moe_tactic) override { FusedMoeLauncher::prepare_moe_common(moe_tactic); + // Blackwell TMA may select BASE_128KB address generation from the logical tensor-map shape. + // Keep at least 128 KiB mapped from each activation base, as the quantized launchers do below. int32_t max_num_padded_tokens = workspace.total_max_padded_tokens; - gemm1_output = alloc_tensor({max_num_padded_tokens, args->intermediate_size}, dl_bfloat16, + int32_t max_num_padded_tokens_gemm1 = + tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + max_num_padded_tokens, args->intermediate_size, + btg::dtypeGetNumBits(btg::Dtype::Bfloat16)); + int32_t max_num_padded_tokens_gemm2 = + tensorrt_llm::kernels::trtllmgen_moe::Routing::maybeGetMinTokenCount( + max_num_padded_tokens, args->hidden_size, btg::dtypeGetNumBits(btg::Dtype::Bfloat16)); + gemm1_output = alloc_tensor({max_num_padded_tokens_gemm1, args->intermediate_size}, dl_bfloat16, hidden_states.device()); - gemm2_output = alloc_tensor({max_num_padded_tokens, args->hidden_size}, dl_bfloat16, + gemm2_output = alloc_tensor({max_num_padded_tokens_gemm2, args->hidden_size}, dl_bfloat16, hidden_states.device()); workspace.hidden_states_scale_linear = nullptr;