QMoE (CPU)#

  • Domain: com.microsoft

  • Device: CPU

  • Kernel: onnx_light_cpu::QMoE

  • Supported types: FLOAT, FLOAT16, BFLOAT16

  • Opset bounds: since opset 1, until opset 1

Operator support#

Capability

Implementation

Shape inference

onnx_light_cpu::ComputeShapeQMoE()

Peak memory

onnx_light_cpu::ComputePeakMemoryQMoE()

Fusion patterns

None

Gradient

Not provided

LightOpSchema#

Version: 1

CPU quantized mixture of experts. Weights use the mbext CPU layout, not CUDA packing.

Inputs

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

  • router_probs (T): Rank-2 token, expert router logits.

  • fc1_experts_weights (Q): Packed expert projection.

  • fc1_scales (S): Per-channel or block-wise scales.

  • fc1_experts_bias (T): Optional first projection bias.

  • fc2_experts_weights (Q): Packed output projection.

  • fc2_scales (S): Per-channel or block-wise scales.

  • fc2_experts_bias (T): Optional output projection bias.

  • fc3_experts_weights (Q): Optional gate projection.

  • fc3_scales (S): Optional gate scales.

  • fc3_experts_bias (T): Optional gate bias.

  • fc1_zero_points (Q): Unsupported optional zero points.

  • fc2_zero_points (Q): Unsupported optional zero points.

  • fc3_zero_points (Q): Unsupported optional zero points.

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)

  • block_size (INT): Zero for tensor scales, or a power of two of at least 16. (optional, default: 0)

  • expert_weight_bits (INT): 4 or 8 bits per weight. (optional, default: 4)

  • 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 mode. (optional, default: 0)

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

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

  • weights_prepacked (INT): Only zero (CPU layout) is supported. (optional, default: 0)

Type constraints

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

  • S: FLOAT, FLOAT16 or BFLOAT16 scales. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • Q: Packed unsigned bytes. Allowed types: tensor(uint8).