fix: make the cutlass MoE gemm profiler MXFP8-aware (autotune crash on MXFP8xMXFP8) - #3614
Conversation
The autotuner tuning pass crashed for cutlass_fused_moe with MXFP8 activation scaling (illegal instruction / null TMA SF buffer on the reporter's SM121; !use_block_scaling kernel assert on current main). MXFP8xMXFP8 stores activations and weights as plain FP8, so the gemm profiler could not tell it apart from per-tensor FP8 and prepared the wrong tuning state in three coordinated places: - GemmProfilerBackend::init derived mScalingType from dtypes alone -> NONE -> getOffsetActivationSF sized the activation-SF workspace to zero -> the TMA warp-specialized MXFP8 kernels got a null SF pointer. - getProfilerWorkspaces only reserved per-tensor float scalars for FP8 weights; MXFP8 needs per-expert weight block SFs + global scales. - prepareQuantParams fell into the per-tensor QuantParams::FP8 branch; MXFP8 needs QuantParams::MXFP8MXFP8. Additionally CutlassMoeFCRunner::getTactics now filters out non-TMA (SM80/SM89-style) fallback configs for MXFP8 instantiations: those paths hard-assert !use_block_scaling, so offering them as tactics guarantees a tuning-pass crash on any arch where they are enumerated. Adds use_autotune parametrization to test_moe_mxfp8_mxfp8 (SM100-gated) as the regression test. Validated on SM120 (RTX PRO 6000): all changes JIT-compile; pure-FP8 cutlass_fused_moe with autotune produces bit-identical outputs pre/post change (the filter is a no-op for non-MXFP8); MXFP8 on SM120 (not an enabled arch on main) now fails fast at init with 'No valid tactics available' instead of the deep kernel assert. MXFP8 e2e runs on the SM100-gated CI test added here. Fixes the original-report half of flashinfer-ai#3558 (the kErrorInternal tactic failures from the follow-up comment are tracked by flashinfer-ai#3568). AI-assisted (Claude), reviewed and validated by the submitter.
📝 WalkthroughWalkthroughAdds MXFP8×MXFP8 profiler support: new init flag to enable MXFPX scaling, tactic filtering for MXFP8, profiler workspace and QuantParams for per-expert MXFPX scales, runtime wiring of the flag, and tests parametrized to exercise autotuning. ChangesMXFP8x MXFP8 Fused MoE Quantization
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses issue #3558 by correctly supporting MXFP8xMXFP8 block-scaled quantization in the GEMM profiler, preventing crashes during the autotuning pass. The changes include calculating the correct workspace sizes for MXFP8, preparing the appropriate quantization parameters, filtering out incompatible non-TMA tactics for MXFP8, and adding regression test coverage. The review feedback suggests improving type safety by using MXFPXElementSF instead of ElementSF when calculating quantization sizes, and explicitly including the <algorithm> header to ensure compilation stability when using std::remove_if.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if (is_mxfp8_quant) { | ||
| quant_1_size = 0; | ||
| 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 = 0; | ||
| quant_5_size = getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) * | ||
| sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF); | ||
| quant_6_size = num_experts_per_node * sizeof(float); | ||
| } |
There was a problem hiding this comment.
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);
}
There was a problem hiding this comment.
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.
| 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(), |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done in ddeb61d — explicit #include <algorithm> added (it was indeed only transitively available).
…consistency AI-assisted (Claude).
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu (1)
676-690: 💤 Low valueOSS vs non-OSS profiler initialization divergence.
The OSS Cutlass path (line 682) now receives
mUseMxfp8ActScaling, while the non-OSS path (lines 684-689) does not. This divergence is intentional per the PR scope, which targets OSS Cutlass MXFP8 support. However, a brief inline comment explaining why only the OSS path receives this parameter would improve maintainability.📝 Suggested comment
int64_t const unpadded_hidden_size_profiler = hidden_size; // HACK no padding by default `#ifdef` USING_OSS_CUTLASS_MOE_GEMM + // OSS Cutlass profiler needs mUseMxfp8ActScaling to allocate MXFP8-specific workspace 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); `#else`🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu` around lines 676 - 690, Add a brief inline comment near the conditional that calls mProfiler->init under the USING_OSS_CUTLASS_MOE_GEMM branch explaining why mUseMxfp8ActScaling is passed only in the OSS path (e.g., "MXFP8 activation scaling is supported/enabled only in the OSS Cutlass build, so this flag is not applicable to the non-OSS profiler init"), referencing the conditional macro USING_OSS_CUTLASS_MOE_GEMM and the mUseMxfp8ActScaling parameter so future maintainers understand the intentional divergence in the mProfiler->init calls.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh`:
- Around line 4677-4691: When constructing mQuantParams via
QuantParams::MXFP8MXFP8(...) inside the MXFP8MXFP8 branch, also populate the
common FP8 dequantation aliases mQuantParams.fp8.dequant_fc1 and
mQuantParams.fp8.dequant_fc2 (the same pointers used by runProfiler(),
prepareTmaWsInputs(), and the TMA/GEMM paths) by remapping the MXFP8-specific
element scales/float pointers into those alias fields so downstream code
(runProfiler, prepareTmaWsInputs, runMoe) sees valid dequant pointers; make this
change in the MXFP8 x MXFP8 branch alongside the existing
QuantParams::MXFP8MXFP8(...) construction.
In `@csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu`:
- Around line 677-682: The mUseMxfp8ActScaling boolean must be passed as the
parameter immediately following enable_alltoall to match
GemmProfilerBackend::init(..., bool const enable_alltoall, bool
use_mxfp8_act_scaling = false); confirm the call to mProfiler->init(...) uses
mUseMxfp8ActScaling in that position (as it currently does), and if any
reorderings were made ensure the argument order is restored so the final two
arguments are enable_alltoall then mUseMxfp8ActScaling (referencing
mProfiler->init, enable_alltoall, and mUseMxfp8ActScaling).
---
Nitpick comments:
In `@csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu`:
- Around line 676-690: Add a brief inline comment near the conditional that
calls mProfiler->init under the USING_OSS_CUTLASS_MOE_GEMM branch explaining why
mUseMxfp8ActScaling is passed only in the OSS path (e.g., "MXFP8 activation
scaling is supported/enabled only in the OSS Cutlass build, so this flag is not
applicable to the non-OSS profiler init"), referencing the conditional macro
USING_OSS_CUTLASS_MOE_GEMM and the mUseMxfp8ActScaling parameter so future
maintainers understand the intentional divergence in the mProfiler->init calls.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 539cffb8-5164-4051-b7e1-f10b91c7570b
📒 Files selected for processing (4)
csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuhcsrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cucsrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.htests/moe/test_trtllm_cutlass_fused_moe.py
… branch prepareTmaWsInputs() reads mQuantParams.fp8.dequant_fc1/fc2 directly for the TMA strides dispatch; mirror the WMXFP8AMXFP8 remap that runMoe() performs so the profiler hands real per-expert global scales instead of null pointers. AI-assisted (Claude).
|
Confirming the split from the GB10 side: it matches. Our #3568 only touches the dense MXFP8 GEMM path ( Diff read-through from our side: the One non-blocking question: the profiler's dummy activation SF content is still random bits — valid allocation now, but garbage E8M0 exponents. For pure timing that should be fine (the TMA descriptor is valid; decoding garbage scales can't fault), but worth confirming it's intentional vs. filling a neutral scale (all-ones exponent), in case any tactic's runtime is data-dependent on scale magnitude. Running the definitive e2e check now on the GB10: #3463 + #3568 + this PR, vLLM |
|
GB10 (SM121) validation results — PASS. Setup:
This was the kernel-level repro of the original #3558 crash on the affected hardware. We'll fold the fix into our next production vLLM image rebuild ( |
|
Addendum — full in-vLLM data from the GB10 (SM121), autotune now enabled end-to-end. Beyond the kernel-level validation above, we ran the complete vLLM A/B (gemma4-26B-A4B MXFP8 W8A8, 1. The crash fix holds in vLLM. Boot-time autotuning runs to completion, serves correct output. No IMA, no asserts — the original #3558 repro is gone at the serving layer too. 2. But the tactics the autotuner selects regress high-batch throughput. Concurrency sweep (200 tok/req, ignore_eos), same image/flags, only the autotune flag differs:
Autotune ON helps bs=1 (+22%) but selects tactics that are up to 40% slower at high concurrency. (With MTP speculative decoding at the bs=1 serving regime the picture is the same direction: 96.9 tok/s autotuned vs 75.1 with default tactics — a +29% win there.) This pattern is consistent with the dummy-input question above: if the profiler's dummy activations/SFs aren't representative (random E8M0 exponents, and/or the M-bucketing used for tuning), the measured timings can rank tactics differently than real traffic does, and the selection generalizes badly to large M. Happy to re-run the sweep against any revision — the harness is a one-command run on this box. (For completeness: this doesn't change our production choice — MARLIN W8A16 still wins our interactive MTP regime 102.6 vs 96.9 tok/s on identical same-day config — but it makes the autotuner usable, which is the point of this PR. The tactic-quality question is a separate, smaller follow-up.) |
The profiler sizes the weight/activation block-SF workspaces but never fills them, so the timed kernels read uninitialized E8M0 exponents — a nondeterministic scale pattern that can skew tactic rankings vs real traffic (observed as high-batch tactic regressions in GB10 validation). Fill all three SF regions with biased exponent 0x7F (= 1.0) on the MXFP8 path. AI-assisted (Claude).
|
@tgmerritt Outstanding validation — thank you. The A/B table is exactly the data this needed, and your read of the mechanism matches mine. Three responses: 1. One concrete cause of unrepresentative dummies is now fixed in 054ec99. You guessed right with "random E8M0 exponents": the profiler sized the SF workspaces (this PR's earlier commits) but never filled them — the timed kernels were reading uninitialized exponent bytes. That's worse than random: it's nondeterministic across allocations, so two boots can rank tactics differently. The new commit memsets all three SF regions (fc1/fc2 weight block-SFs + the activation SF buffer) to biased exponent 2. The remaining suspects are profiler-structural and probably the bigger half: (a) routing distribution — the profiler benchmarks under NUM_ROUTING_SAMPLES=16 uniform-random expert assignments, while real gemma4 traffic is skewed; tile-config rankings at high M are sensitive to per-expert token counts; (b) M-bucketing on the vLLM side — which bucket M values boot-time tuning ran at, and how runtime M maps onto them, is outside this PR entirely. Both deserve a dedicated issue rather than riding this crash fix — I'll file it with your table as the evidence base once we see the post-054ec99 numbers, unless you'd rather file from the GB10 side with the harness details. 3. Agreed on scope: this PR's job is "autotune doesn't crash and doesn't lie nondeterministically"; tactic quality is the follow-up. Your bs=1 (+22%) and MTP (+29%) wins suggest the autotuner is already worth having at interactive regimes even before the quality work. |
|
Re-ran the sweep on 054ec99 (GB10/SM121): the SF-init fix compiles, serves correctly — but does not move the high-concurrency regression. Same image/flags/harness as the table above, only the kernel header updated to 054ec99 (incremental JIT rebuild, 49s):
All deltas vs pre-054 are within run-to-run noise (the C=16 +9% is the largest and that band is noisy on this box). Correctness intact ("Paris" / "42" smoke on chat completions). So: the uninitialized-SF read was real and worth fixing for determinism alone, but SF content is not what's mispricing the tactics — the bs=1-good/high-M-bad selection survives a clean 1.0-exponent fill. That leaves your structural suspects holding the bag: the uniform-random routing distribution (NUM_ROUTING_SAMPLES=16) and the vLLM-side M-bucketing. Both consistent with what we see: the selected tactics behave like they were ranked at small effective per-expert M. Please go ahead and file the follow-up issue from your side — you have the mechanism scoped better than we do. Feel free to embed both tables; harness details for the repro: gemma4-26B-A4B MXFP8 W8A8 checkpoint, vLLM dev50-era build + #3463 + #3568 + this PR, |
|
@tgmerritt Thanks for the fast re-run — clean experimental result: SF content ruled out, determinism fix kept. Follow-up filed as #3622 with both your tables, the per-expert-M arithmetic behind your "ranked at small effective per-expert M" read (uniform routing with E=128/top_k=8 gives every expert ≈ M/16 — balanced and small, exactly the regime your selected tactics look tuned for), and the M-bucketing half flagged as possibly needing a vLLM-side companion. Please paste the harness script there when you get a chance — your GB10 same-day loop + this SM120 box gives the follow-up the same validation setup that worked here. |
| use_mxfp8_act_scaling=True, | ||
| output=flash_output, | ||
| ) | ||
| with autotune(True) if use_autotune else nullcontext(): |
There was a problem hiding this comment.
skip_ops is also available for this purpose
flashinfer/flashinfer/autotuner.py
Line 523 in af979cc
cc @qiching
There was a problem hiding this comment.
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.
|
/bot run tests/moe |
|
[FAILED] Pipeline #57191307: 12/20 passed |
5 participants
Description
Fixes the original-report half of #3558: the autotuner tuning pass crashes for
cutlass_fused_moewith MXFP8 activation scaling (CUDA illegal instruction / null TMA SF buffer on the reporter's SM121 build;!use_block_scalingkernel assert on current main).Root cause (one level deeper than my scoping comment on the issue —
input_sfis not even a tuning input on this path; the bug is entirely in the C++ gemm profiler): MXFP8×MXFP8 stores activations and weights as plain FP8, so the profiler cannot tell it apart from per-tensor FP8 by dtypes and prepared wrong tuning state in three coordinated places:GemmProfilerBackend::initderivedmScalingTypefrom dtypes alone →NONE→getOffsetActivationSFsized the activation-SF workspace to zero → the TMA warp-specialized MXFP8 kernels were launched with a null SF pointer ("gmem address 0 / null TMA buffer" in the report).getProfilerWorkspacesreserved only per-tensor float scalars for FP8 weights; MXFP8 needs per-expert weight block SFs + global scales.prepareQuantParamsfell into per-tensorQuantParams::FP8; MXFP8 needsQuantParams::MXFP8MXFP8.Plus tactic hygiene at the source:
CutlassMoeFCRunner::getTacticsnow filters out non-TMA (SM80/SM89-style) fallback configs for MXFP8 instantiations — those paths hard-assert!use_block_scaling, so offering them as tuning tactics guarantees a crash on any arch where they are enumerated. This composes with future SM12x enablement (#3463-style): onceisValidSM120MOESpecialisationadmits fp8×fp8, the SM120 TMA configs flow through this filter untouched.The binding passes
mUseMxfp8ActScalingintoGemmProfilerBackend::init(new defaulted parameter — other call sites unaffected).Validation
This box is SM120 (RTX PRO 6000), where MXFP8×MXFP8 is not an enabled arch on main (
test_moe_mxfp8_mxfp8is SM100-gated), so the honest validation matrix is:cutlass_fused_moeFP8 per-tensor withautotune(True)runs the full tuning pass and produces bit-identical outputs across {stock main, this PR} × {autotune on, off} (cos identical to 4+ digits on a 128×1024, 8-expert config).test_moe_fp8passes.No valid tactics available for fused moe op...instead of the deepAssertion failed: !use_block_scaling (cutlass_fused_moe_kernels.cuh:3153)mid-tuning.use_autotuneparametrization to the SM100-gatedtest_moe_mxfp8_mxfp8— on a pre-fix tree theTruecases crash the tuning pass; with this PR they should pass on SM100 CI.@tgmerritt — if you can ride this onto your #3463-based branch on the GB10, your original repro (vLLM
--moe-backend flashinfer_cutlass, autotune enabled) is the definitive e2e check for the IMA half. ThekErrorInternaltactic-init failures from your follow-up remain #3568's territory.Tests
use_autotuneparametrization ontest_moe_mxfp8_mxfp8)test_moe_fp8on SM120; MXFP8 cases are SM100-gated)Disclosure
AI-assisted (Claude); reviewed, compiled, and validated by the submitter.
Summary by CodeRabbit
New Features
Tests