fix(moe): prepare MXFP8 MXFP4 profiler inputs

[ormandj]

📌 Description

GemmProfilerBackend::init() receives use_mxfp8_act_scaling but does not retain it, so prepareQuantParams() always constructs QuantParams::FP8MXFP4 on the FP8-activation/MXFP4-weight path. That leaves the mxfp8_mxfp4 weight block-scale pointers null, so the MXFPX scale-factor descriptor setup is skipped during profiling.

VictoriaLogs captured the affected DeepSeek-V4-Flash-0731 TP2 startup reaching fused-MoE profiling and then exiting after this error surfaced:

[AutoTuner]: Tuning trtllm::fused_moe::gemm1
...
torch.AcceleratorError: CUDA error: an illegal instruction was encountered
...
[2026-07-31 15:42:01] Received sigquit from a child process. It usually means the child failed.

This change:

• stores use_mxfp8_act_scaling in the profiler;
• sizes the fc1 weight scale-factor workspace for the doubled gated-activation output;
• initializes the profiler's fc1/fc2 weight block scale-factor workspaces to UE8M0 unity on the MXFP8-activation path;
• constructs QuantParams::MXFP8MXFP4 for MXFP8-activation/MXFP4-weight profiling; and
• runs the existing numerical test with autotuning disabled and enabled.

When use_mxfp8_act_scaling is false, the existing QuantParams::FP8MXFP4 branch remains in use. The shared fc1 scale-factor workspace sizing is corrected for the gated output in both branches. PR #4066 skips profiling and selects fallback tactics for MXFP8×MXFP4 on SM120/SM121; this change repairs the profiler inputs and retains tactic profiling. Both changes update the same regression test.

If this change is accepted, PR #4066's SM120/SM121 fallback-tactic override should be removed, or #4066 closed. That override keeps SM120/SM121 on fallback tactics for this mode, which makes this profiler fix inert there.

🔍 Related Issues

Fixes #4049.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I used uvx pre-commit.
• I installed the hooks with pre-commit install. Not used; hooks were run directly with uvx.
• Focused pre-commit checks pass for all three changed paths: csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh, csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h, and tests/moe/test_trtllm_cutlass_fused_moe.py.

🧪 Tests

• Tests have been added or updated as needed.
• The complete repository test suite was not run.
• uv run --no-project python -m pytest -q tests/moe/test_trtllm_cutlass_fused_moe.py -k mxfp8_mxfp4 on SM120: 12 passed, 6 skipped, 132 deselected.
• With this change applied, a DeepSeek-V4-Flash-0731 TP2 startup completed all 10 gemm1 and 10 gemm2 profiler candidates on both ranks, completed FlashInfer autotuning, captured the target and draft CUDA graphs, and reached ready state with zero restarts.

Test limitations: the autotuned numerical test exercises this path and checks numerical parity, but it does not directly inspect the fabricated workspace contents. The non-MXFP8 activation branch touched by the shared fc1 scale-factor sizing has no focused autotune regression case in this validation.

Reviewer Notes

The table below compares heuristic fallback tactics with profiler-selected tactics. It is not a direct patched-versus-unpatched throughput comparison; the patch is what allows the profiled-tactic side to start successfully. Both sides used the same SGLang revision, model, TP2 configuration, FP8 KV cache, DSpark width 5, hardware, and workload. The control skipped trtllm::fused_moe::gemm1 and gemm2; the candidate profiled them.

Methodology: each decode cell is one unreplicated 30-second run at C1, C2, C4, C8, C16, and C32. The coding row is the median of five sequential requests with a 2,000-token cap. The prefill rows are the median client-observed throughput from standalone cold-prefill runs at exact 8K, 64K, and 128K prompts with 7, 2, and 1 samples, respectively. All rows are single comparison runs with no variance estimate, and no effect size is claimed from any row.

