Skip to main content
Log in

Mojo function

moe_create_indices_kernel

moe_create_indices_kernel[input_type: DType, num_threads: Int, token_expert_order_layout: Layout, expert_start_indices_layout: Layout, restore_token_order_layout: Layout, expert_ids_layout: Layout, expert_usage_stats_layout: Layout, indices_padded_layout: Layout, padded_input_layout: Layout, topk_ids_layout: Layout](token_expert_order: LayoutTensor[uint32, token_expert_order_layout, MutableAnyOrigin], expert_start_indices: LayoutTensor[uint32, expert_start_indices_layout, MutableAnyOrigin], restore_token_order: LayoutTensor[uint32, restore_token_order_layout, MutableAnyOrigin], expert_ids: LayoutTensor[uint32, expert_ids_layout, MutableAnyOrigin], expert_usage_stats: LayoutTensor[uint32, expert_usage_stats_layout, MutableAnyOrigin], indices_padded: LayoutTensor[uint32, indices_padded_layout, MutableAnyOrigin], topk_ids_padded: LayoutTensor[input_type, padded_input_layout, MutableAnyOrigin], topk_ids: LayoutTensor[input_type, topk_ids_layout, MutableAnyOrigin])