Skip to content

fix: make the cutlass MoE gemm profiler MXFP8-aware (autotune crash on MXFP8xMXFP8) #3614

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kahyunnam merged 4 commits into from
Jul 9, 2026
Merged

fix: make the cutlass MoE gemm profiler MXFP8-aware (autotune crash on MXFP8xMXFP8) #3614

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6dfbdb2
fix: make the cutlass MoE gemm profiler MXFP8-aware (autotune crash)
waynehacking8 Jun 12, 2026
ddeb61d
address review: explicit <algorithm> include + MXFPXElementSF sizeof …
waynehacking8 Jun 12, 2026
9a553af
address review: backfill shared fp8 dequant aliases in profiler MXFP8…
waynehacking8 Jun 12, 2026
054ec99
profiler: deterministic representative init for MXFP8 SF dummy buffers
waynehacking8 Jun 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 59 additions & 1 deletion csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4450,6 +4450,26 @@ std::map<std::string, std::pair<size_t, size_t>> 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);
}
Comment on lines +4461 to +4471

@gemini-code-assist gemini-code-assist Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To ensure type safety and consistency with prepareQuantParams (where quant_2 and quant_5 are cast to TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF const*), it is safer to use sizeof(TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF) instead of sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF) when calculating quant_2_size and quant_5_size.

  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);
  }

@waynehacking8 waynehacking8 Jun 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in ddeb61d — both sizeofs now use MXFPXElementSF for consistency with the prepareQuantParams casts (currently an alias of ElementSF, but this keeps the branch type-coherent if they ever diverge). Recompiled + reran the FP8 no-regression check, bit-identical.


size_t tma_ws_input_workspace_size = 0;
if (is_tma_ws_input) {
tma_ws_input_workspace_size =
Expand Down Expand Up @@ -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 \
Expand Down Expand Up @@ -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<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::MXFP8MXFP8(
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));
// 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<float const*>(quant_3);
mQuantParams.fp8.dequant_fc2 = static_cast<float const*>(quant_6);
#else
TLLM_CHECK_WITH_INFO(false, "MXFP8 x MXFP8 profiling requires OSS Cutlass MoE GEMM");
#endif
Comment thread
coderabbitai[bot] marked this conversation as resolved.
} else if (mWType == nvinfer1::DataType::kFP8) {
TLLM_CHECK(quant_1 && quant_2 && quant_3);
mQuantParams =
Expand Down Expand Up @@ -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);

Expand Down
12 changes: 6 additions & 6 deletions csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(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<int>(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);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
#else
mProfiler->init(*mKernelRunner.get(), mProfiler->mGemmToProfile,
DtypeUtils::dataType(activation_dtype), DtypeUtils::dataType(mWeightDtype),
Expand Down
36 changes: 32 additions & 4 deletions csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
*/

#pragma once
#include <algorithm>
#include <cstdint>

#include "cutlass/gemm/gemm.h"
Expand Down Expand Up @@ -635,7 +636,8 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface {
}

std::vector<cutlass_extensions::CutlassGemmConfig> 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 {
Expand All @@ -644,8 +646,25 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface {

static std::vector<cutlass_extensions::CutlassGemmConfig> 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<cutlass_extensions::CutlassGemmConfig> filterMxfp8Tactics(
std::vector<cutlass_extensions::CutlassGemmConfig> configs) {
if constexpr (use_mxfp8) {
configs.erase(std::remove_if(configs.begin(), configs.end(),

@gemini-code-assist gemini-code-assist Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The use of std::remove_if requires the <algorithm> header. While it might be transitively included by other headers in some compilation units, it is highly recommended to explicitly include <algorithm> at the top of moe_kernels.h to prevent potential compilation failures on other toolchains or standard library implementations.

@waynehacking8 waynehacking8 Jun 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in ddeb61d — explicit #include <algorithm> added (it was indeed only transitively available).

[](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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}

Expand Down
36 changes: 21 additions & 15 deletions tests/moe/test_trtllm_cutlass_fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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:
Expand Down Expand Up @@ -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():

@aleozlx aleozlx Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

skip_ops is also available for this purpose

flashinfer/flashinfer/autotuner.py

Line 523 in af979cc

skip_ops: Optional set of ``custom_op`` names to exclude from

cc @qiching

@waynehacking8 waynehacking8 Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good pointer, thanks - kept the parametrized context here since the crash lived in the profiling path itself, so the test wants autotune fully on/off per case, but skip_ops is the right tool when only a specific op should stay on heuristics.

_ = 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(
Expand Down
Loading