Gemma 4 MoE: re-enable fused com.microsoft::MoE op - #326
Merged
Conversation
ORT main now plumbs the SwiGLU schema attributes (`swiglu_fusion`, `activation_alpha`, `activation_beta`, `swiglu_limit`) through to the kernel via microsoft/onnxruntime#28467 (QMoE CUDA EP + MoE GEMM Refactor), so the fused MoE op now correctly implements standard SwiGLU (`y = silu(gate) * up`) rather than only GPT-OSS-style SwiGLU. This was the original blocker tracked in microsoft/onnxruntime-genai#2062. Switch Gemma 4's MoE block back to `com.microsoft::MoE` when the EP advertises `supports_fused_moe`, with the explicit attribute set required by standard SwiGLU: activation_type = 'swiglu' activation_alpha = 1.0 (no GPT-OSS 1.702 multiplier) activation_beta = 0.0 (no GPT-OSS "+1" bias on the up branch) swiglu_limit = inf (no clipping) swiglu_fusion = 1 (interleaved) normalize_routing_weights = 1 k = top_k The CPU MoE kernel still only supports interleaved layout (`contrib_ops/cpu/moe/moe_cpu.cc:27`), and the new CUDA kernel accepts either, so emit interleaved (`swiglu_fusion=1`) for maximum portability. HuggingFace stores `experts.gate_up_proj` chunked as `[E, 2*inter, H]` (first `inter` rows = gate, next `inter` = up). Convert at graph-emit time via Reshape→Transpose→Reshape on the initializer; ORT folds the chain to a single static tensor at session load. The static-unroll `_dispatch_moe_fallback` is kept verbatim for EPs that don't expose the fused op. Validation on H200 with ORT 1.27.0.dev20260511001 (which contains #28467): * fp16 build of google/gemma-4-26b-a4b-it: 30 MoE nodes emitted, all with the expected attribute set. * `InferenceSession` on CUDAExecutionProvider loads in 12.3s (vs 959s with the previous fully-unrolled fallback, a ~78x speedup on session creation alone). * Prefill (B=1, S=4) runs in 0.65s; logits are well-behaved (no NaN or Inf, top-k token IDs land in the valid Gemma 4 vocab range). All 15 `gemma4` graph-construction tests pass, lintrunner clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Re-enables the fused com.microsoft::MoE op for the Gemma 4 MoE block (gated by ep_capabilities().supports_fused_moe), reverting the previous always-fallback workaround now that ORT main (post microsoft/onnxruntime#28467) supports standard SwiGLU semantics. The static-unroll path is retained for EPs without fused MoE support.
Changes:
- Cache
moe_intermediate_sizeandhidden_sizeon the block for reshape constants. - Emit
com.microsoft::MoEwith standard-SwiGLU attributes (alpha=1.0,beta=0.0,swiglu_limit=inf,swiglu_fusion=1) and a constant-foldableReshape→Transpose→Reshapeto interleavefc1_experts_weights. - Wrap the fused op in
CastLiketo restore input dtype and keep the legacy fallback underelse.
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-enables the fused
com.microsoft::MoEop for Gemma 4 now that ORT main supports standard SwiGLU. Replaces the static-unroll fallback that was shipped in #324 as a workaround.Why now
microsoft/onnxruntime#28467 (QMoE CUDA EP + MoE GEMM Refactor, merged on main) plumbs the existing schema attributes (
swiglu_fusion,activation_alpha,activation_beta,swiglu_limit) all the way through to the kernel. Before that PR, the CUDA MoE kernel hardcoded the GPT-OSS values (alpha=1.702, beta=1.0, limit=7.0, interleaved) and silently produced wrong output for any model using standard SwiGLU. That was the root cause tracked in microsoft/onnxruntime-genai#2062.Attribute set
Weight layout
HuggingFace stores
experts.gate_up_projchunked as[E, 2*inter, H](firstinterrows are gate, nextinterare up). The CPU MoE kernel still only accepts interleaved layout (contrib_ops/cpu/moe/moe_cpu.cc:27); the new CUDA kernel accepts either. We emitswiglu_fusion=1(interleaved) for maximum portability and reshape at graph-emit time viaReshape → Transpose → Reshapeon the initializer — ORT folds the chain to a single static tensor at session load.Fallback
_dispatch_moe_fallback(static per-expert unroll) is kept verbatim for EPs that don't advertisesupports_fused_moe. The fused path is only taken whenep_capabilities().supports_fused_moeis true (default for the CUDA / DML / default EPs).Validation (H200, ORT
1.27.0.dev20260511001which contains #28467)gemma4graph-construction testsgoogle/gemma-4-26b-a4b-itInferenceSessionload on CUDA EPCompatibility