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 function

moe_router_group_limited

moe_router_group_limited()

max.experimental.nn.common_layers.functional_kernels.moe_router_group_limited(expert_scores, expert_bias, n_routed_experts, n_experts_per_tok, n_groups, topk_group, norm_weights, routed_scaling_factor)

source

Routes tokens with the group-limited MoE router.

When n_groups > 1, selects up to topk_group expert groups, then picks n_experts_per_tok experts within those groups (DeepSeek-V3 style). When n_groups == 1, there is only one group, so group selection is skipped and routing uses the dedicated GPU single-group path (mo.moe.single.group.router, implemented as single_group_router in Mojo). In that case topk_group is not used by the kernel.

Reference: https://github.com/deepseek-ai/DeepSeek-V3/blob/9b4e9788e4a3a731f7567338ed15d3ec549ce03b/inference/model.py#L566

Parameters:

• expert_scores (TensorValue) – The scores for each expert for each token. Shape: [num_tokens, n_routed_experts].
• expert_bias (TensorValue) – The bias for each expert. Shape: [n_routed_experts].
• n_routed_experts (int) – The total number of experts. Must be divisible by n_groups.
• n_experts_per_tok (int) – The number of experts to be selected per token.
• n_groups (int) – The total number of expert groups. n_routed_experts must be divisible by this.
• topk_group (int) – The maximum number of expert groups that a token will be routed to.
• norm_weights (bool) – Whether to normalize the selected expert weights when n_groups > 1. When n_groups == 1, normalization is currently always enabled (norm_weights is treated as True) so behavior matches the graph path that always divided weights by their sum per token.
• routed_scaling_factor (float)

Returns:

• expert_indices: The indices of the routed experts for each token. Shape: [num_tokens, n_experts_per_tok].
• expert_weights: The weights of the routed experts for each token. Shape: [num_tokens, n_experts_per_tok].

Return type:

A tuple of two tensors