From 4f6321805051733b10815b6e0810529f35cd1764 Mon Sep 17 00:00:00 2001
From: DOCCA0 <2947323341@qq.com>
Date: Sat, 5 Sep 2026 18:56:25 +0800
Subject: [PATCH 1/6] Add contention-aware EPLB migration batching

Signed-off-by: DOCCA0 <2947323341@qq.com>
---
 .../test_eplb_migration_scheduler.py          | 247 ++++++++++++++++++
 vllm/config/parallel.py                       |   6 +
 vllm/distributed/eplb/async_worker.py         |   3 +
 vllm/distributed/eplb/eplb_state.py           |   3 +
 vllm/distributed/eplb/migration_scheduler.py  |  87 ++++++
 vllm/distributed/eplb/rebalance_execute.py    | 143 +++++++++-
 6 files changed, 481 insertions(+), 8 deletions(-)
 create mode 100644 tests/distributed/test_eplb_migration_scheduler.py
 create mode 100644 vllm/distributed/eplb/migration_scheduler.py

diff --git a/tests/distributed/test_eplb_migration_scheduler.py b/tests/distributed/test_eplb_migration_scheduler.py
new file mode 100644
index 000000000000..0c1753864848
--- /dev/null
+++ b/tests/distributed/test_eplb_migration_scheduler.py
@@ -0,0 +1,247 @@
+# 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,
+)
+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_enabled_by_default_and_can_be_disabled() -> None:
+    assert EPLBConfig().enable_migration_batching
+    assert not EPLBConfig(enable_migration_batching=False).enable_migration_batching
+
+
+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_move_to_buffer_uses_multiple_batches() -> 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()
+
+    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)],
+        cuda_stream=None,
+        ep_rank=0,
+        communicator=communicator,
+        layer_idx=7,
+        enable_migration_batching=True,
+    )
+
+    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],
+        cuda_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
diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
index 26b27622c8bf..b1948a032d30 100644
--- a/vllm/config/parallel.py
+++ b/vllm/config/parallel.py
@@ -99,6 +99,12 @@ class EPLBConfig:
     - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")
     """
 
+    enable_migration_batching: bool = True
+    """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. Set to False to use one
+    communication step for all migrations."""
+
     @model_validator(mode="after")
     def _validate_eplb_config(self) -> Self:
         if self.use_async and self.policy != "default":
diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py
index ff872c8c2220..5f77041b823a 100644
--- a/vllm/distributed/eplb/async_worker.py
+++ b/vllm/distributed/eplb/async_worker.py
@@ -136,6 +136,9 @@ def transfer_run_periodically(
                     is_profile=is_profile,
                     cuda_stream=cuda_stream,
                     layer_idx=layer_idx,
+                    enable_migration_batching=(
+                        state.parallel_config.eplb_config.enable_migration_batching
+                    ),
                 )
 
                 # Wait until all writes to expert_buffer have finished before making the
diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py
index 580477ba239e..c878f38c7056 100644
--- a/vllm/distributed/eplb/eplb_state.py
+++ b/vllm/distributed/eplb/eplb_state.py
@@ -899,6 +899,9 @@ def rank_load_imbalance(
                         eplb_model_state.communicator,
                         is_profile,
                         rank_mapping,
+                        enable_migration_batching=(
+                            self.parallel_config.eplb_config.enable_migration_batching
+                        ),
                     )
 
                     if not is_profile:
diff --git a/vllm/distributed/eplb/migration_scheduler.py b/vllm/distributed/eplb/migration_scheduler.py
new file mode 100644
index 000000000000..72e7651f9669
--- /dev/null
+++ b/vllm/distributed/eplb/migration_scheduler.py
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Schedule EPLB expert migrations into conflict-free batches."""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass(frozen=True, slots=True)
+class MigrationFlow:
+    """Expert transfers between one directed rank pair."""
+
+    src_rank: int
+    dst_rank: int
+    expert_ids: tuple[int, ...]
+
+
+def schedule_migration_batches(
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+) -> list[list[MigrationFlow]]:
+    """Build and greedily group rank-pair-disjoint migration flows."""
+    assert old_indices.shape == new_indices.shape
+    recv_ranks_by_expert: dict[int, list[int]] = {}
+    old_experts_by_rank: list[set[int]] = []
+    old_by_rank = old_indices.reshape(-1, num_local_experts)
+    new_by_rank = new_indices.reshape(-1, num_local_experts)
+    for rank, (old_local, new_local) in enumerate(zip(old_by_rank, new_by_rank)):
+        old_experts = set(old_local.tolist())
+        old_experts.discard(-1)
+        old_experts_by_rank.append(old_experts)
+
+        new_experts = set(new_local.tolist())
+        new_experts.discard(-1)
+        for expert_id in new_experts - old_experts:
+            recv_ranks_by_expert.setdefault(expert_id, []).append(rank)
+
+    old_ranks_by_expert: dict[int, list[int]] = {}
+    needed_experts = recv_ranks_by_expert.keys()
+    for rank, old_experts in enumerate(old_experts_by_rank):
+        for expert_id in old_experts & needed_experts:
+            old_ranks_by_expert.setdefault(expert_id, []).append(rank)
+
+    expert_ids_by_pair: dict[tuple[int, int], list[int]] = {}
+    for expert_id in sorted(recv_ranks_by_expert):
+        send_ranks = old_ranks_by_expert.get(expert_id, [])
+        if not send_ranks:
+            continue
+        recv_ranks = recv_ranks_by_expert[expert_id]
+
+        num_per_sender, remainder = divmod(len(recv_ranks), len(send_ranks))
+        remainder_start = len(send_ranks) * num_per_sender
+        for sender_idx, src_rank in enumerate(send_ranks):
+            start = sender_idx * num_per_sender
+            for dst_rank in recv_ranks[start : start + num_per_sender]:
+                expert_ids_by_pair.setdefault((src_rank, dst_rank), []).append(
+                    expert_id
+                )
+            if sender_idx < remainder:
+                dst_rank = recv_ranks[remainder_start + sender_idx]
+                expert_ids_by_pair.setdefault((src_rank, dst_rank), []).append(
+                    expert_id
+                )
+
+    batches: list[list[MigrationFlow]] = []
+    endpoints_used: list[int] = []
+    for (src_rank, dst_rank), expert_ids in expert_ids_by_pair.items():
+        flow = MigrationFlow(src_rank, dst_rank, tuple(expert_ids))
+        endpoints = (1 << src_rank) | (1 << dst_rank)
+        for batch_idx, (batch, used) in enumerate(zip(batches, endpoints_used)):
+            if not endpoints & used:
+                batch.append(flow)
+                endpoints_used[batch_idx] |= endpoints
+                break
+        else:
+            batches.append([flow])
+            endpoints_used.append(endpoints)
+
+    return batches
+
+
+__all__ = [
+    "MigrationFlow",
+    "schedule_migration_batches",
+]
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index cbd513294565..14cadd6d7998 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -15,6 +15,10 @@
 
 from vllm.distributed.eplb.eplb_communicator import EplbCommunicator
 from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
