From 660f358a60667a680aa4fa2e9078b91900ee6f67 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 27 Oct 2026 00:31:27 +0000 Subject: [PATCH 1/3] Enable tanh-GEGLU in CUDA MoE and integer QMoE Expose the existing gated activation through the public operator parser and validate fused gate/up layouts without enabling SwiGLU-only fast paths. Use tanh-GELU and preserve chunked gate/up ordering for GEGLU, while retaining the existing GELU and SwiGLU behavior. Add independent numerical/layout regression tests and document the CUDA contract, precision limits, and unsupported quantization families. Preserve the current-main INT3 schema and shape inference during the port. Validation: scoped lintrunner and whitespace checks pass. Earlier SM80 build executed 12 selected CUDA provider tests with no failures or skips. The current-main port still requires a fresh build/runtime qualification; full-model numerical and performance attribution remain unestablished. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/moe_qmoe.md | 29 +++- .../moe_gemm/moe_gemm_activation_kernels.cuh | 19 ++- onnxruntime/contrib_ops/cuda/moe/moe.cc | 12 +- onnxruntime/contrib_ops/cuda/moe/moe_base.h | 9 +- .../contrib_ops/cuda/moe/moe_quantization.cc | 12 +- .../core/graph/contrib_ops/contrib_defs.cc | 25 +++- onnxruntime/test/contrib_ops/moe_test.cc | 138 ++++++++++++++++++ 7 files changed, 215 insertions(+), 29 deletions(-) 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..581ff7e4ba805 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -256,6 +256,144 @@ 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 = "") { + SCOPED_TRACE(::testing::Message() << "fusion=" << fusion << ", rows=" << rows << ", top_k=" << top_k); + 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 qfc1(fc1.size() / 2, 0x88); + std::vector qfc2(fc2.size() / 2, 0x88); + auto set_weight = [](auto& dense, auto& packed, size_t index, int value) { + dense[index] = ElementType(static_cast(value)); + const int shift = static_cast(index % 2) * 4; + const auto mask = static_cast(0xF << shift); + packed[index / 2] = static_cast( + (packed[index / 2] & ~mask) | ((value + 8) << 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, expert + 1); + set_weight(fc1, qfc1, (expert * 2 * inter + up_row) * hidden + channel, 3 - expert); + set_weight(fc2, qfc2, (expert * hidden + channel) * inter + channel, 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); + const float gate = value * static_cast(expert + 1); + const float gelu = 0.5f * gate * + (1.0f + std::tanh(0.7978845608f * (gate + 0.044715f * gate * gate * gate))); + result += routing_weight * gelu * value * static_cast(3 - expert); + const float erf_gelu = 0.5f * gate * (1.0f + std::erf(gate / std::sqrt(2.0f))); + erf_result += routing_weight * erf_gelu * value * static_cast(3 - expert); + } + 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", 4); + tester.AddAttribute("block_size", block); + tester.AddAttribute("weights_prepacked", 0); + tester.AddInput("fc1_experts_weights", {experts, 2 * inter, hidden / 2}, qfc1, true); + tester.AddInput("fc1_scales", {experts, 2 * inter, hidden / block}, + std::vector(experts * 2 * inter * hidden / block, ElementType(1.0f)), true); + tester.AddOptionalInputEdge(); + tester.AddInput("fc2_experts_weights", {experts, hidden, inter / 2}, 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); + 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")); + 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"); +} + static void RunQMoETest(const std::vector& input, const std::vector& router_probs, const std::vector& fc1_experts_weights, const std::vector& fc2_experts_weights, From e93f30907a81b7df1fea0d0f134c84227894c0e2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:10:32 +0000 Subject: [PATCH 2/3] Update generated MoE and QMoE GEGLU documentation Synchronize generated documentation with the GPU doc-gen artifact. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) 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
From 92741826f04921e7c995b291a4bdb5d3b6ae3078 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:54:34 +0000 Subject: [PATCH 3/3] test: cover biased GEGLU and mixed-width QMoE Expand test-only GEGLU and QMoE coverage with nonzero independent gate/up bias, fusion1/2 layouts, FP32/FP16/BF16 types, raw INT2, both mixed-width quantization directions, and an explicitly guarded SM90 TMA case. Validation: 11 passed and 1 explicit SM90 build skip against the existing older custom runtime, not the full current HEAD. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- onnxruntime/test/contrib_ops/moe_test.cc | 128 +++++++++++++++++++---- 1 file changed, 108 insertions(+), 20 deletions(-) diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 581ff7e4ba805..b3aa286dd6957 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -258,8 +258,12 @@ TEST(MoETest, CudaReluSquaredQuantizedInt8) { template static void RunCudaGegluTest(bool quantized, int fusion, int rows = 2, int top_k = 2, - const std::string& expected_failure = "") { - SCOPED_TRACE(::testing::Message() << "fusion=" << fusion << ", rows=" << rows << ", top_k=" << top_k); + 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"; } @@ -274,22 +278,30 @@ static void RunCudaGegluTest(bool quantized, int fusion, int rows = 2, int top_k 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 qfc1(fc1.size() / 2, 0x88); - std::vector qfc2(fc2.size() / 2, 0x88); - auto set_weight = [](auto& dense, auto& packed, size_t index, int value) { + 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 shift = static_cast(index % 2) * 4; - const auto mask = static_cast(0xF << shift); - packed[index / 2] = static_cast( - (packed[index / 2] & ~mask) | ((value + 8) << shift)); + 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, expert + 1); - set_weight(fc1, qfc1, (expert * 2 * inter + up_row) * hidden + channel, 3 - expert); - set_weight(fc2, qfc2, (expert * hidden + channel) * inter + channel, 1); + 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) { @@ -305,12 +317,16 @@ static void RunCudaGegluTest(bool quantized, int fusion, int rows = 2, int top_k 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); - const float gate = value * static_cast(expert + 1); + // 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 * value * static_cast(3 - expert); + 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 * value * static_cast(3 - expert); + 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)); @@ -324,19 +340,29 @@ static void RunCudaGegluTest(bool quantized, int fusion, int rows = 2, int top_k tester.AddInput("input", {rows, hidden}, input); tester.AddInput("router_probs", {rows, experts}, router_probs); if (quantized) { - tester.AddAttribute("expert_weight_bits", 4); + 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 / 2}, qfc1, true); + 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); - tester.AddOptionalInputEdge(); - tester.AddInput("fc2_experts_weights", {experts, hidden, inter / 2}, qfc2, 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); - tester.AddOptionalInputEdge(); + 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); @@ -351,6 +377,10 @@ static void RunCudaGegluTest(bool quantized, int fusion, int rows = 2, int top_k } 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, @@ -394,6 +424,64 @@ TEST(MoETest, CudaGegluRequiresFusedWeights) { 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,