vllm.distributed.device_communicators.all2all

Classes:

• AgRsAll2AllManager –

  An implementation of all2all communication based on

• DeepEPAll2AllManagerBase –

  All2All communication based on DeepEP High-Throughput kernels.

• DeepEPHTAll2AllManager –

  All2All communication based on DeepEP High-Throughput kernels.

• DeepEPLLAll2AllManager –

  All2All communication based on DeepEP Low-Latency kernels.

• DeepEPV2All2AllManager –

  All2All communication based on DeepEP v2 ElasticBuffer (unified API).

• FlashInferNVLinkOneSidedManager –

  MoE EP communication using FlashInfer's NVLink one-sided wrapper.

• FlashInferNVLinkTwoSidedManager –

  All2All communication based on flashinfer all2allv/two-sided NVLink kernels.

• MoonEPAll2AllManager –

  All2All communication based on MoonEP

• NixlEPAll2AllManager –

  All2All communication based on NIXL EP kernels.

• PassThroughAll2AllManager –

  Placeholder for all2all_backend="passthrough".

AgRsAll2AllManager

Bases: All2AllManagerBase

An implementation of all2all communication based on all-gather (dispatch) and reduce-scatter (combine).

Methods:

• combine –

  Reduce-scatter hidden_states across all dp ranks.

• dispatch –

  Gather hidden_states and router_logits from all dp ranks.

• dispatch_router_logits –

  Gather hidden_states and router_logits from all dp ranks.

Source code in vllm/distributed/device_communicators/all2all.py

class AgRsAll2AllManager(All2AllManagerBase):
    """An implementation of all2all communication based on
    all-gather (dispatch) and reduce-scatter (combine).
    """

    def __init__(self, cpu_group, tcp_store_group=None):
        super().__init__(cpu_group, tcp_store_group)
        self.use_ep = get_current_vllm_config().parallel_config.enable_expert_parallel

    def _get_comm_group(self, is_sequence_parallel: bool) -> Any:
        if is_sequence_parallel:
            return get_ep_group()
        if self.dp_world_size > 1:
            if self.use_ep and get_pcp_group().world_size > 1:
                assert self.tp_group.world_size == 1, (
                    "DP+PCP with TP>1 requires sequence-parallel MoE inputs"
                )
                return get_ep_group()
            return get_dp_group()
        return get_pcp_group()

    def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]:
        if self.dp_world_size == 1:
            return [num_local_tokens] * comm_group.world_size

        dp_metadata = get_forward_context().dp_metadata
        assert dp_metadata is not None
        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
        assert sizes is not None
        assert len(sizes) == comm_group.world_size
        return sizes

    def dispatch_router_logits(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        is_sequence_parallel: bool = False,
        extra_tensors: list[torch.Tensor] | None = None,
    ) -> (
        tuple[torch.Tensor, torch.Tensor]
        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
    ):
        """Gather hidden_states and router_logits from all dp ranks."""
        dist_group = self._get_comm_group(is_sequence_parallel)
        sizes = self._get_sizes(hidden_states.shape[0], dist_group)
        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

        tensors_to_gather = [hidden_states, router_logits]
        if extra_tensors is not None:
            tensors_to_gather.extend(extra_tensors)

        gathered_tensors = dist_group.all_gatherv(
            tensors_to_gather,
            dim=0,
            sizes=sizes,
        )

        if extra_tensors is not None:
            return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
        return gathered_tensors[0], gathered_tensors[1]

    def dispatch(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        is_sequence_parallel: bool = False,
        extra_tensors: list[torch.Tensor] | None = None,
    ) -> (
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
    ):
        """Gather hidden_states and router_logits from all dp ranks."""
        dist_group = self._get_comm_group(is_sequence_parallel)
        sizes = self._get_sizes(hidden_states.shape[0], dist_group)
        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

        tensors_to_gather = [hidden_states, topk_weights, topk_ids]
        if extra_tensors is not None:
            tensors_to_gather.extend(extra_tensors)

        gathered_tensors = dist_group.all_gatherv(
            tensors_to_gather,
            dim=0,
            sizes=sizes,
        )

        hidden_states = gathered_tensors[0]
        topk_weights = gathered_tensors[1]
        topk_ids = gathered_tensors[2]

        if extra_tensors is None:
            return hidden_states, topk_weights, topk_ids

        return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]

    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
        """Reduce-scatter hidden_states across all dp ranks."""
        dist_group = self._get_comm_group(is_sequence_parallel)
        sizes = self._get_sizes(
            hidden_states.shape[0] // dist_group.world_size,
            dist_group,
        )
        hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
        return hidden_states

    def allocate_combine_input(
        self,
        shape: tuple[int, ...],
        dtype: torch.dtype,
        device: torch.device,
        is_sequence_parallel: bool = False,
    ) -> torch.Tensor | None:
        dist_group = self._get_comm_group(is_sequence_parallel)
        sizes = self._get_sizes(shape[0] // dist_group.world_size, dist_group)
        if sum(sizes) != shape[0] or any(size != sizes[0] for size in sizes):
            return None
        device_communicator = dist_group.device_communicator
        if device_communicator is None:
            return None
        return device_communicator.get_symmetric_memory_buffer(
            "moe_ag_rs_combine", shape, dtype, device
        )

    def combine_into_output(
        self,
        hidden_states: torch.Tensor,
        output: torch.Tensor,
        is_sequence_parallel: bool = False,
    ) -> torch.Tensor:
        dist_group = self._get_comm_group(is_sequence_parallel)
        sizes = self._get_sizes(
            hidden_states.shape[0] // dist_group.world_size,
            dist_group,
        )
        return dist_group.reduce_scatterv_into_output(
            hidden_states,
            output,
            dim=0,
            sizes=sizes,
        )

    def destroy(self):
        pass

combine(hidden_states, is_sequence_parallel=False)

Reduce-scatter hidden_states across all dp ranks.

Source code in vllm/distributed/device_communicators/all2all.py

def combine(
    self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
) -> torch.Tensor:
    """Reduce-scatter hidden_states across all dp ranks."""
    dist_group = self._get_comm_group(is_sequence_parallel)
    sizes = self._get_sizes(
        hidden_states.shape[0] // dist_group.world_size,
        dist_group,
    )
    hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
    return hidden_states

dispatch(hidden_states, topk_weights, topk_ids, is_sequence_parallel=False, extra_tensors=None)

Gather hidden_states and router_logits from all dp ranks.

Source code in vllm/distributed/device_communicators/all2all.py

def dispatch(
    self,
    hidden_states: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    is_sequence_parallel: bool = False,
    extra_tensors: list[torch.Tensor] | None = None,
) -> (
    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
    | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
):
    """Gather hidden_states and router_logits from all dp ranks."""
    dist_group = self._get_comm_group(is_sequence_parallel)
    sizes = self._get_sizes(hidden_states.shape[0], dist_group)
    assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

    tensors_to_gather = [hidden_states, topk_weights, topk_ids]
    if extra_tensors is not None:
        tensors_to_gather.extend(extra_tensors)

    gathered_tensors = dist_group.all_gatherv(
        tensors_to_gather,
        dim=0,
        sizes=sizes,
    )

    hidden_states = gathered_tensors[0]
    topk_weights = gathered_tensors[1]
    topk_ids = gathered_tensors[2]

    if extra_tensors is None:
        return hidden_states, topk_weights, topk_ids

    return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]

