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

def single_group_router_eplb[scores_type: DType, bias_type: DType, //, n_routed_experts: Int, n_experts_per_tok: Int, norm_weights: Bool, num_log: Int, max_replicas: Int, hash_decorrelate: Bool, target: StringSlice[ImmStaticOrigin], scores_input_fn: OptionalReg[def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[scores_type, width]] = None](expert_indices: TileTensor[DType.int32, Storage=expert_indices.Storage, address_space=expert_indices.address_space, linear_idx_type=expert_indices.linear_idx_type], expert_indices_log: TileTensor[DType.int32, Storage=expert_indices_log.Storage, address_space=expert_indices_log.address_space, linear_idx_type=expert_indices_log.linear_idx_type], expert_weights: TileTensor[scores_type, Storage=expert_weights.Storage, address_space=expert_weights.address_space, linear_idx_type=expert_weights.linear_idx_type], expert_scores: TileTensor[scores_type, Storage=expert_scores.Storage, address_space=expert_scores.address_space, linear_idx_type=expert_scores.linear_idx_type], expert_bias: TileTensor[bias_type, Storage=expert_bias.Storage, address_space=expert_bias.address_space, linear_idx_type=expert_bias.linear_idx_type], logcnt: TileTensor[DType.int32, Storage=logcnt.Storage, address_space=logcnt.address_space, linear_idx_type=logcnt.linear_idx_type], log2phy: TileTensor[DType.int32, Storage=log2phy.Storage, address_space=log2phy.address_space, linear_idx_type=log2phy.linear_idx_type], layer_idx: TileTensor[DType.int32, Storage=layer_idx.Storage, address_space=layer_idx.address_space, linear_idx_type=layer_idx.linear_idx_type], routed_scaling_factor: Float32, context: DeviceContext)

Launches the single-group MoE router with EPLB log->phy remap on GPU.

Selects the top n_experts_per_tok experts per token using warp-bitonic sort (2 or 3 phases depending on warp size), then remaps each selected logical expert ID to a physical expert ID via the per-layer logcnt and log2phy tables. One block is launched per token.

Inputs: expert_indices: Output physical expert IDs. Shape: [num_tokens, n_experts_per_tok]. expert_indices_log: Output logical expert IDs for EPLB histogram. Shape: [num_tokens, n_experts_per_tok]. expert_weights: Output expert weights. Shape: [num_tokens, n_experts_per_tok]. expert_scores: Input routing scores. Shape: [num_tokens, n_routed_experts]. expert_bias: Per-expert correction bias used for selection only. logcnt: Per-(layer, logical) replica count. Shape: [num_layers, num_log]. log2phy: Per-(layer, logical, replica) physical-ID table. Shape: [num_layers, num_log, max_replicas]. layer_idx: Rank-1 scalar tensor carrying the current MoE layer index. routed_scaling_factor: Scalar multiplied into every output weight. context: The device context.

Parameters:

scores_type (DType): DType of the routing scores and output weights.
bias_type (DType): DType of the expert correction bias.
n_routed_experts (Int): Total number of routed experts.
n_experts_per_tok (Int): Experts selected per token, must be a power of two.
norm_weights (Bool): If True, normalize selected weights to sum to 1 before applying routed_scaling_factor.
num_log (Int): Number of logical experts per layer.
max_replicas (Int): Maximum number of physical replicas per logical expert.
hash_decorrelate (Bool): If True, xor-hash the flat position with a Knuth multiplicative hash before the modulo to break structured-position bias in replica selection.
target (StringSlice[ImmStaticOrigin]): The target device to run the kernel on.
scores_input_fn (OptionalReg[def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[scores_type, width]]): Optional fused input lambda to load scores. If None, scores are loaded directly from expert_scores.