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

SigmoidTopKRouter

SigmoidTopKRouter

class max.nn.moe.SigmoidTopKRouter(devices, hidden_dim, num_experts, num_experts_per_token, dtype, norm_topk_prob=True, gate_dtype=None, correction_bias_dtype=float32, routed_scaling_factor=1.0, linear_cls=<class 'max.nn.linear.Linear'>, is_sharding=False)

source

Bases: MoEGate

MoE gate with sigmoid scores and a correction bias (noaux_tc).

Routes as DeepSeek-V3’s auxiliary-loss-free noaux_tc gate does with a single expert group:

1. Computes the gate logits with a linear projection.
2. Applies a sigmoid, in correction_bias_dtype.
3. Selects the top-k experts by the sigmoid score plus the learned e_score_correction_bias.
4. Weights each selected expert by its unbiased sigmoid score, optionally renormalized to sum to 1, times routed_scaling_factor.

For a float32 router, pass dtype=DType.float32 and a float32 input: the gate projection and the scores then stay in float32.

Parameters:

• devices (list[DeviceRef]) – The devices to place the gate on.
• hidden_dim (int) – The dimension of the hidden state.
• num_experts (int) – The number of routed experts.
• num_experts_per_token (int) – The number of experts each token selects.
• dtype (DType) – The gate projection’s dtype, unless gate_dtype is set.
• norm_topk_prob (bool) – Whether to renormalize the selected weights to sum to 1. Defaults to True.
• gate_dtype (DType | None) – The gate projection’s dtype. Defaults to dtype.
• correction_bias_dtype (DType) – The dtype of the correction bias and of the sigmoid scores. Defaults to DType.float32.
• routed_scaling_factor (float) – The factor applied to the routing weights. Defaults to 1.0.
• linear_cls (Callable[..., Linear]) – The linear class for the gate projection.
• is_sharding (bool) – Whether shard() is constructing this gate, which then assigns the weights itself.

Args: devices: List of devices to use for the MoEGate. hidden_dim: The dimension of the hidden state. num_experts: The number of experts. num_experts_per_token: The number of experts per token. dtype: The data type of the MoEGate.

shard()

shard(devices)

source

Creates one replica of this gate per device.

Parameters:

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

Returns:

One gate per device, of this gate’s class.

Raises:

ValueError – If no sharding strategy has been set.

Return type:

list[Self]