dispatch_router_logits(hidden_states, router_logits, is_sequence_parallel=False, extra_tensors=None)

Gather hidden_states and router_logits from all dp ranks.

Source code in vllm/distributed/device_communicators/all2all.py

def dispatch_router_logits(
    self,
    hidden_states: torch.Tensor,
    router_logits: torch.Tensor,
    is_sequence_parallel: bool = False,
    extra_tensors: list[torch.Tensor] | None = None,
) -> (
    tuple[torch.Tensor, torch.Tensor]
    | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
):
    """Gather hidden_states and router_logits from all dp ranks."""
    dist_group = self._get_comm_group(is_sequence_parallel)
    sizes = self._get_sizes(hidden_states.shape[0], dist_group)
    assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

    tensors_to_gather = [hidden_states, router_logits]
    if extra_tensors is not None:
        tensors_to_gather.extend(extra_tensors)

    gathered_tensors = dist_group.all_gatherv(
        tensors_to_gather,
        dim=0,
        sizes=sizes,
    )

    if extra_tensors is not None:
        return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
    return gathered_tensors[0], gathered_tensors[1]

DeepEPAll2AllManagerBase

Bases: All2AllManagerBase

All2All communication based on DeepEP High-Throughput kernels.

Source code in vllm/distributed/device_communicators/all2all.py

class DeepEPAll2AllManagerBase(All2AllManagerBase):
    """All2All communication based on DeepEP High-Throughput kernels."""

    def __init__(self, cpu_group, tcp_store_group=None):
        assert has_deep_ep(), (
            "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md"
            " to install DeepEP kernels."
        )  # noqa
        super().__init__(cpu_group, tcp_store_group)
        self.handle_cache = Cache()

        # This is the DeepEP default. Stick to it till we can establish
        # reasonable defaults based on profiling.
        self.num_sms = 20

    def get_handle(self, kwargs):
        raise NotImplementedError

    def dispatch_router_logits(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        is_sequence_parallel: bool = False,
        extra_tensors: list[torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        raise NotImplementedError

    def dispatch(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        is_sequence_parallel: bool = False,
        extra_tensors: list[torch.Tensor] | None = None,
    ) -> (
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
    ):
        raise NotImplementedError

    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
        raise NotImplementedError

    def destroy(self):
        with self.handle_cache._lock:
            for _, handle in self.handle_cache._cache.items():
                handle.destroy()
            self.handle_cache._cache.clear()

DeepEPHTAll2AllManager

Bases: DeepEPAll2AllManagerBase

All2All communication based on DeepEP High-Throughput kernels.

Source code in vllm/distributed/device_communicators/all2all.py

class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
    """All2All communication based on DeepEP High-Throughput kernels."""

    def __init__(self, cpu_group, tcp_store_group=None):
        super().__init__(cpu_group, tcp_store_group)

    def _make_all2all_kwargs(self) -> dict[Any, Any]:
        # Defaults for internode and intranode are taken from DeepEP tests.
        num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
        num_rdma_bytes = None
        num_qps_per_rank = None

        if self.internode and not envs.VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE:
            num_rdma_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
            num_qps_per_rank = self.num_sms // 2
        else:
            num_rdma_bytes = 0
            num_qps_per_rank = 1

        assert num_rdma_bytes is not None
        assert num_qps_per_rank is not None
        # TODO: remove platform-specific logic
        # once ROCm DeepEP is updated with the latest APIs.
        kwargs = dict(
            group=self.cpu_group,
            num_nvl_bytes=num_nvl_bytes,
            num_rdma_bytes=num_rdma_bytes,
            low_latency_mode=False,
            num_qps_per_rank=num_qps_per_rank,
            explicitly_destroy=True,
        )
        return kwargs

    def get_handle(self, kwargs):
        assert len(kwargs) == 0, (
            "DeepEPHTAll2AllManager expects no arguments. All the required "
            "args are computed in the Manager itself."
        )

        import deep_ep  # type: ignore[import-not-found]

        buffer_kwargs = self._make_all2all_kwargs()
        logger.debug("DeepEP all2all args %s", buffer_kwargs)
        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
            buffer_kwargs, deep_ep.Buffer
        )
        return handle

    def set_num_sms(self, num_sms: int):
        import deep_ep  # type: ignore[import-not-found]

        # Right now the buffers are sized for only what the kernels were
        # created with. So we can only reduce the number of SMS used
        # but not increase it.
        if num_sms > self.num_sms:
            num_sms = self.num_sms
        deep_ep.Buffer.set_num_sms(num_sms)

DeepEPLLAll2AllManager

Bases: DeepEPAll2AllManagerBase

All2All communication based on DeepEP Low-Latency kernels.

Methods:

• get_handle –

  The kwargs for DeepEPLLAll2AllManager is dictated by

Source code in vllm/distributed/device_communicators/all2all.py

class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
    """All2All communication based on DeepEP Low-Latency kernels."""

    _buffer: Any = None
    _mask: torch.Tensor | None = None
    _last_mask: torch.Tensor | None = None

    def __init__(self, cpu_group, tcp_store_group=None):
        super().__init__(cpu_group, tcp_store_group)
        self.support_fault_tolerance = (
            get_current_vllm_config().parallel_config.enable_fault_tolerance
        )

    def _make_all2all_kwargs(
        self,
        max_num_tokens_per_dp_rank: int,
        token_hidden_size: int,
        num_ep_ranks: int,
        num_global_experts: int,
        num_local_experts: int,
    ) -> dict[Any, Any]:
        """max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
          can dispatch all the ranks must hold the same value.
        token_hidden_size: the hidden dimension of each token.
        num_ep_ranks: the number of EP group ranks.
        num_global_experts: Number of experts in the model.
        num_local_experts: Number of experts in an EP rank.
        """
        import deep_ep  # type: ignore[import-not-found]

        # Defaults for internode and intranode are taken from DeepEP tests.
        num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
        num_qps_per_rank = num_local_experts
        num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
            num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
            hidden=token_hidden_size,
            num_ranks=num_ep_ranks,
            num_experts=num_global_experts,
        )

        assert num_rdma_bytes is not None
        # TODO: remove platform-specific logic
        # once ROCm DeepEP is updated with the latest APIs.
        kwargs = dict(
            group=self.cpu_group,
            num_nvl_bytes=num_nvl_bytes,
            num_rdma_bytes=num_rdma_bytes,
            low_latency_mode=True,
            num_qps_per_rank=num_qps_per_rank,
            allow_nvlink_for_low_latency_mode=True,
            allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
            explicitly_destroy=True,
            enable_shrink=self.support_fault_tolerance,
        )
        return kwargs

    def get_handle(self, kwargs):
        """The kwargs for DeepEPLLAll2AllManager is dictated by
        _make_all2all_kwargs.
        """
        import deep_ep  # type: ignore[import-not-found]

        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
        logger.debug("DeepEP all2all args %s", buffer_kwargs)
        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
            buffer_kwargs, deep_ep.Buffer
        )
        DeepEPLLAll2AllManager._buffer = handle
        return handle

    # DeepEP LL uses RDMA so no SMs are used for communication
    def max_sms_used(self) -> int | None:
        return 0

    def query_active_mask(self) -> torch.Tensor:
        buf = DeepEPLLAll2AllManager._buffer
        assert buf is not None
        if DeepEPLLAll2AllManager._mask is None:
            DeepEPLLAll2AllManager._mask = torch.zeros(
                self.world_size, device="cuda", dtype=torch.int32
            )
        buf.low_latency_query_mask_buffer(DeepEPLLAll2AllManager._mask)
        return DeepEPLLAll2AllManager._mask

    def query_fault(self) -> torch.Tensor:
        current = self.query_active_mask()
        if DeepEPLLAll2AllManager._last_mask is None:
            DeepEPLLAll2AllManager._last_mask = torch.zeros_like(current)
        has_fault = (current != DeepEPLLAll2AllManager._last_mask).any()
        return has_fault

    def clean_buffers(self) -> None:
        buf = DeepEPLLAll2AllManager._buffer
        if buf is None:
            return
        buf.get_local_buffer_tensor(dtype=torch.int8, use_rdma_buffer=True).zero_()
        torch.accelerator.synchronize()
        buf.low_latency_clean_mask_buffer()
        torch.accelerator.synchronize()
        DeepEPLLAll2AllManager._last_mask = None

    def update_mask(self, rank: int, masked: bool = True) -> None:
        buf = DeepEPLLAll2AllManager._buffer
        assert buf is not None
        buf.low_latency_update_mask_buffer(rank, masked)
        torch.accelerator.synchronize()
        if DeepEPLLAll2AllManager._last_mask is None:
            DeepEPLLAll2AllManager._last_mask = torch.zeros(
                self.world_size, device="cuda", dtype=torch.int32
            )
        DeepEPLLAll2AllManager._last_mask[rank] = int(masked)

