Fix QMoE SM80 FP4 init order: don't let the tactic seed throw

[Gopalakrishnan Nallasamy (GopalakrishnanN)]

Problem

The SM80 FP4 grouped-GEMM path is unreachable on Ampere. Any MXFP4 MoE model (e.g. openai/gpt-oss-20b) throws at session initialization:

TMA WS grouped MoE GEMM for SM80 is not compiled, and this QMoE configuration has no SM80 fallback
  onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h:687

Root cause — an ordering problem

CutlassMoeFCRunner's constructor seeds a default GEMM tactic:

// moe_kernels.cu
auto tactics = getTactics(sm_);

That resolves to the static overload getConfigs(int sm, bool use_sm80_fp4 = false) — the flag defaults to false.

On SM80 with wfp4a16, that provisional view is wrong. getTmaWarpSpecializedConfigs() guards its early exit on the flag:

if constexpr (use_wfp4a16) {
  if (moeUseSm80Fp4(sm, use_sm80_fp4)) {
    return {};          // taken when the flag is TRUE -- no throw
  }
}
if (!isTmaWarpSpecializedGroupedGemmCompiledForSm(config_sm)) {
  if constexpr (use_w4afp8 || use_wfp4a16 || ...) {
    ORT_THROW("TMA WS grouped MoE GEMM for SM%d is not compiled, ...");
  }
}

With the flag false the early return {} is skipped and the throw fires. It fires during construction, before QMoE can push the real decision:

// moe_quantization.cc
m_moe_runner = std::make_unique<CutlassMoeFCRunner<half, __nv_fp4_e2m1, half>>(...);  // L428 - throws here
...
m_moe_runner->setUseSm80Fp4(enable_fp4_sm80_gemm_);                                   // L449 - never reached

So the decision made at L449 can never take effect, and users must fall back to the dequant path via ORT_FP4_SM80_GEMM=0.

Fix

The constructor's query is only a best-effort default seed — the surrounding code already tolerates an empty result via if (!tactics.empty()). Throwing there is the defect, so the seed is now non-fatal.

setUseSm80Fp4() re-queries immediately afterward and re-picks. With use_sm80_fp4=true, getTmaWarpSpecializedConfigs() returns {} by design and getAmpereConfigs() supplies real candidate configs, so both gemm1_config_ and gemm2_config_ get valid tactics. A genuinely unsupported configuration still throws later at dispatch.

I kept this minimal rather than plumbing the flag through the constructor signature, which would be tidier but touches every CutlassMoeFCRunner construction site.

Validation

A100-80GB (SM80), openai/gpt-oss-20b MXFP4 exported for paged attention, run through ONNX Runtime GenAI's continuous-batching engine.

	Before (ORT_FP4_SM80_GEMM=0 fallback)	After (SM80 path)
Session init	throws on default path	succeeds
Peak GPU memory	~49,800 MiB	24,149 MiB
Load time	~5.9 s	4.56 s
Output parity	exact	exact

2.1x memory reduction, with exact output parity between concurrent and isolated runs (0 mismatches) and stable output across repeats.

Note for reviewers

I validated this empirically as above, but I have not run ONNX Runtime's own MoE/QMoE test suites against the change — worth confirming in CI. I'd also welcome a view on whether the non-fatal seed is the preferred shape here versus passing the SM80 decision into the constructor.

[Copilot]

Pull request overview

Prevents provisional tactic selection from blocking SM80 FP4 QMoE initialization.

Changes:

• Makes constructor tactic seeding non-fatal.
• Allows setUseSm80Fp4() to select valid Ampere tactics afterward.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

This shared constructor is also instantiated by the standard MoE path and by QMoE's dense fallback runners; those callers never invoke setUseSm80Fp4(). The comment's "immediately ... re-picks" invariant therefore holds only for the primary WFP4A16 QMoE runner. That makes it important to scope the bypass to this provisional WFP4A16/SM80 query (or pass the decision into construction) instead of suppressing failures in shared initialization.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

Suppressed comments (2)

onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu:2122

• Please add automated SM80 session-initialization coverage for this path. The existing TestQMoEFP4Sm80SingleWeightCopy._skip_unless_sm80_regime() in onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py:946-951 skips all devices below SM90, so it skips A100/SM80 and would not detect the constructor throw this change is intended to fix.

  try {
    auto tactics = getTactics(sm_);

onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu:2127

• This suppresses every failure from tactic generation, not just the expected provisional SM80/WFP4 mismatch. getConfigs() allocates and copies vectors and can also throw for genuinely unsupported configurations, so an OOM or unrelated configuration defect now leaves both tactic optionals unset and converts a precise initialization failure into a delayed failure. Avoid broad exception swallowing; pass the SM80-FP4 decision into construction (or otherwise bypass only this known WFP4/Ampere seed) so all unrelated exceptions still propagate.

  } catch (const std::exception&) {

[Tianlei Wu (tianleiwu)]

One additional maintainability issue is noted inline. The two existing open threads already cover the broader correctness and regression-test concerns, so I have not duplicated them here.

[Tianlei Wu (tianleiwu)]

Nit: please include <exception> directly now that this translation unit names std::exception. It is currently available only through core/common/common.h -> core/common/exceptions.h, so a change to that internal include chain could break this file unexpectedly.