chore: import upstream snapshot with attribution

python/sglang/srt/eplb/eplb_algorithms/__init__.py

@@ -0,0 +1,87 @@
+from enum import Enum, auto
+from typing import Optional
+
+import torch
+
+from sglang.srt.eplb.eplb_algorithms import deepseek, deepseek_vec, elasticity_aware
+
+
+class EplbAlgorithm(Enum):
+    deepseek = auto()
+    deepseek_hierarchical = auto()
+    deepseek_vec = auto()
+    deepseek_vec_hierarchical = auto()
+    elasticity_aware = auto()
+    elasticity_aware_hierarchical = auto()
+    # TODO may have more algorithm later
+
+
+def rebalance_experts(
+    tokens_per_expert: torch.Tensor,
+    num_physical_experts: int,
+    num_local_physical_experts: int,
+    num_groups: Optional[int],
+    num_nodes: int,
+    algorithm: EplbAlgorithm,
+):
+    if algorithm in [EplbAlgorithm.deepseek, EplbAlgorithm.deepseek_hierarchical]:
+        return deepseek.rebalance_experts(
+            weight=tokens_per_expert.sum(dim=0),
+            num_replicas=num_physical_experts,
+            num_groups=num_groups,
+            num_nodes=num_nodes,
+            num_gpus=num_physical_experts // num_local_physical_experts,
+            enable_hierarchical=algorithm == EplbAlgorithm.deepseek_hierarchical,
+        )
+
+    if algorithm in [
+        EplbAlgorithm.deepseek_vec,
+        EplbAlgorithm.deepseek_vec_hierarchical,
+    ]:
+        return deepseek_vec.rebalance_experts(
+            tokens_per_expert=tokens_per_expert,
+            num_physical_experts=num_physical_experts,
+            num_local_physical_experts=num_local_physical_experts,
+            num_groups=num_groups,
+            num_nodes=num_nodes,
+            enable_hierarchical=algorithm == EplbAlgorithm.deepseek_vec_hierarchical,
+        )
+
+    if algorithm in [
+        EplbAlgorithm.elasticity_aware,
+        EplbAlgorithm.elasticity_aware_hierarchical,
+    ]:
+        from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
+
+        return elasticity_aware.rebalance_experts(
+            weight=tokens_per_expert.sum(dim=0),
+            num_replicas=num_physical_experts,
+            num_groups=num_groups,
+            num_nodes=num_nodes,
+            num_gpus=num_physical_experts // num_local_physical_experts,
+            enable_hierarchical=(
+                algorithm == EplbAlgorithm.elasticity_aware_hierarchical
+            ),
+            active_ranks=(
+                ElasticEPStateManager.instance().active_ranks
+                if ElasticEPStateManager.instance() is not None
+                else ElasticEPStateManager.healthy_rank_state()
+            ),
+        )
+
+    raise NotImplementedError
+
+
+def compute_algorithm(
+    raw_algorithm: str,
+    num_groups: Optional[int],
+    num_nodes: int,
+) -> EplbAlgorithm:
+    if raw_algorithm != "auto":
+        return EplbAlgorithm[raw_algorithm]
+
+    # TODO test on real scenarios and know which ones perform better
+    if (num_groups is not None) and (num_groups % num_nodes == 0):
+        return EplbAlgorithm.deepseek_hierarchical
+    else:
+        return EplbAlgorithm.deepseek

python/sglang/srt/eplb/eplb_algorithms/deepseek.py

