[Ti-Tai Wang (titaiwangms)]

Summary

Expose the existing internal GEGLU activation through the public CUDA MoE and integer QMoE operators, with explicit tanh-GELU semantics and fused gate/up layouts.

The current upstream main has an internal ActivationType::Geglu, but MoEBase rejects activation_type="geglu" and restricts nonzero swiglu_fusion to SwiGLU. This prevents an ONNX caller from using the existing gated implementation for this activation.

Changes

• Accept activation_type="geglu" with explicit fused FC1 layouts: mode 1 is interleaved [gate0, up0, ...], mode 2 is chunked [gate | up]. Reject separate FC3 inputs and unfused mode 0.
• Select CUTLASS GELU_taylor for GELU(gate, approximate="tanh") * up in both gated activation dispatch paths. Keep ordinary GELU and existing SwiGLU behavior unchanged.
• Correct the chunked GEGLU gate/up and bias offsets while preserving the legacy SwiGLU ordering.
• Treat GEGLU as gated for MoE/QMoE shape validation and integer FC1 geometry, but do not route it through fast paths that hardcode SwiGLU.
• Restrict QMoE GEGLU to integer quantization; explicitly reject FP4/NVFP4/FP8/WFP4AFP8 rather than silently selecting the wrong activation.
• Add independent numerical and invalid-layout tests, and document the CUDA-only contract and precision limitations.

The port is based on 3b4a0fff56c9b89f644acfc6613ae4dad3aeb5ff. It preserves the newer QMoE INT3 schema factory and its type/shape inference implementation; schema edits are documentation-only.

Validation

• Current-main port: scoped lintrunner/clang-format and git diff --check passed.
• Earlier patched build based on 942356fce6bdc78c016e4b902e49492fdf695453: 12 selected CUDA provider tests executed, 12 passed, 0 skipped on A100/SM80.
• Copilot follow-up regression harness: freshly compiled the updated GEGLU test translation unit and linked it against that immutable earlier custom runtime. On A100/SM80, MoETest.CudaGeglu* matched 12 tests: 11 passed, 1 explicitly skipped (Hopper TMA kernels are not compiled in that build), 0 failed. This is not a rebuild or execution of the current PR head.
• The GEGLU numerical tests cover float32/float16/bfloat16 MoE and float16/bfloat16 INT4 QMoE, fusion modes 1/2, token counts 1/17, and top-k 1/2. An invalid fusion-0 case checks explicit rejection.
• The follow-up adds distinct, nonzero per-expert/per-channel gate and up biases, plus raw uniform INT2 and mixed INT2/INT4 cases in both width directions, for both layouts. The integer cases preserve canonical raw tensors and disable prepacking to exercise dense-dequantization fallback.
• A separate test-only mutation swapped serialized gate/up biases while leaving the reference unchanged. The FP32 bias test failed as expected, demonstrating that the fixture detects this layout error; no production runtime artifacts were modified.
• The reference is computed independently with std::tanh, with inputs that distinguish it from erf-GELU. FP32 uses absolute tolerance 3e-5, relative tolerance zero, and asserts erf/tanh reference separation greater than ten times that tolerance. FP16/BF16 have explicit dtype-specific tolerances.
• CPU fallback is disabled in the CUDA operator tests. The BF16 INT4 case is guarded for the relevant non-quick-build configuration.

Validation boundary: the clean current-main port has not been rebuilt or executed locally. Neither earlier-runtime result is presented as execution of this PR head. A dedicated SM90 bias test pins the first TMA tactic when Hopper kernels and an SM90 GPU are available, but SM90/TMA execution, additional GPU architectures, and a fresh full-model qualification remain unverified.

Accuracy and performance limitations

These changes enable a previously unreachable public operator path; they are not a claim that the complete Gemma4 model is numerically qualified or faster.

• Fused activation/product computation uses float before the final output conversion. Eager PyTorch can materialize GELU in FP16/BF16 before multiplication, so bitwise equivalence is not promised.
• Specialized SwiGLU-only GEMV/prefill implementations remain excluded. This PR adds no GEGLU-specific performance kernel and makes no throughput claim.
• Existing full-model HF logit/KV parity failures remain failed and unwaived. MMLU-Pro generation scores do not clear those gates.
• The observed cross-runtime accuracy/performance differences have not been attributed to this patch. Matching-input activation/kernel comparisons and CUDA profiling are needed to distinguish kernel/rounding effects from quantization, attention/export differences, and generation truncation.

