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..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 @@ -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. @@ -21,6 +21,7 @@ #include #include #include +#include #include // non-persistent-cooperative GEMM @@ -55,6 +56,10 @@ class CutlassFp8BlockScaleGemmRunnerInterface float const* scales_b = nullptr) = 0; + 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) @@ -117,6 +122,16 @@ 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_; + } + + 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/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index 8bed9c16b58e..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, @@ -1703,6 +1704,173 @@ 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. +// 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). +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_leading_dim) +{ + return fp8BlockScaleByteOffset(rows, cols) + + 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. +__device__ inline int64_t fp8BlockScaleExpertPad(int64_t num_tokens_before_expert, int64_t 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 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_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; +#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_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, + 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; + 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; + 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) + { + 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_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) + { + 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, + 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; + 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, 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 { NO_SCALE = 0, @@ -2070,14 +2238,15 @@ 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, Fp8BlockScaleActOutput fp8_block_scale_out = {}) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v @@ -2140,8 +2309,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 +2353,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 +2461,23 @@ __global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKern /* input_sf */ nullptr); // Pass nullptr input_sf so we write 0 } } + else if constexpr (WriteFp8BlockScale && ACTIVATION_ELEM_PER_THREAD == 8) + { + // 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) + 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_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 { // Use storeVec to force STG.128 vectorized store @@ -2502,7 +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) + GemmOutputType const* prequant_scale = nullptr, Fp8BlockScaleActOutput fp8_block_scale_out = {}) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v; @@ -2528,11 +2715,24 @@ 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*, 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 (fp8_block_scale_out.fp8_out != nullptr) + { + return static_cast(&doActivationKernel, + TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, + num_rows_per_cta_v, false, true>); + } + } + auto fn = [&](auto block_scaling_type) -> KernelFnPtr { switch (activation_type.activation_type) { @@ -2631,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); + (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 @@ -2729,7 +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); + (GemmOutputType const*) nullptr, dynamic_amax, bf16_intermediate, Fp8BlockScaleActOutput{}); sync_check_cuda_error(stream); } @@ -2936,10 +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>>()} + // 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>())} { } @@ -3055,15 +3278,38 @@ 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); - 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); + // 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; @@ -3236,7 +3482,10 @@ void CutlassMoeFCRunnerconfigureWorkspace(getWsPtr(char{}, "deepseek_fc_workspace")); + // 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) @@ -3270,16 +3519,38 @@ 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 (!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), + 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 quant into the SwiGLU epilogue: write fp8 + scales into the FC2-input buffer + // (`output`, bf16-sized so the scales fit above the fp8), removing FC2's standalone scale_1x128. + 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); + auto* fp8_block_scales = reinterpret_cast( + reinterpret_cast(output) + fp8BlockScaleByteOffset(expanded_num_rows, inter_size)); + 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, Fp8BlockScaleActOutput{fp8_block_output, fp8_block_scales, scale_leading_dim}); + } sync_check_cuda_error(stream); } @@ -3299,9 +3570,12 @@ 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, Fp8BlockScaleActOutput{}); sync_check_cuda_error(stream); } @@ -4087,6 +4361,8 @@ void CutlassMoeFCRunnerisActivationPrequantized(); TLLM_CHECK(input_activations); TLLM_CHECK(token_selected_experts); @@ -4302,12 +4578,33 @@ 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 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 (use_fused_block_scale_quant) + { + 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, + 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; + } + } + 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); + } auto const* gemm1_input = gemm1_input_expand; sync_check_cuda_error(stream); @@ -4374,7 +4671,10 @@ void CutlassMoeFCRunner 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: + + * pre-FC1: folded into the row-expansion kernel (expandInputRowsFp8BlockScaleKernel). + * pre-FC2: folded into the FC1 SwiGLU epilogue (doActivationKernel WriteFp8BlockScale). + +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 +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_FP8_BLOCK_SCALE_ACT_FUSION" + + +# 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 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_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). + """ + 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, + ) + + weights = quantize_util.create_weights(**quant_kwargs) + + ref_fused_moe = quantize_util.create_ref_module(routing_method) + ref_fused_moe.load_weights([weights]) + ref_fused_moe.cuda() + + 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, + ) + 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(): + ref_output = ref_fused_moe.forward(x, router_logits) + + # 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 activation 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})" + )