_make_all2all_kwargs(max_num_tokens_per_dp_rank, token_hidden_size, num_ep_ranks, num_global_experts, num_local_experts)

the maximum number of tokens a DP rank

can dispatch all the ranks must hold the same value.

token_hidden_size: the hidden dimension of each token. num_ep_ranks: the number of EP group ranks. num_global_experts: Number of experts in the model. num_local_experts: Number of experts in an EP rank.

Source code in vllm/distributed/device_communicators/all2all.py

def _make_all2all_kwargs(
    self,
    max_num_tokens_per_dp_rank: int,
    token_hidden_size: int,
    num_ep_ranks: int,
    num_global_experts: int,
    num_local_experts: int,
) -> dict[Any, Any]:
    """max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
      can dispatch all the ranks must hold the same value.
    token_hidden_size: the hidden dimension of each token.
    num_ep_ranks: the number of EP group ranks.
    num_global_experts: Number of experts in the model.
    num_local_experts: Number of experts in an EP rank.
    """
    import deep_ep  # type: ignore[import-not-found]

    # Defaults for internode and intranode are taken from DeepEP tests.
    num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
    num_qps_per_rank = num_local_experts
    num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
        num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
        hidden=token_hidden_size,
        num_ranks=num_ep_ranks,
        num_experts=num_global_experts,
    )

    assert num_rdma_bytes is not None
    # TODO: remove platform-specific logic
    # once ROCm DeepEP is updated with the latest APIs.
    kwargs = dict(
        group=self.cpu_group,
        num_nvl_bytes=num_nvl_bytes,
        num_rdma_bytes=num_rdma_bytes,
        low_latency_mode=True,
        num_qps_per_rank=num_qps_per_rank,
        allow_nvlink_for_low_latency_mode=True,
        allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
        explicitly_destroy=True,
        enable_shrink=self.support_fault_tolerance,
    )
    return kwargs

get_handle(kwargs)

The kwargs for DeepEPLLAll2AllManager is dictated by _make_all2all_kwargs.

Source code in vllm/distributed/device_communicators/all2all.py

def get_handle(self, kwargs):
    """The kwargs for DeepEPLLAll2AllManager is dictated by
    _make_all2all_kwargs.
    """
    import deep_ep  # type: ignore[import-not-found]

    buffer_kwargs = self._make_all2all_kwargs(**kwargs)
    logger.debug("DeepEP all2all args %s", buffer_kwargs)
    handle: deep_ep.Buffer = self.handle_cache.get_or_create(
        buffer_kwargs, deep_ep.Buffer
    )
    DeepEPLLAll2AllManager._buffer = handle
    return handle

DeepEPV2All2AllManager

Bases: All2AllManagerBase

All2All communication based on DeepEP v2 ElasticBuffer (unified API). Uses NCCL Gin backend with analytical SM calculation.

Source code in vllm/distributed/device_communicators/all2all.py

