diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index 6dcb758a234fe..2a37bcc10e985 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -4597,6 +4597,10 @@ This version of the operator has been available since version 1 of the 'com.micr
When swiglu_fusion=0, two GEMMs are not fused, and they are FC1 and FC3 in the inputs.
When swiglu_fusion=1, two GEMMs are fused so that g and l are computed in a single GEMM (FC1), and g and l are interleaved on each row of size 2 * inter_size.
When swiglu_fusion=2, two GEMMs are fused, and g and l are concatenated on each row.
+ CUDA also supports GEGLU: geglu = GELU(g, approximate='tanh') * l.
+ GEGLU requires fused FC1 projections: swiglu_fusion=1 stores [g0,l0,g1,l1,...],
+ and swiglu_fusion=2 stores [g | l]. Separate FC3 inputs are not supported.
+ activation_alpha, activation_beta and swiglu_limit do not affect GEGLU.
#### Version
@@ -4611,15 +4615,15 @@ This version of the operator has been available since version 1 of the 'com.micr
activation_beta : float
Beta parameter used in activation function.
activation_type : string
-Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu
+Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU) and geglu (GELU-tanh(gate) * up). GEGLU requires fused FC1 gate/up weights, swiglu_fusion=1 or 2, and no FC3 inputs. Default is relu
k : int
Number of top experts to select from expert pool
normalize_routing_weights : int
Whether to normalize routing weights
swiglu_fusion : int
-0: not fused, 1: fused and interleaved. 2: fused and not interleaved.
+Gate/up layout for swiglu or geglu. 0: not fused (swiglu only), 1: fused and interleaved, 2: fused and not interleaved ([gate | up] for geglu).
swiglu_limit : float
-The limit used to clamp in SwiGLU. No clamp when limit is not provided.
+The limit used to clamp in SwiGLU. No clamp when limit is not provided. Does not affect GEGLU.
use_sparse_mixer : int
Whether to use sparse mixer
@@ -4632,7 +4636,7 @@ This version of the operator has been available since version 1 of the 'com.micr
router_probs : T
2D input tensor with shape (total_tokens, num_experts), where total_tokens must match the flattened token count of input
fc1_experts_weights : T
-3D input tensor with shape (num_experts, fusion_size * inter_size, hidden_size), where fusion_size is 2 for fused swiglu, and 1 otherwise
+3D input tensor with shape (num_experts, fusion_size * inter_size, hidden_size), where fusion_size is 2 for fused swiglu or geglu, and 1 otherwise
fc1_experts_bias (optional) : T
2D optional input tensor with shape (num_experts, fusion_size * inter_size)
fc2_experts_weights : T
@@ -6452,6 +6456,10 @@ This version of the operator has been available since version 1 of the 'com.micr
When swiglu_fusion=0, two GEMMs are not fused, and they are FC1 and FC3 in the inputs.
When swiglu_fusion=1, two GEMMs are fused so that g and l are computed in a single GEMM (FC1), and g and l are interleaved on each row of size 2 * inter_size.
When swiglu_fusion=2, two GEMMs are fused, and g and l are concatenated on each row.
+ CUDA also supports GEGLU with quant_type='int': geglu = GELU(g, approximate='tanh') * l.
+ GEGLU requires fused FC1 projections: swiglu_fusion=1 stores [g0,l0,g1,l1,...],
+ and swiglu_fusion=2 stores [g | l]. Separate FC3 weights are not supported.
+ activation_alpha, activation_beta and swiglu_limit do not affect GEGLU.
The raw INT3 profile extends version 1; execution support is provider-specific.
For INT3, the requirements below supersede the legacy format descriptions above.
@@ -6483,7 +6491,7 @@ This version of the operator has been available since version 1 of the 'com.micr
activation_beta : float
Beta parameter used in activation function.
activation_type : string
-Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu
+Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU) and geglu (GELU-tanh(gate) * up). GEGLU requires quant_type='int', fused FC1 gate/up weights, swiglu_fusion=1 or 2, and no separate FC3 weights. Default is relu
block_size : int
For integer quantization, size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). For integer and FP4 quantization, both hidden_size and inter_size must be divisible by the block size. FP8 with block_size=128 supports partial 128x128 tiles. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. Without a positive value, FP8 uses per-expert global scales. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor.
expert_weight_bits : int
@@ -6501,9 +6509,9 @@ This version of the operator has been available since version 1 of the 'com.micr
quant_type : string
Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, 'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weight-only quantization, or 'wfp4afp8' for MXFP4 weight with FP8 activation. When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. 'fp4' uses Float8E8M0 block scales with block_size 32; 'nvfp4' uses Float8E4M3FN block scales with block_size 16. 'fp8' uses float8e4m3fn weights and, when block_size > 0, float32, float16, or bfloat16 square-block scales in fc*_scales instead of fc*_global_scale. The INT3 profile requires integer quantization.
swiglu_fusion : int
-0: not fused, 1: fused and interleaved. 2: fused and not interleaved.
+Gate/up layout for swiglu or geglu. 0: not fused (swiglu only), 1: fused and interleaved, 2: fused and not interleaved ([gate | up] for geglu).
swiglu_limit : float
-The limit used to clamp inputs in SwiGLU. It is infinite when limit is not provided.
+The limit used to clamp inputs in SwiGLU. It is infinite when limit is not provided. Does not affect GEGLU.
use_sparse_mixer : int
Whether to use sparse mixer
weights_prepacked : int
@@ -6518,7 +6526,7 @@ This version of the operator has been available since version 1 of the 'com.micr
router_probs : T
2D tensor with shape (total_tokens, num_experts), where total_tokens must match the flattened token count of input
fc1_experts_weights : T1
-3D tensor with shape (num_experts, fusion_size * inter_size, hidden_size * effective_fc1_bits / 8). The last dimension must be byte-aligned. The fusion_size is 2 for fused swiglu, or 1 otherwise. effective_fc1_bits is fc1_expert_weight_bits when provided, otherwise expert_weight_bits.
+3D tensor with shape (num_experts, fusion_size * inter_size, hidden_size * effective_fc1_bits / 8). The last dimension must be byte-aligned. The fusion_size is 2 for fused swiglu or geglu, or 1 otherwise. effective_fc1_bits is fc1_expert_weight_bits when provided, otherwise expert_weight_bits.
fc1_scales (optional) : T2
Optional weight scales. For quant_type='int', this is a 2D tensor with shape (num_experts, fusion_size * inter_size), or a 3D tensor with shape (num_experts, fusion_size * inter_size, hidden_size / block_size) when block_size is provided. For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape (num_experts, fusion_size * inter_size, hidden_size / 32). For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape (num_experts, fusion_size * inter_size, hidden_size / 16). For quant_type='fp8' and block_size > 0, required float32, float16, or bfloat16 scales with shape (num_experts, ceil(fusion_size * inter_size / block_size), ceil(hidden_size / block_size)).
fc1_experts_bias (optional) : T
diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md
index 089f28a05f48e..4b9961538913f 100644
--- a/docs/contrib_ops/cuda/moe_qmoe.md
+++ b/docs/contrib_ops/cuda/moe_qmoe.md
@@ -59,7 +59,7 @@ The execution pipeline is:
```
input tokens → router (top-k softmax) → permute by expert
- → GEMM1 (per-expert) → activation (SiLU/GeLU/ReLU/SwiGLU)
+ → GEMM1 (per-expert) → activation (SiLU/GeLU/ReLU/SwiGLU/GEGLU)
→ GEMM2 (per-expert) → un-permute → weighted sum
```
@@ -72,7 +72,7 @@ input tokens → router (top-k softmax) → permute by expert
| Attribute | Type | Default | Description |
|-----------|------|---------|-------------|
| `k` | int | 1 | Top-K experts selected per token. |
-| `activation_type` | string | `"relu"` | One of `"relu"`, `"gelu"`, `"silu"`, `"swiglu"`, `"identity"`. These are the only values accepted end-to-end (attribute parsing); other kernel-internal types are not reachable from ONNX. |
+| `activation_type` | string | `"relu"` | One of `"relu"`, `"relu2"`, `"gelu"`, `"silu"`, `"swiglu"`, `"geglu"`, `"identity"`. CUDA GEGLU computes `GELU(gate, approximate="tanh") * up`, requires fused FC1 and `swiglu_fusion=1` or `2`; QMoE additionally requires `quant_type="int"`. |
| `normalize_routing_weights` | int | 0 | Re-normalize the top-k weights to sum to 1. |
| `use_sparse_mixer` | int | 0 | Enable sparse-mixer routing variant. |
| `swiglu_fusion` | int | 0 | 0=no fusion, 1=interleaved (Gate/Value), 2=block (Gate;Value). See [§8](#8-swiglu-fusion). |
@@ -123,12 +123,13 @@ to the selected `quant_type` are simply omitted (most are `Optional`).
| 20 | `fc2_act_block_scale` | T2 (Opt, float8e8m0) | `(E, M_pad, inter/32)` | wfp4afp8 (Variant B) |
`E = num_experts`. `pack = 8 / expert_weight_bits` for INT/MXFP4 weights; `pack = 1`
-for FP8 weights. `fusion = 2` for `swiglu_fusion=1`, otherwise `1`.
+for FP8 weights. `fusion = 2` for fused SwiGLU or GEGLU with
+`swiglu_fusion=1` or `2`, otherwise `1`.
> **CUDA EP restriction for `quant_type="int"`:** `fc1_scales` and `fc2_scales`
> must use the same FP16 or BF16 type as the activation input. Although the
> operator schema also permits float scales, the CUDA INT kernel does not
-> currently cast float scales internally. `fusion = 2` for `swiglu_fusion=1`, otherwise `1`.
+> currently cast float scales internally.
For packed input, tokens from variable-length sequences may be concatenated
without padding. MoE routing is token-local, so no sequence-offset input is
@@ -636,6 +637,26 @@ easy to feed Mixtral-style models without offline fusion.
---
+### GEGLU
+
+CUDA `MoE` and integer `QMoE` also accept `activation_type="geglu"`.
+They evaluate `GELU(gate, approximate="tanh") * up`, not a scaled-SwiGLU
+approximation or erf-based GELU. The fused FC1 output has `2 * inter_size`
+columns: fusion mode 1 stores `[gate0, up0, gate1, up1, ...]`, and mode 2
+stores `[gate | up]`. GEGLU requires one of these explicit fusion modes and
+does not support separate FC3 weights. `activation_alpha`, `activation_beta`,
+and `swiglu_limit` apply to SwiGLU only.
+
+GEGLU uses the generic gated activation path; specialized kernels that
+hardcode SwiGLU remain excluded. Integer INT4/INT8 still use weight-only
+CUTLASS runners. GEGLU is rejected for FP4, NVFP4, FP8, and WFP4AFP8 rather
+than silently selecting a different activation. CPU and other execution
+providers do not gain GEGLU support from this CUDA implementation.
+
+The activation and product are computed in float before conversion to the
+activation output type. This is not a promise of bitwise parity with eager
+PyTorch code that materializes GELU in FP16/BF16 before multiplication.
+
## 9. FP4 (MXFP4) Details
The QMoE operator supports **MXFP4** quantized weights with FP16/BF16 activations
diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh
index db01d095c3e64..3e2393597b639 100644
--- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh
+++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_activation_kernels.cuh
@@ -131,7 +131,7 @@ void doGatedActivation(ActivationOutputType* output, const GemmOutputType* gemm_
case ActivationType::SwigluBias:
return &doGatedActivationKernel;
case ActivationType::Geglu:
- return &doGatedActivationKernel;
+ return &doGatedActivationKernel;
case ActivationType::Silu:
return &doGatedActivationKernel;
case ActivationType::Gelu:
@@ -226,6 +226,10 @@ __global__ void doActivationKernel(T* output, const GemmOutputType* gemm_result,
const int64_t num_elems_in_col = inter_size / ACTIVATION_ELEM_PER_THREAD;
assert(gated_off % ACTIVATION_ELEM_PER_THREAD == 0);
const int64_t gated_off_vec = gated_off / ACTIVATION_ELEM_PER_THREAD;
+ // Chunked GEGLU is [gate | up]; preserve the legacy SwiGLU chunk order.
+ const int64_t gate_offset =
+ gated && activation_params.activation_type == ActivationType::Geglu ? 0 : gated_off_vec;
+ const int64_t linear_offset = gated_off_vec - gate_offset;
ActFn fn{};
constexpr bool IsSiLu = std::is_same_v, cutlass::epilogue::thread::SiLu>;
@@ -244,11 +248,9 @@ __global__ void doActivationKernel(T* output, const GemmOutputType* gemm_result,
linear_part[i] = static_cast(scalar_gemm[2 * global_elem + 1]);
}
} else {
- // If not gated, gate_part reads from elem_index (gated_off_vec is 0).
- // If gated, gate_part reads from elem_index + inter_size_vec (chunk 1).
- gate_part = arrayConvert(gemm_result_vec[elem_index + gated_off_vec]);
+ gate_part = arrayConvert(gemm_result_vec[elem_index + gate_offset]);
if (gated) {
- linear_part = arrayConvert(gemm_result_vec[elem_index]);
+ linear_part = arrayConvert(gemm_result_vec[elem_index + linear_offset]);
}
}
@@ -261,9 +263,9 @@ __global__ void doActivationKernel(T* output, const GemmOutputType* gemm_result,
linear_part[i] += static_cast(scalar_bias[2 * global_elem + 1]);
}
} else {
- gate_part = gate_part + arrayConvert(bias_ptr_vec[elem_index + gated_off_vec]);
+ gate_part = gate_part + arrayConvert(bias_ptr_vec[elem_index + gate_offset]);
if (gated) {
- linear_part = linear_part + arrayConvert(bias_ptr_vec[elem_index]);
+ linear_part = linear_part + arrayConvert(bias_ptr_vec[elem_index + linear_offset]);
}
}
}
@@ -390,8 +392,9 @@ void doActivation(T* output, const GemmOutputType* gemm_result, const float* fp8
case ActivationType::Identity:
return &doActivationKernel;
case ActivationType::Gelu:
- case ActivationType::Geglu:
return &doActivationKernel;
+ case ActivationType::Geglu:
+ return &doActivationKernel;
case ActivationType::Relu:
return &doActivationKernel;
case ActivationType::Relu2:
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe.cc b/onnxruntime/contrib_ops/cuda/moe/moe.cc
index f102aba72971e..f11b5902df68c 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe.cc
+++ b/onnxruntime/contrib_ops/cuda/moe/moe.cc
@@ -72,9 +72,13 @@ Status MoE::ComputeInternal(OpKernelContext* context) const {
LogSwigluFusionRemapOnce();
}
- bool is_fused_swiglu = (activation_type_ == ActivationType::Swiglu) &&
- (swiglu_fusion != 0) &&
- (fc3_experts_weights_optional == nullptr);
+ ORT_RETURN_IF(activation_type_ == ActivationType::Geglu &&
+ (fc3_experts_weights_optional != nullptr || fc3_experts_bias_optional != nullptr),
+ "GEGLU requires gate/up projections fused into FC1, without FC3 inputs.");
+ bool is_fused_gated = (activation_type_ == ActivationType::Swiglu ||
+ activation_type_ == ActivationType::Geglu) &&
+ (swiglu_fusion != 0) &&
+ (fc3_experts_weights_optional == nullptr);
MoEParameters moe_params;
ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs(
@@ -83,7 +87,7 @@ Status MoE::ComputeInternal(OpKernelContext* context) const {
fc2_experts_weights, fc2_experts_bias_optional, nullptr, nullptr,
fc3_experts_weights_optional, fc3_experts_bias_optional, nullptr, nullptr,
1, // no quantization so pack size is 1
- is_fused_swiglu,
+ is_fused_gated,
0)); // no block-wise quantization for regular MoE
ORT_RETURN_IF_NOT(k_ > 0 && k_ <= moe_params.num_experts,
"MoE requires 0 < k <= num_experts, got k=", k_,
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_base.h b/onnxruntime/contrib_ops/cuda/moe/moe_base.h
index 3fca3b6153903..c7e809fa55f46 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe_base.h
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_base.h
@@ -57,6 +57,8 @@ class MoEBase {
activation_type_ = ActivationType::Silu;
} else if (activation_type_str == "swiglu") {
activation_type_ = ActivationType::Swiglu;
+ } else if (activation_type_str == "geglu") {
+ activation_type_ = ActivationType::Geglu;
} else if (activation_type_str == "identity") {
activation_type_ = ActivationType::Identity;
} else {
@@ -81,8 +83,11 @@ class MoEBase {
swiglu_fusion_ = static_cast(op_kernel_info.GetAttrOrDefault("swiglu_fusion", 0));
ORT_ENFORCE(swiglu_fusion_ >= 0 && swiglu_fusion_ <= 2,
"swiglu_fusion must be 0, 1, or 2, but got ", swiglu_fusion_);
- ORT_ENFORCE(activation_type_ == ActivationType::Swiglu || swiglu_fusion_ == 0,
- "swiglu_fusion is only valid when activation_type is 'swiglu'.");
+ ORT_ENFORCE(activation_type_ == ActivationType::Swiglu ||
+ activation_type_ == ActivationType::Geglu || swiglu_fusion_ == 0,
+ "swiglu_fusion is only valid when activation_type is 'swiglu' or 'geglu'.");
+ ORT_ENFORCE(activation_type_ != ActivationType::Geglu || swiglu_fusion_ != 0,
+ "GEGLU requires fused FC1 gate/up weights and swiglu_fusion=1 or 2.");
// SwiGLU limit for clamping (optional, use infinity if not provided)
swiglu_limit_ = op_kernel_info.GetAttrOrDefault("swiglu_limit", std::numeric_limits::infinity());
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
index 028761186bad7..6f405362694c1 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
@@ -381,6 +381,9 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE
block_size_ = op_kernel_info.GetAttrOrDefault("block_size", -1);
this->quant_type_ = op_kernel_info.GetAttrOrDefault("quant_type", "int");
+ ORT_ENFORCE(activation_type_ != onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Geglu ||
+ quant_type_ == "int",
+ "CUDA QMoE GEGLU currently supports only quant_type='int'.");
ORT_ENFORCE(quant_type_ == "int" || quant_type_ == "fp4" || quant_type_ == "nvfp4" ||
quant_type_ == "fp8" || quant_type_ == "wfp4afp8",
"quant_type must be 'int', 'fp4', 'nvfp4', 'fp8', or 'wfp4afp8', but got '", quant_type_, "'");
@@ -951,6 +954,9 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
}
bool is_fused_swiglu = activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Swiglu;
+ // GEGLU shares gated shapes, not the SwiGLU-only GEMV/prefill implementations.
+ const bool is_fused_gated = is_fused_swiglu ||
+ activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Geglu;
MoEParameters moe_params;
// Prefer the cached shapes when PrePack consumed the source initializer.
const TensorShape& fc1_shape = weights_consumed_by_prepack ? fc1_weights_shape_ : fc1_experts_weights->Shape();
@@ -963,7 +969,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
moe_helper::MoEWeightBits{fc1_expert_weight_bits_,
fc2_expert_weight_bits_,
fc3_expert_weight_bits_},
- is_fused_swiglu, block_size_));
+ is_fused_gated, block_size_));
if (is_block_fp8) {
const int64_t fc1_n = is_fused_swiglu ? 2 * moe_params.inter_size : moe_params.inter_size;
ORT_RETURN_IF_NOT(fc1_shape == TensorShape({moe_params.num_experts, fc1_n, moe_params.hidden_size}) &&
@@ -1000,7 +1006,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
"INT2 or mixed-width CUDA QMoE requires inter_size to be divisible by block_size, got inter_size=",
moe_params.inter_size, " and block_size=", block_size_, ".");
- const int64_t fc1_n = is_fused_swiglu ? moe_params.inter_size * 2 : moe_params.inter_size;
+ const int64_t fc1_n = is_fused_gated ? moe_params.inter_size * 2 : moe_params.inter_size;
const int64_t fc1_k = moe_params.hidden_size;
const int64_t fc2_n = moe_params.hidden_size;
const int64_t fc2_k = moe_params.inter_size;
@@ -2695,7 +2701,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
fc1_scales != nullptr && fc2_scales != nullptr,
"INT2 or mixed-width CUDA QMoE requires raw weights and scales for FC1 and FC2.");
- const int fc1_n = static_cast(is_fused_swiglu ? moe_params.inter_size * 2 : moe_params.inter_size);
+ const int fc1_n = static_cast(is_fused_gated ? moe_params.inter_size * 2 : moe_params.inter_size);
const int fc1_k = static_cast(moe_params.hidden_size);
const int fc2_n = static_cast(moe_params.hidden_size);
const int fc2_k = static_cast(moe_params.inter_size);
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index e05119bd05dfa..2e1a42e6c7c5f 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -1430,15 +1430,19 @@ constexpr const char* MoE_ver1_doc = R"DOC(
When swiglu_fusion=0, two GEMMs are not fused, and they are FC1 and FC3 in the inputs.
When swiglu_fusion=1, two GEMMs are fused so that g and l are computed in a single GEMM (FC1), and g and l are interleaved on each row of size 2 * inter_size.
When swiglu_fusion=2, two GEMMs are fused, and g and l are concatenated on each row.
+ CUDA also supports GEGLU: geglu = GELU(g, approximate='tanh') * l.
+ GEGLU requires fused FC1 projections: swiglu_fusion=1 stores [g0,l0,g1,l1,...],
+ and swiglu_fusion=2 stores [g | l]. Separate FC3 inputs are not supported.
+ activation_alpha, activation_beta and swiglu_limit do not affect GEGLU.
)DOC";
ONNX_MS_OPERATOR_SET_SCHEMA(
MoE, 1,
OpSchema()
.SetDoc(MoE_ver1_doc)
- .Attr("activation_type", "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu", AttributeProto::STRING, std::string("relu"))
- .Attr("swiglu_fusion", "0: not fused, 1: fused and interleaved. 2: fused and not interleaved.", AttributeProto::INT, static_cast(0))
- .Attr("swiglu_limit", "The limit used to clamp in SwiGLU. No clamp when limit is not provided.", AttributeProto::FLOAT, OPTIONAL_VALUE)
+ .Attr("activation_type", "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU) and geglu (GELU-tanh(gate) * up). GEGLU requires fused FC1 gate/up weights, swiglu_fusion=1 or 2, and no FC3 inputs. Default is relu", AttributeProto::STRING, std::string("relu"))
+ .Attr("swiglu_fusion", "Gate/up layout for swiglu or geglu. 0: not fused (swiglu only), 1: fused and interleaved, 2: fused and not interleaved ([gate | up] for geglu).", AttributeProto::INT, static_cast(0))
+ .Attr("swiglu_limit", "The limit used to clamp in SwiGLU. No clamp when limit is not provided. Does not affect GEGLU.", AttributeProto::FLOAT, OPTIONAL_VALUE)
.Attr("activation_alpha", "Alpha parameter used in activation function.", AttributeProto::FLOAT, 1.0f)
.Attr("activation_beta", "Beta parameter used in activation function.", AttributeProto::FLOAT, 0.0f)
.Attr("k", "Number of top experts to select from expert pool", AttributeProto::INT, static_cast(1))
@@ -1453,7 +1457,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"2D input tensor with shape (total_tokens, num_experts), where total_tokens "
"must match the flattened token count of input",
"T")
- .Input(2, "fc1_experts_weights", "3D input tensor with shape (num_experts, fusion_size * inter_size, hidden_size), where fusion_size is 2 for fused swiglu, and 1 otherwise", "T")
+ .Input(2, "fc1_experts_weights", "3D input tensor with shape (num_experts, fusion_size * inter_size, hidden_size), where fusion_size is 2 for fused swiglu or geglu, and 1 otherwise", "T")
.Input(3, "fc1_experts_bias", "2D optional input tensor with shape (num_experts, fusion_size * inter_size)", "T", OpSchema::Optional)
.Input(4, "fc2_experts_weights", "3D input tensor with shape (num_experts, hidden_size, inter_size)", "T")
.Input(5, "fc2_experts_bias", "2D optional input tensor with shape (num_experts, hidden_size)", "T", OpSchema::Optional)
@@ -1517,6 +1521,10 @@ constexpr const char* qMoE_ver1_doc = R"DOC(
When swiglu_fusion=0, two GEMMs are not fused, and they are FC1 and FC3 in the inputs.
When swiglu_fusion=1, two GEMMs are fused so that g and l are computed in a single GEMM (FC1), and g and l are interleaved on each row of size 2 * inter_size.
When swiglu_fusion=2, two GEMMs are fused, and g and l are concatenated on each row.
+ CUDA also supports GEGLU with quant_type='int': geglu = GELU(g, approximate='tanh') * l.
+ GEGLU requires fused FC1 projections: swiglu_fusion=1 stores [g0,l0,g1,l1,...],
+ and swiglu_fusion=2 stores [g | l]. Separate FC3 weights are not supported.
+ activation_alpha, activation_beta and swiglu_limit do not affect GEGLU.
)DOC";
static void QMoETypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
@@ -1728,7 +1736,7 @@ static OpSchema QMoESchema() {
auto schema = OpSchema()
.SetDoc(std::string(qMoE_ver1_doc) + qMoE_int3_doc)
.Attr("activation_type",
- "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu",
+ "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU) and geglu (GELU-tanh(gate) * up). GEGLU requires quant_type='int', fused FC1 gate/up weights, swiglu_fusion=1 or 2, and no separate FC3 weights. Default is relu",
AttributeProto::STRING,
std::string("relu"))
.Attr("k",
@@ -1761,11 +1769,12 @@ static OpSchema QMoESchema() {
AttributeProto::INT,
OPTIONAL_VALUE)
.Attr("swiglu_fusion",
- "0: not fused, 1: fused and interleaved. 2: fused and not interleaved.",
+ "Gate/up layout for swiglu or geglu. 0: not fused (swiglu only), "
+ "1: fused and interleaved, 2: fused and not interleaved ([gate | up] for geglu).",
AttributeProto::INT,
static_cast(0))
.Attr("swiglu_limit",
- "The limit used to clamp inputs in SwiGLU. It is infinite when limit is not provided.",
+ "The limit used to clamp inputs in SwiGLU. It is infinite when limit is not provided. Does not affect GEGLU.",
AttributeProto::FLOAT,
OPTIONAL_VALUE)
.Attr("activation_alpha",
@@ -1819,7 +1828,7 @@ static OpSchema QMoESchema() {
"fc1_experts_weights",
"3D tensor with shape (num_experts, fusion_size * inter_size, "
"hidden_size * effective_fc1_bits / 8). The last dimension must be byte-aligned. "
- "The fusion_size is 2 for fused swiglu, or 1 otherwise. effective_fc1_bits is "
+ "The fusion_size is 2 for fused swiglu or geglu, or 1 otherwise. effective_fc1_bits is "
"fc1_expert_weight_bits when provided, otherwise expert_weight_bits.",
"T1")
.Input(3,
diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc
index b58462f784636..b3aa286dd6957 100644
--- a/onnxruntime/test/contrib_ops/moe_test.cc
+++ b/onnxruntime/test/contrib_ops/moe_test.cc
@@ -256,6 +256,232 @@ TEST(MoETest, CudaReluSquaredQuantizedInt8) {
RunCudaReluSquaredMoETest(750, 8);
}
+template
+static void RunCudaGegluTest(bool quantized, int fusion, int rows = 2, int top_k = 2,
+ const std::string& expected_failure = "", bool has_bias = false,
+ int fc1_bits = 4, int fc2_bits = 4) {
+ SCOPED_TRACE(::testing::Message() << "fusion=" << fusion << ", rows=" << rows << ", top_k=" << top_k
+ << ", bias=" << has_bias << ", fc1_bits=" << fc1_bits << ", fc2_bits=" << fc2_bits);
+ ASSERT_TRUE(fc1_bits == 2 || fc1_bits == 4);
+ ASSERT_TRUE(fc2_bits == 2 || fc2_bits == 4);
+ if (!HasCudaEnvironment(std::is_same_v ? 800 : 750)) {
+ GTEST_SKIP() << "CUDA architecture does not support the test element type";
+ }
+ constexpr int experts = 2;
+ constexpr int hidden = 128;
+ constexpr int inter = 128;
+ constexpr int block = 32;
+ constexpr float fp32_abs_tolerance = 3e-5f;
+ float max_erf_tanh_separation = 0.0f;
+ std::vector input(rows * hidden);
+ std::vector expected(rows * hidden);
+ std::vector router_probs(rows * experts);
+ std::vector fc1(experts * 2 * inter * hidden, ElementType(0.0f));
+ std::vector fc2(experts * hidden * inter, ElementType(0.0f));
+ std::vector fc1_bias(experts * 2 * inter);
+ std::vector qfc1(fc1.size() / (8 / fc1_bits), fc1_bits == 2 ? 0xAA : 0x88);
+ std::vector qfc2(fc2.size() / (8 / fc2_bits), fc2_bits == 2 ? 0xAA : 0x88);
+ auto set_weight = [](auto& dense, auto& packed, size_t index, int value, int bits) {
+ dense[index] = ElementType(static_cast(value));
+ const int pack_size = 8 / bits;
+ const int shift = static_cast(index % pack_size) * bits;
+ const auto mask = static_cast(((1 << bits) - 1) << shift);
+ packed[index / pack_size] = static_cast(
+ (packed[index / pack_size] & ~mask) | ((value + (1 << (bits - 1))) << shift));
+ };
+ for (int expert = 0; expert < experts; ++expert) {
+ for (int channel = 0; channel < hidden; ++channel) {
+ const int gate_row = fusion == 1 ? channel * 2 : channel;
+ const int up_row = fusion == 1 ? channel * 2 + 1 : inter + channel;
+ set_weight(fc1, qfc1, (expert * 2 * inter + gate_row) * hidden + channel,
+ fc1_bits == 2 ? 1 - 2 * expert : expert + 1, fc1_bits);
+ set_weight(fc1, qfc1, (expert * 2 * inter + up_row) * hidden + channel,
+ fc1_bits == 2 ? -1 - expert : 3 - expert, fc1_bits);
+ set_weight(fc2, qfc2, (expert * hidden + channel) * inter + channel, 1, fc2_bits);
+ fc1_bias[expert * 2 * inter + gate_row] =
+ ElementType(0.125f * (expert + 1) + 0.03125f * (channel % 5 - 2));
+ fc1_bias[expert * 2 * inter + up_row] =
+ ElementType(-0.375f * (expert + 1) + 0.0625f * (channel % 3 - 1));
+ }
+ }
+ for (int row = 0; row < rows; ++row) {
+ // Despite the schema input name, CUDA consumes logits and applies softmax.
+ router_probs[row * experts] = ElementType(row % 2 == 0 ? 1.0f : 0.0f);
+ router_probs[row * experts + 1] = ElementType(row % 2 == 0 ? 0.0f : 1.0f);
+ for (int channel = 0; channel < hidden; ++channel) {
+ const float value = static_cast((channel + row) % 17 - 8) * 0.25f;
+ input[row * hidden + channel] = ElementType(value);
+ float result = 0.0f;
+ float erf_result = 0.0f;
+ for (int expert = 0; expert < experts; ++expert) {
+ const float logit = static_cast(router_probs[row * experts + expert]);
+ const float routing_weight = top_k == 1 ? (logit > 0.5f ? 1.0f : 0.0f)
+ : std::exp(logit) / (std::exp(1.0f) + 1.0f);
+ // Compute logical gate/up values without reading the serialized FC1 layout.
+ const float gate = value * static_cast(fc1_bits == 2 ? 1 - 2 * expert : expert + 1) +
+ (has_bias ? 0.125f * (expert + 1) + 0.03125f * (channel % 5 - 2) : 0.0f);
+ const float up = value * static_cast(fc1_bits == 2 ? -1 - expert : 3 - expert) +
+ (has_bias ? -0.375f * (expert + 1) + 0.0625f * (channel % 3 - 1) : 0.0f);
+ const float gelu = 0.5f * gate *
+ (1.0f + std::tanh(0.7978845608f * (gate + 0.044715f * gate * gate * gate)));
+ result += routing_weight * gelu * up;
+ const float erf_gelu = 0.5f * gate * (1.0f + std::erf(gate / std::sqrt(2.0f)));
+ erf_result += routing_weight * erf_gelu * up;
+ }
+ expected[row * hidden + channel] = ElementType(result);
+ max_erf_tanh_separation = std::max(max_erf_tanh_separation, std::abs(result - erf_result));
+ }
+ }
+ OpTester tester(quantized ? "QMoE" : "MoE", 1, onnxruntime::kMSDomain);
+ tester.AddAttribute("k", top_k);
+ tester.AddAttribute("activation_type", "geglu");
+ tester.AddAttribute("swiglu_fusion", fusion);
+ tester.AddAttribute("normalize_routing_weights", 1);
+ tester.AddInput("input", {rows, hidden}, input);
+ tester.AddInput("router_probs", {rows, experts}, router_probs);
+ if (quantized) {
+ tester.AddAttribute("expert_weight_bits", fc2_bits);
+ tester.AddAttribute("fc1_expert_weight_bits", fc1_bits);
+ tester.AddAttribute("fc3_expert_weight_bits", fc1_bits);
+ tester.AddAttribute("block_size", block);
+ tester.AddAttribute("weights_prepacked", 0);
+ tester.AddInput("fc1_experts_weights", {experts, 2 * inter, hidden / (8 / fc1_bits)}, qfc1, true);
+ tester.AddInput("fc1_scales", {experts, 2 * inter, hidden / block},
+ std::vector(experts * 2 * inter * hidden / block, ElementType(1.0f)), true);
+ if (has_bias) {
+ tester.AddInput("fc1_experts_bias", {experts, 2 * inter}, fc1_bias);
+ } else {
+ tester.AddOptionalInputEdge();
+ }
+ tester.AddInput("fc2_experts_weights", {experts, hidden, inter / (8 / fc2_bits)}, qfc2, true);
+ tester.AddInput("fc2_scales", {experts, hidden, inter / block},
+ std::vector(experts * hidden * inter / block, ElementType(1.0f)), true);
+ } else {
+ tester.AddInput("fc1_experts_weights", {experts, 2 * inter, hidden}, fc1);
+ if (has_bias) {
+ tester.AddInput("fc1_experts_bias", {experts, 2 * inter}, fc1_bias);
+ } else {
+ tester.AddOptionalInputEdge();
+ }
+ tester.AddInput("fc2_experts_weights", {experts, hidden, inter}, fc2);
+ }
+ tester.AddOutput("output", {rows, hidden}, expected);
+ tester.SetOutputTolerance(std::is_same_v ? 0.02f : 0.001f);
+ if constexpr (std::is_same_v) {
+ // CUTLASS uses hardware fast-tanh: its activation error is amplified by
+ // the up projection (2.3e-5 observed on SM80 for this input sweep).
+ // Independently require ample separation from an erf-GELU implementation.
+ ASSERT_GT(max_erf_tanh_separation, 10.0f * fp32_abs_tolerance);
+ tester.SetOutputAbsErr("output", fp32_abs_tolerance);
+ tester.SetOutputRelErr("output", 0.0f);
+ }
+ SessionOptions options;
+ ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1"));
+ if (quantized && (fc1_bits == 2 || fc2_bits == 2)) {
+ // Keep raw [E,N,K/pack] tensors live for the integer dense-dequantization fallback.
+ ORT_THROW_IF_ERROR(options.config_options.AddConfigEntry("session.disable_prepacking", "1"));
+ }
+ std::vector> providers;
+ providers.push_back(DefaultCudaExecutionProvider());
+ tester.Run(options, expected_failure.empty() ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure,
+ expected_failure, {}, nullptr, &providers);
+}
+
+TEST(MoETest, CudaGegluFloat32) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(false, fusion, rows, top_k);
+}
+
+TEST(MoETest, CudaGegluFloat16) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(false, fusion, rows, top_k);
+}
+
+TEST(MoETest, CudaGegluBFloat16) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(false, fusion, rows, top_k);
+}
+
+TEST(MoETest, CudaGegluQuantizedInt4Float16) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(true, fusion, rows, top_k);
+}
+
+#if !defined(ORT_QUICK_BUILD) && defined(ENABLE_BF16)
+TEST(MoETest, CudaGegluQuantizedInt4BFloat16) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(true, fusion, rows, top_k);
+}
+#endif
+
+TEST(MoETest, CudaGegluRequiresFusedWeights) {
+ RunCudaGegluTest(false, 0, 1, 1, "GEGLU requires fused FC1 gate/up weights");
+ RunCudaGegluTest(true, 0, 1, 1, "GEGLU requires fused FC1 gate/up weights");
+}
+
+TEST(MoETest, CudaGegluBiasFloat32) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(false, fusion, rows, top_k, "", true);
+}
+
+TEST(MoETest, CudaGegluBiasFloat16) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(false, fusion, rows, top_k, "", true);
+}
+
+TEST(MoETest, CudaGegluBiasBFloat16) {
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2}) RunCudaGegluTest(false, fusion, rows, top_k, "", true);
+}
+
+TEST(MoETest, CudaGegluBiasTmaSm90) {
+#if defined(USE_CUDA) && defined(HAS_SM90_OR_LATER) && !defined(EXCLUDE_SM_90) && !defined(_MSC_VER)
+ if (!HasCudaEnvironment(900) || GetCudaArchitecture() != 900) {
+ GTEST_SKIP() << "SM90 is required for the Hopper TMA activation regression";
+ }
+ // Non-MSVC SM90 builds provide TMA tactics first; pin that tactic instead of profiling.
+ ScopedEnvironmentVariables env{{{"ORT_FORCE_DETERMINISTIC_MOE", "1"}}};
+ for (int fusion : {1, 2}) {
+ for (int rows : {1, 17}) {
+ for (int top_k : {1, 2}) {
+ RunCudaGegluTest(false, fusion, rows, top_k, "", true);
+#if defined(ENABLE_BF16)
+ RunCudaGegluTest(false, fusion, rows, top_k, "", true);
+#endif
+ }
+ }
+ }
+#else
+ GTEST_SKIP() << "Hopper TMA grouped GEMMs are not compiled in this build";
+#endif
+}
+
+TEST(MoETest, CudaGegluQuantizedInt2AndMixedWidthFloat16) {
+ for (const auto& bits : {std::pair{2, 2}, std::pair{2, 4}, std::pair{4, 2}})
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2})
+ RunCudaGegluTest(true, fusion, rows, top_k, "", true, bits.first, bits.second);
+}
+
+#if defined(ENABLE_BF16)
+TEST(MoETest, CudaGegluQuantizedInt2AndMixedWidthBFloat16) {
+ for (const auto& bits : {std::pair{2, 2}, std::pair{2, 4}, std::pair{4, 2}})
+ for (int fusion : {1, 2})
+ for (int rows : {1, 17})
+ for (int top_k : {1, 2})
+ RunCudaGegluTest(true, fusion, rows, top_k, "", true, bits.first, bits.second);
+}
+#endif
+
static void RunQMoETest(const std::vector& input, const std::vector& router_probs,
const std::vector& fc1_experts_weights,
const std::vector& fc2_experts_weights,