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

moe_create_indices

def moe_create_indices[input_type: DType, //, target: StringSlice[ImmStaticOrigin], expected_count: Int = Int(8192)](token_expert_order: TileTensor[DType.uint32, Storage=token_expert_order.Storage, address_space=token_expert_order.address_space, linear_idx_type=token_expert_order.linear_idx_type], expert_start_indices: TileTensor[DType.uint32, Storage=expert_start_indices.Storage, address_space=expert_start_indices.address_space, linear_idx_type=expert_start_indices.linear_idx_type], restore_token_order: TileTensor[DType.uint32, Storage=restore_token_order.Storage, address_space=restore_token_order.address_space, linear_idx_type=restore_token_order.linear_idx_type], expert_ids: TileTensor[DType.int32, Storage=expert_ids.Storage, address_space=expert_ids.address_space, linear_idx_type=expert_ids.linear_idx_type], expert_usage_stats: TileTensor[DType.uint32, Storage=expert_usage_stats.Storage, address_space=expert_usage_stats.address_space, linear_idx_type=expert_usage_stats.linear_idx_type], topk_ids: TileTensor[input_type, Storage=topk_ids.Storage, address_space=topk_ids.address_space, linear_idx_type=topk_ids.linear_idx_type], context: DeviceContext, scales_offset_p: Optional[Pointer[UInt32, MutAnyOrigin, _safe=False]] = None)

Launches the MoE index creation kernel on GPU.

Groups tokens by their assigned expert using a bucket sort algorithm so that downstream kernels such as grouped matmul can process each expert's tokens contiguously. Allocates and zero-initializes the atomic lock buffer and expert usage stats, reshapes topk_ids to 2D, and launches one block per expert.

Parameters:

input_type (DType): DType of the topk_ids tensor.
target (StringSlice[ImmStaticOrigin]): The target device to run the kernel on.
expected_count (Int): Maximum number of token indices cached per expert in shared memory before spilling to global memory.

Args:

token_expert_order (TileTensor[DType.uint32, Storage=token_expert_order.Storage, 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[DType.uint32, Storage=expert_start_indices.Storage, 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[DType.uint32, Storage=restore_token_order.Storage, 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[DType.int32, Storage=expert_ids.Storage, address_space=expert_ids.address_space, linear_idx_type=expert_ids.linear_idx_type]): Output 1D tensor of the active expert IDs in output order.
expert_usage_stats (TileTensor[DType.uint32, Storage=expert_usage_stats.Storage, 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 count of active experts.
topk_ids (TileTensor[input_type, Storage=topk_ids.Storage, 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, _safe=False]]): Optional pointer receiving the aligned scale offsets for FP8/block-scaled grouped matmul.