vllm.distributed.eplb.eplb_communicator

EPLB communicator implementations and factory.

Classes:

• EplbCommunicator –

  Abstract EPLB communicator for expert weight transfers.

• NixlEplbCommunicator –

  EPLB communicator backed by NIXL READ transfers.

• PyNcclEplbCommunicator –

  EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv.

• SymmMemEplbCommunicator –

  Push expert rows into peers' receive buffers over symmetric memory.

• TorchDistGlooStagedEplbCommunicator –

  EPLB communicator using gloo P2P with CPU staging.

• TorchDistNcclEplbCommunicator –

  EPLB communicator backed by torch.distributed isend/irecv.

• TorchDistXCCLStagedEplbCommunicator –

  EPLB communicator using XCCL device-to-device P2P on XPU.

Functions:

• create_eplb_communicator –

  Create an EPLB communicator for the given backend.

• has_nixl –

  Whether the optional NIXL package is available.

EplbCommunicator

Bases: ABC

Abstract EPLB communicator for expert weight transfers.

Methods:

• execute –

  Complete all enqueued transfers.

• set_transfer_context –

  Pre-set layer context before add_send/add_recv calls.

Attributes:

• needs_profile_buffer_reservation (bool) –

  Whether the profile path must run a dummy collective operation to reserve

• needs_recv_rows (bool) –

  Whether set_transfer_context must be given every rank's recv rows.

• owned_expert_buffer (list[Tensor] | None) –

  Backend-allocated receive buffer, or None to use the caller's.

Source code in vllm/distributed/eplb/eplb_communicator.py

class EplbCommunicator(ABC):
    """Abstract EPLB communicator for expert weight transfers."""

    @abstractmethod
    def add_send(
        self,
        tensors: list[torch.Tensor],
        dst_rank: int,
        expert_id: int,
    ) -> None:
        pass

    @abstractmethod
    def add_recv(
        self,
        tensors: list[torch.Tensor],
        src_rank: int,
        expert_id: int,
    ) -> None:
        pass

    @abstractmethod
    def execute(self) -> None:
        """Complete all enqueued transfers.

        Some backends perform communication here; others (e.g. NIXL)
        issue transfers eagerly in add_recv and only wait here.
        On return, all data is available in the destination buffers.
        """

    def set_transfer_context(  # noqa: B027
        self,
        old_indices: np.ndarray,
        layer_idx: int,
        recv_rows: Sequence[dict[int, int]] | None = None,
    ) -> None:
        """Pre-set layer context before add_send/add_recv calls.

        Default is a no-op; overridden by backends (e.g. NIXL) that need
        layer-level context to issue transfers inside add_recv.

        Args:
            old_indices: Global physical-to-logical map before the move.
            layer_idx: Index of the MoE layer being transferred.
            recv_rows: Per EP rank, the ``expert_buffer`` row each expert
                received by that rank lands in. Only push backends need it.

        """

    @property
    def needs_recv_rows(self) -> bool:
        """Whether set_transfer_context must be given every rank's recv rows."""
        return False

    @property
    def owned_expert_buffer(self) -> list[torch.Tensor] | None:
        """Backend-allocated receive buffer, or ``None`` to use the caller's."""
        return None

    @property
    def needs_profile_buffer_reservation(self) -> bool:
        """Whether the profile path must run a dummy collective operation to reserve
        communication buffers."""
        return True

    def set_stream(self, stream: torch.Stream | None) -> None:
        self._stream = stream

    def _log_initialized(self) -> None:
        if is_local_first_rank():
            logger.info("Initialized EPLB communicator: %s.", self.__class__.__name__)

needs_profile_buffer_reservation property

Whether the profile path must run a dummy collective operation to reserve communication buffers.

needs_recv_rows property

Whether set_transfer_context must be given every rank's recv rows.

owned_expert_buffer property

Backend-allocated receive buffer, or None to use the caller's.

execute() abstractmethod

Complete all enqueued transfers.

Some backends perform communication here; others (e.g. NIXL) issue transfers eagerly in add_recv and only wait here. On return, all data is available in the destination buffers.

Source code in vllm/distributed/eplb/eplb_communicator.py

@abstractmethod
def execute(self) -> None:
    """Complete all enqueued transfers.

    Some backends perform communication here; others (e.g. NIXL)
    issue transfers eagerly in add_recv and only wait here.
    On return, all data is available in the destination buffers.
    """

set_transfer_context(old_indices, layer_idx, recv_rows=None)

Pre-set layer context before add_send/add_recv calls.

Default is a no-op; overridden by backends (e.g. NIXL) that need layer-level context to issue transfers inside add_recv.

Parameters:

• old_indices

  (ndarray) –

  Global physical-to-logical map before the move.

• layer_idx

  (int) –

  Index of the MoE layer being transferred.

• recv_rows

  (Sequence[dict[int, int]] | None, default: None ) –

  Per EP rank, the expert_buffer row each expert received by that rank lands in. Only push backends need it.

Source code in vllm/distributed/eplb/eplb_communicator.py

def set_transfer_context(  # noqa: B027
    self,
    old_indices: np.ndarray,
    layer_idx: int,
    recv_rows: Sequence[dict[int, int]] | None = None,
) -> None:
    """Pre-set layer context before add_send/add_recv calls.

    Default is a no-op; overridden by backends (e.g. NIXL) that need
    layer-level context to issue transfers inside add_recv.

    Args:
        old_indices: Global physical-to-logical map before the move.
        layer_idx: Index of the MoE layer being transferred.
        recv_rows: Per EP rank, the ``expert_buffer`` row each expert
            received by that rank lands in. Only push backends need it.

    """

