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_create_indices
moe_create_indices()
max.experimental.nn.common_layers.functional_kernels.moe_create_indices(topk_ids, num_local_experts, *, needs_scales_offset=False)
Creates indices for the MoE layer.
-
Parameters:
-
- topk_ids (TensorValue) – The expert assignments for each token from the router.
- num_local_experts (int) – The number of experts on this device.
- needs_scales_offset (bool)
-
Returns:
-
- token_expert_order: The reordered token indices, grouped by assigned expert.
- expert_start_indices: The starting index for each expert’s token group in the reordered sequence.
- restore_token_order: The indices that restore the original token ordering after expert computation.
- expert_ids: The IDs of the active experts selected for tokens.
- expert_usage_stats: The maximum number of tokens assigned to any expert, and the number of active experts.
-
Return type:
-
A tuple of five tensors