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

def moe_create_indices[input_type: DType, //, target: StringSpan[ImmStaticOrigin]](token_expert_order: TileTensor[.uint32, Engine=token_expert_order.Engine, address_space=token_expert_order.address_space, linear_idx_type=token_expert_order.linear_idx_type], expert_start_indices: TileTensor[.uint32, Engine=expert_start_indices.Engine, address_space=expert_start_indices.address_space, linear_idx_type=expert_start_indices.linear_idx_type], restore_token_order: TileTensor[.uint32, Engine=restore_token_order.Engine, address_space=restore_token_order.address_space, linear_idx_type=restore_token_order.linear_idx_type], expert_ids: TileTensor[.int32, Engine=expert_ids.Engine, address_space=expert_ids.address_space, linear_idx_type=expert_ids.linear_idx_type], expert_usage_stats: TileTensor[.uint32, Engine=expert_usage_stats.Engine, address_space=expert_usage_stats.address_space, linear_idx_type=expert_usage_stats.linear_idx_type], topk_ids: TileTensor[input_type, Engine=topk_ids.Engine, address_space=topk_ids.address_space, linear_idx_type=topk_ids.linear_idx_type], context: DeviceContext, scales_offset_p: Optional[Pointer[UInt32, MutAnyOrigin]] = None)

Launches the MoE index creation kernel on GPU.

Groups tokens by their assigned expert so that downstream kernels such as grouped matmul can process each expert's tokens contiguously. One CTA histograms the expert ids in shared memory, scans the histogram into CSR offsets, and scatters the tokens.

Parameters:

• ​input_type (DType): DType of the topk_ids tensor.
• ​target (StringSpan[ImmStaticOrigin]): The target device to run the kernel on.

Args:

• ​token_expert_order (TileTensor[.uint32, Engine=token_expert_order.Engine, address_space=token_expert_order.address_space, linear_idx_type=token_expert_order.linear_idx_type]): Output 1D tensor of token indices grouped by expert.
• ​expert_start_indices (TileTensor[.uint32, Engine=expert_start_indices.Engine, address_space=expert_start_indices.address_space, linear_idx_type=expert_start_indices.linear_idx_type]): Output 1D tensor of CSR-style start offsets for each expert in token_expert_order.
• ​restore_token_order (TileTensor[.uint32, Engine=restore_token_order.Engine, address_space=restore_token_order.address_space, linear_idx_type=restore_token_order.linear_idx_type]): Output 1D tensor mapping each token to its new position in token_expert_order.
• ​expert_ids (TileTensor[.int32, Engine=expert_ids.Engine, address_space=expert_ids.address_space, linear_idx_type=expert_ids.linear_idx_type]): Output 1D tensor of the expert IDs in output order, one slot per expert in ascending ID order.
• ​expert_usage_stats (TileTensor[.uint32, Engine=expert_usage_stats.Engine, address_space=expert_usage_stats.address_space, linear_idx_type=expert_usage_stats.linear_idx_type]): Output 1D tensor holding the maximum tokens assigned to any expert and the expert count.
• ​topk_ids (TileTensor[input_type, Engine=topk_ids.Engine, address_space=topk_ids.address_space, linear_idx_type=topk_ids.linear_idx_type]): Input 1D tensor of expert IDs, one per token.
• ​context (DeviceContext): The device context.
• ​scales_offset_p (Optional[Pointer[UInt32, MutAnyOrigin]]): Optional pointer receiving the aligned scale offsets for FP8/block-scaled grouped matmul.