This PR does not change CPU or other execution providers, rebuild llama.cpp/GenAI, or change model weights/quantization recipes.

[Copilot]

🟡 Changes recommended

Generated schema documentation is stale, and the changed bias indexing and integer fallback geometry lack regression coverage.

3 open issues

• Add biased GEGLU regression coverage for fusion modes and TMA · New
• Test QMoE GEGLU sizing with INT2 and mixed-width weights · New
• Regenerate ContribOperators.md with GEGLU schema contracts · New

What changed in this PR

This PR exposes tanh-based GEGLU through CUDA MoE and integer QMoE, reusing existing gated activation infrastructure while preserving ordinary GELU and SwiGLU behavior.

Changes:

• Accept GEGLU with interleaved or chunked fused FC1 layouts and reject unsupported configurations.
• Update activation dispatch, gate/up offsets, and integer FC1 geometry.
• Add independent numerical tests and document CUDA support and precision limitations.

File	Description
onnxruntime/​test/​contrib_ops/​moe_test.cc	Adds GEGLU numerical and invalid-layout tests.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Documents GEGLU attributes and fused layouts in operator schemas.
onnxruntime/​contrib_ops/​cuda/​moe/​moe.cc	Validates fused GEGLU inputs and gated shapes.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Restricts GEGLU to integer quantization and adjusts FC1 geometry.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_base.h	Accepts GEGLU and requires an explicit fusion mode.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_activation_kernels.cuh	Selects tanh-GELU and corrects chunked gate/up and bias offsets.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Explains GEGLU layouts, supported paths, and precision limitations.

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

I traced the three GEMM_1 routes for GEGLU (small-M GEMV, grouped GEMM + doGatedActivation, Hopper TMA + doActivation) for FP16/BF16/FP32 and INT4/INT8/INT2/mixed weights, plus validation, docs and tests. Semantics are consistent: GELU_taylor(gate) * up, swiglu_fusion 1 = interleaved, 2 = chunked [gate | up], alpha/beta/limit stay SwiGLU-only, legacy SwiGLU behaviour is untouched, and unsupported combinations (FC3, FP4/NVFP4/FP8/WFP4AFP8, INT3, CPU/WebGPU) are rejected up front. I found no High-Priority issue. The three open threads are addressed at this head and are not repeated.

Inline Suggestion: in CI the FP16 + INT4 GEMM tactic profiler finds no valid tactic for GEGLU (eight 'No valid GEMM config found' warnings, only in that test), so an unprofiled default tactic is used.

Test note: recomputing the test's formulas, the gap between tanh-GELU and erf-GELU outputs exceeds the allowed error (0.001 + 0.001*|expected|) only in the FP16 rows=17 / top_k=1 cases (about 1.2-1.7x; the other FP16 cases are at 0.3-0.85x), and it is under 0.1x of the allowed error in every BF16 case. So only the FP32 tests (which assert a 10x separation) pin the tanh approximation, and the sm_90 TMA test (FP16/BF16) would not reliably catch the Geglu case in doActivation being switched to the erf GELU functor; I checked that mapping by inspection only (moe_gemm_activation_kernels.cuh L396-397).

Recommend approving after verification on the impacted devices: sm_90 (H100/H200) with TMA grouped GEMMs compiled (MoETest.CudaGegluBiasTmaSm90 plus the GEGLU FP16/BF16 tests) and an sm_80 (A100) / sm_89 spot-check; sm_100/sm_120 TMA is unexercised.

Validation: CI: all checks green on 9274182; the A10 (sm_86) job ran 11 of 12 MoETest.CudaGeglu* tests, the TMA test was skipped. Local: not run (the A100 (sm_80) available for this review cannot execute the sm_90 TMA path, and CI had already built and run this head on sm_86, so I relied on that evidence plus a static review). Not verified: sm_90/TMA, Blackwell, sm_89, INT8 GEGLU, performance.

[Copilot]

🟢 Approval recommended

No active review issues remain for the added GEGLU support in CUDA MoE and integer QMoE.

No open issues 🎉

3 resolved since last review

• Test QMoE GEGLU sizing with INT2 and mixed-width weights
• Add biased GEGLU regression coverage for fusion modes and TMA
• Regenerate ContribOperators.md with GEGLU schema contracts

_{Review effort: Balanced}