class DeepEPV2All2AllManager(All2AllManagerBase):
    """All2All communication based on DeepEP v2 ElasticBuffer (unified API).
    Uses NCCL Gin backend with analytical SM calculation.
    """

    def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
        assert has_deep_ep_v2(), deep_ep_v2_unavailable_reason()
        super().__init__(cpu_group, tcp_store_group)
        self._device_group = device_group
        self.handle_cache = Cache()
        self._num_sms: int | None = None
        self._gin_checked = False

    def _make_all2all_kwargs(
        self,
        num_max_tokens_per_rank: int,
        hidden: int,
        num_topk: int,
        use_fp8_dispatch: bool,
    ) -> dict:
        allow_hybrid_mode = envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE
        return dict(
            group=self._device_group
            if self._device_group is not None
            else self.cpu_group,
            num_max_tokens_per_rank=num_max_tokens_per_rank,
            hidden=hidden,
            num_topk=num_topk,
            use_fp8_dispatch=use_fp8_dispatch,
            allow_hybrid_mode=allow_hybrid_mode,
            # Direct-mode automatic selection uses at most nine QPs.
            num_allocated_qps=0 if allow_hybrid_mode else 9,
            prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP,
            allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION),
            explicitly_destroy=True,
        )

    def _check_gin_support(self, group) -> None:
        from vllm.utils.nccl import query_nccl_gin_type

        # ProcessGroupNCCL creates communicators lazily. Initialize this exact
        # group before querying so a null comm pointer is not mistaken for
        # missing GIN support.
        probe = torch.zeros(1, device="cuda")
        torch.distributed.all_reduce(probe, group=group)
        # DeepEPv2 respects EP_DISABLE_GIN, so skip the GIN requirement check.
        if os.environ.get("EP_DISABLE_GIN", "0") != "0":
            return

        gin_type = query_nccl_gin_type(
            group, railed=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE
        )
        if gin_type is None:
            raise RuntimeError(
                "DeepEPv2 communicator properties query failed; "
                "networking capability could not be determined."
            )
        if gin_type == 0:
            raise RuntimeError(
                "DeepEPv2 requires NCCL GIN (GPU-Initiated Networking). "
                "This usually means IBGDA-capable InfiniBand NICs or drivers "
                "are not available. See tools/ep_kernels/README.md for "
                "requirements."
            )

    def get_handle(self, kwargs):
        import deep_ep  # type: ignore[import-not-found]

        num_experts = kwargs.pop("num_experts", 256)
        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
        if not self._gin_checked:
            self._check_gin_support(buffer_kwargs["group"])
            self._gin_checked = True
        logger.debug("DeepEP v2 all2all args %s", buffer_kwargs)
        handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create(
            buffer_kwargs, deep_ep.ElasticBuffer
        )
        if self._num_sms is None:
            self._num_sms = handle.get_theoretical_num_sms(
                num_experts=num_experts,
                num_topk=kwargs["num_topk"],
            )
        return handle

    def max_sms_used(self) -> int | None:
        return self._num_sms

    def destroy(self):
        with self.handle_cache._lock:
            for _, handle in self.handle_cache._cache.items():
                handle.destroy()
            self.handle_cache._cache.clear()

FlashInferNVLinkOneSidedManager

Bases: All2AllManagerBase

MoE EP communication using FlashInfer's NVLink one-sided wrapper.

Methods:

• initialize –

  Share the largest token, dispatch-row and combine-row capacities.

Source code in vllm/distributed/device_communicators/all2all.py

class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
    """MoE EP communication using FlashInfer's NVLink one-sided wrapper."""

    def __init__(self, cpu_group):
        assert has_flashinfer_nvlink_one_sided(), (
            "FlashInfer MoEEpCommunication NVLink one-sided API not found. "
            "Please install FlashInfer 0.7.2 or later."
        )
        super().__init__(cpu_group)
        self.initialized = False
        self.communication: NVLinkOneSidedAlltoAll | None = None
        self.top_k = 0
        self.num_experts = 0

    def initialize(
        self,
        max_num_tokens: int,
        top_k: int,
        num_experts: int,
        hidden_size: int,
        dispatch_format: "QuantFormat | None" = None,
        extra_payload_bytes_per_token: int = 0,
    ) -> None:
        """Share the largest token, dispatch-row and combine-row capacities."""
        from flashinfer.moe_ep import (
            BootstrapConfig,
            MoEEpCommParams,
            NVLinkOneSidedAlltoAll,
            NVLinkOneSidedConfig,
            create_communication,
        )

        if self.initialized:
            assert top_k == self.top_k, (
                "FlashInfer one-sided communication does not support "
                f"heterogeneous top_k (got {top_k}, was built with {self.top_k})"
            )
            assert num_experts == self.num_experts, (
                "FlashInfer one-sided communication does not support "
                f"heterogeneous num_experts (got {num_experts}, "
                f"was built with {self.num_experts})"
            )

        if extra_payload_bytes_per_token < 0:
            raise ValueError("extra_payload_bytes_per_token must be non-negative")
        params = MoEEpCommParams(
            num_experts=num_experts,
            top_k=top_k,
            max_tokens_per_rank=max_num_tokens,
            hidden_size=hidden_size,
            dtype=torch.bfloat16,
            dispatch_format=dispatch_format,
        )
        communication = self.communication
        if communication is not None:
            previous_max = communication.params.max_tokens_per_rank
            previous_extra = communication.config.extra_payload_bytes_per_token
            previous_bytes = (
                communication.params.dispatch_bytes_per_token + previous_extra
            )
            dispatch_bytes = (
                params.dispatch_bytes_per_token + extra_payload_bytes_per_token
            )
            if (
                max_num_tokens <= previous_max
                and hidden_size <= communication.params.hidden_size
                and dispatch_bytes <= previous_bytes
            ):
                return
            # Use the larger hidden size for combine capacity. Its dispatch
            # format may need fewer bytes than another layer's actual payload.
            if hidden_size < communication.params.hidden_size or (
                hidden_size == communication.params.hidden_size
                and dispatch_bytes < previous_bytes
            ):
                params = communication.params
            params = replace(
                params, max_tokens_per_rank=max(max_num_tokens, previous_max)
            )
            extra_payload_bytes_per_token = (
                max(dispatch_bytes, previous_bytes) - params.dispatch_bytes_per_token
            )

        from vllm.distributed.device_communicators.mnnvl_compat import (
            CustomCommunicator,
        )

        # MNNVL handle exchange must cover EP (= DP*PCP*TP), including with PP.
        bootstrap = BootstrapConfig(
            world_size=self.world_size,
            rank=self.rank,
            device=torch.accelerator.current_device_index(),
            process_group=self.cpu_group,
        )
        config = NVLinkOneSidedConfig(
            extra_payload_bytes_per_token=extra_payload_bytes_per_token,
            use_low_precision_combine=envs.VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE,
            comm_backend=CustomCommunicator(self.cpu_group),
            # Keep the fence transport; FlashInfer otherwise picks CFT counted
            # writes or the Lamport transport where the platform supports them.
            cft=False,
            lamport=False,
        )
        if communication is not None:
            # Growth happens during model initialization, before graph capture.
            # Destroy the old wrapper on every rank before replacing it;
            # FlashInfer owns the underlying workspace cache.
            torch.accelerator.synchronize()
            dist.barrier(group=self.cpu_group)
            self.cleanup()
            dist.barrier(group=self.cpu_group)
        torch.accelerator.empty_cache()
        # Bootstrap metadata must stay on the CPU even when model layers are
        # constructed under a CUDA default-device context.
        with torch.device("cpu"):
            replacement = create_communication(bootstrap, params, backend=config)
        assert isinstance(replacement, NVLinkOneSidedAlltoAll)
        self.communication = replacement
        self.top_k = top_k
        self.num_experts = num_experts
        self.initialized = True
        logger.info(
            "FlashInfer One-sided NVLink initialized for rank %s, size %s, hidden %s",
            self.rank,
            self.world_size,
            params.hidden_size,
        )
        dist.barrier(group=self.cpu_group)

    def get_communication(self) -> "NVLinkOneSidedAlltoAll":
        # Resolve at use time: another layer can grow the shared capacity during
        # model construction, replacing the wrapper created for an earlier layer.
        assert self.communication is not None, "One-sided communication not initialized"
        return self.communication

    def get_handle(self, kwargs):
        return self

    def cleanup(self):
        if self.communication is not None:
            self.communication.destroy()
            self.communication = None
        self.initialized = False

    def checkpoint_prepare(self) -> None:
        if self.communication is not None:
            communication = self.communication
            if hasattr(communication, "checkpoint_prepare"):
                communication.checkpoint_prepare()
            else:
                # FlashInfer 0.7.2 exposes checkpoints on the wrapped backend.
                communication.alltoall.checkpoint_prepare()

    def checkpoint_restore(self) -> None:
        if self.initialized:
            from vllm.distributed.device_communicators.mnnvl_compat import (
                CustomCommunicator,
            )

            communicator = CustomCommunicator(self.cpu_group)
            communication = self.get_communication()
            if hasattr(communication, "checkpoint_restore"):
                communication.checkpoint_restore(communicator)
            else:
                communication.alltoall.checkpoint_restore(communicator)