@@ -0,0 +1,224 @@
+# This file is copied from https://github.com/deepseek-ai/EPLB/blob/main/eplb.py since that one is not a pypi package
+from typing import Tuple
+
+import torch
+
+
+def balanced_packing(
+    weight: torch.Tensor, num_packs: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Pack n weighted objects to m packs, such that each bin contains exactly n/m objects and the weights of all packs
+    are as balanced as possible.
+
+    Parameters:
+        weight: [X, n], the weight of each item
+        num_packs: number of packs
+
+    Returns:
+        pack_index: [X, n], the pack index of each item
+        rank_in_pack: [X, n], the rank of the item in the pack
+    """
+    num_layers, num_groups = weight.shape
+    assert num_groups % num_packs == 0
+    groups_per_pack = num_groups // num_packs
+
+    if groups_per_pack == 1:
+        pack_index = torch.arange(
+            weight.size(-1), dtype=torch.int64, device=weight.device
+        ).expand(weight.shape)
+        rank_in_pack = torch.zeros_like(weight, dtype=torch.int64)
+        return pack_index, rank_in_pack
+
+    indices_list = weight.float().sort(-1, descending=True).indices.tolist()
+    weight_list = weight.tolist()
+    pack_index_list = [[-1] * num_groups for _ in range(num_layers)]
+    rank_in_pack_list = [[-1] * num_groups for _ in range(num_layers)]
+    for i in range(num_layers):
+        pack_weights = [0] * num_packs
+        pack_items = [0] * num_packs
+        for group in indices_list[i]:
+            pack = min(
+                (j for j in range(num_packs) if pack_items[j] < groups_per_pack),
+                key=pack_weights.__getitem__,
+            )
+            assert pack_items[pack] < groups_per_pack
+            pack_index_list[i][group] = pack
+            rank_in_pack_list[i][group] = pack_items[pack]
+            pack_weights[pack] += weight_list[i][group]
+            pack_items[pack] += 1
+    pack_index = torch.tensor(pack_index_list, dtype=torch.int64, device="cpu")
+    rank_in_pack = torch.tensor(rank_in_pack_list, dtype=torch.int64, device="cpu")
+    return pack_index, rank_in_pack
+
+
+def replicate_experts(
+    weight: torch.Tensor, num_phy: int
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Replicate `num_log` experts to `num_phy` replicas, such that the maximum load of all replicas is minimized.
+
+    Parameters:
+        weight: [X, num_log]
+        num_phy: total number of experts after replication
+
+    Returns:
+        phy2log: [X, num_phy], logical expert id of each physical expert
+        rank: [X, num_phy], the replica rank
+        logcnt: [X, num_log], number of replicas for each logical expert
+    """
+    n, num_log = weight.shape
+    num_redundant = num_phy - num_log
+    assert num_redundant >= 0
+    device = weight.device
+    phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)
+    rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
+    logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
+    arangen = torch.arange(n, dtype=torch.int64, device=device)
+    for i in range(num_log, num_phy):
+        redundant_indices = (weight / logcnt).max(dim=-1).indices
+        phy2log[:, i] = redundant_indices
+        rank[:, i] = logcnt[arangen, redundant_indices]
+        logcnt[arangen, redundant_indices] += 1
+    return phy2log, rank, logcnt
+
+
+def rebalance_experts_hierarchical(
+    weight: torch.Tensor,
+    num_physical_experts: int,
+    num_groups: int,
+    num_nodes: int,
+    num_gpus: int,
+):
+    """
+    Parameters:
+        weight: [num_moe_layers, num_logical_experts]
+        num_physical_experts: number of physical experts after replication
+        num_groups: number of expert groups
+        num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
+        num_gpus: number of GPUs, must be a multiple of `num_nodes`
+
+    Returns:
+        physical_to_logical_map: [num_moe_layers, num_physical_experts]
+        logical_to_physical_map: [num_moe_layers, num_logical_experts, X]
+        logical_count: [num_moe_layers, num_logical_experts]
+    """
+    num_layers, num_logical_experts = weight.shape
+    assert num_logical_experts % num_groups == 0
+    group_size = num_logical_experts // num_groups
+    assert num_groups % num_nodes == 0
+    groups_per_node = num_groups // num_nodes
+    assert num_gpus % num_nodes == 0
+    assert num_physical_experts % num_gpus == 0
+    phy_experts_per_gpu = num_physical_experts // num_gpus
+
+    def inverse(perm: torch.Tensor) -> torch.Tensor:
+        inv = torch.empty_like(perm)
+        inv.scatter_(
+            1,
+            perm,
+            torch.arange(perm.size(1), dtype=torch.int64, device=perm.device).expand(
+                perm.shape
+            ),
+        )
+        return inv
+
+    # Step 1: pack groups to nodes
+    tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
+    group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes)
+    log2mlog = (
+        (
+            (group_pack_index * groups_per_node + group_rank_in_pack) * group_size
+        ).unsqueeze(-1)
+        + torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device)
+    ).flatten(-2)
+    mlog2log = inverse(log2mlog)
+
+    # Step 2: construct redundant experts within nodes
+    # [num_layers * num_nodes, num_logical_experts // num_nodes]
+    tokens_per_mlog = weight.gather(-1, mlog2log).view(
+        -1, num_logical_experts // num_nodes
+    )
+    phy2mlog, phyrank, mlogcnt = replicate_experts(
+        tokens_per_mlog, num_physical_experts // num_nodes
+    )
+
+    # Step 3: pack physical_experts to GPUs
+    # [num_layers * num_nodes, num_physical_experts // num_nodes]
+    tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
+    pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes)
+    phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
+    pphy2phy = inverse(phy2pphy)
+
+    pphy2mlog = phy2mlog.gather(
+        -1, pphy2phy
+    )  # [num_layers * num_nodes, num_log_per_nodes]
+    pphy2mlog = (
+        pphy2mlog.view(num_layers, num_nodes, -1)
+        + torch.arange(
+            0,
+            num_logical_experts,
+            num_logical_experts // num_nodes,
+            device=group_pack_index.device,
+        ).view(1, -1, 1)
+    ).flatten(-2)
+    pphy2log = mlog2log.gather(-1, pphy2mlog)
+    pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
+    logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
+    return pphy2log, pphyrank, logcnt
+
+
+def rebalance_experts(
+    weight: torch.Tensor,
+    num_replicas: int,
+    num_groups: int,
+    num_nodes: int,
+    num_gpus: int,
+    enable_hierarchical: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Entry point for expert-parallelism load balancer.
+
+    Parameters:
+        weight: [layers, num_logical_experts], the load statistics for all logical experts
+        num_replicas: number of physical experts, must be a multiple of `num_gpus`
+        num_groups: number of expert groups
+        num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
+        num_gpus: number of GPUs, must be a multiple of `num_nodes`
+
+    Returns:
+        physical_to_logical_map: [layers, num_replicas], the expert index of each replica
+        logical_to_physical_map: [layers, num_logical_experts, X], the replica indices for each expert
+        expert_count: [layers, num_logical_experts], number of physical replicas for each logical expert
+    """
+
+    num_layers, num_logical_experts = weight.shape
+    weight = weight.float().cpu()
+    if enable_hierarchical:
+        # use hierarchical load-balance policy
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight, num_replicas, num_groups, num_nodes, num_gpus
+        )
+    else:
+        # use global load-balance policy
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight, num_replicas, 1, 1, num_gpus
+        )
+    maxlogcnt = logcnt.max().item()
+    log2phy: torch.Tensor = torch.full(
+        (num_layers, num_logical_experts, maxlogcnt),
+        -1,
+        dtype=torch.int64,
+        device=logcnt.device,
+    )
+    log2phy.view(num_layers, -1).scatter_(
+        -1,
+        phy2log * maxlogcnt + phyrank,
+        torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(
+            num_layers, -1
+        ),
+    )
+    return phy2log, log2phy, logcnt
+
+
+__all__ = ["rebalance_experts"]

python/sglang/srt/eplb/eplb_algorithms/deepseek_vec.py

@@ -0,0 +1,276 @@
+# This file is copied from https://github.com/deepseek-ai/EPLB/blob/main/eplb.py since that one is not a pypi package
+from typing import Optional, Tuple
+
+import torch
+
+
+def pack_groups(tokens_per_group: torch.Tensor, num_nodes: int) -> torch.Tensor:
+    num_layers, num_groups = tokens_per_group.shape
+    assert num_groups % num_nodes == 0
+    groups_per_rank = num_groups // num_nodes
+
+    indices = tokens_per_group.float().sort(-1, descending=True).indices.cpu()
+    ret = torch.full_like(
+        tokens_per_group, fill_value=-1, dtype=torch.int64, device="cpu"
+    )
+    for layer in range(num_layers):
+        node_tokens = [0] * num_nodes
+        node_groups = [0] * num_nodes
+        for group in indices[layer]:
+
+            def key_func(rank: int) -> int:
+                if node_groups[rank] >= groups_per_rank:
+                    return 1, 0
+                else:
+                    return 0, node_tokens[rank]
+
+            rank = min(range(num_nodes), key=key_func)
+            assert node_groups[rank] < groups_per_rank
+            ret[layer, group] = rank * groups_per_rank + node_groups[rank]
+            node_tokens[rank] += tokens_per_group[layer, group]
+            node_groups[rank] += 1
+    return ret
+
+
+def make_redundant_experts_chunkwise(
+    tokens_per_expert: torch.Tensor,
+    num_physical_experts: int,
+    num_local_physical_experts: int,
+    num_physical_experts_per_chunk: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    num_steps, num_moe_layers, num_logical_experts = tokens_per_expert.shape
+    num_redundancy_experts = num_physical_experts - num_logical_experts
+
+    physical_to_logical_map = torch.empty(
+        num_moe_layers,
+        num_physical_experts,
+        dtype=torch.int,
+        device=tokens_per_expert.device,
+    )
+    logical_to_physical_map = torch.full(
+        (num_moe_layers, num_logical_experts, num_redundancy_experts + 1),
+        -1,
+        dtype=torch.int,
+        device=tokens_per_expert.device,
+    )
+    logical_count = torch.ones(
+        num_moe_layers,
+        num_logical_experts,
+        dtype=torch.int,
+        device=tokens_per_expert.device,
+    )
+
+    assert num_physical_experts % num_physical_experts_per_chunk == 0
+    num_chunks = num_physical_experts // num_physical_experts_per_chunk
+    assert num_logical_experts % num_chunks == 0
+    num_logical_experts_per_group = num_logical_experts // num_chunks
+    assert num_redundancy_experts % num_chunks == 0
+    num_redundancy_experts_per_group = num_redundancy_experts // num_chunks
+
+    arange_num_moe_layers_num_groups = torch.arange(
+        num_moe_layers * num_chunks, dtype=torch.int, device=tokens_per_expert.device
+    )
+    arange_num_logical_experts = torch.arange(
+        num_logical_experts, dtype=torch.int, device=tokens_per_expert.device
+    )
+    arange_num_logical_experts_per_group = torch.arange(
+        num_logical_experts_per_group, dtype=torch.int, device=tokens_per_expert.device
+    )
+    arange_num_groups = torch.arange(
+        num_chunks, dtype=torch.int, device=tokens_per_expert.device
+    )
+    physical_to_logical_map.view(
+        num_moe_layers, num_chunks, num_physical_experts_per_chunk
+    )[:, :, :num_logical_experts_per_group] = arange_num_logical_experts.view(
+        num_chunks, num_logical_experts_per_group
+    )
+    logical_to_physical_map[:, :, 0] = (
+        arange_num_logical_experts_per_group.expand(
+            num_chunks, num_logical_experts_per_group
+        )
+        + arange_num_groups[:, None] * num_physical_experts_per_chunk
+    ).view(num_logical_experts)
+
+    tokens_per_expert_all_diff = tokens_per_expert + arange_num_logical_experts * 1e-4
+    for i in range(num_redundancy_experts_per_group):
+        score = (
+            tokens_per_expert_all_diff / logical_count
+        )  # NOTE: Values in score must be different from each other
+        score1 = tokens_per_expert / (logical_count + 1)
+        score = score.view(
+            num_steps, num_moe_layers, num_chunks, num_logical_experts_per_group
+        )
+        score1 = score1.view_as(score)
+        values, indices = score.max(-1, keepdim=True)
+        values = values.expand_as(score).contiguous()
+        score.scatter_(-1, indices, score1.gather(-1, indices))
+        values.scatter_(-1, indices, score.max(-1, keepdim=True).values)
+        redundancy_indices = values.sum(0).argmin(-1)
+        physical_to_logical_map.view(
+            num_moe_layers, num_chunks, num_physical_experts_per_chunk
+        )[:, :, num_logical_experts_per_group + i] = (
+            redundancy_indices + arange_num_groups * num_logical_experts_per_group
+        )
+        redundancy_count = (
+            logical_count.view(
+                num_moe_layers * num_chunks, num_logical_experts_per_group
+            )
+            .gather(-1, redundancy_indices.view(num_moe_layers * num_chunks, 1))
+            .squeeze(1)
+        )
+        physical_redundancy_indices = (
+            (
+                arange_num_groups * num_physical_experts_per_chunk
+                + num_logical_experts_per_group
+                + i
+            )
+            .expand(num_moe_layers, num_chunks)
+            .flatten()
+        )
+        logical_to_physical_map.view(
+            num_moe_layers * num_chunks,
+            num_logical_experts_per_group,
+            num_redundancy_experts + 1,
+        )[
+            arange_num_moe_layers_num_groups,
+            redundancy_indices.view(num_moe_layers * num_chunks),
+            redundancy_count,
+        ] = physical_redundancy_indices
+        logical_count.view(num_moe_layers * num_chunks, num_logical_experts_per_group)[
+            arange_num_moe_layers_num_groups,
+            redundancy_indices.view(num_moe_layers * num_chunks),
+        ] += 1
+
+    if num_local_physical_experts > 1:
+        # Load-balancing between GPUs
+        physical_to_logical_map_int64 = physical_to_logical_map.to(torch.int64)
+        counts = logical_count.gather(-1, physical_to_logical_map_int64)
+        score = tokens_per_expert.sum(0).gather(-1, physical_to_logical_map_int64)
+        score = score / counts
+        score = score.view(num_moe_layers, num_chunks, num_physical_experts_per_chunk)
+        indices = score.argsort(-1, descending=True)
+        indices += torch.arange(
+            0,
+            num_physical_experts,
+            num_physical_experts_per_chunk,
+            dtype=indices.dtype,
+            device=indices.device,
+        )[None, :, None]
+
+        assert num_physical_experts_per_chunk % num_local_physical_experts == 0
+        num_local_groups = num_physical_experts_per_chunk // num_local_physical_experts
+        indices = indices.view(
+            num_moe_layers, num_chunks, num_local_physical_experts, num_local_groups
+        )
+        indices[:, :, 1::2, :] = indices[:, :, 1::2, :].flip(-1)
+        indices = indices.transpose(2, 3)
+        indices = indices.reshape(num_moe_layers, num_physical_experts)
+        physical_to_logical_map = physical_to_logical_map.gather(-1, indices)
+        mask = logical_to_physical_map == -1
+        logical_to_physical_map[mask] = 0
+        logical_to_physical_map = (
+            indices.argsort(-1)
+            .gather(
+                -1, logical_to_physical_map.view(num_moe_layers, -1).to(torch.int64)
+            )
+            .view_as(logical_to_physical_map)
+            .to(torch.int)
+        )
+        logical_to_physical_map[mask] = -1
+
+    return physical_to_logical_map, logical_to_physical_map, logical_count
+
+
+def decode_rebalance_experts(
+    tokens_per_expert: torch.Tensor,
+    num_physical_experts: int,
+    num_local_physical_experts: int,
+):
+    return make_redundant_experts_chunkwise(
+        tokens_per_expert,
+        num_physical_experts,
+        num_local_physical_experts,
+        num_physical_experts,
+    )
+
+
+def prefill_rebalance_experts(
+    tokens_per_expert: torch.Tensor,
+    num_physical_experts: int,
+    num_local_physical_experts: int,
+    num_groups: int,
+    num_nodes: int,
+):
+    tokens_per_expert = tokens_per_expert.float().cpu()
+
+    num_steps, _, num_logical_experts = tokens_per_expert.shape
+    assert num_logical_experts % num_groups == 0
+    group_size = num_logical_experts // num_groups
+    assert num_groups % num_nodes == 0, f"{num_groups=} {num_nodes=}"
+
+    tokens_per_group = tokens_per_expert.sum(0).unflatten(-1, (num_groups, -1)).sum(-1)
+    group_perm = pack_groups(
+        tokens_per_group, num_nodes
+    )  # [num_moe_layers, num_groups] => [num_moe_layers, num_nodes]
+
+    # log2mlog [layers, #logexp] -> [layers, #logexp]
+    log2mlog = (
+        (group_perm * group_size).unsqueeze(-1)
+        + torch.arange(group_size, dtype=torch.int64, device=group_perm.device)
+    ).flatten(-2)
+
+    # mlog2log [layers, #logexp] -> [layers, #logexp], inverse of log2mlog
+    mlog2log = torch.empty_like(log2mlog)
+    arange = torch.arange(
+        num_logical_experts, dtype=torch.int64, device=mlog2log.device
+    )
+    mlog2log.scatter_(1, log2mlog, arange.expand(log2mlog.size(0), -1))
+
+    # tokens_per_mlog[i][j][k] = tokens_per_expert[i][j][mlog2log[j][k]]
+    tokens_per_mlog = tokens_per_expert.gather(
+        2, mlog2log.unsqueeze(0).expand(num_steps, -1, -1)
+    )
+
+    phy2mlog, mlog2phy, mlog_count = make_redundant_experts_chunkwise(
+        tokens_per_mlog,
+        num_physical_experts,
+        num_local_physical_experts,
+        num_physical_experts // num_nodes,
+    )
+
+    # phy2log[i][j] = mlog2log[i][phy2mlog[i][j]]
+    phy2log = mlog2log.gather(1, phy2mlog.to(torch.int64))
+
+    # mlog2phy: [num_moe_layers, num_logical_experts, ...]
+    # log2phy[i][j][k] = mlog2phy[i][log2mlog[i][j]][k]
+    log2phy = mlog2phy.gather(
+        1, log2mlog.unsqueeze(-1).expand(-1, -1, mlog2phy.size(-1)).to(torch.int64)
+    )
+
+    # log_count[i][j] = mlog_count[i][log2mlog[i][j]]
+    log_count = mlog_count.gather(1, log2mlog)
+    return phy2log, log2phy, log_count
+
+
+def rebalance_experts(
+    tokens_per_expert: torch.Tensor,
+    num_physical_experts: int,
+    num_local_physical_experts: int,
+    num_groups: Optional[int],
+    num_nodes: int,
+    enable_hierarchical: bool,
+):
+    if enable_hierarchical:
+        return prefill_rebalance_experts(
+            tokens_per_expert=tokens_per_expert,
+            num_physical_experts=num_physical_experts,
+            num_local_physical_experts=num_local_physical_experts,
+            num_groups=num_groups,
+            num_nodes=num_nodes,
+        )
+    else:
+        return decode_rebalance_experts(
+            tokens_per_expert=tokens_per_expert,
+            num_physical_experts=num_physical_experts,
+            num_local_physical_experts=num_local_physical_experts,
+        )

python/sglang/srt/eplb/eplb_algorithms/elasticity_aware.py

@@ -0,0 +1,87 @@
+from typing import Tuple
+
+import torch
+
+from sglang.srt.eplb.eplb_algorithms.deepseek import rebalance_experts_hierarchical
+
+
+def rebalance_experts(
+    weight: torch.Tensor,
+    num_replicas: int,
+    num_groups: int,
+    num_nodes: int,
+    num_gpus: int,
+    enable_hierarchical: bool,
+    active_ranks: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Entry point for expert-parallelism load balancer.
+
+    Parameters:
+        weight: [layers, num_logical_experts], the load statistics for all logical experts
+        num_replicas: number of physical experts, must be a multiple of `num_gpus`
+        num_groups: number of expert groups
+        num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
+        num_gpus: number of GPUs, must be a multiple of `num_nodes`
+
+    Returns:
+        physical_to_logical_map: [layers, num_replicas], the expert index of each replica
+        logical_to_physical_map: [layers, num_logical_experts, X], the replica indices for each expert
+        expert_count: [layers, num_logical_experts], number of physical replicas for each logical expert
+    """
+
+    num_layers, num_logical_experts = weight.shape
+    weight = weight.float().cpu()
+    num_active_ranks = active_ranks.sum().item()
+    num_local_experts = num_replicas // num_gpus
+    if num_active_ranks < num_gpus:
+        # Must fall back to global load-balance policy
+        # and fix some params
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight,
+            num_local_experts * num_active_ranks,
+            1,
+            1,
+            num_active_ranks,
+        )
+    elif enable_hierarchical:
+        # use hierarchical load-balance policy
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight, num_replicas, num_groups, num_nodes, num_gpus
+        )
+    else:
+        # use global load-balance policy
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight, num_replicas, 1, 1, num_gpus
+        )
+    maxlogcnt = logcnt.max().item()
+    log2phy: torch.Tensor = torch.full(
+        (num_layers, num_logical_experts, maxlogcnt),
+        -1,
+        dtype=torch.int64,
+        device=logcnt.device,
+    )
+    log2phy.view(num_layers, -1).scatter_(
+        -1,
+        phy2log * maxlogcnt + phyrank,
+        torch.arange(
+            num_local_experts * num_active_ranks,
+            dtype=torch.int64,
+            device=log2phy.device,
+        ).expand(num_layers, -1),
+    )
+    if num_active_ranks < num_gpus:
+        phy2log_slices = list(
+            phy2log.view(num_layers, num_active_ranks, -1).unbind(dim=1)
+        )
+        active_ranks_list = active_ranks.tolist()
+        for idx, active_rank in enumerate(active_ranks_list):
+            if not active_rank:
+                phy2log_slices.insert(idx, torch.zeros_like(phy2log_slices[0]))
+                log2phy = torch.where(
+                    log2phy >= idx * num_local_experts,
+                    log2phy + num_local_experts,
+                    log2phy,
+                )
+        phy2log = torch.stack(phy2log_slices, dim=1).contiguous().view(num_layers, -1)
+    return phy2log, log2phy, logcnt

