For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

MoE

MoE

class max.nn.MoE(devices, hidden_dim, num_experts, num_experts_per_token, moe_dim, num_logical_experts=None, gate_cls=<class 'max.nn.moe.moe.MoEGate'>, mlp_cls=<class 'max.nn.linear.MLP'>, shared_mlp_cls=None, has_shared_experts=False, shared_experts_dim=0, ep_size=1, dtype=bfloat16, apply_router_weight_first=False, use_swigluoai=False, swiglu_alpha=0.0, swiglu_limit=0.0, gated_activation_fn=None, pre_expert_norm_cls=None, ep_batch_manager=None, quant_config=None, shared_experts_dtype=None, shared_experts_quant_config=None, is_sharding=False)

source

Bases: Module, Shardable

Implementation of Mixture of Experts (MoE).

Parameters:

• devices (list[DeviceRef]) – The list of devices to use for the MoE.
• hidden_dim (int) – The dimension of the hidden state.
• num_experts (int) – The number of experts.
• num_experts_per_token (int) – The number of experts per token.
• moe_dim (int) – The intermediate dimension of each expert.
• gate_cls (Callable[..., MoEGate]) – The model-specific gate implementation. Defaults to MoEGate.
• mlp_cls (Callable[..., MLP]) – The MLP class to use for experts. Defaults to MLP.
• has_shared_experts (bool) – Whether to use shared experts. Defaults to False.
• shared_experts_dim (int) – The dimension of the shared experts. Defaults to 0.
• ep_size (int) – The expert parallelism size. Defaults to 1.
• dtype (DType) – The data type of the MoE. Defaults to DType.bfloat16.
• apply_router_weight_first (bool) – Whether to apply the router weight first. Defaults to False.
• ep_batch_manager (EPBatchManager | None) – The expert parallel batch manager. Defaults to None.
• quant_config (QuantConfig | None) – The scaled quantization configuration. Defaults to None.
• use_swigluoai (bool) – Whether to use the OAI-style clamped SwiGLU activation function. Defaults to False.
• swiglu_alpha (float) – The alpha value for the clamped SwiGLU activation function. Defaults to 0.0.
• swiglu_limit (float) – The limit value for the clamped SwiGLU activation function. Defaults to 0.0.
• gated_activation_fn (Callable[[TensorValue, int], TensorValue] | None) – Activation applied to the concatenated [gate | up] projection. None (default) uses a fused SiLU kernel; use make_concatenated_gated_activation_fn() for custom activations.
• shared_experts_dtype (DType | None) – Weight storage dtype for shared-expert MLPs. When equal to dtype (routed experts) and quant_config is set, shared experts use the same quantization as routed experts. When different (e.g. BF16 shared weights with packed NVFP4 routed experts), shared linears omit quant_config unless shared_experts_quant_config is set. Defaults to dtype.
• shared_experts_quant_config (QuantConfig | None) – Optional separate QuantConfig for shared-expert MLPs when their storage dtype differs from routed experts (e.g. MXFP8 shared with NVFP4 routed). Defaults to None.
• pre_expert_norm_cls (Callable[[], Module] | None) – A callable that returns a normalization module to apply before expert computation. Defaults to None.
• is_sharding (bool) – Whether the constructor is being called during sharding. Defaults to False.
• num_logical_experts (int | None)
• shared_mlp_cls (Callable[..., MLP] | None)

configure_ep_scale_fusion()

configure_ep_scale_fusion(dispatch_supports_fold)

source

Configure any EP dispatch-scale fusion before the dispatch op.

No-op on the base class; MoEQuantized overrides it to enable the MXFP4 up-proj A-scale preshuffle fold. Defined here (rather than duck-typed) so the EP forward driver can call it on any MoE shard: non-quantized subclasses inherit this no-op and consistently skip the fold (no fusion, no corruption).

Parameters:

dispatch_supports_fold (bool) – Whether the selected dispatch path wires the A-scale fold params. Ignored by this base no-op.

Return type:

None

down_proj

property down_proj: TensorValue

source

ep_batch_manager

property ep_batch_manager: EPBatchManager

source

Get the expert parallel batch manager.

experts

experts: LayerList

source

The list of experts.

gate_up_proj

property gate_up_proj: TensorValue

source

layer_idx

layer_idx: int | None = None

source

The index of the MoE layer.

shard()

shard(devices)

source

Create sharded views of this MoE module across multiple devices.

Parameters:

devices (Iterable[DeviceRef]) – Iterable of devices to place the shards on.

Returns:

List of sharded MoE instances, one for each device.

Return type:

list[Self]

shard_devices

shard_devices: list[DeviceRef] = []

source

The list of devices the MoE layer was sharded to.

shard_index

shard_index: int = 0

source

The index of the current shard (if the MoE layer was sharded).

sharding_strategy

property sharding_strategy: ShardingStrategy | None

source

Get the sharding strategy for the module.