vllm-cpu-avx512bf16 0.14.0__cp313-cp313-manylinux_2_28_x86_64.whl

vllm/distributed/eplb/rebalance_execute.py

@@ -0,0 +1,699 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+The actual execution of the rearrangement.
+
+This involves the exchange of expert weights between GPUs.
+"""
+
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch.distributed import (
+    P2POp,
+    ProcessGroup,
+    all_gather,
+    batch_isend_irecv,
+    get_global_rank,
+)
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class RecvMetadata:
+    """Metadata describing remote receives during EPLB rebalancing."""
+
+    recv_primary_mask: np.ndarray
+    """Mask of (num_local_experts,) indicating primary experts received."""
+    recv_count: int
+    """Number of received experts for the layer."""
+    recv_expert_ids: np.ndarray
+    """Expert ids (num_local_experts,) of remote primary experts."""
+    recv_dst_rows: np.ndarray
+    """Target expert indices (num_local_experts,) in local tensors to send."""
+
+
+# Type alias for the result of move_to_buffer or transfer_layer
+MoveToBufferResult = tuple[np.ndarray, np.ndarray, RecvMetadata]
+
+
+def get_ep_ranks_with_experts_batch(
+    expert_ids: np.ndarray,
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+) -> tuple[dict[int, list[int]], dict[int, list[int]]]:
+    """
+    Get the ranks of the experts that need to be exchanged.
+
+    Args:
+        expert_ids: 1D array of expert indices to query.
+        num_local_experts: The number of local experts.
+        old_indices: The old indices of the experts.
+        new_indices: The new indices of the experts.
+
+    Returns:
+        A tuple of two dictionaries mapping expert_id to:
+        - ranks_to_send: The ranks that have this expert and need to send.
+        - ranks_to_recv: The ranks that need to receive this expert.
+    """
+    ranks_to_send_map: dict[int, list[int]] = {}
+    ranks_to_recv_map: dict[int, list[int]] = {}
+
+    # Fast path: if no experts, return empty dicts
+    if expert_ids.size == 0:
+        return ranks_to_send_map, ranks_to_recv_map
+
+    unique_experts = np.unique(expert_ids)
+    num_positions = len(old_indices)
+    position_indices = np.arange(num_positions, dtype=np.int32)
+
+    # Vectorized approach: find all positions matching any query expert in one pass
+    # Use np.isin to get boolean masks for all relevant positions at once
+    old_relevant_mask = np.isin(old_indices, unique_experts)
+    new_relevant_mask = np.isin(new_indices, unique_experts)
+
+    # Process old_indices (send ranks)
+    if np.any(old_relevant_mask):
+        old_relevant_positions = position_indices[old_relevant_mask]
+        old_relevant_experts = old_indices[old_relevant_mask]
+        old_relevant_ranks = old_relevant_positions // num_local_experts
+
+        # Sort by expert first, then by position (to maintain first-appearance order)
+        sort_order = np.lexsort((old_relevant_positions, old_relevant_experts))
+        sorted_experts = old_relevant_experts[sort_order]
+        sorted_ranks = old_relevant_ranks[sort_order]
+
+        # Find boundaries where expert changes
+        expert_boundaries = np.concatenate(
+            [[0], np.where(np.diff(sorted_experts) != 0)[0] + 1, [len(sorted_experts)]]
+        )
+
+        # For each expert, extract unique ranks in order of first appearance
+        for i in range(len(expert_boundaries) - 1):
+            start, end = expert_boundaries[i], expert_boundaries[i + 1]
+            expert = int(sorted_experts[start])
+            expert_ranks = sorted_ranks[start:end]
+
+            # Get unique ranks preserving order
+            _, unique_idx = np.unique(expert_ranks, return_index=True)
+            unique_ranks = expert_ranks[np.sort(unique_idx)]
+            ranks_to_send_map[expert] = unique_ranks.tolist()
+
+    # Process new_indices (recv ranks)
+    if np.any(new_relevant_mask):
+        new_relevant_positions = position_indices[new_relevant_mask]
+        new_relevant_experts = new_indices[new_relevant_mask]
+        new_relevant_ranks = new_relevant_positions // num_local_experts
+
+        # Sort by expert first, then by position
+        sort_order = np.lexsort((new_relevant_positions, new_relevant_experts))
+        sorted_experts = new_relevant_experts[sort_order]
+        sorted_ranks = new_relevant_ranks[sort_order]
+
+        # Find boundaries where expert changes
+        expert_boundaries = np.concatenate(
+            [[0], np.where(np.diff(sorted_experts) != 0)[0] + 1, [len(sorted_experts)]]
+        )
+
+        # For each expert, extract unique ranks and exclude local copies
+        for i in range(len(expert_boundaries) - 1):
+            start, end = expert_boundaries[i], expert_boundaries[i + 1]
+            expert = int(sorted_experts[start])
+            expert_ranks = sorted_ranks[start:end]
+
+            # Get unique ranks preserving order
+            _, unique_idx = np.unique(expert_ranks, return_index=True)
+            unique_ranks = expert_ranks[np.sort(unique_idx)]
+
+            # Remove ranks that have local copies (in send map)
+            send_ranks_set = set(ranks_to_send_map.get(expert, []))
+            recv_ranks_actual = [
+                int(r) for r in unique_ranks if r not in send_ranks_set
+            ]
+            ranks_to_recv_map[expert] = recv_ranks_actual
+
+    # Handle experts that only appear in old (send only) or new (recv only)
+    for expert in unique_experts:
+        expert = int(expert)
+        if expert not in ranks_to_send_map:
+            ranks_to_send_map[expert] = []
+        if expert not in ranks_to_recv_map:
+            ranks_to_recv_map[expert] = []
+
+    return ranks_to_send_map, ranks_to_recv_map
+
+
+def move_to_buffer(
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+    expert_weights: Iterable[torch.Tensor],
+    expert_weights_buffers: Sequence[torch.Tensor],
+    cuda_stream: torch.cuda.Stream | None,
+    ep_group: ProcessGroup,
+) -> MoveToBufferResult:
+    """
+    Rearranges expert weights during EPLB rebalancing.
+
+    Args:
+        num_local_experts: Number of local experts.
+        old_indices: (num_experts_total,) ndarray of current (old)
+            global-to-local expert assignments.
+        new_indices: (num_experts_total,) ndarray of desired (new)
+            global-to-local assignments after rebalance.
+        expert_weights: Original expert weights for the layer.
+        expert_weights_buffers: Intermediate buffers (one per tensor).
+        cuda_stream: CUDA stream for async copies (can be None for sync mode).
+        ep_group: Distributed process group for expert parallel comms.
+
+    Returns:
+        is_unchanged (np.ndarray): (num_local_experts,), True where an expert row
+            is unchanged after rebalance.
+        is_received_locally (np.ndarray): (num_local_experts,), True where a row
+            can be updated from local data.
+        RecvMetadata: Metadata needed for completing remote weight transfers.
+    """
+    assert old_indices.shape == new_indices.shape
+    ep_rank = ep_group.rank()
+
+    recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_)
+    send_expert_ids = np.full((num_local_experts,), -1, dtype=np.int64)
+    send_src_rows = np.full((num_local_experts,), -1, dtype=np.int32)
+    recv_expert_ids = np.full((num_local_experts,), -1, dtype=np.int64)
+    recv_dst_rows = np.full((num_local_experts,), -1, dtype=np.int32)
+
+    base = ep_rank * num_local_experts
+    local_rows = np.arange(num_local_experts, dtype=np.int32)
+    local_global = base + local_rows
+
+    old_local_expert_ids = old_indices[local_global]
+    new_local_expert_ids = new_indices[local_global]
+
+    # Unchanged mask
+    is_unchanged = old_local_expert_ids == new_local_expert_ids
+
+    # Local receive eligibility
+    new_valid = new_local_expert_ids != -1
+    can_recv_local = np.isin(
+        new_local_expert_ids, old_local_expert_ids, assume_unique=False
+    )
+    is_received_locally = np.logical_or(
+        is_unchanged, np.logical_and(new_valid, can_recv_local)
+    )
+
+    # Send map: first src row per unique expert present locally in old mapping
+    send_count = 0
+    valid_old = old_local_expert_ids != -1
+    if np.any(valid_old):
+        uniq_experts, first_idx = np.unique(
+            old_local_expert_ids[valid_old], return_index=True
+        )
+        filtered_rows = local_rows[valid_old]
+        src_rows = filtered_rows[first_idx]
+        send_count = int(uniq_experts.shape[0])
+        send_expert_ids[:send_count] = uniq_experts
+        send_src_rows[:send_count] = src_rows
+
+    # Recv map: primary dst per unique expert needed remotely
+    recv_count = 0
+    need_recv_mask = np.logical_and(~is_received_locally, new_valid)
+    if np.any(need_recv_mask):
+        desired_experts = new_local_expert_ids[need_recv_mask]
+        desired_dsts = local_rows[need_recv_mask]
+        uniq_recv_experts, uniq_indices = np.unique(desired_experts, return_index=True)
+        dst_rows = desired_dsts[uniq_indices]
+        recv_count = int(uniq_recv_experts.shape[0])
+        recv_expert_ids[:recv_count] = uniq_recv_experts
+        recv_dst_rows[:recv_count] = dst_rows
+        recv_primary_mask[dst_rows] = True
+
+    eligible_local_buffer_mask = np.logical_and(~is_unchanged, is_received_locally)
+
+    # 1. Local moves into tmp buffers
+    if bool(eligible_local_buffer_mask.any()) and send_count > 0:
+        dest_indices = np.nonzero(eligible_local_buffer_mask)[0].tolist()
+        expert_to_src_map = dict(
+            zip(send_expert_ids[:send_count], send_src_rows[:send_count])
+        )
+        for dst in dest_indices:
+            expert = new_local_expert_ids[dst]
+            src_local = expert_to_src_map.get(expert, -1)
+            if src_local != -1:
+                for w, b in zip(expert_weights, expert_weights_buffers):
+                    b[dst].copy_(w[src_local], non_blocking=True)
+
+    p2p_ops: list[P2POp] = []
+
+    # Pre-compute global ranks mapping
+    ep_size = ep_group.size()
+    rank_to_global = {rank: get_global_rank(ep_group, rank) for rank in range(ep_size)}
+
+    # 2. Post sends
+    if send_count > 0:
+        experts = send_expert_ids[:send_count]
+        srcs = send_src_rows[:send_count]
+        order = np.argsort(experts, kind="stable")
+        experts = experts[order]
+        srcs = srcs[order]
+
+        send_map, recv_map = get_ep_ranks_with_experts_batch(
+            experts,
+            num_local_experts,
+            old_indices,
+            new_indices,
+        )
+
+        for expert, src in zip(experts.tolist(), srcs.tolist()):
+            ranks_to_send = send_map[expert]
+            ranks_to_recv = recv_map[expert]
+            if not ranks_to_send or not ranks_to_recv:
+                continue
+            num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)
+            sender_pos = ranks_to_send.index(ep_rank)
+            recv_begin = sender_pos * num_dst_per_sender
+            recv_end = recv_begin + num_dst_per_sender
+            recv_ranks = ranks_to_recv[recv_begin:recv_end]
+            remainder_start = len(ranks_to_send) * num_dst_per_sender
+            recver_pos = remainder_start + sender_pos
+            if recver_pos < len(ranks_to_recv):
+                recv_ranks.append(ranks_to_recv[recver_pos])
+            for dst in recv_ranks:
+                dst_global = rank_to_global[dst]
+                p2p_ops += [
+                    P2POp(
+                        torch.distributed.isend,
+                        w[src],
+                        dst_global,
+                    )
+                    for w in expert_weights
+                ]
+
+    # 3. Post recvs
+    if recv_count > 0:
+        experts = recv_expert_ids[:recv_count]
+        dsts = recv_dst_rows[:recv_count]
+        order = np.argsort(experts, kind="stable")
+        experts = experts[order]
+        dsts = dsts[order]
+
+        send_map, recv_map = get_ep_ranks_with_experts_batch(
+            experts,
+            num_local_experts,
+            old_indices,
+            new_indices,
+        )
+
+        for expert, dst in zip(experts.tolist(), dsts.tolist()):
+            ranks_to_send = send_map[expert]
+            ranks_to_recv = recv_map[expert]
+            if not ranks_to_send or not ranks_to_recv:
+                continue
+            num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)
+            recver_pos = ranks_to_recv.index(ep_rank)
+            remainder_start = len(ranks_to_send) * num_dst_per_sender
+            if recver_pos < remainder_start:
+                src = ranks_to_send[recver_pos // num_dst_per_sender]
+            else:
+                src = ranks_to_send[recver_pos - remainder_start]
+            src_global = rank_to_global[src]
+            p2p_ops += [
+                P2POp(
+                    torch.distributed.irecv,
+                    b[dst],
+                    src_global,
+                )
+                for b in expert_weights_buffers
+            ]
+
+    # 4. Execute the P2P operations. The real communication happens here.
+    if p2p_ops and cuda_stream is not None:
+        with torch.cuda.stream(cuda_stream):
+            reqs = batch_isend_irecv(p2p_ops)
+            for req in reqs:
+                req.wait()
+    elif p2p_ops:
+        reqs = batch_isend_irecv(p2p_ops)
+        for req in reqs:
+            req.wait()
+    # wait for the communication to finish
+    return (
+        is_unchanged,
+        is_received_locally,
+        RecvMetadata(
+            recv_primary_mask=recv_primary_mask,
+            recv_count=recv_count,
+            recv_expert_ids=recv_expert_ids,
+            recv_dst_rows=recv_dst_rows,
+        ),
+    )
+
+
+def move_from_buffer(
+    expert_weights: Iterable[torch.Tensor],
+    expert_weights_buffers: list[torch.Tensor],
+    is_unchanged: np.ndarray,
+    is_received_locally: np.ndarray,
+    recv_metadata: RecvMetadata,
+    new_indices: np.ndarray,
+    ep_rank: int,
+) -> None:
+    """
+    Copies expert weights from communication buffers back to the target weight tensors
+    after EPLB rebalancing.
+
+    Args:
+        expert_weights: List of the actual MoE layer weights used in the execution.
+        expert_weights_buffers: Intermediate buffers containing the experts weights
+            after the transfer is completed.
+        is_unchanged: (num_local_experts,), True where an expert row is unchanged.
+        is_received_locally: (num_local_experts,), True where a row is updated locally.
+        recv_metadata: RecvMetadata containing remote receive metadata.
+        new_indices: (num_experts_total,) mapping from local rows to desired
+            (possibly global) expert id, after rebalance.
+        ep_rank: Rank of the process in the expert parallel group.
+    """
+    recv_primary_mask = recv_metadata.recv_primary_mask
+    recv_count = recv_metadata.recv_count
+    recv_expert_ids = recv_metadata.recv_expert_ids
+    recv_dst_rows = recv_metadata.recv_dst_rows
+    num_local_experts = is_unchanged.shape[0]
+
+    # Mask for rows to copy back from buffers:
+    # copy if locally received OR remote primary recv
+    copy_mask = np.logical_or(is_received_locally, recv_primary_mask)
+    dest_mask_np = np.logical_and(~is_unchanged, copy_mask)
+    if bool(dest_mask_np.any()):
+        dest_indices = np.nonzero(dest_mask_np)[0].tolist()
+        for dst in dest_indices:
+            for w, b in zip(expert_weights, expert_weights_buffers):
+                w[dst].copy_(b[dst], non_blocking=True)
+
+    if recv_count == 0:
+        return
+
+    # Duplicate remote received rows to non-primary duplicate dsts
+    base = ep_rank * num_local_experts
+    local_experts = new_indices[base + np.arange(num_local_experts, dtype=np.int32)]
+    duplicate_mask = np.logical_and(
+        np.logical_and(~is_unchanged, ~is_received_locally),
+        np.logical_and(~recv_primary_mask, local_experts != -1),
+    )
+    # All received experts are unique in the destination, so no need to copy duplicates
+    if not bool(duplicate_mask.any()):
+        return
+
+    dup_dst_rows = np.nonzero(duplicate_mask)[0]
+    dup_experts = local_experts[dup_dst_rows]
+
+    prim_experts = recv_expert_ids[:recv_count]
+    prim_dsts = recv_dst_rows[:recv_count]
+    order = np.argsort(prim_experts, kind="stable")
+    prim_experts_sorted = prim_experts[order]
+    prim_dsts_sorted = prim_dsts[order]
+    pos = np.searchsorted(prim_experts_sorted, dup_experts)
+    valid = np.logical_and(
+        pos < prim_experts_sorted.shape[0],
+        prim_experts_sorted[np.minimum(pos, prim_experts_sorted.shape[0] - 1)]
+        == dup_experts,
+    )
+    if not bool(valid.any()):
+        return
+
+    matched_dst_rows = dup_dst_rows[valid]
+    matched_src_rows = prim_dsts_sorted[pos[valid]]
+
+    for dst, src in zip(matched_dst_rows.tolist(), matched_src_rows.tolist()):
+        for w in expert_weights:
+            w[dst].copy_(w[src], non_blocking=True)
+
+
+async def transfer_layer(
+    old_global_expert_indices: torch.Tensor,
+    new_global_expert_indices: torch.Tensor,
+    expert_weights: Sequence[Iterable[torch.Tensor]],
+    expert_weights_buffer: Sequence[torch.Tensor],
+    ep_group: ProcessGroup,
+    is_profile: bool = False,
+    layer: int = 0,
+    cuda_stream: torch.cuda.Stream | None = None,
+    rank_mapping: dict[int, int] | None = None,
+) -> MoveToBufferResult:
+    """
+    Rearranges the expert weights in place according to the new expert indices.
+
+    The value of the indices arguments are logical indices of the experts,
+    while keys are physical.
+
+    Args:
+        old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
+        new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
+        expert_weights: A sequence of shape (num_moe_layers)(weight_count)
+            of tensors of shape (num_local_physical_experts, hidden_size_i).
+            For example, a linear layer may have up and down projection,
+            so weight_count = 2. Each weight's hidden size can be different.
+        ep_group: The device process group for expert parallelism.
+        is_profile (bool): If `True`, do not perform any actual weight copy.
+            This is used during profile run, where we only perform dummy
+            communications to reserve enough memory for the buffers.
+
+    Returns:
+        is_unchanged (np.ndarray): (1, num_local_experts), True where expert
+            is left unchanged.
+        is_received_locally (np.ndarray): (1, num_local_experts), True where expert
+            can be received locally.
+        RecvMetadata: Metadata needed for completing remote weight transfers.
+    """
+    ep_size = ep_group.size()
+    if rank_mapping is not None:
+        if len(rank_mapping) == ep_group.size():
+            # scale down
+            new_global_expert_indices = _map_new_expert_indices_with_rank_mapping(
+                new_global_expert_indices,
+                rank_mapping,
+            )
+        else:
+            # scale up
+            old_global_expert_indices = _map_old_expert_indices_with_rank_mapping(
+                old_global_expert_indices,
+                rank_mapping,
+                ep_group.size(),
+            )
+
+    assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
+    num_moe_layers, num_physical_experts = old_global_expert_indices.shape
+    assert len(expert_weights) == num_moe_layers
+    num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
+    assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
+    assert num_physical_experts == ep_size * num_local_physical_experts
+
+    old_global_expert_indices_np = old_global_expert_indices.cpu().numpy()
+    new_global_expert_indices_np = new_global_expert_indices.cpu().numpy()
+
+    is_unchanged, is_received_locally, recv_metadata = move_to_buffer(
+        num_local_experts=num_local_physical_experts,
+        old_indices=old_global_expert_indices_np[layer],
+        new_indices=new_global_expert_indices_np[layer],
+        expert_weights=expert_weights[layer],
+        expert_weights_buffers=expert_weights_buffer,
+        cuda_stream=cuda_stream,
+        ep_group=ep_group,
+    )
+    return is_unchanged, is_received_locally, recv_metadata
+
+
+def rearrange_expert_weights_inplace(
+    old_global_expert_indices: torch.Tensor,
+    new_global_expert_indices: torch.Tensor,
+    expert_weights: Sequence[Iterable[torch.Tensor]],
+    ep_group: ProcessGroup,
+    is_profile: bool = False,
+    rank_mapping: dict[int, int] | None = None,
+) -> None:
+    """
+    Rearranges the expert weights in place according to the new expert indices.
+
+    The value of the indices arguments are logical indices of the experts,
+    while keys are physical.
+
+    Args:
+        old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
+        new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
+        expert_weights: A sequence of shape (num_moe_layers)(weight_count)
+            of tensors of shape (num_local_physical_experts, hidden_size_i).
+            For example, a linear layer may have up and down projection,
+            so weight_count = 2. Each weight's hidden size can be different.
+        ep_group: The device process group for expert parallelism.
+        is_profile (bool): If `True`, do not perform any actual weight copy.
+            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.
+    """
+    if rank_mapping is not None:
+        if len(rank_mapping) == ep_group.size():
+            # scale down
+            new_global_expert_indices = _map_new_expert_indices_with_rank_mapping(
+                new_global_expert_indices,
+                rank_mapping,
+            )
+        else:
+            # scale up
+            old_global_expert_indices = _map_old_expert_indices_with_rank_mapping(
+                old_global_expert_indices,
+                rank_mapping,
+                ep_group.size(),
+            )
+
+    assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
+
+    num_moe_layers, num_physical_experts = old_global_expert_indices.shape
+    assert len(expert_weights) == num_moe_layers
+
+    num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
+    assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
+
+    ep_size = ep_group.size()
+    assert num_physical_experts == ep_size * num_local_physical_experts
+
+    first_layer_weights = list(expert_weights[0])
+    # Buffers to hold the expert weights during the exchange.
+    # NOTE: Currently we assume the same weights across different layers
+    # have the same shape.
+    weights_buffer: list[torch.Tensor] = [
+        torch.empty_like(w) for w in first_layer_weights
+    ]
+    if is_profile:
+        # Reserve communication buffers via a minimal dummy all_gather on first layer
+        for weight, buffer in zip(expert_weights[0], weights_buffer):
+            dummy_recv_buffer = [buffer for _ in range(ep_size)]
+            torch.distributed.barrier()
+            all_gather(
+                dummy_recv_buffer,
+                weight,
+                group=ep_group,
+            )
+        return
+
+    # NOTE(bowen): We need this synchronize to run, but I don't know why.
+    # If you figure out the reason, please let me know -- thank you!
+    torch.cuda.synchronize()
+
+    old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()
+    new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()
+
+    for layer_idx in range(num_moe_layers):
+        is_unchanged, is_received_locally, recv_metadata = move_to_buffer(
+            num_local_experts=num_local_physical_experts,
+            old_indices=old_global_expert_indices_cpu[layer_idx],
+            new_indices=new_global_expert_indices_cpu[layer_idx],
+            expert_weights=expert_weights[layer_idx],
+            expert_weights_buffers=weights_buffer,
+            cuda_stream=None,
+            ep_group=ep_group,
+        )
+
+        move_from_buffer(
+            expert_weights=expert_weights[layer_idx],
+            expert_weights_buffers=weights_buffer,
+            is_unchanged=is_unchanged,
+            is_received_locally=is_received_locally,
+            recv_metadata=recv_metadata,
+            new_indices=new_global_expert_indices_cpu[layer_idx],
+            ep_rank=ep_group.rank(),
+        )
+
+
+def _map_old_expert_indices_with_rank_mapping(
+    old_global_expert_indices: torch.Tensor,
+    rank_mapping: dict[int, int],
+    new_ep_size: int,
+) -> torch.Tensor:
+    """
+    Map the old global expert indices to the new global expert indices.
+
+    Args:
+        old_global_expert_indices:
+            Shape (num_layers, old_ep_size * num_local_physical_experts).
+        rank_mapping: Mapping from old rank to new rank.
+        new_ep_size: New expert parallelism size.
+
+    Returns:
+        Mapped expert indices with shape
+        (num_layers, new_ep_size * num_local_physical_experts).
+    """
+    num_layers, old_num_physical_experts = old_global_expert_indices.shape
+    assert rank_mapping, "Rank mapping is required"
+
+    # Get sizes from parameters and rank_mapping
+    old_ep_size = len(rank_mapping)
+    num_local_physical_experts = old_num_physical_experts // old_ep_size
+    new_num_physical_experts = new_ep_size * num_local_physical_experts
+
+    # Create mapped tensor with new shape, initialized to -1
+    mapped_expert_indices = torch.full(
+        (num_layers, new_num_physical_experts),
+        fill_value=-1,
+        dtype=old_global_expert_indices.dtype,
+        device=old_global_expert_indices.device,
+    )
+
+    # Handle rank mapping (scale up/down with rank changes)
+    for old_rank in range(old_ep_size):
+        new_rank = rank_mapping.get(old_rank)
+        if new_rank is not None and new_rank >= 0 and new_rank < new_ep_size:
+            # This old rank exists in the new configuration
+            old_start_idx = old_rank * num_local_physical_experts
+            old_end_idx = (old_rank + 1) * num_local_physical_experts
+            new_start_idx = new_rank * num_local_physical_experts
+            new_end_idx = (new_rank + 1) * num_local_physical_experts
+
+            mapped_expert_indices[:, new_start_idx:new_end_idx] = (
+                old_global_expert_indices[:, old_start_idx:old_end_idx]
+            )
+        # If new_rank is None or >= new_ep_size, the experts remain -1
+        # (scale down case)
+
+    return mapped_expert_indices
+
+
+def _map_new_expert_indices_with_rank_mapping(
+    new_global_expert_indices: torch.Tensor,
+    rank_mapping: dict[int, int],
+) -> torch.Tensor:
+    num_layers, new_num_physical_experts = new_global_expert_indices.shape
+    assert rank_mapping, "Rank mapping is required"
+
+    # Get sizes from parameters and rank_mapping
+    old_ep_size = len(rank_mapping)
+    new_ep_size = sum(new_rank != -1 for new_rank in rank_mapping.values())
+    num_local_physical_experts = new_num_physical_experts // new_ep_size
+    old_num_physical_experts = old_ep_size * num_local_physical_experts
+
+    mapped_expert_indices = torch.full(
+        (num_layers, old_num_physical_experts),
+        fill_value=-1,
+        dtype=new_global_expert_indices.dtype,
+        device=new_global_expert_indices.device,
+    )
+
+    for old_rank in range(old_ep_size):
+        new_rank = rank_mapping[old_rank]
+        if new_rank >= 0 and new_rank < new_ep_size:
+            old_start_idx = old_rank * num_local_physical_experts
+            old_end_idx = (old_rank + 1) * num_local_physical_experts
+            new_start_idx = new_rank * num_local_physical_experts
+            new_end_idx = (new_rank + 1) * num_local_physical_experts
+
+            mapped_expert_indices[:, old_start_idx:old_end_idx] = (
+                new_global_expert_indices[:, new_start_idx:new_end_idx]
+            )
+
+    return mapped_expert_indices
+
+
+__all__ = ["transfer_layer", "move_from_buffer", "RecvMetadata"]