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

eplb_remap_kernel

def eplb_remap_kernel[PhyIdxLayoutType: TensorLayout, RouterIdxLayoutType: TensorLayout, LogcntLayoutType: TensorLayout, Log2phyLayoutType: TensorLayout, LayerIdxLayoutType: TensorLayout, num_log: Int, max_replicas: Int, K: Int, tile_tokens: Int, hash_decorrelate: Bool](phy_idx: TileTensor[DType.int32, PhyIdxLayoutType, MutAnyOrigin], router_idx: TileTensor[DType.int32, RouterIdxLayoutType, ImmutAnyOrigin], logcnt: TileTensor[DType.int32, LogcntLayoutType, ImmutAnyOrigin], log2phy: TileTensor[DType.int32, Log2phyLayoutType, ImmutAnyOrigin], layer_idx: TileTensor[DType.int32, LayerIdxLayoutType, ImmutAnyOrigin])

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.

Portable across all hardwares.

Optimality of choosing : hash_decorrelate=True xor-hashes the flat position with a Knuth multiplicative hash of the logical id before the modulo, breaking structured position-vs-cnt alignment without warp ops.

Parameters:

PhyIdxLayoutType (TensorLayout): TensorLayout of the phy_idx output tensor.
RouterIdxLayoutType (TensorLayout): TensorLayout of the router_idx input tensor.
LogcntLayoutType (TensorLayout): TensorLayout of the logcnt input tensor.
Log2phyLayoutType (TensorLayout): TensorLayout of the log2phy input tensor.
LayerIdxLayoutType (TensorLayout): TensorLayout of the layer_idx input tensor.
num_log (Int): Number of logical experts per MoE layer.
max_replicas (Int): Maximum number of physical replicas per logical expert.
K (Int): Number of top-K experts selected per token. Must be a power of two so tid % K is a bitmask.
tile_tokens (Int): Number of router_idx rows processed per block. Block size is tile_tokens * K threads.
hash_decorrelate (Bool): If True, xor-hash the flat position with a Knuth multiplicative hash of the logical id before the modulo to break structured-position bias. If False, use plain pos % cnt.

Args:

phy_idx (TileTensor[DType.int32, PhyIdxLayoutType, MutAnyOrigin]): Output [num_tokens, K] tensor of physical expert IDs after EPLB remap.
router_idx (TileTensor[DType.int32, RouterIdxLayoutType, ImmutAnyOrigin]): Input [num_tokens, K] tensor of logical expert IDs from the gate.
logcnt (TileTensor[DType.int32, LogcntLayoutType, ImmutAnyOrigin]): Input [num_moe_layers, num_log] tensor of replica counts per (layer, logical expert).
log2phy (TileTensor[DType.int32, Log2phyLayoutType, ImmutAnyOrigin]): Input [num_moe_layers, num_log, max_replicas] tensor of physical-id lookup table entries.
layer_idx (TileTensor[DType.int32, LayerIdxLayoutType, ImmutAnyOrigin]): Input rank-1 [1] scalar tensor carrying the current MoE layer index.