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

def eplb_remap[num_log: Int, max_replicas: Int, K: Int, hash_decorrelate: Bool, target: StringSlice[ImmStaticOrigin]](phy_idx: TileTensor[DType.int32, Storage=phy_idx.Storage, address_space=phy_idx.address_space, linear_idx_type=phy_idx.linear_idx_type], router_idx: TileTensor[DType.int32, Storage=router_idx.Storage, address_space=router_idx.address_space, linear_idx_type=router_idx.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], context: DeviceContext)

Launch the fused EPLB log->phy remap on GPU.

One block per tile_tokens rows of router_idx; one thread per (n, k) element.

Parameters:

num_log (Int): Number of logical experts per layer.
max_replicas (Int): Maximum physical replicas per logical expert.
K (Int): Top-K experts per token. Must be a power of two.
hash_decorrelate (Bool): If True, xor-hash the position before the modulo to break structured-position bias in replica selection. If False, preserves the exact pos % cnt semantics of the legacy chain.
target (StringSlice[ImmStaticOrigin]): The target device to run the kernel on.

Args:

phy_idx (TileTensor[DType.int32, Storage=phy_idx.Storage, address_space=phy_idx.address_space, linear_idx_type=phy_idx.linear_idx_type]): Output physical expert ids. Shape: [num_tokens, K].
router_idx (TileTensor[DType.int32, Storage=router_idx.Storage, address_space=router_idx.address_space, linear_idx_type=router_idx.linear_idx_type]): Input logical expert ids from the gate. Shape: [num_tokens, K].
logcnt (TileTensor[DType.int32, Storage=logcnt.Storage, address_space=logcnt.address_space, linear_idx_type=logcnt.linear_idx_type]): Per-(layer, logical) replica count. Shape: [num_moe_layers, num_log].
log2phy (TileTensor[DType.int32, Storage=log2phy.Storage, address_space=log2phy.address_space, linear_idx_type=log2phy.linear_idx_type]): Per-(layer, logical, replica) physical-id table. Shape: [num_moe_layers, num_log, max_replicas].
layer_idx (TileTensor[DType.int32, Storage=layer_idx.Storage, address_space=layer_idx.address_space, linear_idx_type=layer_idx.linear_idx_type]): Rank-1 scalar tensor of shape [1] carrying the current MoE layer index. Sits on the same device as router_idx.
context (DeviceContext): DeviceContext.