[EPLB] Add the MoE Load Balancer policy: pluggable expert placement and replica routing

.buildkite/test_areas/expert_parallelism.yaml

@@ -9,13 +9,22 @@ steps:
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/distributed/eplb
+  - vllm/model_executor/layers/fused_moe/router/base_router.py
   - tests/distributed/test_eplb_algo.py
   - tests/distributed/test_eplb_migration_scheduler.py
   - tests/distributed/test_eplb_utils.py
+  - tests/distributed/test_eplb_mlb_policy.py
+  - tests/distributed/test_eplb_mlb_routing.py
+  - tests/distributed/test_eplb_replica_shares_kernel.py
   commands:
   - pytest -v -s distributed/test_eplb_algo.py
   - pytest -v -s distributed/test_eplb_migration_scheduler.py
   - pytest -v -s distributed/test_eplb_utils.py
+  # The MoE Load Balancer tests skip themselves when the optional
+  # moe_load_balancer package is not installed.
+  - pytest -v -s distributed/test_eplb_mlb_policy.py
+  - pytest -v -s distributed/test_eplb_mlb_routing.py
+  - pytest -v -s distributed/test_eplb_replica_shares_kernel.py
   mirror:
     amd:
       label: ":amd: (MI355 DPX) EPLB Algorithm"
@@ -26,8 +35,12 @@ steps:
       - image-build-amd
       source_file_dependencies:
       - vllm/distributed/eplb
+      - vllm/model_executor/layers/fused_moe/router/base_router.py
       - tests/distributed/test_eplb_algo.py
       - tests/distributed/test_eplb_utils.py
+      - tests/distributed/test_eplb_mlb_policy.py
+      - tests/distributed/test_eplb_mlb_routing.py
+      - tests/distributed/test_eplb_replica_shares_kernel.py
       - vllm/platforms/rocm.py
 
 - label: ":nvidia: (L4) EPLB Execution"

docs/serving/expert_parallel_deployment.md

@@ -155,7 +155,10 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
 | `log_balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` |
 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
 | `use_async` | Use non-blocking EPLB for reduced latency overhead | `true` |
-| `policy` | The policy type for expert parallel load balancing | `"default"` |
+| `policy` | The policy type for expert parallel load balancing: `"default"` (built-in) or `"mlb"` (MoE Load Balancer, see below) | `"default"` |
+| `l2_algorithm` | MoE Load Balancer replica-routing expression (only with `policy: "mlb"`); empty keeps vLLM's built-in replica choice | `""` |
+| `init_placement_path` | Load a saved expert placement at start-up instead of the trivial layout | `null` |
+| `save_placement_path` | Save the expert placement to this path after every rearrangement | `null` |
 | `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`,  or `null` (auto) | `null` |
 
 For example:
@@ -177,6 +180,50 @@ vllm serve Qwen/Qwen3-30B-A3B \
             --eplb-config.log_balancedness true
     ```
 
+### MoE Load Balancer policy
+
+`policy: "mlb"` delegates expert placement, and optionally per-token replica routing, to the
+[MoE Load Balancer](https://github.com/xutizhou/moe_load_balancer) (`moe_load_balancer`), a
+framework-neutral library of expert-parallel load-balancing policies. vLLM imports it lazily, so
+nothing changes unless the policy is selected; install it into the serving environment first.
+
+Two stages are configurable:
+
+- **L1, placement** (`VLLM_MLB_L1_ALGORITHM`): which logical expert lives in which physical slot,
+  decided at every rearrangement. `auto` (default) selects the balancer's default algorithm for the
+  model's expert-group layout; `ultraep` re-solves placement and per-replica quotas online every
+  `MLB_ULTRAEP_REFRESH_INTERVAL` representative batches and moves the affected expert weights itself.
+- **L2, replica routing** (`--eplb-config.l2_algorithm` or `VLLM_MLB_L2_ALGORITHM`): which replica
+  of a logical expert serves each token. Expressions combine a routed-expert policy (`static`,
+  `lplb`, `ultraep`) with an optional shared-expert policy (`+waterfill`), e.g. `lplb+waterfill`.
+  `waterfill` dispatches the shared expert through the expert-parallel all-to-all so that it can be
+  assigned to the least-loaded rank; this widens the expert space by one slot per rank and can be
+  enabled on its own with `VLLM_FUSE_SHARED_EXPERTS=1`. L2 routing is skipped automatically when
+  there are no redundant experts, where it could only reproduce the built-in choice.
+
+```bash
+# Placement fitted by the balancer, online UltraEP refresh, Waterfill for the shared expert
+VLLM_MLB_L1_ALGORITHM=ultraep MLB_ULTRAEP_REFRESH_INTERVAL=8 \
+vllm serve deepseek-ai/DeepSeek-V3 \
+    --tensor-parallel-size 1 --data-parallel-size 8 --enable-expert-parallel \
+    --enable-eplb \
+    --eplb-config '{"policy":"mlb","use_async":false,"num_redundant_experts":16,"l2_algorithm":"ultraep+waterfill","log_balancedness":true}'
+```
+
+`use_async` must be `false` with this policy: the balancer plans synchronously on the caller's
+thread. To fit a placement once and serve with it held, run a short profiling serve with
+`"save_placement_path": "/path/placement.pt"` and start production with
+`"init_placement_path": "/path/placement.pt"` and a `step_interval` that is never reached.
+
+Tuning knobs for the online UltraEP refresh (all read through `vllm.envs`):
+
+| Variable | Default | Meaning |
+| -------- | ------- | ------- |
+| `MLB_ULTRAEP_REFRESH_INTERVAL` | `64` | Re-solve every N representative batches; `8` is a good production value, `1` maximizes balance at the cost of one refresh per step |
+| `MLB_ULTRAEP_REFRESH_MIN_TOKENS` | `512` | A batch counts as representative when it routes at least this many tokens |
+| `MLB_ULTRAEP_MOVER` | `direct` | How re-solved replica weights move between ranks: `direct` (send/recv over the EPLB communicator), `symm` (torch symmetric memory; one host or a multi-node NVLink fabric), `auto` (symm when every rank agrees it can, otherwise direct) |
+| `MLB_ULTRAEP_LAGGED_APPLY` | `0` | `1` lands a plan on the layer's next forward, avoiding a per-layer device sync at the price of one step of staleness; worth about +3% at interval 1 on Blackwell-class step times, nothing at interval 8 |
+
 ### Expert Distribution Formula
 
 - **Default**: Each EP rank has `NUM_TOTAL_EXPERTS ÷ NUM_EP_RANKS` experts

tests/distributed/test_eplb_execute.py

@@ -3,6 +3,7 @@
 
 import random
 
+import numpy as np
 import pytest
 import torch
 import torch.distributed
@@ -15,6 +16,8 @@
     has_nixl,
 )
 from vllm.distributed.eplb.rebalance_execute import (
+    experts_being_acquired,
+    get_ep_ranks_with_experts_batch,
     move_from_buffer,
     rearrange_expert_weights_inplace,
     transfer_layer,
@@ -855,6 +858,71 @@ def test_rearrange_expert_weights_profile_mode(world_size):
     )
 
 
+@pytest.mark.parametrize("seed", range(20))
+def test_narrowed_send_query_drops_only_experts_with_nowhere_to_send(seed: int):
+    """Experts move_to_buffer stops asking about could not have sent anything.
+
+    move_to_buffer asks get_ep_ranks_with_experts_batch only about experts that
+    occupy a position where the new placement differs from the old, rather than
+    about every expert resident on the rank. That is sound only because the
+    send loop skips any expert whose ranks_to_recv is empty, and an expert
+    acquired by no rank has exactly that. The narrowing and the skip live in
+    different functions with nothing tying them together, so assert the
+    implication against the real mapping instead of trusting the argument for
+    it -- a change to how recv ranks are derived would otherwise turn this into
+    silently dropped transfers.
+    """
+    rng = np.random.default_rng(seed)
+    ep_size, num_local = 8, 34
+    num_physical = ep_size * num_local
+    num_logical = 256
+
+    old_indices = np.concatenate(
+        [
+            np.arange(num_logical),
+            rng.choice(num_logical, num_physical - num_logical, replace=False),
+        ]
+    )
+    rng.shuffle(old_indices)
+    new_indices = old_indices.copy()
+    moved = rng.choice(num_physical, int(rng.integers(1, 24)), replace=False)
+    new_indices[moved] = rng.choice(num_logical, moved.size, replace=False)
+
+    arriving = experts_being_acquired(old_indices, new_indices)
+
+    dropped_any = False
+    kept_with_work = False
+    for ep_rank in range(ep_size):
+        local = slice(ep_rank * num_local, (ep_rank + 1) * num_local)
+        resident = np.unique(old_indices[local])
+
+        dropped = resident[~np.isin(resident, arriving)]
+        if dropped.size:
+            dropped_any = True
+            _, recv_map = get_ep_ranks_with_experts_batch(
+                dropped, num_local, old_indices, new_indices
+            )
+            for expert in dropped.tolist():
+                assert not recv_map[int(expert)], (
+                    f"rank {ep_rank} expert {expert} was dropped by the "
+                    f"narrowing but has receivers {recv_map[int(expert)]}"
+                )
+
+        kept = resident[np.isin(resident, arriving)]
+        if kept.size:
+            _, recv_map = get_ep_ranks_with_experts_batch(
+                kept, num_local, old_indices, new_indices
+            )
+            kept_with_work |= any(recv_map[int(e)] for e in kept.tolist())
+
+    # Both guard against a vacuous pass: nothing dropped means the narrowing
+    # was never exercised, and nothing kept with receivers means the placement
+    # asked for no transfers at all, in which case an empty recv_map says
+    # nothing about whether the narrowing is selective.
+    assert dropped_any, "narrowing dropped nothing -- the case is not exercised"
+    assert kept_with_work, "no expert had receivers -- recv_map is empty for all"
+
+
 def _test_eplb_communicator_pipeline_parallel(env, world_size, backend):
     set_env_vars_and_device(env)
     config = VllmConfig()

tests/distributed/test_eplb_mlb_policy.py

@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the MLB-backed EPLB placement policy.
+
+These run on CPU and need no distributed group: the placement contract is a
+pure function of the load statistics and the topology.
+
+They are skipped unless the ``moe-load-balancer`` package is importable.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.distributed.eplb.policy import EPLB_POLICIES, DefaultEplbPolicy
+
+pytest.importorskip("moe_load_balancer")
+
+from vllm.distributed.eplb.policy.mlb import MlbEplbPolicy  # noqa: E402
+
+# DeepSeek-R1 geometry: 58 MoE layers, 256 routed experts, EP=8, 32 redundant.
+NUM_LAYERS = 58
+NUM_LOGICAL = 256
+NUM_RANKS = 8
+NUM_REPLICAS = NUM_LOGICAL + 32
+NUM_GROUPS = 8
+
+
+def _skewed_load(seed: int = 0) -> torch.Tensor:
+    """A realistic per-layer expert load: a hot head plus a long tail."""
+    rng = np.random.default_rng(seed)
+    base = rng.pareto(1.6, size=(NUM_LAYERS, NUM_LOGICAL)) + 1.0
+    # Shuffle which experts are hot per layer, as real routing does.
+    for layer in range(NUM_LAYERS):
+        rng.shuffle(base[layer])
+    return torch.from_numpy((base * 1000).astype(np.float32))
+
+
+def _rank_loads(phy2log: torch.Tensor, weight: torch.Tensor) -> np.ndarray:
+    """Per-layer, per-rank load if each replica takes an equal share."""
+    phy = phy2log.numpy()
+    w = weight.numpy()
+    replicas = np.zeros_like(w)
+    np.add.at(
+        replicas,
+        (np.arange(phy.shape[0])[:, None].repeat(phy.shape[1], 1), phy),
+        1,
+    )
+    per_replica = w / np.maximum(replicas, 1)
+    slots_per_rank = phy.shape[1] // NUM_RANKS
+    loads = np.take_along_axis(per_replica, phy, axis=1)
+    return loads.reshape(phy.shape[0], NUM_RANKS, slots_per_rank).sum(axis=2)
+
+
+def _imbalance(phy2log: torch.Tensor, weight: torch.Tensor) -> float:
+    loads = _rank_loads(phy2log, weight)
+    return float(np.mean(loads.max(axis=1) / np.maximum(loads.mean(axis=1), 1e-9)))
+
+
+def _assert_valid_placement(phy2log: torch.Tensor) -> None:
+    assert phy2log.shape == (NUM_LAYERS, NUM_REPLICAS)
+    assert phy2log.dtype == torch.int64
+    assert phy2log.device.type == "cpu"
+    arr = phy2log.numpy()
+    assert arr.min() >= 0 and arr.max() < NUM_LOGICAL
+    # Every logical expert must own at least one physical slot in every layer,
+    # otherwise its tokens have nowhere to go.
+    for layer in range(NUM_LAYERS):
+        assert len(np.unique(arr[layer])) == NUM_LOGICAL
+
+
+def test_mlb_policy_is_registered():
+    assert EPLB_POLICIES["mlb"] is MlbEplbPolicy
+
+
+def test_mlb_placement_is_valid():
+    weight = _skewed_load()
+    phy2log = MlbEplbPolicy.rebalance_experts(
+        weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS
+    )
+    _assert_valid_placement(phy2log)
+
+
+def test_mlb_placement_is_deterministic():
+    weight = _skewed_load()
+    a = MlbEplbPolicy.rebalance_experts(weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS)
+    b = MlbEplbPolicy.rebalance_experts(weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS)
+    assert torch.equal(a, b), "placement must be reproducible across EP ranks"
+
+
+def test_mlb_balances_at_least_as_well_as_default():
+    """MLB's DeepSeek EPLB and vLLM's are ports of the same algorithm, so the
+    balance quality should be equivalent -- not merely 'not much worse'."""
+    weight = _skewed_load()
+    mlb = MlbEplbPolicy.rebalance_experts(
+        weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS
+    )
+    default = DefaultEplbPolicy.rebalance_experts(
+        weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS
+    )
+    mlb_imb = _imbalance(mlb, weight)
+    default_imb = _imbalance(default, weight)
+    assert mlb_imb <= default_imb * 1.02, (
+        f"MLB imbalance {mlb_imb:.4f} materially worse than "
+        f"vLLM default {default_imb:.4f}"
+    )
+
+
+def test_slot_preservation_reduces_weight_movement():
+    """The built-in policy preserves intra-GPU slots to avoid weight copies.
+    MLB's L1 policies do not consume a previous placement, so the glue applies
+    the same pass; without it, a re-plan churns far more slots."""
+    weight = _skewed_load(seed=0)
+    first = MlbEplbPolicy.rebalance_experts(
+        weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS
+    )
+
+    drifted = _skewed_load(seed=1) * 0.25 + weight * 0.75
+    without = MlbEplbPolicy.rebalance_experts(
+        drifted, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS
+    )
+    with_preserve = MlbEplbPolicy.rebalance_experts(
+        drifted, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS, old_global_expert_indices=first
+    )
+
+    _assert_valid_placement(with_preserve)
+    moved_without = int((without != first).sum())
+    moved_with = int((with_preserve != first).sum())
+    assert moved_with < moved_without, (
+        f"slot preservation did not reduce churn: {moved_with} vs {moved_without}"
+    )
+
+    # Preservation must only permute slots inside a rank, never move an expert
+    # to a different rank than MLB placed it on.
+    slots = NUM_REPLICAS // NUM_RANKS
+    a = without.numpy().reshape(NUM_LAYERS, NUM_RANKS, slots)
+    b = with_preserve.numpy().reshape(NUM_LAYERS, NUM_RANKS, slots)
+    assert np.array_equal(np.sort(a, axis=2), np.sort(b, axis=2)), (
+        "slot preservation changed which experts live on which rank"
+    )
+
+
+@pytest.mark.parametrize(
+    "algorithm",
+    ["auto", "deepseek", "deepseek_hierarchical", "deepseek_vec"],
+)
+def test_mlb_algorithms_selectable(monkeypatch, algorithm):
+    monkeypatch.setenv("VLLM_MLB_L1_ALGORITHM", algorithm)
+    weight = _skewed_load()
+    phy2log = MlbEplbPolicy.rebalance_experts(
+        weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS
+    )
+    _assert_valid_placement(phy2log)
+
+
+def test_no_expert_groups():
+    """VLLM passes num_groups=0 for models without expert groups."""
+    weight = _skewed_load()
+    phy2log = MlbEplbPolicy.rebalance_experts(weight, NUM_REPLICAS, 0, 1, NUM_RANKS)
+    _assert_valid_placement(phy2log)
+
+
+def test_ep_rank_reaches_the_placement_request(monkeypatch):
+    """Ultraep's placement kernel solves once per EP rank and validates the
+    value is in range, so the caller's rank must reach the request -- every
+    other algorithm ignores it, but the plumbing is shared."""
+    import vllm.distributed.eplb.mlb_runtime as mlb_runtime
+
+    real_placement_request = mlb_runtime.placement_request
+    captured = {}
+
+    def spy(*args, **kwargs):
+        captured["ep_rank"] = kwargs.get("ep_rank")
+        return real_placement_request(*args, **kwargs)
+
+    monkeypatch.setattr(mlb_runtime, "placement_request", spy)
+    weight = _skewed_load()
+    MlbEplbPolicy.rebalance_experts(
+        weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS, ep_rank=5
+    )
+    assert captured["ep_rank"] == 5