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).

Mojo function

moe_create_indices_kernel

def moe_create_indices_kernel[input_type: DType, TokenExpertOrderLayoutType: TensorLayout, ExpertStartIndicesLayoutType: TensorLayout, RestoreTokenOrderLayoutType: TensorLayout, ExpertIdsLayoutType: TensorLayout, ExpertUsageStatsLayoutType: TensorLayout, TopkIdsLayoutType: TensorLayout, _scale_alignment: UInt32 = UInt32(128)](token_expert_order: TileTensor[.uint32, TokenExpertOrderLayoutType, MutAnyOrigin], expert_start_indices: TileTensor[.uint32, ExpertStartIndicesLayoutType, MutAnyOrigin], restore_token_order: TileTensor[.uint32, RestoreTokenOrderLayoutType, MutAnyOrigin], expert_ids: TileTensor[.int32, ExpertIdsLayoutType, MutAnyOrigin], expert_usage_stats: TileTensor[.uint32, ExpertUsageStatsLayoutType, MutAnyOrigin], topk_ids: TileTensor[input_type, TopkIdsLayoutType, ImmutAnyOrigin], scales_offset_p: Optional[Pointer[UInt32, MutAnyOrigin]])

Builds MoE routing indices in one CTA using a block-wide scan.

Groups tokens by their assigned expert so downstream kernels such as grouped matmul can process each expert's tokens contiguously. Every expert in [0, num_experts) gets a slot, including ones with zero tokens; expert_usage_stats holds the largest per-expert token count and num_experts.

Slot order is ascending expert id; token order within a slot follows the shared cursor. Neither is load-bearing: consumers pair expert_start_indices[g] and [g + 1] with expert_ids[g], and invert through restore_token_order.