NixlEplbCommunicator

Bases: EplbCommunicator

EPLB communicator backed by NIXL READ transfers.

Methods:

• __init__ –

  Create a NIXL-backed EPLB communicator.

Source code in vllm/distributed/eplb/eplb_communicator.py

class NixlEplbCommunicator(EplbCommunicator):
    """EPLB communicator backed by NIXL READ transfers."""

    def __init__(
        self,
        cpu_group: ProcessGroup,
        all_expert_weights: Sequence[Sequence[torch.Tensor]],
        expert_buffer: Sequence[torch.Tensor],
    ) -> None:
        """Create a NIXL-backed EPLB communicator.

        Args:
            cpu_group: CPU process group for metadata exchange.
            all_expert_weights: Expert weight tensors for all MoE layers.
            expert_buffer: Pre-allocated receive buffer tensors.

        """
        assert all_expert_weights, (
            "NixlEplbCommunicator requires non-empty all_expert_weights."
        )
        assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer."
        nixl_wrapper_cls = nixl_utils.NixlWrapper
        if nixl_wrapper_cls is None:
            raise RuntimeError("NIXL is unavailable.")

        self._cpu_group = cpu_group
        self._world_size = cpu_group.size()
        self._rank = cpu_group.rank()

        self._all_expert_weights = all_expert_weights
        self._expert_buffer = expert_buffer
        self._num_local_experts: int = all_expert_weights[0][0].shape[0]
        self._device = all_expert_weights[0][0].device

        for layer_tensors in all_expert_weights:
            for tensor in layer_tensors:
                assert is_weak_contiguous(tensor), (
                    "Expert weight tensors must be contiguous in memory"
                )
                assert tensor.device == self._device, (
                    "All local EPLB tensors are expected to be on the same "
                    f"device: expected={self._device}, got={tensor.device}"
                )
        for tensor in expert_buffer:
            assert is_weak_contiguous(tensor), (
                "expert_buffer tensors must be contiguous in memory"
            )

        # (local_dlist, remote_dlist, xfer_handle) for in-flight READs;
        # accumulated by add_recv, drained by execute.
        self._xfer_entries: list[tuple[int, int, int]] = []
        # Per-rank expert_id -> physical row; set by set_transfer_context.
        self._expert_to_src_row: list[dict[int, int]] | None = None
        self._layer_idx: int | None = None
        self.set_stream(None)

        nixl_agent_config = nixl_utils.nixl_agent_config
        config = (
            nixl_agent_config(capture_telemetry=False)
            if nixl_agent_config is not None
            else None
        )
        self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config)
        self._nixl_memory_type = "VRAM"
        # NIXL registration handles; deregistered in __del__.
        self._registered_descs: list[object] = []
        self._remote_agents: dict[int, str] = {}
        # peer -> (layer, tensor) -> (base_ptr, bytes_per_expert, dev_id).
        self._remote_send_meta: dict[
            int, dict[tuple[int, int], tuple[int, int, int]]
        ] = {}

        self._cuda_device_id = int(self._device.index or 0)
        self._init_step("buffers", self._init_registered_buffers)
        self._init_remote_state()
        self._log_initialized()

    def _init_remote_state(self) -> None:
        """Exchange NIXL agent metadata and RDMA pointer info with all peers.

        This is a collective operation (uses ``all_gather_object`` twice).
        """
        self._init_step("agents", self._init_remote_agents)
        self._init_step("send meta", self._exchange_remote_send_meta)

    @property
    def needs_profile_buffer_reservation(self) -> bool:
        return False

    @staticmethod
    def _init_step(name: str, fn: object, *args: object, **kwargs: object) -> None:
        try:
            fn(*args, **kwargs)  # type: ignore[operator]
        except Exception as exc:
            raise RuntimeError(f"NIXL EPLB init failed: {name}") from exc

    def _make_agent_name(self) -> str:
        """Build a deployment-unique nixl agent name."""
        pp_size = get_pp_group().world_size
        pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else ""
        uid = uuid.uuid4().hex[:8]
        return f"eplb-{self._rank}{pp_suffix}-{uid}"

    def set_stream(self, stream: torch.Stream | None) -> None:
        self._stream = stream

    def add_send(
        self,
        tensors: list[torch.Tensor],
        dst_rank: int,
        expert_id: int,
    ) -> None:
        # No-op: NIXL READ is receiver-initiated. The sender's expert
        # weights are pre-registered and always readable in-place.
        pass

    def set_transfer_context(
        self,
        old_indices: np.ndarray,
        layer_idx: int,
        recv_rows: Sequence[dict[int, int]] | None = None,
    ) -> None:
        assert not self._xfer_entries, (
            f"set_transfer_context() called with {len(self._xfer_entries)} "
            f"pending transfers from layer {self._layer_idx}; "
            f"execute() was not called after previous add_recv() calls"
        )
        # NIXL READs bypass device streams; wait for pending copies out of
        # the receive buffer before add_recv overwrites it.
        if current_platform.is_rocm():
            with gpu_sync_allowed():
                torch.accelerator.current_stream().synchronize()
                if self._stream is not None:
                    self._stream.synchronize()
        self._layer_idx = layer_idx
        n = self._num_local_experts
        rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n)
        self._expert_to_src_row = [
            {int(eid): i for i, eid in enumerate(row) if eid != -1}
            for row in rank_experts
        ]

    def add_recv(
        self,
        tensors: list[torch.Tensor],
        src_rank: int,
        expert_id: int,
    ) -> None:
        # Build NIXL descriptors and issue the RDMA READ immediately,
        # overlapping the transfer with the remaining Python loop in
        # move_to_buffer.
        assert self._expert_to_src_row is not None and self._layer_idx is not None, (
            "set_transfer_context() must be called before add_recv()"
        )
        src_row = self._expert_to_src_row[src_rank][expert_id]
        layer_idx = self._layer_idx

        local_descs: list[tuple[int, int, int]] = []
        remote_descs: list[tuple[int, int, int]] = []
        for t_idx, t in enumerate(tensors):
            send_base, send_stride, remote_dev = self._remote_send_meta[src_rank][
                (layer_idx, t_idx)
            ]
            assert t.nbytes == send_stride, (
                f"tensor {t_idx} size {t.nbytes} != remote stride {send_stride}"
            )
            local_descs.append(
                (
                    t.data_ptr(),
                    t.nbytes,
                    self._cuda_device_id,
                )
            )
            remote_descs.append(
                (
                    send_base + src_row * send_stride,
                    send_stride,
                    remote_dev,
                )
            )

        local_h, remote_h, xfer_h = self._create_peer_xfer(
            src_rank, local_descs, remote_descs
        )
        self._nixl_wrapper.transfer(xfer_h)
        self._xfer_entries.append((local_h, remote_h, xfer_h))

    def _init_remote_agents(self) -> None:
        local_metadata = self._nixl_wrapper.get_agent_metadata()
        gathered_metadata: list[bytes | None] = [None] * self._world_size
        torch.distributed.all_gather_object(
            gathered_metadata, local_metadata, group=self._cpu_group
        )
        for peer in range(self._world_size):
            if peer == self._rank:
                continue
            peer_metadata = gathered_metadata[peer]
            assert peer_metadata is not None
            self._remote_agents[peer] = self._nixl_wrapper.add_remote_agent(
                peer_metadata
            )

    def _init_registered_buffers(self) -> None:
        all_tensors: list[torch.Tensor] = []
        for layer_tensors in self._all_expert_weights:
            all_tensors.extend(layer_tensors)
        all_tensors.extend(self._expert_buffer)

        descs = self._nixl_wrapper.get_reg_descs(all_tensors)
        self._nixl_wrapper.register_memory(descs)
        self._registered_descs.append(descs)

    def _exchange_remote_send_meta(self) -> None:
        """Exchange per-layer per-tensor metadata so receivers can compute
        remote RDMA addresses at transfer time."""
        local_meta: dict[tuple[int, int], tuple[int, int, int]] = {}
        for layer_idx, layer_tensors in enumerate(self._all_expert_weights):
            for t_idx, t in enumerate(layer_tensors):
                nbytes_per_expert = t.nbytes // self._num_local_experts
                local_meta[(layer_idx, t_idx)] = (
                    t.data_ptr(),
                    nbytes_per_expert,
                    self._cuda_device_id,
                )

        # Per-rank map: (layer_idx, tensor_idx) -> (base_ptr, bytes_per_expert, dev_id).
        # add_recv uses base_ptr + src_row * bytes_per_expert to compute
        # the remote RDMA address for each expert.
        gathered_meta: list[dict[tuple[int, int], tuple[int, int, int]] | None] = [
            None
        ] * self._world_size
        torch.distributed.all_gather_object(
            gathered_meta, local_meta, group=self._cpu_group
        )

        local_keys = set(local_meta.keys())
        for peer in self._remote_agents:
            peer_meta = gathered_meta[peer]
            assert peer_meta is not None
            peer_keys = set(peer_meta.keys())
            if peer_keys != local_keys:
                raise RuntimeError(
                    f"NIXL EPLB metadata key mismatch with rank {peer}: "
                    f"local={sorted(local_keys)}, peer={sorted(peer_keys)}"
                )
            for key in local_keys:
                _, local_stride, _ = local_meta[key]
                _, peer_stride, _ = peer_meta[key]
                if local_stride != peer_stride:
                    raise RuntimeError(
                        f"NIXL EPLB nbytes_per_expert mismatch for {key} "
                        f"with rank {peer}: "
                        f"local={local_stride}, peer={peer_stride}"
                    )
            self._remote_send_meta[peer] = peer_meta

    def _wait_for_all_transfers(self, handles: list[int]) -> None:
        pending = set(handles)
        while pending:
            completed: list[int] = []
            for handle in pending:
                state = self._nixl_wrapper.check_xfer_state(handle)
                if state == "DONE":
                    completed.append(handle)
                    continue
                if state != "PROC":
                    raise RuntimeError(f"NIXL transfer failed with state={state}")
            for handle in completed:
                pending.remove(handle)
            if pending:
                time.sleep(0.0005)

    def _create_peer_xfer(
        self,
        src: int,
        local_descs: list[tuple[int, int, int]],
        remote_descs: list[tuple[int, int, int]],
    ) -> tuple[int, int, int]:
        """Create a batched xfer for multiple descriptors from one peer.

        Each element in *local_descs* / *remote_descs* is an
        ``(address, size, device_id)`` tuple.

        Returns ``(local_dlist, remote_dlist, xfer_handle)``.
        """
        local_desc = self._nixl_wrapper.get_xfer_descs(
            local_descs, self._nixl_memory_type
        )
        local_handle = self._nixl_wrapper.prep_xfer_dlist(
            "NIXL_INIT_AGENT",
            local_desc,
        )

        remote_desc = self._nixl_wrapper.get_xfer_descs(
            remote_descs, self._nixl_memory_type
        )
        remote_handle = self._nixl_wrapper.prep_xfer_dlist(
            self._remote_agents[src],
            remote_desc,
        )

        indices = np.arange(len(local_descs), dtype=np.int32)
        xfer_handle = self._nixl_wrapper.make_prepped_xfer(
            "READ",
            local_handle,
            indices,
            remote_handle,
            indices,
        )
        return (local_handle, remote_handle, xfer_handle)

    def _post_read_barrier(self) -> None:
        """Correctness fence: prevents overwrite-while-remote-read race.

        We avoid ``torch.distributed.monitored_barrier`` because it
        calls ``get_backend(group)`` which fails for stateless groups
        (elastic EP).  An async ``all_reduce`` + ``wait(timeout)``
        works with both regular and stateless groups and provides
        equivalent timeout detection.
        """
        _dummy = torch.zeros(1, dtype=torch.int32)
        work = torch.distributed.all_reduce(
            _dummy, group=self._cpu_group, async_op=True
        )
        work.wait(timeout=timedelta(minutes=5))

    def execute(self) -> None:
        assert self._layer_idx is not None or not self._xfer_entries, (
            "set_transfer_context() must be called before execute() "
            "if any add_recv() calls were made"
        )
        try:
            self._wait_for_all_transfers([x[2] for x in self._xfer_entries])

            self._post_read_barrier()
        finally:
            for local_h, remote_h, xfer_h in self._xfer_entries:
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.release_xfer_handle(xfer_h)
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.release_dlist_handle(local_h)
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.release_dlist_handle(remote_h)
            self._xfer_entries.clear()
            self._expert_to_src_row = None
            self._layer_idx = None

    def __del__(self) -> None:
        with contextlib.suppress(Exception):
            for local_h, remote_h, xfer_h in self._xfer_entries:
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.release_xfer_handle(xfer_h)
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.release_dlist_handle(local_h)
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.release_dlist_handle(remote_h)
        with contextlib.suppress(Exception):
            for descs in self._registered_descs:
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.deregister_memory(descs)
            self._registered_descs.clear()
        with contextlib.suppress(Exception):
            for agent_name in self._remote_agents.values():
                with contextlib.suppress(Exception):
                    self._nixl_wrapper.remove_remote_agent(agent_name)
            self._remote_agents.clear()