python/sglang/srt/eplb/eplb_manager.py

@@ -0,0 +1,186 @@
+import logging
+import time
+from typing import TYPE_CHECKING, List
+
+import torch.cuda
+
+from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
+from sglang.srt.eplb.expert_location import (
+    ExpertLocationMetadata,
+    format_expert_location_layout,
+    format_expert_location_layout_diff,
+    get_global_expert_location_metadata,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class EPLBManager:
+    def __init__(self, model_runner: "ModelRunner"):
+        super().__init__()
+        self._model_runner = model_runner
+        self._server_args = model_runner.server_args
+        self._rebalance_layers_per_chunk = (
+            self._server_args.eplb_rebalance_layers_per_chunk
+        )
+        self._rebalance_num_iterations = self._server_args.eplb_rebalance_num_iterations
+
+        # Otherwise, the circular buffer will contain stale data. If the case is needed, it can be implemented.
+        assert (
+            self._server_args.eplb_rebalance_num_iterations
+            >= self._server_args.expert_distribution_recorder_buffer_size
+        ), "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size"
+
+        if not get_global_expert_distribution_recorder().recording:
+            get_global_expert_distribution_recorder().start_record()
+
+        logger.info(
+            f"[EPLBManager] system started, will rebalance per {self._rebalance_num_iterations} iterations."
+        )
+
+        self._main_generator = self._entrypoint()
+
+    def on_forward_pass_end(self):
+        next(self._main_generator)
+
+    def reset_generator(self):
+        self._main_generator = self._entrypoint()
+
+    # can be more complex if needed
+    def _entrypoint(self):
+        while True:
+            for _ in range(self._rebalance_num_iterations):
+                yield
+
+            yield from self.rebalance()
+
+    def rebalance(self):
+        logger.info("[EPLBManager] rebalance start")
+
+        enable_timing = self._rebalance_layers_per_chunk is None
+
+        if enable_timing:
+            torch.get_device_module().synchronize()
+            time_start = time.time()
+
+        dump_record_output = get_global_expert_distribution_recorder().dump_record(
+            output_mode="object"
+        )
+        logical_count = dump_record_output["logical_count"]
+        average_utilization_rate_over_window = dump_record_output[
+            "average_utilization_rate_over_window"
+        ]
+
+        # Check whether rebalancing is needed
+        if not self._check_rebalance_needed(average_utilization_rate_over_window):
+            return
+
+        expert_location_metadata = ExpertLocationMetadata.init_by_eplb(
+            self._server_args, self._model_runner.model_config, logical_count
+        )
+
+        update_layer_ids_chunks = self._compute_update_layer_ids_chunks()
+        all_update_layer_ids = [
+            layer_id for chunk in update_layer_ids_chunks for layer_id in chunk
+        ]
+        self._log_rebalance_layout_before_update(
+            expert_location_metadata,
+            update_layer_ids=all_update_layer_ids,
+        )
+        for chunk_layer_ids in update_layer_ids_chunks:
+            if len(update_layer_ids_chunks) > 1:
+                yield
+            self._model_runner.update_expert_location(
+                expert_location_metadata,
+                update_layer_ids=chunk_layer_ids,
+            )
+
+        self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)
+
+        msg = f"[EPLBManager] rebalance end"
+        if enable_timing:
+            torch.get_device_module().synchronize()
+            time_end = time.time()
+            msg += f" time={time_end - time_start:.3f}s"
+        logger.info(msg)
+
+    def _check_rebalance_needed(self, average_utilization_rate_over_window):
+        if average_utilization_rate_over_window is None:
+            return True
+
+        if (
+            average_utilization_rate_over_window
+            > self._server_args.eplb_min_rebalancing_utilization_threshold
+        ):
+            logger.info(
+                f"[EPLBManager] Skipped ep rebalancing: current GPU utilization {average_utilization_rate_over_window:.2f} > minimum rebalance threshold {self._server_args.eplb_min_rebalancing_utilization_threshold:.2f}"
+            )
+            return False
+
+        return True
+
+    def _compute_update_layer_ids_chunks(self) -> List[List[int]]:
+        all_layer_ids = sorted(
+            list(self._model_runner.model.routed_experts_weights_of_layer.keys())
+        )
+        chunk_size = self._rebalance_layers_per_chunk or 1000000
+        return list(_chunk_list(all_layer_ids, chunk_size=chunk_size))
+
+    def _should_log_expert_location_metadata(self) -> bool:
+        return (
+            self._model_runner.tp_rank == 0
+            and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get()
+        )
+
+    def _log_rebalance_layout_before_update(
+        self,
+        new_expert_location_metadata: ExpertLocationMetadata,
+        update_layer_ids: List[int],
+    ):
+        if not self._should_log_expert_location_metadata():
+            return
+
+        old_expert_location_metadata = get_global_expert_location_metadata()
+        logger.info(
+            "[EPLBManager] rebalance layout before:\n%s",
+            format_expert_location_layout(
+                old_expert_location_metadata,
+                layer_ids=update_layer_ids,
+            ),
+        )
+        logger.info(
+            "[EPLBManager] rebalance layout target:\n%s",
+            format_expert_location_layout(
+                new_expert_location_metadata,
+                layer_ids=update_layer_ids,
+            ),
+        )
+        logger.info(
+            "[EPLBManager] rebalance layout diff:\n%s",
+            format_expert_location_layout_diff(
+                old_expert_location_metadata,
+                new_expert_location_metadata,
+                layer_ids=update_layer_ids,
+            ),
+        )
+
+    def _log_rebalance_layout_after_update(self, update_layer_ids: List[int]):
+        if not self._should_log_expert_location_metadata():
+            return
+
+        logger.info(
+            "[EPLBManager] rebalance layout after:\n%s",
+            format_expert_location_layout(
+                get_global_expert_location_metadata(),
+                layer_ids=update_layer_ids,
+            ),
+        )
+
+
+def _chunk_list(items: List, chunk_size):
+    for start_index in range(0, len(items), chunk_size):
+        yield items[start_index : start_index + chunk_size]

python/sglang/srt/eplb/eplb_simulator/__init__.py

@@ -0,0 +1 @@
+from . import reader

python/sglang/srt/eplb/eplb_simulator/reader.py

@@ -0,0 +1,51 @@
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+
+from sglang.srt.eplb.expert_distribution import (
+    _convert_global_physical_count_to_logical_count,
+)
+
+convert_global_physical_count_to_logical_count = (
+    _convert_global_physical_count_to_logical_count
+)
+
+
+def read_mode_per_pass(dir_data: Path):
+    """Read data from ExpertDistributionRecorder when recorded with mode `per_pass`"""
+
+    # gpc := global_physical_count
+    gpc_of_forward_pass_and_rank = defaultdict(lambda: defaultdict())
+    for path in tqdm(list(dir_data.glob("*.pt"))):
+        data_pack = torch.load(path, weights_only=True)
+        last_physical_to_logical_map = data_pack["last_physical_to_logical_map"]
+        for record in data_pack["records"]:
+            forward_pass_id = record["forward_pass_id"]
+            rank = record["rank"]
+            assert (
+                gpc_of_forward_pass_and_rank[forward_pass_id].get(rank) is None
+            ), f"Duplicated {forward_pass_id=} {rank=}"
+            gpc_of_forward_pass_and_rank[forward_pass_id][rank] = record[
+                "global_physical_count"
+            ]
+
+    forward_pass_ids = sorted(gpc_of_forward_pass_and_rank.keys())
+    print(f"Make {forward_pass_ids=} into array")
+
+    items = []
+    for forward_pass_id, gpc_of_rank in sorted(gpc_of_forward_pass_and_rank.items()):
+        gpc_of_rank_tensor = torch.stack(
+            [gpc for rank, gpc in sorted(gpc_of_rank.items())]
+        ).sum(dim=0)
+        items.append(gpc_of_rank_tensor)
+
+    gpc_of_forward_pass = torch.stack(items)
+    print(f"{gpc_of_forward_pass.shape=}")
+
+    return dict(
+        global_physical_count_of_forward_pass=gpc_of_forward_pass,
+        last_physical_to_logical_map=last_physical_to_logical_map,
+        forward_pass_ids=forward_pass_ids,
+    )

python/sglang/srt/eplb/expert_location.py

