Skip to main content

Mojo function

moe_create_indices

moe_create_indices[input_type: DType, //, target: StringSlice[StaticConstantOrigin], expected_count: Int = 8192](token_expert_order: TileTensor[DType.uint32, token_expert_order.LayoutType, token_expert_order.origin, address_space=token_expert_order.address_space, linear_idx_type=token_expert_order.linear_idx_type, element_size=token_expert_order.element_size], expert_start_indices: TileTensor[DType.uint32, expert_start_indices.LayoutType, expert_start_indices.origin, address_space=expert_start_indices.address_space, linear_idx_type=expert_start_indices.linear_idx_type, element_size=expert_start_indices.element_size], restore_token_order: TileTensor[DType.uint32, restore_token_order.LayoutType, restore_token_order.origin, address_space=restore_token_order.address_space, linear_idx_type=restore_token_order.linear_idx_type, element_size=restore_token_order.element_size], expert_ids: TileTensor[DType.int32, expert_ids.LayoutType, expert_ids.origin, address_space=expert_ids.address_space, linear_idx_type=expert_ids.linear_idx_type, element_size=expert_ids.element_size], expert_usage_stats: TileTensor[DType.uint32, expert_usage_stats.LayoutType, expert_usage_stats.origin, address_space=expert_usage_stats.address_space, linear_idx_type=expert_usage_stats.linear_idx_type, element_size=expert_usage_stats.element_size], topk_ids: TileTensor[input_type, topk_ids.LayoutType, topk_ids.origin, address_space=topk_ids.address_space, linear_idx_type=topk_ids.linear_idx_type, element_size=topk_ids.element_size], context: DeviceContextPtr)

Was this page helpful?