initialize(max_num_tokens, top_k, num_experts, hidden_size, dispatch_format=None, extra_payload_bytes_per_token=0)

Share the largest token, dispatch-row and combine-row capacities.

Source code in vllm/distributed/device_communicators/all2all.py

def initialize(
    self,
    max_num_tokens: int,
    top_k: int,
    num_experts: int,
    hidden_size: int,
    dispatch_format: "QuantFormat | None" = None,
    extra_payload_bytes_per_token: int = 0,
) -> None:
    """Share the largest token, dispatch-row and combine-row capacities."""
    from flashinfer.moe_ep import (
        BootstrapConfig,
        MoEEpCommParams,
        NVLinkOneSidedAlltoAll,
        NVLinkOneSidedConfig,
        create_communication,
    )

    if self.initialized:
        assert top_k == self.top_k, (
            "FlashInfer one-sided communication does not support "
            f"heterogeneous top_k (got {top_k}, was built with {self.top_k})"
        )
        assert num_experts == self.num_experts, (
            "FlashInfer one-sided communication does not support "
            f"heterogeneous num_experts (got {num_experts}, "
            f"was built with {self.num_experts})"
        )

    if extra_payload_bytes_per_token < 0:
        raise ValueError("extra_payload_bytes_per_token must be non-negative")
    params = MoEEpCommParams(
        num_experts=num_experts,
        top_k=top_k,
        max_tokens_per_rank=max_num_tokens,
        hidden_size=hidden_size,
        dtype=torch.bfloat16,
        dispatch_format=dispatch_format,
    )
    communication = self.communication
    if communication is not None:
        previous_max = communication.params.max_tokens_per_rank
        previous_extra = communication.config.extra_payload_bytes_per_token
        previous_bytes = (
            communication.params.dispatch_bytes_per_token + previous_extra
        )
        dispatch_bytes = (
            params.dispatch_bytes_per_token + extra_payload_bytes_per_token
        )
        if (
            max_num_tokens <= previous_max
            and hidden_size <= communication.params.hidden_size
            and dispatch_bytes <= previous_bytes
        ):
            return
        # Use the larger hidden size for combine capacity. Its dispatch
        # format may need fewer bytes than another layer's actual payload.
        if hidden_size < communication.params.hidden_size or (
            hidden_size == communication.params.hidden_size
            and dispatch_bytes < previous_bytes
        ):
            params = communication.params
        params = replace(
            params, max_tokens_per_rank=max(max_num_tokens, previous_max)
        )
        extra_payload_bytes_per_token = (
            max(dispatch_bytes, previous_bytes) - params.dispatch_bytes_per_token
        )

    from vllm.distributed.device_communicators.mnnvl_compat import (
        CustomCommunicator,
    )

    # MNNVL handle exchange must cover EP (= DP*PCP*TP), including with PP.
    bootstrap = BootstrapConfig(
        world_size=self.world_size,
        rank=self.rank,
        device=torch.accelerator.current_device_index(),
        process_group=self.cpu_group,
    )
    config = NVLinkOneSidedConfig(
        extra_payload_bytes_per_token=extra_payload_bytes_per_token,
        use_low_precision_combine=envs.VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE,
        comm_backend=CustomCommunicator(self.cpu_group),
        # Keep the fence transport; FlashInfer otherwise picks CFT counted
        # writes or the Lamport transport where the platform supports them.
        cft=False,
        lamport=False,
    )
    if communication is not None:
        # Growth happens during model initialization, before graph capture.
        # Destroy the old wrapper on every rank before replacing it;
        # FlashInfer owns the underlying workspace cache.
        torch.accelerator.synchronize()
        dist.barrier(group=self.cpu_group)
        self.cleanup()
        dist.barrier(group=self.cpu_group)
    torch.accelerator.empty_cache()
    # Bootstrap metadata must stay on the CPU even when model layers are
    # constructed under a CUDA default-device context.
    with torch.device("cpu"):
        replacement = create_communication(bootstrap, params, backend=config)
    assert isinstance(replacement, NVLinkOneSidedAlltoAll)
    self.communication = replacement
    self.top_k = top_k
    self.num_experts = num_experts
    self.initialized = True
    logger.info(
        "FlashInfer One-sided NVLink initialized for rank %s, size %s, hidden %s",
        self.rank,
        self.world_size,
        params.hidden_size,
    )
    dist.barrier(group=self.cpu_group)

FlashInferNVLinkTwoSidedManager

Bases: All2AllManagerBase

All2All communication based on flashinfer all2allv/two-sided NVLink kernels.

Methods:

• get_communication –

  Return the communication for these sizes, creating it on first use.

