MoE (CPU)#

  • Domain: com.microsoft

  • Device: CPU

  • Kernel: onnx_light_cpu::MoE

  • Supported types: FLOAT, FLOAT16, BFLOAT16

  • Opset bounds: since opset 1, until opset 1

Operator support#

Capability

Implementation

Shape inference

onnx_light_cpu::ComputeShapeMoE()

Peak memory

onnx_light_cpu::ComputePeakMemoryMoE()

Fusion patterns

None

Gradient

Not provided

LightOpSchema#

Version: 1

ONNX Runtime CPU MoE: softmax router logits, select top-k experts and combine their outputs. FC1 and FC2 weights are expert-major, output-major.

Inputs

  • input (T): Rank-2 (tokens, hidden) or rank-3 (batch, sequence, hidden).

  • router_probs (T): Rank-2 (tokens, experts) logits.

  • fc1_experts_weights (T): (experts, intermediate * fusion, hidden).

  • fc1_experts_bias (T): Optional (experts, intermediate * fusion).

  • fc2_experts_weights (T): (experts, hidden, intermediate).

  • fc2_experts_bias (T): Optional (experts, hidden).

  • fc3_experts_weights (T): Unsupported optional FC3 slot.

  • fc3_experts_bias (T): Unsupported optional FC3 slot.

Outputs

  • output (T): Same shape and type as input.

Attributes

  • activation_alpha (FLOAT): SwiGLU sigmoid scale. (optional, default: 1)

  • activation_beta (FLOAT): SwiGLU linear offset. (optional, default: 0)

  • activation_type (STRING): relu, gelu, silu, identity or swiglu. (optional, default: relu)

  • k (INT): Number of selected experts. (optional, default: 1)

  • normalize_routing_weights (INT): Renormalize selected softmax weights. (optional, default: 0)

  • swiglu_fusion (INT): Interleaved SwiGLU requires 1. (optional, default: 0)

  • swiglu_limit (FLOAT): Optional SwiGLU clamp limit. (optional)

  • use_sparse_mixer (INT): Only zero is supported. (optional, default: 0)

Type constraints

  • T: Matching FLOAT, FLOAT16 or BFLOAT16 tensors. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).