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

StackedMoE

StackedMoE

class max.nn.moe.StackedMoE(devices, hidden_dim, num_experts, num_experts_per_token, moe_dim, gate_cls, dtype=bfloat16, gate_up_format=GateUpFormat.CONCATENATED, gated_activation_fn=None, has_bias=False, has_shared_experts=False, shared_experts_dim=0, quant_config=None, apply_router_weight_first=False, router_dtype=None, mxfp8_activations=False, is_sharding=False)

source

Bases: Module, Shardable

Stacked Mixture of Experts layer with configurable components.

This class consolidates MoE implementations from multiple architectures (Llama4, Qwen3VL, GptOss) into a single base layer. All expert weights are stored in stacked format rather than as individual MLP experts.

Weight tensor shapes:

• gate_up_proj: [num_experts, hidden_dim, 2 * moe_dim]
• down_proj: [num_experts, moe_dim, hidden_dim]
• Optional FP8 scales: [num_experts, scaled_rows, scaled_cols]

Supported configurations:

• Gate/up formats: concatenated or interleaved.
• Activation functions: configurable (default: SiLU).
• Optional bias support for projections.
• Optional FP8 quantization with block scaling.
• Optional MXFP4 weights, run W4A16, or W4A8 with MXFP8 activations.
• Optional float32 router input.
• Optional shared experts.

from max.driver import Accelerator, CPU, accelerator_count
from max.graph import DeviceRef
from max.nn.moe import MoEGate, StackedMoE

device = Accelerator() if accelerator_count() > 0 else CPU()
device_ref = DeviceRef.from_device(device)

moe = StackedMoE(
    devices=[device_ref],
    hidden_dim=4096,
    num_experts=8,
    num_experts_per_token=2,
    moe_dim=14336,
    gate_cls=MoEGate,
)

Pass gate_up_format=GateUpFormat.INTERLEAVED, a custom gated_activation_fn, has_bias=True, or a quant_config to match other architectures (for example, GptOss style with interleaved format and bias, or FP8 quantization).

Parameters:

• devices (list[DeviceRef]) – A list of devices to use for the MoE.
• hidden_dim (int) – The dimension of the hidden state.
• num_experts (int) – The total number of experts.
• num_experts_per_token (int) – The number of experts per token (top-k).
• moe_dim (int) – The intermediate dimension of each expert.
• gate_cls (Callable[..., MoEGate]) – The model-specific gate implementation class.
• dtype (DType) – The data type of the MoE weights. Defaults to DType.bfloat16.
• gate_up_format (GateUpFormat) – The format of the combined gate/up weights. Defaults to GateUpFormat.CONCATENATED.
• gated_activation_fn (Callable[[TensorValue, TensorValue], TensorValue] | None) – Activation applied to the split (gate, up) projections. None (default) uses SiLU gating; use make_stacked_gated_activation_fn() for custom activations.
• has_bias (bool) – Whether to include bias for projections. Defaults to False.
• 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.
• quant_config (QuantConfig | None) – The configuration for scaled quantization. Defaults to None.
• apply_router_weight_first (bool) – Whether to apply router weights before expert computation. Defaults to False.
• router_dtype (DType | None) – The dtype the gate is built in and its input is cast to. None (default) builds a BF16 gate and passes the input through unchanged. Use DType.float32 for a router whose selection is sensitive to BF16 rounding.
• mxfp8_activations (bool) – Whether MXFP4 experts run W4A8: activations are quantized to MXFP8 per expert, with the routed-row gather fused into the quantize, for the SM100 block-scaled grouped matmul. The weight scales are then declared in that kernel’s interleaved layout, [num_experts, N / 128, K / 128, 32, 4, 4] (see interleaved_block_scales_shape()), and tensor parallelism splits every expert into whole scale granules. Requires an MXFP4 quant_config. Defaults to False.
• is_sharding (bool) – Whether this instance is being created for sharding. Set by shard() to skip weight initialization for sharded instances. Defaults to False.

down_proj_transposed

property down_proj_transposed: TensorValue

source

The down weights transposed to [num_experts, out_features, in_features] layout.

gate_up_proj_transposed

property gate_up_proj_transposed: TensorValue

source

The gate/up weights transposed to [num_experts, out_features, in_features] layout.

shard()

shard(devices)

source

Creates sharded views of this MoE module across multiple devices.

Parameters:

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

Returns:

A sequence of sharded instances, one for each device.

Raises:

ValueError – If no sharding strategy has been set.

Return type:

Sequence[Self]

sharding_strategy

property sharding_strategy: ShardingStrategy | None

source

The sharding strategy for this module.