diff --git a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh index e3a9562d03a..4c5a6aedbbb 100644 --- a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh +++ b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh @@ -4450,6 +4450,26 @@ std::map> GemmProfilerBackend::getProfile : 0; size_t quant_6_size = is_fp4_w_quant ? num_experts_per_node * sizeof(float) : 0; + // MXFP8xMXFP8 sizes: the FP8 branch above only reserves per-tensor float + // scalars, but block-scaled MXFP8 needs per-expert weight block SFs and + // global scales (issue #3558). Layout mirrors QuantParams::MXFP8MXFP8: + // quant_2/quant_5 = fc1/fc2 weight block SFs, quant_3/quant_6 = fc1/fc2 + // per-expert global scales. fc1 N uses fc1_out_size (doubled when gated). + bool is_mxfp8_quant = + is_fp8_act_quant && is_fp8_w_quant && + mScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; + if (is_mxfp8_quant) { + quant_1_size = 0; + quant_2_size = + getOffsetWeightSF(num_experts_per_node, fc1_out_size, hidden_size, mScalingType) * + sizeof(TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF); + quant_3_size = num_experts_per_node * sizeof(float); + quant_4_size = 0; + quant_5_size = getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) * + sizeof(TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF); + quant_6_size = num_experts_per_node * sizeof(float); + } + size_t tma_ws_input_workspace_size = 0; if (is_tma_ws_input) { tma_ws_input_workspace_size = @@ -4622,7 +4642,7 @@ void GemmProfilerBackend::prepareRouting(int num_tokens, char* workspace_ptr_cha } void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr_char, - cudaStream_t) { + cudaStream_t stream) { auto workspaces = getProfilerWorkspaces(num_tokens, mSM >= 90); #define GET_WS_PTR(type, name) \ auto* name = (workspaces.at(#name).first \ @@ -4654,6 +4674,34 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr mQuantParams = QuantParams::GroupWise(mGroupSize, quant_1, quant_2, nullptr, nullptr, quant_3, quant_4); } + } else if (mDType == nvinfer1::DataType::kFP8 && mWType == nvinfer1::DataType::kFP8 && + mScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX) { +#ifdef USING_OSS_CUTLASS_MOE_GEMM + // MXFP8xMXFP8 must be matched before the generic FP8 branch below: the + // storage dtypes are identical and only mScalingType distinguishes them + // (issue #3558). + TLLM_CHECK(quant_2 && quant_3 && quant_5 && quant_6); + // Initialize the weight block-SF buffers to E8M0 1.0 (biased exponent + // 0x7F): the profiler never fills them with real data, and uninitialized + // exponents make the timed kernels read nondeterministic scale patterns, + // which can skew tactic rankings vs real traffic. + TLLM_CUDA_CHECK( + cudaMemsetAsync(const_cast(quant_2), 0x7F, workspaces.at("quant_2").first, stream)); + TLLM_CUDA_CHECK( + cudaMemsetAsync(const_cast(quant_5), 0x7F, workspaces.at("quant_5").first, stream)); + mQuantParams = QuantParams::MXFP8MXFP8( + static_cast(quant_2), + static_cast(quant_3), + static_cast(quant_5), + static_cast(quant_6)); + // Backfill the shared FP8 dequant aliases: prepareTmaWsInputs() and the + // common GEMM/TMA setup read mQuantParams.fp8.dequant_fc1/fc2 directly, + // mirroring the WMXFP8AMXFP8 remap that runMoe() performs. + mQuantParams.fp8.dequant_fc1 = static_cast(quant_3); + mQuantParams.fp8.dequant_fc2 = static_cast(quant_6); +#else + TLLM_CHECK_WITH_INFO(false, "MXFP8 x MXFP8 profiling requires OSS Cutlass MoE GEMM"); +#endif } else if (mWType == nvinfer1::DataType::kFP8) { TLLM_CHECK(quant_1 && quant_2 && quant_3); mQuantParams = @@ -4735,6 +4783,16 @@ void GemmProfilerBackend::prepareTmaWsInputs( #undef GET_WS_PTR + // For MXFP8, fill the activation block-SF buffer with E8M0 1.0 (biased + // exponent 0x7F). The profiler never quantizes real activations into it, + // and uninitialized exponents give the timed kernels a nondeterministic + // scale pattern, skewing tactic rankings relative to real traffic. + if (fp4_act_scale_flat != nullptr && + mScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX) { + TLLM_CUDA_CHECK(cudaMemsetAsync(fp4_act_scale_flat, 0x7F, + workspaces.at("fp4_act_scale_flat").first, stream)); + } + size_t tma_ws_size = TmaWarpSpecializedGroupedGemmInput::workspaceSize(mNumExpertsPerNode, mScalingType); diff --git a/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu b/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu index fdf8dbb6dab..582039b77ac 100644 --- a/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu +++ b/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu @@ -674,12 +674,12 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj { activation_dtype = isNvfp4Quant() ? dl_int64 : activation_dtype; int64_t const unpadded_hidden_size_profiler = hidden_size; // HACK no padding by default #ifdef USING_OSS_CUTLASS_MOE_GEMM - mProfiler->init(*mKernelRunner.get(), mProfiler->mGemmToProfile, - DtypeUtils::dataType(activation_dtype), DtypeUtils::dataType(mWeightDtype), - DtypeUtils::dataType(mOutputDtype), num_experts, static_cast(top_k), - hidden_size, unpadded_hidden_size_profiler, inter_size, group_size, - activation_type, USE_BIAS, USE_LORA, min_latency_mode, - /*need_weights*/ false, parallelism_config, enable_alltoall); + mProfiler->init( + *mKernelRunner.get(), mProfiler->mGemmToProfile, DtypeUtils::dataType(activation_dtype), + DtypeUtils::dataType(mWeightDtype), DtypeUtils::dataType(mOutputDtype), num_experts, + static_cast(top_k), hidden_size, unpadded_hidden_size_profiler, inter_size, + group_size, activation_type, USE_BIAS, USE_LORA, min_latency_mode, + /*need_weights*/ false, parallelism_config, enable_alltoall, mUseMxfp8ActScaling); #else mProfiler->init(*mKernelRunner.get(), mProfiler->mGemmToProfile, DtypeUtils::dataType(activation_dtype), DtypeUtils::dataType(mWeightDtype), diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 3b408dc24e8..429626bd8c6 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -15,6 +15,7 @@ */ #pragma once +#include #include #include "cutlass/gemm/gemm.h" @@ -635,7 +636,8 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { } std::vector getTactics(MoeGemmId gemm_id) override { - return moe_gemm_runner_.getConfigs(gemm_id == MoeGemmId::GEMM_2 && mayHaveFinalizeFused()); + return filterMxfp8Tactics( + moe_gemm_runner_.getConfigs(gemm_id == MoeGemmId::GEMM_2 && mayHaveFinalizeFused())); } int queryOccupancyForConfig(cutlass_extensions::CutlassGemmConfig const& config) override { @@ -644,8 +646,25 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { static std::vector getTactics(int sm, MoeGemmId gemm_id) { using RunnerType = decltype(moe_gemm_runner_); - return RunnerType::getConfigs(sm, - gemm_id == MoeGemmId::GEMM_2 && Self::mayHaveFinalizeFused(sm)); + return filterMxfp8Tactics( + RunnerType::getConfigs(sm, gemm_id == MoeGemmId::GEMM_2 && Self::mayHaveFinalizeFused(sm))); + } + + // MXFP8 shares the FP8 activation/weight types, so the underlying gemm + // runner also reports the non-TMA (SM80/SM89-style) fallback configs. + // Those paths cannot do FpX block scaling (they hard-assert + // !use_block_scaling), so they must never be offered as tactics for an + // MXFP8 runner instantiation. + static std::vector filterMxfp8Tactics( + std::vector configs) { + if constexpr (use_mxfp8) { + configs.erase(std::remove_if(configs.begin(), configs.end(), + [](cutlass_extensions::CutlassGemmConfig const& config) { + return !config.is_tma_warp_specialized; + }), + configs.end()); + } + return configs; } void runMoe(void const* input_activations, void const* input_sf, bool const swizzled_input_sf, @@ -1011,7 +1030,8 @@ struct GemmProfilerBackend { int num_experts, int k, int64_t hidden_size, int64_t unpadded_hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, - MOEParallelismConfig parallelism_config, bool const enable_alltoall) { + MOEParallelismConfig parallelism_config, bool const enable_alltoall, + bool use_mxfp8_act_scaling = false) { mInterface = &runner; mGemmToProfile = gemm_to_profile; mDType = dtype; @@ -1040,6 +1060,14 @@ struct GemmProfilerBackend { } else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; + } else if (use_mxfp8_act_scaling && dtype == nvinfer1::DataType::kFP8 && + wtype == nvinfer1::DataType::kFP8) { + // MXFP8xMXFP8: same storage dtypes as plain FP8, so the data types + // alone cannot identify it. Without this the profiler prepares + // per-tensor FP8 quant params and a null activation-SF buffer, and the + // TMA warp-specialized MXFP8 kernels fault on the null SF descriptor + // during the tuning pass (issue #3558). + mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } } diff --git a/tests/moe/test_trtllm_cutlass_fused_moe.py b/tests/moe/test_trtllm_cutlass_fused_moe.py index 5fc76cef86c..f92a15b19c5 100644 --- a/tests/moe/test_trtllm_cutlass_fused_moe.py +++ b/tests/moe/test_trtllm_cutlass_fused_moe.py @@ -1558,6 +1558,10 @@ def test_moe_mxfp8_mxfp4( @pytest.mark.parametrize( ("alpha", "beta", "limit"), [(None, None, None), (0.5, 0.0, 7.0), (1.702, 1.0, 7.0)] ) +# use_autotune=True is the regression coverage for issue #3558: the gemm +# profiler used to prepare per-tensor FP8 quant params and a null +# activation-SF buffer for MXFP8xMXFP8, crashing the tuning pass. +@pytest.mark.parametrize("use_autotune", [False, True]) @pytest.mark.skipif( torch.cuda.get_device_capability()[0] not in [10], reason="MXFP8xMXFP8 is only supported on SM100 for now", @@ -1572,6 +1576,7 @@ def test_moe_mxfp8_mxfp8( alpha, beta, limit, + use_autotune, ): """Test MoE with MXFP8 activations and MXFP8 weights.""" if top_k > num_experts: @@ -1617,21 +1622,22 @@ def test_moe_mxfp8_mxfp8( limit_t = None beta_t = None - _ = fused_moe.cutlass_fused_moe( - mxfp8_x, - selected_experts.to(torch.int), - routing_weights, - mxfp8_w1.contiguous(), - mxfp8_w2.contiguous(), - otype, - swiglu_alpha=alpha_t, - swiglu_limit=limit_t, - swiglu_beta=beta_t, - quant_scales=quant_scales, - input_sf=mxfp8_x_sf, - use_mxfp8_act_scaling=True, - output=flash_output, - ) + with autotune(True) if use_autotune else nullcontext(): + _ = fused_moe.cutlass_fused_moe( + mxfp8_x, + selected_experts.to(torch.int), + routing_weights, + mxfp8_w1.contiguous(), + mxfp8_w2.contiguous(), + otype, + swiglu_alpha=alpha_t, + swiglu_limit=limit_t, + swiglu_beta=beta_t, + quant_scales=quant_scales, + input_sf=mxfp8_x_sf, + use_mxfp8_act_scaling=True, + output=flash_output, + ) dq_mxfp8_x = ( mxfp8_dequantize_host(