Source code in vllm/distributed/device_communicators/all2all.py

class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
    """All2All communication based on flashinfer all2allv/two-sided NVLink kernels."""

    # This type lint could be removed after all of the work in
    # https://github.com/vllm-project/vllm/issues/26533 done.
    rank: int
    world_size: int

    def __init__(self, cpu_group, tcp_store_group=None):
        assert has_flashinfer_nvlink_two_sided(), (
            "flashinfer all2all module not found. Please install/check flashinfer"
        )  # noqa
        super().__init__(cpu_group, tcp_store_group)
        logger.debug(
            "Initialize for flashinfer All2All rank=%d, world size=%d",
            self.rank,
            self.world_size,
        )
        self.communications: dict[tuple[int, ...], NVLinkTwoSidedAlltoAll] = {}

    def get_communication(
        self,
        max_num_tokens: int,
        top_k: int,
        num_experts: int,
        hidden_size: int,
    ) -> "NVLinkTwoSidedAlltoAll":
        """Return the communication for these sizes, creating it on first use.

        Creation is collective over the EP group. FlashInfer keeps one
        workspace per process and shares it across communications.
        """
        key = (max_num_tokens, top_k, num_experts, hidden_size)
        communication = self.communications.get(key)
        if communication is not None:
            return communication

        from flashinfer.moe_ep import (
            BootstrapConfig,
            MoEEpCommParams,
            NVLinkTwoSidedAlltoAll,
            NVLinkTwoSidedConfig,
            create_communication,
        )

        from vllm.distributed.device_communicators.mnnvl_compat import (
            CustomCommunicator,
        )

        # MNNVL workspace is allocated per rank in the comm_backend's group; the
        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
        # must span the EP group (= DP*PCP*TP), not the DP group.
        bootstrap = BootstrapConfig(
            world_size=self.world_size,
            rank=self.rank,
            device=torch.accelerator.current_device_index(),
            process_group=self.cpu_group,
        )
        params = MoEEpCommParams(
            num_experts=num_experts,
            top_k=top_k,
            max_tokens_per_rank=max_num_tokens,
            hidden_size=hidden_size,
            # Received rows past a source rank's token count keep expert id
            # num_experts, which no rank owns.
            invalid_expert_id=num_experts,
        )
        config = NVLinkTwoSidedConfig(comm_backend=CustomCommunicator(self.cpu_group))
        # Bootstrap metadata must stay on the CPU even when model layers are
        # constructed under a CUDA default-device context.
        with torch.device("cpu"):
            communication = create_communication(bootstrap, params, backend=config)
        assert isinstance(communication, NVLinkTwoSidedAlltoAll)
        self.communications[key] = communication
        logger.info(
            "FlashInfer All2All initialized for rank %s, size %s",
            self.rank,
            self.world_size,
        )
        return communication

    def get_handle(self, kwargs):
        return self

    def cleanup(self):
        for communication in self.communications.values():
            communication.destroy()
        self.communications.clear()

get_communication(max_num_tokens, top_k, num_experts, hidden_size)

Return the communication for these sizes, creating it on first use.

Creation is collective over the EP group. FlashInfer keeps one workspace per process and shares it across communications.

Source code in vllm/distributed/device_communicators/all2all.py

def get_communication(
    self,
    max_num_tokens: int,
    top_k: int,
    num_experts: int,
    hidden_size: int,
) -> "NVLinkTwoSidedAlltoAll":
    """Return the communication for these sizes, creating it on first use.

    Creation is collective over the EP group. FlashInfer keeps one
    workspace per process and shares it across communications.
    """
    key = (max_num_tokens, top_k, num_experts, hidden_size)
    communication = self.communications.get(key)
    if communication is not None:
        return communication

    from flashinfer.moe_ep import (
        BootstrapConfig,
        MoEEpCommParams,
        NVLinkTwoSidedAlltoAll,
        NVLinkTwoSidedConfig,
        create_communication,
    )

    from vllm.distributed.device_communicators.mnnvl_compat import (
        CustomCommunicator,
    )

    # MNNVL workspace is allocated per rank in the comm_backend's group; the
    # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
    # must span the EP group (= DP*PCP*TP), not the DP group.
    bootstrap = BootstrapConfig(
        world_size=self.world_size,
        rank=self.rank,
        device=torch.accelerator.current_device_index(),
        process_group=self.cpu_group,
    )
    params = MoEEpCommParams(
        num_experts=num_experts,
        top_k=top_k,
        max_tokens_per_rank=max_num_tokens,
        hidden_size=hidden_size,
        # Received rows past a source rank's token count keep expert id
        # num_experts, which no rank owns.
        invalid_expert_id=num_experts,
    )
    config = NVLinkTwoSidedConfig(comm_backend=CustomCommunicator(self.cpu_group))
    # Bootstrap metadata must stay on the CPU even when model layers are
    # constructed under a CUDA default-device context.
    with torch.device("cpu"):
        communication = create_communication(bootstrap, params, backend=config)
    assert isinstance(communication, NVLinkTwoSidedAlltoAll)
    self.communications[key] = communication
    logger.info(
        "FlashInfer All2All initialized for rank %s, size %s",
        self.rank,
        self.world_size,
    )
    return communication

MoonEPAll2AllManager

Bases: All2AllManagerBase

