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.legacy.MoE(devices, hidden_dim, num_experts, num_experts_per_token, moe_dim, num_logical_experts=None, gate_cls=<class 'max.nn.legacy.moe.moe.MoEGate'>, mlp_cls=<class 'max.nn.legacy.linear.MLP'>, shared_mlp_cls=None, has_shared_experts=False, shared_experts_dim=0, ep_size=1, dtype=bfloat16, activation=None, interleaved_gate_up=False, 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.
• ep_batch_manager (EPBatchManager | None) – The expert parallel batch manager. Defaults to None.
• quant_config (QuantConfig | None) – The scaled quantization configuration. Defaults to None.
• activation (GatedActivation | None) – The gated activation between the gate-up and down projections; see GatedActivation. Defaults to a fused SiLU.
• interleaved_gate_up (bool) – Whether each row of the gate-up projection’s output is [gate_0, up_0, gate_1, up_1, ...] rather than [gate | up]. The default gate_up_proj builds concatenated rows; a subclass whose gate_up_proj returns interleaved rows sets this. Defaults to False.
• 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 (such as 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 (such as 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)

down_proj

property down_proj: TensorValue

source

The down projection.

ep_batch_manager

property ep_batch_manager: EPBatchManager

source

Get the expert parallel batch manager.

eplb_router_weights()

eplb_router_weights()

source

Returns the router weights whose leading axis indexes experts.

Reordering these rows by a permutation placement makes the router emit physical ids directly. Returns None when that is unsafe: the router selects within expert groups, which a permutation can break, or some router weight has an expert-sized dimension off axis 0, so a row reorder would miss it.

Returns:

FQNs relative to this module, or None.

Return type:

list[str] | None

eplb_weight_sources()

eplb_weight_sources(phy2log)

source

Maps each physical-slot weight to the logical expert weight it loads.

Parameters:

phy2log (Sequence[int]) – The logical expert id of every physical slot in this layer.

Returns:

A map from destination FQN to source FQN, both relative to this module.

Raises:

ValueError – If phy2log does not cover every slot, or this MoE does not keep one module per expert.

Return type:

dict[str, str]

experts

experts: LayerList

source

The physical expert slots, one per num_experts.

Without an EPLB placement slot i holds logical expert i. With one, apply_eplb_placement() decides which logical expert’s checkpoint weights each slot loads.

gate_up_proj

property gate_up_proj: TensorValue

source

The fused gate and up projection.

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.