__init__(cpu_group, all_expert_weights, expert_buffer)

Create a NIXL-backed EPLB communicator.

Parameters:

• cpu_group

  (ProcessGroup) –

  CPU process group for metadata exchange.

• all_expert_weights

  (Sequence[Sequence[Tensor]]) –

  Expert weight tensors for all MoE layers.

• expert_buffer

  (Sequence[Tensor]) –

  Pre-allocated receive buffer tensors.

Source code in vllm/distributed/eplb/eplb_communicator.py

def __init__(
    self,
    cpu_group: ProcessGroup,
    all_expert_weights: Sequence[Sequence[torch.Tensor]],
    expert_buffer: Sequence[torch.Tensor],
) -> None:
    """Create a NIXL-backed EPLB communicator.

    Args:
        cpu_group: CPU process group for metadata exchange.
        all_expert_weights: Expert weight tensors for all MoE layers.
        expert_buffer: Pre-allocated receive buffer tensors.

    """
    assert all_expert_weights, (
        "NixlEplbCommunicator requires non-empty all_expert_weights."
    )
    assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer."
    nixl_wrapper_cls = nixl_utils.NixlWrapper
    if nixl_wrapper_cls is None:
        raise RuntimeError("NIXL is unavailable.")

    self._cpu_group = cpu_group
    self._world_size = cpu_group.size()
    self._rank = cpu_group.rank()

    self._all_expert_weights = all_expert_weights
    self._expert_buffer = expert_buffer
    self._num_local_experts: int = all_expert_weights[0][0].shape[0]
    self._device = all_expert_weights[0][0].device

    for layer_tensors in all_expert_weights:
        for tensor in layer_tensors:
            assert is_weak_contiguous(tensor), (
                "Expert weight tensors must be contiguous in memory"
            )
            assert tensor.device == self._device, (
                "All local EPLB tensors are expected to be on the same "
                f"device: expected={self._device}, got={tensor.device}"
            )
    for tensor in expert_buffer:
        assert is_weak_contiguous(tensor), (
            "expert_buffer tensors must be contiguous in memory"
        )

    # (local_dlist, remote_dlist, xfer_handle) for in-flight READs;
    # accumulated by add_recv, drained by execute.
    self._xfer_entries: list[tuple[int, int, int]] = []
    # Per-rank expert_id -> physical row; set by set_transfer_context.
    self._expert_to_src_row: list[dict[int, int]] | None = None
    self._layer_idx: int | None = None
    self.set_stream(None)

    nixl_agent_config = nixl_utils.nixl_agent_config
    config = (
        nixl_agent_config(capture_telemetry=False)
        if nixl_agent_config is not None
        else None
    )
    self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config)
    self._nixl_memory_type = "VRAM"
    # NIXL registration handles; deregistered in __del__.
    self._registered_descs: list[object] = []
    self._remote_agents: dict[int, str] = {}
    # peer -> (layer, tensor) -> (base_ptr, bytes_per_expert, dev_id).
    self._remote_send_meta: dict[
        int, dict[tuple[int, int], tuple[int, int, int]]
    ] = {}

    self._cuda_device_id = int(self._device.index or 0)
    self._init_step("buffers", self._init_registered_buffers)
    self._init_remote_state()
    self._log_initialized()

