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 module

moe

Implements Mixture-of-Experts (MoE) routing, token dispatch, and expert computation kernels.

Functions

calculate_warp_offset: Computes warp-level write counts and per-thread offsets using warp voting.
eplb_remap: Launch the fused EPLB log->phy remap on GPU.
eplb_remap_kernel: Fused EPLB per tile_token rows of router idx; one thread per (n,k) element. Each block cooperatively caces the current layer's logcnt and log2phy slices in SMEM, then every thread does: HBM-load logical id -> SMEM-looup cnt -> int mod -> SMEM-Lookup phy id -> HBM-store.
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: Launches the MoE index creation kernel on GPU.
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_eplb: Launches the single-group MoE router with EPLB log->phy remap on GPU.
single_group_router_eplb_kernel: Single-group MoE router fused with EPLB log->phy remap.
single_group_router_kernel: Single-group MoE router kernel. One block per token, one thread per expert.

Functions​

Functions