+from vllm.distributed.eplb.migration_scheduler import (
+    MigrationFlow,
+    schedule_migration_batches,
+)
 from vllm.logger import init_logger
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 
@@ -170,6 +174,103 @@ def get_ep_ranks_with_experts_batch(
     return ranks_to_send_map, ranks_to_recv_map
 
 
+def _execute_migration_batch(
+    batch: list[MigrationFlow],
+    old_indices: np.ndarray,
+    old_rows: dict[int, int],
+    new_rows: dict[int, int],
+    ep_rank: int,
+    expert_weights: Sequence[torch.Tensor],
+    expert_weights_buffers: Sequence[torch.Tensor],
+    communicator: EplbCommunicator,
+    layer_idx: int,
+) -> None:
+    """Execute one batch of rank-pair-disjoint expert transfers."""
+    communicator.set_transfer_context(old_indices, layer_idx)
+    for flow in batch:
+        if flow.src_rank == ep_rank:
+            for expert_id in flow.expert_ids:
+                src_row = _get_migration_row(
+                    old_rows, expert_id, flow, ep_rank, mapping_name="old"
+                )
+                communicator.add_send(
+                    [weight[src_row] for weight in expert_weights],
+                    flow.dst_rank,
+                    expert_id=expert_id,
+                )
+        elif flow.dst_rank == ep_rank:
+            for expert_id in flow.expert_ids:
+                dst_row = _get_migration_row(
+                    new_rows, expert_id, flow, ep_rank, mapping_name="new"
+                )
+                communicator.add_recv(
+                    [buffer[dst_row] for buffer in expert_weights_buffers],
+                    flow.src_rank,
+                    expert_id=expert_id,
+                )
+    communicator.execute()
+
+
+def _execute_migration_batches(
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+    ep_rank: int,
+    expert_weights: Sequence[torch.Tensor],
+    expert_weights_buffers: Sequence[torch.Tensor],
+    communicator: EplbCommunicator,
+    layer_idx: int,
+) -> None:
+    """Execute all contention-aware migration batches for one layer."""
+    with torch.profiler.record_function("eplb: schedule migration batches"):
+        batches = schedule_migration_batches(
+            num_local_experts, old_indices, new_indices
+        )
+    base = ep_rank * num_local_experts
+    old_local = old_indices[base : base + num_local_experts]
+    new_local = new_indices[base : base + num_local_experts]
+    old_rows: dict[int, int] = {}
+    new_rows: dict[int, int] = {}
+    for row, expert_id in enumerate(old_local):
+        if expert_id != -1:
+            old_rows.setdefault(int(expert_id), row)
+    for row, expert_id in enumerate(new_local):
+        if expert_id != -1:
+            new_rows.setdefault(int(expert_id), row)
+    for batch in batches:
+        _execute_migration_batch(
+            batch=batch,
+            old_indices=old_indices,
+            old_rows=old_rows,
+            new_rows=new_rows,
+            ep_rank=ep_rank,
+            expert_weights=expert_weights,
+            expert_weights_buffers=expert_weights_buffers,
+            communicator=communicator,
+            layer_idx=layer_idx,
+        )
+
+
+def _get_migration_row(
+    rows: dict[int, int],
+    expert_id: int,
+    flow: MigrationFlow,
+    ep_rank: int,
+    mapping_name: str,
+) -> int:
+    row = rows.get(expert_id)
+    if row is not None:
+        return row
+
+    message = (
+        f"Expert {expert_id} is missing from the {mapping_name} "
+        f"mapping on rank {ep_rank} for migration "
+        f"{flow.src_rank} -> {flow.dst_rank}"
+    )
+    logger.error(message)
+    raise RuntimeError(message)
+
+
 def move_to_buffer(
     num_local_experts: int,
     old_indices: np.ndarray,
@@ -180,6 +281,7 @@ def move_to_buffer(
     ep_rank: int,
     communicator: EplbCommunicator,
     layer_idx: int = 0,
+    enable_migration_batching: bool = False,
 ) -> TransferMetadata:
     """
     Rearranges expert weights during EPLB rebalancing.
@@ -196,6 +298,8 @@ def move_to_buffer(
         ep_rank: Rank of this process in expert parallel group.
         communicator: EplbCommunicator instance for P2P communication.
         layer_idx: Index of the MoE layer being transferred.
+        enable_migration_batching: Schedule remote transfers in batches where
+            each rank communicates with at most one peer.
 
     Returns:
         TransferMetadata: Metadata needed for completing remote weight transfers.
@@ -268,6 +372,28 @@ def move_to_buffer(
                     for w, b in zip(expert_weights, expert_weights_buffers):
                         b[dst].copy_(w[src_local], non_blocking=True)
 
+    transfer_metadata = TransferMetadata(
+        is_unchanged=is_unchanged,
+        is_received_locally=is_received_locally,
+        recv_primary_mask=recv_primary_mask,
+        recv_count=recv_count,
+        recv_expert_ids=recv_expert_ids,
+        recv_dst_rows=recv_dst_rows,
+    )
+
+    if enable_migration_batching:
+        _execute_migration_batches(
+            num_local_experts=num_local_experts,
+            old_indices=old_indices,
+            new_indices=new_indices,
+            ep_rank=ep_rank,
+            expert_weights=expert_weights,
+            expert_weights_buffers=expert_weights_buffers,
+            communicator=communicator,
+            layer_idx=layer_idx,
+        )
+        return transfer_metadata
+
     communicator.set_transfer_context(old_indices, layer_idx)
 
     # 2. Post sends
@@ -338,14 +464,7 @@ def move_to_buffer(
 
     # 4. Execute transfers and wait for completion.
     communicator.execute()
-    return TransferMetadata(
-        is_unchanged=is_unchanged,
-        is_received_locally=is_received_locally,
-        recv_primary_mask=recv_primary_mask,
-        recv_count=recv_count,
-        recv_expert_ids=recv_expert_ids,
-        recv_dst_rows=recv_dst_rows,
-    )
+    return transfer_metadata
 
 
 def move_from_buffer(
@@ -436,6 +555,7 @@ def transfer_layer(
     cuda_stream: torch.cuda.Stream | None = None,
     rank_mapping: dict[int, int] | None = None,
     layer_idx: int = 0,
+    enable_migration_batching: bool = False,
 ) -> TransferMetadata:
     """
     Rearranges the expert weights in place according to the new expert indices.
@@ -458,6 +578,8 @@ def transfer_layer(
         cuda_stream: CUDA stream for async copies (can be None for sync mode).
         rank_mapping: Optional rank mapping for elastic expert parallelism.
         layer_idx: Index of the MoE layer being transferred.
+        enable_migration_batching: Schedule remote transfers in batches where
+            each rank communicates with at most one peer.
 
     Returns:
         TransferMetadata: Metadata needed for completing remote weight transfers,
@@ -506,6 +628,7 @@ def transfer_layer(
         ep_rank=ep_group.rank(),
         communicator=communicator,
         layer_idx=layer_idx,
+        enable_migration_batching=enable_migration_batching,
     )
 
 
@@ -518,6 +641,7 @@ def rearrange_expert_weights_inplace(
     communicator: EplbCommunicator,
     is_profile: bool = False,
     rank_mapping: dict[int, int] | None = None,
+    enable_migration_batching: bool = False,
 ) -> None:
     """
     Rearranges the expert weights in place according to the new expert indices.
@@ -540,6 +664,8 @@ def rearrange_expert_weights_inplace(
             This is used during profile run, where we only perform dummy
             communications to reserve enough memory for the buffers.
         rank_mapping: A dictionary mapping old rank to new rank.
+        enable_migration_batching: Schedule remote transfers in batches where
+            each rank communicates with at most one peer.
     """
     if rank_mapping is not None:
         if len(rank_mapping) == ep_group.size():
@@ -608,6 +734,7 @@ def rearrange_expert_weights_inplace(
             ep_rank=ep_rank,
             communicator=communicator,
             layer_idx=layer_idx,
+            enable_migration_batching=enable_migration_batching,
         )
 
         move_from_buffer(

From 35cdfc1c25c92071c99f1c7c9780647bb1e00854 Mon Sep 17 00:00:00 2001
From: DOCCA0 <2947323341@qq.com>
Date: Wed, 9 Sep 2026 21:32:12 +0800
Subject: [PATCH 2/6] Explain EPLB migration batching invariants

Signed-off-by: DOCCA0 <2947323341@qq.com>
---
 vllm/distributed/eplb/migration_scheduler.py | 15 +++++++++++++++
 vllm/distributed/eplb/rebalance_execute.py   |  2 ++
 2 files changed, 17 insertions(+)

diff --git a/vllm/distributed/eplb/migration_scheduler.py b/vllm/distributed/eplb/migration_scheduler.py
index 72e7651f9669..aaa2ef29edea 100644
--- a/vllm/distributed/eplb/migration_scheduler.py
+++ b/vllm/distributed/eplb/migration_scheduler.py
@@ -27,6 +27,9 @@ def schedule_migration_batches(
     old_experts_by_rank: list[set[int]] = []
     old_by_rank = old_indices.reshape(-1, num_local_experts)
     new_by_rank = new_indices.reshape(-1, num_local_experts)
+
+    # First find the ranks that gain each expert. Existing local copies need no
+    # transfer, even if their physical row changes.
     for rank, (old_local, new_local) in enumerate(zip(old_by_rank, new_by_rank)):
         old_experts = set(old_local.tolist())
         old_experts.discard(-1)
@@ -39,6 +42,8 @@ def schedule_migration_batches(
 
     old_ranks_by_expert: dict[int, list[int]] = {}
     needed_experts = recv_ranks_by_expert.keys()
+
+    # Only index source copies for experts that need a remote transfer.
     for rank, old_experts in enumerate(old_experts_by_rank):
         for expert_id in old_experts & needed_experts:
             old_ranks_by_expert.setdefault(expert_id, []).append(rank)
@@ -52,6 +57,9 @@ def schedule_migration_batches(
 
         num_per_sender, remainder = divmod(len(recv_ranks), len(send_ranks))
         remainder_start = len(send_ranks) * num_per_sender
+
+        # Preserve the unbatched path's balanced source assignment. Batching
+        # changes transfer order, not which existing replica supplies a copy.
         for sender_idx, src_rank in enumerate(send_ranks):
             start = sender_idx * num_per_sender
             for dst_rank in recv_ranks[start : start + num_per_sender]:
@@ -64,10 +72,17 @@ def schedule_migration_batches(
                     expert_id
                 )
 
+    # Transfers with the same directed rank pair have identical contention
+    # constraints, so coalesce them before grouping. The loop below is a
+    # deterministic first-fit edge coloring of the rank communication graph.
     batches: list[list[MigrationFlow]] = []
     endpoints_used: list[int] = []
     for (src_rank, dst_rank), expert_ids in expert_ids_by_pair.items():
         flow = MigrationFlow(src_rank, dst_rank, tuple(expert_ids))
+
+        # Direction is irrelevant for NIC contention. This endpoint mask keeps
+        # each rank connected to at most one peer in a batch. The conservative
+        # limit leaves bandwidth headroom for inference during async migration.
         endpoints = (1 << src_rank) | (1 << dst_rank)
         for batch_idx, (batch, used) in enumerate(zip(batches, endpoints_used)):
             if not endpoints & used:
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index 14cadd6d7998..b83c9de3484d 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -186,6 +186,7 @@ def _execute_migration_batch(
     layer_idx: int,
 ) -> None:
     """Execute one batch of rank-pair-disjoint expert transfers."""
+    # NIXL drains its transfer context in execute(), so each batch restores it.
     communicator.set_transfer_context(old_indices, layer_idx)
     for flow in batch:
         if flow.src_rank == ep_rank:
@@ -208,6 +209,7 @@ def _execute_migration_batch(
                     flow.src_rank,
                     expert_id=expert_id,
                 )
+    # Combining execute calls would erase the scheduler's batch boundaries.
     communicator.execute()
 
 

From b5ad4d7c192b5fb3e624c26e56967dec08e74251 Mon Sep 17 00:00:00 2001
From: DOCCA0 <2947323341@qq.com>
Date: Sat, 12 Sep 2026 20:06:45 +0800
Subject: [PATCH 3/6] Precompute EPLB migration schedules per rebalance cycle

Signed-off-by: DOCCA0 <2947323341@qq.com>
---
 .buildkite/test_areas/expert_parallelism.yaml |  2 +
 .../test_eplb_migration_scheduler.py          | 39 +++++++++++++++++--
 vllm/config/parallel.py                       | 11 ++++--
 vllm/distributed/eplb/async_worker.py         | 25 +++++++++++-
 vllm/distributed/eplb/migration_scheduler.py  | 29 +++++++++++++-
 vllm/distributed/eplb/rebalance_execute.py    | 17 ++++++--
 6 files changed, 110 insertions(+), 13 deletions(-)

diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml
index f74b1adf2e2a..69ce3197cc1d 100644
--- a/.buildkite/test_areas/expert_parallelism.yaml
+++ b/.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:
diff --git a/tests/distributed/test_eplb_migration_scheduler.py b/tests/distributed/test_eplb_migration_scheduler.py
index 0c1753864848..6746de6603b3 100644
--- a/tests/distributed/test_eplb_migration_scheduler.py
+++ b/tests/distributed/test_eplb_migration_scheduler.py
@@ -10,6 +10,7 @@
 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
 
@@ -38,9 +39,22 @@ def execute(self) -> None:
         self.execute_count += 1
 
 
-def test_migration_batching_is_enabled_by_default_and_can_be_disabled() -> None:
-    assert EPLBConfig().enable_migration_batching
-    assert not EPLBConfig(enable_migration_batching=False).enable_migration_batching
+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:
@@ -139,10 +153,26 @@ def test_schedule_migration_batches_balances_replicas() -> None:
     }
 
 
-def test_move_to_buffer_uses_multiple_batches() -> None:
+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,
@@ -155,6 +185,7 @@ def test_move_to_buffer_uses_multiple_batches() -> None:
         communicator=communicator,
         layer_idx=7,
         enable_migration_batching=True,
+        migration_batches=migration_batches,
     )
 
     assert communicator.context_calls == 2
diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py
index 756a43f29754..2dbaefc921b1 100644
--- a/vllm/config/parallel.py
+++ b/vllm/config/parallel.py
@@ -99,11 +99,16 @@ class EPLBConfig:
     - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")
     """
 
-    enable_migration_batching: bool = True
+    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. Set to False to use one
-    communication step for all migrations."""
+    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:
diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py
index 5f77041b823a..504ea4c4366d 100644
--- a/vllm/distributed/eplb/async_worker.py
+++ b/vllm/distributed/eplb/async_worker.py
@@ -13,6 +13,7 @@
 from vllm.logger import init_logger
 
 from .eplb_utils import CpuGpuEvent
+from .migration_scheduler import schedule_migration_batches_for_layers
 from .rebalance_execute import AsyncEplbLayerResult, transfer_layer
 
 if TYPE_CHECKING:
@@ -101,6 +102,21 @@ def transfer_run_periodically(
                 model_state, state, physical_to_logical_map_cpu, cuda_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
@@ -137,7 +153,14 @@ def transfer_run_periodically(
                     cuda_stream=cuda_stream,
                     layer_idx=layer_idx,
                     enable_migration_batching=(
-                        state.parallel_config.eplb_config.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
                     ),
                 )
 
diff --git a/vllm/distributed/eplb/migration_scheduler.py b/vllm/distributed/eplb/migration_scheduler.py
index aaa2ef29edea..1844d79b0b03 100644
--- a/vllm/distributed/eplb/migration_scheduler.py
+++ b/vllm/distributed/eplb/migration_scheduler.py
@@ -13,6 +13,7 @@ class MigrationFlow:
 
     src_rank: int
     dst_rank: int
+    # Experts being transferred.
     expert_ids: tuple[int, ...]
 
 
@@ -21,7 +22,13 @@ def schedule_migration_batches(
     old_indices: np.ndarray,
     new_indices: np.ndarray,
 ) -> list[list[MigrationFlow]]:
-    """Build and greedily group rank-pair-disjoint migration flows."""
+    """Schedule expert migrations without per-rank communication contention.
+
+    Experts assigned to the same source and destination ranks form one flow.
+    Each flow is placed in the first batch that does not already use either
+    endpoint, so a rank communicates with at most one peer in each batch while
+    independent rank pairs can transfer concurrently.
+    """
     assert old_indices.shape == new_indices.shape
     recv_ranks_by_expert: dict[int, list[int]] = {}
     old_experts_by_rank: list[set[int]] = []
@@ -96,7 +103,27 @@ def schedule_migration_batches(
     return batches
 
 
+def schedule_migration_batches_for_layers(
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+) -> list[list[list[MigrationFlow]]]:
+    """Precompute the independent migration schedule for every MoE layer.
+
+    This reduces scheduling entry points from once per layer to once per
+    rebalance cycle. Only CPU scheduling is grouped; transfers remain ordered
+    by layer and preserve each layer's batch boundaries.
+    """
+    assert old_indices.shape == new_indices.shape
+    assert old_indices.ndim == 2
+    return [
+        schedule_migration_batches(num_local_experts, old_layer, new_layer)
+        for old_layer, new_layer in zip(old_indices, new_indices)
+    ]
+
+
 __all__ = [
     "MigrationFlow",
     "schedule_migration_batches",
+    "schedule_migration_batches_for_layers",
 ]
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index b83c9de3484d..2fe896445700 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -222,12 +222,14 @@ def _execute_migration_batches(
     expert_weights_buffers: Sequence[torch.Tensor],
     communicator: EplbCommunicator,
     layer_idx: int,
+    batches: list[list[MigrationFlow]] | None = None,
 ) -> None:
     """Execute all contention-aware migration batches for one layer."""
-    with torch.profiler.record_function("eplb: schedule migration batches"):
-        batches = schedule_migration_batches(
-            num_local_experts, old_indices, new_indices
-        )
+    if batches is None:
+        with torch.profiler.record_function("eplb: schedule migration batches"):
+            batches = schedule_migration_batches(
+                num_local_experts, old_indices, new_indices
+            )
     base = ep_rank * num_local_experts
     old_local = old_indices[base : base + num_local_experts]
     new_local = new_indices[base : base + num_local_experts]
@@ -284,6 +286,7 @@ def move_to_buffer(
     communicator: EplbCommunicator,
     layer_idx: int = 0,
     enable_migration_batching: bool = False,
+    migration_batches: list[list[MigrationFlow]] | None = None,
 ) -> TransferMetadata:
     """
     Rearranges expert weights during EPLB rebalancing.
@@ -302,6 +305,7 @@ def move_to_buffer(
         layer_idx: Index of the MoE layer being transferred.
         enable_migration_batching: Schedule remote transfers in batches where
             each rank communicates with at most one peer.
+        migration_batches: Optional batches precomputed for this layer.
 
     Returns:
         TransferMetadata: Metadata needed for completing remote weight transfers.
@@ -393,9 +397,11 @@ def move_to_buffer(
             expert_weights_buffers=expert_weights_buffers,
             communicator=communicator,
             layer_idx=layer_idx,
+            batches=migration_batches,
         )
         return transfer_metadata
 
+    # By default, submit all expert transfers in one communication step.
     communicator.set_transfer_context(old_indices, layer_idx)
 
     # 2. Post sends
@@ -558,6 +564,7 @@ def transfer_layer(
     rank_mapping: dict[int, int] | None = None,
     layer_idx: int = 0,
     enable_migration_batching: bool = False,
+    migration_batches: list[list[MigrationFlow]] | None = None,
 ) -> TransferMetadata:
     """
     Rearranges the expert weights in place according to the new expert indices.
@@ -582,6 +589,7 @@ def transfer_layer(
         layer_idx: Index of the MoE layer being transferred.
         enable_migration_batching: Schedule remote transfers in batches where
             each rank communicates with at most one peer.
+        migration_batches: Optional batches precomputed for this layer.
 
     Returns:
         TransferMetadata: Metadata needed for completing remote weight transfers,
@@ -631,6 +639,7 @@ def transfer_layer(
         communicator=communicator,
         layer_idx=layer_idx,
         enable_migration_batching=enable_migration_batching,
+        migration_batches=migration_batches,
     )
 
 

From 009a5c0228250337118b5df4ad57fae98c5a946c Mon Sep 17 00:00:00 2001
From: DOCCA0 <2947323341@qq.com>
Date: Sat, 12 Sep 2026 20:59:27 +0800
Subject: [PATCH 4/6] Clarify default EPLB migration path

Signed-off-by: DOCCA0 <2947323341@qq.com>
---
 vllm/distributed/eplb/rebalance_execute.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index 2fe896445700..39e6acd52c6d 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -387,6 +387,7 @@ def move_to_buffer(
         recv_dst_rows=recv_dst_rows,
     )
 
+    # Batching is opt-in; by default, submit all transfers in one step.
     if enable_migration_batching:
         _execute_migration_batches(
             num_local_experts=num_local_experts,
@@ -401,7 +402,6 @@ def move_to_buffer(
         )
         return transfer_metadata
 
-    # By default, submit all expert transfers in one communication step.
     communicator.set_transfer_context(old_indices, layer_idx)
 
     # 2. Post sends

From b49207d13c31a984ee0232f5a034887c497bb853 Mon Sep 17 00:00:00 2001
From: DOCCA0 <2947323341@qq.com>
Date: Sat, 12 Sep 2026 21:01:49 +0800
Subject: [PATCH 5/6] State EPLB batching default explicitly

Signed-off-by: DOCCA0 <2947323341@qq.com>
---
 vllm/distributed/eplb/rebalance_execute.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index 39e6acd52c6d..8f6ceadfea98 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -387,7 +387,7 @@ def move_to_buffer(
         recv_dst_rows=recv_dst_rows,
     )
 
-    # Batching is opt-in; by default, submit all transfers in one step.
+    # Migration batching is disabled by default and must be explicitly enabled.
     if enable_migration_batching:
         _execute_migration_batches(
             num_local_experts=num_local_experts,

From 940a0c81225a8933149febfc9774400c36f22b8b Mon Sep 17 00:00:00 2001
From: DOCCA0 <2947323341@qq.com>
Date: Fri, 25 Sep 2026 08:30:35 +0800
Subject: [PATCH 6/6] Keep migration batching off the sync EPLB path

Signed-off-by: DOCCA0 <2947323341@qq.com>
---
 vllm/distributed/eplb/eplb_state.py        | 3 ---
 vllm/distributed/eplb/rebalance_execute.py | 4 ----
 2 files changed, 7 deletions(-)

diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py
index 8e42203c9dfa..e4d805b9fead 100644
--- a/vllm/distributed/eplb/eplb_state.py
+++ b/vllm/distributed/eplb/eplb_state.py
@@ -915,9 +915,6 @@ def rank_load_imbalance(
                         eplb_model_state.communicator,
                         is_profile,
                         rank_mapping,
-                        enable_migration_batching=(
-                            self.parallel_config.eplb_config.enable_migration_batching
-                        ),
                     )
 
                     if not is_profile:
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index 54b8f87cd158..5479ef01df4e 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -648,7 +648,6 @@ def rearrange_expert_weights_inplace(
     communicator: EplbCommunicator,
     is_profile: bool = False,
     rank_mapping: dict[int, int] | None = None,
-    enable_migration_batching: bool = False,
 ) -> None:
     """Rearranges the expert weights in place according to the new expert indices.
 
@@ -670,8 +669,6 @@ def rearrange_expert_weights_inplace(
             This is used during profile run, where we only perform dummy
             communications to reserve enough memory for the buffers.
         rank_mapping: A dictionary mapping old rank to new rank.
-        enable_migration_batching: Schedule remote transfers in batches where
-            each rank communicates with at most one peer.
 
     """
     if rank_mapping is not None:
@@ -741,7 +738,6 @@ def rearrange_expert_weights_inplace(
             ep_rank=ep_rank,
             communicator=communicator,
             layer_idx=layer_idx,
-            enable_migration_batching=enable_migration_batching,
         )
 
         move_from_buffer(