_create_peer_xfer(src, local_descs, remote_descs)

Create a batched xfer for multiple descriptors from one peer.

Each element in local_descs / remote_descs is an (address, size, device_id) tuple.

Returns (local_dlist, remote_dlist, xfer_handle).

Source code in vllm/distributed/eplb/eplb_communicator.py

def _create_peer_xfer(
    self,
    src: int,
    local_descs: list[tuple[int, int, int]],
    remote_descs: list[tuple[int, int, int]],
) -> tuple[int, int, int]:
    """Create a batched xfer for multiple descriptors from one peer.

    Each element in *local_descs* / *remote_descs* is an
    ``(address, size, device_id)`` tuple.

    Returns ``(local_dlist, remote_dlist, xfer_handle)``.
    """
    local_desc = self._nixl_wrapper.get_xfer_descs(
        local_descs, self._nixl_memory_type
    )
    local_handle = self._nixl_wrapper.prep_xfer_dlist(
        "NIXL_INIT_AGENT",
        local_desc,
    )

    remote_desc = self._nixl_wrapper.get_xfer_descs(
        remote_descs, self._nixl_memory_type
    )
    remote_handle = self._nixl_wrapper.prep_xfer_dlist(
        self._remote_agents[src],
        remote_desc,
    )

    indices = np.arange(len(local_descs), dtype=np.int32)
    xfer_handle = self._nixl_wrapper.make_prepped_xfer(
        "READ",
        local_handle,
        indices,
        remote_handle,
        indices,
    )
    return (local_handle, remote_handle, xfer_handle)

