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

single_group_router_eplb_kernel

def single_group_router_eplb_kernel[scores_type: DType, bias_type: DType, ExpertIndicesLayoutType: TensorLayout, ExpertIndicesLogLayoutType: TensorLayout, ExpertWeightsLayoutType: TensorLayout, ExpertScoresLayoutType: TensorLayout, ExpertBiasLayoutType: TensorLayout, LogcntLayoutType: TensorLayout, Log2phyLayoutType: TensorLayout, LayerIdxLayoutType: TensorLayout, n_routed_experts: Int, n_experts_per_tok: Int, norm_weights: Bool, num_threads: Int, num_log: Int, max_replicas: Int, hash_decorrelate: Bool, scores_input_fn: OptionalReg[def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[scores_type, width]] = None](expert_indices: TileTensor[DType.int32, ExpertIndicesLayoutType, MutAnyOrigin], expert_indices_log: TileTensor[DType.int32, ExpertIndicesLogLayoutType, MutAnyOrigin], expert_weights: TileTensor[scores_type, ExpertWeightsLayoutType, MutAnyOrigin], expert_scores: TileTensor[scores_type, ExpertScoresLayoutType, ImmutAnyOrigin], expert_bias: TileTensor[bias_type, ExpertBiasLayoutType, ImmutAnyOrigin], logcnt: TileTensor[DType.int32, LogcntLayoutType, ImmutAnyOrigin], log2phy: TileTensor[DType.int32, Log2phyLayoutType, ImmutAnyOrigin], layer_idx: TileTensor[DType.int32, LayerIdxLayoutType, ImmutAnyOrigin], routed_scaling_factor: Float32)

Single-group MoE router fused with EPLB log->phy remap.

Mirrors single_group_router_kernel exactly through phase 3, then at the K writers performs the EPLB lookup using a per-block SMEM cache of the current layer's logcnt/log2phy slice.

Backend specialization:

NVIDIA: cp.async issues the table fetch up front; sort hides the latency.
AMD/Apple: plain ld_global into registers up front, ds_write later.