@@ -0,0 +1,737 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import json
+import logging
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Iterable, List, Optional
+
+import torch
+import torch.distributed
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExpertLocationMetadata:
+    physical_to_logical_map: torch.Tensor  # (layers, num_physical_experts)
+    physical_to_logical_map_cpu: torch.Tensor
+    logical_to_all_physical_map: torch.Tensor  # (layers, num_logical_experts, X)
+    logical_to_all_physical_map_cpu: torch.Tensor  # CPU copy for performance
+    logical_to_all_physical_map_num_valid: torch.Tensor  # (layers, num_logical_experts)
+    # (layers, num_logical_experts)
+    logical_to_rank_dispatch_physical_map: Optional[torch.Tensor]
+
+    # -------------------------------- properties ------------------------------------
+
+    @property
+    def num_layers(self) -> int:
+        return self.physical_to_logical_map.shape[0]
+
+    @property
+    def num_physical_experts(self) -> int:
+        return self.physical_to_logical_map.shape[1]
+
+    @property
+    def num_local_physical_experts(self) -> int:
+        ans, remainder = divmod(self.num_physical_experts, self.ep_size)
+        assert remainder == 0
+        return ans
+
+    @property
+    def num_logical_experts(self) -> int:
+        return self.logical_to_all_physical_map.shape[1]
+
+    @property
+    def ep_size(self):
+        # TODO change when EP size != world size
+        return torch.distributed.get_world_size()
+
+    def __post_init__(self):
+        num_layers_0, num_physical_experts_0 = self.physical_to_logical_map.shape
+        num_layers_1, num_logical_experts_0, num_physical_experts_1 = (
+            self.logical_to_all_physical_map.shape
+        )
+        num_layers_2, num_logical_experts_1 = (
+            self.logical_to_all_physical_map_num_valid.shape
+        )
+        assert num_layers_0 == num_layers_1 == num_layers_2
+        assert num_logical_experts_0 == num_logical_experts_1
+        assert num_physical_experts_0 == num_physical_experts_1
+
+    # -------------------------------- construction ------------------------------------
+
+    @staticmethod
+    def init_trivial(
+        server_args: ServerArgs, model_config: ModelConfig, moe_ep_rank: int
+    ):
+        """Trivial location - logical expert i corresponds to physical expert i"""
+        common = ExpertLocationMetadata._init_common(server_args, model_config)
+
+        if common is None:
+            return None
+
+        num_physical_experts = common["num_physical_experts"]
+        model_config_for_expert_location = common["model_config_for_expert_location"]
+        num_layers = model_config_for_expert_location.num_layers
+        num_logical_experts = model_config_for_expert_location.num_logical_experts
+
+        physical_to_logical_map = (
+            torch.arange(0, num_physical_experts).repeat(num_layers, 1)
+            % num_logical_experts
+        )
+
+        return ExpertLocationMetadata.init_by_mapping(
+            server_args,
+            model_config,
+            physical_to_logical_map=physical_to_logical_map,
+            moe_ep_rank=moe_ep_rank,
+        )
+
+    @staticmethod
+    def init_by_mapping(
+        server_args: ServerArgs,
+        model_config: ModelConfig,
+        physical_to_logical_map,
+        moe_ep_rank: int = None,
+    ):
+        if not isinstance(physical_to_logical_map, torch.Tensor):
+            physical_to_logical_map = torch.tensor(physical_to_logical_map)
+        physical_to_logical_map = physical_to_logical_map.to(server_args.device)
+
+        common = ExpertLocationMetadata._init_common(server_args, model_config)
+
+        if common is None:
+            return None
+
+        model_config_for_expert_location = common["model_config_for_expert_location"]
+        logical_to_all_physical_map = _compute_logical_to_all_physical_map(
+            server_args=server_args,
+            physical_to_logical_map=physical_to_logical_map,
+            num_logical_experts=model_config_for_expert_location.num_logical_experts,
+            ep_size=common["ep_size"],
+            moe_ep_rank=moe_ep_rank,
+        )
+
+        return ExpertLocationMetadata._init_raw(
+            server_args=server_args,
+            ep_size=common["ep_size"],
+            physical_to_logical_map=physical_to_logical_map,
+            logical_to_all_physical_map=logical_to_all_physical_map,
+        )
+
+    @staticmethod
+    def init_by_eplb(
+        server_args: ServerArgs, model_config: ModelConfig, logical_count: torch.Tensor
+    ):
+        if not isinstance(logical_count, torch.Tensor):
+            logical_count = torch.tensor(logical_count)
+        if len(logical_count.shape) == 2:
+            logical_count = logical_count.unsqueeze(0)
+        logical_count = logical_count.to(server_args.device)
+
+        common = ExpertLocationMetadata._init_common(server_args, model_config)
+
+        if common is None:
+            return None
+
+        model_config_for_expert_location = common["model_config_for_expert_location"]
+        num_physical_experts = common["num_physical_experts"]
+        num_groups = model_config_for_expert_location.num_groups
+        num_nodes = server_args.nnodes
+
+        from sglang.srt.eplb import eplb_algorithms
+
+        physical_to_logical_map, logical_to_all_physical_map, expert_count = (
+            eplb_algorithms.rebalance_experts(
+                tokens_per_expert=logical_count,
+                num_physical_experts=num_physical_experts,
+                num_local_physical_experts=num_physical_experts // common["ep_size"],
+                num_groups=num_groups,
+                num_nodes=num_nodes,
+                algorithm=eplb_algorithms.compute_algorithm(
+                    raw_algorithm=server_args.eplb_algorithm,
+                    num_groups=num_groups,
+                    num_nodes=num_nodes,
+                ),
+            )
+        )
+
+        return ExpertLocationMetadata._init_raw(
+            server_args=server_args,
+            ep_size=common["ep_size"],
+            physical_to_logical_map=physical_to_logical_map.to(server_args.device),
+            logical_to_all_physical_map=logical_to_all_physical_map.to(
+                server_args.device
+            ),
+        )
+
+    @staticmethod
+    def _init_common(server_args: ServerArgs, model_config: ModelConfig):
+        model_config_for_expert_location = (
+            ModelConfigForExpertLocation.from_model_config(model_config)
+        )
+
+        if model_config_for_expert_location is None:
+            return None
+
+        num_physical_experts = (
+            model_config_for_expert_location.num_logical_experts
+            + server_args.ep_num_redundant_experts
+        )
+        ep_size = server_args.ep_size
+        assert num_physical_experts % ep_size == 0
+        num_local_physical_experts = num_physical_experts // ep_size
+
+        return dict(
+            model_config_for_expert_location=model_config_for_expert_location,
+            num_physical_experts=num_physical_experts,
+            num_local_physical_experts=num_local_physical_experts,
+            ep_size=ep_size,
+        )
+
+    @staticmethod
+    def _init_raw(
+        server_args: ServerArgs,
+        ep_size: int,
+        physical_to_logical_map: torch.Tensor,
+        logical_to_all_physical_map: torch.Tensor,
+    ):
+        _, num_physical_experts = physical_to_logical_map.shape
+
+        logical_to_all_physical_map_padded = F.pad(
+            logical_to_all_physical_map,
+            (0, num_physical_experts - logical_to_all_physical_map.shape[-1]),
+            value=-1,
+        )
+
+        logical_to_all_physical_map_num_valid = torch.count_nonzero(
+            logical_to_all_physical_map != -1, dim=-1
+        )
+
+        return ExpertLocationMetadata(
+            physical_to_logical_map=physical_to_logical_map,
+            physical_to_logical_map_cpu=physical_to_logical_map.cpu(),
+            logical_to_all_physical_map=logical_to_all_physical_map_padded,
+            logical_to_all_physical_map_cpu=logical_to_all_physical_map_padded.cpu(),
+            logical_to_all_physical_map_num_valid=logical_to_all_physical_map_num_valid,
+            logical_to_rank_dispatch_physical_map=(
+                compute_logical_to_rank_dispatch_physical_map(
+                    server_args=server_args,
+                    logical_to_all_physical_map=logical_to_all_physical_map,
+                    ep_size=ep_size,
+                    num_physical_experts=num_physical_experts,
+                    # TODO improve when we have real EP rank
+                    ep_rank=torch.distributed.get_rank() % ep_size,
+                )
+                if server_args.ep_dispatch_algorithm == "static"
+                else None
+            ),
+        )
+
+    # -------------------------------- mutation ------------------------------------
+
+    def update(
+        self,
+        other: ExpertLocationMetadata,
+        update_layer_ids: List[int],
+    ):
+        for field in [
+            "ep_size",
+        ]:
+            assert getattr(self, field) == getattr(other, field)
+
+        for field in [
+            "physical_to_logical_map",
+            "physical_to_logical_map_cpu",
+            "logical_to_all_physical_map",
+            "logical_to_all_physical_map_cpu",
+            "logical_to_all_physical_map_num_valid",
+            "logical_to_rank_dispatch_physical_map",
+        ]:
+            other_field = getattr(other, field)
+            self_field = getattr(self, field)
+            assert (other_field is not None) == (self_field is not None)
+            if self_field is not None:
+                mask_update = torch.tensor(
+                    [i in update_layer_ids for i in range(self.num_layers)]
+                )
+                mask_update = mask_update.view(*([-1] + [1] * (self_field.dim() - 1)))
+                mask_update = mask_update.to(self_field.device, non_blocking=True)
+                self_field[...] = torch.where(mask_update, other_field, self_field)
+
+    # -------------------------------- usage ------------------------------------
+
+    def logical_to_all_physical(
+        self,
+        layer_id: int,
+        logical_expert_id: int,
+        require_global_experts: bool = False,
+    ) -> List[int]:
+        # Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario
+        cpu_map = self.logical_to_all_physical_map_cpu
+        # Draft workers can query MoE layers whose layer_id lies beyond the
+        # target-sized expert map; fall back to the identity mapping (no EPLB
+        # rebalancing for those layers) instead of indexing out of range.
+        if layer_id >= cpu_map.shape[0]:
+            if require_global_experts:
+                num_physical_experts = cpu_map.shape[-1]
+                return list(
+                    range(
+                        logical_expert_id,
+                        num_physical_experts,
+                        self.num_logical_experts,
+                    )
+                )
+            return [logical_expert_id]
+        if require_global_experts:
+            num_physical_experts = cpu_map[layer_id].shape[-1]
+            return list(
+                range(logical_expert_id, num_physical_experts, self.num_logical_experts)
+            )
+        return [
+            physical_expert_id
+            for physical_expert_id in cpu_map[layer_id, logical_expert_id].tolist()
+            if physical_expert_id != -1
+        ]
+
+
+def format_expert_location_layout(
+    metadata: Optional[ExpertLocationMetadata],
+    layer_ids: Optional[Iterable[int]] = None,
+) -> str:
+    if metadata is None:
+        return "<none>"
+
+    return format_physical_to_logical_map(
+        metadata.physical_to_logical_map_cpu,
+        ep_size=metadata.ep_size,
+        layer_ids=layer_ids,
+    )
+
+
+def format_expert_location_layout_diff(
+    old_metadata: Optional[ExpertLocationMetadata],
+    new_metadata: Optional[ExpertLocationMetadata],
+    layer_ids: Optional[Iterable[int]] = None,
+) -> str:
+    if old_metadata is None or new_metadata is None:
+        return "<none>"
+
+    old_map = old_metadata.physical_to_logical_map_cpu
+    new_map = new_metadata.physical_to_logical_map_cpu
+    if old_map.shape != new_map.shape:
+        return f"shape_changed old_shape={tuple(old_map.shape)} new_shape={tuple(new_map.shape)}"
+
+    layer_ids = _normalize_layer_ids(layer_ids, num_layers=old_map.shape[0])
+    num_physical_experts = old_map.shape[1]
+
+    changed_by_layer = []
+    for layer_id in layer_ids:
+        num_changed = torch.count_nonzero(old_map[layer_id] != new_map[layer_id]).item()
+        if num_changed > 0:
+            changed_by_layer.append((layer_id, num_changed))
+
+    total_changed = sum(num_changed for _, num_changed in changed_by_layer)
+    total_slots = len(layer_ids) * num_physical_experts
+    lines = [f"changed_physical_slots={total_changed}/{total_slots}"]
+    if not changed_by_layer:
+        lines.append("changed_layers=[]")
+        return "\n".join(lines)
+
+    for layer_id, num_changed in changed_by_layer:
+        lines.append(f"layer={layer_id}: changed={num_changed}/{num_physical_experts}")
+    return "\n".join(lines)
+
+
+def format_physical_to_logical_map(
+    physical_to_logical_map: torch.Tensor,
+    ep_size: int,
+    layer_ids: Optional[Iterable[int]] = None,
+) -> str:
+    physical_to_logical_map = physical_to_logical_map.cpu()
+    if physical_to_logical_map.numel() == 0:
+        return "<empty>"
+
+    layer_ids = _normalize_layer_ids(
+        layer_ids, num_layers=physical_to_logical_map.shape[0]
+    )
+    num_physical_experts = physical_to_logical_map.shape[1]
+    num_local_physical_experts, remainder = divmod(num_physical_experts, ep_size)
+
+    lines = [
+        "physical_to_logical_map "
+        f"num_layers={physical_to_logical_map.shape[0]} "
+        f"num_physical_experts={num_physical_experts} "
+        f"ep_size={ep_size}"
+    ]
+    for layer_id in layer_ids:
+        row = physical_to_logical_map[layer_id].tolist()
+        if remainder != 0:
+            lines.append(
+                f"layer={layer_id}: "
+                f"physical={json.dumps(row, separators=(',', ':'))}"
+            )
+            continue
+
+        rank_chunks = []
+        for ep_rank in range(ep_size):
+            start = ep_rank * num_local_physical_experts
+            end = start + num_local_physical_experts
+            rank_chunks.append(
+                f"ep{ep_rank}={json.dumps(row[start:end], separators=(',', ':'))}"
+            )
+        lines.append(f"layer={layer_id}: " + " ".join(rank_chunks))
+
+    return "\n".join(lines)
+
+
+def _normalize_layer_ids(
+    layer_ids: Optional[Iterable[int]],
+    num_layers: int,
+) -> List[int]:
+    if layer_ids is None:
+        return list(range(num_layers))
+
+    normalized_layer_ids = [int(layer_id) for layer_id in layer_ids]
+    for layer_id in normalized_layer_ids:
+        assert 0 <= layer_id < num_layers, f"{layer_id=} {num_layers=}"
+    return normalized_layer_ids
+
+
+def get_global_expert_location_metadata():
+    from sglang.srt.runtime_context import get_resources
+
+    return get_resources().expert_location_metadata
+
+
+def set_global_expert_location_metadata(value):
+    from sglang.srt.runtime_context import get_resources
+
+    resources = get_resources()
+    assert resources.expert_location_metadata is None
+    resources.expert_location_metadata = value
+
+
+def broadcast_global_expert_location_metadata(
+    src_rank: int = 0, group: Optional[torch.distributed.ProcessGroup] = None
+):
+    """Broadcast the global ExpertLocationMetadata from src_rank to all ranks.
+
+    This is used in Elastic EP rank recovery to ensure that all ranks (including
+    newly recovered ones) share exactly the same expert location metadata.
+
+    Note: The caller must ensure src_rank is a healthy rank. In recovery scenarios,
+    this function is called after try_recover_ranks succeeds, at which point all
+    ranks (including src_rank=0) have recovered and are ready.
+    """
+    metadata = get_global_expert_location_metadata()
+    assert metadata is not None
+
+    # Ensure device tensors are contiguous before broadcasting in-place
+    metadata.physical_to_logical_map = metadata.physical_to_logical_map.contiguous()
+    metadata.logical_to_all_physical_map = (
+        metadata.logical_to_all_physical_map.contiguous()
+    )
+    metadata.logical_to_all_physical_map_num_valid = (
+        metadata.logical_to_all_physical_map_num_valid.contiguous()
+    )
+    if metadata.logical_to_rank_dispatch_physical_map is not None:
+        metadata.logical_to_rank_dispatch_physical_map = (
+            metadata.logical_to_rank_dispatch_physical_map.contiguous()
+        )
+
+    device_tensors = [
+        metadata.physical_to_logical_map,
+        metadata.logical_to_all_physical_map,
+        metadata.logical_to_all_physical_map_num_valid,
+    ]
+    if metadata.logical_to_rank_dispatch_physical_map is not None:
+        device_tensors.append(metadata.logical_to_rank_dispatch_physical_map)
+
+    for tensor in device_tensors:
+        torch.distributed.broadcast(tensor, src=src_rank, group=group)
+
+    # After broadcasting device tensors, refresh corresponding CPU copies
+    metadata.physical_to_logical_map_cpu = metadata.physical_to_logical_map.cpu()
+    metadata.logical_to_all_physical_map_cpu = (
+        metadata.logical_to_all_physical_map.cpu()
+    )
+
+
+def _compute_logical_to_all_physical_map(
+    server_args: ServerArgs,
+    physical_to_logical_map: torch.Tensor,
+    num_logical_experts: int,
+    ep_size: int,
+    moe_ep_rank: int,
+):
+    # This is rarely called, so we use for loops for maximum clarity
+
+    num_layers, num_physical_experts = physical_to_logical_map.shape
+
+    logical_to_all_physical_map = [
+        [[] for _ in range(num_logical_experts)] for _ in range(num_layers)
+    ]
+
+    # Find out the candidate physical experts for each logical expert on each layer
+    for layer_id in range(num_layers):
+        for physical_expert_id in range(num_physical_experts):
+            logical_expert_id = physical_to_logical_map[
+                layer_id, physical_expert_id
+            ].item()
+            logical_to_all_physical_map[layer_id][logical_expert_id].append(
+                physical_expert_id
+            )
+
+    # Replace by the physical expert on local GPU or node if possible
+    if moe_ep_rank is not None:
+        num_gpus_per_node = server_args.ep_size // server_args.nnodes
+        num_local_gpu_physical_experts = num_physical_experts // ep_size
+        num_local_node_physical_experts = (
+            num_local_gpu_physical_experts * num_gpus_per_node
+        )
+        for layer_id in range(num_layers):
+            for logical_expert_id in range(num_logical_experts):
+                # Try to find the nearest physical expert
+                nearest_expert = _find_nearest_expert(
+                    candidate_physical_expert_ids=logical_to_all_physical_map[layer_id][
+                        logical_expert_id
+                    ],
+                    num_local_gpu_physical_experts=num_local_gpu_physical_experts,
+                    moe_ep_rank=moe_ep_rank,
+                    num_gpus_per_node=num_gpus_per_node,
+                    num_local_node_physical_experts=num_local_node_physical_experts,
+                )
+
+                # Replace by the nearest physical expert
+                if nearest_expert != -1:
+                    logical_to_all_physical_map[layer_id][logical_expert_id] = [
+                        nearest_expert
+                    ]
+
+    logical_to_all_physical_map = _pad_nested_array(
+        logical_to_all_physical_map, pad_value=-1
+    )
+
+    return torch.tensor(
+        logical_to_all_physical_map, device=physical_to_logical_map.device
+    )
+
+
+def _pad_nested_array(arr, pad_value):
+    max_len = max(len(inner) for outer in arr for inner in outer)
+    padded = [
+        [inner + [pad_value] * (max_len - len(inner)) for inner in outer]
+        for outer in arr
+    ]
+    return padded
+
+
+# TODO optimize performance (rewrite and/or run in separate process with overlap)
+def compute_logical_to_rank_dispatch_physical_map(
+    server_args: ServerArgs,
+    logical_to_all_physical_map: torch.Tensor,
+    ep_size: int,
+    num_physical_experts: int,
+    ep_rank: int,
+    seed: int = 42,
+):
+    r = random.Random(seed)
+
+    device = logical_to_all_physical_map.device
+    logical_to_all_physical_map = logical_to_all_physical_map.cpu()
+
+    num_local_gpu_physical_experts = num_physical_experts // ep_size
+    num_gpus_per_node = server_args.ep_size // server_args.nnodes
+    num_local_node_physical_experts = num_local_gpu_physical_experts * num_gpus_per_node
+    num_layers, num_logical_experts, _ = logical_to_all_physical_map.shape
+    dtype = logical_to_all_physical_map.dtype
+
+    result_list = [
+        [[-1] * num_logical_experts for _ in range(num_layers)] for _ in range(ep_size)
+    ]
+
+    for layer_id in range(num_layers):
+        for logical_expert_id in range(num_logical_experts):
+            candidate_physical_expert_ids = _logical_to_all_physical_raw(
+                logical_to_all_physical_map, layer_id, logical_expert_id
+            )
+
+            remaining_ranks = []
+            for moe_ep_rank in range(ep_size):
+                val = _find_nearest_expert(
+                    candidate_physical_expert_ids=candidate_physical_expert_ids,
+                    num_local_gpu_physical_experts=num_local_gpu_physical_experts,
+                    moe_ep_rank=moe_ep_rank,
+                    num_gpus_per_node=num_gpus_per_node,
+                    num_local_node_physical_experts=num_local_node_physical_experts,
+                )
+
+                result_list[moe_ep_rank][layer_id][logical_expert_id] = val
+                if val == -1:
+                    remaining_ranks.append(moe_ep_rank)
+
+            if remaining_ranks:
+                choices = _fair_choices(
+                    candidate_physical_expert_ids, k=len(remaining_ranks), r=r
+                )
+                for moe_ep_rank, choice in zip(remaining_ranks, choices, strict=True):
+                    result_list[moe_ep_rank][layer_id][logical_expert_id] = choice
+
+    logical_to_rank_dispatch_physical_map = torch.tensor(result_list, dtype=dtype)
+    assert torch.all(logical_to_rank_dispatch_physical_map != -1)
+
+    return logical_to_rank_dispatch_physical_map[ep_rank, :, :].to(device)
+
+
+def _logical_to_all_physical_raw(
+    logical_to_all_physical_map, layer_id: int, logical_expert_id: int
+) -> List[int]:
+    return [
+        physical_expert_id
+        for physical_expert_id in logical_to_all_physical_map[
+            layer_id, logical_expert_id
+        ].tolist()
+        if physical_expert_id != -1
+    ]
+
+
+def _compute_gpu_id_of_physical_expert(
+    physical_expert_id: int, num_local_gpu_physical_experts: int
+) -> int:
+    return physical_expert_id // num_local_gpu_physical_experts
+
+
+def _compute_node_id_of_physical_expert(
+    physical_expert_id: int, num_local_host_physical_experts: int
+) -> int:
+    return physical_expert_id // num_local_host_physical_experts
+
+
+def _find_nearest_expert(
+    candidate_physical_expert_ids: List[int],
+    num_local_gpu_physical_experts: int,
+    moe_ep_rank: int,
+    num_gpus_per_node: int,
+    num_local_node_physical_experts: int,
+) -> int:
+    # 1. If only one candidate, return it directly
+    if len(candidate_physical_expert_ids) == 1:
+        return candidate_physical_expert_ids[0]
+
+    # 2. Prefer same-GPU experts
+    same_gpu_physical_expert_ids = [
+        physical_expert_id
+        for physical_expert_id in candidate_physical_expert_ids
+        if _compute_gpu_id_of_physical_expert(
+            physical_expert_id, num_local_gpu_physical_experts
+        )
+        == moe_ep_rank
+    ]
+    if len(same_gpu_physical_expert_ids) > 0:
+        return same_gpu_physical_expert_ids[0]
+
+    # 3. Otherwise, prefer same-node experts
+    node_rank = moe_ep_rank // num_gpus_per_node
+    same_node_physical_expert_ids = [
+        physical_expert_id
+        for physical_expert_id in candidate_physical_expert_ids
+        if _compute_node_id_of_physical_expert(
+            physical_expert_id, num_local_node_physical_experts
+        )
+        == node_rank
+    ]
+    if len(same_node_physical_expert_ids) > 0:
+        return same_node_physical_expert_ids[0]
+
+    # 4. At last, leave it as -1 to indicate not found.
+    return -1
+
+
+def _fair_choices(arr: List, k: int, r: random.Random) -> List:
+    quotient, remainder = divmod(k, len(arr))
+    ans = arr * quotient + r.sample(arr, k=remainder)
+    r.shuffle(ans)
+    return ans
+
+
+@dataclass
+class ModelConfigForExpertLocation:
+    num_layers: int
+    num_logical_experts: int
+    num_groups: Optional[int] = None
+
+    @staticmethod
+    def from_model_config(model_config: ModelConfig):
+        from sglang.srt.model_loader import get_model_architecture
+
+        model_class, _ = get_model_architecture(model_config)
+        if hasattr(model_class, "get_model_config_for_expert_location"):
+            return model_class.get_model_config_for_expert_location(
+                model_config.hf_config
+            )
+        else:
+            return None
+
+
+def compute_initial_expert_location_metadata(
+    server_args: ServerArgs,
+    model_config: ModelConfig,
+    moe_ep_rank: int,
+) -> Optional[ExpertLocationMetadata]:
+    data = server_args.init_expert_location
+    if data == "trivial":
+        return ExpertLocationMetadata.init_trivial(
+            server_args, model_config, moe_ep_rank
+        )
+
+    # TODO unify with the utils function
+    if data.endswith(".pt"):
+        data_dict = torch.load(data, weights_only=True)
+    elif data.endswith(".json"):
+        data_dict = json.loads(Path(data).read_text())
+    else:
+        data_dict = json.loads(data)
+
+    if "physical_to_logical_map" in data_dict:
+        logger.info(
+            "init_expert_location from init_by_mapping using ServerArgs.init_expert_location"
+        )
+        return ExpertLocationMetadata.init_by_mapping(
+            server_args,
+            model_config,
+            **data_dict,
+            moe_ep_rank=moe_ep_rank,
+        )
+    elif "logical_count" in data_dict:
+        logger.info(
+            "init_expert_location from init_by_eplb using ServerArgs.init_expert_location"
+        )
+        return ExpertLocationMetadata.init_by_eplb(
+            server_args, model_config, logical_count=data_dict["logical_count"]
+        )
+    else:
+        raise NotImplementedError(
+            f"Unknown init_expert_location format ({list(data_dict.keys())=})"
+        )

