Mojo module

moe

Functions

bitonic_sort_step:
calculate_warp_offset:
fill_zero_kernel:
group_limited_router_kernel: A manually fused MoE router with the group-limited strategy. It divides all the experts into n_groups groups and then finds the top topk_group groups with the highest scores. The final experts for each token are selected from the experts in the selected groups. The bias will be applied to the scores during the selection process, but the final weights will not include the bias.
moe_create_indices:
moe_create_indices_bucket_group_kernel: Create indices for MoE routing using bucket sort algorithm.
router_group_limited: A manually fused MoE router with the group-limited strategy.
single_group_router: Launch the single-group MoE router on GPU.
single_group_router_kernel: Single-group MoE router kernel. One block per token, one thread per expert.