All2All communication based on MoonEP (https://github.com/MoonshotAI/MoonEP).

MoonEP keeps token loads perfectly balanced across EP ranks by planning a small number of dynamically redundant experts online and prefetching their weights before expert compute. Every rank receives exactly S x K token slots regardless of router skew, so all communication and compute shapes are static.

Requires NVLink symmetric-memory / multicast capable topologies (single node NVSwitch, e.g. H100/H200/B200/GB300 class).

Source code in vllm/distributed/device_communicators/all2all.py

class MoonEPAll2AllManager(All2AllManagerBase):
    """All2All communication based on MoonEP
    (https://github.com/MoonshotAI/MoonEP).

    MoonEP keeps token loads perfectly balanced across EP ranks by planning
    a small number of dynamically redundant experts online and prefetching
    their weights before expert compute. Every rank receives exactly
    S x K token slots regardless of router skew, so all communication and
    compute shapes are static.

    Requires NVLink symmetric-memory / multicast capable topologies
    (single node NVSwitch, e.g. H100/H200/B200/GB300 class).
    """

    def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
        check_moonep_system_support()
        super().__init__(cpu_group, tcp_store_group)
        self._device_group = device_group
        self.handle_cache = Cache()

    def _make_buffer_kwargs(
        self,
        max_num_tokens_per_dp_rank: int,
        token_hidden_size: int,
        num_topk: int,
        num_global_experts: int,
        num_prefetch_slots: int,
        token_padding: int,
        num_sms: int,
    ) -> dict:
        return dict(
            S=max_num_tokens_per_dp_rank,
            H=token_hidden_size,
            K=num_topk,
            E=num_global_experts,
            num_ep_ranks=self.world_size,
            num_sms=num_sms,
            token_padding=token_padding,
            B=num_prefetch_slots,
            group=self._device_group
            if self._device_group is not None
            else self.cpu_group,
            explicitly_destroy=True,
        )

    def get_handle(self, kwargs):
        from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import (
            MoonEPBufferPool,
        )

        buffer_kwargs = self._make_buffer_kwargs(**kwargs)
        logger.debug("MoonEP all2all args %s", buffer_kwargs)

        def make_pool(**kw):
            return MoonEPBufferPool(kw, max_tokens_per_rank=kw["S"])

        handle: MoonEPBufferPool = self.handle_cache.get_or_create(
            buffer_kwargs, make_pool
        )
        return handle

    def destroy(self):
        with self.handle_cache._lock:
            for _, handle in self.handle_cache._cache.items():
                handle.destroy()
            self.handle_cache._cache.clear()

NixlEPAll2AllManager

Bases: All2AllManagerBase

All2All communication based on NIXL EP kernels. This backend supports elastic EP with dynamic rank connection/disconnection.

Methods:

• commit_ep_size –

  Commit the staged EP size to the active communication set.

Source code in vllm/distributed/device_communicators/all2all.py

class NixlEPAll2AllManager(All2AllManagerBase):
    """All2All communication based on NIXL EP kernels.
    This backend supports elastic EP with dynamic rank connection/disconnection.
    """

    _buffer: _NixlEPBufferState | None = None
    _lock = threading.RLock()
    _mask: torch.Tensor | None = None
    _last_mask: torch.Tensor | None = None

    def __init__(self, cpu_group, tcp_store_group=None):
        if tcp_store_group is None:
            tcp_store_group = StatelessProcessGroup(
                rank=cpu_group.rank(),
                world_size=cpu_group.size(),
                store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()),
            )
        super().__init__(cpu_group, tcp_store_group)
        self.support_fault_tolerance = True

    @property
    def max_num_ep_ranks(self) -> int:
        assert NixlEPAll2AllManager._buffer is not None
        return NixlEPAll2AllManager._buffer.max_num_ep_ranks

    def _init_buffer(
        self,
        max_num_tokens_per_dp_rank: int,
        token_hidden_size: int,
        num_experts_per_rank: int,
    ) -> None:
        from nixl_ep import Buffer  # type: ignore[import-not-found]

        parallel_config = get_current_vllm_config().parallel_config
        max_num_ep_ranks = (
            self.world_size
            * parallel_config.elastic_ep_max_dp_size
            // parallel_config.data_parallel_size
            if parallel_config.enable_elastic_ep
            else self.world_size
        )
        max_num_global_experts = max_num_ep_ranks * num_experts_per_rank
        num_rdma_bytes = Buffer.get_rdma_size_hint(
            num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
            hidden=token_hidden_size,
            num_ranks=max_num_ep_ranks,
            num_experts=max_num_global_experts,
        )
        assert NixlEPAll2AllManager._buffer is None, (
            "NIXL EP buffer already initialized"
        )
        buffer = Buffer(
            rank=self.rank,
            tcp_store_group=self.tcp_store_group.store,
        )
        buffer.update_memory_buffers(
            num_ranks=max_num_ep_ranks,
            num_experts_per_rank=num_experts_per_rank,
            num_rdma_bytes=num_rdma_bytes,
        )
        ranks_to_connect = list(range(self.world_size))
        buffer.connect_ranks(ranks_to_connect)
        NixlEPAll2AllManager._buffer = _NixlEPBufferState(
            buffer=buffer,
            max_num_ep_ranks=max_num_ep_ranks,
            connected_ep_size=self.world_size,
            active_ep_size=self.world_size,
        )

    def _connect_to_ep_size(self, ep_size: int, *, make_active: bool) -> None:
        assert NixlEPAll2AllManager._buffer is not None
        state = NixlEPAll2AllManager._buffer
        if ep_size <= state.connected_ep_size:
            return

        state.buffer.set_tcp_store_group(self.tcp_store_group.store)
        ranks_to_connect = list(range(state.connected_ep_size, ep_size))
        state.buffer.connect_ranks(ranks_to_connect, activate=make_active)
        state.connected_ep_size = ep_size
        if make_active:
            state.active_ep_size = ep_size

    def _disconnect_to_ep_size(self, ep_size: int) -> None:
        assert NixlEPAll2AllManager._buffer is not None
        state = NixlEPAll2AllManager._buffer
        if ep_size >= state.connected_ep_size:
            return

        state.buffer.set_tcp_store_group(self.tcp_store_group.store)
        ranks_to_disconnect = list(range(ep_size, state.connected_ep_size))
        state.buffer.disconnect_ranks(ranks_to_disconnect)
        state.connected_ep_size = ep_size
        state.active_ep_size = min(state.active_ep_size, ep_size)

    def _unmask_connected_ranks(self, target_ep_size: int) -> None:
        assert NixlEPAll2AllManager._buffer is not None
        state = NixlEPAll2AllManager._buffer
        state.buffer.set_tcp_store_group(self.tcp_store_group.store)
        if target_ep_size <= state.active_ep_size:
            return
        assert state.connected_ep_size >= target_ep_size

        for rank in range(state.active_ep_size, target_ep_size):
            state.buffer.update_mask_buffer(rank, mask=False)
        torch.accelerator.synchronize()
        state.active_ep_size = target_ep_size

    def _stage_ep_size(self) -> None:
        assert NixlEPAll2AllManager._buffer is not None
        state = NixlEPAll2AllManager._buffer
        target_ep_size = self.world_size

        # Scale-up can safely connect standby ranks while leaving them masked.
        # Scale-down must not disconnect active ranks until commit.
        if target_ep_size > state.connected_ep_size:
            self._connect_to_ep_size(target_ep_size, make_active=False)

    def stage_ep_size(self) -> None:
        with NixlEPAll2AllManager._lock:
            self._stage_ep_size()

    def commit_ep_size(self) -> None:
        """Commit the staged EP size to the active communication set."""
        with NixlEPAll2AllManager._lock:
            assert NixlEPAll2AllManager._buffer is not None
            state = NixlEPAll2AllManager._buffer
            target_ep_size = self.world_size

            if target_ep_size < state.connected_ep_size:
                self._disconnect_to_ep_size(target_ep_size)
            elif target_ep_size > state.connected_ep_size:
                self._connect_to_ep_size(target_ep_size, make_active=True)

            self._unmask_connected_ranks(target_ep_size)

    @contextmanager
    def mask_remote_ranks(self) -> Iterator[None]:
        peers = [rank for rank in range(self.world_size) if rank != self.rank]
        assert NixlEPAll2AllManager._buffer is not None
        buffer = NixlEPAll2AllManager._buffer.buffer
        for rank in peers:
            buffer.update_mask_buffer(rank, mask=True)
        torch.accelerator.synchronize()
        try:
            yield
        finally:
            for rank in peers:
                buffer.update_mask_buffer(rank, mask=False)
            torch.accelerator.synchronize()

    def get_handle(self, kwargs):
        with NixlEPAll2AllManager._lock:
            state = NixlEPAll2AllManager._buffer
            if state is None:
                max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"]
                self._init_buffer(
                    max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank,
                    token_hidden_size=kwargs["token_hidden_size"],
                    num_experts_per_rank=kwargs["num_local_experts"],
                )

            assert NixlEPAll2AllManager._buffer is not None
            handle = NixlEPAll2AllManager._buffer.buffer
            return handle

    def dispatch(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        is_sequence_parallel: bool = False,
        extra_tensors: list[torch.Tensor] | None = None,
    ) -> (
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
    ):
        raise NotImplementedError

    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
        raise NotImplementedError

    def destroy(self):
        # NOTE(yongji): NIXLEPAll2AllManager instance is recreated during
        # scale-up/down, so we cannot destroy the persistent buffer here.
        assert NixlEPAll2AllManager._buffer is not None
        buffer = NixlEPAll2AllManager._buffer.buffer
        buffer.set_tcp_store_group(None)

    # NIXL EP uses RDMA so no SMs are used for communication
    def max_sms_used(self) -> int | None:
        return 0

    def query_active_mask(self) -> torch.Tensor:
        state = NixlEPAll2AllManager._buffer
        assert state is not None
        if NixlEPAll2AllManager._mask is None:
            NixlEPAll2AllManager._mask = torch.zeros(
                state.max_num_ep_ranks, device="cuda", dtype=torch.int32
            )
        state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask)
        return NixlEPAll2AllManager._mask[: state.active_ep_size]

    def query_fault(self) -> torch.Tensor:
        current = self.query_active_mask()
        last = NixlEPAll2AllManager._last_mask
        if last is None or last.shape != current.shape:
            NixlEPAll2AllManager._last_mask = torch.zeros_like(current)
            last = NixlEPAll2AllManager._last_mask
        has_fault = (current != last).any()
        return has_fault

    def clean_buffers(self) -> None:
        if NixlEPAll2AllManager._buffer is None:
            return
        state = NixlEPAll2AllManager._buffer
        state.buffer.get_local_buffer_tensor(
            dtype=torch.int8, use_rdma_buffer=True
        ).zero_()
        torch.accelerator.synchronize()
        state.buffer.clean_mask_buffer()
        torch.accelerator.synchronize()
        NixlEPAll2AllManager._last_mask = None

    def update_mask(self, rank: int, masked: bool = True) -> None:
        state = NixlEPAll2AllManager._buffer
        assert state is not None
        state.buffer.update_mask_buffer(rank, mask=masked)
        torch.accelerator.synchronize()
        last = NixlEPAll2AllManager._last_mask
        if last is None or last.shape[0] != state.active_ep_size:
            NixlEPAll2AllManager._last_mask = torch.zeros(
                state.active_ep_size, device="cuda", dtype=torch.int32
            )
            last = NixlEPAll2AllManager._last_mask
        last[rank] = int(masked)