Cell	Fallback tactics	Profiled tactics
C1 decode (tok/s)	187.7	240.5
C2 decode (tok/s)	241.4	349.5
C4 decode (tok/s)	347.1	515.2
C8 decode (tok/s)	456.5	662.9
C16 decode (tok/s)	683.8	992.8
C32 decode (tok/s)	1030.7	1413.6
Coding median (tok/s)	218.1	264.8
8K prefill (tok/s)	7444	7638
64K prefill (tok/s)	8557	8351
128K prefill (tok/s)	7922	7744

The pinned 1,319-row GSM8K gate was not counted as passing because each run had one response finish at the 1,024-token cap. Two profiled-tactic runs scored 1,241 and 1,239 correct; the fallback-tactic control scored 1,240 correct. No quality difference is claimed from these runs.

AI was used to assist with implementation.

Summary by CodeRabbit

• Bug Fixes

  • Improved profiling for fused mixture-of-experts workloads using MXFP8 activations and MXFP4 weights.
  • Corrected workspace sizing and scaling behavior for quantized weight formats.
  • Added validation for unsupported scaling configurations.

• Tests

  • Expanded coverage to verify both autotuned and default execution paths for quantized MoE workloads.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 77d1531d-1788-4df8-a927-76b41de180ea

📥 Commits

Reviewing files that changed from the base of the PR and between 2e97a60 and 4ec7f23.

📒 Files selected for processing (3)

• csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh
• csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h
• tests/moe/test_trtllm_cutlass_fused_moe.py

---

📝 Walkthrough

Walkthrough

The profiler now preserves MXFP8 activation-scaling state, sizes gated FC1 scale storage correctly, and selects MXFP8×MXFP4 parameters with initialized scales. Tests cover both default and autotuned fused MoE execution.

Changes

MXFP8/MXFP4 fused MoE profiling

Layer / File(s)	Summary
Retain profiler scaling state csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h	GemmProfilerBackend stores the MXFP8 activation-scaling option for later profiler configuration.
Correct profiler sizing and quantization csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh	Profiler workspace sizing uses the gated FC1 output width. MXFP8 activation scaling selects MXFP8MXFP4 parameters and initializes weight block scales to E8M0 unity under OSS Cutlass.
Validate autotuned execution tests/moe/test_trtllm_cutlass_fused_moe.py	The MXFP8×MXFP4 test covers autotuning disabled and enabled while retaining the existing inputs and correctness comparison.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related issues

• flashinfer-ai/flashinfer#4003 — Both changes adjust native FP4 profiler workspace allocation for MoE autotuning.
• flashinfer-ai/flashinfer#3558 — Both changes handle MXFP8 activation scaling during Cutlass fused-MoE profiling and test autotuned execution.

Possibly related PRs

• flashinfer-ai/flashinfer#4130 — Both changes address gated MoE GEMM sizing and MXFP8/MXFP4 execution paths.
• flashinfer-ai/flashinfer#4159 — Both changes modify TRTLLM MXFP4/MXFP8 profiling and autotune coverage.

Suggested labels: run-ci

Suggested reviewers: aleozlx, samuellees, zianglih

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the primary change: preparing MXFP8-activation and MXFP4-weight profiler inputs.
Description check	✅ Passed	The description includes the required sections, explains the fix, links issue #4049, documents tests, and notes validation limits.
Linked Issues check	✅ Passed	The changes address issue #4049 by repairing MXFP8×MXFP4 profiler inputs and preserving tactic profiling without skipping the affected path.
Out of Scope Changes check	✅ Passed	The code and test changes are directly related to the profiler fix and its regression coverage; no unrelated changes are evident.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[ormandj]

@waynehacking8 @nv-yunzheq @jimmyzho, could you review this as the profiler-repair alternative to #4066? This initializes the MXFP8×MXFP4 profiler inputs and retains tactic profiling; #4066 selects fallback tactics for this mode on SM120/SM121.

[bkryu]

/bot run tests/moe

[flashinfer-bot]

GitLab MR !1178 has been created, and the CI pipeline #62008314 is currently running. I'll report back once the pipeline job completes.