@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 93bada494f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 12:15 UTC

[MoE Refactor] Split of DefaultMoERunner class (#35326)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed Apr 6, 2026 at 12:41 UTC 93bada494f78b274867772ff337a3e3fb15976d6
8 files changed +868 -705
vllm/model_executor/layers/fused_moe/layer.py
+4 -4
@@ -39,8 +39,8 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
39 from vllm.model_executor.layers.fused_moe.router.router_factory import (
40 create_fused_moe_router,
41 )
42 -from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (
43 - DefaultMoERunner,
42 +from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import (
43 + create_moe_runner,
44 )
45 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
46 SharedExperts,
@@ -572,8 +572,8 @@ class FusedMoE(CustomOp):
572 # Storing the runner in the FusedMoE is an intermediate state, eventually
573 # the runner will own the FusedMoE layer and provide the execution interface
574 # for MoE ops.
575 - self.runner = DefaultMoERunner(
576 - layer=self,
575 + self.runner = create_moe_runner(
576 + layer_name=self.layer_name,
577 moe_config=self.moe_config,
578 router=self.router,
579 routed_input_transform=self._routed_input_transform,
vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py new
+243
@@ -0,0 +1,243 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.forward_context import (
7 + get_forward_context,
8 +)
9 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
10 + FusedMoEMethodBase,
11 +)
12 +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase
13 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
14 + SharedExperts,
15 +)
16 +from vllm.utils.math_utils import cdiv
17 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id
18 +from vllm.v1.worker.workspace import current_workspace_manager
19 +
20 +
21 +class ChunkingMoERunner(MoERunnerBase):
22 + """
23 + MoE runner wrapper that adds chunked processing to any MoERunnerBase.
24 +
25 + This runner wraps an inner MoERunnerBase and overrides _forward_impl to
26 + process large batches by breaking them into smaller chunks. Each chunk
27 + is delegated to the inner runner's _forward_impl, making chunking
28 + composable with any runner implementation.
29 +
30 + All MoERunnerBase state (moe_config, router, quant_method, etc.) is
31 + transparently delegated to the inner runner via __getattr__.
32 + ChunkingMoERunner only owns chunking-specific state: the pre-allocated
33 + workspace buffers and the reduce_results override.
34 +
35 + Key behaviors:
36 + - Pre-allocates workspace tensors for CUDA graph compatibility
37 + - Processes chunks via inner._forward_impl per chunk
38 + - Never reduces results (reduce_results always returns False)
39 + """
40 +
41 + def __init__(self, inner: MoERunnerBase):
42 + # Assert that _maybe_dispatch/_maybe_combine will be nops.
43 + assert inner.moe_config.pcp_size == 1
44 +
45 + # Skip MoERunnerBase.__init__ — all state is delegated to inner
46 + # via __getattr__. Only chunking-specific state lives here.
47 + self._inner = inner
48 +
49 + # Pre-allocated staging buffers. These need to exist ahead of time
50 + # due to CUDA graph construction needing fixed buffer addresses.
51 + self.batched_hidden_states, self.batched_router_logits = (
52 + self._init_dp_chunking()
53 + )
54 +
55 + def __getattr__(self, name):
56 + # Delegate attribute access to the inner runner. This is only
57 + # called when normal lookup (instance __dict__, class MRO) fails,
58 + # so ChunkingMoERunner's own attributes and methods take priority.
59 + return getattr(self._inner, name)
60 +
61 + @property
62 + def shared_experts(self) -> SharedExperts | None:
63 + return self._inner.shared_experts
64 +
65 + # TODO(bnell): temporary hack, do not call this method.
66 + def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
67 + self._inner._replace_quant_method(quant_method)
68 + self.quant_method = quant_method
69 +
70 + def is_internal_router(self) -> bool:
71 + return self._inner.gate is not None
72 +
73 + # Reducing results when chunking is handled by the MK finalize operations
74 + # when DP chunking is enabled..
75 + # This will be removed by #35949
76 + @property
77 + def reduce_results(self) -> bool:
78 + return False
79 +
80 + def _init_dp_chunking(self) -> list[torch.Tensor]:
81 + states_shape: tuple[int, ...]
82 + logits_shape: tuple[int, ...]
83 +
84 + moe = self.moe_config
85 +
86 + if self.enable_dbo:
87 + states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim)
88 + logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts)
89 + else:
90 + states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)
91 + logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)
92 +
93 + # Does this need some kind of profiling run check like modular_kernel.py?
94 + return current_workspace_manager().get_simultaneous(
95 + (states_shape, moe.in_dtype),
96 + (logits_shape, moe.router_logits_dtype),
97 + )
98 +
99 + def _allocate_dp_chunking_outputs(
100 + self,
101 + hidden_states: torch.Tensor,
102 + router_logits: torch.Tensor,
103 + shared_experts_input: torch.Tensor | None,
104 + ) -> tuple[torch.Tensor | None, torch.Tensor]:
105 + # Assert the inputs are of the proper type and shape.
106 + assert self.batched_hidden_states is not None
107 + assert self.batched_router_logits is not None
108 +
109 + assert self.batched_hidden_states.dtype == hidden_states.dtype, (
110 + f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}"
111 + )
112 + assert self.batched_router_logits.dtype == router_logits.dtype, (
113 + f"{self.batched_router_logits.dtype} == {router_logits.dtype}"
114 + )
115 +
116 + # Check size compatibility.
117 + assert self.batched_hidden_states.size(-1) == hidden_states.size(-1)
118 + assert self.batched_router_logits.size(-1) == router_logits.size(-1)
119 +
120 + final_fused_hidden_states = torch.empty_like(hidden_states)
121 + if self.shared_experts is not None:
122 + if shared_experts_input is not None:
123 + final_shared_hidden_states = torch.empty_like(shared_experts_input)
124 + else:
125 + final_shared_hidden_states = torch.empty_like(hidden_states)
126 + else:
127 + final_shared_hidden_states = None
128 +
129 + return final_shared_hidden_states, final_fused_hidden_states
130 +
131 + def _slice_and_copy_input(
132 + self,
133 + out_slice: torch.Tensor,
134 + orig: torch.Tensor | None,
135 + start: int,
136 + end: int,
137 + ) -> torch.Tensor:
138 + assert orig is not None
139 + slice_size = end - start
140 + orig_slice = orig[start:end, :]
141 + if self.enable_dbo:
142 + assert out_slice.dim() == 3
143 + batch_buffer_idx = dbo_current_ubatch_id()
144 + out_slice = out_slice[batch_buffer_idx, :]
145 +
146 + assert out_slice.size(0) >= slice_size
147 + out_slice = out_slice[:slice_size, :]
148 + out_slice.copy_(orig_slice, non_blocking=True)
149 + return out_slice
150 +
151 + def _forward_impl(
152 + self,
153 + layer: torch.nn.Module,
154 + hidden_states: torch.Tensor,
155 + router_logits: torch.Tensor,
156 + shared_experts_input: torch.Tensor | None,
157 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
158 + final_shared_hidden_states, final_fused_hidden_states = (
159 + self._allocate_dp_chunking_outputs(
160 + hidden_states, router_logits, shared_experts_input
161 + )
162 + )
163 +
164 + ctx = get_forward_context()
165 + # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
166 + max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu
167 + moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens
168 +
169 + # If the input to the MoE is sequence parallel then divide by sp_size
170 + # to find the maximum number of tokens for any individual dispatcher.
171 + if self.moe_config.is_sequence_parallel:
172 + max_tokens_across_dispatchers = cdiv(
173 + max_tokens_across_dispatchers, self.moe_config.sp_size
174 + )
175 +
176 + num_tokens = hidden_states.size(0)
177 + for chunk_idx, chunk_start_ in enumerate(
178 + range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)
179 + ):
180 + chunk_start = chunk_start_
181 + chunk_end = min(
182 + chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers
183 + )
184 + # clamp start and end
185 + chunk_start = min(chunk_start, num_tokens - 1)
186 + chunk_end = min(chunk_end, num_tokens)
187 + chunk_sizes = ctx.dp_metadata.chunked_sizes(
188 + self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx
189 + )
190 + with chunk_sizes:
191 + hidden_states_chunk = self._slice_and_copy_input(
192 + self.batched_hidden_states,
193 + hidden_states,
194 + chunk_start,
195 + chunk_end,
196 + )
197 +
198 + router_logits_chunk = self._slice_and_copy_input(
199 + self.batched_router_logits,
200 + router_logits,
201 + chunk_start,
202 + chunk_end,
203 + )
204 +
205 + shared_experts_input_chunk = (
206 + shared_experts_input[chunk_start:chunk_end, :]
207 + if shared_experts_input is not None
208 + else None
209 + )
210 +
211 + # Delegate per-chunk computation to the inner runner.
212 + chunk_result = self._inner._forward_impl(
213 + layer=layer,
214 + hidden_states=hidden_states_chunk,
215 + router_logits=router_logits_chunk,
216 + shared_experts_input=shared_experts_input_chunk,
217 + )
218 +
219 + # Store outputs
220 + # TODO(bnell): document when chunk_start >= num_tokens
221 + if chunk_start < num_tokens:
222 + if self.shared_experts is not None:
223 + assert isinstance(chunk_result, tuple)
224 + shared_output_chunk, hidden_states_chunk = chunk_result
225 + final_fused_hidden_states[chunk_start:chunk_end, :].copy_(
226 + hidden_states_chunk, non_blocking=True
227 + )
228 + assert shared_output_chunk is not None
229 + assert final_shared_hidden_states is not None
230 + final_shared_hidden_states[chunk_start:chunk_end, :].copy_(
231 + shared_output_chunk, non_blocking=True
232 + )
233 + else:
234 + assert isinstance(chunk_result, torch.Tensor)
235 + final_fused_hidden_states[chunk_start:chunk_end, :].copy_(
236 + chunk_result, non_blocking=True
237 + )
238 +
239 + if self.shared_experts is None:
240 + return final_fused_hidden_states
241 + else:
242 + assert final_shared_hidden_states is not None
243 + return (final_shared_hidden_states, final_fused_hidden_states)
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+22 -682
@@ -1,516 +1,45 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Callable
4 -from contextlib import nullcontext
5 -from typing import TYPE_CHECKING
3
4 import torch
8 -import torch.nn.functional as F
5
6 from vllm.distributed import (
7 get_ep_group,
8 get_pcp_group,
13 - tensor_model_parallel_all_reduce,
14 -)
15 -from vllm.forward_context import (
16 - ForwardContext,
17 - get_forward_context,
18 - is_forward_context_available,
19 -)
20 -from vllm.logger import init_logger
21 -from vllm.model_executor.layers.fused_moe.config import (
22 - FusedMoEConfig,
23 -)
24 -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
25 - FusedMoEMethodBase,
26 -)
27 -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
28 - FusedMoERouter,
29 -)
30 -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
31 -from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
32 - SharedExperts,
33 - SharedExpertsOrder,
34 -)
35 -from vllm.platforms import current_platform
36 -from vllm.utils.math_utils import cdiv
37 -from vllm.utils.torch_utils import (
38 - HAS_OPAQUE_TYPE,
39 - ModuleName,
40 - direct_register_custom_op,
41 -)
42 -from vllm.v1.worker.ubatching import dbo_current_ubatch_id
43 -
44 -logger = init_logger(__name__)
45 -
46 -
47 -def get_layer_from_name(layer_name: str) -> torch.nn.Module:
48 - forward_context: ForwardContext = get_forward_context()
49 - if layer_name == "from_forward_context":
50 - all_moe_layers = forward_context.all_moe_layers
51 - assert all_moe_layers is not None
52 - moe_layer_index = forward_context.moe_layer_index
53 - if moe_layer_index >= len(all_moe_layers):
54 - raise AssertionError(
55 - "We expected the number of MOE layers in `all_moe_layers` "
56 - "to be equal to the number of "
57 - "{vllm.moe_forward, vllm.moe_forward_shared} calls."
58 - )
59 - layer_name = all_moe_layers[moe_layer_index]
60 - forward_context.moe_layer_index += 1
61 - return forward_context.no_compile_layers[layer_name]
62 -
63 -
64 -# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object;
65 -# on older versions it remains a plain str.
66 -if TYPE_CHECKING:
67 - from typing import TypeAlias
68 -
69 - _layer_name_type: TypeAlias = str | ModuleName
70 -else:
71 - _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str
72 -
73 -
74 -def _resolve_layer_name(layer_name: str | ModuleName) -> str:
75 - return layer_name.value if isinstance(layer_name, ModuleName) else layer_name
76 -
77 -
78 -# Note: _moe_forward and _moe_forward_shared should not contain any
79 -# implementation details, They should merely pass along control to
80 -# the runner's 'forward_dispatch' method.
81 -def _moe_forward(
82 - hidden_states: torch.Tensor,
83 - router_logits: torch.Tensor,
84 - shared_experts_input: torch.Tensor | None,
85 - layer_name: _layer_name_type,
86 -) -> torch.Tensor:
87 - layer = get_layer_from_name(_resolve_layer_name(layer_name))
88 - return layer.runner.forward_dispatch(
89 - layer,
90 - hidden_states,
91 - router_logits,
92 - shared_experts_input,
93 - )
94 -
95 -
96 -def _moe_forward_fake(
97 - hidden_states: torch.Tensor,
98 - router_logits: torch.Tensor,
99 - shared_experts_input: torch.Tensor | None,
100 - layer_name: _layer_name_type,
101 -) -> torch.Tensor:
102 - return torch.empty_like(hidden_states)
103 -
104 -
105 -def _moe_forward_shared(
106 - hidden_states: torch.Tensor,
107 - router_logits: torch.Tensor,
108 - shared_experts_input: torch.Tensor | None,
109 - layer_name: _layer_name_type,
110 -) -> tuple[torch.Tensor, torch.Tensor]:
111 - layer = get_layer_from_name(_resolve_layer_name(layer_name))
112 - return layer.runner.forward_dispatch(
113 - layer,
114 - hidden_states,
115 - router_logits,
116 - shared_experts_input,
117 - )
118 -
119 -
120 -def _moe_forward_shared_fake(
121 - hidden_states: torch.Tensor,
122 - router_logits: torch.Tensor,
123 - shared_experts_input: torch.Tensor | None,
124 - layer_name: _layer_name_type,
125 -) -> tuple[torch.Tensor, torch.Tensor]:
126 - # Output shapes:
127 - # - fused_out: same as hidden_states (routed experts use transformed size)
128 - # - shared_out: same as shared_experts_input if provided, else same as
129 - # hidden_states
130 - # (For latent MoE: shared experts use original hidden_size, not latent size)
131 - fused_out = torch.empty_like(hidden_states)
132 - if shared_experts_input is not None:
133 - shared_out = torch.empty_like(shared_experts_input)
134 - else:
135 - shared_out = torch.empty_like(hidden_states)
136 - return shared_out, fused_out
137 -
138 -
139 -direct_register_custom_op(
140 - op_name="moe_forward",
141 - op_func=_moe_forward,
142 - mutates_args=["hidden_states"], # is this still true?
143 - fake_impl=_moe_forward_fake,
144 - tags=(torch.Tag.needs_fixed_stride_order,),
145 -)
146 -
147 -
148 -direct_register_custom_op(
149 - op_name="moe_forward_shared",
150 - op_func=_moe_forward_shared,
151 - fake_impl=_moe_forward_shared_fake,
152 - tags=(torch.Tag.needs_fixed_stride_order,),
9 )
10 +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase
11
12
156 -class DefaultMoERunner(MoERunner):
13 +class DefaultMoERunner(MoERunnerBase):
14 """
158 - Default implementation of the MoE runner for executing Mixture of Experts layers.
15 + Standard MoE runner implementation for executing Mixture of Experts layers.
16
160 - This class provides a comprehensive implementation for running MoE computations
161 - with support for:
162 - - Expert routing and token dispatching
17 + This is the primary concrete implementation of MoE execution logic, providing
18 + comprehensive support for standard MoE operations. It handles:
19 + - Expert routing and token dispatching using various routing strategies
20 - Shared experts computation with optional parallel execution using CUDA streams
164 - - Data parallel (DP) chunking for large batch processing
21 - Tensor model parallel and expert parallel operations
166 - - Various quantization methods and custom operators
22 + - Multiple quantization methods and optimized kernel selection
23 - Both monolithic and decomposed expert execution paths
24 + - Integration with various parallel execution modes (TP, EP, DP)
25
169 - The runner handles the complete MoE forward pass including routing tokens to
170 - experts, executing expert computations, and combining results. It supports
171 - advanced features like overlapped execution of shared experts and optimized
172 - kernels for different parallel execution modes.
173 -
174 - Eventually, this class will be split up and specialized for different
175 - configurations, e.g. the presence or absence of shared experts, a gate, etc.
176 - """
177 -
178 - def __init__(
179 - self,
180 - layer: torch.nn.Module,
181 - moe_config: FusedMoEConfig,
182 - router: FusedMoERouter,
183 - routed_input_transform: torch.nn.Module | None,
184 - gate: torch.nn.Module | None,
185 - shared_experts: torch.nn.Module | None,
186 - quant_method: FusedMoEMethodBase,
187 - reduce_results: bool,
188 - enable_dbo: bool,
189 - ):
190 - super().__init__()
191 - self.moe_config = moe_config
192 - self.router = router
193 - self.routed_input_transform = routed_input_transform
194 - self.gate = gate
195 - self.quant_method = quant_method
196 - self.reduce_results = reduce_results
197 - self.enable_dbo = enable_dbo
198 -
199 - self.shared_experts: SharedExperts | None = None
200 - if shared_experts is not None:
201 - self.shared_experts = SharedExperts(
202 - shared_experts,
203 - moe_config=moe_config,
204 - # Note: For now we must pass quant_method along to SharedExperts so it
205 - # can property determine where the shared experts are supposed to be
206 - # called, i.e. by a MK or by the MoERunner.
207 - # Once the MK can be created upfront, we can just pass in the proper
208 - # flags derived from the quant_method's MK.
209 - reduce_results=reduce_results,
210 - quant_method=quant_method,
211 - enable_dbo=enable_dbo,
212 - )
213 -
214 - # Chunked all2all staging tensor
215 - # These need to exist ahead of time due to CUDAgraph construction
216 - # needing a fixed buffer address.
217 - self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking
218 - self.batched_hidden_states: torch.Tensor | None = None
219 - self.batched_router_logits: torch.Tensor | None = None
220 - self._maybe_init_dp_chunking()
221 -
222 - # Needed for string -> FusedMoE layer lookup in custom ops.
223 - self.layer_name = layer.layer_name
224 -
225 - self.forward_entry, self.forward_impl = self._select_forward(layer)
226 -
227 - def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]:
228 - # Select implementation based on presence of DP chunking.
229 - forward_impl_fn = (
230 - self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl
231 - )
232 -
233 - if current_platform.is_tpu() or current_platform.is_cpu():
234 - # TODO: Once the OOM issue for the TPU backend is resolved, we
235 - # will switch to using the moe_forward custom op.
236 - # Note: CPU doesn't require wrapped forward_impl.
237 - return (
238 - _moe_forward if self.shared_experts is None else _moe_forward_shared,
239 - forward_impl_fn,
240 - )
241 -
242 - return (
243 - torch.ops.vllm.moe_forward
244 - if self.shared_experts is None
245 - else torch.ops.vllm.moe_forward_shared,
246 - forward_impl_fn,
247 - )
248 -
249 - # TODO(bnell): temporary hack, do not call this method.
250 - def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
251 - if self.shared_experts is not None:
252 - self.shared_experts._quant_method = quant_method
253 - self.quant_method = quant_method
254 -
255 - def is_internal_router(self) -> bool:
256 - return self.gate is not None
257 -
258 - def _maybe_init_dp_chunking(self):
259 - if not self.use_dp_chunking:
260 - return
261 -
262 - assert self.batched_hidden_states is None
263 - states_shape: tuple[int, ...]
264 - logits_shape: tuple[int, ...]
265 -
266 - moe = self.moe_config
267 -
268 - if self.enable_dbo:
269 - states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim)
270 - logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts)
271 - else:
272 - states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)
273 - logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)
274 -
275 - device = torch.accelerator.current_device_index()
276 - self.batched_hidden_states = torch.zeros(
277 - states_shape,
278 - dtype=moe.in_dtype,
279 - device=device,
280 - )
281 -
282 - self.batched_router_logits = torch.zeros(
283 - logits_shape,
284 - dtype=moe.router_logits_dtype,
285 - device=device,
286 - )
287 -
288 - def must_reduce_shared_expert_outputs(self) -> bool:
289 - """
290 - The shared_experts are typically computed using the RowParallelLinear
291 - layer. The result of this function is typically used as
292 - the reduce_results argument to the module.
293 - When just tensor-parallel is used, it is not required to reduce
294 - the shared_experts results immediately. Instead we reduce at the
295 - once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
296 - With EP and all2all kernels - this is no longer viable as all
297 - GPU ranks in DP, produce the complete set of hidden_states.
298 - Therefore it is required that we reduce the shared_experts output
299 - early.
300 - """
301 - return (
302 - self.quant_method.moe_kernel is not None
303 - and self.quant_method.moe_kernel.output_is_reduced()
304 - )
305 -
306 - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):
307 - """
308 - Some combine kernels reduce across GPU ranks by default.
309 - """
310 - if self.must_reduce_shared_expert_outputs():
311 - return final_hidden_states
312 - else:
313 - return tensor_model_parallel_all_reduce(final_hidden_states)
314 -
315 - def apply_routed_input_transform(
316 - self, hidden_states: torch.Tensor
317 - ) -> tuple[torch.Tensor, torch.Tensor | None]:
318 - """Apply transform for routed experts (e.g., latent projection).
319 -
320 - This is called by FusedMoE.forward_native. The original hidden_states
321 - is saved separately so shared experts get [S, hidden_size] while
322 - routed experts get the transformed [S, moe_latent_size].
323 -
324 - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be
325 - moved inside SharedFusedMoE to all-reduce on the smaller latent
326 - dimension.
327 -
328 - Returns (possibly transformed) hidden states and the input for shared
329 - experts (or None if there are no shared experts).
330 - """
331 - if self.routed_input_transform is not None:
332 - result = self.routed_input_transform(hidden_states)
333 - # ReplicatedLinear returns (output, extra_bias) tuple.
334 - # We only need the output tensor; extra_bias is not used here.
335 - if isinstance(result, tuple):
336 - return result[0], hidden_states
337 - return result, hidden_states
338 -
339 - return (
340 - hidden_states,
341 - hidden_states if self.shared_experts is not None else None,
342 - )
343 -
344 - def _maybe_reduce_output(
345 - self,
346 - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
347 - trunc_sizes: list[int],
348 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
349 - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
350 - return x[..., :trunc_size]
351 -
352 - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
353 - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size)
354 -
355 - if (
356 - not self.moe_config.is_sequence_parallel
357 - and not self.use_dp_chunking
358 - and self.reduce_results
359 - and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
360 - ):
361 - func = reduce_and_trunc
362 - else:
363 - func = trunc
364 -
365 - if isinstance(states, tuple):
366 - return tuple(
367 - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)]
368 - )
369 - else:
370 - assert len(trunc_sizes) == 1
371 - return func(states, trunc_sizes[0])
372 -
373 - def _encode_layer_name(self) -> str | ModuleName:
374 - if HAS_OPAQUE_TYPE:
375 - return ModuleName(self.layer_name)
376 - # Can be unavailable or None in unittests
377 - if (
378 - is_forward_context_available()
379 - and get_forward_context().all_moe_layers is not None
380 - ):
381 - return "from_forward_context"
382 - return self.layer_name
383 -
384 - def _maybe_pad_hidden_states(
385 - self,
386 - shared_experts_input: torch.Tensor | None,
387 - hidden_states: torch.Tensor,
388 - ) -> tuple[torch.Tensor, list[int]]:
389 - shared_experts_hidden_dim = (
390 - shared_experts_input.shape[-1] if shared_experts_input is not None else 0
391 - )
392 - transformed_hidden_dim = hidden_states.shape[-1]
393 - if (
394 - not self.quant_method.skip_forward_padding
395 - and self.moe_config.hidden_dim != transformed_hidden_dim
396 - ):
397 - hidden_states = F.pad(
398 - hidden_states,
399 - (0, self.moe_config.hidden_dim - transformed_hidden_dim),
400 - mode="constant",
401 - value=0.0,
402 - )
403 -
404 - if self.shared_experts is not None:
405 - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim]
406 - else:
407 - orig_hidden_dims = [transformed_hidden_dim]
408 -
409 - return hidden_states, orig_hidden_dims
26 + The runner orchestrates the complete MoE forward pass including routing tokens
27 + to experts, executing expert computations in parallel, and combining results.
28 + It supports advanced features like overlapped execution of shared experts,
29 + optimized kernels for different parallel configurations, and seamless
30 + integration with vLLM's distributed execution framework.
31
411 - def _maybe_apply_shared_experts(
412 - self,
413 - shared_experts_input: torch.Tensor | None,
414 - order: SharedExpertsOrder,
415 - ):
416 - if self.shared_experts is not None:
417 - assert shared_experts_input is not None
418 - self.shared_experts.apply(shared_experts_input, order)
32 + This implementation is suitable for most standard MoE use cases. For specialized
33 + scenarios like large batch chunking, alternative runners like ChunkingMoERunner
34 + may be more appropriate.
35
420 - def _apply_quant_method(
421 - self,
422 - layer: torch.nn.Module,
423 - hidden_states: torch.Tensor,
424 - router_logits: torch.Tensor,
425 - shared_experts_input: torch.Tensor | None,
426 - ) -> tuple[torch.Tensor | None, torch.Tensor]:
427 - # Run this before quant_method to avoid inplace issues.
428 - # TODO(bnell): probably not needed anymore since inplace is
429 - # disabled when shared experts are present.
430 - self._maybe_apply_shared_experts(
431 - shared_experts_input, SharedExpertsOrder.NO_OVERLAP
432 - )
433 -
434 - if self.quant_method.is_monolithic:
435 - fused_out = self.quant_method.apply_monolithic(
436 - layer=layer,
437 - x=hidden_states,
438 - router_logits=router_logits,
439 - )
440 - else:
441 - topk_weights, topk_ids = self.router.select_experts(
442 - hidden_states=hidden_states,
443 - router_logits=router_logits,
444 - )
445 -
446 - # Passing shared_experts_input in case SharedExpertsOrder is
447 - # NO_OVERLAP or MK_INTERNAL_OVERLAPPED.
448 - fused_out = self.quant_method.apply(
449 - layer=layer,
450 - x=hidden_states,
451 - topk_weights=topk_weights,
452 - topk_ids=topk_ids,
453 - shared_experts_input=shared_experts_input,
454 - )
455 -
456 - self._maybe_apply_shared_experts(
457 - shared_experts_input,
458 - SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
459 - )
460 -
461 - return (
462 - self.shared_experts.output if self.shared_experts is not None else None,
463 - fused_out,
464 - )
465 -
466 - def _sequence_parallel_context(self):
467 - ctx = get_forward_context()
468 - return (
469 - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
470 - if ctx.dp_metadata
471 - else nullcontext()
472 - )
473 -
474 - def _allocate_dp_chunking_outputs(
475 - self,
476 - hidden_states: torch.Tensor,
477 - router_logits: torch.Tensor,
478 - ) -> tuple[torch.Tensor | None, torch.Tensor]:
479 - assert self.use_dp_chunking
480 -
481 - # Assert the inputs are of the proper type and shape.
482 - assert self.batched_hidden_states is not None
483 - assert self.batched_router_logits is not None
484 -
485 - assert self.batched_hidden_states.dtype == hidden_states.dtype, (
486 - f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}"
487 - )
488 - assert self.batched_router_logits.dtype == router_logits.dtype, (
489 - f"{self.batched_router_logits.dtype} == {router_logits.dtype}"
490 - )
491 -
492 - # Check size compatibility.
493 - assert self.batched_hidden_states.size(-1) == hidden_states.size(-1)
494 - assert self.batched_router_logits.size(-1) == router_logits.size(-1)
495 -
496 - final_fused_hidden_states = torch.empty_like(hidden_states)
497 - if self.shared_experts is not None:
498 - final_shared_hidden_states = torch.empty_like(hidden_states)
499 - else:
500 - final_shared_hidden_states = None
501 -
502 - return final_shared_hidden_states, final_fused_hidden_states
36 + Eventually, this class may be split into more specialized implementations
37 + for different configurations (e.g., with/without shared experts, gates, etc.).
38 + """
39
504 - def _maybe_sync_shared_experts_stream(
505 - self,
506 - shared_experts_input: torch.Tensor | None,
507 - ):
508 - # If router/gate provided, then apply it here.
509 - # (Note: This code runs only when "overlapped mode" is on to allow
510 - # parallel execution of shared experts with the FusedMoE via
511 - # separate cuda stream)
512 - if self.shared_experts is not None:
513 - self.shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
40 + @property
41 + def reduce_results(self) -> bool:
42 + return self._reduce_results
43
44 @property
45 def do_naive_dispatch_combine(self) -> bool:
@@ -572,195 +101,6 @@ class DefaultMoERunner(MoERunner):
101 else:
102 return hidden_states
103
575 - def forward(
576 - self,
577 - hidden_states: torch.Tensor,
578 - router_logits: torch.Tensor,
579 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
580 - """Invoke the fused moe layer.
581 -
582 - Input:
583 - - hidden_states
584 - - router_logits
585 -
586 - Output:
587 - - The new hidden_states.
588 - or
589 - - A tuple of (shared experts output, new hidden_states).
590 -
591 - Calling sequence
592 - - forward
593 - - self.forward_entry (_moe_forward or _moe_forward_shared custom op)
594 - - forward_dispatch
595 - - forward_impl (_forward_impl or _forward_impl_chunked)
596 -
597 - Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
598 - to the following reasons:
599 - 1. the chunking loop in _forward_impl_chunked cannot be compiled by
600 - torch.compile
601 - 2. pytorch cannot handle union types in custom op signatures so _moe_forward
602 - and _moe_forward_shared must be split.
603 -
604 - If _forward_impl_chunked can be implemented via torch.scan we can potentially
605 - get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence
606 - into the 'forward' method.
607 - """
608 -
609 - # Apply transform for routed experts (e.g., latent projection for latent MoE)
610 - hidden_states, shared_experts_input = self.apply_routed_input_transform(
611 - hidden_states
612 - )
613 -
614 - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(
615 - shared_experts_input,
616 - hidden_states,
617 - )
618 -
619 - fused_output = self.forward_entry(
620 - hidden_states,
621 - router_logits,
622 - shared_experts_input,
623 - self._encode_layer_name(),
624 - )
625 -
626 - return self._maybe_reduce_output(fused_output, og_hidden_dims)
627 -
628 - def forward_dispatch(
629 - self,
630 - layer: torch.nn.Module,
631 - hidden_states: torch.Tensor,
632 - router_logits: torch.Tensor,
633 - shared_experts_input: torch.Tensor | None,
634 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
635 - # TODO(bnell): this can be removed after MK migration is complete.
636 - layer.ensure_moe_quant_config_init()
637 -
638 - # Sync aux and main stream for shared expert multi-stream overlap.
639 - self._maybe_sync_shared_experts_stream(shared_experts_input)
640 -
641 - # If the Runner holds the gate, apply it after the stream sync,
642 - # so it can run overlapped with the
643 - # NOTE: in future PR, MoE runner will always hold the gate.
644 - if self.gate is not None:
645 - router_logits, _ = self.gate(hidden_states)
646 -
647 - self._maybe_apply_shared_experts(
648 - shared_experts_input,
649 - SharedExpertsOrder.EXTERNAL,
650 - )
651 -
652 - with self._sequence_parallel_context():
653 - return self.forward_impl(
654 - layer,
655 - hidden_states,
656 - router_logits,
657 - shared_experts_input,
658 - )
659 -
660 - def _slice_and_copy_input(
661 - self,
662 - out_slice: torch.Tensor,
663 - orig: torch.Tensor | None,
664 - start: int,
665 - end: int,
666 - ) -> torch.Tensor:
667 - assert orig is not None
668 - slice_size = end - start
669 - orig_slice = orig[start:end, :]
670 - if self.enable_dbo:
671 - assert out_slice.dim() == 3
672 - batch_buffer_idx = dbo_current_ubatch_id()
673 - out_slice = out_slice[batch_buffer_idx, :]
674 -
675 - assert out_slice.size(0) >= slice_size
676 - out_slice = out_slice[:slice_size, :]
677 - out_slice.copy_(orig_slice, non_blocking=True)
678 - return out_slice
679 -
680 - def _forward_impl_chunked(
681 - self,
682 - layer: torch.nn.Module,
683 - hidden_states: torch.Tensor,
684 - router_logits: torch.Tensor,
685 - shared_experts_input: torch.Tensor | None,
686 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
687 - final_shared_hidden_states, final_fused_hidden_states = (
688 - self._allocate_dp_chunking_outputs(hidden_states, router_logits)
689 - )
690 -
691 - ctx = get_forward_context()
692 - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
693 - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu
694 - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens
695 -
696 - # If the input to the MoE is sequence parallel then divide by sp_size
697 - # to find the maximum number of tokens for any individual dispatcher.
698 - if self.moe_config.is_sequence_parallel:
699 - max_tokens_across_dispatchers = cdiv(
700 - max_tokens_across_dispatchers, self.moe_config.sp_size
701 - )
702 -
703 - num_tokens = hidden_states.size(0)
704 - for chunk_idx, chunk_start_ in enumerate(
705 - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)
706 - ):
707 - chunk_start = chunk_start_
708 - chunk_end = min(
709 - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers
710 - )
711 - # clamp start and end
712 - chunk_start = min(chunk_start, num_tokens - 1)
713 - chunk_end = min(chunk_end, num_tokens)
714 - chunk_sizes = ctx.dp_metadata.chunked_sizes(
715 - self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx
716 - )
717 - with chunk_sizes:
718 - hidden_states_chunk = self._slice_and_copy_input(
719 - self.batched_hidden_states,
720 - hidden_states,
721 - chunk_start,
722 - chunk_end,
723 - )
724 -
725 - router_logits_chunk = self._slice_and_copy_input(
726 - self.batched_router_logits,
727 - router_logits,
728 - chunk_start,
729 - chunk_end,
730 - )
731 -
732 - shared_experts_input_chunk = (
733 - shared_experts_input[chunk_start:chunk_end, :]
734 - if shared_experts_input is not None
735 - else None
736 - )
737 -
738 - shared_output_chunk, hidden_states_chunk = self._apply_quant_method(
739 - layer=layer,
740 - hidden_states=hidden_states_chunk,
741 - router_logits=router_logits_chunk,
742 - shared_experts_input=shared_experts_input_chunk,
743 - )
744 -
745 - # Store outputs
746 - # TODO(bnell): document when chunk_start >= num_tokens
747 - if chunk_start < num_tokens:
748 - final_fused_hidden_states[chunk_start:chunk_end, :].copy_(
749 - hidden_states_chunk, non_blocking=True
750 - )
751 - if self.shared_experts is not None:
752 - assert shared_output_chunk is not None
753 - assert final_shared_hidden_states is not None
754 - final_shared_hidden_states[chunk_start:chunk_end, :].copy_(
755 - shared_output_chunk, non_blocking=True
756 - )
757 -
758 - if self.shared_experts is None:
759 - return final_fused_hidden_states
760 - else:
761 - assert final_shared_hidden_states is not None
762 - return (final_shared_hidden_states, final_fused_hidden_states)
763 -
104 def _forward_impl(
105 self,
106 layer: torch.nn.Module,
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+17
@@ -4,6 +4,13 @@ from abc import ABC, abstractmethod
4
5 import torch
6
7 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
8 + FusedMoEMethodBase,
9 +)
10 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
11 + SharedExperts,
12 +)
13 +
14
15 class MoERunner(ABC):
16 """
@@ -36,3 +43,13 @@ class MoERunner(ABC):
43 @abstractmethod
44 def is_internal_router(self) -> bool:
45 raise NotImplementedError
46 +
47 + @property
48 + @abstractmethod
49 + def shared_experts(self) -> SharedExperts | None:
50 + raise NotImplementedError
51 +
52 + # TODO(bnell): temporary hack, do not call this method.
53 + @abstractmethod
54 + def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
55 + raise NotImplementedError
vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py new
+527
@@ -0,0 +1,527 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import abstractmethod
4 +from collections.abc import Callable
5 +from contextlib import nullcontext
6 +from typing import TYPE_CHECKING
7 +
8 +import torch
9 +import torch.nn.functional as F
10 +
11 +from vllm.distributed import (
12 + tensor_model_parallel_all_reduce,
13 +)
14 +from vllm.forward_context import (
15 + ForwardContext,
16 + get_forward_context,
17 + is_forward_context_available,
18 +)
19 +from vllm.model_executor.layers.fused_moe.config import (
20 + FusedMoEConfig,
21 +)
22 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
23 + FusedMoEMethodBase,
24 +)
25 +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
26 + FusedMoERouter,
27 +)
28 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
29 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
30 + SharedExperts,
31 + SharedExpertsOrder,
32 +)
33 +from vllm.platforms import current_platform
34 +from vllm.utils.torch_utils import (
35 + HAS_OPAQUE_TYPE,
36 + ModuleName,
37 + direct_register_custom_op,
38 +)
39 +
40 +
41 +def get_layer_from_name(layer_name: str) -> torch.nn.Module:
42 + forward_context: ForwardContext = get_forward_context()
43 + if layer_name == "from_forward_context":
44 + all_moe_layers = forward_context.all_moe_layers
45 + assert all_moe_layers is not None
46 + moe_layer_index = forward_context.moe_layer_index
47 + if moe_layer_index >= len(all_moe_layers):
48 + raise AssertionError(
49 + "We expected the number of MOE layers in `all_moe_layers` "
50 + "to be equal to the number of "
51 + "{vllm.moe_forward, vllm.moe_forward_shared} calls."
52 + )
53 + layer_name = all_moe_layers[moe_layer_index]
54 + forward_context.moe_layer_index += 1
55 + return forward_context.no_compile_layers[layer_name]
56 +
57 +
58 +# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object;
59 +# on older versions it remains a plain str.
60 +if TYPE_CHECKING:
61 + from typing import TypeAlias
62 +
63 + _layer_name_type: TypeAlias = str | ModuleName
64 +else:
65 + _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str
66 +
67 +
68 +def _resolve_layer_name(layer_name: str | ModuleName) -> str:
69 + return layer_name.value if isinstance(layer_name, ModuleName) else layer_name
70 +
71 +
72 +# Note: _moe_forward and _moe_forward_shared should not contain any
73 +# implementation details, They should merely pass along control to
74 +# the runner's 'forward_dispatch' method.
75 +def _moe_forward(
76 + hidden_states: torch.Tensor,
77 + router_logits: torch.Tensor,
78 + shared_experts_input: torch.Tensor | None,
79 + layer_name: _layer_name_type,
80 +) -> torch.Tensor:
81 + layer = get_layer_from_name(_resolve_layer_name(layer_name))
82 + return layer.runner.forward_dispatch(
83 + layer,
84 + hidden_states,
85 + router_logits,
86 + shared_experts_input,
87 + )
88 +
89 +
90 +def _moe_forward_fake(
91 + hidden_states: torch.Tensor,
92 + router_logits: torch.Tensor,
93 + shared_experts_input: torch.Tensor | None,
94 + layer_name: _layer_name_type,
95 +) -> torch.Tensor:
96 + return torch.empty_like(hidden_states)
97 +
98 +
99 +def _moe_forward_shared(
100 + hidden_states: torch.Tensor,
101 + router_logits: torch.Tensor,
102 + shared_experts_input: torch.Tensor | None,
103 + layer_name: _layer_name_type,
104 +) -> tuple[torch.Tensor, torch.Tensor]:
105 + layer = get_layer_from_name(_resolve_layer_name(layer_name))
106 + return layer.runner.forward_dispatch(
107 + layer,
108 + hidden_states,
109 + router_logits,
110 + shared_experts_input,
111 + )
112 +
113 +
114 +def _moe_forward_shared_fake(
115 + hidden_states: torch.Tensor,
116 + router_logits: torch.Tensor,
117 + shared_experts_input: torch.Tensor | None,
118 + layer_name: _layer_name_type,
119 +) -> tuple[torch.Tensor, torch.Tensor]:
120 + # Output shapes:
121 + # - fused_out: same as hidden_states (routed experts use transformed size)
122 + # - shared_out: same as shared_experts_input if provided, else same as
123 + # hidden_states
124 + # (For latent MoE: shared experts use original hidden_size, not latent size)
125 + fused_out = torch.empty_like(hidden_states)
126 + if shared_experts_input is not None:
127 + shared_out = torch.empty_like(shared_experts_input)
128 + else:
129 + shared_out = torch.empty_like(hidden_states)
130 + return shared_out, fused_out
131 +
132 +
133 +direct_register_custom_op(
134 + op_name="moe_forward",
135 + op_func=_moe_forward,
136 + mutates_args=["hidden_states"], # is this still true?
137 + fake_impl=_moe_forward_fake,
138 + tags=(torch.Tag.needs_fixed_stride_order,),
139 +)
140 +
141 +
142 +direct_register_custom_op(
143 + op_name="moe_forward_shared",
144 + op_func=_moe_forward_shared,
145 + fake_impl=_moe_forward_shared_fake,
146 + tags=(torch.Tag.needs_fixed_stride_order,),
147 +)
148 +
149 +
150 +class MoERunnerBase(MoERunner):
151 + """
152 + Abstract base class providing common functionality for MoE runner implementations.
153 +
154 + This class serves as the foundation for concrete MoE runner implementations by
155 + providing shared state management and common utilities. It handles:
156 + - Common initialization and configuration management
157 + - Shared expert output reduction logic for tensor parallel scenarios
158 + - Base methods for tensor model parallel reductions
159 + - Common properties and utility functions used across different runner types
160 +
161 + Concrete subclasses must implement the abstract methods to define their specific
162 + execution strategies, such as standard execution, chunked processing, or other
163 + specialized approaches. The base class provides the infrastructure while
164 + allowing flexibility in the actual MoE computation implementation.
165 +
166 + Key abstract methods that subclasses must implement:
167 + - reduce_results: Determines whether results should be reduced across ranks
168 + - _forward_impl: The core MoE computation logic specific to each runner type
169 + """
170 +
171 + def __init__(
172 + self,
173 + layer_name: str,
174 + moe_config: FusedMoEConfig,
175 + router: FusedMoERouter,
176 + routed_input_transform: torch.nn.Module | None,
177 + gate: torch.nn.Module | None,
178 + shared_experts: torch.nn.Module | None,
179 + quant_method: FusedMoEMethodBase,
180 + reduce_results: bool,
181 + enable_dbo: bool,
182 + ):
183 + super().__init__()
184 + self.moe_config = moe_config
185 + self.router = router
186 + self.routed_input_transform = routed_input_transform
187 + self.gate = gate
188 + self.quant_method = quant_method
189 + self._reduce_results = reduce_results
190 + self.enable_dbo = enable_dbo
191 +
192 + self._shared_experts: SharedExperts | None = None
193 + if shared_experts is not None:
194 + self._shared_experts = SharedExperts(
195 + shared_experts,
196 + moe_config=moe_config,
197 + # Note: For now we must pass quant_method along to SharedExperts so it
198 + # can property determine where the shared experts are supposed to be
199 + # called, i.e. by a MK or by the MoERunner.
200 + # Once the MK can be created upfront, we can just pass in the proper
201 + # flags derived from the quant_method's MK.
202 + reduce_results=reduce_results,
203 + quant_method=quant_method,
204 + enable_dbo=enable_dbo,
205 + )
206 +
207 + # Needed for string -> FusedMoE layer lookup in custom ops.
208 + self.layer_name = layer_name
209 +
210 + self.forward_entry = self._select_forward()
211 +
212 + def _select_forward(self) -> Callable:
213 + if current_platform.is_tpu() or current_platform.is_cpu():
214 + # TODO: Once the OOM issue for the TPU backend is resolved, we
215 + # will switch to using the moe_forward custom op.
216 + # Note: CPU doesn't require wrapped _forward_impl.
217 + return _moe_forward if self._shared_experts is None else _moe_forward_shared
218 +
219 + return (
220 + torch.ops.vllm.moe_forward
221 + if self._shared_experts is None
222 + else torch.ops.vllm.moe_forward_shared
223 + )
224 +
225 + @property
226 + def shared_experts(self) -> SharedExperts | None:
227 + return self._shared_experts
228 +
229 + # TODO(bnell): temporary hack, do not call this method.
230 + def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
231 + if self._shared_experts is not None:
232 + self._shared_experts._quant_method = quant_method
233 + self.quant_method = quant_method
234 +
235 + def is_internal_router(self) -> bool:
236 + return self.gate is not None
237 +
238 + @property
239 + @abstractmethod
240 + def reduce_results(self) -> bool:
241 + raise NotImplementedError
242 +
243 + def must_reduce_shared_expert_outputs(self) -> bool:
244 + """
245 + The shared_experts are typically computed using the RowParallelLinear
246 + layer. The result of this function is typically used as
247 + the reduce_results argument to the module.
248 + When just tensor-parallel is used, it is not required to reduce
249 + the shared_experts results immediately. Instead we reduce at the
250 + once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
251 + With EP and all2all kernels - this is no longer viable as all
252 + GPU ranks in DP, produce the complete set of hidden_states.
253 + Therefore it is required that we reduce the shared_experts output
254 + early.
255 + """
256 + return (
257 + self.quant_method.moe_kernel is not None
258 + and self.quant_method.moe_kernel.output_is_reduced()
259 + )
260 +
261 + def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):
262 + """
263 + Some combine kernels reduce across GPU ranks by default.
264 + """
265 + if self.must_reduce_shared_expert_outputs():
266 + return final_hidden_states
267 + else:
268 + return tensor_model_parallel_all_reduce(final_hidden_states)
269 +
270 + def apply_routed_input_transform(
271 + self, hidden_states: torch.Tensor
272 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
273 + """Apply transform for routed experts (e.g., latent projection).
274 +
275 + This is called by FusedMoE.forward_native. The original hidden_states
276 + is saved separately so shared experts get [S, hidden_size] while
277 + routed experts get the transformed [S, moe_latent_size].
278 +
279 + TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be
280 + moved inside SharedFusedMoE to all-reduce on the smaller latent
281 + dimension.
282 +
283 + Returns (possibly transformed) hidden states and the input for shared
284 + experts (or None if there are no shared experts).
285 + """
286 + if self.routed_input_transform is not None:
287 + result = self.routed_input_transform(hidden_states)
288 + # ReplicatedLinear returns (output, extra_bias) tuple.
289 + # We only need the output tensor; extra_bias is not used here.
290 + if isinstance(result, tuple):
291 + return result[0], hidden_states
292 + return result, hidden_states
293 +
294 + return (
295 + hidden_states,
296 + hidden_states if self._shared_experts is not None else None,
297 + )
298 +
299 + def _maybe_reduce_output(
300 + self,
301 + states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
302 + trunc_sizes: list[int],
303 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
304 + def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
305 + return x[..., :trunc_size]
306 +
307 + def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
308 + return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size)
309 +
310 + if (
311 + not self.moe_config.is_sequence_parallel
312 + and self.reduce_results
313 + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
314 + ):
315 + func = reduce_and_trunc
316 + else:
317 + func = trunc
318 +
319 + if isinstance(states, tuple):
320 + return tuple(
321 + [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)]
322 + )
323 + else:
324 + assert len(trunc_sizes) == 1
325 + return func(states, trunc_sizes[0])
326 +
327 + def _encode_layer_name(self) -> str | ModuleName:
328 + if HAS_OPAQUE_TYPE:
329 + return ModuleName(self.layer_name)
330 + # Can be unavailable or None in unittests
331 + if (
332 + is_forward_context_available()
333 + and get_forward_context().all_moe_layers is not None
334 + ):
335 + return "from_forward_context"
336 + return self.layer_name
337 +
338 + def _maybe_pad_hidden_states(
339 + self,
340 + shared_experts_input: torch.Tensor | None,
341 + hidden_states: torch.Tensor,
342 + ) -> tuple[torch.Tensor, list[int]]:
343 + shared_experts_hidden_dim = (
344 + shared_experts_input.shape[-1] if shared_experts_input is not None else 0
345 + )
346 + transformed_hidden_dim = hidden_states.shape[-1]
347 + if (
348 + not self.quant_method.skip_forward_padding
349 + and self.moe_config.hidden_dim != transformed_hidden_dim
350 + ):
351 + hidden_states = F.pad(
352 + hidden_states,
353 + (0, self.moe_config.hidden_dim - transformed_hidden_dim),
354 + mode="constant",
355 + value=0.0,
356 + )
357 +
358 + if self._shared_experts is not None:
359 + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim]
360 + else:
361 + orig_hidden_dims = [transformed_hidden_dim]
362 +
363 + return hidden_states, orig_hidden_dims
364 +
365 + def _maybe_apply_shared_experts(
366 + self,
367 + shared_experts_input: torch.Tensor | None,
368 + order: SharedExpertsOrder,
369 + ):
370 + if self._shared_experts is not None:
371 + assert shared_experts_input is not None
372 + self._shared_experts.apply(shared_experts_input, order)
373 +
374 + def _apply_quant_method(
375 + self,
376 + layer: torch.nn.Module,
377 + hidden_states: torch.Tensor,
378 + router_logits: torch.Tensor,
379 + shared_experts_input: torch.Tensor | None,
380 + ) -> tuple[torch.Tensor | None, torch.Tensor]:
381 + # Run this before quant_method to avoid inplace issues.
382 + # TODO(bnell): probably not needed anymore since inplace is
383 + # disabled when shared experts are present.
384 + self._maybe_apply_shared_experts(
385 + shared_experts_input, SharedExpertsOrder.NO_OVERLAP
386 + )
387 +
388 + if self.quant_method.is_monolithic:
389 + fused_out = self.quant_method.apply_monolithic(
390 + layer=layer,
391 + x=hidden_states,
392 + router_logits=router_logits,
393 + )
394 + else:
395 + topk_weights, topk_ids = self.router.select_experts(
396 + hidden_states=hidden_states,
397 + router_logits=router_logits,
398 + )
399 +
400 + # Passing shared_experts_input in case SharedExpertsOrder is
401 + # NO_OVERLAP or MK_INTERNAL_OVERLAPPED.
402 + fused_out = self.quant_method.apply(
403 + layer=layer,
404 + x=hidden_states,
405 + topk_weights=topk_weights,
406 + topk_ids=topk_ids,
407 + shared_experts_input=shared_experts_input,
408 + )
409 +
410 + self._maybe_apply_shared_experts(
411 + shared_experts_input,
412 + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
413 + )
414 +
415 + return (
416 + self._shared_experts.output if self._shared_experts is not None else None,
417 + fused_out,
418 + )
419 +
420 + def _sequence_parallel_context(self):
421 + ctx = get_forward_context()
422 + return (
423 + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
424 + if ctx.dp_metadata
425 + else nullcontext()
426 + )
427 +
428 + def _maybe_sync_shared_experts_stream(
429 + self,
430 + shared_experts_input: torch.Tensor | None,
431 + ):
432 + # If router/gate provided, then apply it here.
433 + # (Note: This code runs only when "overlapped mode" is on to allow
434 + # parallel execution of shared experts with the FusedMoE via
435 + # separate cuda stream)
436 + if self._shared_experts is not None:
437 + self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
438 +
439 + def forward(
440 + self,
441 + hidden_states: torch.Tensor,
442 + router_logits: torch.Tensor,
443 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
444 + """Invoke the fused moe layer.
445 +
446 + Input:
447 + - hidden_states
448 + - router_logits
449 +
450 + Output:
451 + - The new hidden_states.
452 + or
453 + - A tuple of (shared experts output, new hidden_states).
454 +
455 + Calling sequence
456 + - forward
457 + - self.forward_entry (_moe_forward or _moe_forward_shared custom op)
458 + - forward_dispatch
459 + - _forward_impl
460 +
461 + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
462 + to the following reasons:
463 + 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by
464 + torch.compile
465 + 2. pytorch cannot handle union types in custom op signatures so _moe_forward
466 + and _moe_forward_shared must be split.
467 +
468 + If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can
469 + potentially get rid of _moe_forward and _moe_forward_shared and collapse the
470 + whole sequence into the 'forward' method.
471 + """
472 +
473 + # Apply transform for routed experts (e.g., latent projection for latent MoE)
474 + hidden_states, shared_experts_input = self.apply_routed_input_transform(
475 + hidden_states
476 + )
477 +
478 + hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(
479 + shared_experts_input,
480 + hidden_states,
481 + )
482 +
483 + fused_output = self.forward_entry(
484 + hidden_states,
485 + router_logits,
486 + shared_experts_input,
487 + self._encode_layer_name(),
488 + )
489 +
490 + return self._maybe_reduce_output(fused_output, og_hidden_dims)
491 +
492 + def forward_dispatch(
493 + self,
494 + layer: torch.nn.Module,
495 + hidden_states: torch.Tensor,
496 + router_logits: torch.Tensor,
497 + shared_experts_input: torch.Tensor | None,
498 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
499 + # TODO(bnell): this can be removed after MK migration is complete.
500 + layer.ensure_moe_quant_config_init()
501 +
502 + # Sync aux and main stream for shared expert multi-stream overlap.
503 + self._maybe_sync_shared_experts_stream(shared_experts_input)
504 +
505 + # If the Runner holds the gate, apply it after the stream sync,
506 + # so it can run overlapped with the
507 + # NOTE: in future PR, MoE runner will always hold the gate.
508 + if self.gate is not None:
509 + router_logits, _ = self.gate(hidden_states)
510 +
511 + with self._sequence_parallel_context():
512 + return self._forward_impl(
513 + layer,
514 + hidden_states,
515 + router_logits,
516 + shared_experts_input,
517 + )
518 +
519 + @abstractmethod
520 + def _forward_impl(
521 + self,
522 + layer: torch.nn.Module,
523 + hidden_states: torch.Tensor,
524 + router_logits: torch.Tensor,
525 + shared_experts_input: torch.Tensor | None,
526 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
527 + raise NotImplementedError
vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py new
+51
@@ -0,0 +1,51 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.model_executor.layers.fused_moe.config import (
7 + FusedMoEConfig,
8 +)
9 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
10 + FusedMoEMethodBase,
11 +)
12 +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
13 + FusedMoERouter,
14 +)
15 +from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import (
16 + ChunkingMoERunner,
17 +)
18 +from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (
19 + DefaultMoERunner,
20 +)
21 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
22 +from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
23 + SharedExperts,
24 +)
25 +
26 +
27 +def create_moe_runner(
28 + layer_name: str,
29 + moe_config: FusedMoEConfig,
30 + router: FusedMoERouter,
31 + routed_input_transform: torch.nn.Module | None,
32 + gate: torch.nn.Module | None,
33 + shared_experts: SharedExperts | None,
34 + quant_method: FusedMoEMethodBase,
35 + reduce_results: bool,
36 + enable_dbo: bool,
37 +) -> MoERunner:
38 + runner = DefaultMoERunner(
39 + layer_name,
40 + moe_config,
41 + router,
42 + routed_input_transform,
43 + gate,
44 + shared_experts,
45 + quant_method,
46 + reduce_results,
47 + enable_dbo,
48 + )
49 + if moe_config.moe_parallel_config.use_dp_chunking:
50 + return ChunkingMoERunner(runner)
51 + return runner
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+3 -19
@@ -32,19 +32,14 @@ class SharedExpertsOrder(IntEnum):
32 # No shared experts.
33 NONE = (0,)
34
35 - # Get rid of this one? combine with BEFORE?
36 - # Note: this might be important for torch.compile reasons. Can
37 - # get rid of it after _moe_forward is undone.
38 - EXTERNAL = (1,)
39 -
35 # No overlap - defensively called before MK.
41 - NO_OVERLAP = (2,)
36 + NO_OVERLAP = (1,)
37
38 # Overlapped with dispatch/combine in DP/EP - called by the MK.
44 - MK_INTERNAL_OVERLAPPED = (3,)
39 + MK_INTERNAL_OVERLAPPED = (2,)
40
41 # Overlapped with the gate, router, experts in aux stream.
47 - MULTI_STREAM_OVERLAPPED = (4,)
42 + MULTI_STREAM_OVERLAPPED = (3,)
43
44
45 class SharedExperts:
@@ -110,9 +105,6 @@ class SharedExperts:
105 self,
106 hidden_states: torch.Tensor,
107 ) -> SharedExpertsOrder:
113 - if self._use_external_experts:
114 - return SharedExpertsOrder.EXTERNAL
115 -
108 if self._quant_method.mk_owns_shared_expert:
109 return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
110
@@ -205,12 +197,4 @@ class SharedExperts:
197 else:
198 self._output[self._output_idx] = self._layer(shared_experts_input)
199
208 - if order == SharedExpertsOrder.EXTERNAL:
209 - # TODO: figure out how to combine this with maybe_reduce_output?
210 - # or get rid of it completely.
211 - assert self._output[self._output_idx] is not None
212 - self._output[self._output_idx] = self._maybe_reduce_shared_out(
213 - self._output[self._output_idx]
214 - )
215 -
200 assert self._output[self._output_idx] is not None
vllm/model_executor/models/nemotron_h.py
+1
@@ -231,6 +231,7 @@ class NemotronHMoE(nn.Module):
231 num_redundant_experts=self.n_redundant_experts,
232 is_sequence_parallel=self.is_sequence_parallel,
233 routed_input_transform=self.fc1_latent_proj,
234 + router_logits_dtype=self.gate.out_dtype,
235 )
236
237 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: