[EPLB] Add contention-aware expert migration batching

.buildkite/test_areas/expert_parallelism.yaml

@@ -10,9 +10,11 @@ steps:
   source_file_dependencies:
   - vllm/distributed/eplb
   - tests/distributed/test_eplb_algo.py
+  - tests/distributed/test_eplb_migration_scheduler.py
   - tests/distributed/test_eplb_utils.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
   mirror:
     amd:

tests/distributed/test_eplb_migration_scheduler.py

@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config.parallel import EPLBConfig
+from vllm.distributed.eplb.eplb_communicator import EplbCommunicator
+from vllm.distributed.eplb.migration_scheduler import (
+    MigrationFlow,
+    schedule_migration_batches,
+    schedule_migration_batches_for_layers,
+)
+from vllm.distributed.eplb.rebalance_execute import move_from_buffer, move_to_buffer
+
+
+class _MockEplbCommunicator(EplbCommunicator):
+    def __init__(self) -> None:
+        self.send_calls: list[tuple[list[torch.Tensor], int, int]] = []
+        self.recv_calls: list[tuple[list[torch.Tensor], int, int]] = []
+        self.context_calls = 0
+        self.execute_count = 0
+
+    def add_send(
+        self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int
+    ) -> None:
+        self.send_calls.append((tensors, dst_rank, expert_id))
+
+    def add_recv(
+        self, tensors: list[torch.Tensor], src_rank: int, expert_id: int
+    ) -> None:
+        self.recv_calls.append((tensors, src_rank, expert_id))
+
+    def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
+        self.context_calls += 1
+
+    def execute(self) -> None:
+        self.execute_count += 1
+
+
+def test_migration_batching_is_disabled_by_default_and_can_be_enabled() -> None:
+    assert not EPLBConfig().enable_migration_batching
+    assert EPLBConfig(enable_migration_batching=True).enable_migration_batching
+
+
+def test_migration_batching_only_applies_to_async_eplb() -> None:
+    assert not EPLBConfig(use_async=True).migration_batching_enabled
+    assert EPLBConfig(
+        use_async=True, enable_migration_batching=True
+    ).migration_batching_enabled
+    assert not EPLBConfig(
+        use_async=True, enable_migration_batching=False
+    ).migration_batching_enabled
+    assert not EPLBConfig(
+        use_async=False, enable_migration_batching=True
+    ).migration_batching_enabled
+
+
+def test_schedule_migration_batches_is_deterministic() -> None:
+    transfers = [
+        (1, 3, 0),
+        (2, 4, 1),
+        (0, 1, 2),
+        (0, 2, 3),
+        (0, 3, 4),
+        (0, 4, 5),
+    ]
+    num_local_experts, old_indices, new_indices = _placements_for_transfers(
+        5, transfers
+    )
+
+    batches = schedule_migration_batches(num_local_experts, old_indices, new_indices)
+
+    assert [_endpoints(batch) for batch in batches] == [
+        [(1, 3), (2, 4)],
+        [(0, 1)],
+        [(0, 2)],
+        [(0, 3)],
+        [(0, 4)],
+    ]
+    for batch in batches:
+        _assert_no_endpoint_conflict(batch)
+
+
+def test_schedule_migration_batches_coalesces_rank_pair() -> None:
+    transfers = [
+        (0, 1, 0),
+        (0, 1, 1),
+        (2, 3, 2),
+    ]
+    num_local_experts, old_indices, new_indices = _placements_for_transfers(
+        4, transfers
+    )
+
+    assert schedule_migration_batches(num_local_experts, old_indices, new_indices) == [
+        [
+            MigrationFlow(0, 1, expert_ids=(0, 1)),
+            MigrationFlow(2, 3, expert_ids=(2,)),
+        ]
+    ]
+
+
+def test_schedule_migration_batches_covers_random_instructions() -> None:
+    rng = np.random.default_rng(42)
+    transfers = []
+    for expert_id in range(80):
+        src_rank, dst_rank = rng.integers(0, 8, size=2).tolist()
+        if src_rank != dst_rank:
+            transfers.append((src_rank, dst_rank, expert_id))
+    num_local_experts, old_indices, new_indices = _placements_for_transfers(
+        8, transfers
+    )
+
+    batches = schedule_migration_batches(num_local_experts, old_indices, new_indices)
+
+    scheduled = _flatten_transfers(batches)
+    assert len(scheduled) == len(transfers)
+    assert set(scheduled) == set(transfers)
+    for batch in batches:
+        _assert_no_endpoint_conflict(batch)
+
+
+def test_schedule_migration_batches_excludes_local_copies() -> None:
+    # Old: rank 0 [0, 1], rank 1 [1, 2]. New: rank 1 needs expert 0.
+    old_indices = np.array([0, 1, 1, 2], dtype=np.int64)
+    new_indices = np.array([0, 1, 0, 2], dtype=np.int64)
+
+    assert schedule_migration_batches(2, old_indices, new_indices) == [
+        [MigrationFlow(0, 1, expert_ids=(0,))]
+    ]
+
+
+def test_schedule_migration_batches_requires_matching_shapes() -> None:
+    old_indices = np.array([0, 1], dtype=np.int64)
+    new_indices = np.array([0], dtype=np.int64)
+
+    with pytest.raises(AssertionError):
+        schedule_migration_batches(1, old_indices, new_indices)
+
+
+def test_schedule_migration_batches_balances_replicas() -> None:
+    # Rank 0 and rank 1 hold expert 0; ranks 2, 3, and 4 need it.
+    old_indices = np.array([0, 0, 1, 1, 1], dtype=np.int64)
+    new_indices = np.array([0, 0, 0, 0, 0], dtype=np.int64)
+
+    batches = schedule_migration_batches(1, old_indices, new_indices)
+
+    assert set(_flatten_transfers(batches)) == {
+        (0, 2, 0),
+        (0, 4, 0),
+        (1, 3, 0),
+    }
+
+
+def test_schedule_migration_batches_for_layers_matches_individual_calls() -> None:
+    old_indices = np.array([[0, 1, 2, 3], [3, 2, 1, 0]], dtype=np.int64)
+    new_indices = np.array([[1, 2, 3, 0], [0, 3, 2, 1]], dtype=np.int64)
+
+    assert schedule_migration_batches_for_layers(1, old_indices, new_indices) == [
+        schedule_migration_batches(1, old_indices[0], new_indices[0]),
+        schedule_migration_batches(1, old_indices[1], new_indices[1]),
+    ]
+
+
+@pytest.mark.parametrize("use_precomputed", [False, True])
+def test_move_to_buffer_uses_multiple_batches(use_precomputed: bool) -> None:
+    old_indices = np.array([0, 1, 2, 3], dtype=np.int64)
+    new_indices = np.array([1, 2, 3, 0], dtype=np.int64)
+    communicator = _MockEplbCommunicator()
+    migration_batches = (
+        schedule_migration_batches(1, old_indices, new_indices)
+        if use_precomputed
+        else None
+    )
+
+    move_to_buffer(
+        num_local_experts=1,
+        old_indices=old_indices,
+        new_indices=new_indices,
+        expert_weights=[torch.zeros(1, 1)],
+        expert_weights_buffers=[torch.zeros(1, 1)],
+        stream=None,
+        ep_rank=0,
+        communicator=communicator,
+        layer_idx=7,
+        enable_migration_batching=True,
+        migration_batches=migration_batches,
+    )
+
+    assert communicator.context_calls == 2
+    assert communicator.execute_count == 2
+    assert [(dst, expert) for _, dst, expert in communicator.send_calls] == [(3, 0)]
+    assert [(src, expert) for _, src, expert in communicator.recv_calls] == [(1, 1)]
+
+
+def test_move_to_buffer_uses_primary_duplicate_destination() -> None:
+    old_indices = np.array([0, 1, 2, 3], dtype=np.int64)
+    new_indices = np.array([2, 2, 0, 3], dtype=np.int64)
+    expert_weights = [torch.zeros(2, 1)]
+    recv_buffer = torch.zeros(2, 1)
+    communicator = _MockEplbCommunicator()
+
+    transfer_metadata = move_to_buffer(
+        num_local_experts=2,
+        old_indices=old_indices,
+        new_indices=new_indices,
+        expert_weights=expert_weights,
+        expert_weights_buffers=[recv_buffer],
+        stream=None,
+        ep_rank=0,
+        communicator=communicator,
+        enable_migration_batching=True,
+    )
+
+    assert len(communicator.recv_calls) == 1
+    recv_tensors, src_rank, expert_id = communicator.recv_calls[0]
+    assert (src_rank, expert_id) == (1, 2)
+    assert recv_tensors[0].data_ptr() == recv_buffer[0].data_ptr()
+
+    recv_buffer[0].fill_(42)
+    move_from_buffer(
+        expert_weights=expert_weights,
+        expert_weights_buffers=[recv_buffer],
+        transfer_metadata=transfer_metadata,
+        new_indices=new_indices,
+        ep_rank=0,
+    )
+    torch.testing.assert_close(expert_weights[0], torch.full((2, 1), 42.0))
+
+
+def _endpoints(
+    flows: list[MigrationFlow],
+) -> list[tuple[int, int]]:
+    return [(item.src_rank, item.dst_rank) for item in flows]
+
+
+def _assert_no_endpoint_conflict(batch: list[MigrationFlow]) -> None:
+    endpoints: set[int] = set()
+    for flow in batch:
+        assert flow.src_rank not in endpoints
+        assert flow.dst_rank not in endpoints
+        endpoints.update((flow.src_rank, flow.dst_rank))
+
+
+def _flatten_transfers(
+    batches: list[list[MigrationFlow]],
+) -> list[tuple[int, int, int]]:
+    return [
+        (flow.src_rank, flow.dst_rank, expert_id)
+        for batch in batches
+        for flow in batch
+        for expert_id in flow.expert_ids
+    ]
+
+
+def _placements_for_transfers(
+    num_ranks: int,
+    transfers: list[tuple[int, int, int]],
+) -> tuple[int, np.ndarray, np.ndarray]:
+    old_counts = [0] * num_ranks
+    new_counts = [0] * num_ranks
+    for src_rank, dst_rank, _ in transfers:
+        old_counts[src_rank] += 1
+        new_counts[dst_rank] += 1
+    num_local_experts = max(1, *old_counts, *new_counts)
+    old_indices = np.full(num_ranks * num_local_experts, -1, dtype=np.int64)
+    new_indices = np.full_like(old_indices, -1)
+    old_counts = [0] * num_ranks
+    new_counts = [0] * num_ranks
+    for src_rank, dst_rank, expert_id in transfers:
+        old_offset = src_rank * num_local_experts + old_counts[src_rank]
+        new_offset = dst_rank * num_local_experts + new_counts[dst_rank]
+        old_indices[old_offset] = expert_id
+        new_indices[new_offset] = expert_id
+        old_counts[src_rank] += 1
+        new_counts[dst_rank] += 1
+    return num_local_experts, old_indices, new_indices

vllm/config/parallel.py

@@ -106,6 +106,17 @@ class EPLBConfig:
       on CUDA, falls back to "torch_gloo")
     """
 
+    enable_migration_batching: bool = False
+    """Schedule expert migrations in batches where each rank communicates with
+    at most one peer. This reduces per-rank network contention at the cost of
+    additional sequential communication steps. This option only applies to
+    async EPLB and is disabled by default."""
+
+    @property
+    def migration_batching_enabled(self) -> bool:
+        """Whether contention-aware batching is active for this configuration."""
+        return self.use_async and self.enable_migration_batching
+
     @model_validator(mode="after")
     def _validate_eplb_config(self) -> Self:
         if self.use_async and self.policy != "default":

vllm/distributed/eplb/async_worker.py

@@ -11,6 +11,7 @@
 from vllm.logger import init_logger
 
 from .eplb_utils import CpuGpuEvent, device_stream
+from .migration_scheduler import schedule_migration_batches_for_layers
 from .rebalance_execute import AsyncEplbLayerResult, transfer_layer
 
 if TYPE_CHECKING:
@@ -99,6 +100,21 @@ def transfer_run_periodically(
                 model_state, state, physical_to_logical_map_cpu, stream
             )
 
+            migration_batches_by_layer = None
+            if state.parallel_config.eplb_config.migration_batching_enabled:
+                profile_name = "eplb: schedule migration batches"
+                num_local_experts = model_state.model.expert_weights[0][0].shape[0]
+                # The placement snapshots stay fixed for this rebalance cycle.
+                # Precompute all layer schedules once, but execute transfers one
+                # layer at a time below to preserve the migration lifecycle. This
+                # profiling range therefore covers scheduling for every MoE layer.
+                with torch.profiler.record_function(profile_name):
+                    migration_batches_by_layer = schedule_migration_batches_for_layers(
+                        num_local_experts,
+                        physical_to_logical_map_cpu.numpy(),
+                        new_physical_to_logical_map.numpy(),
+                    )
+
             # Execute one EPLB layer transfer per model forward pass. Each iteration
             # of this loop will copy the new set of expert weights into
             # model_state.expert_buffer, which will be consumed by the main thread in
@@ -134,6 +150,16 @@ def transfer_run_periodically(
                     is_profile=is_profile,
                     stream=stream,
                     layer_idx=layer_idx,
+                    enable_migration_batching=(
+                        state.parallel_config.eplb_config.migration_batching_enabled
+                    ),
+                    migration_batches=(
+                        # This is only the current layer's schedule; it does not
+                        # combine transfers from different layers.
+                        migration_batches_by_layer[layer_idx]
+                        if migration_batches_by_layer is not None
+                        else None
+                    ),
                 )
 
                 # Wait until all writes to expert_buffer have finished before making the