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

FusedMLP

FusedMLP

class max.nn.FusedMLP(dtype, hidden_dim, feed_forward_length, devices, activation_function='silu', swiglu_limit=0.0, is_sharding=False)

source

Bases: Module, Shardable

Stores the gate and up projections as one pre-fused weight.

The checkpoint provides gate_up_proj_fused with shape [2 * feed_forward_length, hidden_dim], with gate rows followed by up rows. Only bias-free, unquantized projections are supported.

Initializes the fused MLP layer.

Parameters:

• dtype (DType) – DType for the layer weights.
• hidden_dim (int) – The last dimension of the layer input.
• feed_forward_length (int) – Size of the intermediate projection.
• devices (Sequence[DeviceRef]) – Devices to place the weights on. Single-device only.
• activation_function (str) – Activation applied to the gate output.
• swiglu_limit (float) – Optional SwiGLU clamp limit (0 disables clamping).
• is_sharding (bool) – Disable weight creation during sharding.

shard()

shard(devices)

source

Creates sharded views of this fused MLP (single-device only).

Parameters:

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

Returns:

List of sharded FusedMLP instances, one per device.

Return type:

list[FusedMLP]

sharding_strategy

property sharding_strategy: ShardingStrategy | None

source

Get the fused MLP sharding strategy.