fix(moe): prepare MXFP8 MXFP4 profiler inputs

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

@@ -4730,8 +4730,9 @@ std::map<std::string, std::pair<size_t, size_t>> GemmProfilerBackend::getProfile
     quant_6_size = num_experts_per_node * sizeof(float);
   } else if (is_native_wfp4afp8_family) {
     quant_1_size = sizeof(float);
-    quant_2_size = getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, mScalingType) *
-                   sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF);
+    quant_2_size =
+        getOffsetWeightSF(num_experts_per_node, fc1_out_size, hidden_size, mScalingType) *
+        sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF);
     quant_3_size = num_experts_per_node * sizeof(float);
     quant_4_size = sizeof(float);
     quant_5_size = getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) *
@@ -5034,12 +5035,30 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr
   } else if (mDType == nvinfer1::DataType::kFP8 &&
              (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) {
     TLLM_CHECK(quant_1 && quant_2 && quant_3 && quant_4 && quant_5 && quant_6);
-    mQuantParams = QuantParams::FP8MXFP4(
-        static_cast<float const*>(quant_1),
-        static_cast<TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*>(quant_2),
-        static_cast<float const*>(quant_3), static_cast<float const*>(quant_4),
-        static_cast<TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*>(quant_5),
-        static_cast<float const*>(quant_6));
+    if (mUseMxfp8ActScaling) {
+#ifdef USING_OSS_CUTLASS_MOE_GEMM
+      // Normalize the profiler's fabricated weight block scale factors to UE8M0 unity so tactic
+      // timing is not driven by random exponents from the workspace.
+      TLLM_CUDA_CHECK(cudaMemsetAsync(const_cast<void*>(quant_2), 0x7F,
+                                      workspaces.at("quant_2").first, stream));
+      TLLM_CUDA_CHECK(cudaMemsetAsync(const_cast<void*>(quant_5), 0x7F,
+                                      workspaces.at("quant_5").first, stream));
+      mQuantParams = QuantParams::MXFP8MXFP4(
+          static_cast<TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*>(quant_2),
+          static_cast<float const*>(quant_3),
+          static_cast<TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*>(quant_5),
+          static_cast<float const*>(quant_6));
+#else
+      TLLM_CHECK_WITH_INFO(false, "MXFP8 x MXFP4 profiling requires OSS Cutlass MoE GEMM");
+#endif
+    } else {
+      mQuantParams = QuantParams::FP8MXFP4(
+          static_cast<float const*>(quant_1),
+          static_cast<TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*>(quant_2),
+          static_cast<float const*>(quant_3), static_cast<float const*>(quant_4),
+          static_cast<TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*>(quant_5),
+          static_cast<float const*>(quant_6));
+    }
   } else if ((mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64) &&
              (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) {
     // nvllm still uses int64 because torch doesn't have fp4 yet.

csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h

@@ -1078,6 +1078,7 @@ struct GemmProfilerBackend {
     mNeedWeights = need_weights;
     mParallelismConfig = parallelism_config;
     mEnableAlltoall = enable_alltoall;
+    mUseMxfp8ActScaling = use_mxfp8_act_scaling;
     mSm90Wfp4Afp8Mode = sm90_wfp4afp8_mode;
     mSM = common::getSMVersion();
 
@@ -1144,6 +1145,7 @@ struct GemmProfilerBackend {
   bool mUseLora{};
   bool mMinLatencyMode{};
   bool mNeedWeights{};
+  bool mUseMxfp8ActScaling{};
   Sm90Wfp4Afp8ScaleMode mSm90Wfp4Afp8Mode = Sm90Wfp4Afp8ScaleMode::kDisabled;
 
   TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType mScalingType{};

tests/moe/test_trtllm_cutlass_fused_moe.py

@@ -1498,6 +1498,7 @@ def dequant_mxfp8_batches(
 @pytest.mark.parametrize(
     ("alpha", "beta", "limit"), [(None, None, None), (0.5, 0.0, 7.0), (1.702, 1.0, 7.0)]
 )
+@pytest.mark.parametrize("use_autotune", [False, True])
 @pytest.mark.skipif(
     torch.cuda.get_device_capability()[0] not in [10, 11, 12],
     reason="MXFP8xMXFP4 is only supported on SM100, SM110 and SM120/SM121",
@@ -1512,6 +1513,7 @@ def test_moe_mxfp8_mxfp4(
     alpha,
     beta,
     limit,
+    use_autotune,
 ):
     """
     Test MoE with MXFP8 activations and MXFP4 weights.
@@ -1562,21 +1564,22 @@ def test_moe_mxfp8_mxfp4(
         beta_t = None
 
     # Call cutlass_fused_moe with MXFP8 activations and MXFP4 weights
-    _ = fused_moe.cutlass_fused_moe(
-        mxfp8_x,
-        selected_experts.to(torch.int),
-        routing_weights,
-        mxfp4_w1.contiguous().view(torch.long),
-        mxfp4_w2.contiguous().view(torch.long),
-        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,
+            mxfp4_w1.contiguous().view(torch.long),
+            mxfp4_w2.contiguous().view(torch.long),
+            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(