From 0250211d60aa9ccf56804b0888b24ffbd7a3adad Mon Sep 17 00:00:00 2001 From: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:37:10 +0000 Subject: [PATCH 1/3] Fuse pre-FC2 FP8 1x128 activation quant into doActivation kernel Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> --- .../fp8_blockscale_gemm.cu | 5 +- .../fp8_blockscale_gemm/fp8_blockscale_gemm.h | 12 +- .../cutlass_kernels/include/moe_kernels.h | 8 +- .../cutlass_kernels/moe_gemm/moe_kernels.cu | 206 ++++++++++++++++-- .../_torch/modules/moe/test_fc2_fusion.py | 163 ++++++++++++++ 5 files changed, 370 insertions(+), 24 deletions(-) create mode 100644 tests/unittest/_torch/modules/moe/test_fc2_fusion.py diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu index cb59a97b6ec8..3eb7f2887987 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -94,6 +94,9 @@ void CutlassFp8BlockScaleGemmRunner::moeGemm(void* void const* mat_b, int64_t const* problem_m_offsets, size_t num_problems, size_t expected_m, size_t shape_n, size_t shape_k, cudaStream_t stream, float const* scales_a, float const* scales_b) { + // ElementA == fp8_e4m3 selects the prequantized-A path: the caller (e.g. the fused FC1 activation + // epilogue) already produced fp8 A + per-token 1x128 scales, so the internal scale_1x128 quant is + // compiled out and mat_a / scales_a are consumed directly. constexpr bool internal_quantize_a = !std::is_same_v; constexpr bool internal_quantize_b = !std::is_same_v; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h index b178c1a1b806..abb3feb8468e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -55,6 +55,11 @@ class CutlassFp8BlockScaleGemmRunnerInterface float const* scales_b = nullptr) = 0; + // Leading dim (stride between K-blocks) of the per-token 1x128 activation scale buffer that moeGemm's + // internal quantization writes / that a prequantized-fp8-A runner () reads as its external + // scales_a. Valid after getWorkspaceSize. + virtual int64_t getActScaleLeadingDim() const = 0; + virtual void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a, __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) @@ -117,6 +122,11 @@ class CutlassFp8BlockScaleGemmRunner : public CutlassFp8BlockScaleGemmRunnerInte size_t num_problems, size_t shape_n, size_t shape_k, cudaStream_t stream, float const* scales_a = nullptr, float const* scales_b = nullptr) override; + int64_t getActScaleLeadingDim() const override + { + return max_shape_m_32_align_padded_; + } + void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a, __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) override; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index ab7ed876257d..905ac51af8e0 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -978,7 +978,13 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface int const num_experts_per_node = 0); MoeGemmRunner moe_gemm_runner_; + // Block-scale runner for FC1 + unfused FC2 (A is bf16 -> quantized internally). std::unique_ptr blockscale_gemm_runner_; + // Second block-scale runner for the fused FC2 path: A arrives pre-quantized to fp8 by the FC1 + // SwiGLU epilogue (doActivation WriteFp8BlockScale), so consumes it directly with no + // internal scale_1x128 quant. Needs no workspace; getWorkspaceSize() is still called on it to set the + // 1x128 scale leading dim it reads. + std::unique_ptr blockscale_gemm_runner_fp8a_; std::optional gemm1_config_; std::optional gemm2_config_; 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 8bed9c16b58e..71fda410bc6c 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 @@ -2070,14 +2070,16 @@ void doGatedActivation(ActivationOutputType* output, GemmOutputType const* gemm_ // ============================== Activation ================================= template + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType BlockScalingType, int kProcessRows, bool DynamicFc2 = false, + bool WriteFp8BlockScale = false> __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKernel(T* output, GemmOutputType const* gemm_result, float const* fp8_quant, ScaleBiasType const* bias_ptr, bool bias_is_broadcast, int64_t const* expert_first_token_offset, int num_experts_per_node, int64_t inter_size, float const* fc2_act_global_scale, bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, ActivationParams activation_params, GemmOutputType const* prequant_scale, float* dynamic_fc2_amax = nullptr, - GemmOutputType* bf16_intermediate_output = nullptr) + GemmOutputType* bf16_intermediate_output = nullptr, __nv_fp8_e4m3* fp8_block_output = nullptr, + float* fp8_block_scales = nullptr, int64_t fp8_block_scale_ld = 0) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v @@ -2140,8 +2142,9 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern // Grid stride loop for activation processing for (int64_t row_offset = blockIdx.x * rows_per_cta; row_offset < num_valid_tokens; row_offset += grid_stride) { - bool const per_expert_aux_input = bias_ptr || IsNVFP4 || IsMXFP8 || use_per_expert_act_scale - || activation_params.swiglu_alpha || activation_params.swiglu_beta || activation_params.swiglu_limit; + bool const per_expert_aux_input = bias_ptr || IsNVFP4 || IsMXFP8 || WriteFp8BlockScale + || use_per_expert_act_scale || activation_params.swiglu_alpha || activation_params.swiglu_beta + || activation_params.swiglu_limit; int32_t expert = 0; if (per_expert_aux_input) { @@ -2183,7 +2186,7 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern // Some globals for FP4 [[maybe_unused]] float global_scale_val = fc2_act_global_scale ? fc2_act_global_scale[act_scale_idx] : 1.0f; [[maybe_unused]] int64_t num_tokens_before_expert - = (IsNVFP4 || IsMXFP8) ? expert_first_token_offset[expert] : 0; + = (IsNVFP4 || IsMXFP8 || WriteFp8BlockScale) ? expert_first_token_offset[expert] : 0; size_t bias_offset = 0; if (bias_ptr) @@ -2291,6 +2294,52 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern /* input_sf */ nullptr); // Pass nullptr input_sf so we write 0 } } + else if constexpr (WriteFp8BlockScale) + { + // Fuse the pre-FC2 1x128 activation quantization into the SwiGLU epilogue, replacing the + // standalone grouped scale_1x128_kernel. Mirror that kernel bit-for-bit: it quantizes the + // *bf16* activation, reduces amax in bf16, and stores the plain-float 1/scale (no UE8M0). + // With ACTIVATION_ELEM_PER_THREAD == 8 (bf16), a 1x128 block spans 16 contiguous lanes. + ComputeElem bf16_round; + float amax = 0.f; +#pragma unroll + for (int k = 0; k < ACTIVATION_ELEM_PER_THREAD; ++k) + { + // Round-trip through GemmOutputType (bf16) to match the standalone path, which reads the + // bf16 activation buffer that doActivation would otherwise have written. + bf16_round[k] = static_cast(static_cast(post_act_val[k])); + amax = fmaxf(amax, fabsf(bf16_round[k])); + } + // Reduce amax across the 16 lanes covering this token's 128-channel block. The lanes of a block + // are 16-aligned within the warp (blockDim == (1, threads, 1)); mask only that half-warp so the + // sync is valid even when the sibling half is inactive (guaranteed disjoint by inter_size % 128 == 0). + unsigned const lane = static_cast(threadIdx.y) & 31u; + unsigned const group_mask = 0xFFFFu << (lane & 16u); +#pragma unroll + for (int m = 1; m < 16; m <<= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(group_mask, amax, m, 32)); + } + // Truncate amax to GemmOutputType (bf16) to match the standalone scale_1x128 path, then floor. + float const amax_bf = fmaxf(static_cast(static_cast(amax)), 1e-10f); + float const scale = 448.f / amax_bf; + // Per-element float -> fp8_e4m3 conversion, matching the standalone kernel's `OutputType(value)`. + __nv_fp8_e4m3* out_row = fp8_block_output + output_offset + elem_index * ACTIVATION_ELEM_PER_THREAD; +#pragma unroll + for (int k = 0; k < ACTIVATION_ELEM_PER_THREAD; ++k) + { + out_row[k] = static_cast<__nv_fp8_e4m3>(bf16_round[k] * scale); + } + if ((col_offset & 15) == 0) + { + // Transposed / per-expert padded scale layout consumed by the grouped GEMM + // (deep_gemm::compute_padded_offset with alignment 32). + int64_t const kb = col_offset >> 4; + int64_t const pad = (num_tokens_before_expert + static_cast(expert) * 31) / 32 * 32 + - num_tokens_before_expert; + fp8_block_scales[kb * fp8_block_scale_ld + token + pad] = 1.f / scale; + } + } else { // Use storeVec to force STG.128 vectorized store @@ -2502,7 +2551,8 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 bool bias_is_broadcast, int64_t const* expert_first_token_offset, int num_experts_per_node, int64_t inter_size, int64_t expanded_num_tokens, ActivationParams activation_type, QuantParams const& quant_params, bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, cudaStream_t stream, - GemmOutputType const* prequant_scale = nullptr) + GemmOutputType const* prequant_scale = nullptr, __nv_fp8_e4m3* fp8_block_output = nullptr, + float* fp8_block_scales = nullptr, int64_t fp8_block_scale_ld = 0, bool write_fp8_block_scale = false) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v; @@ -2528,11 +2578,25 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 { // IMPORTANT: Keep the order of the activation functions in the same order as the ActivationType enum in // common.h - auto fn - = [&](auto block_scaling_type) -> void (*)(T*, GemmOutputType const*, float const*, - ScaleBiasType const*, bool, int64_t const*, int, int64_t, - float const*, bool, TmaWarpSpecializedGroupedGemmInput::ElementSF*, - ActivationParams, GemmOutputType const*, float*, GemmOutputType*) + using KernelFnPtr = void (*)(T*, GemmOutputType const*, float const*, ScaleBiasType const*, bool, + int64_t const*, int, int64_t, float const*, bool, TmaWarpSpecializedGroupedGemmInput::ElementSF*, + ActivationParams, GemmOutputType const*, float*, GemmOutputType*, __nv_fp8_e4m3*, float*, int64_t); + // Fuse the pre-FC2 1x128 activation quantization (WriteFp8BlockScale=true). Only Swiglu + a 16-bit T + // (bf16/half) + NONE block-scaling is used on the DeepSeek FP8 block-scale MoE path. Guard the + // instantiation so it is never emitted for fp4/fp8 T (where NONE block-scaling would compile the fp8 + // branch with an illegal wide vector width). + if constexpr (!IsNVFP4 && !IsMXFP8) + { + if (write_fp8_block_scale) + { + return static_cast( + &doActivationKernel, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, num_rows_per_cta_v, + /*DynamicFc2=*/false, /*WriteFp8BlockScale=*/true>); + } + } + auto fn = [&](auto block_scaling_type) -> KernelFnPtr { switch (activation_type.activation_type) { @@ -2631,7 +2695,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 cudaLaunchKernelEx(&config, fn, output, gemm_result, fp8_quant, bias, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, quant_params.fp4.fc2.act_global_scale, use_per_expert_act_scale, fc2_act_sf_flat, activation_type, prequant_scale, (float*) nullptr, - (GemmOutputType*) nullptr); + (GemmOutputType*) nullptr, fp8_block_output, fp8_block_scales, fp8_block_scale_ld); }; // end lambda doActivationKernelLauncher // 256 threads per block * 256 blocks / 1 rows per block can be handled by 1-2 waves depending on SM arch @@ -2729,7 +2793,8 @@ void doActivationDynamic(T* output, GemmOutputType const* gemm_result, float con fn<<>>(output, gemm_result, fp8_quant, bias, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, quant_params.fp4.fc2.act_global_scale, use_per_expert_act_scale, fc2_act_sf_flat, activation_type, - (GemmOutputType const*) nullptr, dynamic_amax, bf16_intermediate); + (GemmOutputType const*) nullptr, dynamic_amax, bf16_intermediate, /*fp8_block_output=*/nullptr, + /*fp8_block_scales=*/nullptr, /*fp8_block_scale_ld=*/0); sync_check_cuda_error(stream); } @@ -2940,6 +3005,8 @@ template ::CutlassMoeFCRunner() : blockscale_gemm_runner_{std::make_unique< kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>>()} + , blockscale_gemm_runner_fp8a_{std::make_unique< + kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, __nv_bfloat16>>()} { } @@ -3055,7 +3122,10 @@ CutlassMoeFCRunner:: 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); - overlapped_gemm1_gemm2_outputs_size = blockscale_fc1_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); auto* blockscale_gemm_runner = getDeepSeekBlockScaleGemmRunner(); TLLM_CHECK(blockscale_gemm_runner != nullptr); @@ -3064,6 +3134,25 @@ CutlassMoeFCRunner:: auto deepseek_fc2_workspace_size = blockscale_gemm_runner->getWorkspaceSize( num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); deepseek_fc_workspace_size = std::max(deepseek_fc1_workspace_size, deepseek_fc2_workspace_size); + + // The fused-FC2 runner () needs no workspace (both operands pre-quantized), but we + // must call getWorkspaceSize on it with the FC2 shape so its 1x128 scale leading dim matches the + // one BlockScaleFC1 wrote scales at (getActScaleLeadingDim() below, on the bf16 runner). Same + // (num_rows, top_k, num_experts) -> identical compute_padded_offset, so the strides line up. + blockscale_gemm_runner_fp8a_->getWorkspaceSize( + num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); + + // The fused pre-FC2 activation quant (BlockScaleFC1/FC2) writes the fp8 FC2 input into the low part of + // the overlapped inputs buffer and the per-token 1x128 scales just above it. Size that buffer for + // fp8 + the *padded* scale layout: the scale leading dim is getActScaleLeadingDim() (== the per-expert + // padded row count, ~num_experts*32), which for many experts and few tokens dwarfs the token count, so + // the bf16-activation size (fc1_result_size) underestimates it. getActScaleLeadingDim() is valid here + // because getWorkspaceSize() above set the padded member. + size_t const fused_fc2_fp8_bytes = ((num_moe_inputs * inter_size * sizeof(__nv_fp8_e4m3)) + 15) / 16 * 16; + size_t const fused_fc2_scale_bytes = static_cast(blockscale_gemm_runner->getActScaleLeadingDim()) + * tensorrt_llm::common::ceilDiv(inter_size, static_cast(128)) * sizeof(float); + overlapped_gemm1_gemm2_inputs_size + = std::max(overlapped_gemm1_gemm2_inputs_size, fused_fc2_fp8_bytes + fused_fc2_scale_bytes); } size_t map_offset = 0; @@ -3237,6 +3326,8 @@ void CutlassMoeFCRunnerconfigureWorkspace(getWsPtr(char{}, "deepseek_fc_workspace")); + // The fused-FC2 runner reads pre-quantized A + external scales, so it needs no internal workspace. + blockscale_gemm_runner_fp8a_->configureWorkspace(nullptr); } if (use_awq) @@ -3256,6 +3347,20 @@ CutlassMoeFCRunner: return blockscale_gemm_runner_.get(); } +// TEMP/DEBUG: set TLLM_MOE_DISABLE_FUSED_FC2_QUANT=1 to fall back to the standalone pre-FC2 scale_1x128 quant +// (unfused path) instead of fusing it into the FC1 activation epilogue. Read live so a single process can A/B both. +// The fused path is currently wired only for the Hopper (SM90) bf16xfp8 grouped GEMM; other arches (e.g. SM120) +// keep the standalone quant. +static inline bool isFusedFc2QuantDisabled() +{ + if (tensorrt_llm::common::getSMVersion() != 90) + { + return true; + } + char const* env = std::getenv("TLLM_MOE_DISABLE_FUSED_FC2_QUANT"); + return env != nullptr && env[0] == '1'; +} + template void CutlassMoeFCRunner::BlockScaleFC1( DeepSeekBlockScaleGemmRunner& gemm_runner, T const* const input, T* const output, void* const gemm_output, @@ -3277,9 +3382,32 @@ void CutlassMoeFCRunner(output, static_cast(gemm_output), - fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, - inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, nullptr, stream); + if (isFusedFc2QuantDisabled()) + { + // Unfused path: write the bf16 activation; FC2's moeGemm will quantize it via the standalone scale_1x128. + doActivation(output, static_cast(gemm_output), + fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, + inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, nullptr, + stream); + } + else + { + // Fuse the pre-FC2 1x128 activation quantization into the SwiGLU epilogue: write fp8 into the low half of the + // FC2-input buffer (`output`, sized for bf16 == 2x) and the per-token 1x128 scales into its freed upper half. + // This removes the standalone grouped scale_1x128_kernel that FC2's moeGemm would otherwise launch. + TLLM_CHECK_WITH_INFO( + inter_size % 128 == 0, "Fused FC2 activation quant requires inter_size to be a multiple of 128."); + auto* fp8_block_output = reinterpret_cast<__nv_fp8_e4m3*>(output); + size_t const fp8_bytes = static_cast(expanded_num_rows) * inter_size * sizeof(__nv_fp8_e4m3); + size_t const scale_off = ((fp8_bytes + 15) / 16) * 16; // 16B-align the float scale region + auto* fp8_block_scales = reinterpret_cast(reinterpret_cast(output) + scale_off); + int64_t const fp8_block_scale_ld = gemm_runner.getActScaleLeadingDim(); + doActivation(output, static_cast(gemm_output), + fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, + inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, nullptr, stream, + /*prequant_scale=*/nullptr, fp8_block_output, fp8_block_scales, fp8_block_scale_ld, + /*write_fp8_block_scale=*/true); + } sync_check_cuda_error(stream); } @@ -3299,9 +3427,34 @@ void CutlassMoeFCRunner runner (consumes pre-quantized A, no internal quant), + // and `gemm_output` points at glu_inter_result_ (the outputs buffer), NOT the aliased + // fc2_result_, so the GEMM output does not clobber the fp8 A / 1x128 scales that `input` + // (== fc1_result_) still holds. + // - unfused: `gemm_runner` is the runner (quantizes A internally into its own + // workspace), and `gemm_output` is fc2_result_. // NOTE: we assume gemm_runner.configureWorkspace has already been called. - gemm_runner.moeGemm(gemm_output, input, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, - expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, quant_params.fp8_block_scaling.fc2_scales_ptrs); + if (isFusedFc2QuantDisabled()) + { + // Unfused path: `input` (fc1_result_) holds the bf16 activation; moeGemm quantizes it internally. + gemm_runner.moeGemm(gemm_output, input, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, + expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, + quant_params.fp8_block_scaling.fc2_scales_ptrs); + } + else + { + // Fused path: `input` (fc1_result_) was populated by the FC1 activation epilogue: fp8 activations in the low + // half, per-token 1x128 scales in the 16B-aligned upper half. The runner consumes them + // directly (internal_quantize_a is compile-time false), so no a_is_prequantized flag is needed. + auto const* fp8_a = reinterpret_cast<__nv_fp8_e4m3 const*>(input); + size_t const fp8_bytes = static_cast(expanded_num_rows) * inter_size * sizeof(__nv_fp8_e4m3); + size_t const scale_off = ((fp8_bytes + 15) / 16) * 16; + auto const* scales_a = reinterpret_cast(reinterpret_cast(input) + scale_off); + gemm_runner.moeGemm(gemm_output, fp8_a, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, + expected_tokens_per_expert, shape_n, shape_k, stream, scales_a, + quant_params.fp8_block_scaling.fc2_scales_ptrs); + } sync_check_cuda_error(stream); @@ -3655,7 +3808,8 @@ void CutlassMoeFCRunner(gemm_output), nullptr, static_cast(fc2_lora), false, expert_first_token_offset, num_experts_per_node, hidden_size, expanded_num_rows, ActivationParams(ActivationType::Identity), {}, false, nullptr, stream, - /*prequant_scale=*/nullptr); + /*prequant_scale=*/nullptr, /*fp8_block_output=*/nullptr, /*fp8_block_scales=*/nullptr, + /*fp8_block_scale_ld=*/0, /*write_fp8_block_scale=*/false); sync_check_cuda_error(stream); } @@ -4374,7 +4528,17 @@ void CutlassMoeFCRunner runner, which consumes the fp8 A + 1x128 scales + // the FC1 SwiGLU epilogue already wrote in place into fc1_result_. Since the deepseek workspace + // aliases fc2_result_ onto fc1_result_, the GEMM output must go to glu_inter_result_ (already + // consumed by that epilogue) so input and output are disjoint. + // - unfused: use the runner (quantizes A into its own workspace first) and write to + // fc2_result_ as before. + bool const fc2_use_fused = (blockscale_gemm_runner != nullptr) && !isFusedFc2QuantDisabled(); + auto* fc2_blockscale_runner = fc2_use_fused ? blockscale_gemm_runner_fp8a_.get() : blockscale_gemm_runner; + void* const fc2_gemm_output = fc2_use_fused ? glu_inter_result_ : fc2_result_; + Self::gemm2(moe_gemm_runner_, fc2_blockscale_runner, gemm2_input, fc2_gemm_output, final_output, expert_first_token_offset_, gemm2_tma_ws_input, fc2_expert_weights, fc2_expert_biases, fc2_int_scales, fc2_fp8_dequant, fc2_fp4_act_scale_, quant_params, token_topk_unpermuted_scales, permuted_token_final_scales_, unpermuted_row_to_permuted_row, permuted_row_to_unpermuted_row_, diff --git a/tests/unittest/_torch/modules/moe/test_fc2_fusion.py b/tests/unittest/_torch/modules/moe/test_fc2_fusion.py new file mode 100644 index 000000000000..07ced968aa66 --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_fc2_fusion.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit test for the CUTLASS DeepSeek FP8 block-scale MoE FC2 activation-quant fusion. + +The block-scale MoE runs FC1 -> SwiGLU activation -> FC2 with bf16 (non-fp8) inputs to FC1. +FC2 requires the SwiGLU activation to be quantized to fp8 with per-token 1x128 scales. Two code +paths produce that fp8 input: + + * fused (default): the 1x128 quantization is folded into the FC1 SwiGLU epilogue + (doActivationKernel WriteFp8BlockScale), so FC2's grouped GEMM consumes + pre-quantized fp8 activations. + * unfused (TLLM_MOE_DISABLE_FUSED_FC2_QUANT=1): FC1 writes bf16, and FC2's grouped GEMM + quantizes it with the standalone scale_1x128 kernel. + +The fused path is designed to be numerically identical to the unfused one, so this test asserts +they match, and that both are close to a bf16 reference. Hopper (SM90) only. +""" + +import pytest +import torch +from _torch.modules.moe.moe_test_utils import MoeBackendType +from _torch.modules.moe.quantize_utils import get_test_quant_params +from _torch.modules.moe.test_moe_backend import create_test_backend, run_backend_moe + +from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod +from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm._utils import get_sm_version, mpi_rank +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo + +DISABLE_FUSION_ENV = "TLLM_MOE_DISABLE_FUSED_FC2_QUANT" + + +# Qwen3.5-35B-A3B MoE shapes (text_config): hidden_size=2048, moe_intermediate_size=512. +QWEN35_35B_A3B_HIDDEN_SIZE = 2048 +QWEN35_35B_A3B_INTERMEDIATE_SIZE = 512 + + +@pytest.mark.skipif( + get_sm_version() != 90, + reason="CUTLASS FP8 block-scale FC2 fusion is wired for Hopper (SM90) only.", +) +@pytest.mark.parametrize("seq_len", [1, 8, 64], ids=lambda s: f"seq{s}") +def test_cutlass_fp8_block_scale_fc2_fusion(seq_len, monkeypatch): + """Fused FC2 activation-quant must match the unfused standalone scale_1x128 path. + + Uses the Qwen3.5-35B-A3B MoE hidden/intermediate sizes (2048 / 512). + """ + dtype = torch.bfloat16 + hidden_size = QWEN35_35B_A3B_HIDDEN_SIZE + intermediate_size = QWEN35_35B_A3B_INTERMEDIATE_SIZE + num_experts = 256 + top_k = 8 + + mapping = Mapping() + mapping.rank = mpi_rank() + + with torch.device(f"cuda:{mapping.rank}"): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + AutoTuner.get().setup_distributed_state(mapping) + + routing_method = RenormalizeMoeRoutingMethod(top_k=top_k) + x = torch.randn((seq_len, hidden_size), dtype=dtype, device="cuda") + router_logits = torch.randn((seq_len, num_experts), dtype=dtype, device="cuda") + + # CUTLASS FP8_BLOCK_SCALES uses plain-float 128x128 weight scales (FP8BlockScalesQuantizeUtil). + quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( + QuantAlgo.FP8_BLOCK_SCALES, x, MoeBackendType.CUTLASS + ) + quantize_util = quantize_util_cls( + num_experts=num_experts, + dtype=dtype, + intermediate_size=intermediate_size, + hidden_size=hidden_size, + quant_config=quant_config, + activation_type=ActivationType.Swiglu, + ) + + backend = create_test_backend( + backend_type=MoeBackendType.CUTLASS, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + quant_config=quant_config, + mapping=mapping, + activation_type=ActivationType.Swiglu, + ) + weights = quantize_util.create_weights(**quant_kwargs) + backend.load_weights([weights]) + backend.post_load_weights() + backend.cuda() + + ref_fused_moe = quantize_util.create_ref_module(routing_method) + ref_fused_moe.load_weights([weights]) + ref_fused_moe.cuda() + + def run_moe(): + token_selected_experts, token_final_scales = routing_method.apply(router_logits) + x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False) + return run_backend_moe( + backend, + MoeBackendType.CUTLASS, + x_quantized, + x_sf, + token_selected_experts, + token_final_scales, + dtype, + router_logits, + ) + + # Autotune once (with fusion on) so both A/B runs replay the same FC2 GEMM tactics. + monkeypatch.delenv(DISABLE_FUSION_ENV, raising=False) + AutoTuner.get().clear_cache() + with torch.inference_mode(), autotune(): + run_moe() + + with torch.inference_mode(): + out_fused = run_moe() + + # Unfused path: FC1 writes bf16, FC2 quantizes with the standalone scale_1x128 kernel. + monkeypatch.setenv(DISABLE_FUSION_ENV, "1") + with torch.inference_mode(): + out_unfused = run_moe() + monkeypatch.delenv(DISABLE_FUSION_ENV, raising=False) + + with torch.inference_mode(): + ref_output = ref_fused_moe.forward(x, router_logits) + + # Primary check: the fused epilogue reproduces the standalone quant bit-for-bit (same bf16 + # activation, same amax, same fp8 rounding, same GEMM), so outputs should be ~identical. + torch.testing.assert_close( + out_fused, + out_unfused, + rtol=5e-3, + atol=5e-3, + msg="Fused FC2 quant diverges from the unfused standalone path", + ) + + # Sanity: the fused output is a valid MoE result (finite and well-correlated with the bf16 + # reference), i.e. the fused==unfused match above is not two identical garbage tensors. + assert torch.isfinite(out_fused).all(), "Fused output contains non-finite values" + cos = torch.nn.functional.cosine_similarity( + out_fused.float().flatten(), ref_output.float().flatten(), dim=0 + ) + assert cos > 0.9, ( + f"Fused output poorly correlated with bf16 reference (cos={cos.item():.4f})" + ) From 4806536c2d4effb00b468767ab6aa0f96c52db47 Mon Sep 17 00:00:00 2001 From: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:23:29 +0000 Subject: [PATCH 2/3] Fuse pre-FC1 FP8 1x128 activation quant into the row-expansion kernel Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> --- .../fp8_blockscale_gemm.cu | 5 +- .../fp8_blockscale_gemm/fp8_blockscale_gemm.h | 4 +- .../cutlass_kernels/include/moe_kernels.h | 6 +- .../cutlass_kernels/moe_gemm/moe_kernels.cu | 350 ++++++++++++------ ...sion.py => test_fp8_block_scale_fusion.py} | 37 +- 5 files changed, 263 insertions(+), 139 deletions(-) rename tests/unittest/_torch/modules/moe/{test_fc2_fusion.py => test_fp8_block_scale_fusion.py} (77%) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu index 3eb7f2887987..cb59a97b6ec8 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -94,9 +94,6 @@ void CutlassFp8BlockScaleGemmRunner::moeGemm(void* void const* mat_b, int64_t const* problem_m_offsets, size_t num_problems, size_t expected_m, size_t shape_n, size_t shape_k, cudaStream_t stream, float const* scales_a, float const* scales_b) { - // ElementA == fp8_e4m3 selects the prequantized-A path: the caller (e.g. the fused FC1 activation - // epilogue) already produced fp8 A + per-token 1x128 scales, so the internal scale_1x128 quant is - // compiled out and mat_a / scales_a are consumed directly. constexpr bool internal_quantize_a = !std::is_same_v; constexpr bool internal_quantize_b = !std::is_same_v; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h index abb3feb8468e..0530f7082475 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h @@ -55,9 +55,7 @@ class CutlassFp8BlockScaleGemmRunnerInterface float const* scales_b = nullptr) = 0; - // Leading dim (stride between K-blocks) of the per-token 1x128 activation scale buffer that moeGemm's - // internal quantization writes / that a prequantized-fp8-A runner () reads as its external - // scales_a. Valid after getWorkspaceSize. + // Stride between K-blocks of per-token 1x128 activation scale buffer virtual int64_t getActScaleLeadingDim() const = 0; virtual void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 905ac51af8e0..ab5e9affe761 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -978,12 +978,8 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface int const num_experts_per_node = 0); MoeGemmRunner moe_gemm_runner_; - // Block-scale runner for FC1 + unfused FC2 (A is bf16 -> quantized internally). std::unique_ptr blockscale_gemm_runner_; - // Second block-scale runner for the fused FC2 path: A arrives pre-quantized to fp8 by the FC1 - // SwiGLU epilogue (doActivation WriteFp8BlockScale), so consumes it directly with no - // internal scale_1x128 quant. Needs no workspace; getWorkspaceSize() is still called on it to set the - // 1x128 scale leading dim it reads. + // Second block-scale runner with pre-quantized fp8 activations std::unique_ptr blockscale_gemm_runner_fp8a_; std::optional gemm1_config_; 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 71fda410bc6c..823346f6eb64 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 @@ -1703,6 +1703,147 @@ INSTANTIATE_EXPAND_INPUT_ROWS(half, half); INSTANTIATE_EXPAND_INPUT_ROWS(__nv_bfloat16, __nv_bfloat16); #endif +// ---- DeepSeek FP8 block-scale MoE: shared fused activation-quant helpers (SM90, bf16) ---- +// +// FC1 and FC2 each consume fp8 activations + per-token 1x128 scales. Two fusions produce that input +// in-place instead of a standalone scale_1x128 kernel: pre-FC1 in the row expansion, pre-FC2 in the +// SwiGLU epilogue. Both pack fp8 in the low part of a buffer and the scales just above, and both must +// match the standalone quant bit-for-bit. These helpers hold the shared layout and quant so the two +// call sites stay identical by construction. + +// Byte offset from the buffer base to the 1x128 scale region that follows `rows` x `cols` fp8 +// activations (16B-aligned so the float scales are aligned). +static inline size_t fp8BlockScaleByteOffset(int64_t rows, int64_t cols) +{ + return ((static_cast(rows) * cols * sizeof(__nv_fp8_e4m3)) + 15) / 16 * 16; +} + +// Total bytes of the packed fp8-activations + 1x128-scales region (scale leading dim == grouped-GEMM +// padded M, so shape_k-independent). +static inline size_t fp8BlockScaleRegionBytes(int64_t rows, int64_t cols, int64_t scale_ld) +{ + return fp8BlockScaleByteOffset(rows, cols) + + static_cast(scale_ld) * tensorrt_llm::common::ceilDiv(cols, static_cast(128)) + * sizeof(float); +} + +// Row offset of an expert's scales in the transposed / per-expert padded layout the grouped GEMM reads +// (deep_gemm::compute_padded_offset, alignment 32). +__device__ inline int64_t fp8BlockScaleExpertPad(int64_t num_tokens_before_expert, int64_t expert) +{ + return (num_tokens_before_expert + expert * 31) / 32 * 32 - num_tokens_before_expert; +} + +// Quantize one thread's 8 already-bf16-rounded channels to fp8 with a per-token 1x128 scale, bit-for-bit +// with the standalone scale_1x128 (amax reduced and truncated in bf16, plain-float 1/scale). The 8 +// channels span 16 warp lanes per 128-block; only bit 4 of `lane` is used, to mask that half-warp so a +// disjoint inactive sibling is safe (guaranteed by dim % 128 == 0). The block leader writes 1/scale. +template +__device__ inline void fp8BlockScaleQuantize(float const (&vals)[N], __nv_fp8_e4m3* out, float* block_scales, + int64_t kb, int64_t scale_ld, int64_t scale_row, bool write_scale, unsigned lane) +{ + static_assert(N == 8, "1x128 block quant assumes 8 activation channels per thread (16 lanes/block)."); + float amax = 0.f; +#pragma unroll + for (int e = 0; e < N; ++e) + amax = fmaxf(amax, fabsf(vals[e])); + unsigned const group_mask = 0xFFFFu << (lane & 16u); +#pragma unroll + for (int m = 1; m < 16; m <<= 1) + amax = fmaxf(amax, __shfl_xor_sync(group_mask, amax, m, 32)); + float const scale = 448.f / fmaxf(static_cast(static_cast<__nv_bfloat16>(amax)), 1e-10f); +#pragma unroll + for (int e = 0; e < N; ++e) + out[e] = static_cast<__nv_fp8_e4m3>(vals[e] * scale); + if (write_scale) + block_scales[kb * scale_ld + scale_row] = 1.f / scale; +} + +// Fuse the pre-FC1 1x128 activation quant into the row expansion: permute the bf16 input rows like +// expandInputRowsKernel, but write fp8 + scales the FC1 GEMM consumes directly. +template +__global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* unpermuted_input, + __nv_fp8_e4m3* permuted_output, float* fp8_block_scales, int64_t const fp8_block_scale_ld, + float const* unpermuted_scales, float* permuted_scales, int const* permuted_row_to_unpermuted_row, + int64_t const num_tokens, int64_t const hidden_size, int64_t const k, int64_t const* expert_first_token_offset, + int64_t const num_experts_per_node) +{ + // 16-bit input => 8 channels per thread => a 1x128 block spans 16 contiguous lanes. + constexpr int64_t ELEM_PER_THREAD = 128 / sizeof_bits::value; + using DataElem = cutlass::Array; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; + int64_t const num_elems_in_col = hidden_size / ELEM_PER_THREAD; + + for (int64_t permuted_row = blockIdx.x; permuted_row < num_valid_tokens; permuted_row += gridDim.x) + { + int64_t const unpermuted_row = permuted_row_to_unpermuted_row[permuted_row]; + int64_t const source_k_rank = unpermuted_row / num_tokens; + int64_t const source_row = unpermuted_row % num_tokens; + + auto const* source_row_ptr = reinterpret_cast(unpermuted_input + source_row * hidden_size); + auto* dest_row_ptr = permuted_output + permuted_row * hidden_size; + + int64_t const expert + = findTotalEltsLessThanTarget(expert_first_token_offset, num_experts_per_node, permuted_row + 1) - 1; + int64_t const scale_row = permuted_row + fp8BlockScaleExpertPad(expert_first_token_offset[expert], expert); + + for (int64_t elem_index = threadIdx.x; elem_index < num_elems_in_col; elem_index += EXPAND_THREADS_PER_BLOCK) + { + DataElem const in_vec = source_row_ptr[elem_index]; + float vals[ELEM_PER_THREAD]; +#pragma unroll + for (int e = 0; e < ELEM_PER_THREAD; ++e) + vals[e] = static_cast(in_vec[e]); // bf16 input is already exact + fp8BlockScaleQuantize(vals, dest_row_ptr + elem_index * ELEM_PER_THREAD, fp8_block_scales, elem_index >> 4, + fp8_block_scale_ld, scale_row, (elem_index & 15) == 0, threadIdx.x); + } + + if (permuted_scales && threadIdx.x == 0) + { + int64_t const source_k_idx = source_row * k + source_k_rank; + permuted_scales[permuted_row] = unpermuted_scales ? unpermuted_scales[source_k_idx] : 1.0f; + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +void expandInputRowsFp8BlockScaleKernelLauncher(InputActivationsType const* unpermuted_input, + __nv_fp8_e4m3* permuted_output, float* fp8_block_scales, int64_t const fp8_block_scale_ld, + float const* unpermuted_scales, float* permuted_scales, int const* permuted_row_to_unpermuted_row, + int64_t const num_rows, int64_t const hidden_size, int const k, int const num_experts_per_node, + int64_t* expert_first_token_offset, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(hidden_size % 128 == 0, "Fused FC1 activation quant requires hidden_size %% 128 == 0."); + auto* func = &expandInputRowsFp8BlockScaleKernel; + static int32_t const smCount = tensorrt_llm::common::getMultiProcessorCount(); + int32_t const maxBlocksPerSM = tensorrt_llm::common::getMaxActiveBlocksPerSM(func, EXPAND_THREADS_PER_BLOCK, 0); + int32_t const blocks = std::min(smCount * maxBlocksPerSM, static_cast(std::max(num_rows * k, 1))); + int32_t const threads = EXPAND_THREADS_PER_BLOCK; + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, func, unpermuted_input, permuted_output, fp8_block_scales, fp8_block_scale_ld, + unpermuted_scales, permuted_scales, permuted_row_to_unpermuted_row, num_rows, hidden_size, + static_cast(k), expert_first_token_offset, static_cast(num_experts_per_node)); +} + enum class ScaleMode : int { NO_SCALE = 0, @@ -2294,51 +2435,21 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern /* input_sf */ nullptr); // Pass nullptr input_sf so we write 0 } } - else if constexpr (WriteFp8BlockScale) + else if constexpr (WriteFp8BlockScale && ACTIVATION_ELEM_PER_THREAD == 8) { - // Fuse the pre-FC2 1x128 activation quantization into the SwiGLU epilogue, replacing the - // standalone grouped scale_1x128_kernel. Mirror that kernel bit-for-bit: it quantizes the - // *bf16* activation, reduces amax in bf16, and stores the plain-float 1/scale (no UE8M0). - // With ACTIVATION_ELEM_PER_THREAD == 8 (bf16), a 1x128 block spans 16 contiguous lanes. - ComputeElem bf16_round; - float amax = 0.f; -#pragma unroll - for (int k = 0; k < ACTIVATION_ELEM_PER_THREAD; ++k) - { - // Round-trip through GemmOutputType (bf16) to match the standalone path, which reads the - // bf16 activation buffer that doActivation would otherwise have written. - bf16_round[k] = static_cast(static_cast(post_act_val[k])); - amax = fmaxf(amax, fabsf(bf16_round[k])); - } - // Reduce amax across the 16 lanes covering this token's 128-channel block. The lanes of a block - // are 16-aligned within the warp (blockDim == (1, threads, 1)); mask only that half-warp so the - // sync is valid even when the sibling half is inactive (guaranteed disjoint by inter_size % 128 == 0). - unsigned const lane = static_cast(threadIdx.y) & 31u; - unsigned const group_mask = 0xFFFFu << (lane & 16u); -#pragma unroll - for (int m = 1; m < 16; m <<= 1) - { - amax = fmaxf(amax, __shfl_xor_sync(group_mask, amax, m, 32)); - } - // Truncate amax to GemmOutputType (bf16) to match the standalone scale_1x128 path, then floor. - float const amax_bf = fmaxf(static_cast(static_cast(amax)), 1e-10f); - float const scale = 448.f / amax_bf; - // Per-element float -> fp8_e4m3 conversion, matching the standalone kernel's `OutputType(value)`. - __nv_fp8_e4m3* out_row = fp8_block_output + output_offset + elem_index * ACTIVATION_ELEM_PER_THREAD; + // Fuse the pre-FC2 1x128 activation quant into the SwiGLU epilogue (see fp8BlockScaleQuantize). + // Only the bf16 activation (8 channels/thread == 16 lanes/block) is wired; other T never reach + // this at runtime, so their instantiations fall through to the plain store below. + // Round the activation through GemmOutputType (bf16) first, matching the standalone path that + // reads the bf16 buffer doActivation would otherwise have written. + float rounded[ACTIVATION_ELEM_PER_THREAD]; #pragma unroll for (int k = 0; k < ACTIVATION_ELEM_PER_THREAD; ++k) - { - out_row[k] = static_cast<__nv_fp8_e4m3>(bf16_round[k] * scale); - } - if ((col_offset & 15) == 0) - { - // Transposed / per-expert padded scale layout consumed by the grouped GEMM - // (deep_gemm::compute_padded_offset with alignment 32). - int64_t const kb = col_offset >> 4; - int64_t const pad = (num_tokens_before_expert + static_cast(expert) * 31) / 32 * 32 - - num_tokens_before_expert; - fp8_block_scales[kb * fp8_block_scale_ld + token + pad] = 1.f / scale; - } + rounded[k] = static_cast(static_cast(post_act_val[k])); + int64_t const scale_row = token + fp8BlockScaleExpertPad(num_tokens_before_expert, expert); + fp8BlockScaleQuantize(rounded, + fp8_block_output + output_offset + elem_index * ACTIVATION_ELEM_PER_THREAD, fp8_block_scales, + col_offset >> 4, fp8_block_scale_ld, scale_row, (col_offset & 15) == 0, threadIdx.y); } else { @@ -3135,24 +3246,20 @@ CutlassMoeFCRunner:: num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); deepseek_fc_workspace_size = std::max(deepseek_fc1_workspace_size, deepseek_fc2_workspace_size); - // The fused-FC2 runner () needs no workspace (both operands pre-quantized), but we - // must call getWorkspaceSize on it with the FC2 shape so its 1x128 scale leading dim matches the - // one BlockScaleFC1 wrote scales at (getActScaleLeadingDim() below, on the bf16 runner). Same - // (num_rows, top_k, num_experts) -> identical compute_padded_offset, so the strides line up. + // Prime the runner's 1x128 scale leading dim (getActScaleLeadingDim()) so it matches + // what the fused quant writes. It needs no workspace (both operands pre-quantized); the leading dim + // depends only on (num_rows, top_k, num_experts), so it is shape_k-independent and shared by FC1/FC2. blockscale_gemm_runner_fp8a_->getWorkspaceSize( num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); - // The fused pre-FC2 activation quant (BlockScaleFC1/FC2) writes the fp8 FC2 input into the low part of - // the overlapped inputs buffer and the per-token 1x128 scales just above it. Size that buffer for - // fp8 + the *padded* scale layout: the scale leading dim is getActScaleLeadingDim() (== the per-expert - // padded row count, ~num_experts*32), which for many experts and few tokens dwarfs the token count, so - // the bf16-activation size (fc1_result_size) underestimates it. getActScaleLeadingDim() is valid here - // because getWorkspaceSize() above set the padded member. - size_t const fused_fc2_fp8_bytes = ((num_moe_inputs * inter_size * sizeof(__nv_fp8_e4m3)) + 15) / 16 * 16; - size_t const fused_fc2_scale_bytes = static_cast(blockscale_gemm_runner->getActScaleLeadingDim()) - * tensorrt_llm::common::ceilDiv(inter_size, static_cast(128)) * sizeof(float); - overlapped_gemm1_gemm2_inputs_size - = std::max(overlapped_gemm1_gemm2_inputs_size, fused_fc2_fp8_bytes + fused_fc2_scale_bytes); + // The fused quant 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. + int64_t const fused_scale_ld = blockscale_gemm_runner->getActScaleLeadingDim(); + overlapped_gemm1_gemm2_inputs_size = std::max( + {overlapped_gemm1_gemm2_inputs_size, fp8BlockScaleRegionBytes(num_moe_inputs, inter_size, fused_scale_ld), + fp8BlockScaleRegionBytes(num_moe_inputs, hidden_size, fused_scale_ld)}); } size_t map_offset = 0; @@ -3347,18 +3454,18 @@ CutlassMoeFCRunner: return blockscale_gemm_runner_.get(); } -// TEMP/DEBUG: set TLLM_MOE_DISABLE_FUSED_FC2_QUANT=1 to fall back to the standalone pre-FC2 scale_1x128 quant -// (unfused path) instead of fusing it into the FC1 activation epilogue. Read live so a single process can A/B both. -// The fused path is currently wired only for the Hopper (SM90) bf16xfp8 grouped GEMM; other arches (e.g. SM120) -// keep the standalone quant. -static inline bool isFusedFc2QuantDisabled() +// Both FP8 block-scale activation-quant fusions (pre-FC1 row expansion and pre-FC2 SwiGLU epilogue) are +// enabled automatically wherever the format runs on supported hardware (Hopper / SM90), mirroring how the +// nvfp4 / mxfp8 formats fold their activation quant into those kernels by capability rather than a flag. +// The env override forces the legacy standalone scale_1x128 for A/B validation only (read live). +static inline bool useFp8BlockScaleActFusion() { if (tensorrt_llm::common::getSMVersion() != 90) { - return true; + return false; } - char const* env = std::getenv("TLLM_MOE_DISABLE_FUSED_FC2_QUANT"); - return env != nullptr && env[0] == '1'; + char const* env = std::getenv("TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION"); + return !(env != nullptr && env[0] == '1'); } template @@ -3375,14 +3482,30 @@ void CutlassMoeFCRunner runner with + // `input` (permuted_data_) already holding fp8 A + 1x128 scales from the fused expand; unfused uses the + // runner with `input` a bf16 activation moeGemm quantizes internally. configureWorkspace + // is assumed already called. + if (!useFp8BlockScaleActFusion()) + { + gemm_runner.moeGemm(gemm_output, input, fc1_expert_weights, expert_first_token_offset, num_experts_per_node, + expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, + quant_params.fp8_block_scaling.fc1_scales_ptrs); + } + else + { + auto const* fp8_a = reinterpret_cast<__nv_fp8_e4m3 const*>(input); + auto const* scales_a = reinterpret_cast( + reinterpret_cast(input) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); + gemm_runner.moeGemm(gemm_output, fp8_a, fc1_expert_weights, expert_first_token_offset, num_experts_per_node, + expected_tokens_per_expert, shape_n, shape_k, stream, scales_a, + quant_params.fp8_block_scaling.fc1_scales_ptrs); + } sync_check_cuda_error(stream); constexpr bool bias_is_broadcast = true; constexpr bool use_per_expert_act_scale = false; - if (isFusedFc2QuantDisabled()) + if (!useFp8BlockScaleActFusion()) { // Unfused path: write the bf16 activation; FC2's moeGemm will quantize it via the standalone scale_1x128. doActivation(output, static_cast(gemm_output), @@ -3392,15 +3515,13 @@ void CutlassMoeFCRunner(output); - size_t const fp8_bytes = static_cast(expanded_num_rows) * inter_size * sizeof(__nv_fp8_e4m3); - size_t const scale_off = ((fp8_bytes + 15) / 16) * 16; // 16B-align the float scale region - auto* fp8_block_scales = reinterpret_cast(reinterpret_cast(output) + scale_off); + auto* fp8_block_scales = reinterpret_cast( + reinterpret_cast(output) + fp8BlockScaleByteOffset(expanded_num_rows, inter_size)); int64_t const fp8_block_scale_ld = gemm_runner.getActScaleLeadingDim(); doActivation(output, static_cast(gemm_output), fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, @@ -3427,30 +3548,21 @@ void CutlassMoeFCRunner runner (consumes pre-quantized A, no internal quant), - // and `gemm_output` points at glu_inter_result_ (the outputs buffer), NOT the aliased - // fc2_result_, so the GEMM output does not clobber the fp8 A / 1x128 scales that `input` - // (== fc1_result_) still holds. - // - unfused: `gemm_runner` is the runner (quantizes A internally into its own - // workspace), and `gemm_output` is fc2_result_. - // NOTE: we assume gemm_runner.configureWorkspace has already been called. - if (isFusedFc2QuantDisabled()) - { - // Unfused path: `input` (fc1_result_) holds the bf16 activation; moeGemm quantizes it internally. + // `gemm_runner` and `gemm_output` are chosen by the caller (runMoe): fused uses the runner + // with `input` (fc1_result_) already holding fp8 A + 1x128 scales from the SwiGLU epilogue, and writes to + // glu_inter_result_ so the GEMM output does not clobber that fp8 A (fc1_result_ is aliased onto + // fc2_result_); unfused uses the runner with `input` a bf16 activation moeGemm quantizes. + if (!useFp8BlockScaleActFusion()) + { gemm_runner.moeGemm(gemm_output, input, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, quant_params.fp8_block_scaling.fc2_scales_ptrs); } else { - // Fused path: `input` (fc1_result_) was populated by the FC1 activation epilogue: fp8 activations in the low - // half, per-token 1x128 scales in the 16B-aligned upper half. The runner consumes them - // directly (internal_quantize_a is compile-time false), so no a_is_prequantized flag is needed. auto const* fp8_a = reinterpret_cast<__nv_fp8_e4m3 const*>(input); - size_t const fp8_bytes = static_cast(expanded_num_rows) * inter_size * sizeof(__nv_fp8_e4m3); - size_t const scale_off = ((fp8_bytes + 15) / 16) * 16; - auto const* scales_a = reinterpret_cast(reinterpret_cast(input) + scale_off); + auto const* scales_a = reinterpret_cast( + reinterpret_cast(input) + fp8BlockScaleByteOffset(expanded_num_rows, inter_size)); gemm_runner.moeGemm(gemm_output, fp8_a, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, expected_tokens_per_expert, shape_n, shape_k, stream, scales_a, quant_params.fp8_block_scaling.fc2_scales_ptrs); @@ -4456,12 +4568,37 @@ void CutlassMoeFCRunner(smoothed_act_) : reinterpret_cast(permuted_data_); - // Expand input and maybe apply prequant scale for AWQ - expandInputRowsKernelLauncher(input_activations, gemm1_input_expand, token_topk_unpermuted_scales, - permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, experts_per_token, - num_experts_per_node, quant_params, use_per_expert_act_scale, expert_first_token_offset_, - fc1_fp4_act_scale_, input_sf, swizzled_input_sf, - (use_w4afp8 && !use_fp8_input) ? quant_params.groupwise.fc1.act_scales : nullptr, stream); + // Fuse the pre-FC1 quant into the row expansion when enabled: write fp8 A + 1x128 scales into + // permuted_data_ (matching the layout BlockScaleFC1 reads) so FC1 uses the runner. The + // if constexpr keeps the bf16->fp8 quant out of the fp8/fp4 InputType runner instantiations, which + // never take this path at runtime anyway. + bool const fc1_use_fused = (blockscale_gemm_runner != nullptr) && useFp8BlockScaleActFusion(); + bool fused_fc1_expand_done = false; + if constexpr (std::is_same_v) + { + if (fc1_use_fused) + { + int64_t const fc1_scale_ld = blockscale_gemm_runner_fp8a_->getActScaleLeadingDim(); + auto* fp8_a = reinterpret_cast<__nv_fp8_e4m3*>(permuted_data_); + auto* fp8_scales = reinterpret_cast( + reinterpret_cast(permuted_data_) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); + expandInputRowsFp8BlockScaleKernelLauncher(input_activations, fp8_a, fp8_scales, fc1_scale_ld, + token_topk_unpermuted_scales, permuted_token_final_scales_, permuted_row_to_unpermuted_row_, + num_rows, hidden_size, experts_per_token, num_experts_per_node, expert_first_token_offset_, stream); + fused_fc1_expand_done = true; + } + } + if (!fused_fc1_expand_done) + { + // Expand input and maybe apply prequant scale for AWQ + expandInputRowsKernelLauncher(input_activations, gemm1_input_expand, token_topk_unpermuted_scales, + permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, experts_per_token, + num_experts_per_node, quant_params, use_per_expert_act_scale, expert_first_token_offset_, + fc1_fp4_act_scale_, input_sf, swizzled_input_sf, + (use_w4afp8 && !use_fp8_input) ? quant_params.groupwise.fc1.act_scales : nullptr, stream); + } + // For the fused FC1 path gemm1_input still points at permuted_data_ (the fp8-A base); BlockScaleFC1 + // reinterprets it as fp8 + scales. auto const* gemm1_input = gemm1_input_expand; sync_check_cuda_error(stream); @@ -4502,7 +4639,10 @@ void CutlassMoeFCRunner(smoothed_act_) : fc1_result_; - Self::gemm1(moe_gemm_runner_, blockscale_gemm_runner, gemm1_input, gemm1_output, glu_inter_result_, + // Fused FC1 uses the runner (consumes the fp8 A + scales the fused expand wrote); unfused + // uses the runner (quantizes A internally). BlockScaleFC1 reads the input accordingly. + auto* fc1_blockscale_runner = fc1_use_fused ? blockscale_gemm_runner_fp8a_.get() : blockscale_gemm_runner; + Self::gemm1(moe_gemm_runner_, fc1_blockscale_runner, gemm1_input, gemm1_output, glu_inter_result_, expert_first_token_offset_, gemm1_tma_ws_input, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, fc1_fp8_dequant, use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, fc1_fp4_act_scale_, fc2_fp4_act_scale_, quant_params, num_rows, expanded_num_rows, @@ -4528,14 +4668,10 @@ void CutlassMoeFCRunner runner, which consumes the fp8 A + 1x128 scales - // the FC1 SwiGLU epilogue already wrote in place into fc1_result_. Since the deepseek workspace - // aliases fc2_result_ onto fc1_result_, the GEMM output must go to glu_inter_result_ (already - // consumed by that epilogue) so input and output are disjoint. - // - unfused: use the runner (quantizes A into its own workspace first) and write to - // fc2_result_ as before. - bool const fc2_use_fused = (blockscale_gemm_runner != nullptr) && !isFusedFc2QuantDisabled(); + // Fused FC2 uses the runner (consumes the fp8 A + scales the SwiGLU epilogue wrote into + // fc1_result_) and writes to glu_inter_result_ so the output does not clobber that fp8 A (the workspace + // aliases fc2_result_ onto fc1_result_); unfused uses the runner and writes fc2_result_. + bool const fc2_use_fused = (blockscale_gemm_runner != nullptr) && useFp8BlockScaleActFusion(); auto* fc2_blockscale_runner = fc2_use_fused ? blockscale_gemm_runner_fp8a_.get() : blockscale_gemm_runner; void* const fc2_gemm_output = fc2_use_fused ? glu_inter_result_ : fc2_result_; Self::gemm2(moe_gemm_runner_, fc2_blockscale_runner, gemm2_input, fc2_gemm_output, final_output, diff --git a/tests/unittest/_torch/modules/moe/test_fc2_fusion.py b/tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py similarity index 77% rename from tests/unittest/_torch/modules/moe/test_fc2_fusion.py rename to tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py index 07ced968aa66..e551471c6695 100644 --- a/tests/unittest/_torch/modules/moe/test_fc2_fusion.py +++ b/tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py @@ -12,20 +12,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit test for the CUTLASS DeepSeek FP8 block-scale MoE FC2 activation-quant fusion. +"""Unit test for the CUTLASS DeepSeek FP8 block-scale MoE activation-quant fusions. -The block-scale MoE runs FC1 -> SwiGLU activation -> FC2 with bf16 (non-fp8) inputs to FC1. -FC2 requires the SwiGLU activation to be quantized to fp8 with per-token 1x128 scales. Two code -paths produce that fp8 input: +The block-scale MoE runs FC1 -> SwiGLU -> FC2; both GEMMs need their bf16 input quantized to fp8 with +per-token 1x128 scales. Two fusions produce those inputs in-place instead of a standalone scale_1x128: - * fused (default): the 1x128 quantization is folded into the FC1 SwiGLU epilogue - (doActivationKernel WriteFp8BlockScale), so FC2's grouped GEMM consumes - pre-quantized fp8 activations. - * unfused (TLLM_MOE_DISABLE_FUSED_FC2_QUANT=1): FC1 writes bf16, and FC2's grouped GEMM - quantizes it with the standalone scale_1x128 kernel. + * pre-FC1: folded into the row-expansion kernel (expandInputRowsFp8BlockScaleKernel). + * pre-FC2: folded into the FC1 SwiGLU epilogue (doActivationKernel WriteFp8BlockScale). -The fused path is designed to be numerically identical to the unfused one, so this test asserts -they match, and that both are close to a bf16 reference. Hopper (SM90) only. +Both are enabled automatically on SM90; setting TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION forces the +legacy standalone path for validation. The fusions are designed to be numerically identical to that +path, so this test asserts fused == unfused and that both are close to a bf16 reference. SM90 only. """ import pytest @@ -41,7 +38,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo -DISABLE_FUSION_ENV = "TLLM_MOE_DISABLE_FUSED_FC2_QUANT" +DISABLE_FUSION_ENV = "TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION" # Qwen3.5-35B-A3B MoE shapes (text_config): hidden_size=2048, moe_intermediate_size=512. @@ -51,11 +48,11 @@ @pytest.mark.skipif( get_sm_version() != 90, - reason="CUTLASS FP8 block-scale FC2 fusion is wired for Hopper (SM90) only.", + reason="CUTLASS FP8 block-scale activation-quant fusion is wired for Hopper (SM90) only.", ) @pytest.mark.parametrize("seq_len", [1, 8, 64], ids=lambda s: f"seq{s}") -def test_cutlass_fp8_block_scale_fc2_fusion(seq_len, monkeypatch): - """Fused FC2 activation-quant must match the unfused standalone scale_1x128 path. +def test_cutlass_fp8_block_scale_act_quant_fusion(seq_len, monkeypatch): + """Fused FC1+FC2 activation-quant must match the unfused standalone scale_1x128 path. Uses the Qwen3.5-35B-A3B MoE hidden/intermediate sizes (2048 / 512). """ @@ -124,7 +121,7 @@ def run_moe(): router_logits, ) - # Autotune once (with fusion on) so both A/B runs replay the same FC2 GEMM tactics. + # Autotune once (with fusion on) so both A/B runs replay the same FC1/FC2 GEMM tactics. monkeypatch.delenv(DISABLE_FUSION_ENV, raising=False) AutoTuner.get().clear_cache() with torch.inference_mode(), autotune(): @@ -133,7 +130,7 @@ def run_moe(): with torch.inference_mode(): out_fused = run_moe() - # Unfused path: FC1 writes bf16, FC2 quantizes with the standalone scale_1x128 kernel. + # Unfused path: FC1/FC2 quantize with the standalone scale_1x128 kernel. monkeypatch.setenv(DISABLE_FUSION_ENV, "1") with torch.inference_mode(): out_unfused = run_moe() @@ -142,14 +139,14 @@ def run_moe(): with torch.inference_mode(): ref_output = ref_fused_moe.forward(x, router_logits) - # Primary check: the fused epilogue reproduces the standalone quant bit-for-bit (same bf16 - # activation, same amax, same fp8 rounding, same GEMM), so outputs should be ~identical. + # Primary check: the fused quant reproduces the standalone quant bit-for-bit (same bf16 activation, + # same amax, same fp8 rounding, same GEMM), so outputs should be ~identical. torch.testing.assert_close( out_fused, out_unfused, rtol=5e-3, atol=5e-3, - msg="Fused FC2 quant diverges from the unfused standalone path", + msg="Fused activation quant diverges from the unfused standalone path", ) # Sanity: the fused output is a valid MoE result (finite and well-correlated with the bf16 From 4dd38b4458ed73b714268154808fd75c4b5e6f51 Mon Sep 17 00:00:00 2001 From: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:08:09 +0000 Subject: [PATCH 3/3] Build a single FP8 block-scale runner Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> --- .../fp8_blockscale_gemm/fp8_blockscale_gemm.h | 9 +- .../cutlass_kernels/include/moe_kernels.h | 4 +- .../cutlass_kernels/moe_gemm/moe_kernels.cu | 308 +++++++++--------- .../moe/test_fp8_block_scale_fusion.py | 84 ++--- 4 files changed, 207 insertions(+), 198 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h index 0530f7082475..87daa8a1d661 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h @@ -21,6 +21,7 @@ #include #include #include +#include #include // non-persistent-cooperative GEMM @@ -55,9 +56,10 @@ class CutlassFp8BlockScaleGemmRunnerInterface float const* scales_b = nullptr) = 0; - // Stride between K-blocks of per-token 1x128 activation scale buffer virtual int64_t getActScaleLeadingDim() const = 0; + virtual bool isActivationPrequantized() const = 0; + virtual void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a, __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) @@ -125,6 +127,11 @@ class CutlassFp8BlockScaleGemmRunner : public CutlassFp8BlockScaleGemmRunnerInte return max_shape_m_32_align_padded_; } + bool isActivationPrequantized() const override + { + return std::is_same_v; + } + void strideBatchGemm(__nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a, __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n, int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) override; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index ab5e9affe761..ab7ed876257d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -979,8 +979,6 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface MoeGemmRunner moe_gemm_runner_; std::unique_ptr blockscale_gemm_runner_; - // Second block-scale runner with pre-quantized fp8 activations - std::unique_ptr blockscale_gemm_runner_fp8a_; std::optional gemm1_config_; std::optional gemm2_config_; 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 823346f6eb64..230168ac27e9 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 @@ -59,6 +59,7 @@ #include "tensorrt_llm/kernels/quantization.cuh" #include "tensorrt_llm/common/tllmDataType.h" +#include "tensorrt_llm/deep_gemm/scheduler.cuh" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h" // NOTE: the grouped-GEMM dispatch (cudaGraph(SplitK)GroupedGemm, @@ -1707,9 +1708,15 @@ INSTANTIATE_EXPAND_INPUT_ROWS(__nv_bfloat16, __nv_bfloat16); // // FC1 and FC2 each consume fp8 activations + per-token 1x128 scales. Two fusions produce that input // in-place instead of a standalone scale_1x128 kernel: pre-FC1 in the row expansion, pre-FC2 in the -// SwiGLU epilogue. Both pack fp8 in the low part of a buffer and the scales just above, and both must -// match the standalone quant bit-for-bit. These helpers hold the shared layout and quant so the two -// call sites stay identical by construction. +// SwiGLU epilogue. Both pack fp8 in the low part of a buffer and the scales just above. +// These helpers hold the shared layout and quant. + +struct Fp8BlockScaleActOutput +{ + __nv_fp8_e4m3* fp8_out = nullptr; + float* scales = nullptr; + int64_t scale_leading_dim = 0; +}; // Byte offset from the buffer base to the 1x128 scale region that follows `rows` x `cols` fp8 // activations (16B-aligned so the float scales are aligned). @@ -1720,27 +1727,27 @@ static inline size_t fp8BlockScaleByteOffset(int64_t rows, int64_t cols) // Total bytes of the packed fp8-activations + 1x128-scales region (scale leading dim == grouped-GEMM // padded M, so shape_k-independent). -static inline size_t fp8BlockScaleRegionBytes(int64_t rows, int64_t cols, int64_t scale_ld) +static inline size_t fp8BlockScaleRegionBytes(int64_t rows, int64_t cols, int64_t scale_leading_dim) { return fp8BlockScaleByteOffset(rows, cols) - + static_cast(scale_ld) * tensorrt_llm::common::ceilDiv(cols, static_cast(128)) + + static_cast(scale_leading_dim) * tensorrt_llm::common::ceilDiv(cols, static_cast(128)) * sizeof(float); } -// Row offset of an expert's scales in the transposed / per-expert padded layout the grouped GEMM reads -// (deep_gemm::compute_padded_offset, alignment 32). +// Row offset of an expert's scales in the transposed / per-expert padded layout the grouped GEMM reads. __device__ inline int64_t fp8BlockScaleExpertPad(int64_t num_tokens_before_expert, int64_t expert) { - return (num_tokens_before_expert + expert * 31) / 32 * 32 - num_tokens_before_expert; + return deep_gemm::compute_padded_offset(num_tokens_before_expert, expert) - num_tokens_before_expert; } // Quantize one thread's 8 already-bf16-rounded channels to fp8 with a per-token 1x128 scale, bit-for-bit -// with the standalone scale_1x128 (amax reduced and truncated in bf16, plain-float 1/scale). The 8 -// channels span 16 warp lanes per 128-block; only bit 4 of `lane` is used, to mask that half-warp so a -// disjoint inactive sibling is safe (guaranteed by dim % 128 == 0). The block leader writes 1/scale. +// with the standalone scale_1x128: amax is reduced in float, then truncated to bf16 once before scale = +// 448/amax (float). The 8 channels span 16 warp lanes per 128-block; only bit 4 of `lane` is used, to mask +// that half-warp so a disjoint inactive sibling is safe (guaranteed by dim % 128 == 0). The block leader +// writes 1/scale. template __device__ inline void fp8BlockScaleQuantize(float const (&vals)[N], __nv_fp8_e4m3* out, float* block_scales, - int64_t kb, int64_t scale_ld, int64_t scale_row, bool write_scale, unsigned lane) + int64_t kb, int64_t scale_leading_dim, int64_t scale_row, bool write_scale, unsigned lane) { static_assert(N == 8, "1x128 block quant assumes 8 activation channels per thread (16 lanes/block)."); float amax = 0.f; @@ -1756,17 +1763,37 @@ __device__ inline void fp8BlockScaleQuantize(float const (&vals)[N], __nv_fp8_e4 for (int e = 0; e < N; ++e) out[e] = static_cast<__nv_fp8_e4m3>(vals[e] * scale); if (write_scale) - block_scales[kb * scale_ld + scale_row] = 1.f / scale; + block_scales[kb * scale_leading_dim + scale_row] = 1.f / scale; +} + +// Run one block-scale grouped GEMM (FC1 or FC2). When the runner is prequantized, `input` packs the fp8 A +// followed by its per-token 1x128 scales (see fp8BlockScaleByteOffset), where `cols` is the activation +// width; otherwise `input` is a bf16 activation the runner quantizes internally. +static inline void runBlockScaleMoeGemm(kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunnerInterface& runner, + void* gemm_output, void const* input, void const* weights, int64_t const* expert_first_token_offset, + int64_t num_experts_per_node, int64_t expected_tokens_per_expert, int shape_n, int shape_k, + int64_t expanded_num_rows, int64_t cols, float const* weight_scales, cudaStream_t stream) +{ + if (!runner.isActivationPrequantized()) + { + runner.moeGemm(gemm_output, input, weights, expert_first_token_offset, num_experts_per_node, + expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, weight_scales); + return; + } + auto const* fp8_a = reinterpret_cast<__nv_fp8_e4m3 const*>(input); + auto const* scales_a = reinterpret_cast( + reinterpret_cast(input) + fp8BlockScaleByteOffset(expanded_num_rows, cols)); + runner.moeGemm(gemm_output, fp8_a, weights, expert_first_token_offset, num_experts_per_node, + expected_tokens_per_expert, shape_n, shape_k, stream, scales_a, weight_scales); } // Fuse the pre-FC1 1x128 activation quant into the row expansion: permute the bf16 input rows like // expandInputRowsKernel, but write fp8 + scales the FC1 GEMM consumes directly. template __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* unpermuted_input, - __nv_fp8_e4m3* permuted_output, float* fp8_block_scales, int64_t const fp8_block_scale_ld, - float const* unpermuted_scales, float* permuted_scales, int const* permuted_row_to_unpermuted_row, - int64_t const num_tokens, int64_t const hidden_size, int64_t const k, int64_t const* expert_first_token_offset, - int64_t const num_experts_per_node) + Fp8BlockScaleActOutput fp8_block_scale_out, float const* unpermuted_scales, float* permuted_scales, + int const* permuted_row_to_unpermuted_row, int64_t const num_tokens, int64_t const hidden_size, int64_t const k, + int64_t const* expert_first_token_offset, int64_t const num_experts_per_node) { // 16-bit input => 8 channels per thread => a 1x128 block spans 16 contiguous lanes. constexpr int64_t ELEM_PER_THREAD = 128 / sizeof_bits::value; @@ -1778,6 +1805,7 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; int64_t const num_elems_in_col = hidden_size / ELEM_PER_THREAD; + auto* const permuted_output = fp8_block_scale_out.fp8_out; for (int64_t permuted_row = blockIdx.x; permuted_row < num_valid_tokens; permuted_row += gridDim.x) { @@ -1799,8 +1827,8 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u #pragma unroll for (int e = 0; e < ELEM_PER_THREAD; ++e) vals[e] = static_cast(in_vec[e]); // bf16 input is already exact - fp8BlockScaleQuantize(vals, dest_row_ptr + elem_index * ELEM_PER_THREAD, fp8_block_scales, elem_index >> 4, - fp8_block_scale_ld, scale_row, (elem_index & 15) == 0, threadIdx.x); + fp8BlockScaleQuantize(vals, dest_row_ptr + elem_index * ELEM_PER_THREAD, fp8_block_scale_out.scales, + elem_index >> 4, fp8_block_scale_out.scale_leading_dim, scale_row, (elem_index & 15) == 0, threadIdx.x); } if (permuted_scales && threadIdx.x == 0) @@ -1817,10 +1845,9 @@ __global__ void expandInputRowsFp8BlockScaleKernel(InputActivationsType const* u template void expandInputRowsFp8BlockScaleKernelLauncher(InputActivationsType const* unpermuted_input, - __nv_fp8_e4m3* permuted_output, float* fp8_block_scales, int64_t const fp8_block_scale_ld, - float const* unpermuted_scales, float* permuted_scales, int const* permuted_row_to_unpermuted_row, - int64_t const num_rows, int64_t const hidden_size, int const k, int const num_experts_per_node, - int64_t* expert_first_token_offset, cudaStream_t stream) + Fp8BlockScaleActOutput fp8_block_scale_out, float const* unpermuted_scales, float* permuted_scales, + int const* permuted_row_to_unpermuted_row, int64_t const num_rows, int64_t const hidden_size, int const k, + int const num_experts_per_node, int64_t* expert_first_token_offset, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(hidden_size % 128 == 0, "Fused FC1 activation quant requires hidden_size %% 128 == 0."); auto* func = &expandInputRowsFp8BlockScaleKernel; @@ -1839,9 +1866,9 @@ void expandInputRowsFp8BlockScaleKernelLauncher(InputActivationsType const* unpe attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, func, unpermuted_input, permuted_output, fp8_block_scales, fp8_block_scale_ld, - unpermuted_scales, permuted_scales, permuted_row_to_unpermuted_row, num_rows, hidden_size, - static_cast(k), expert_first_token_offset, static_cast(num_experts_per_node)); + cudaLaunchKernelEx(&config, func, unpermuted_input, fp8_block_scale_out, unpermuted_scales, permuted_scales, + permuted_row_to_unpermuted_row, num_rows, hidden_size, static_cast(k), expert_first_token_offset, + static_cast(num_experts_per_node)); } enum class ScaleMode : int @@ -2219,8 +2246,7 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern float const* fc2_act_global_scale, bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, ActivationParams activation_params, GemmOutputType const* prequant_scale, float* dynamic_fc2_amax = nullptr, - GemmOutputType* bf16_intermediate_output = nullptr, __nv_fp8_e4m3* fp8_block_output = nullptr, - float* fp8_block_scales = nullptr, int64_t fp8_block_scale_ld = 0) + GemmOutputType* bf16_intermediate_output = nullptr, Fp8BlockScaleActOutput fp8_block_scale_out = {}) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v @@ -2448,8 +2474,9 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern rounded[k] = static_cast(static_cast(post_act_val[k])); int64_t const scale_row = token + fp8BlockScaleExpertPad(num_tokens_before_expert, expert); fp8BlockScaleQuantize(rounded, - fp8_block_output + output_offset + elem_index * ACTIVATION_ELEM_PER_THREAD, fp8_block_scales, - col_offset >> 4, fp8_block_scale_ld, scale_row, (col_offset & 15) == 0, threadIdx.y); + fp8_block_scale_out.fp8_out + output_offset + elem_index * ACTIVATION_ELEM_PER_THREAD, + fp8_block_scale_out.scales, col_offset >> 4, fp8_block_scale_out.scale_leading_dim, scale_row, + (col_offset & 15) == 0, threadIdx.y); } else { @@ -2662,8 +2689,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 bool bias_is_broadcast, int64_t const* expert_first_token_offset, int num_experts_per_node, int64_t inter_size, int64_t expanded_num_tokens, ActivationParams activation_type, QuantParams const& quant_params, bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, cudaStream_t stream, - GemmOutputType const* prequant_scale = nullptr, __nv_fp8_e4m3* fp8_block_output = nullptr, - float* fp8_block_scales = nullptr, int64_t fp8_block_scale_ld = 0, bool write_fp8_block_scale = false) + GemmOutputType const* prequant_scale = nullptr, Fp8BlockScaleActOutput fp8_block_scale_out = {}) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v; @@ -2691,20 +2717,19 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 // common.h using KernelFnPtr = void (*)(T*, GemmOutputType const*, float const*, ScaleBiasType const*, bool, int64_t const*, int, int64_t, float const*, bool, TmaWarpSpecializedGroupedGemmInput::ElementSF*, - ActivationParams, GemmOutputType const*, float*, GemmOutputType*, __nv_fp8_e4m3*, float*, int64_t); - // Fuse the pre-FC2 1x128 activation quantization (WriteFp8BlockScale=true). Only Swiglu + a 16-bit T - // (bf16/half) + NONE block-scaling is used on the DeepSeek FP8 block-scale MoE path. Guard the - // instantiation so it is never emitted for fp4/fp8 T (where NONE block-scaling would compile the fp8 - // branch with an illegal wide vector width). + ActivationParams, GemmOutputType const*, float*, GemmOutputType*, Fp8BlockScaleActOutput); + + // Fuse the pre-FC2 1x128 activation quantization. Only Swiglu + a 16-bit T (bf16/half) + NONE block-scaling + // is used on the DeepSeek FP8 block-scale MoE path. Guard the instantiation so it is never emitted for + // fp4/fp8 T if constexpr (!IsNVFP4 && !IsMXFP8) { - if (write_fp8_block_scale) + if (fp8_block_scale_out.fp8_out != nullptr) { - return static_cast( - &doActivationKernel, - TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, num_rows_per_cta_v, - /*DynamicFc2=*/false, /*WriteFp8BlockScale=*/true>); + return static_cast(&doActivationKernel, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, + num_rows_per_cta_v, false, true>); } } auto fn = [&](auto block_scaling_type) -> KernelFnPtr @@ -2806,7 +2831,7 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 cudaLaunchKernelEx(&config, fn, output, gemm_result, fp8_quant, bias, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, quant_params.fp4.fc2.act_global_scale, use_per_expert_act_scale, fc2_act_sf_flat, activation_type, prequant_scale, (float*) nullptr, - (GemmOutputType*) nullptr, fp8_block_output, fp8_block_scales, fp8_block_scale_ld); + (GemmOutputType*) nullptr, fp8_block_scale_out); }; // end lambda doActivationKernelLauncher // 256 threads per block * 256 blocks / 1 rows per block can be handled by 1-2 waves depending on SM arch @@ -2904,8 +2929,7 @@ void doActivationDynamic(T* output, GemmOutputType const* gemm_result, float con fn<<>>(output, gemm_result, fp8_quant, bias, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, quant_params.fp4.fc2.act_global_scale, use_per_expert_act_scale, fc2_act_sf_flat, activation_type, - (GemmOutputType const*) nullptr, dynamic_amax, bf16_intermediate, /*fp8_block_output=*/nullptr, - /*fp8_block_scales=*/nullptr, /*fp8_block_scale_ld=*/0); + (GemmOutputType const*) nullptr, dynamic_amax, bf16_intermediate, Fp8BlockScaleActOutput{}); sync_check_cuda_error(stream); } @@ -3112,12 +3136,33 @@ void dequantFP8(OutputType* output, InputType const* input, int64_t const* num_v <<>>(output, input, num_valid_tokens_ptr, inter_size, scale, scale_is_dequant); } +// The DeepSeek FP8 block-scale MoE folds the pre-FC1 and pre-FC2 1x128 activation quant into the +// row-expansion and SwiGLU-epilogue kernels, enabled automatically on supported hardware (Hopper / SM90) -- +// mirroring how the nvfp4 / mxfp8 formats fold their activation quant by capability rather than a flag. The +// env override forces the legacy standalone scale_1x128 path. Evaluated once at construction to pick the +// block-scale runner, so it is fixed for the runner's lifetime (set the env before building it). +static inline bool useFp8BlockScaleActFusion() +{ + if (tensorrt_llm::common::getSMVersion() != 90) + { + return false; + } + char const* env = std::getenv("TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION"); + return !(env != nullptr && env[0] == '1'); +} + template CutlassMoeFCRunner::CutlassMoeFCRunner() - : blockscale_gemm_runner_{std::make_unique< - kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>>()} - , blockscale_gemm_runner_fp8a_{std::make_unique< - kernels::fp8_blockscale_gemm::CutlassFp8BlockScaleGemmRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, __nv_bfloat16>>()} + // Build only the runner the enabled path needs: fused consumes pre-quantized fp8 A (, no + // internal scale_1x128 and no deepseek_fc_workspace); unfused quantizes bf16 A internally (, + // also the only non-Hopper-capable variant). + : blockscale_gemm_runner_{useFp8BlockScaleActFusion() + ? std::unique_ptr( + std::make_unique>()) + : std::unique_ptr( + std::make_unique>())} { } @@ -3240,26 +3285,31 @@ CutlassMoeFCRunner:: auto* blockscale_gemm_runner = getDeepSeekBlockScaleGemmRunner(); TLLM_CHECK(blockscale_gemm_runner != nullptr); - auto deepseek_fc1_workspace_size = blockscale_gemm_runner->getWorkspaceSize( - num_rows, factor * inter_size, hidden_size, experts_per_token, num_experts_per_node); - auto deepseek_fc2_workspace_size = blockscale_gemm_runner->getWorkspaceSize( - num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); - deepseek_fc_workspace_size = std::max(deepseek_fc1_workspace_size, deepseek_fc2_workspace_size); - - // Prime the runner's 1x128 scale leading dim (getActScaleLeadingDim()) so it matches - // what the fused quant writes. It needs no workspace (both operands pre-quantized); the leading dim - // depends only on (num_rows, top_k, num_experts), so it is shape_k-independent and shared by FC1/FC2. - blockscale_gemm_runner_fp8a_->getWorkspaceSize( - num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); - - // The fused quant 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. - int64_t const fused_scale_ld = blockscale_gemm_runner->getActScaleLeadingDim(); - overlapped_gemm1_gemm2_inputs_size = std::max( - {overlapped_gemm1_gemm2_inputs_size, fp8BlockScaleRegionBytes(num_moe_inputs, inter_size, fused_scale_ld), - fp8BlockScaleRegionBytes(num_moe_inputs, hidden_size, fused_scale_ld)}); + // getWorkspaceSize also sets the runner's 1x128 scale leading dim (getActScaleLeadingDim()); the dim + // 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()) + { + // 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. + 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)}); + } + else + { + // 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); + auto deepseek_fc2_workspace_size = blockscale_gemm_runner->getWorkspaceSize( + num_rows, hidden_size, inter_size, experts_per_token, num_experts_per_node); + deepseek_fc_workspace_size = std::max(deepseek_fc1_workspace_size, deepseek_fc2_workspace_size); + } } size_t map_offset = 0; @@ -3432,9 +3482,10 @@ void CutlassMoeFCRunnerconfigureWorkspace(getWsPtr(char{}, "deepseek_fc_workspace")); - // The fused-FC2 runner reads pre-quantized A + external scales, so it needs no internal workspace. - blockscale_gemm_runner_fp8a_->configureWorkspace(nullptr); + // Fused runner reads pre-quantized A + external scales, so it needs no internal workspace; the unfused + // runner quantizes A into deepseek_fc_workspace (only allocated in that case). + blockscale_gemm_runner->configureWorkspace( + blockscale_gemm_runner->isActivationPrequantized() ? nullptr : getWsPtr(char{}, "deepseek_fc_workspace")); } if (use_awq) @@ -3454,20 +3505,6 @@ CutlassMoeFCRunner: return blockscale_gemm_runner_.get(); } -// Both FP8 block-scale activation-quant fusions (pre-FC1 row expansion and pre-FC2 SwiGLU epilogue) are -// enabled automatically wherever the format runs on supported hardware (Hopper / SM90), mirroring how the -// nvfp4 / mxfp8 formats fold their activation quant into those kernels by capability rather than a flag. -// The env override forces the legacy standalone scale_1x128 for A/B validation only (read live). -static inline bool useFp8BlockScaleActFusion() -{ - if (tensorrt_llm::common::getSMVersion() != 90) - { - return false; - } - char const* env = std::getenv("TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION"); - return !(env != nullptr && env[0] == '1'); -} - template void CutlassMoeFCRunner::BlockScaleFC1( DeepSeekBlockScaleGemmRunner& gemm_runner, T const* const input, T* const output, void* const gemm_output, @@ -3482,30 +3519,16 @@ void CutlassMoeFCRunner runner with - // `input` (permuted_data_) already holding fp8 A + 1x128 scales from the fused expand; unfused uses the - // runner with `input` a bf16 activation moeGemm quantizes internally. configureWorkspace - // is assumed already called. - if (!useFp8BlockScaleActFusion()) - { - gemm_runner.moeGemm(gemm_output, input, fc1_expert_weights, expert_first_token_offset, num_experts_per_node, - expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, - quant_params.fp8_block_scaling.fc1_scales_ptrs); - } - else - { - auto const* fp8_a = reinterpret_cast<__nv_fp8_e4m3 const*>(input); - auto const* scales_a = reinterpret_cast( - reinterpret_cast(input) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); - gemm_runner.moeGemm(gemm_output, fp8_a, fc1_expert_weights, expert_first_token_offset, num_experts_per_node, - expected_tokens_per_expert, shape_n, shape_k, stream, scales_a, - quant_params.fp8_block_scaling.fc1_scales_ptrs); - } + // When the runner is prequantized, `input` (permuted_data_) already holds fp8 A + 1x128 scales from the + // fused expand; otherwise it is a bf16 activation that moeGemm quantizes internally. + runBlockScaleMoeGemm(gemm_runner, gemm_output, input, fc1_expert_weights, expert_first_token_offset, + num_experts_per_node, expected_tokens_per_expert, shape_n, shape_k, expanded_num_rows, hidden_size, + quant_params.fp8_block_scaling.fc1_scales_ptrs, stream); sync_check_cuda_error(stream); constexpr bool bias_is_broadcast = true; constexpr bool use_per_expert_act_scale = false; - if (!useFp8BlockScaleActFusion()) + if (!gemm_runner.isActivationPrequantized()) { // Unfused path: write the bf16 activation; FC2's moeGemm will quantize it via the standalone scale_1x128. doActivation(output, static_cast(gemm_output), @@ -3522,12 +3545,11 @@ void CutlassMoeFCRunner(output); auto* fp8_block_scales = reinterpret_cast( reinterpret_cast(output) + fp8BlockScaleByteOffset(expanded_num_rows, inter_size)); - int64_t const fp8_block_scale_ld = gemm_runner.getActScaleLeadingDim(); + int64_t const scale_leading_dim = gemm_runner.getActScaleLeadingDim(); doActivation(output, static_cast(gemm_output), fc2_fp8_quant, fc1_expert_biases, bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size, expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale, nullptr, stream, - /*prequant_scale=*/nullptr, fp8_block_output, fp8_block_scales, fp8_block_scale_ld, - /*write_fp8_block_scale=*/true); + /*prequant_scale=*/nullptr, Fp8BlockScaleActOutput{fp8_block_output, fp8_block_scales, scale_leading_dim}); } sync_check_cuda_error(stream); @@ -3548,25 +3570,12 @@ void CutlassMoeFCRunner runner - // with `input` (fc1_result_) already holding fp8 A + 1x128 scales from the SwiGLU epilogue, and writes to - // glu_inter_result_ so the GEMM output does not clobber that fp8 A (fc1_result_ is aliased onto - // fc2_result_); unfused uses the runner with `input` a bf16 activation moeGemm quantizes. - if (!useFp8BlockScaleActFusion()) - { - gemm_runner.moeGemm(gemm_output, input, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, - expected_tokens_per_expert, shape_n, shape_k, stream, nullptr, - quant_params.fp8_block_scaling.fc2_scales_ptrs); - } - else - { - auto const* fp8_a = reinterpret_cast<__nv_fp8_e4m3 const*>(input); - auto const* scales_a = reinterpret_cast( - reinterpret_cast(input) + fp8BlockScaleByteOffset(expanded_num_rows, inter_size)); - gemm_runner.moeGemm(gemm_output, fp8_a, fc2_expert_weights, expert_first_token_offset, num_experts_per_node, - expected_tokens_per_expert, shape_n, shape_k, stream, scales_a, - quant_params.fp8_block_scaling.fc2_scales_ptrs); - } + // When the runner is prequantized, `input` (fc1_result_) already holds fp8 A + 1x128 scales from the SwiGLU + // epilogue; the caller routes `gemm_output` to glu_inter_result_ so it does not clobber that fp8 A + // (fc1_result_ is aliased onto fc2_result_). Otherwise `input` is a bf16 activation moeGemm quantizes. + runBlockScaleMoeGemm(gemm_runner, gemm_output, input, fc2_expert_weights, expert_first_token_offset, + num_experts_per_node, expected_tokens_per_expert, shape_n, shape_k, expanded_num_rows, inter_size, + quant_params.fp8_block_scaling.fc2_scales_ptrs, stream); sync_check_cuda_error(stream); @@ -3920,8 +3929,7 @@ void CutlassMoeFCRunner(gemm_output), nullptr, static_cast(fc2_lora), false, expert_first_token_offset, num_experts_per_node, hidden_size, expanded_num_rows, ActivationParams(ActivationType::Identity), {}, false, nullptr, stream, - /*prequant_scale=*/nullptr, /*fp8_block_output=*/nullptr, /*fp8_block_scales=*/nullptr, - /*fp8_block_scale_ld=*/0, /*write_fp8_block_scale=*/false); + /*prequant_scale=*/nullptr, Fp8BlockScaleActOutput{}); sync_check_cuda_error(stream); } @@ -4353,6 +4361,8 @@ void CutlassMoeFCRunnerisActivationPrequantized(); TLLM_CHECK(input_activations); TLLM_CHECK(token_selected_experts); @@ -4568,23 +4578,21 @@ void CutlassMoeFCRunner(smoothed_act_) : reinterpret_cast(permuted_data_); - // Fuse the pre-FC1 quant into the row expansion when enabled: write fp8 A + 1x128 scales into - // permuted_data_ (matching the layout BlockScaleFC1 reads) so FC1 uses the runner. The - // if constexpr keeps the bf16->fp8 quant out of the fp8/fp4 InputType runner instantiations, which - // never take this path at runtime anyway. - bool const fc1_use_fused = (blockscale_gemm_runner != nullptr) && useFp8BlockScaleActFusion(); + // Fuse the pre-FC1 quant into the row expansion when the runner is prequantized: write fp8 A + 1x128 + // scales into permuted_data_ in the layout BlockScaleFC1 reads. bool fused_fc1_expand_done = false; if constexpr (std::is_same_v) { - if (fc1_use_fused) + if (use_fused_block_scale_quant) { - int64_t const fc1_scale_ld = blockscale_gemm_runner_fp8a_->getActScaleLeadingDim(); + int64_t const fc1_scale_leading_dim = blockscale_gemm_runner->getActScaleLeadingDim(); auto* fp8_a = reinterpret_cast<__nv_fp8_e4m3*>(permuted_data_); auto* fp8_scales = reinterpret_cast( reinterpret_cast(permuted_data_) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); - expandInputRowsFp8BlockScaleKernelLauncher(input_activations, fp8_a, fp8_scales, fc1_scale_ld, - token_topk_unpermuted_scales, permuted_token_final_scales_, permuted_row_to_unpermuted_row_, - num_rows, hidden_size, experts_per_token, num_experts_per_node, expert_first_token_offset_, stream); + expandInputRowsFp8BlockScaleKernelLauncher(input_activations, + Fp8BlockScaleActOutput{fp8_a, fp8_scales, fc1_scale_leading_dim}, token_topk_unpermuted_scales, + permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, + experts_per_token, num_experts_per_node, expert_first_token_offset_, stream); fused_fc1_expand_done = true; } } @@ -4597,8 +4605,6 @@ void CutlassMoeFCRunner(smoothed_act_) : fc1_result_; - // Fused FC1 uses the runner (consumes the fp8 A + scales the fused expand wrote); unfused - // uses the runner (quantizes A internally). BlockScaleFC1 reads the input accordingly. - auto* fc1_blockscale_runner = fc1_use_fused ? blockscale_gemm_runner_fp8a_.get() : blockscale_gemm_runner; - Self::gemm1(moe_gemm_runner_, fc1_blockscale_runner, gemm1_input, gemm1_output, glu_inter_result_, + Self::gemm1(moe_gemm_runner_, blockscale_gemm_runner, gemm1_input, gemm1_output, glu_inter_result_, expert_first_token_offset_, gemm1_tma_ws_input, fc1_expert_weights, fc1_expert_biases, num_valid_tokens_ptr, fc1_int_scales, fc1_fp8_dequant, use_wfp4afp8 ? fc2_wfp4afp8_quant_scale : fc2_fp8_quant, fc1_fp4_act_scale_, fc2_fp4_act_scale_, quant_params, num_rows, expanded_num_rows, @@ -4668,13 +4671,10 @@ void CutlassMoeFCRunner runner (consumes the fp8 A + scales the SwiGLU epilogue wrote into - // fc1_result_) and writes to glu_inter_result_ so the output does not clobber that fp8 A (the workspace - // aliases fc2_result_ onto fc1_result_); unfused uses the runner and writes fc2_result_. - bool const fc2_use_fused = (blockscale_gemm_runner != nullptr) && useFp8BlockScaleActFusion(); - auto* fc2_blockscale_runner = fc2_use_fused ? blockscale_gemm_runner_fp8a_.get() : blockscale_gemm_runner; - void* const fc2_gemm_output = fc2_use_fused ? glu_inter_result_ : fc2_result_; - Self::gemm2(moe_gemm_runner_, fc2_blockscale_runner, gemm2_input, fc2_gemm_output, final_output, + // Fused FC2 writes glu_inter_result_ instead of fc2_result_: its input (fc1_result_) holds the fp8 A + // from the SwiGLU epilogue and fc2_result_ is aliased onto fc1_result_, so writing there would clobber it. + void* const fc2_gemm_output = use_fused_block_scale_quant ? glu_inter_result_ : fc2_result_; + Self::gemm2(moe_gemm_runner_, blockscale_gemm_runner, gemm2_input, fc2_gemm_output, final_output, expert_first_token_offset_, gemm2_tma_ws_input, fc2_expert_weights, fc2_expert_biases, fc2_int_scales, fc2_fp8_dequant, fc2_fp4_act_scale_, quant_params, token_topk_unpermuted_scales, permuted_token_final_scales_, unpermuted_row_to_permuted_row, permuted_row_to_unpermuted_row_, diff --git a/tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py b/tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py index e551471c6695..f559b079bf9e 100644 --- a/tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py +++ b/tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py @@ -87,53 +87,57 @@ def test_cutlass_fp8_block_scale_act_quant_fusion(seq_len, monkeypatch): activation_type=ActivationType.Swiglu, ) - backend = create_test_backend( - backend_type=MoeBackendType.CUTLASS, - routing_method=routing_method, - num_experts=num_experts, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - dtype=dtype, - quant_config=quant_config, - mapping=mapping, - activation_type=ActivationType.Swiglu, - ) weights = quantize_util.create_weights(**quant_kwargs) - backend.load_weights([weights]) - backend.post_load_weights() - backend.cuda() ref_fused_moe = quantize_util.create_ref_module(routing_method) ref_fused_moe.load_weights([weights]) ref_fused_moe.cuda() - def run_moe(): - token_selected_experts, token_final_scales = routing_method.apply(router_logits) - x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False) - return run_backend_moe( - backend, - MoeBackendType.CUTLASS, - x_quantized, - x_sf, - token_selected_experts, - token_final_scales, - dtype, - router_logits, + def run_config(disable_fusion): + # The block-scale runner picks fused () vs unfused () once, at + # construction, from the env; so set the env BEFORE building the backend for this config. The runner + # is a per-instance member of the C++ FusedMoeRunner, so two backends give two independent runners. + if disable_fusion: + monkeypatch.setenv(DISABLE_FUSION_ENV, "1") + else: + monkeypatch.delenv(DISABLE_FUSION_ENV, raising=False) + backend = create_test_backend( + backend_type=MoeBackendType.CUTLASS, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + quant_config=quant_config, + mapping=mapping, + activation_type=ActivationType.Swiglu, ) - - # Autotune once (with fusion on) so both A/B runs replay the same FC1/FC2 GEMM tactics. - monkeypatch.delenv(DISABLE_FUSION_ENV, raising=False) - AutoTuner.get().clear_cache() - with torch.inference_mode(), autotune(): - run_moe() - - with torch.inference_mode(): - out_fused = run_moe() - - # Unfused path: FC1/FC2 quantize with the standalone scale_1x128 kernel. - monkeypatch.setenv(DISABLE_FUSION_ENV, "1") - with torch.inference_mode(): - out_unfused = run_moe() + backend.load_weights([weights]) + backend.post_load_weights() + backend.cuda() + + def run_moe(): + token_selected_experts, token_final_scales = routing_method.apply(router_logits) + x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False) + return run_backend_moe( + backend, + MoeBackendType.CUTLASS, + x_quantized, + x_sf, + token_selected_experts, + token_final_scales, + dtype, + router_logits, + ) + + AutoTuner.get().clear_cache() + with torch.inference_mode(), autotune(): + run_moe() + with torch.inference_mode(): + return run_moe() + + out_fused = run_config(disable_fusion=False) + out_unfused = run_config(disable_fusion=True) monkeypatch.delenv(DISABLE_FUSION_ENV, raising=False) with torch.inference_mode():