python/sglang/srt/eplb/expert_location_dispatch.py

@@ -0,0 +1,145 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+import torch
+
+from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
+from sglang.srt.runtime_context import get_server_args
+
+
+@dataclass
+class ExpertLocationDispatchInfo:
+    ep_dispatch_algorithm: Literal["static", "random"]
+    # (num_logical_experts,)
+    partial_logical_to_rank_dispatch_physical_map: Optional[torch.Tensor]
+    # (num_logical_experts, X)
+    partial_logical_to_all_physical_map: torch.Tensor
+    # (num_logical_experts,)
+    partial_logical_to_all_physical_map_num_valid: torch.Tensor
+    num_physical_experts: int
+
+    @classmethod
+    def init_new(cls, layer_id: int):
+        ep_dispatch_algorithm = get_server_args().ep_dispatch_algorithm
+        expert_location_metadata = get_global_expert_location_metadata()
+        assert expert_location_metadata is not None
+
+        if ep_dispatch_algorithm is None:
+            return None
+
+        return cls(
+            ep_dispatch_algorithm=ep_dispatch_algorithm,
+            partial_logical_to_rank_dispatch_physical_map=(
+                expert_location_metadata.logical_to_rank_dispatch_physical_map[
+                    layer_id, :
+                ]
+                if expert_location_metadata.logical_to_rank_dispatch_physical_map
+                is not None
+                else None
+            ),
+            partial_logical_to_all_physical_map=expert_location_metadata.logical_to_all_physical_map[
+                layer_id, :
+            ],
+            partial_logical_to_all_physical_map_num_valid=expert_location_metadata.logical_to_all_physical_map_num_valid[
+                layer_id, :
+            ],
+            num_physical_experts=expert_location_metadata.num_physical_experts,
+        )
+
+
+def transform_select_experts_inputs(
+    router_logits: torch.Tensor,
+    correction_bias: Optional[torch.Tensor],
+    info: Optional[ExpertLocationDispatchInfo],
+):
+    if (info is not None) and (info.ep_dispatch_algorithm == "fake"):
+        router_logits.uniform_(5, 10)
+        if correction_bias is not None:
+            correction_bias = torch.zeros_like(correction_bias)
+    return router_logits, correction_bias
+
+
+def topk_ids_logical_to_physical(
+    topk_ids: torch.Tensor,
+    info: Optional[ExpertLocationDispatchInfo],
+    log2phy_prob: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if info is None:
+        return topk_ids
+
+    if info.ep_dispatch_algorithm == "static":
+        return _topk_ids_logical_to_physical_static(topk_ids, info)
+    if info.ep_dispatch_algorithm in ["dynamic", "fake"]:
+        return _topk_ids_logical_to_physical_dynamic(topk_ids, info)
+    if info.ep_dispatch_algorithm == "lp":
+        if log2phy_prob is None:
+            raise RuntimeError(
+                "ep_dispatch_algorithm='lp' but log2phy_prob is None at dispatch "
+                f"time (topk_ids.shape={tuple(topk_ids.shape)})."
+            )
+        return _topk_ids_logical_to_physical_probability(topk_ids, info, log2phy_prob)
+    raise NotImplementedError(f"Unknown algorithm {info.ep_dispatch_algorithm}")
+
+
+def _topk_ids_logical_to_physical_static(
+    topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
+) -> torch.Tensor:
+    physical_topk_ids = info.partial_logical_to_rank_dispatch_physical_map[topk_ids]
+    if physical_topk_ids.dtype != topk_ids.dtype:
+        physical_topk_ids = physical_topk_ids.to(topk_ids.dtype)
+    return physical_topk_ids
+
+
+def _topk_ids_logical_to_physical_dynamic(
+    topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
+) -> torch.Tensor:
+    topk_ids_original_shape = topk_ids.shape
+    original_dtype = topk_ids.dtype
+    device = topk_ids.device
+    topk_ids = topk_ids.flatten()
+
+    chosen_dispatch_index = (
+        torch.randint(0, 65536, topk_ids.shape, dtype=torch.int32, device=device)
+        % info.partial_logical_to_all_physical_map_num_valid[topk_ids]
+    )
+    topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index]
+    if topk_ids.dtype != original_dtype:
+        topk_ids = topk_ids.to(original_dtype)
+
+    topk_ids = topk_ids.view(topk_ids_original_shape)
+    return topk_ids
+
+
+def _topk_ids_logical_to_physical_probability(
+    topk_ids: torch.Tensor,
+    info: ExpertLocationDispatchInfo,
+    log2phy_prob: torch.Tensor,
+) -> torch.Tensor:
+    """Select physical experts via the JIT-compiled CUDA dispatch kernel.
+
+    Raises if ``topk_ids`` isn't on CUDA — the LP path requires the fused
+    kernel and there is no torch reference fallback at runtime.
+    """
+    if not topk_ids.is_cuda:
+        raise RuntimeError(
+            "LP dispatch requires CUDA tensors; got topk_ids on " f"{topk_ids.device}."
+        )
+    from sglang.jit_kernel.lplb import cuda_solver
+
+    return cuda_solver.dispatch_probability(
+        topk_ids, log2phy_prob, info.partial_logical_to_all_physical_map
+    )

