| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from enum import IntEnum |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.envs as envs |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.model_executor.layers.fused_moe.config import ( |
| 11 | FusedMoEConfig, |
| 12 | ) |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.torch_utils import ( |
| 15 | aux_stream, |
| 16 | current_stream, |
| 17 | ) |
| 18 | from vllm.v1.worker.ubatching import ( |
| 19 | dbo_current_ubatch_id, |
| 20 | ) |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | |
| 25 | class SharedExpertsOrder(IntEnum): |
| 26 | # No shared experts. |
| 27 | NONE = (0,) |
| 28 | |
| 29 | # No overlap - defensively called before MK. |
| 30 | NO_OVERLAP = (1,) |
| 31 | |
| 32 | # Overlapped with dispatch/combine in DP/EP - called by the MK. |
| 33 | MK_INTERNAL_OVERLAPPED = (2,) |
| 34 | |
| 35 | # Overlapped with the gate, router, experts in aux stream. |
| 36 | MULTI_STREAM_OVERLAPPED = (3,) |
| 37 | |
| 38 | |
| 39 | class SharedExperts(torch.nn.Module): |
| 40 | def __init__( |
| 41 | self, |
| 42 | layer: torch.nn.Module, |
| 43 | moe_config: FusedMoEConfig, |
| 44 | enable_dbo: bool, |
| 45 | mk_can_overlap_shared_experts: Callable[[], bool], |
| 46 | ): |
| 47 | super().__init__() |
| 48 | |
| 49 | # The SharedExperts need to handle DBO since they can be called from |
| 50 | # an MK's finalize method. We keep a list of outputs indexed by current |
| 51 | # DBO ubatch id to handle this case. If DBO is not enabled, the |
| 52 | # index is always 0 and the second output list element is ignored. |
| 53 | self.enable_dbo = enable_dbo |
| 54 | self._output: list[torch.Tensor | None] = [None, None] |
| 55 | self._layer = layer |
| 56 | self._moe_config = moe_config |
| 57 | |
| 58 | self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts |
| 59 | |
| 60 | # Allow disabling of the separate shared experts stream for |
| 61 | # debug purposes. |
| 62 | # TODO: Remove this after more extensive testings with TP/DP |
| 63 | # and other execution modes |
| 64 | if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: |
| 65 | logger.debug_once("Disabling MoE shared_experts cuda stream") |
| 66 | self._stream = None |
| 67 | else: |
| 68 | self._stream = aux_stream() |
| 69 | if self._stream is not None: |
| 70 | logger.debug_once("Enabled separate cuda stream for MoE shared_experts") |
| 71 | # One pair per DBO ubatch id to sync aux and main stream. |
| 72 | self._input_ready_event = [torch.cuda.Event(), torch.cuda.Event()] |
| 73 | self._output_ready_event = [torch.cuda.Event(), torch.cuda.Event()] |
| 74 | |
| 75 | # TODO(bnell): Hack for elastic_ep. Get rid of this |
| 76 | def _set_moe_config(self, new_moe_config: FusedMoEConfig): |
| 77 | self.moe_config = new_moe_config |
| 78 | |
| 79 | @property |
| 80 | def _disable_shared_experts_overlap(self) -> bool: |
| 81 | # Disable shared expert overlap if: |
| 82 | # - we are using eplb with non-safe backend, because of correctness issues |
| 83 | # - we are using flashinfer with DP, since there nothing to gain |
| 84 | |
| 85 | # Both these comm backends have been shown to be safe for shared expert overlap. |
| 86 | _EPLB_OVERLAP_SAFE_BACKENDS = ( |
| 87 | "allgather_reducescatter", |
| 88 | "deepep_v2", |
| 89 | "flashinfer_nvlink_one_sided", |
| 90 | ) |
| 91 | |
| 92 | parallel_config = self._moe_config.moe_parallel_config |
| 93 | if getattr(self._layer, "shard_sequence_parallel", False): |
| 94 | # TODO: we may enable this to optimize further |
| 95 | return True |
| 96 | return ( |
| 97 | parallel_config.enable_eplb |
| 98 | and parallel_config.all2all_backend not in _EPLB_OVERLAP_SAFE_BACKENDS |
| 99 | ) or parallel_config.use_fi_nvl_two_sided_kernels |
| 100 | |
| 101 | def _determine_shared_experts_order( |
| 102 | self, |
| 103 | hidden_states: torch.Tensor, |
| 104 | ) -> SharedExpertsOrder: |
| 105 | if self._disable_shared_experts_overlap: |
| 106 | return SharedExpertsOrder.NO_OVERLAP |
| 107 | |
| 108 | if self._mk_can_overlap_shared_experts(): |
| 109 | return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED |
| 110 | |
| 111 | should_run_shared_in_aux_stream = ( |
| 112 | current_platform.is_cuda_alike() |
| 113 | and self._stream is not None |
| 114 | and hidden_states.shape[0] |
| 115 | <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD |
| 116 | ) |
| 117 | |
| 118 | if should_run_shared_in_aux_stream: |
| 119 | return SharedExpertsOrder.MULTI_STREAM_OVERLAPPED |
| 120 | else: |
| 121 | return SharedExpertsOrder.NO_OVERLAP |
| 122 | |
| 123 | def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool: |
| 124 | """Enqueue shared experts on the aux stream without waiting for them. |
| 125 | |
| 126 | Returns true if the shared experts were enqueued, false otherwise. Call |
| 127 | `wait` to wait for the shared experts to finish if this returns true. |
| 128 | """ |
| 129 | if ( |
| 130 | self._determine_shared_experts_order(shared_experts_input) |
| 131 | != SharedExpertsOrder.MULTI_STREAM_OVERLAPPED |
| 132 | ): |
| 133 | return False |
| 134 | assert self._stream is not None |
| 135 | idx = self._output_idx |
| 136 | assert self._output[idx] is None |
| 137 | self._input_ready_event[idx].record(current_stream()) |
| 138 | with torch.cuda.stream(self._stream): |
| 139 | self._input_ready_event[idx].wait(self._stream) |
| 140 | self._output[idx] = self._layer(shared_experts_input) |
| 141 | self._output_ready_event[idx].record(self._stream) |
| 142 | return True |
| 143 | |
| 144 | def wait(self) -> None: |
| 145 | """Block the main stream until `maybe_forward_async` output is ready.""" |
| 146 | assert self._stream is not None |
| 147 | self._output_ready_event[self._output_idx].wait(current_stream()) |
| 148 | |
| 149 | @property |
| 150 | def _output_idx(self) -> int: |
| 151 | return dbo_current_ubatch_id() if self.enable_dbo else 0 |
| 152 | |
| 153 | @property |
| 154 | def output(self) -> torch.Tensor: |
| 155 | assert self._output[self._output_idx] is not None |
| 156 | output = self._output[self._output_idx] |
| 157 | self._output[self._output_idx] = None |
| 158 | return output |
| 159 | |
| 160 | def forward( |
| 161 | self, |
| 162 | shared_experts_input: torch.Tensor, |
| 163 | order: SharedExpertsOrder, |
| 164 | ): |
| 165 | experts_order = self._determine_shared_experts_order(shared_experts_input) |
| 166 | |
| 167 | if order != experts_order: |
| 168 | return None |
| 169 | |
| 170 | assert self._output[self._output_idx] is None |
| 171 | |
| 172 | self._output[self._output_idx] = self._layer(shared_experts_input) |
| 173 | |
| 174 | assert self._output[self._output_idx] is not None |