For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

group_limited_router_kernel

def group_limited_router_kernel[scores_type: DType, bias_type: DType, ExpertIndicesLayoutType: TensorLayout, ExpertWeightsLayoutType: TensorLayout, ExpertScoresLayoutType: TensorLayout, ExpertBiasLayoutType: TensorLayout, n_routed_experts: Int, n_experts_per_tok: Int, n_groups: Int, topk_group: Int, norm_weights: Bool, num_threads: Int, scores_input_fn: OptionalReg[def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[scores_type, width]] = None](expert_indices: TileTensor[DType.int32, ExpertIndicesLayoutType, MutAnyOrigin], expert_weights: TileTensor[scores_type, ExpertWeightsLayoutType, MutAnyOrigin], expert_scores: TileTensor[scores_type, ExpertScoresLayoutType, ImmutAnyOrigin], expert_bias: TileTensor[bias_type, ExpertBiasLayoutType, ImmutAnyOrigin], routed_scaling_factor: Float32)

A manually fused MoE router with the group-limited strategy. It divides all the experts into n_groups groups and then finds the top topk_group groups with the highest scores. The final experts for each token are selected from the experts in the selected groups. The bias will be applied to the scores during the selection process, but the final weights will not include the bias.

Parameters:

scores_type (DType): DType of the routing scores and the output expert weights.
bias_type (DType): DType of the per-expert bias added to scores during selection.
ExpertIndicesLayoutType (TensorLayout): TensorLayout of the expert_indices output tensor.
ExpertWeightsLayoutType (TensorLayout): TensorLayout of the expert_weights output tensor.
ExpertScoresLayoutType (TensorLayout): TensorLayout of the expert_scores input tensor.
ExpertBiasLayoutType (TensorLayout): TensorLayout of the expert_bias input tensor.
n_routed_experts (Int): Total number of routed experts scored per token. Also equals the thread count per block.
n_experts_per_tok (Int): Number of experts selected per token (the top-k value).
n_groups (Int): Number of groups the routed experts are partitioned into. n_routed_experts must be divisible by this.
topk_group (Int): Number of highest-scoring groups from which the final experts are selected.
norm_weights (Bool): Whether to normalize the selected weights to sum to one before applying the scaling factor.
num_threads (Int): Threads per block; must equal n_routed_experts so each thread scores one expert.
scores_input_fn (OptionalReg[def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[scores_type, width]]): Optional lambda that loads scores given a (token, expert) index; when None, scores load from expert_scores.

Args:

expert_indices (TileTensor[DType.int32, ExpertIndicesLayoutType, MutAnyOrigin]): Output tensor holding the selected expert index per token. Shape [num_tokens, n_experts_per_tok].
expert_weights (TileTensor[scores_type, ExpertWeightsLayoutType, MutAnyOrigin]): Output tensor holding the routing weight per selected expert per token, excluding the bias. Shape [num_tokens, n_experts_per_tok].
expert_scores (TileTensor[scores_type, ExpertScoresLayoutType, ImmutAnyOrigin]): Input tensor of routing scores for every expert per token. Shape [num_tokens, n_routed_experts].
expert_bias (TileTensor[bias_type, ExpertBiasLayoutType, ImmutAnyOrigin]): Input tensor of per-expert bias added to scores during selection but excluded from the final weights. Shape [n_routed_experts].
routed_scaling_factor (Float32): Factor multiplied into the final, optionally normalized, expert weights.