_exchange_remote_send_meta()

Exchange per-layer per-tensor metadata so receivers can compute remote RDMA addresses at transfer time.

Source code in vllm/distributed/eplb/eplb_communicator.py

def _exchange_remote_send_meta(self) -> None:
    """Exchange per-layer per-tensor metadata so receivers can compute
    remote RDMA addresses at transfer time."""
    local_meta: dict[tuple[int, int], tuple[int, int, int]] = {}
    for layer_idx, layer_tensors in enumerate(self._all_expert_weights):
        for t_idx, t in enumerate(layer_tensors):
            nbytes_per_expert = t.nbytes // self._num_local_experts
            local_meta[(layer_idx, t_idx)] = (
                t.data_ptr(),
                nbytes_per_expert,
                self._cuda_device_id,
            )

    # Per-rank map: (layer_idx, tensor_idx) -> (base_ptr, bytes_per_expert, dev_id).
    # add_recv uses base_ptr + src_row * bytes_per_expert to compute
    # the remote RDMA address for each expert.
    gathered_meta: list[dict[tuple[int, int], tuple[int, int, int]] | None] = [
        None
    ] * self._world_size
    torch.distributed.all_gather_object(
        gathered_meta, local_meta, group=self._cpu_group
    )

    local_keys = set(local_meta.keys())
    for peer in self._remote_agents:
        peer_meta = gathered_meta[peer]
        assert peer_meta is not None
        peer_keys = set(peer_meta.keys())
        if peer_keys != local_keys:
            raise RuntimeError(
                f"NIXL EPLB metadata key mismatch with rank {peer}: "
                f"local={sorted(local_keys)}, peer={sorted(peer_keys)}"
            )
        for key in local_keys:
            _, local_stride, _ = local_meta[key]
            _, peer_stride, _ = peer_meta[key]
            if local_stride != peer_stride:
                raise RuntimeError(
                    f"NIXL EPLB nbytes_per_expert mismatch for {key} "
                    f"with rank {peer}: "
                    f"local={local_stride}, peer={peer_stride}"
                )
        self._remote_send_meta[peer] = peer_meta

_init_remote_state()

Exchange NIXL agent metadata and RDMA pointer info with all peers.

This is a collective operation (uses all_gather_object twice).

Source code in vllm/distributed/eplb/eplb_communicator.py

def _init_remote_state(self) -> None:
    """Exchange NIXL agent metadata and RDMA pointer info with all peers.

    This is a collective operation (uses ``all_gather_object`` twice).
    """
    self._init_step("agents", self._init_remote_agents)
    self._init_step("send meta", self._exchange_remote_send_meta)

_make_agent_name()

Build a deployment-unique nixl agent name.

Source code in vllm/distributed/eplb/eplb_communicator.py

def _make_agent_name(self) -> str:
    """Build a deployment-unique nixl agent name."""
    pp_size = get_pp_group().world_size
    pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else ""
    uid = uuid.uuid4().hex[:8]
    return f"eplb-{self._rank}{pp_suffix}-{uid}"

_post_read_barrier()

Correctness fence: prevents overwrite-while-remote-read race.

We avoid torch.distributed.monitored_barrier because it calls get_backend(group) which fails for stateless groups (elastic EP). An async all_reduce + wait(timeout) works with both regular and stateless groups and provides equivalent timeout detection.

Source code in vllm/distributed/eplb/eplb_communicator.py

def _post_read_barrier(self) -> None:
    """Correctness fence: prevents overwrite-while-remote-read race.

    We avoid ``torch.distributed.monitored_barrier`` because it
    calls ``get_backend(group)`` which fails for stateless groups
    (elastic EP).  An async ``all_reduce`` + ``wait(timeout)``
    works with both regular and stateless groups and provides
    equivalent timeout detection.
    """
    _dummy = torch.zeros(1, dtype=torch.int32)
    work = torch.distributed.all_reduce(
        _dummy, group=self._cpu_group, async_op=True
    )
    work.wait(timeout=timedelta(minutes=5))

PyNcclEplbCommunicator

Bases: EplbCommunicator

EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv.

Source code in vllm/distributed/eplb/eplb_communicator.py

