diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index e32b7d461175..50c0f990e402 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -3332,12 +3332,8 @@ CutlassMoeFCRunner:: size_t factor = is_gated_activation ? 2 : 1; size_t blockscale_fc1_output_size = factor * interbuf_elems * gemm_output_dtype; size_t blockscale_fc2_output_size = permuted_elems * gemm_output_dtype; - overlapped_gemm1_gemm2_inputs_size - = std::max(std::max(permuted_data_size, fc1_result_size), blockscale_fc2_output_size); - // The fused pre-FC2 path writes the FC2 GEMM output into the outputs buffer (glu_inter_result_) - // instead of the aliased fc2_result_, so size it for the larger of the FC1 raw output and the - // FC2 output. - overlapped_gemm1_gemm2_outputs_size = std::max(blockscale_fc1_output_size, blockscale_fc2_output_size); + overlapped_gemm1_gemm2_inputs_size = std::max(permuted_data_size, fc1_result_size); + overlapped_gemm1_gemm2_outputs_size = blockscale_fc1_output_size; auto* blockscale_gemm_runner = getDeepSeekBlockScaleGemmRunner(); TLLM_CHECK(blockscale_gemm_runner != nullptr); @@ -3345,20 +3341,26 @@ CutlassMoeFCRunner:: // depends only on (num_rows, top_k, num_experts) so it is shape_k-independent and shared by FC1/FC2. if (blockscale_gemm_runner->isActivationPrequantized()) { + // FC2's input (fc1_result_) is aliased onto fc2_result_ in the inputs buffer, so FC2 writes + // its output into the outputs buffer (glu_inter_result_). + overlapped_gemm1_gemm2_outputs_size = std::max(blockscale_fc1_output_size, blockscale_fc2_output_size); // Fused: the runner needs no internal workspace (both operands pre-quantized). The fused quant // instead packs fp8 activations + the padded 1x128 scales into the overlapped inputs buffer // (fc1_result_ for FC2, permuted_data_ for FC1); size it for both (hidden_size > inter_size makes - // FC1 dominant). The scale leading dim (~num_experts*32) can dwarf the token count, so it is not - // covered by the bf16-activation size. + // FC1 dominant). The scale leading dim (~num_experts*32) can dwarf the token count. The buffer + // holds these fp8 regions only, never the bf16 activations. blockscale_gemm_runner->getWorkspaceSize( num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); int64_t const scale_leading_dim = blockscale_gemm_runner->getActScaleLeadingDim(); - overlapped_gemm1_gemm2_inputs_size = std::max({overlapped_gemm1_gemm2_inputs_size, - fp8BlockScaleRegionBytes(num_moe_inputs, inter_size, scale_leading_dim), - fp8BlockScaleRegionBytes(num_moe_inputs, hidden_size, scale_leading_dim)}); + overlapped_gemm1_gemm2_inputs_size + = std::max(fp8BlockScaleRegionBytes(num_moe_inputs, inter_size, scale_leading_dim), + fp8BlockScaleRegionBytes(num_moe_inputs, hidden_size, scale_leading_dim)); } else { + // FC2 writes fc2_result_, which is aliased into the inputs buffer. + overlapped_gemm1_gemm2_inputs_size + = std::max(overlapped_gemm1_gemm2_inputs_size, blockscale_fc2_output_size); // Unfused: the runner quantizes A internally into deepseek_fc_workspace. auto deepseek_fc1_workspace_size = blockscale_gemm_runner->getWorkspaceSize( num_rows, factor * inter_size, hidden_size, experts_per_token, num_experts_per_node);