@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 809d83c2dc
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 4, 20:00 UTC

[MoE Refactor] Combine MoERunnerBase + DefaultMoERunner (#40560)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Apr 22, 2026 at 10:43 UTC 809d83c2dc2f74d52fe852b09d0e38c91c7d8334
6 files changed +729 -835
vllm/model_executor/layers/fused_moe/layer.py
+6 -3
@@ -38,8 +38,11 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
38 from vllm.model_executor.layers.fused_moe.router.router_factory import (
39 create_fused_moe_router,
40 )
41 -from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import (
42 - create_moe_runner,
41 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import (
42 + MoERunner,
43 +)
44 +from vllm.model_executor.layers.fused_moe.runner.moe_runner_interface import (
45 + MoERunnerInterface,
46 )
47 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
48 SharedExperts,
@@ -586,7 +589,7 @@ class FusedMoE(PluggableLayer):
589 # Storing the runner in the FusedMoE is an intermediate state, eventually
590 # the runner will own the FusedMoE layer and provide the execution interface
591 # for MoE ops.
589 - self.runner = create_moe_runner(
592 + self.runner: MoERunnerInterface = MoERunner(
593 layer_name=self.layer_name,
594 moe_config=self.moe_config,
595 router=self.router,
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py deleted
-128
@@ -1,128 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch
5 -
6 -from vllm.distributed import (
7 - get_ep_group,
8 - get_pcp_group,
9 -)
10 -from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase
11 -
12 -
13 -class DefaultMoERunner(MoERunnerBase):
14 - """
15 - Standard MoE runner implementation for executing Mixture of Experts layers.
16 -
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
21 - - Tensor model parallel and expert parallel operations
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 -
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 -
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 -
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 -
40 - @property
41 - def do_naive_dispatch_combine(self) -> bool:
42 - return (
43 - self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk
44 - )
45 -
46 - def _maybe_dispatch(
47 - self,
48 - layer: torch.nn.Module,
49 - hidden_states: torch.Tensor,
50 - router_logits: torch.Tensor,
51 - ) -> tuple[torch.Tensor, torch.Tensor]:
52 - # For naive dispatch/combine Dp/Ep, dispatch the hidden states and
53 - # router logits to all experts.
54 - # NOTE: this will be removed once all kernels are migrated into the
55 - # MoEKernel framework.
56 - if self.do_naive_dispatch_combine:
57 - res = get_ep_group().dispatch_router_logits(
58 - hidden_states,
59 - router_logits,
60 - self.moe_config.is_sequence_parallel,
61 - )
62 - assert len(res) == 2
63 - hidden_states, router_logits = res
64 -
65 - # NOTE: Similar with DP, PCP also needs dispatch and combine. For
66 - # simplicity, AgRsAll2All was added separately for PCP here. Maybe
67 - # we should modify All2AllManager abstraction to better support PCP.
68 - if self.moe_config.pcp_size > 1:
69 - hidden_states = get_pcp_group().all_gather(
70 - hidden_states,
71 - dim=0,
72 - )
73 - router_logits = get_pcp_group().all_gather(
74 - router_logits,
75 - dim=0,
76 - )
77 -
78 - return hidden_states, router_logits
79 -
80 - def _maybe_combine(
81 - self,
82 - shared_output: torch.Tensor | None,
83 - hidden_states: torch.Tensor,
84 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
85 - if self.do_naive_dispatch_combine:
86 - hidden_states = get_ep_group().combine(
87 - hidden_states, self.moe_config.is_sequence_parallel
88 - )
89 -
90 - if self.moe_config.pcp_size > 1:
91 - hidden_states = get_pcp_group().reduce_scatter(
92 - hidden_states,
93 - dim=0,
94 - )
95 -
96 - if self.shared_experts is not None:
97 - assert shared_output is not None
98 - return shared_output, hidden_states
99 - else:
100 - return hidden_states
101 -
102 - def _forward_impl(
103 - self,
104 - layer: torch.nn.Module,
105 - hidden_states: torch.Tensor,
106 - router_logits: torch.Tensor,
107 - shared_experts_input: torch.Tensor | None,
108 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
109 - # TODO(bnell): parts of the dispatch/combine steps will go away once
110 - # #32567 lands and the remaining kernels are made MKs. The PCP
111 - # code will probably remain
112 - hidden_states, router_logits = self._maybe_dispatch(
113 - layer,
114 - hidden_states,
115 - router_logits,
116 - )
117 -
118 - shared_output, hidden_states = self._apply_quant_method(
119 - layer=layer,
120 - hidden_states=hidden_states,
121 - router_logits=router_logits,
122 - shared_experts_input=shared_experts_input,
123 - )
124 -
125 - return self._maybe_combine(
126 - shared_output,
127 - hidden_states,
128 - )
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+679 -18
@@ -1,44 +1,705 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from abc import ABC, abstractmethod
3 +from collections.abc import Callable
4 +from contextlib import nullcontext
5 +from typing import TYPE_CHECKING
6
7 import torch
8 +import torch.nn.functional as F
9
10 +from vllm.distributed import (
11 + get_ep_group,
12 + 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.model_executor.layers.fused_moe.config import (
21 + FusedMoEConfig,
22 +)
23 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
24 FusedMoEMethodBase,
25 )
26 +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
27 + FusedMoERouter,
28 +)
29 +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
30 + ZeroExpertRouter,
31 +)
32 +from vllm.model_executor.layers.fused_moe.runner.moe_runner_interface import (
33 + MoERunnerInterface,
34 +)
35 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
36 SharedExperts,
37 + SharedExpertsOrder,
38 +)
39 +from vllm.platforms import current_platform
40 +from vllm.utils.torch_utils import (
41 + _USE_LAYERNAME,
42 + LayerName,
43 + direct_register_custom_op,
44 +)
45 +
46 +
47 +def get_layer_from_name(layer_name: str) -> torch.nn.Module:
48 + forward_context: ForwardContext = get_forward_context()
49 + if not _USE_LAYERNAME and 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 LayerName 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 | LayerName
70 +else:
71 + _layer_name_type = LayerName if _USE_LAYERNAME else str
72 +
73 +
74 +@torch.compiler.assume_constant_result
75 +def _resolve_layer_name(layer_name: str | LayerName) -> str:
76 + from torch._library.fake_class_registry import FakeScriptObject
77 +
78 + if isinstance(layer_name, LayerName):
79 + return layer_name.value
80 + elif isinstance(layer_name, FakeScriptObject):
81 + return layer_name.real_obj.value
82 + return layer_name
83 +
84 +
85 +# Note: _moe_forward and _moe_forward_shared should not contain any
86 +# implementation details, They should merely pass along control to
87 +# the runner's '_forward_impl' method.
88 +# These functions should never be called directly since they do not
89 +# include all the functionality of the MoE layer.
90 +def _moe_forward(
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 + layer = get_layer_from_name(_resolve_layer_name(layer_name))
97 + return layer.runner._forward_impl(
98 + layer,
99 + hidden_states,
100 + router_logits,
101 + shared_experts_input,
102 + )
103 +
104 +
105 +def _moe_forward_fake(
106 + hidden_states: torch.Tensor,
107 + router_logits: torch.Tensor,
108 + shared_experts_input: torch.Tensor | None,
109 + layer_name: _layer_name_type,
110 +) -> torch.Tensor:
111 + return torch.empty_like(hidden_states)
112 +
113 +
114 +def _moe_forward_shared(
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 + layer = get_layer_from_name(_resolve_layer_name(layer_name))
121 + return layer.runner._forward_impl(
122 + layer,
123 + hidden_states,
124 + router_logits,
125 + shared_experts_input,
126 + )
127 +
128 +
129 +def _moe_forward_shared_fake(
130 + hidden_states: torch.Tensor,
131 + router_logits: torch.Tensor,
132 + shared_experts_input: torch.Tensor | None,
133 + layer_name: _layer_name_type,
134 +) -> tuple[torch.Tensor, torch.Tensor]:
135 + # Output shapes:
136 + # - fused_out: same as hidden_states (routed experts use transformed size)
137 + # - shared_out: same as shared_experts_input if provided, else same as
138 + # hidden_states
139 + # (For latent MoE: shared experts use original hidden_size, not latent size)
140 + fused_out = torch.empty_like(hidden_states)
141 + if shared_experts_input is not None:
142 + shared_out = torch.empty_like(shared_experts_input)
143 + else:
144 + shared_out = torch.empty_like(hidden_states)
145 + return shared_out, fused_out
146 +
147 +
148 +direct_register_custom_op(
149 + op_name="moe_forward",
150 + op_func=_moe_forward,
151 + mutates_args=["hidden_states"],
152 + fake_impl=_moe_forward_fake,
153 + tags=(torch.Tag.needs_fixed_stride_order,),
154 )
155
156
15 -class MoERunner(ABC):
157 +direct_register_custom_op(
158 + op_name="moe_forward_shared",
159 + op_func=_moe_forward_shared,
160 + fake_impl=_moe_forward_shared_fake,
161 + tags=(torch.Tag.needs_fixed_stride_order,),
162 +)
163 +
164 +
165 +def _unpack(
166 + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
167 +) -> tuple[torch.Tensor | None, torch.Tensor]:
168 + if isinstance(result, tuple):
169 + return result
170 + else:
171 + return (None, result)
172 +
173 +
174 +class MoERunner(MoERunnerInterface):
175 """
17 - Abstract base class for Mixture of Experts (MoE) runners.
176 + Standard MoE runner implementation for executing Mixture of Experts layers.
177 +
178 + This is the primary concrete implementation of MoE execution logic, providing
179 + comprehensive support for standard MoE operations. It handles:
180 + - Expert routing and token dispatching using various routing strategies
181 + - Shared experts computation with optional parallel execution using CUDA streams
182 + - Tensor model parallel and expert parallel operations
183 + - Multiple quantization methods and optimized kernel selection
184 + - Both monolithic and decomposed expert execution paths
185 + - Integration with various parallel execution modes (TP, EP, DP)
186
19 - This class defines the interface that all MoE runner implementations must follow.
20 - MoE runners are responsible for executing the forward pass of MoE layers, handling
21 - expert routing, and managing tensor parallel operations.
187 + The runner orchestrates the complete MoE forward pass including routing tokens
188 + to experts, executing expert computations in parallel, and combining results.
189 + It supports advanced features like overlapped execution of shared experts,
190 + optimized kernels for different parallel configurations, and seamless
191 + integration with vLLM's distributed execution framework.
192 +
193 + Eventually, this class may be split into more specialized implementations
194 + for different configurations (e.g., with/without shared experts, gates, etc.).
195 """
196
24 - @abstractmethod
197 + def __init__(
198 + self,
199 + layer_name: str,
200 + moe_config: FusedMoEConfig,
201 + router: FusedMoERouter,
202 + routed_input_transform: torch.nn.Module | None,
203 + gate: torch.nn.Module | None,
204 + shared_experts: torch.nn.Module | None,
205 + quant_method: FusedMoEMethodBase,
206 + enable_dbo: bool,
207 + routed_output_transform: torch.nn.Module | None = None,
208 + routed_scaling_factor: float = 1.0,
209 + ):
210 + super().__init__()
211 + self.moe_config = moe_config
212 + self.router = router
213 + self.routed_input_transform = routed_input_transform
214 + self.routed_output_transform = routed_output_transform
215 + self.routed_scaling_factor = routed_scaling_factor
216 + self.gate = gate
217 + self.quant_method = quant_method
218 + self.enable_dbo = enable_dbo
219 +
220 + self._shared_experts: SharedExperts | None = None
221 + if shared_experts is not None:
222 + self._shared_experts = SharedExperts(
223 + shared_experts,
224 + moe_config=moe_config,
225 + # Note: For now we must pass quant_method along to SharedExperts so it
226 + # can property determine where the shared experts are supposed to be
227 + # called, i.e. by a MK or by the MoERunner.
228 + # Once the MK can be created upfront, we can just pass in the proper
229 + # flags derived from the quant_method's MK.
230 + quant_method=quant_method,
231 + enable_dbo=enable_dbo,
232 + )
233 +
234 + # Needed for string -> FusedMoE layer lookup in custom ops.
235 + self.layer_name = layer_name
236 +
237 + self._forward_entry = self._select_forward()
238 +
239 + def _select_forward(self) -> Callable:
240 + if current_platform.is_tpu() or current_platform.is_cpu():
241 + # TODO: Once the OOM issue for the TPU backend is resolved, we
242 + # will switch to using the moe_forward custom op.
243 + # Note: CPU doesn't require wrapped _forward_impl.
244 + return _moe_forward if self._shared_experts is None else _moe_forward_shared
245 +
246 + return (
247 + torch.ops.vllm.moe_forward
248 + if self._shared_experts is None
249 + else torch.ops.vllm.moe_forward_shared
250 + )
251 +
252 + @property
253 + def shared_experts(self) -> SharedExperts | None:
254 + return self._shared_experts
255 +
256 + # TODO(bnell): temporary hack, do not call this method.
257 + def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
258 + if self._shared_experts is not None:
259 + self._shared_experts._quant_method = quant_method
260 + self.quant_method = quant_method
261 +
262 + def is_internal_router(self) -> bool:
263 + return self.gate is not None
264 +
265 + def apply_routed_input_transform(
266 + self, hidden_states: torch.Tensor
267 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
268 + """Apply transform for routed experts (e.g., latent projection).
269 +
270 + This is called by FusedMoE.forward_native. The original hidden_states
271 + is saved separately so shared experts get [S, hidden_size] while
272 + routed experts get the transformed [S, moe_latent_size].
273 +
274 + Returns (possibly transformed) hidden states and the input for shared
275 + experts (or None if there are no shared experts).
276 + """
277 + if self.routed_input_transform is not None:
278 + result = self.routed_input_transform(hidden_states)
279 + # ReplicatedLinear returns (output, extra_bias) tuple.
280 + # We only need the output tensor; extra_bias is not used here.
281 + if isinstance(result, tuple):
282 + return result[0], hidden_states
283 + return result, hidden_states
284 +
285 + return (
286 + hidden_states,
287 + hidden_states if self._shared_experts is not None else None,
288 + )
289 +
290 + def apply_routed_output_transform(
291 + self,
292 + fused_output: torch.Tensor,
293 + ) -> torch.Tensor:
294 + """Apply transform to routed expert output (e.g., latent to full dim).
295 +
296 + Used by latent MoE models (e.g., NemotronH) where routed experts
297 + operate in a compressed latent space and need projection back to
298 + the full hidden dimension before combining with shared expert output.
299 + """
300 + if self.routed_output_transform is not None:
301 + r = self.routed_output_transform(fused_output)
302 + fused_output = r[0] if isinstance(r, tuple) else r
303 + return fused_output
304 +
305 + def _maybe_apply_routed_scale_to_output(
306 + self,
307 + shared_output: torch.Tensor | None,
308 + fused_output: torch.Tensor,
309 + ) -> tuple[torch.Tensor | None, torch.Tensor]:
310 + """Apply routed_scaling_factor to the output with FP16 overflow
311 + protection.
312 +
313 + Scale the fused expert output by routed_scaling_factor. For FP16,
314 + avoid overflow by dividing shared_output by the scale instead
315 + (the decoder layer compensates with matching divisions).
316 + """
317 + if self.routed_scaling_factor != 1.0:
318 + if fused_output.dtype != torch.float16 or shared_output is None:
319 + fused_output *= self.routed_scaling_factor
320 + elif shared_output is not None:
321 + shared_output *= 1.0 / self.routed_scaling_factor
322 + return shared_output, fused_output
323 +
324 + @property
325 + def _fused_output_is_reduced(self) -> bool:
326 + return (
327 + self.quant_method.moe_kernel is not None
328 + and self.quant_method.moe_kernel.output_is_reduced()
329 + )
330 +
331 + def _maybe_reduce_shared_expert_output(
332 + self,
333 + shared_output: torch.Tensor | None,
334 + ) -> torch.Tensor | None:
335 + """All-reduce shared expert output when the combine kernel already
336 + reduced fused output.
337 +
338 + This is the "early" all-reduce path. When the combine kernel produces
339 + already-reduced fused output, shared output must be reduced separately
340 + to match.
341 + """
342 + if shared_output is not None and self._fused_output_is_reduced:
343 + shared_output = tensor_model_parallel_all_reduce(shared_output)
344 + return shared_output
345 +
346 + def _maybe_reduce_final_output(
347 + self,
348 + states: torch.Tensor,
349 + trunc_size: int,
350 + ) -> torch.Tensor:
351 + """Truncate padded dimensions and all-reduce the combined output.
352 +
353 + This is the "late" all-reduce path. When neither fused nor shared
354 + output was individually reduced, the combined sum is all-reduced
355 + here. Skipped when sequence-parallel is active (SP handles its
356 + own reduction) or when the early path already reduced both outputs.
357 + """
358 + # We don't need to reduce the final output if:
359 + # - We are not running with TP or DP
360 + # - The MK already reduced the fused output itself.
361 + if (
362 + not self.moe_config.is_sequence_parallel
363 + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
364 + and not self._fused_output_is_reduced
365 + ):
366 + states = tensor_model_parallel_all_reduce(states)
367 +
368 + return states[..., :trunc_size]
369 +
370 + def _encode_layer_name(self) -> str | LayerName:
371 + if _USE_LAYERNAME:
372 + return LayerName(self.layer_name)
373 + # Can be unavailable or None in unittests
374 + if (
375 + is_forward_context_available()
376 + and get_forward_context().all_moe_layers is not None
377 + ):
378 + return "from_forward_context"
379 + return self.layer_name
380 +
381 + def _maybe_pad_hidden_states(
382 + self,
383 + shared_experts_input: torch.Tensor | None,
384 + hidden_states: torch.Tensor,
385 + ) -> tuple[torch.Tensor, int]:
386 + """Pad hidden_states to moe_config.hidden_dim and compute the
387 + original dimension for later truncation.
388 +
389 + For latent MoE, the routed hidden_states may be smaller than
390 + hidden_dim. Padding ensures uniform tensor sizes through the
391 + fused MoE kernel. The returned trunc_size is used by
392 + _maybe_reduce_final_output to strip the padding from the result.
393 + """
394 + shared_experts_hidden_dim = (
395 + shared_experts_input.shape[-1] if shared_experts_input is not None else 0
396 + )
397 + transformed_hidden_dim = hidden_states.shape[-1]
398 + if (
399 + not self.quant_method.skip_forward_padding
400 + and self.moe_config.hidden_dim != transformed_hidden_dim
401 + ):
402 + hidden_states = F.pad(
403 + hidden_states,
404 + (0, self.moe_config.hidden_dim - transformed_hidden_dim),
405 + mode="constant",
406 + value=0.0,
407 + )
408 +
409 + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0:
410 + orig_hidden_dims = shared_experts_hidden_dim
411 + else:
412 + orig_hidden_dims = transformed_hidden_dim
413 +
414 + return hidden_states, orig_hidden_dims
415 +
416 + def _maybe_apply_shared_experts(
417 + self,
418 + shared_experts_input: torch.Tensor | None,
419 + order: SharedExpertsOrder,
420 + ):
421 + if self._shared_experts is not None:
422 + assert shared_experts_input is not None
423 + self._shared_experts.apply(shared_experts_input, order)
424 +
425 + def _apply_quant_method(
426 + self,
427 + layer: torch.nn.Module,
428 + hidden_states: torch.Tensor,
429 + router_logits: torch.Tensor,
430 + shared_experts_input: torch.Tensor | None,
431 + ) -> tuple[torch.Tensor | None, torch.Tensor]:
432 + """Run expert routing and the fused MoE kernel via the quant method.
433 +
434 + Orchestrates shared expert execution (before/after), expert selection
435 + via the router, and the actual fused MoE computation. Returns
436 + (shared_expert_output, fused_expert_output).
437 + """
438 + self._maybe_apply_shared_experts(
439 + shared_experts_input, SharedExpertsOrder.NO_OVERLAP
440 + )
441 +
442 + if self.quant_method.is_monolithic:
443 + fused_out = self.quant_method.apply_monolithic(
444 + layer=layer,
445 + x=hidden_states,
446 + router_logits=router_logits,
447 + )
448 + else:
449 + topk_weights, topk_ids = self.router.select_experts(
450 + hidden_states=hidden_states,
451 + router_logits=router_logits,
452 + )
453 +
454 + # Passing shared_experts_input in case SharedExpertsOrder is
455 + # MK_INTERNAL_OVERLAPPED.
456 + fused_out = self.quant_method.apply(
457 + layer=layer,
458 + x=hidden_states,
459 + topk_weights=topk_weights,
460 + topk_ids=topk_ids,
461 + shared_experts_input=shared_experts_input,
462 + )
463 +
464 + self._maybe_apply_shared_experts(
465 + shared_experts_input,
466 + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
467 + )
468 +
469 + return (
470 + self._shared_experts.output if self._shared_experts is not None else None,
471 + fused_out,
472 + )
473 +
474 + def _sequence_parallel_context(self):
475 + """Return a context manager for sequence-parallel token
476 + redistribution.
477 +
478 + When sequence parallelism is active, returns a context that handles
479 + local size tracking for proper token scatter/gather. Otherwise
480 + returns a no-op context.
481 + """
482 + ctx = get_forward_context()
483 + return (
484 + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
485 + if ctx.dp_metadata
486 + else nullcontext()
487 + )
488 +
489 + def _maybe_sync_shared_experts_stream(
490 + self,
491 + shared_experts_input: torch.Tensor | None,
492 + ):
493 + # If router/gate provided, then apply it here.
494 + # (Note: This code runs only when "overlapped mode" is on to allow
495 + # parallel execution of shared experts with the FusedMoE via
496 + # separate cuda stream)
497 + if self._shared_experts is not None:
498 + assert shared_experts_input is not None
499 + self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
500 +
501 + def _maybe_add_zero_expert_output(
502 + self,
503 + result: torch.Tensor,
504 + ) -> torch.Tensor:
505 + """Add the zero expert's contribution to the final result.
506 +
507 + When a ZeroExpertRouter is used, it computes a bias-like output
508 + from the "zero expert" that is added to the combined routed+shared
509 + expert output.
510 + """
511 + if isinstance(self.router, ZeroExpertRouter):
512 + zero_expert_output = self.router.zero_expert_output
513 + assert zero_expert_output is not None
514 + result = result + zero_expert_output
515 + return result
516 +
517 def forward(
518 self,
519 hidden_states: torch.Tensor,
520 router_logits: torch.Tensor,
521 ) -> torch.Tensor:
30 - raise NotImplementedError
522 + """Invoke the fused moe layer.
523
32 - @abstractmethod
33 - def is_internal_router(self) -> bool:
34 - raise NotImplementedError
524 + Input:
525 + - hidden_states
526 + - router_logits
527 +
528 + Output:
529 + - The new hidden_states.
530 +
531 + Calling sequence
532 + - forward
533 + - self._forward_entry (_moe_forward or _moe_forward_shared custom op)
534 + - _forward_impl
535 +
536 + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
537 + to the following reason:
538 + 1. pytorch cannot handle union types in custom op signatures so
539 + _moe_forward and _moe_forward_shared must be split.
540 + """
541 +
542 + # Apply transform for routed experts (e.g., latent projection
543 + # for latent MoE)
544 + hidden_states, shared_experts_input = self.apply_routed_input_transform(
545 + hidden_states
546 + )
547 +
548 + hidden_states, og_hidden_dim = self._maybe_pad_hidden_states(
549 + shared_experts_input,
550 + hidden_states,
551 + )
552 +
553 + result = self._forward_entry(
554 + hidden_states,
555 + router_logits,
556 + shared_experts_input,
557 + self._encode_layer_name(),
558 + )
559 +
560 + #
561 + # Note: there are two all-reduce points below. They are mutually
562 + # exclusive, controlled by _fused_output_is_reduced
563 + # - When True: the combine kernel already reduced fused_output,
564 + # so we reduce shared_output here to match, then skip the
565 + # all-reduce in _maybe_reduce_final_output.
566 + # - When False: neither output is reduced yet, so we combine
567 + # them first and all-reduce the sum in _maybe_reduce_final_output.
568 +
569 + # Extract outputs from result
570 + shared_output, fused_output = _unpack(result)
571 +
572 + # If combine kernel already reduced fused, reduce shared to match.
573 + # See note above re: the two all-reduce points.
574 + shared_output = self._maybe_reduce_shared_expert_output(shared_output)
575 +
576 + shared_output, fused_output = self._maybe_apply_routed_scale_to_output(
577 + shared_output, fused_output
578 + )
579 +
580 + # Apply output transform (e.g. latent -> full dim)
581 + fused_output = self.apply_routed_output_transform(fused_output)
582 +
583 + if shared_output is not None:
584 + result = shared_output + fused_output
585 + else:
586 + result = fused_output
587 +
588 + result = self._maybe_reduce_final_output(result, og_hidden_dim)
589 +
590 + return self._maybe_add_zero_expert_output(result)
591
592 @property
37 - @abstractmethod
38 - def shared_experts(self) -> SharedExperts | None:
39 - raise NotImplementedError
593 + def do_naive_dispatch_combine(self) -> bool:
594 + return (
595 + self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk
596 + )
597
41 - # TODO(bnell): temporary hack, do not call this method.
42 - @abstractmethod
43 - def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
44 - raise NotImplementedError
598 + def _maybe_dispatch(
599 + self,
600 + layer: torch.nn.Module,
601 + hidden_states: torch.Tensor,
602 + router_logits: torch.Tensor,
603 + ) -> tuple[torch.Tensor, torch.Tensor]:
604 + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and
605 + # router logits to all experts.
606 + # NOTE: this will be removed once all kernels are migrated into the
607 + # MoEKernel framework.
608 + if self.do_naive_dispatch_combine:
609 + result = get_ep_group().dispatch_router_logits(
610 + hidden_states,
611 + router_logits,
612 + self.moe_config.is_sequence_parallel,
613 + )
614 + assert len(result) == 2
615 + hidden_states, router_logits = result
616 +
617 + # NOTE: Similar with DP, PCP also needs dispatch and combine. For
618 + # simplicity, AgRsAll2All was added separately for PCP here. Maybe
619 + # we should modify All2AllManager abstraction to better support PCP.
620 + if self.moe_config.pcp_size > 1:
621 + hidden_states = get_pcp_group().all_gather(
622 + hidden_states,
623 + dim=0,
624 + )
625 + router_logits = get_pcp_group().all_gather(
626 + router_logits,
627 + dim=0,
628 + )
629 +
630 + return hidden_states, router_logits
631 +
632 + def _maybe_combine(
633 + self,
634 + shared_output: torch.Tensor | None,
635 + hidden_states: torch.Tensor,
636 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
637 + if self.do_naive_dispatch_combine:
638 + hidden_states = get_ep_group().combine(
639 + hidden_states, self.moe_config.is_sequence_parallel
640 + )
641 +
642 + if self.moe_config.pcp_size > 1:
643 + hidden_states = get_pcp_group().reduce_scatter(
644 + hidden_states,
645 + dim=0,
646 + )
647 +
648 + if self.shared_experts is not None:
649 + assert shared_output is not None
650 + return shared_output, hidden_states
651 + else:
652 + return hidden_states
653 +
654 + def _forward_impl(
655 + self,
656 + layer: torch.nn.Module,
657 + hidden_states: torch.Tensor,
658 + router_logits: torch.Tensor,
659 + shared_experts_input: torch.Tensor | None,
660 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
661 + """Entry point called by the custom op to run the MoE computation.
662 +
663 + Handles pre-dispatch setup (gate application, external shared expert
664 + triggering, quant config init) then performs the following steps
665 + within the sequence-parallel context.
666 +
667 + - Performs expert routing
668 + - fused MoE kernel execution
669 + - shared expert computation.
670 +
671 + Returns a single tensor of combined fused and shared output (if present).
672 + """
673 + # TODO(bnell): this can be removed after MK migration is complete.
674 + layer.ensure_moe_quant_config_init()
675 +
676 + # Sync aux and main stream for shared expert multi-stream overlap.
677 + self._maybe_sync_shared_experts_stream(shared_experts_input)
678 +
679 + # If the Runner holds the gate, apply it after the stream sync,
680 + # so it can run overlapped with the
681 + # NOTE: in future PR, MoE runner will always hold the gate.
682 + if self.gate is not None:
683 + router_logits, _ = self.gate(hidden_states)
684 +
685 + with self._sequence_parallel_context():
686 + # TODO(bnell): parts of the dispatch/combine steps will go away once
687 + # #32567 lands and the remaining kernels are made MKs. The PCP
688 + # code will probably remain
689 + hidden_states, router_logits = self._maybe_dispatch(
690 + layer,
691 + hidden_states,
692 + router_logits,
693 + )
694 +
695 + shared_output, hidden_states = self._apply_quant_method(
696 + layer=layer,
697 + hidden_states=hidden_states,
698 + router_logits=router_logits,
699 + shared_experts_input=shared_experts_input,
700 + )
701 +
702 + return self._maybe_combine(
703 + shared_output,
704 + hidden_states,
705 + )
vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py deleted
-639
@@ -1,639 +0,0 @@
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.router.zero_expert_router import (
29 - ZeroExpertRouter,
30 -)
31 -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
32 -from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
33 - SharedExperts,
34 - SharedExpertsOrder,
35 -)
36 -from vllm.platforms import current_platform
37 -from vllm.utils.torch_utils import (
38 - _USE_LAYERNAME,
39 - LayerName,
40 - direct_register_custom_op,
41 -)
42 -
43 -
44 -def get_layer_from_name(layer_name: str) -> torch.nn.Module:
45 - forward_context: ForwardContext = get_forward_context()
46 - if not _USE_LAYERNAME and layer_name == "from_forward_context":
47 - all_moe_layers = forward_context.all_moe_layers
48 - assert all_moe_layers is not None
49 - moe_layer_index = forward_context.moe_layer_index
50 - if moe_layer_index >= len(all_moe_layers):
51 - raise AssertionError(
52 - "We expected the number of MOE layers in `all_moe_layers` "
53 - "to be equal to the number of "
54 - "{vllm.moe_forward, vllm.moe_forward_shared} calls."
55 - )
56 - layer_name = all_moe_layers[moe_layer_index]
57 - forward_context.moe_layer_index += 1
58 - return forward_context.no_compile_layers[layer_name]
59 -
60 -
61 -# On torch >= 2.11, layer_name is a hoisted LayerName opaque object;
62 -# on older versions it remains a plain str.
63 -if TYPE_CHECKING:
64 - from typing import TypeAlias
65 -
66 - _layer_name_type: TypeAlias = str | LayerName
67 -else:
68 - _layer_name_type = LayerName if _USE_LAYERNAME else str
69 -
70 -
71 -@torch.compiler.assume_constant_result
72 -def _resolve_layer_name(layer_name: str | LayerName) -> str:
73 - from torch._library.fake_class_registry import FakeScriptObject
74 -
75 - if isinstance(layer_name, LayerName):
76 - return layer_name.value
77 - elif isinstance(layer_name, FakeScriptObject):
78 - return layer_name.real_obj.value
79 - return layer_name
80 -
81 -
82 -# Note: _moe_forward and _moe_forward_shared should not contain any
83 -# implementation details, They should merely pass along control to
84 -# the runner's '_forward_dispatch' method.
85 -# These functions should never be called directly since they do not
86 -# include all the functionality of the MoE layer.
87 -def _moe_forward(
88 - hidden_states: torch.Tensor,
89 - router_logits: torch.Tensor,
90 - shared_experts_input: torch.Tensor | None,
91 - layer_name: _layer_name_type,
92 -) -> torch.Tensor:
93 - layer = get_layer_from_name(_resolve_layer_name(layer_name))
94 - return layer.runner._forward_dispatch(
95 - layer,
96 - hidden_states,
97 - router_logits,
98 - shared_experts_input,
99 - )
100 -
101 -
102 -def _moe_forward_fake(
103 - hidden_states: torch.Tensor,
104 - router_logits: torch.Tensor,
105 - shared_experts_input: torch.Tensor | None,
106 - layer_name: _layer_name_type,
107 -) -> torch.Tensor:
108 - return torch.empty_like(hidden_states)
109 -
110 -
111 -def _moe_forward_shared(
112 - hidden_states: torch.Tensor,
113 - router_logits: torch.Tensor,
114 - shared_experts_input: torch.Tensor | None,
115 - layer_name: _layer_name_type,
116 -) -> tuple[torch.Tensor, torch.Tensor]:
117 - layer = get_layer_from_name(_resolve_layer_name(layer_name))
118 - return layer.runner._forward_dispatch(
119 - layer,
120 - hidden_states,
121 - router_logits,
122 - shared_experts_input,
123 - )
124 -
125 -
126 -def _moe_forward_shared_fake(
127 - hidden_states: torch.Tensor,
128 - router_logits: torch.Tensor,
129 - shared_experts_input: torch.Tensor | None,
130 - layer_name: _layer_name_type,
131 -) -> tuple[torch.Tensor, torch.Tensor]:
132 - # Output shapes:
133 - # - fused_out: same as hidden_states (routed experts use transformed size)
134 - # - shared_out: same as shared_experts_input if provided, else same as
135 - # hidden_states
136 - # (For latent MoE: shared experts use original hidden_size, not latent size)
137 - fused_out = torch.empty_like(hidden_states)
138 - if shared_experts_input is not None:
139 - shared_out = torch.empty_like(shared_experts_input)
140 - else:
141 - shared_out = torch.empty_like(hidden_states)
142 - return shared_out, fused_out
143 -
144 -
145 -direct_register_custom_op(
146 - op_name="moe_forward",
147 - op_func=_moe_forward,
148 - mutates_args=["hidden_states"],
149 - fake_impl=_moe_forward_fake,
150 - tags=(torch.Tag.needs_fixed_stride_order,),
151 -)
152 -
153 -
154 -direct_register_custom_op(
155 - op_name="moe_forward_shared",
156 - op_func=_moe_forward_shared,
157 - fake_impl=_moe_forward_shared_fake,
158 - tags=(torch.Tag.needs_fixed_stride_order,),
159 -)
160 -
161 -
162 -def _unpack(
163 - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
164 -) -> tuple[torch.Tensor | None, torch.Tensor]:
165 - if isinstance(result, tuple):
166 - return result
167 - else:
168 - return (None, result)
169 -
170 -
171 -class MoERunnerBase(MoERunner):
172 - """
173 - Abstract base class providing common functionality for MoE runner implementations.
174 -
175 - This class serves as the foundation for concrete MoE runner implementations by
176 - providing shared state management and common utilities. It handles:
177 - - Common initialization and configuration management
178 - - Shared expert output reduction logic for tensor parallel scenarios
179 - - Base methods for tensor model parallel reductions
180 - - Common properties and utility functions used across different runner types
181 -
182 - Concrete subclasses must implement the abstract methods to define their specific
183 - execution strategies, such as standard execution, chunked processing, or other
184 - specialized approaches. The base class provides the infrastructure while
185 - allowing flexibility in the actual MoE computation implementation.
186 -
187 - Key abstract methods that subclasses must implement:
188 - - _forward_impl: The core MoE computation logic specific to each runner type
189 - """
190 -
191 - def __init__(
192 - self,
193 - layer_name: str,
194 - moe_config: FusedMoEConfig,
195 - router: FusedMoERouter,
196 - routed_input_transform: torch.nn.Module | None,
197 - gate: torch.nn.Module | None,
198 - shared_experts: torch.nn.Module | None,
199 - quant_method: FusedMoEMethodBase,
200 - enable_dbo: bool,
201 - routed_output_transform: torch.nn.Module | None = None,
202 - routed_scaling_factor: float = 1.0,
203 - ):
204 - super().__init__()
205 - self.moe_config = moe_config
206 - self.router = router
207 - self.routed_input_transform = routed_input_transform
208 - self.routed_output_transform = routed_output_transform
209 - self.routed_scaling_factor = routed_scaling_factor
210 - self.gate = gate
211 - self.quant_method = quant_method
212 - self.enable_dbo = enable_dbo
213 - self._fused_output_is_reduced = (
214 - self.quant_method.moe_kernel is not None
215 - and self.quant_method.moe_kernel.output_is_reduced()
216 - )
217 -
218 - self._shared_experts: SharedExperts | None = None
219 - if shared_experts is not None:
220 - self._shared_experts = SharedExperts(
221 - shared_experts,
222 - moe_config=moe_config,
223 - # Note: For now we must pass quant_method along to SharedExperts so it
224 - # can property determine where the shared experts are supposed to be
225 - # called, i.e. by a MK or by the MoERunner.
226 - # Once the MK can be created upfront, we can just pass in the proper
227 - # flags derived from the quant_method's MK.
228 - quant_method=quant_method,
229 - enable_dbo=enable_dbo,
230 - )
231 -
232 - # Needed for string -> FusedMoE layer lookup in custom ops.
233 - self.layer_name = layer_name
234 -
235 - self._forward_entry = self._select_forward()
236 -
237 - def _select_forward(self) -> Callable:
238 - if current_platform.is_tpu() or current_platform.is_cpu():
239 - # TODO: Once the OOM issue for the TPU backend is resolved, we
240 - # will switch to using the moe_forward custom op.
241 - # Note: CPU doesn't require wrapped _forward_impl.
242 - return _moe_forward if self._shared_experts is None else _moe_forward_shared
243 -
244 - return (
245 - torch.ops.vllm.moe_forward
246 - if self._shared_experts is None
247 - else torch.ops.vllm.moe_forward_shared
248 - )
249 -
250 - @property
251 - def shared_experts(self) -> SharedExperts | None:
252 - return self._shared_experts
253 -
254 - # TODO(bnell): temporary hack, do not call this method.
255 - def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
256 - if self._shared_experts is not None:
257 - self._shared_experts._quant_method = quant_method
258 - self.quant_method = quant_method
259 -
260 - def is_internal_router(self) -> bool:
261 - return self.gate is not None
262 -
263 - def apply_routed_input_transform(
264 - self, hidden_states: torch.Tensor
265 - ) -> tuple[torch.Tensor, torch.Tensor | None]:
266 - """Apply transform for routed experts (e.g., latent projection).
267 -
268 - This is called by FusedMoE.forward_native. The original hidden_states
269 - is saved separately so shared experts get [S, hidden_size] while
270 - routed experts get the transformed [S, moe_latent_size].
271 -
272 - Returns (possibly transformed) hidden states and the input for shared
273 - experts (or None if there are no shared experts).
274 - """
275 - if self.routed_input_transform is not None:
276 - result = self.routed_input_transform(hidden_states)
277 - # ReplicatedLinear returns (output, extra_bias) tuple.
278 - # We only need the output tensor; extra_bias is not used here.
279 - if isinstance(result, tuple):
280 - return result[0], hidden_states
281 - return result, hidden_states
282 -
283 - return (
284 - hidden_states,
285 - hidden_states if self._shared_experts is not None else None,
286 - )
287 -
288 - def apply_routed_output_transform(
289 - self,
290 - fused_output: torch.Tensor,
291 - ) -> torch.Tensor:
292 - """Apply transform to routed expert output (e.g., latent to full dim).
293 -
294 - Used by latent MoE models (e.g., NemotronH) where routed experts
295 - operate in a compressed latent space and need projection back to
296 - the full hidden dimension before combining with shared expert output.
297 - """
298 - if self.routed_output_transform is not None:
299 - r = self.routed_output_transform(fused_output)
300 - fused_output = r[0] if isinstance(r, tuple) else r
301 - return fused_output
302 -
303 - def _maybe_apply_routed_scale_to_output(
304 - self,
305 - shared_output: torch.Tensor | None,
306 - fused_output: torch.Tensor,
307 - ) -> tuple[torch.Tensor | None, torch.Tensor]:
308 - """Apply routed_scaling_factor to the output with FP16 overflow
309 - protection.
310 -
311 - Scale the fused expert output by routed_scaling_factor. For FP16,
312 - avoid overflow by dividing shared_output by the scale instead
313 - (the decoder layer compensates with matching divisions).
314 - """
315 - if self.routed_scaling_factor != 1.0:
316 - if fused_output.dtype != torch.float16:
317 - fused_output *= self.routed_scaling_factor
318 - elif shared_output is not None:
319 - shared_output *= 1.0 / self.routed_scaling_factor
320 - return shared_output, fused_output
321 -
322 - def _maybe_reduce_shared_expert_output(
323 - self,
324 - shared_output: torch.Tensor | None,
325 - ) -> torch.Tensor | None:
326 - """All-reduce shared expert output when the combine kernel already
327 - reduced fused output.
328 -
329 - This is the "early" all-reduce path. When the combine kernel produces
330 - already-reduced fused output, shared output must be reduced separately
331 - to match.
332 - """
333 - if self._fused_output_is_reduced:
334 - assert shared_output is not None
335 - shared_output = tensor_model_parallel_all_reduce(shared_output)
336 - return shared_output
337 -
338 - def _maybe_reduce_final_output(
339 - self,
340 - states: torch.Tensor,
341 - trunc_size: int,
342 - ) -> torch.Tensor:
343 - """Truncate padded dimensions and all-reduce the combined output.
344 -
345 - This is the "late" all-reduce path. When neither fused nor shared
346 - output was individually reduced, the combined sum is all-reduced
347 - here. Skipped when sequence-parallel is active (SP handles its
348 - own reduction) or when the early path already reduced both outputs.
349 - """
350 - # We don't need to reduce the final output if:
351 - # - We are not running with TP or DP
352 - # - The MK already reduced the fused output itself.
353 - if (
354 - not self.moe_config.is_sequence_parallel
355 - and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
356 - and not self._fused_output_is_reduced
357 - ):
358 - states = tensor_model_parallel_all_reduce(states)
359 -
360 - return states[..., :trunc_size]
361 -
362 - def _encode_layer_name(self) -> str | LayerName:
363 - if _USE_LAYERNAME:
364 - return LayerName(self.layer_name)
365 - # Can be unavailable or None in unittests
366 - if (
367 - is_forward_context_available()
368 - and get_forward_context().all_moe_layers is not None
369 - ):
370 - return "from_forward_context"
371 - return self.layer_name
372 -
373 - def _maybe_pad_hidden_states(
374 - self,
375 - shared_experts_input: torch.Tensor | None,
376 - hidden_states: torch.Tensor,
377 - ) -> tuple[torch.Tensor, int]:
378 - """Pad hidden_states to moe_config.hidden_dim and compute the
379 - original dimension for later truncation.
380 -
381 - For latent MoE, the routed hidden_states may be smaller than
382 - hidden_dim. Padding ensures uniform tensor sizes through the
383 - fused MoE kernel. The returned trunc_size is used by
384 - _maybe_reduce_final_output to strip the padding from the result.
385 - """
386 - shared_experts_hidden_dim = (
387 - shared_experts_input.shape[-1] if shared_experts_input is not None else 0
388 - )
389 - transformed_hidden_dim = hidden_states.shape[-1]
390 - if (
391 - not self.quant_method.skip_forward_padding
392 - and self.moe_config.hidden_dim != transformed_hidden_dim
393 - ):
394 - hidden_states = F.pad(
395 - hidden_states,
396 - (0, self.moe_config.hidden_dim - transformed_hidden_dim),
397 - mode="constant",
398 - value=0.0,
399 - )
400 -
401 - if self.routed_output_transform is not None and shared_experts_hidden_dim > 0:
402 - orig_hidden_dims = shared_experts_hidden_dim
403 - else:
404 - orig_hidden_dims = transformed_hidden_dim
405 -
406 - return hidden_states, orig_hidden_dims
407 -
408 - def _maybe_apply_shared_experts(
409 - self,
410 - shared_experts_input: torch.Tensor | None,
411 - order: SharedExpertsOrder,
412 - ):
413 - if self._shared_experts is not None:
414 - assert shared_experts_input is not None
415 - self._shared_experts.apply(shared_experts_input, order)
416 -
417 - def _apply_quant_method(
418 - self,
419 - layer: torch.nn.Module,
420 - hidden_states: torch.Tensor,
421 - router_logits: torch.Tensor,
422 - shared_experts_input: torch.Tensor | None,
423 - ) -> tuple[torch.Tensor | None, torch.Tensor]:
424 - """Run expert routing and the fused MoE kernel via the quant method.
425 -
426 - Orchestrates shared expert execution (before/after), expert selection
427 - via the router, and the actual fused MoE computation. Returns
428 - (shared_expert_output, fused_expert_output).
429 - """
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 - # 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 - """Return a context manager for sequence-parallel token
468 - redistribution.
469 -
470 - When sequence parallelism is active, returns a context that handles
471 - local size tracking for proper token scatter/gather. Otherwise
472 - returns a no-op context.
473 - """
474 - ctx = get_forward_context()
475 - return (
476 - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
477 - if ctx.dp_metadata
478 - else nullcontext()
479 - )
480 -
481 - def _maybe_sync_shared_experts_stream(
482 - self,
483 - shared_experts_input: torch.Tensor | None,
484 - ):
485 - # If router/gate provided, then apply it here.
486 - # (Note: This code runs only when "overlapped mode" is on to allow
487 - # parallel execution of shared experts with the FusedMoE via
488 - # separate cuda stream)
489 - if self._shared_experts is not None:
490 - assert shared_experts_input is not None
491 - self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
492 -
493 - def _maybe_add_zero_expert_output(
494 - self,
495 - result: torch.Tensor,
496 - ) -> torch.Tensor:
497 - """Add the zero expert's contribution to the final result.
498 -
499 - When a ZeroExpertRouter is used, it computes a bias-like output
500 - from the "zero expert" that is added to the combined routed+shared
501 - expert output.
502 - """
503 - if isinstance(self.router, ZeroExpertRouter):
504 - zero_expert_output = self.router.zero_expert_output
505 - assert zero_expert_output is not None
506 - result = result + zero_expert_output
507 - return result
508 -
509 - def forward(
510 - self,
511 - hidden_states: torch.Tensor,
512 - router_logits: torch.Tensor,
513 - ) -> torch.Tensor:
514 - """Invoke the fused moe layer.
515 -
516 - Input:
517 - - hidden_states
518 - - router_logits
519 -
520 - Output:
521 - - The new hidden_states.
522 -
523 - Calling sequence
524 - - forward
525 - - self._forward_entry (_moe_forward or _moe_forward_shared custom op)
526 - - _forward_dispatch
527 - - _forward_impl
528 -
529 - Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
530 - to the following reasons:
531 - 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by
532 - torch.compile
533 - 2. pytorch cannot handle union types in custom op signatures so
534 - _moe_forward and _moe_forward_shared must be split.
535 -
536 - If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can
537 - potentially get rid of _moe_forward and _moe_forward_shared and collapse the
538 - whole sequence into the 'forward' method.
539 - """
540 -
541 - # Apply transform for routed experts (e.g., latent projection
542 - # for latent MoE)
543 - hidden_states, shared_experts_input = self.apply_routed_input_transform(
544 - hidden_states
545 - )
546 -
547 - hidden_states, og_hidden_dim = self._maybe_pad_hidden_states(
548 - shared_experts_input,
549 - hidden_states,
550 - )
551 -
552 - result = self._forward_entry(
553 - hidden_states,
554 - router_logits,
555 - shared_experts_input,
556 - self._encode_layer_name(),
557 - )
558 -
559 - #
560 - # Note: there are two all-reduce points below. They are mutually
561 - # exclusive, controlled by _fused_output_is_reduced
562 - # - When True: the combine kernel already reduced fused_output,
563 - # so we reduce shared_output here to match, then skip the
564 - # all-reduce in _maybe_reduce_final_output.
565 - # - When False: neither output is reduced yet, so we combine
566 - # them first and all-reduce the sum in _maybe_reduce_final_output.
567 -
568 - # Extract outputs from result
569 - shared_output, fused_output = _unpack(result)
570 -
571 - # If combine kernel already reduced fused, reduce shared to match.
572 - # See note above re: the two all-reduce points.
573 - shared_output = self._maybe_reduce_shared_expert_output(shared_output)
574 -
575 - shared_output, fused_output = self._maybe_apply_routed_scale_to_output(
576 - shared_output, fused_output
577 - )
578 -
579 - # Apply output transform (e.g. latent -> full dim)
580 - fused_output = self.apply_routed_output_transform(fused_output)
581 -
582 - if shared_output is not None:
583 - result = shared_output + fused_output
584 - else:
585 - result = fused_output
586 -
587 - result = self._maybe_reduce_final_output(result, og_hidden_dim)
588 -
589 - return self._maybe_add_zero_expert_output(result)
590 -
591 - def _forward_dispatch(
592 - self,
593 - layer: torch.nn.Module,
594 - hidden_states: torch.Tensor,
595 - router_logits: torch.Tensor,
596 - shared_experts_input: torch.Tensor | None,
597 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
598 - """Entry point called by the custom op to run the MoE computation.
599 -
600 - Handles pre-dispatch setup (gate application, external shared expert
601 - triggering, quant config init) then delegates to _forward_impl within
602 - the sequence-parallel context.
603 - """
604 - # TODO(bnell): this can be removed after MK migration is complete.
605 - layer.ensure_moe_quant_config_init()
606 -
607 - # Sync aux and main stream for shared expert multi-stream overlap.
608 - self._maybe_sync_shared_experts_stream(shared_experts_input)
609 -
610 - # If the Runner holds the gate, apply it after the stream sync,
611 - # so it can run overlapped with the
612 - # NOTE: in future PR, MoE runner will always hold the gate.
613 - if self.gate is not None:
614 - router_logits, _ = self.gate(hidden_states)
615 -
616 - with self._sequence_parallel_context():
617 - return self._forward_impl(
618 - layer,
619 - hidden_states,
620 - router_logits,
621 - shared_experts_input,
622 - )
623 -
624 - @abstractmethod
625 - def _forward_impl(
626 - self,
627 - layer: torch.nn.Module,
628 - hidden_states: torch.Tensor,
629 - router_logits: torch.Tensor,
630 - shared_experts_input: torch.Tensor | None,
631 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
632 - """Core MoE computation to be implemented by subclasses.
633 -
634 - Performs expert routing, fused MoE kernel execution, and shared
635 - expert computation. Returns a single tensor (fused output only)
636 - or a tuple of (shared_output, fused_output) when shared experts
637 - are present.
638 - """
639 - raise NotImplementedError
vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py deleted
-47
@@ -1,47 +0,0 @@
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.default_moe_runner import (
16 - DefaultMoERunner,
17 -)
18 -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
19 -from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
20 - SharedExperts,
21 -)
22 -
23 -
24 -def create_moe_runner(
25 - layer_name: str,
26 - moe_config: FusedMoEConfig,
27 - router: FusedMoERouter,
28 - routed_input_transform: torch.nn.Module | None,
29 - gate: torch.nn.Module | None,
30 - shared_experts: SharedExperts | None,
31 - quant_method: FusedMoEMethodBase,
32 - enable_dbo: bool,
33 - routed_output_transform: torch.nn.Module | None = None,
34 - routed_scaling_factor: float = 1.0,
35 -) -> MoERunner:
36 - return DefaultMoERunner(
37 - layer_name,
38 - moe_config,
39 - router,
40 - routed_input_transform,
41 - gate,
42 - shared_experts,
43 - quant_method,
44 - enable_dbo,
45 - routed_output_transform=routed_output_transform,
46 - routed_scaling_factor=routed_scaling_factor,
47 - )
vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +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 MoERunnerInterface(ABC):
16 + """
17 + Abstract base class for Mixture of Experts (MoE) runners.
18 +
19 + This class defines the interface that all MoE runner implementations must follow.
20 + MoE runners are responsible for executing the forward pass of MoE layers, handling
21 + expert routing, and managing tensor parallel operations.
22 + """
23 +
24 + @abstractmethod
25 + def forward(
26 + self,
27 + hidden_states: torch.Tensor,
28 + router_logits: torch.Tensor,
29 + ) -> torch.Tensor:
30 + raise NotImplementedError
31 +
32 + @abstractmethod
33 + def is_internal_router(self) -> bool:
34 + raise NotImplementedError
35 +
36 + @property
37 + @abstractmethod
38 + def shared_experts(self) -> SharedExperts | None:
39 + raise NotImplementedError
40 +
41 + # TODO(bnell): temporary hack, do not call this method.
42 + @abstractmethod
43 + def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
44 + raise NotImplementedError