commit_ep_size()

Commit the staged EP size to the active communication set.

Source code in vllm/distributed/device_communicators/all2all.py

def commit_ep_size(self) -> None:
    """Commit the staged EP size to the active communication set."""
    with NixlEPAll2AllManager._lock:
        assert NixlEPAll2AllManager._buffer is not None
        state = NixlEPAll2AllManager._buffer
        target_ep_size = self.world_size

        if target_ep_size < state.connected_ep_size:
            self._disconnect_to_ep_size(target_ep_size)
        elif target_ep_size > state.connected_ep_size:
            self._connect_to_ep_size(target_ep_size, make_active=True)

        self._unmask_connected_ranks(target_ep_size)

PassThroughAll2AllManager

Bases: All2AllManagerBase

Placeholder for all2all_backend="passthrough".

The MoE backend dispatches and combines itself, so there is no all2all to manage.

Source code in vllm/distributed/device_communicators/all2all.py

class PassThroughAll2AllManager(All2AllManagerBase):
    """Placeholder for ``all2all_backend="passthrough"``.

    The MoE backend dispatches and combines itself, so there is no all2all to
    manage.
    """

    def get_handle(self, kwargs):
        raise RuntimeError(
            "passthrough has no all2all handle: dispatch and combine run "
            "inside the MoE backend."
        )

_mori_epv2_arena_bytes(cfg)

Upper bound of a MORI EPv2 gather-mode op's symmetric arena.

Mirrors the intranode regions of MORI's hip_backend: per received row disp_out, out_tok, out_scales, out_idx + out_wts (topk * 4 each) and recv_to_src_token (4); per rank recv_num (8 since MORI 1.2.4) and cross_device_barrier (8); and tok_off (4). MORI allocates the arena inside the window itself, so an underestimate fails there rather than overlapping.

Source code in vllm/distributed/device_communicators/all2all.py

def _mori_epv2_arena_bytes(cfg) -> int:
    """Upper bound of a MORI EPv2 gather-mode op's symmetric arena.

    Mirrors the intranode regions of MORI's hip_backend: per received row
    disp_out, out_tok, out_scales, out_idx + out_wts (topk * 4 each) and
    recv_to_src_token (4); per rank recv_num (8 since MORI 1.2.4) and
    cross_device_barrier (8); and tok_off (4). MORI allocates the arena inside
    the window itself, so an underestimate fails there rather than overlapping.
    """
    # The hip backend pads scale rows to 128 B; FlyDSL uses the unpadded row.
    scale_row = -(-cfg.scale_dim * cfg.scale_type_size // 128) * 128
    topk = cfg.num_experts_per_token
    per_recv_row = cfg.token_nbytes + cfg.combine_token_nbytes + scale_row
    per_recv_row += topk * 8 + 4
    per_rank = cfg.world_size * 16
    # The arena aligns each of its (at most 12) regions to 256 B.
    return cfg.effective_max_recv * per_recv_row + per_rank + 4 + 12 * 256