Xenova HF Staff

sync 2e7068faf55e

3ab8080 verified 6 days ago

5.93 kB

library_name: kernels
license: apache-2.0
tags:
  - kernel
  - webgpu
  - wgsl

com.microsoft.MoE

com.microsoft · ONNX Runtime contrib operator · contrib since_version 1

Description

Mixture of Experts: applies softmax to router_probs, routes each token to the top-k experts, applies FC1 and activation_type, projects through FC2, then sums the selected outputs using their routing probabilities. SwiGLU takes its operands from a separate FC3 (swiglu_fusion 0) or a fused FC1 in interleaved (1) or concatenated (2) order; SiLU may also use FC3 as its multiplicative linear projection. This inference package supports float32 and dense routing (use_sparse_mixer = 0); float16, bfloat16, and sparse mixing are not implemented. Quantized weights use com.microsoft.QMoE.

See the ONNX Runtime MoE contrib-operator spec for the reference semantics.

Inputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`input`	`inputT`	`T`	—	—	Token activations, either 2D `(num_tokens, hidden_size)` or 3D `(batch_size, sequence_length, hidden_size)`.	required
`router_probs`	`routerT`	`T`	`2`	—	2D router logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of every leading dimension of `input`. Despite the historical port name, the operator applies a full softmax before top-k selection.	required
`fc1_experts_weights`	`fc1T`	`T`	`3`	—	3D first-layer expert weights of shape `(num_experts, fusion_size * inter_size, hidden_size)`, where `fusion_size` is 2 for fused SwiGLU (`swiglu_fusion` 1 or 2) and 1 otherwise.	required
`fc1_experts_bias`	`fc1BiasT`	`T`	`2`	—	Optional 2D FC1 bias of shape `(num_experts, fusion_size * inter_size)`.	optional
`fc2_experts_weights`	`fc2T`	`T`	`3`	—	3D second-layer expert weights of shape `(num_experts, hidden_size, inter_size)`.	required
`fc2_experts_bias`	`fc2BiasT`	`T`	`2`	—	Optional 2D FC2 bias of shape `(num_experts, hidden_size)`, added per expert before that expert's routing weight is applied.	optional
`fc3_experts_weights`	`fc3T`	`T`	`3`	—	Optional 3D third-layer expert weights of shape `(num_experts, inter_size, hidden_size)`. It supplies the separate linear operand for SwiGLU when `swiglu_fusion` is 0, or the multiplicative linear projection for SiLU gating. Other activations do not consume FC3.	optional
`fc3_experts_bias`	`fc3BiasT`	`T`	`2`	—	Optional 2D FC3 bias of shape `(num_experts, inter_size)`.	optional

Outputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`output`	`outputT`	`T`	same as `input`	same as `input`	Routed expert output with the same shape as `input`.	required

Attributes

Attributes and default values (overridable per request):

Attribute	Default	Description
`activation_alpha`	`1`	Alpha parameter used by the activation; the schema default is 1.
`activation_beta`	`0`	Beta parameter used by the activation; the schema default is 0.
`activation_type`	`"relu"`	Activation applied to the FC1 projection: `relu`, `gelu`, `silu`, `swiglu`, or `identity`. The schema default is `relu`.
`k`	`1`	Number of experts selected per token; the schema default is 1.
`normalize_routing_weights`	`0`	Whether to normalize the selected routing weights; the schema default is 0.
`swiglu_fusion`	`0`	0 keeps the SwiGLU operands in separate FC1/FC3 GEMMs, 1 interleaves them in one FC1 row, and 2 concatenates them. The schema default is 0.
`use_sparse_mixer`	`0`	Whether to use sparse-mixer routing. The standard default and only supported value is 0.
`swiglu_limit`	—	Optional SwiGLU clamp limit; omission means no clamp.

Type constraints

Variable	Allowed dtypes
`T`	`float32`

Files

metadata.json — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
expert-group-slots.wgsl.jinja
moe-ffn-gemv.wgsl.jinja
moe-ffn-grouped.wgsl.jinja
moe-ffn-stage.wgsl.jinja
moe-mix-stage.wgsl.jinja
moe-output-gemv.wgsl.jinja
moe-output-grouped.wgsl.jinja
moe-output-stage.wgsl.jinja
moe-route-stage.wgsl.jinja

Use with `@huggingface/kernels`

The loader derives every required output's shape and logical dtype from the manifest contract and this call. It then allocates the result tensors automatically.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model version.

Replace each *Data placeholder with a typed array containing the corresponding input data.

import { getKernel } from "@huggingface/kernels";

const kernel = await getKernel("webgpu-kernels/com.microsoft.MoE", { version: 1 });
const { outputT } = await kernel({
  inputT: { data: inputTData, shape: [1, 1] },
  routerT: { data: routerTData, shape: [1, 2] },
  fc1T: { data: fc1TData, shape: [2, 1, 1] },
  fc2T: { data: fc2TData, shape: [2, 1, 1] },
});