python/sglang/srt/eplb/expert_location_updater.py

@@ -0,0 +1,633 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import einops
+import torch
+import torch.distributed
+from torch.distributed import P2POp
+
+from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
+from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_location import (
+    ExpertLocationMetadata,
+    get_global_expert_location_metadata,
+)
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils import get_bool_env_var
+
+logger = logging.getLogger(__name__)
+
+
+_LOG_INPUT = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_INPUT")
+
+
+class ExpertLocationUpdater:
+    def __init__(self):
+        self._first_execution = True
+
+    def update(
+        self,
+        routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]],
+        new_expert_location_metadata: ExpertLocationMetadata,
+        update_layer_ids: List[int],
+        nnodes: int,
+        rank: int,
+    ):
+        """
+        Update experts' physical location after EPLB.
+
+        Returns a map of layer_id to expert_ids that are missing due to rank
+        failures during fault conditions when elastic EP is enabled.
+        """
+        if self._first_execution:
+            self._first_execution = False
+            torch.get_device_module().empty_cache()
+
+        old_expert_location_metadata = get_global_expert_location_metadata()
+        assert old_expert_location_metadata is not None
+
+        missing_logical_experts_by_layers = _update_expert_weights(
+            routed_experts_weights_of_layer=routed_experts_weights_of_layer,
+            old_expert_location_metadata=old_expert_location_metadata,
+            new_expert_location_metadata=new_expert_location_metadata,
+            update_layer_ids=update_layer_ids,
+            nnodes=nnodes,
+            rank=rank,
+        )
+        old_expert_location_metadata.update(
+            new_expert_location_metadata,
+            update_layer_ids=update_layer_ids,
+        )
+
+        return missing_logical_experts_by_layers
+
+
+def _update_expert_weights(**kwargs):
+    if get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_CANARY"):
+        return _update_expert_weights_with_canary(**kwargs)
+    else:
+        return _update_expert_weights_raw(**kwargs)
+
+
+# can add watchdog as well
+def _update_expert_weights_with_canary(
+    routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]],
+    old_expert_location_metadata: ExpertLocationMetadata,
+    new_expert_location_metadata: ExpertLocationMetadata,
+    update_layer_ids: List[int],
+    nnodes: int,
+    rank: int,
+):
+    num_local_physical_experts = old_expert_location_metadata.num_local_physical_experts
+
+    def _get_canary_value(meta: ExpertLocationMetadata, layer_id: int):
+        return meta.physical_to_logical_map_cpu[
+            layer_id,
+            num_local_physical_experts * rank : num_local_physical_experts * (rank + 1),
+        ]
+
+    routed_experts_weights_of_layer = {
+        k: [x for x in v] for k, v in routed_experts_weights_of_layer.items()
+    }
+    for layer_id in update_layer_ids:
+        canary_tensor = (
+            _get_canary_value(old_expert_location_metadata, layer_id)
+            .clone()
+            .to(device=get_server_args().device, non_blocking=True)
+        )
+        routed_experts_weights_of_layer[layer_id].append(canary_tensor)
+
+    missing_logical_experts_by_layers = _update_expert_weights_raw(
+        routed_experts_weights_of_layer=routed_experts_weights_of_layer,
+        old_expert_location_metadata=old_expert_location_metadata,
+        new_expert_location_metadata=new_expert_location_metadata,
+        update_layer_ids=update_layer_ids,
+        nnodes=nnodes,
+        rank=rank,
+    )
+
+    for layer_id in update_layer_ids:
+        # can optimize speed if needed
+        expect_value = _get_canary_value(new_expert_location_metadata, layer_id)
+        actual_value = routed_experts_weights_of_layer[layer_id][-1].cpu()
+        assert torch.all(expect_value == actual_value), (
+            f"{expect_value=} {actual_value=} {layer_id=} "
+            f"{old_expert_location_metadata.physical_to_logical_map_cpu.tolist()=} "
+            f"{new_expert_location_metadata.physical_to_logical_map_cpu.tolist()=} "
+        )
+
+    return missing_logical_experts_by_layers
+
+
+def _update_expert_weights_raw(
+    routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]],
+    old_expert_location_metadata: ExpertLocationMetadata,
+    new_expert_location_metadata: ExpertLocationMetadata,
+    update_layer_ids: List[int],
+    nnodes: int,
+    rank: int,
+):
+    log_metrics = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_METRICS")
+
+    temp_buffers = create_temp_buffers(
+        routed_experts_weights_of_layer[update_layer_ids[0]]
+    )
+
+    world_size = torch.distributed.get_world_size()
+    num_local_physical_experts = old_expert_location_metadata.num_local_physical_experts
+    num_gpu_per_node = world_size // nnodes
+
+    missing_logical_experts_by_layers: Dict[int, List[int]] = {}
+
+    for layer_id in update_layer_ids:
+        missing_logical_experts_info: List[int] = []
+        update_expert_weights_single_layer(
+            routed_experts_weights=routed_experts_weights_of_layer[layer_id],
+            temp_buffers=temp_buffers,
+            old_physical_to_logical_map=old_expert_location_metadata.physical_to_logical_map_cpu[
+                layer_id
+            ].tolist(),
+            new_physical_to_logical_map=new_expert_location_metadata.physical_to_logical_map_cpu[
+                layer_id
+            ].tolist(),
+            num_local_physical_experts=num_local_physical_experts,
+            num_gpu_per_node=num_gpu_per_node,
+            rank=rank,
+            world_size=world_size,
+            missing_logical_experts_info=missing_logical_experts_info,
+            log_metrics=log_metrics,
+        )
+        if len(missing_logical_experts_info) > 0:
+            missing_logical_experts_by_layers[layer_id] = missing_logical_experts_info
+    return missing_logical_experts_by_layers
+
+
+def create_temp_buffers(sample_tensors):
+    return [torch.empty_like(tensor) for tensor in sample_tensors]
+
+
+def update_expert_weights_single_layer(
+    routed_experts_weights: List[torch.Tensor],
+    temp_buffers: List[torch.Tensor],
+    old_physical_to_logical_map: List[int],  # (num_physical_Experts,)
+    new_physical_to_logical_map: List[int],  # (num_physical_Experts,)
+    num_local_physical_experts: int,
+    num_gpu_per_node: int,
+    rank: int,
+    world_size: Optional[int] = None,
+    missing_logical_experts_info: Optional[List[int]] = None,
+    debug: bool = False,
+    log_metrics: bool = False,
+):
+    assert all(
+        tensor.shape[0] == num_local_physical_experts
+        for tensor in routed_experts_weights
+    ), f"{num_local_physical_experts=} {[x.shape for x in routed_experts_weights]=}"
+    assert isinstance(old_physical_to_logical_map, list)
+    assert isinstance(new_physical_to_logical_map, list)
+
+    if _LOG_INPUT:
+        logger.info(
+            "update_expert_weights_single_layer "
+            f"{[x.shape for x in routed_experts_weights]=} "
+            f"{[x.shape for x in temp_buffers]=} "
+            f"{old_physical_to_logical_map=} "
+            f"{new_physical_to_logical_map=} "
+            f"{num_local_physical_experts=} "
+            f"{num_gpu_per_node=} "
+            f"{rank=} "
+            f"{world_size=} "
+        )
+
+    output_logs = [] if debug else None
+
+    num_physical_experts = len(old_physical_to_logical_map)
+    num_tensors = len(routed_experts_weights)
+
+    self_node_id = rank // num_gpu_per_node
+
+    local_expert_location_range = (
+        rank * num_local_physical_experts,
+        (rank + 1) * num_local_physical_experts,
+    )
+
+    def _entrypoint():
+        # List[Tuple[logical_expert_id, List[P2POp]]]
+        p2p_op_infos: List[Tuple[int, List[P2POp]]] = []
+        # List[Tuple[temp_buffers_expert_location, routed_experts_weights_expert_location]]
+        buffer2weight_copy_infos: List[Tuple[int, int]] = []
+
+        _handle_recv(buffer2weight_copy_infos, p2p_op_infos)
+        _create_isend_ops(p2p_op_infos)
+        _filter_p2p_ops(p2p_op_infos)
+        _execute_p2p_ops(p2p_op_infos)
+        _execute_buffer2weight_copies(buffer2weight_copy_infos)
+
+        if log_metrics:
+            _log_p2p_op_metrics(
+                p2p_op_infos,
+                world_size=world_size,
+                num_gpu_per_node=num_gpu_per_node,
+                self_node_id=self_node_id,
+            )
+
+        if debug:
+            output_logs.append(f"{p2p_op_infos=}")
+            output_logs.append(f"{buffer2weight_copy_infos=}")
+
+    def _handle_recv(buffer2weight_copy_infos, p2p_op_infos):
+        for dst_expert_location in range(*local_expert_location_range):
+            _handle_recv_of_dst_expert_location(
+                dst_expert_location, buffer2weight_copy_infos, p2p_op_infos
+            )
+
+    def _handle_recv_of_dst_expert_location(
+        dst_expert_location: int, buffer2weight_copy_infos, p2p_op_infos
+    ):
+        logical_expert_id = new_physical_to_logical_map[dst_expert_location]
+
+        # case 1: unchanged
+        if old_physical_to_logical_map[dst_expert_location] == logical_expert_id:
+            if debug:
+                output_logs.append(
+                    f"handle_recv_of_dst_expert_location {dst_expert_location=} case=unchanged"
+                )
+            return
+
+        # case 2: same-gpu
+        for src_expert_location in range(*local_expert_location_range):
+            if old_physical_to_logical_map[src_expert_location] == logical_expert_id:
+                for i in range(num_tensors):
+                    _get_tensor(temp_buffers, i, dst_expert_location).copy_(
+                        _get_tensor(routed_experts_weights, i, src_expert_location)
+                    )
+                buffer2weight_copy_infos.append(
+                    (dst_expert_location, dst_expert_location)
+                )
+                if debug:
+                    output_logs.append(
+                        f"handle_recv_of_dst_expert_location {dst_expert_location=} case=same-gpu {src_expert_location=}"
+                    )
+                return
+
+        # case 3: free-rider
+        for src_expert_location in range(
+            rank * num_local_physical_experts, dst_expert_location
+        ):
+            if new_physical_to_logical_map[src_expert_location] == logical_expert_id:
+                buffer2weight_copy_infos.append(
+                    (src_expert_location, dst_expert_location)
+                )
+                if debug:
+                    output_logs.append(
+                        f"handle_recv_of_dst_expert_location {dst_expert_location=} case=free-rider {src_expert_location=}"
+                    )
+                return
+
+        same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks = (
+            _compute_comm_info(logical_expert_id=logical_expert_id)
+        )
+
+        # case 4: same-node
+        if rank in need_comm_self_node_dst_ranks:
+            chosen_src_rank = same_node_mapping.chunk_value_from_element_value(
+                element_value=rank
+            )
+            _create_p2p_recv_and_buffer2weight_copy(
+                buffer2weight_copy_infos,
+                p2p_op_infos,
+                src_rank=chosen_src_rank,
+                logical_expert_id=logical_expert_id,
+                dst_expert_location=dst_expert_location,
+            )
+            if debug:
+                output_logs.append(
+                    f"handle_recv_of_dst_expert_location {dst_expert_location=} case=same-node {chosen_src_rank=}"
+                )
+            return
+
+        # case 5: cross-node
+        # Future work: can optimize when there are multiple ranks in the same dst node that uses the same logical expert
+        chosen_src_rank = cross_node_mapping.chunk_value_from_element_value(
+            element_value=rank
+        )
+        _create_p2p_recv_and_buffer2weight_copy(
+            buffer2weight_copy_infos,
+            p2p_op_infos,
+            src_rank=chosen_src_rank,
+            logical_expert_id=logical_expert_id,
+            dst_expert_location=dst_expert_location,
+        )
+        if debug:
+            output_logs.append(
+                f"handle_recv_of_dst_expert_location {dst_expert_location=} case=cross-node {chosen_src_rank=}"
+            )
+        return
+
+    def _create_p2p_recv_and_buffer2weight_copy(
+        buffer2weight_copy_infos,
+        p2p_op_infos,
+        *,
+        logical_expert_id: int,
+        src_rank: int,
+        dst_expert_location: int,
+    ):
+        p2p_op_infos.append(
+            (
+                logical_expert_id,
+                [
+                    P2POp(
+                        op=torch.distributed.irecv,
+                        tensor=_get_tensor(temp_buffers, i, dst_expert_location),
+                        peer=src_rank,
+                    )
+                    for i in range(num_tensors)
+                ],
+            )
+        )
+        buffer2weight_copy_infos.append((dst_expert_location, dst_expert_location))
+
+    def _create_isend_ops(p2p_op_infos):
+        handled_logical_expert_ids = set()
+        for src_expert_location in range(*local_expert_location_range):
+            logical_expert_id = old_physical_to_logical_map[src_expert_location]
+
+            if logical_expert_id in handled_logical_expert_ids:
+                continue
+            handled_logical_expert_ids.add(logical_expert_id)
+
+            _create_isend_ops_of_logical_expert_id(
+                logical_expert_id, src_expert_location, p2p_op_infos
+            )
+
+    def _create_isend_ops_of_logical_expert_id(
+        logical_expert_id, src_expert_location, p2p_op_infos
+    ):
+        same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks = (
+            _compute_comm_info(logical_expert_id=logical_expert_id)
+        )
+
+        same_node_dst_ranks = same_node_mapping.element_values_from_chunk_value(
+            chunk_value=rank
+        )
+        cross_node_dst_ranks = cross_node_mapping.element_values_from_chunk_value(
+            chunk_value=rank
+        )
+        all_dst_ranks = same_node_dst_ranks + cross_node_dst_ranks
+
+        if debug:
+            output_logs.append(
+                f"create_isend_ops_of_logical_expert_id {logical_expert_id=} {src_expert_location=} {same_node_dst_ranks=} {cross_node_dst_ranks=}"
+            )
+
+        p2p_op_infos.append(
+            (
+                logical_expert_id,
+                [
+                    P2POp(
+                        op=torch.distributed.isend,
+                        tensor=_get_tensor(
+                            routed_experts_weights, i, src_expert_location
+                        ),
+                        peer=dst_rank,
+                    )
+                    for dst_rank in all_dst_ranks
+                    for i in range(num_tensors)
+                ],
+            )
+        )
+
+    def _compute_comm_info(logical_expert_id: int):
+        all_src_ranks = _deduplicate_ordered(
+            [
+                x // num_local_physical_experts
+                for x in range(num_physical_experts)
+                if old_physical_to_logical_map[x] == logical_expert_id
+            ]
+        )
+        all_src_nodes = [x // num_gpu_per_node for x in all_src_ranks]
+        self_node_src_ranks = [
+            x for x in all_src_ranks if x // num_gpu_per_node == self_node_id
+        ]
+
+        need_comm_dst_ranks = _deduplicate_ordered(
+            [
+                x // num_local_physical_experts
+                for x in range(num_physical_experts)
+                if new_physical_to_logical_map[x] == logical_expert_id
+                and x // num_local_physical_experts not in all_src_ranks
+            ]
+        )
+        need_comm_self_node_dst_ranks = (
+            [x for x in need_comm_dst_ranks if x // num_gpu_per_node == self_node_id]
+            if len(self_node_src_ranks) > 0
+            else []
+        )
+        need_comm_cross_node_dst_ranks = [
+            x
+            for x in need_comm_dst_ranks
+            if (x // num_gpu_per_node) not in all_src_nodes
+        ]
+
+        same_node_mapping = _ChunkUtils(
+            chunk_values=self_node_src_ranks,
+            element_values=need_comm_self_node_dst_ranks,
+        )
+
+        cross_node_mapping = _ChunkUtils(
+            chunk_values=all_src_ranks,
+            element_values=need_comm_cross_node_dst_ranks,
+        )
+
+        return same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks
+
+    def _filter_p2p_ops(p2p_op_infos):
+        elastic_ep_state = ElasticEPStateManager.instance()
+        if elastic_ep_state is not None and missing_logical_experts_info is not None:
+            # Filter out inactive P2P ops and record missing expert IDs in missing_logical_experts_info
+            is_active = elastic_ep_state.active_ranks_cpu
+            for i, (logical_expert_id, ops) in enumerate(p2p_op_infos):
+                has_isend = any(op.op == torch.distributed.isend for op in ops)
+                has_irecv = any(op.op == torch.distributed.irecv for op in ops)
+                assert not (has_isend and has_irecv), (
+                    "Each p2p_op_infos entry is expected to contain only send "
+                    "or only recv ops."
+                )
+
+                if has_isend:
+                    p2p_op_infos[i] = (
+                        logical_expert_id,
+                        [op for op in ops if is_active[op.peer]],
+                    )
+                elif has_irecv:
+                    if any(not is_active[op.peer] for op in ops):
+                        missing_logical_experts_info.append(logical_expert_id)
+                        p2p_op_infos[i] = (logical_expert_id, [])
+
+    def _execute_p2p_ops(p2p_op_infos):
+        sorted_infos = sorted(p2p_op_infos, key=lambda info: info[0])
+        p2p_ops = [op for _, ops in sorted_infos for op in ops]
+        if len(p2p_ops) == 0:
+            return
+
+        # Submit P2P ops in batches to prevent NCCL/RCCL GPU-side accumulation
+        # hangs on large rebalances. All ranks use the same expert_id ranges
+        # (based on num_physical_experts) so matching send/recv pairs land in
+        # the same batch. Set batch_chunk_size >= num_physical_experts to disable.
+        batch_chunk_size = envs.SGLANG_EPLB_P2P_BATCH_CHUNK_SIZE.get()
+        ops_by_expert = {eid: ops for eid, ops in sorted_infos}
+        for start in range(0, num_physical_experts, batch_chunk_size):
+            batch_ops = []
+            for eid in range(
+                start, min(start + batch_chunk_size, num_physical_experts)
+            ):
+                if eid in ops_by_expert:
+                    batch_ops.extend(ops_by_expert[eid])
+            if batch_ops:
+                reqs = torch.distributed.batch_isend_irecv(batch_ops)
+                for req in reqs:
+                    req.wait()
+
+    def _execute_buffer2weight_copies(buffer2weight_copy_infos):
+        for (
+            temp_buffers_expert_location,
+            routed_experts_weights_expert_location,
+        ) in buffer2weight_copy_infos:
+            for i in range(num_tensors):
+                _get_tensor(
+                    routed_experts_weights, i, routed_experts_weights_expert_location
+                ).copy_(_get_tensor(temp_buffers, i, temp_buffers_expert_location))
+
+    def _get_tensor(tensors, tensor_index: int, expert_location: int) -> torch.Tensor:
+        return tensors[tensor_index][_get_local_expert_location(expert_location)]
+
+    def _get_local_expert_location(expert_location: int) -> int:
+        assert (
+            local_expert_location_range[0]
+            <= expert_location
+            < local_expert_location_range[1]
+        )
+        return expert_location % num_local_physical_experts
+
+    _entrypoint()
+
+    return output_logs
+
+
+class _ChunkUtils:
+    def __init__(self, *, chunk_values: List, element_values: List):
+        self.chunk_values = chunk_values
+        self.element_values = element_values
+
+    def chunk_value_from_element_value(self, element_value):
+        chunk_index = self._chunk_index_from_element_index(
+            num_elements=len(self.element_values),
+            num_chunks=len(self.chunk_values),
+            element_index=self.element_values.index(element_value),
+        )
+        return self.chunk_values[chunk_index]
+
+    def element_values_from_chunk_value(self, chunk_value) -> List:
+        if len(self.element_values) == 0:
+            return []
+        element_slice = self._element_slice_from_chunk_index(
+            num_elements=len(self.element_values),
+            num_chunks=len(self.chunk_values),
+            chunk_index=self.chunk_values.index(chunk_value),
+        )
+        return self.element_values[element_slice]
+
+    @staticmethod
+    def _chunk_index_from_element_index(
+        num_elements: int, num_chunks: int, element_index: int
+    ) -> int:
+        short_chunk_size, num_long_chunks = divmod(num_elements, num_chunks)
+        num_elements_for_long_chunks = num_long_chunks * (short_chunk_size + 1)
+        if element_index < num_elements_for_long_chunks:
+            return element_index // (short_chunk_size + 1)
+        else:
+            return (
+                num_long_chunks
+                + (element_index - num_elements_for_long_chunks) // short_chunk_size
+            )
+
+    @staticmethod
+    def _element_slice_from_chunk_index(
+        num_elements: int, num_chunks: int, chunk_index: int
+    ) -> slice:
+        short_chunk_size, num_long_chunks = divmod(num_elements, num_chunks)
+        start = chunk_index * short_chunk_size + min(chunk_index, num_long_chunks)
+        end = start + short_chunk_size + int(chunk_index < num_long_chunks)
+        return slice(start, end)
+
+
+def _deduplicate_ordered(arr: List[int]):
+    output = []
+    for item in arr:
+        if len(output) == 0 or item != output[-1]:
+            output.append(item)
+    return output
+
+
+def _log_p2p_op_metrics(
+    p2p_op_infos: List[Tuple[int, List[P2POp]]],
+    num_gpu_per_node: int,
+    world_size: int,
+    self_node_id: int,
+):
+    text = ""
+    all_ops = [op for _, ops in p2p_op_infos for op in ops]
+
+    for direction, ops in _group_by(all_ops, _get_direction_from_op).items():
+        nbytes_of_gpu = [0] * world_size
+        for op in ops:
+            nbytes_of_gpu[op.peer] += op.tensor.nbytes
+        nbytes_of_gpu = torch.tensor(nbytes_of_gpu, dtype=torch.int64)
+
+        nbytes_of_node = einops.reduce(
+            nbytes_of_gpu,
+            "(num_nodes num_gpu_per_node) -> num_nodes",
+            num_gpu_per_node=num_gpu_per_node,
+            reduction="sum",
+        )
+
+        nbytes_curr_node = nbytes_of_node[self_node_id]
+        nbytes_cross_node = torch.sum(nbytes_of_node) - nbytes_curr_node
+
+        text += (
+            f"{direction}_nbytes_of_gpu={nbytes_of_gpu.tolist()} "
+            f"{direction}_nbytes_of_node={nbytes_of_node.tolist()} "
+            f"{direction}_nbytes_curr_node={nbytes_curr_node.item()} "
+            f"{direction}_nbytes_cross_node={nbytes_cross_node.item()} "
+        )
+
+    logger.info(f"[ExpertLocationUpdater] {text}")
+
+
+def _get_direction_from_op(op: P2POp):
+    if op.op == torch.distributed.isend:
+        return "isend"
+    if op.op == torch.distributed.irecv:
+        return "irecv"
+    raise NotImplementedError
+
+
+def _group_by(items, keyfunc):
+    ans = defaultdict(list)
+    for item in items:
+        ans[keyfunc(item)].append(item)
+    return dict(ans)

python/sglang/srt/eplb/lplb_solver.py

@@ -0,0 +1,285 @@
+"""
+LPLBSolver — Linear-Programming Load Balancer for Expert Parallelism.
+
+Encapsulates LP matrix construction (offline, at init/rebalance) and
+per-batch solving (online, per MoE layer forward pass).
+
+Design for DP-attention:
+    Each EP rank counts its local tokens, then all ranks participate in an
+    all-reduce to obtain identical global counts.  Every rank then solves
+    the same LP independently, producing the same log2phy_prob — no
+    broadcast is needed.  Empty-token ranks contribute zeros in the
+    all-reduce so the collective never deadlocks.
+
+Usage:
+    solver = LPLBSolver(phy2log, log2phy, num_gpus, ep_group)
+    log2phy_prob = solver.solve(topk_ids)  # per batch
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Global per-layer LPLB solvers
+
+
+# LP dispatch requires every EP rank to call solver.solve() on every forward
+# pass (including empty-topk ranks under DP-attention) — the all-reduce inside
+# would otherwise hang. Only the DeepSeek-v2 family and its subclasses route
+# empty-rank paths through solver.solve(); other MoE families would deadlock.
+_LPLB_SUPPORTED_MODEL_ARCHS: frozenset[str] = frozenset(
+    {
+        "DeepseekV2ForCausalLM",
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+        "MistralLarge3ForCausalLM",
+        "MistralLarge3ForCausalLMEagle",
+        "Glm4MoeLiteForCausalLM",
+        "GlmMoeDsaForCausalLM",
+    }
+)
+
+
+def assert_lplb_supported_model(architecture: str) -> None:
+    if architecture not in _LPLB_SUPPORTED_MODEL_ARCHS:
+        supported = ", ".join(sorted(_LPLB_SUPPORTED_MODEL_ARCHS))
+        raise NotImplementedError(
+            f"{architecture} does not support --ep-dispatch-algorithm lp. "
+            f"Validated targets: {supported}. Other MoE families have "
+            "empty-token early returns that don't participate in the EP "
+            "all-reduce inside LPLBSolver.solve(), which would deadlock "
+            "under DP-attention."
+        )
+
+
+def get_global_lplb_solver(layer_id: int) -> Optional[LPLBSolver]:
+    from sglang.srt.runtime_context import get_resources
+
+    return get_resources().lplb_solvers.get(layer_id)
+
+
+def set_global_lplb_solver(layer_id: int, solver: LPLBSolver):
+    from sglang.srt.runtime_context import get_resources
+
+    get_resources().lplb_solvers[layer_id] = solver
+
+
+def clear_global_lplb_solvers():
+    from sglang.srt.runtime_context import get_resources
+
+    get_resources().lplb_solvers.clear()
+
+
+class LPLBSolver:
+    """
+    Per-layer LPLB solver.
+
+    At init: pre-computes LP constraint matrices from expert-to-GPU mapping.
+    At solve: takes topk_ids, counts tokens, all-reduces, runs LP,
+              returns log2phy_prob for probability-based token dispatch.
+    """
+
+    def __init__(
+        self,
+        phy2log: torch.Tensor,
+        log2phy: torch.Tensor,
+        num_gpus: int,
+        ep_group=None,
+        logical_to_all_physical_map_num_valid=None,
+    ):
+        """
+        Args:
+            phy2log: (num_physical_experts,) physical-to-logical expert mapping.
+            log2phy: (num_logical_experts, max_copies) logical-to-physical mapping (-1 padded).
+            num_gpus: Number of GPUs in the EP group.
+            ep_group: GroupCoordinator for EP communication (all-reduce).
+            logical_to_all_physical_map_num_valid: (num_logical_experts,) number of valid physical copies.
+        """
+        device = phy2log.device
+        self.num_gpus = num_gpus
+        self.ep_group = ep_group
+        self._has_redundancy = False
+        if logical_to_all_physical_map_num_valid is not None:
+            self._has_redundancy = bool(
+                (logical_to_all_physical_map_num_valid > 1).any()
+            )
+
+        self.num_logical = log2phy.shape[0]
+        self.max_copies = log2phy.shape[1]
+        self.num_phy = phy2log.shape[0]
+        # B1/B2 GPU-assignment matrices below assume each rank owns a
+        # contiguous block of num_phy // num_gpus physical experts.
+        if self.num_phy % num_gpus != 0:
+            raise ValueError(
+                f"LPLBSolver requires num_phy ({self.num_phy}) to be divisible "
+                f"by num_gpus ({num_gpus}); per-rank-contiguous ownership is "
+                "currently the only supported allocation."
+            )
+        num_phy_per_gpu = self.num_phy // num_gpus
+
+        # Count copies per logical expert
+        logcnt = torch.bincount(phy2log, minlength=self.num_logical)
+
+        # Separate single-copy vs replicated experts.
+        # Stored as int64 so they can be used directly as index tensors in
+        # _solve without per-call .long() casts (Tier 1 optimization).
+        self.log_single = torch.nonzero(logcnt == 1).flatten().to(torch.int64)
+        self.phy_single = log2phy[self.log_single, 0].to(torch.int64)
+        self.log_replicated = torch.nonzero(logcnt > 1).flatten().to(torch.int64)
+        self.phy_replicated = (
+            torch.nonzero(logcnt[phy2log] > 1).flatten().to(torch.int64)
+        )
+
+        self.num_single = len(self.log_single)
+        self.num_red_log = len(self.log_replicated)
+        self.num_red_phy = len(self.phy_replicated)
+
+        # Build GPU assignment matrices
+        B_full = torch.zeros(
+            (num_gpus, self.num_phy), dtype=torch.float32, device=device
+        )
+        for i in range(num_gpus):
+            B_full[i, i * num_phy_per_gpu : (i + 1) * num_phy_per_gpu] = 1
+        self.B1 = B_full[:, self.phy_single].contiguous()
+        B2 = B_full[:, self.phy_replicated]
+
+        # Build C matrix (copy-to-logical mapping)
+        C = torch.zeros(
+            (self.num_red_log, self.num_red_phy), dtype=torch.float32, device=device
+        )
+        phy2log_rep = phy2log[self.phy_replicated]
+        for i in range(self.num_red_log):
+            C[i, phy2log_rep == self.log_replicated[i]] = 1.0
+
+        # Build A_base = [[C, 0, 0], [B2, I, -1]]  (without Big-M column)
+        zeros_top_g = torch.zeros(
+            (self.num_red_log, num_gpus), dtype=torch.float32, device=device
+        )
+        zeros_top_1 = torch.zeros(
+            (self.num_red_log, 1), dtype=torch.float32, device=device
+        )
+        I_g = torch.eye(num_gpus, dtype=torch.float32, device=device)
+        neg_ones = torch.full((num_gpus, 1), -1.0, dtype=torch.float32, device=device)
+
+        A_top = torch.hstack([C, zeros_top_g, zeros_top_1])
+        A_bottom = torch.hstack([B2, I_g, neg_ones])
+        self.A_base = torch.vstack([A_top, A_bottom]).contiguous()
+
+        # Objective: minimize M (second-to-last var), penalize Big-M auxiliary
+        nv = self.A_base.shape[1] + 1  # +1 for Big-M column
+        self.c_vec = torch.zeros(nv, dtype=torch.float32, device=device)
+        self.c_vec[-2] = 1.0
+        self.c_vec[-1] = 1000.0
+
+        # Store log2phy as int64 so it can be used directly as index tensor
+        # without per-call .long() casts (Tier 1 optimization).
+        self.log2phy = log2phy.to(torch.int64).contiguous()
+
+        # Pre-JIT-compile the fused IPM kernel for this (NC, NV) shape so the
+        # 20-40s compile cost happens once at startup rather than on the first
+        # real request. No-op when the fused backend is unavailable.
+        nc = self.A_base.shape[0]
+        nv = self.A_base.shape[1] + 1  # +1 for Big-M column added in solve()
+        from sglang.jit_kernel.lplb.torch_solver import warmup as _ipm_warmup
+
+        _ipm_warmup(nc, nv, num_iters=5, device=device)
+
+        # Pre-compute A_base row sum (used in every prep call).
+        self._A_base_row_sum = self.A_base.sum(dim=1).contiguous()  # (NC,)
+
+        # Pre-allocate the buffers the JIT CUDA prep / IPM / post kernels write
+        # into. All writes are contiguous full-tensor stores (no strided
+        # ``out=`` semantics), so the reuse is safe under high concurrency.
+        # Constructed lazily on the first solve() call (we don't know the
+        # device-side log2phy_prob shape until then) — see _solve.
+        self._A_full = torch.empty(nc, nv, dtype=torch.float32, device=device)
+        self._A_full[:, : nv - 1].copy_(self.A_base)
+        self._b = torch.empty(nc, dtype=torch.float32, device=device)
+        self._t1 = torch.empty(self.num_single, dtype=torch.float32, device=device)
+        self._x = torch.empty(nv, dtype=torch.float32, device=device)
+        self._log2phy_prob = torch.empty(
+            log2phy.shape, dtype=torch.float32, device=device
+        )
+
+    def solve(self, topk_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Full LPLB pipeline: count -> all-reduce -> LP solve -> return log2phy_prob.
+
+        All EP ranks must call this method every MoE layer forward pass,
+        including empty-token ranks (which pass an empty topk_ids tensor).
+        This ensures the all-reduce collective does not deadlock under
+        DP-attention where different ranks may have different token counts.
+
+        Args:
+            topk_ids: (num_tokens, topk) int32 tensor of logical expert IDs.
+                      Can be empty (shape (0, topk)) for idle ranks.
+
+        Returns:
+            log2phy_prob: (num_logical, max_copies) float32 probability tensor.
+        """
+        device = topk_ids.device
+
+        # Step 1: Count local tokens per logical expert.
+        # topk_ids comes from the router and is by construction in
+        # [0, num_logical), so we can scatter_add directly without filtering.
+        # Boolean masking + numel() (the previous defensive form) forced a
+        # GPU->host sync on every forward pass via aten::nonzero and a
+        # tensor-shape read; scatter_add on the flattened tensor is async
+        # and a no-op when topk_ids is empty (DP-attention idle rank case).
+        local_counts = torch.zeros(self.num_logical, dtype=torch.int32, device=device)
+        flat_ids = topk_ids.flatten()
+        local_counts.scatter_add_(
+            0,
+            flat_ids.long(),
+            torch.ones_like(flat_ids, dtype=torch.int32),
+        )
+
+        # Step 2: All-reduce to get global counts across all EP ranks.
+        # All EP ranks must participate — empty-token ranks contribute zeros.
+        # After all-reduce, every rank has identical global_counts and solves
+        # the same LP independently, so no broadcast is needed.
+        # GroupCoordinator.all_reduce may be in-place (pynccl) or out-of-place
+        # (ca_comm / pymscclpp / ...) depending on tensor size; small tensors
+        # like ours (~num_logical * 4 B) typically take the out-of-place path,
+        # so we must capture the return value.
+        global_counts = local_counts.float()
+        if self.ep_group is not None:
+            global_counts = self.ep_group.all_reduce(global_counts)
+
+        # Step 3: Run LP solver
+        return self._solve(global_counts)
+
+    def _solve(self, global_counts: torch.Tensor) -> torch.Tensor:
+        """Three CUDA kernel launches replace ~14 torch ops.
+
+        Pipeline (all writes go into pre-allocated buffers from __init__):
+            prep_lp_inputs → solve_ipm → extract_log2phy_prob
+        Raises if the JIT CUDA backend is unavailable.
+        """
+        from sglang.jit_kernel.lplb import cuda_solver
+
+        cuda_solver.prep_lp_inputs(
+            self._A_full,
+            self._b,
+            self._t1,
+            global_counts,
+            self.log_single,
+            self.log_replicated,
+            self.B1,
+            self._A_base_row_sum,
+        )
+        cuda_solver.solve_ipm(self._A_full, self._b, self.c_vec, result=self._x)
+        cuda_solver.extract_log2phy_prob(
+            self._log2phy_prob,
+            self._x,
+            self._t1,
+            self.phy_single,
+            self.phy_replicated,
+            self.log2phy,
+        )
+        return self._log2phy_prob