class PyNcclEplbCommunicator(EplbCommunicator):
    """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv."""

    def __init__(
        self,
        pynccl_comm: PyNcclCommunicator,
        stream: torch.Stream | None = None,
    ) -> None:
        self._pynccl_comm = pynccl_comm
        self._stream = stream
        self._group_started = False
        self._log_initialized()

    def _ensure_group_started(self) -> None:
        if not self._group_started:
            self._pynccl_comm.group_start()
            self._group_started = True

    def add_send(
        self,
        tensors: list[torch.Tensor],
        dst_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        self._ensure_group_started()
        for tensor in tensors:
            self._pynccl_comm.send(tensor, dst_rank, stream=self._stream)

    def add_recv(
        self,
        tensors: list[torch.Tensor],
        src_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        self._ensure_group_started()
        for tensor in tensors:
            self._pynccl_comm.recv(tensor, src_rank, stream=self._stream)

    def execute(self) -> None:
        if self._group_started:
            self._pynccl_comm.group_end()
            self._group_started = False

SymmMemEplbCommunicator

Bases: EplbCommunicator

Push expert rows into peers' receive buffers over symmetric memory.

Senders copy from their live weights into the receiver's expert_buffer row, which every rank derives from the global placements; execute is collective. CPU barriers only: GPU spin barriers can deadlock MegaMoE cooperative launches.

Source code in vllm/distributed/eplb/eplb_communicator.py

class SymmMemEplbCommunicator(EplbCommunicator):
    """Push expert rows into peers' receive buffers over symmetric memory.

    Senders copy from their live weights into the receiver's ``expert_buffer``
    row, which every rank derives from the global placements; ``execute`` is
    collective. CPU barriers only: GPU spin barriers can deadlock MegaMoE
    cooperative launches.
    """

    def __init__(
        self,
        ep_group: ProcessGroup,
        cpu_group: ProcessGroup,
        layer: Sequence[torch.Tensor],
    ) -> None:
        import torch.distributed._symmetric_memory as symm_mem

        from vllm.distributed.device_communicators.custom_all_reduce import (
            _group_can_attempt_mnnvl,
        )

        device = layer[0].device
        if not current_platform.is_cuda():
            raise ValueError("symm_mem EPLB requires CUDA")
        if all(in_the_same_node_as(cpu_group, source_rank=0)):
            physical_ids = [0] * cpu_group.size()
            torch.distributed.all_gather_object(
                physical_ids,
                current_platform.visible_device_id_to_physical_device_id(device.index),
                group=cpu_group,
            )
            in_nvlink_domain = current_platform.is_fully_connected(physical_ids)
        else:
            in_nvlink_domain = _group_can_attempt_mnnvl(cpu_group, device)
        if not in_nvlink_domain:
            raise ValueError(
                "symm_mem EPLB requires all EP ranks in one NVLink domain "
                "(MNNVL across nodes); use nixl or torch_gloo instead."
            )

        # One buffer and one rendezvous cover all of the layer's tensors.
        offsets = [0, *itertools.accumulate(round_up(w.nbytes, 256) for w in layer)]

        def as_layer(buffer: torch.Tensor) -> list[torch.Tensor]:
            return [
                buffer[o : o + w.nbytes].view(w.dtype).view(w.shape)
                for o, w in zip(offsets, layer)
            ]

        buffer = symm_mem.empty(offsets[-1], dtype=torch.uint8, device=device)
        # For MNNVL, rendezvous is what confirms all ranks share one NVLink domain.
        self._handle = symm_mem.rendezvous(buffer, ep_group.group_name)
        self._peer_buffers = [
            as_layer(self._handle.get_buffer(rank, buffer.shape, torch.uint8))
            for rank in range(ep_group.size())
        ]
        self._buffer = as_layer(buffer)

        self._stream = None
        self._cpu_group = cpu_group
        self._recv_rows: Sequence[dict[int, int]] = []
        self._sends: list[tuple[list[torch.Tensor], int, int]] = []
        self._log_initialized()

    @property
    def needs_profile_buffer_reservation(self) -> bool:
        return False

    @property
    def needs_recv_rows(self) -> bool:
        return True

    @property
    def owned_expert_buffer(self) -> list[torch.Tensor]:
        return self._buffer

    def set_transfer_context(
        self,
        old_indices: np.ndarray,
        layer_idx: int,
        recv_rows: Sequence[dict[int, int]] | None = None,
    ) -> None:
        assert recv_rows is not None
        self._recv_rows = recv_rows

    def add_send(
        self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int
    ) -> None:
        row = self._recv_rows[dst_rank][expert_id]
        self._sends.append((tensors, dst_rank, row))

    def add_recv(
        self, tensors: list[torch.Tensor], src_rank: int, expert_id: int
    ) -> None:
        # No-op: senders push into this rank's buffer in execute().
        pass

    def execute(self) -> None:
        with device_stream(self._stream):
            self._barrier()  # Every peer's buffer rows are free to overwrite.
            for tensors, dst_rank, row in self._sends:
                for src, peer in zip(tensors, self._peer_buffers[dst_rank]):
                    peer[row].copy_(src, non_blocking=True)
            self._barrier()  # Every push has landed in its receiver's buffer.
        self._sends.clear()

    def _barrier(self) -> None:
        with gpu_sync_allowed():
            torch.accelerator.current_stream().synchronize()
        # A CPU all_reduce stays on gloo; barrier() would use NCCL on split groups.
        torch.distributed.all_reduce(
            torch.zeros(1, dtype=torch.int32), group=self._cpu_group
        )

TorchDistGlooStagedEplbCommunicator

Bases: EplbCommunicator

EPLB communicator using gloo P2P with CPU staging.

Source code in vllm/distributed/eplb/eplb_communicator.py

class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
    """EPLB communicator using gloo P2P with CPU staging."""

    def __init__(
        self,
        cpu_group: ProcessGroup,
        stream: torch.Stream | None = None,
    ) -> None:
        self._cpu_group = cpu_group
        self._stream = stream
        self._ops: list[tuple[str, torch.Tensor, int]] = []
        self._log_initialized()

    def add_send(
        self,
        tensors: list[torch.Tensor],
        dst_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        for tensor in tensors:
            self._ops.append(("send", tensor, dst_rank))

    def add_recv(
        self,
        tensors: list[torch.Tensor],
        src_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        for tensor in tensors:
            self._ops.append(("recv", tensor, src_rank))

    def execute(self) -> None:
        if not self._ops:
            return

        p2p_ops: list[P2POp] = []
        recv_staging: list[tuple[torch.Tensor, torch.Tensor]] = []

        def build_ops() -> None:
            for op, tensor, peer_rank in self._ops:
                if op == "send":
                    cpu_tensor = tensor.to(device="cpu", non_blocking=True)
                    p2p_ops.append(
                        P2POp(
                            torch.distributed.isend,
                            cpu_tensor,
                            group=self._cpu_group,
                            group_peer=peer_rank,
                        )
                    )
                    continue
                cpu_tensor = torch.empty_like(
                    tensor, device="cpu", pin_memory=PIN_MEMORY
                )
                p2p_ops.append(
                    P2POp(
                        torch.distributed.irecv,
                        cpu_tensor,
                        group=self._cpu_group,
                        group_peer=peer_rank,
                    )
                )
                recv_staging.append((tensor, cpu_tensor))

        try:
            with device_stream(self._stream):
                build_ops()
        finally:
            self._ops.clear()

        # Wait for all D2H copies to finish
        # before issuing gloo batch_isend_irecv operations.
        with gpu_sync_allowed():
            if self._stream is not None:
                self._stream.synchronize()
            else:
                torch.accelerator.current_stream().synchronize()

        reqs = batch_isend_irecv(p2p_ops)
        for req in reqs:
            req.wait()

        if not recv_staging:
            return
        with device_stream(self._stream):
            for dst_tensor, cpu_tensor in recv_staging:
                dst_tensor.copy_(cpu_tensor, non_blocking=True)

TorchDistNcclEplbCommunicator

Bases: EplbCommunicator

EPLB communicator backed by torch.distributed isend/irecv.

Source code in vllm/distributed/eplb/eplb_communicator.py

class TorchDistNcclEplbCommunicator(EplbCommunicator):
    """EPLB communicator backed by torch.distributed isend/irecv."""

    def __init__(
        self,
        ep_group: ProcessGroup,
        stream: torch.Stream | None = None,
    ) -> None:
        self._ep_group = ep_group
        self._stream = stream
        self._p2p_ops: list[P2POp] = []
        self._log_initialized()

    def add_send(
        self,
        tensors: list[torch.Tensor],
        dst_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        for tensor in tensors:
            self._p2p_ops.append(
                P2POp(
                    torch.distributed.isend,
                    tensor,
                    group=self._ep_group,
                    group_peer=dst_rank,
                )
            )

    def add_recv(
        self,
        tensors: list[torch.Tensor],
        src_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        for tensor in tensors:
            self._p2p_ops.append(
                P2POp(
                    torch.distributed.irecv,
                    tensor,
                    group=self._ep_group,
                    group_peer=src_rank,
                )
            )

    def execute(self) -> None:
        if not self._p2p_ops:
            return
        try:
            with device_stream(self._stream):
                reqs = batch_isend_irecv(self._p2p_ops)
                for req in reqs:
                    req.wait()
        finally:
            self._p2p_ops.clear()

TorchDistXCCLStagedEplbCommunicator

Bases: EplbCommunicator

EPLB communicator using XCCL device-to-device P2P on XPU.

Source code in vllm/distributed/eplb/eplb_communicator.py

class TorchDistXCCLStagedEplbCommunicator(EplbCommunicator):
    """EPLB communicator using XCCL device-to-device P2P on XPU."""

    def __init__(
        self,
        ep_group: ProcessGroup,
        stream: torch.Stream | None = None,
    ) -> None:
        self._ep_group = ep_group
        self._stream = stream
        self._ops: list[tuple[str, torch.Tensor, int]] = []
        self._log_initialized()

    def add_send(
        self,
        tensors: list[torch.Tensor],
        dst_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        for tensor in tensors:
            self._ops.append(("send", tensor, dst_rank))

    def add_recv(
        self,
        tensors: list[torch.Tensor],
        src_rank: int,
        expert_id: int,  # unused by this backend
    ) -> None:
        for tensor in tensors:
            self._ops.append(("recv", tensor, src_rank))

    def execute(self) -> None:
        if not self._ops:
            return

        p2p_ops: list[P2POp] = []
        try:
            for op, tensor, peer_rank in self._ops:
                send_or_recv = (
                    torch.distributed.isend if op == "send" else torch.distributed.irecv
                )
                p2p_ops.append(
                    P2POp(
                        send_or_recv,
                        tensor,
                        peer_rank,
                        self._ep_group,
                    )
                )
        finally:
            self._ops.clear()

        with device_stream(self._stream):
            reqs = batch_isend_irecv(p2p_ops)
            for req in reqs:
                req.wait()

create_eplb_communicator(group_coordinator, backend, expert_weights, expert_buffer)

Create an EPLB communicator for the given backend.

Parameters:

• group_coordinator

  (GroupCoordinator) –

  Process-group coordinator that provides the device and CPU communication groups.

• backend

  (str) –

  Communicator backend name ("torch_nccl", "torch_gloo", "pynccl", "nixl", or "symm_mem"). Falls back to "torch_nccl" when None. Stateless (elastic EP) groups support "torch_nccl", "pynccl", and "nixl"; "torch_nccl" is silently promoted to "pynccl". When tensors reside on CPU, "torch_gloo" or "torch_nccl" are used via the CPU process group.

• expert_weights

  (Sequence[Sequence[Tensor]]) –

  Expert weight tensors for all MoE layers. Shape (num_layers)(num_tensors_per_layer). NixlEplbCommunicator registers all layers with NIXL for zero-copy RDMA reads.

• expert_buffer

  (Sequence[Tensor] | None) –

  Pre-allocated receive buffer tensors (one per weight tensor in a single layer). symm_mem ignores it and uses its own owned_expert_buffer.

Source code in vllm/distributed/eplb/eplb_communicator.py

def create_eplb_communicator(
    group_coordinator: GroupCoordinator,
    backend: str,
    expert_weights: Sequence[Sequence[torch.Tensor]],
    expert_buffer: Sequence[torch.Tensor] | None,
) -> EplbCommunicator:
    """Create an EPLB communicator for the given backend.

    Args:
        group_coordinator: Process-group coordinator that provides the
            device and CPU communication groups.
        backend: Communicator backend name (``"torch_nccl"``,
            ``"torch_gloo"``, ``"pynccl"``, ``"nixl"``, or ``"symm_mem"``).
            Falls back to ``"torch_nccl"`` when *None*.
            Stateless (elastic EP) groups support ``"torch_nccl"``,
            ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently
            promoted to ``"pynccl"``.  When tensors reside on CPU,
            ``"torch_gloo"`` or
            ``"torch_nccl"`` are used via the CPU process group.
        expert_weights: Expert weight tensors for *all* MoE layers.
            Shape ``(num_layers)(num_tensors_per_layer)``.
            NixlEplbCommunicator registers all layers with NIXL for
            zero-copy RDMA reads.
        expert_buffer: Pre-allocated receive buffer tensors (one per
            weight tensor in a single layer). ``symm_mem`` ignores it and
            uses its own ``owned_expert_buffer``.

    """
    first_layer = expert_weights[0] if expert_weights else []
    tensor_device_type = first_layer[0].device.type if first_layer else "cpu"
    torch_group = (
        group_coordinator.cpu_group
        if tensor_device_type == "cpu"
        else group_coordinator.device_group
    )

    def _create_pynccl() -> EplbCommunicator:
        if tensor_device_type == "cpu":
            raise RuntimeError(
                "EPLB communicator 'pynccl' supports only cuda-like devices "
                f"(got {tensor_device_type})."
            )
        unsupported_dtypes = sorted(
            {
                tensor.dtype
                for tensor in first_layer
                if not ncclDataTypeEnum.supports_torch_dtype(tensor.dtype)
            },
            key=str,
        )
        if unsupported_dtypes:
            raise RuntimeError(
                "EPLB communicator 'pynccl' requested but expert weights contain "
                "unsupported dtypes: "
                f"({', '.join(str(dtype) for dtype in unsupported_dtypes)})."
            )

        device_comm = group_coordinator.device_communicator
        pynccl_comm = (
            getattr(device_comm, "pynccl_comm", None)
            if device_comm is not None
            else None
        )
        if pynccl_comm is None or pynccl_comm.disabled or not pynccl_comm.available:
            raise RuntimeError("EPLB communicator 'pynccl' requested but unavailable.")
        try:
            return PyNcclEplbCommunicator(pynccl_comm=pynccl_comm)
        except Exception as exc:
            raise RuntimeError(
                f"Failed to initialize PyNcclEplbCommunicator ({exc})."
            ) from exc

    is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)
    if is_stateless and backend != "nixl":
        if backend not in ("torch_nccl", "pynccl"):
            raise ValueError(
                f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' "
                f"EPLB communicator (got '{backend}')."
            )
        if backend == "torch_nccl":
            logger.warning(
                "Stateless elastic EP requires PyNCCL backend. "
                "Forcing EPLB communicator to 'pynccl'."
            )
            backend = "pynccl"
        return _create_pynccl()

    if backend == "symm_mem":
        return SymmMemEplbCommunicator(
            torch_group, group_coordinator.cpu_group, first_layer
        )
    elif backend == "nixl":
        if not has_nixl():
            raise RuntimeError(
                "EPLB communicator 'nixl' requested but NIXL is unavailable."
            )
        if not (current_platform.is_cuda_alike() and tensor_device_type != "cpu"):
            raise RuntimeError(
                "EPLB communicator 'nixl' supports only cuda-like devices "
                f"(got {tensor_device_type})."
            )
        assert expert_buffer is not None
        try:
            return NixlEplbCommunicator(
                cpu_group=group_coordinator.cpu_group,
                all_expert_weights=expert_weights,
                expert_buffer=expert_buffer,
            )
        except Exception as exc:
            raise RuntimeError(
                f"Failed to initialize NixlEplbCommunicator ({exc})."
            ) from exc
    elif backend == "torch_gloo":
        return TorchDistGlooStagedEplbCommunicator(
            cpu_group=group_coordinator.cpu_group,
        )
    elif backend == "torch_xccl":
        return TorchDistXCCLStagedEplbCommunicator(ep_group=torch_group)
    elif backend == "torch_nccl":
        return TorchDistNcclEplbCommunicator(ep_group=torch_group)
    elif backend == "pynccl":
        return _create_pynccl()
    raise ValueError(f"Unknown EPLB communicator backend: {backend}")

has_nixl()

Whether the optional NIXL package is available.

Source code in vllm/distributed/eplb/eplb_communicator.py

def has_nixl() -> bool:
    """Whether the optional NIXL package is available."""
    return nixl_utils.NixlWrapper is not None