From 4aeabf2aa4272fbf5f58ab82f4e87e174999ffc5 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 13 Jan 2026 18:44:16 -0500 Subject: [PATCH 001/108] initial MoERunner refactor Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 + .../fused_moe/runner/default_moe_runner.py | 701 ++++++++++-------- 2 files changed, 406 insertions(+), 297 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 75283b9bbe39..2f704569209c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -504,6 +504,8 @@ def __init__( self.apply_router_weight_on_input = apply_router_weight_on_input self.activation = MoEActivation.from_str(activation) + # TODO(bnell): we should not have to create a router if the kernel is + # monolithic. self.router = create_fused_moe_router( top_k=top_k, global_num_experts=self.global_num_experts, diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 12b560493fa2..bb1ac465b1c5 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable from contextlib import nullcontext from typing import TYPE_CHECKING @@ -82,9 +83,23 @@ def _moe_forward( layer = get_layer_from_name(_resolve_layer_name(layer_name)) # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() - return layer.runner.forward_impl( - layer, hidden_states, router_logits, shared_experts_input - ) + runner = layer.runner + router_logits = runner._maybe_gate(hidden_states, router_logits) + with runner._sequence_parallel_context(): + if runner.use_dp_chunking: + return runner.forward_impl_chunked( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + else: + return runner.forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) def _moe_forward_fake( @@ -105,9 +120,23 @@ def _moe_forward_shared( layer = get_layer_from_name(_resolve_layer_name(layer_name)) # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() - return layer.runner.forward_impl( - layer, hidden_states, router_logits, shared_experts_input - ) + runner = layer.runner + router_logits = runner._maybe_gate(hidden_states, router_logits) + with runner._sequence_parallel_context(): + if runner.use_dp_chunking: + return runner.forward_impl_chunked( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + else: + return runner.forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) def _moe_forward_shared_fake( @@ -191,10 +220,17 @@ def __init__( self.reduce_results = reduce_results self.enable_dbo = enable_dbo + # Chunked all2all staging tensor + # TODO(bnell) rename these? + self.batched_hidden_states: torch.Tensor | None = None + self.batched_router_logits: torch.Tensor | None = None + self._maybe_init_dp_chunking() + # Allow disabling of the separate shared experts stream for # debug purposes. # TODO: Remove this after more extensive testings with TP/DP # and other execution modes + self.use_shared_experts_stream = False if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") self.shared_experts_stream = None @@ -210,24 +246,22 @@ def __init__( # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name + self.moe_forward = self._select_forward(layer) + + def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. # Note: CPU doesn't require wrapped forward_impl. - if self.shared_experts is None: - self.moe_forward = _moe_forward - else: - self.moe_forward = _moe_forward_shared - else: - if self.shared_experts is None: - self.moe_forward = torch.ops.vllm.moe_forward - else: - self.moe_forward = torch.ops.vllm.moe_forward_shared + return _moe_forward if self.shared_experts is None else _moe_forward_shared - # Chunked all2all staging tensor - self.batched_hidden_states: torch.Tensor | None = None - self.batched_router_logits: torch.Tensor | None = None + return ( + torch.ops.vllm.moe_forward + if self.shared_experts is None + else torch.ops.vllm.moe_forward_shared + ) + # TODO(bnell): make this a member var? @property def use_dp_chunking(self) -> bool: return ( @@ -241,22 +275,9 @@ def _maybe_setup_shared_experts_stream( self, hidden_states: torch.Tensor, shared_input: torch.Tensor | None, - has_separate_shared_experts: bool, - use_chunked_impl: bool, - ) -> tuple[bool, torch.Tensor | None]: - use_shared_experts_stream = ( - current_platform.is_cuda() - and has_separate_shared_experts - and not use_chunked_impl - and self.shared_experts_stream is not None - and ( - hidden_states.shape[0] - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) - ) - - shared_experts_input: torch.Tensor | None = None - if use_shared_experts_stream: + ) -> torch.Tensor | None: + hidden_states_clone: torch.Tensor | None = None + if self.use_shared_experts_stream: assert self.shared_experts_stream is not None assert self.moe_config.disable_inplace @@ -278,12 +299,13 @@ def _maybe_setup_shared_experts_stream( assert self.shared_experts_stream is not None self.shared_experts_stream.wait_stream(current_stream()) - return use_shared_experts_stream, shared_experts_input + return shared_experts_input - def ensure_dp_chunking_init(self): - if not self.use_dp_chunking or self.batched_hidden_states is not None: + def _maybe_init_dp_chunking(self): + if not self.use_dp_chunking: return + assert self.batched_hidden_states is None states_shape: tuple[int, ...] logits_shape: tuple[int, ...] @@ -309,6 +331,38 @@ def ensure_dp_chunking_init(self): device=device, ) + @property + def has_separate_shared_experts(self) -> bool: + return ( + not self.quant_method.mk_owns_shared_expert + and self.shared_experts is not None + ) + + def _apply_shared_experts( + self, + hidden_states: torch.Tensor, + allow_streaming: bool = False, + ) -> torch.Tensor | None: + shared_output: torch.Tensor | None = None + if self.has_separate_shared_experts: + assert self.shared_experts is not None + + if self.use_shared_experts_stream and allow_streaming: + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self.shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream + shared_output = self.shared_experts(hidden_states) + current_stream().wait_stream(self.shared_experts_stream) + else: + shared_output = self.shared_experts(hidden_states) + + return shared_output + def must_reduce_shared_expert_outputs(self) -> bool: """ The shared_experts are typically computed using the RowParallelLinear @@ -322,7 +376,6 @@ def must_reduce_shared_expert_outputs(self) -> bool: Therefore it is required that we reduce the shared_experts output early. """ - assert self.quant_method is not None return ( self.quant_method.moe_kernel is not None and self.quant_method.moe_kernel.output_is_reduced() @@ -357,7 +410,7 @@ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Ten return result return hidden_states - def _reduce_output( + def _maybe_reduce_output( self, states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], trunc_sizes: list[int], @@ -397,23 +450,12 @@ def _encode_layer_name(self) -> str | ModuleName: return "from_forward_context" return self.layer_name - def forward( + def _maybe_pad_hidden_states( self, + original_hidden_states: torch.Tensor | None, hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # For latent MoE: save ORIGINAL hidden_states before transform - # (shared_experts need original dimension, routed experts use transformed) - if self.shared_experts is not None: - original_hidden_states = hidden_states - original_hidden_dim = hidden_states.shape[-1] - else: - original_hidden_states = None - - # Apply transform for routed experts (e.g., latent projection for latent MoE) - hidden_states = self.apply_routed_input_transform(hidden_states) - - # This is the dimension after transform (for routed expert output slicing) + ) -> tuple[torch.Tensor, list[int]]: + original_hidden_dim = original_hidden_states.shape[-1] if original_hidden_states is not None else 0 transformed_hidden_dim = hidden_states.shape[-1] if ( not self.quant_method.skip_forward_padding @@ -426,134 +468,269 @@ def forward( value=0.0, ) - fused_output = self.moe_forward( - hidden_states, - router_logits, - original_hidden_states, - self._encode_layer_name(), - ) - if self.shared_experts is not None: orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim] else: orig_hidden_dims = [transformed_hidden_dim] - return self._reduce_output(fused_output, orig_hidden_dims) + return hidden_states, orig_hidden_dims - def forward_impl_chunked( + def _apply_quant_method( self, layer: torch.nn.Module, - full_hidden_states: torch.Tensor, - full_router_logits: torch.Tensor, - full_shared_input: torch.Tensor | None, - has_separate_shared_experts: bool, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + hidden_states: torch.Tensor, + extra_tensor: torch.Tensor | None, + router_logits: torch.Tensor, + shared_input: torch.Tensor | None, + run_shared_experts_before: bool = True, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + shared_output: torch.Tensor | None = None + + # Run this before quant_method to avoid inplace issues. + if run_shared_experts_before: + shared_input = shared_input if shared_input is not None else hidden_states + shared_output = self._apply_shared_experts( + shared_input, + False, + ) + else: + hidden_states_clone = self._maybe_setup_shared_experts_stream( + hidden_states, + shared_input, + ) + + # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014). + # Figure out nicer way to do this. + x_arg = hidden_states if extra_tensor is None else (hidden_states, extra_tensor) + + if self.quant_method.is_monolithic: + result = self.quant_method.apply_monolithic( + layer=layer, + x=x_arg, + router_logits=router_logits, + ) + else: + topk_weights, topk_ids = self.router.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + + result = self.quant_method.apply( + layer=layer, + x=x_arg, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts_input=shared_input, + ) + + if isinstance(result, tuple): + assert shared_output is None + shared_output, hidden_states = result + else: + hidden_states = result + + if not run_shared_experts_before and self.has_separate_shared_experts: + assert shared_output is None + shared_output = self._apply_shared_experts( + hidden_states_clone, + True, + ) + + return shared_output, hidden_states + + def _sequence_parallel_context(self): + ctx = get_forward_context() + return ( + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) + if ctx.dp_metadata + else nullcontext() + ) + + def _allocate_dp_chunking_outputs( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + assert self.use_dp_chunking + + # Assert the inputs are of the proper type and shape. assert self.batched_hidden_states is not None assert self.batched_router_logits is not None - assert self.batched_hidden_states.dtype == full_hidden_states.dtype, ( - f"{self.batched_hidden_states.dtype} == {full_hidden_states.dtype}" + + assert self.batched_hidden_states.dtype == hidden_states.dtype, ( + f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" ) - assert self.batched_router_logits.dtype == full_router_logits.dtype, ( - f"{self.batched_router_logits.dtype} == {full_router_logits.dtype}" + assert self.batched_router_logits.dtype == router_logits.dtype, ( + f"{self.batched_router_logits.dtype} == {router_logits.dtype}" ) - # Check size compatibility. - assert self.batched_hidden_states.size(-1) == full_hidden_states.size(-1) - assert self.batched_router_logits.size(-1) == full_router_logits.size(-1) - # TODO(bnell): Fix shared_expert_inputs w/chunking. - # assert shared_input is None, ( - # "Routed input transform is not currently supported with DP chunking." - # ) + # Check size compatibility. + assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) + assert self.batched_router_logits.size(-1) == router_logits.size(-1) - full_fused_final_hidden_states = torch.empty_like(full_hidden_states) + final_fused_hidden_states = torch.empty_like(hidden_states) if self.shared_experts is not None: - full_shared_final_hidden_states = torch.empty_like(full_hidden_states) - - def process_chunk(chunk_start, chunk_end, skip_result_store=False): - chunk_size = chunk_end - chunk_start - hidden_states = full_hidden_states[chunk_start:chunk_end, :] - router_logits = full_router_logits[chunk_start:chunk_end, :] - shared_input = ( - full_shared_input[chunk_start:chunk_end, :] - if full_shared_input is not None - else None + final_shared_hidden_states = torch.empty_like(hidden_states) + else: + final_shared_hidden_states = None + + return final_shared_hidden_states, final_fused_hidden_states + + def _maybe_gate( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + # If router/gate provided, then apply it here. + # (Note: This code runs only when "overlapped mode" is on to allow + # parallel execution of shared experts with the FusedMoE via + # separate cuda stream) + if self.gate is not None: + router_logits, _ = self.gate(hidden_states) + return router_logits + + @property + def do_naive_dispatch_combine(self) -> bool: + return ( + self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk + ) + + def _maybe_dispatch( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + extra_tensor: torch.Tensor | None = None + + if self.do_naive_dispatch_combine: + post_quant_allgather = ( + self.moe_config.dp_size > 1 + and self.moe_config.use_ep + and getattr(self.quant_method, "do_post_quant_allgather", False) ) - assert self.batched_hidden_states is not None - assert self.batched_router_logits is not None - # This is only true when DBO has been enabled in the config. - # Both tensors will have an outer dimension for the ubatch id - if self.batched_hidden_states.dim() == 3: - assert self.batched_router_logits.dim() == 3 - batch_buffer_idx = dbo_current_ubatch_id() - batched_hidden_states = self.batched_hidden_states[batch_buffer_idx, :] - batched_router_logits = self.batched_router_logits[batch_buffer_idx, :] + extra_tensors: list[torch.Tensor] | None = None + + if post_quant_allgather: + hidden_states_to_dispatch, extra_tensors = ( + self.quant_method.prepare_dp_allgather_tensor( + layer, hidden_states, router_logits + ) + ) else: - batched_hidden_states = self.batched_hidden_states - batched_router_logits = self.batched_router_logits + hidden_states_to_dispatch = hidden_states - assert ( - batched_hidden_states.size(0) # type: ignore - >= chunk_size + result = get_ep_group().dispatch_router_logits( + hidden_states_to_dispatch, + router_logits, + self.moe_config.is_sequence_parallel, + extra_tensors=extra_tensors, ) - assert ( - batched_router_logits.size(0) # type: ignore - >= chunk_size + + if len(result) == 3: + hidden_states, router_logits, extra_tensors = result + assert isinstance(extra_tensors, list) and len(extra_tensors) == 1 + extra_tensor = extra_tensors[0] + else: + hidden_states, router_logits = result + + # NOTE: Similar with DP, PCP also needs dispatch and combine. For + # simplicity, AgRsAll2All was added separately for PCP here. Maybe + # we should modify All2AllManager abstraction to better support PCP. + # TODO(bnell): see what we can do here + if self.moe_config.pcp_size > 1: + hidden_states = get_pcp_group().all_gather( + hidden_states, + dim=0, ) - staged_hidden_states = batched_hidden_states[:chunk_size, :] # type: ignore - staged_router_logits = batched_router_logits[:chunk_size, :] # type: ignore - staged_hidden_states.copy_(hidden_states, non_blocking=True) - staged_router_logits.copy_(router_logits, non_blocking=True) + router_logits = get_pcp_group().all_gather( + router_logits, + dim=0, + ) + + return hidden_states, router_logits, extra_tensor - shared_input = ( - shared_input if shared_input is not None else staged_hidden_states + def _maybe_combine( + self, + shared_output: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: + if self.do_naive_dispatch_combine: + hidden_states = get_ep_group().combine( + hidden_states, self.moe_config.is_sequence_parallel ) - # Matrix multiply. - if self.quant_method.is_monolithic: - assert has_separate_shared_experts or self.shared_experts is None - final_hidden_states = self.quant_method.apply_monolithic( - layer=layer, - x=staged_hidden_states, - router_logits=staged_router_logits, - ) - else: - topk_weights, topk_ids = self.router.select_experts( - hidden_states=staged_hidden_states, - router_logits=staged_router_logits, - ) + if self.moe_config.pcp_size > 1: + hidden_states = get_pcp_group().reduce_scatter( + hidden_states, + dim=0, + ) + # need RS for shared_output? - final_hidden_states = self.quant_method.apply( - layer=layer, - x=staged_hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_experts_input=shared_input, - ) + if self.shared_experts is not None: + assert shared_output is not None + return shared_output, hidden_states + else: + return hidden_states - if has_separate_shared_experts: - assert not isinstance(final_hidden_states, tuple) - assert self.shared_experts is not None + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # For latent MoE: save ORIGINAL hidden_states before transform + # (shared_experts need original dimension, routed experts use transformed) + original_hidden_states = hidden_states - shared_output = self.shared_experts(shared_input) + # Apply transform for routed experts (e.g., latent projection for latent MoE) + hidden_states = self.apply_routed_input_transform(hidden_states) - final_hidden_states = ( - shared_output, - final_hidden_states, - ) + hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + original_hidden_states, + hidden_states, + ) - if not skip_result_store: - if self.shared_experts is None: - full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_( - final_hidden_states, non_blocking=True - ) - else: - full_shared_final_hidden_states[chunk_start:chunk_end, :].copy_( - final_hidden_states[0], non_blocking=True - ) - full_fused_final_hidden_states[chunk_start:chunk_end, :].copy_( - final_hidden_states[1], non_blocking=True - ) + fused_output = self.moe_forward( + hidden_states, + router_logits, + original_hidden_states, + self._encode_layer_name(), + ) + + return self._maybe_reduce_output(fused_output, og_hidden_dims) + + # TODO: avoid some of the copying by disabling inplace? + def _slice_and_copy_input( + self, + out_slice: torch.Tensor, + orig: torch.Tensor | None, + start: int, + end: int, + ) -> torch.Tensor: + assert orig is not None + slice_size = end - start + orig_slice = orig[start:end, :] + if self.enable_dbo: + assert out_slice.dim() == 3 + batch_buffer_idx = dbo_current_ubatch_id() + out_slice = out_slice[batch_buffer_idx, :] + + assert out_slice.size(0) >= slice_size + out_slice = out_slice[:slice_size, :] + out_slice.copy_(orig_slice, non_blocking=True) + return out_slice + + def forward_impl_chunked( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + final_shared_hidden_states, final_fused_hidden_states = ( + self._allocate_dp_chunking_outputs(hidden_states, router_logits) + ) ctx = get_forward_context() # flashinfer_cutlass_kernels can handle: optional DP + TP/EP @@ -567,7 +744,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False): max_tokens_across_dispatchers, self.moe_config.sp_size ) - num_tokens = full_hidden_states.size(0) + num_tokens = hidden_states.size(0) for chunk_idx, chunk_start_ in enumerate( range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) ): @@ -578,17 +755,52 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False): # clamp start and end chunk_start = min(chunk_start, num_tokens - 1) chunk_end = min(chunk_end, num_tokens) - with ctx.dp_metadata.chunked_sizes( + chunk_sizes = ctx.dp_metadata.chunked_sizes( self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx - ): - process_chunk( - chunk_start, chunk_end, skip_result_store=chunk_start_ >= num_tokens + ) + with chunk_sizes: + hidden_states_chunk = self._slice_and_copy_input( + self.batched_hidden_states, + hidden_states, + chunk_start, + chunk_end, ) + router_logits_chunk = self._slice_and_copy_input( + self.batched_router_logits, + router_logits, + chunk_start, + chunk_end, + ) + + shared_input_chunk = shared_input[chunk_start:chunk_end, :] if shared_input is not None else None + + shared_output_chunk, hidden_states_chunk = self._apply_quant_method( + layer=layer, + hidden_states=hidden_states_chunk, + extra_tensor=None, + router_logits=router_logits_chunk, + shared_input=shared_input_chunk, + ) + + # Store outputs + # TODO(bnell): document when chunk_start >= num_tokens + if chunk_start < num_tokens: + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + hidden_states_chunk, non_blocking=True + ) + if self.shared_experts is not None: + assert shared_output_chunk is not None + assert final_shared_hidden_states is not None + final_shared_hidden_states[chunk_start:chunk_end, :].copy_( + shared_output_chunk, non_blocking=True + ) + if self.shared_experts is None: - return full_fused_final_hidden_states + return final_fused_hidden_states else: - return (full_shared_final_hidden_states, full_fused_final_hidden_states) + assert final_shared_hidden_states is not None + return (final_shared_hidden_states, final_fused_hidden_states) def forward_impl( self, @@ -597,148 +809,43 @@ def forward_impl( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - assert self.quant_method is not None - - self.ensure_dp_chunking_init() - - has_separate_shared_experts = ( - not self.quant_method.mk_owns_shared_expert - and self.shared_experts is not None - ) - - use_chunked_impl = self.use_dp_chunking - - use_shared_experts_stream, shared_experts_input = ( - self._maybe_setup_shared_experts_stream( - hidden_states, - shared_input, - has_separate_shared_experts, - use_chunked_impl, + # TODO(bnell): split this into runtime vs. static parts? + self.use_shared_experts_stream = ( + current_platform.is_cuda() + and self.has_separate_shared_experts + and not self.use_dp_chunking + and self.shared_experts_stream is not None + and ( + hidden_states.shape[0] + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) ) - # If router/gate provided, then apply it here. - # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoE via - # separate cuda stream) - if self.gate is not None: - router_logits, _ = self.gate(hidden_states) - - if use_chunked_impl: - return self.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_input, - has_separate_shared_experts, - ) - - # NOTE(rob): once we finish migrating all the quant methods to use - # MKs, we can remove the naive dispatch/combine path from here. - do_naive_dispatch_combine = ( - self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk + # Check if we need to run shared experts before matrix multiply because + # matrix multiply may modify the hidden_states. + run_shared_experts_before = ( + self.has_separate_shared_experts and not self.use_shared_experts_stream ) - ctx = get_forward_context() - sp_ctx = ( - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) - if ctx.dp_metadata - else nullcontext() + # TODO(bnell): parts of the dispatch/combine steps will go away once + # #32567 lands and the remaining kernels are made MKs. The PCP + # code will probably remain + hidden_states, router_logits, extra_tensor = self._maybe_dispatch( + layer, + hidden_states, + router_logits, ) - with sp_ctx: - # Run shared experts before matrix multiply. - # because matrix multiply maybe modify the hidden_states. - if has_separate_shared_experts and not use_shared_experts_stream: - assert self.shared_experts is not None - shared_input = ( - shared_input if shared_input is not None else hidden_states - ) - shared_output = self.shared_experts(shared_input) - - # For naive dispatch/combine Dp/Ep, dispatch the hidden states and - # router logits to all experts. - # NOTE: this will be removed once all kernels are migrated into the - # MoEKernel framework. - if do_naive_dispatch_combine: - hidden_states, router_logits = get_ep_group().dispatch_router_logits( - hidden_states, - router_logits, - self.moe_config.is_sequence_parallel, - ) - - # NOTE: Similar with DP, PCP also needs dispatch and combine. For - # simplicity, AgRsAll2All was added separately for PCP here. Maybe - # we should modify All2AllManager abstract to better support PCP. - if self.moe_config.pcp_size > 1: - hidden_states = get_pcp_group().all_gather( - hidden_states, - dim=0, - ) - router_logits = get_pcp_group().all_gather( - router_logits, - dim=0, - ) - - # Matrix multiply. - if self.quant_method.is_monolithic: - final_hidden_states = self.quant_method.apply_monolithic( - layer=layer, - x=hidden_states, - router_logits=router_logits, - ) - else: - topk_weights, topk_ids = self.router.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - ) - - final_hidden_states = self.quant_method.apply( - layer=layer, - x=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_experts_input=shared_input, - ) - - if has_separate_shared_experts: - assert self.shared_experts is not None - - if use_shared_experts_stream: - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self.shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - shared_output = self.shared_experts(shared_experts_input) - current_stream().wait_stream(self.shared_experts_stream) - - final_hidden_states = ( - shared_output, - final_hidden_states, - ) - - def combine_output(states: torch.Tensor) -> torch.Tensor: - if do_naive_dispatch_combine: - states = get_ep_group().combine( - states, self.moe_config.is_sequence_parallel - ) - - if self.moe_config.pcp_size > 1: - states = get_pcp_group().reduce_scatter( - states, - dim=0, - ) - - return states + shared_output, hidden_states = self._apply_quant_method( + layer=layer, + hidden_states=hidden_states, + extra_tensor=extra_tensor, + router_logits=router_logits, + shared_input=shared_input, + run_shared_experts_before=run_shared_experts_before, + ) - if self.shared_experts is not None: - return ( - final_hidden_states[0], - combine_output(final_hidden_states[1]), - ) - else: - return combine_output(final_hidden_states) + return self._maybe_combine( + shared_output, + hidden_states, + ) From a4d3acb8469a4355aae6eb89a5146f46caf7c4ef Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 12 Feb 2026 18:03:08 -0500 Subject: [PATCH 002/108] fix lint Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/default_moe_runner.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index bb1ac465b1c5..c7de96e018e4 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -773,7 +773,11 @@ def forward_impl_chunked( chunk_end, ) - shared_input_chunk = shared_input[chunk_start:chunk_end, :] if shared_input is not None else None + shared_input_chunk = ( + shared_input[chunk_start:chunk_end, :] + if shared_input is not None + else None + ) shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, From 5b7f133cab46c2a8e7bad5424f8e54a4671c5c4a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 13:24:40 -0500 Subject: [PATCH 003/108] rebase Signed-off-by: Bill Nell --- .../fused_moe/runner/default_moe_runner.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index c7de96e018e4..5a204bd67ec8 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -261,7 +261,6 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) - # TODO(bnell): make this a member var? @property def use_dp_chunking(self) -> bool: return ( @@ -276,7 +275,7 @@ def _maybe_setup_shared_experts_stream( hidden_states: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | None: - hidden_states_clone: torch.Tensor | None = None + shared_experts_input: torch.Tensor | None = None if self.use_shared_experts_stream: assert self.shared_experts_stream is not None assert self.moe_config.disable_inplace @@ -455,7 +454,11 @@ def _maybe_pad_hidden_states( original_hidden_states: torch.Tensor | None, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, list[int]]: - original_hidden_dim = original_hidden_states.shape[-1] if original_hidden_states is not None else 0 + original_hidden_dim = ( + original_hidden_states.shape[-1] + if original_hidden_states is not None + else 0 + ) transformed_hidden_dim = hidden_states.shape[-1] if ( not self.quant_method.skip_forward_padding @@ -681,7 +684,10 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # For latent MoE: save ORIGINAL hidden_states before transform # (shared_experts need original dimension, routed experts use transformed) - original_hidden_states = hidden_states + if self.shared_experts is not None: + original_hidden_states = hidden_states + else: + original_hidden_states = None # Apply transform for routed experts (e.g., latent projection for latent MoE) hidden_states = self.apply_routed_input_transform(hidden_states) @@ -700,7 +706,6 @@ def forward( return self._maybe_reduce_output(fused_output, og_hidden_dims) - # TODO: avoid some of the copying by disabling inplace? def _slice_and_copy_input( self, out_slice: torch.Tensor, @@ -813,7 +818,6 @@ def forward_impl( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # TODO(bnell): split this into runtime vs. static parts? self.use_shared_experts_stream = ( current_platform.is_cuda() and self.has_separate_shared_experts From fad7f33991fc9209659cd3789a0ae22554efb12c Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 16:03:53 -0500 Subject: [PATCH 004/108] rebase + remove dead code Signed-off-by: Bill Nell --- .../fused_moe/runner/default_moe_runner.py | 53 ++++--------------- 1 file changed, 11 insertions(+), 42 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 5a204bd67ec8..15ccd47667a5 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -482,7 +482,6 @@ def _apply_quant_method( self, layer: torch.nn.Module, hidden_states: torch.Tensor, - extra_tensor: torch.Tensor | None, router_logits: torch.Tensor, shared_input: torch.Tensor | None, run_shared_experts_before: bool = True, @@ -502,14 +501,10 @@ def _apply_quant_method( shared_input, ) - # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014). - # Figure out nicer way to do this. - x_arg = hidden_states if extra_tensor is None else (hidden_states, extra_tensor) - if self.quant_method.is_monolithic: result = self.quant_method.apply_monolithic( layer=layer, - x=x_arg, + x=hidden_states, router_logits=router_logits, ) else: @@ -520,7 +515,7 @@ def _apply_quant_method( result = self.quant_method.apply( layer=layer, - x=x_arg, + x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, shared_experts_input=shared_input, @@ -603,45 +598,21 @@ def _maybe_dispatch( layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - extra_tensor: torch.Tensor | None = None - + ) -> tuple[torch.Tensor, torch.Tensor]: + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and + # router logits to all experts. + # NOTE: this will be removed once all kernels are migrated into the + # MoEKernel framework. if self.do_naive_dispatch_combine: - post_quant_allgather = ( - self.moe_config.dp_size > 1 - and self.moe_config.use_ep - and getattr(self.quant_method, "do_post_quant_allgather", False) - ) - - extra_tensors: list[torch.Tensor] | None = None - - if post_quant_allgather: - hidden_states_to_dispatch, extra_tensors = ( - self.quant_method.prepare_dp_allgather_tensor( - layer, hidden_states, router_logits - ) - ) - else: - hidden_states_to_dispatch = hidden_states - - result = get_ep_group().dispatch_router_logits( - hidden_states_to_dispatch, + hidden_states, router_logits = get_ep_group().dispatch_router_logits( + hidden_states, router_logits, self.moe_config.is_sequence_parallel, - extra_tensors=extra_tensors, ) - if len(result) == 3: - hidden_states, router_logits, extra_tensors = result - assert isinstance(extra_tensors, list) and len(extra_tensors) == 1 - extra_tensor = extra_tensors[0] - else: - hidden_states, router_logits = result - # NOTE: Similar with DP, PCP also needs dispatch and combine. For # simplicity, AgRsAll2All was added separately for PCP here. Maybe # we should modify All2AllManager abstraction to better support PCP. - # TODO(bnell): see what we can do here if self.moe_config.pcp_size > 1: hidden_states = get_pcp_group().all_gather( hidden_states, @@ -652,7 +623,7 @@ def _maybe_dispatch( dim=0, ) - return hidden_states, router_logits, extra_tensor + return hidden_states, router_logits def _maybe_combine( self, @@ -787,7 +758,6 @@ def forward_impl_chunked( shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, hidden_states=hidden_states_chunk, - extra_tensor=None, router_logits=router_logits_chunk, shared_input=shared_input_chunk, ) @@ -838,7 +808,7 @@ def forward_impl( # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain - hidden_states, router_logits, extra_tensor = self._maybe_dispatch( + hidden_states, router_logits = self._maybe_dispatch( layer, hidden_states, router_logits, @@ -847,7 +817,6 @@ def forward_impl( shared_output, hidden_states = self._apply_quant_method( layer=layer, hidden_states=hidden_states, - extra_tensor=extra_tensor, router_logits=router_logits, shared_input=shared_input, run_shared_experts_before=run_shared_experts_before, From 83c1863d91038cc85d7e5b72483cb63601244677 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 4 Feb 2026 10:36:16 -0500 Subject: [PATCH 005/108] wip Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 13 ++-- .../fused_moe/runner/default_moe_runner.py | 64 ++++++++++++++++++- .../layers/fused_moe/shared_fused_moe.py | 41 ++---------- 3 files changed, 74 insertions(+), 44 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 2f704569209c..bf6816fc611e 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -329,6 +329,7 @@ def __init__( is_sequence_parallel=False, expert_mapping: list[tuple[str, str, int, str]] | None = None, n_shared_experts: int | None = None, + shared_experts: torch.nn.Module | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -661,6 +662,7 @@ def _init_runner(self): quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, + enable_eplb=self.enable_eplb, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -708,12 +710,13 @@ def maybe_init_modular_kernel(self) -> None: def shared_experts(self) -> torch.nn.Module | None: return self._shared_experts if self.use_overlapped else None - @property - def layer_id(self): - # Delayed import to avoid circular dependency - from vllm.model_executor.models.utils import extract_layer_index + # TODO(bnell): is this needed? + #@property + #def layer_id(self): + # # Delayed import to avoid circular dependency + # from vllm.model_executor.models.utils import extract_layer_index - return extract_layer_index(self.layer_name) + # return extract_layer_index(self.layer_name) @property def gate(self) -> torch.nn.Module | None: diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 15ccd47667a5..db1c3f006539 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -11,6 +11,7 @@ from vllm.distributed import ( get_ep_group, get_pcp_group, + get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( @@ -139,6 +140,42 @@ def _moe_forward_shared( ) +def _moe_forward_shared_new( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: + # Reduce shared expert outputs if necessary, since the MLP + # should have been created with reduce_results=False. + if ( + self.reduce_results + and get_tensor_model_parallel_world_size() > 1 + and self.must_reduce_shared_expert_outputs() + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out + + if False and not self.overlap_shared_experts: + shared_out = self._shared_experts(hidden_states) + # XXXXXXXXXXXx + shared_experts = None + + router_logits = self._maybe_gate(hidden_states, router_logits) + with self._sequence_parallel_context(): + if self.use_dp_chunking: + shared_out, fused_out = self.forward_impl_chunked( + layer, hidden_states, router_logits + ) + else: + shared_out = fused_out = self.forward_impl( + layer, hidden_states, router_logits + ) + + return reduce_shared_out(shared_out), fused_out + + def _moe_forward_shared_fake( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -209,6 +246,7 @@ def __init__( quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, + enable_eplb: bool, ): super().__init__() self.moe_config = moe_config @@ -219,9 +257,12 @@ def __init__( self.quant_method = quant_method self.reduce_results = reduce_results self.enable_dbo = enable_dbo + self.enable_eplb = enable_eplb # Chunked all2all staging tensor - # TODO(bnell) rename these? + # TODO rename these + # These need to exist ahead of time due to CUDAgraph construction + # needing a fixed buffer address. self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() @@ -270,6 +311,22 @@ def use_dp_chunking(self) -> bool: or self.moe_config.moe_parallel_config.use_nixl_ep_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + # TODO(bnell): better name + @property + def overlap_shared_experts( + self, + shared_experts: torch.nn.Module | None, + ) -> bool: + # Disable shared expert overlap if: + # - we are using eplb with non-default backend, because of correctness issues + # - we are using flashinfer with DP, since there nothing to gain + # - we are using marlin kernels + backend = self.moe_config.moe_parallel_config.all2all_backend + return shared_experts is not None and not ( + (self.enable_eplb and backend != "allgather_reducescatter") + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) + def _maybe_setup_shared_experts_stream( self, hidden_states: torch.Tensor, @@ -481,6 +538,7 @@ def _maybe_pad_hidden_states( def _apply_quant_method( self, layer: torch.nn.Module, + shared_experts: torch.nn.Module | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, @@ -527,7 +585,9 @@ def _apply_quant_method( else: hidden_states = result - if not run_shared_experts_before and self.has_separate_shared_experts: + if not run_shared_experts_before and self.has_separate_shared_experts( + shared_experts + ): assert shared_output is None shared_output = self._apply_shared_experts( hidden_states_clone, diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index 37336df17561..251c05300331 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -3,10 +3,6 @@ import torch -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) from vllm.model_executor.layers.fused_moe.layer import FusedMoE @@ -23,36 +19,7 @@ def forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - if not self.use_overlapped: - if self._shared_experts is not None: - shared_out = self._shared_experts(hidden_states) - - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - self.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - else: - shared_out = None - - fused_out = super().forward( - hidden_states=hidden_states, - router_logits=router_logits, - ) - else: - shared_out, fused_out = super().forward( - hidden_states=hidden_states, - router_logits=router_logits, - ) - # ensure early TP reduction of shared expert outputs when required - if ( - shared_out is not None - and self.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out, fused_out + return super().forward( + hidden_states=hidden_states, + router_logits=router_logits, + ) From 7c7953e5159b4635c12f2761d244caf230c1f967 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 9 Feb 2026 18:27:46 -0500 Subject: [PATCH 006/108] fix Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index bf6816fc611e..c992930c8859 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -329,7 +329,6 @@ def __init__( is_sequence_parallel=False, expert_mapping: list[tuple[str, str, int, str]] | None = None, n_shared_experts: int | None = None, - shared_experts: torch.nn.Module | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -711,8 +710,8 @@ def shared_experts(self) -> torch.nn.Module | None: return self._shared_experts if self.use_overlapped else None # TODO(bnell): is this needed? - #@property - #def layer_id(self): + # @property + # def layer_id(self): # # Delayed import to avoid circular dependency # from vllm.model_executor.models.utils import extract_layer_index From d5b5805aab835a4d28343a3f1bc78931903722d4 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 11 Feb 2026 14:36:30 -0500 Subject: [PATCH 007/108] WIP DOUBLE CHECK THIS Signed-off-by: Bill Nell --- .../fused_moe/router/fused_moe_router.py | 8 +++ .../fused_moe/router/memoizing_router.py | 35 +++++++++++++ .../fused_moe/runner/default_moe_runner.py | 50 +++++++++---------- .../layers/fused_moe/shared_fused_moe.py | 2 +- 4 files changed, 69 insertions(+), 26 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py diff --git a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py index c322a8cd4cd6..d7aed4fdeb2b 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections.abc import Callable import torch @@ -13,6 +14,13 @@ class FusedMoERouter(ABC): method that is used for routing hidden states based on router logits. """ + @abstractmethod + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + raise NotImplementedError + @property @abstractmethod def routing_method_type(self) -> RoutingMethodType: diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py new file mode 100644 index 000000000000..a55bd2f09d6f --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter + + +class MemoizingRouter(FusedMoERouter): + def __init__(self, router: FusedMoERouter): + self.router = router + + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + self.router.set_capture_fn(capture_fn) + self.results: tuple[torch.Tensor, torch.Tensor] | None = None + + @property + def routing_method_type(self) -> RoutingMethodType: + return self.router.routing_method_type + + def select_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.results is None: + self.results = self.router.select_experts(hidden_states, router_logits) + return self.results diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index db1c3f006539..a6c688f887f2 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -146,30 +146,33 @@ def _moe_forward_shared_new( shared_experts_input: torch.Tensor | None, layer_name: str, ) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(layer_name) + runner = layer.runner + def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: # Reduce shared expert outputs if necessary, since the MLP # should have been created with reduce_results=False. if ( - self.reduce_results + runner.reduce_results and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() + and runner.must_reduce_shared_expert_outputs() ): shared_out = tensor_model_parallel_all_reduce(shared_out) return shared_out - if False and not self.overlap_shared_experts: - shared_out = self._shared_experts(hidden_states) - # XXXXXXXXXXXx + if not runner.overlap_shared_experts: + shared_out = runner._shared_experts(hidden_states) + # XXXXXXXXXXXXXXXX shared_experts = None - router_logits = self._maybe_gate(hidden_states, router_logits) - with self._sequence_parallel_context(): - if self.use_dp_chunking: - shared_out, fused_out = self.forward_impl_chunked( + router_logits = runner._maybe_gate(hidden_states, router_logits) + with runner._sequence_parallel_context(): + if runner.use_dp_chunking: + shared_out, fused_out = runner.forward_impl_chunked( layer, hidden_states, router_logits ) else: - shared_out = fused_out = self.forward_impl( + shared_out = fused_out = runner.forward_impl( layer, hidden_states, router_logits ) @@ -284,6 +287,13 @@ def __init__( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) + self.allow_shared_experts_stream = ( + current_platform.is_cuda() + and self.has_separate_shared_experts + and not self.use_dp_chunking + and self.shared_experts_stream is not None + ) + # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name @@ -315,14 +325,13 @@ def use_dp_chunking(self) -> bool: @property def overlap_shared_experts( self, - shared_experts: torch.nn.Module | None, ) -> bool: # Disable shared expert overlap if: # - we are using eplb with non-default backend, because of correctness issues # - we are using flashinfer with DP, since there nothing to gain # - we are using marlin kernels backend = self.moe_config.moe_parallel_config.all2all_backend - return shared_experts is not None and not ( + return self.shared_experts is not None and not ( (self.enable_eplb and backend != "allgather_reducescatter") or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels ) @@ -538,7 +547,6 @@ def _maybe_pad_hidden_states( def _apply_quant_method( self, layer: torch.nn.Module, - shared_experts: torch.nn.Module | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, @@ -585,9 +593,7 @@ def _apply_quant_method( else: hidden_states = result - if not run_shared_experts_before and self.has_separate_shared_experts( - shared_experts - ): + if not run_shared_experts_before and self.has_separate_shared_experts: assert shared_output is None shared_output = self._apply_shared_experts( hidden_states_clone, @@ -848,15 +854,9 @@ def forward_impl( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - self.use_shared_experts_stream = ( - current_platform.is_cuda() - and self.has_separate_shared_experts - and not self.use_dp_chunking - and self.shared_experts_stream is not None - and ( - hidden_states.shape[0] - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) + # TODO(bnell): split this into runtime vs. static parts? + self.use_shared_experts_stream = self.allow_shared_experts_stream and ( + hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) # Check if we need to run shared experts before matrix multiply because diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index 251c05300331..9aa9fe8e0a9a 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -6,7 +6,7 @@ from vllm.model_executor.layers.fused_moe.layer import FusedMoE -# TODO(bnell): Add shared + fused combo function? e.g. + +# TODO(bnell): Remove this entirely class SharedFusedMoE(FusedMoE): """ A FusedMoE operation that also computes the results of shared experts. From 42de827647ef6c7a41f052acccb7717a0f9cd83f Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Feb 2026 19:00:11 -0500 Subject: [PATCH 008/108] wip more refactoring Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 7 +- vllm/model_executor/layers/fused_moe/layer.py | 34 +-- .../fused_moe/runner/default_moe_runner.py | 281 +++++++++++++----- 3 files changed, 231 insertions(+), 91 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 78876ef7c9b0..5aa8a56d4115 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,6 +150,7 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), + self.base_layer.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXXXXX ) if quant_config.use_mxfp4_w4a16: @@ -592,9 +593,9 @@ def forward(self, *args, **kwargs): def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - @property - def _shared_experts(self): - return self.base_layer._shared_experts + # @property + # def _shared_experts(self): + # return self.base_layer._shared_experts @property def quant_method(self): diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c992930c8859..c6c6dab9dd04 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -636,14 +636,14 @@ def _get_quant_method() -> FusedMoEMethodBase: # - we are using eplb with non-default backend, because of correctness issues # - we are using flashinfer with DP, since there nothing to gain # - we are using marlin kernels - backend = self.moe_parallel_config.all2all_backend - self.use_overlapped = ( - not ( - (self.enable_eplb and backend != "allgather_reducescatter") - or self.moe_parallel_config.use_fi_nvl_two_sided_kernels - ) - and self._shared_experts is not None - ) + # backend = self.moe_parallel_config.all2all_backend + # self.use_overlapped = ( + # not ( + # (self.enable_eplb and backend != "allgather_reducescatter") + # or self.moe_parallel_config.use_fi_all2allv_kernels + # ) + # and self._shared_experts is not None + # ) self.runner = self._init_runner() @@ -656,8 +656,8 @@ def _init_runner(self): moe_config=self.moe_config, router=self.router, routed_input_transform=self._routed_input_transform, - gate=self.gate, - shared_experts=self.shared_experts, + gate=self._gate, + shared_experts=self._shared_experts, quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, @@ -700,14 +700,14 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.shared_experts, + self.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXX inplace=not self.moe_config.disable_inplace, ) ) - @property - def shared_experts(self) -> torch.nn.Module | None: - return self._shared_experts if self.use_overlapped else None + # @property + # def shared_experts(self) -> torch.nn.Module | None: + # return self._shared_experts if self.use_overlapped else None # TODO(bnell): is this needed? # @property @@ -717,9 +717,9 @@ def shared_experts(self) -> torch.nn.Module | None: # return extract_layer_index(self.layer_name) - @property - def gate(self) -> torch.nn.Module | None: - return self._gate if self.use_overlapped else None + # @property + # def gate(self) -> torch.nn.Module | None: + # return self._gate if self.use_overlapped else None @property def tp_size(self): diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index a6c688f887f2..8dbf392f430c 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -3,6 +3,7 @@ from collections.abc import Callable from contextlib import nullcontext from typing import TYPE_CHECKING +from enum import IntEnum import torch import torch.nn.functional as F @@ -44,6 +45,23 @@ logger = init_logger(__name__) +class SharedExpertsOrder(IntEnum): + # No shared experts. + NONE = (0,) + + # Get rid of this one? combine with BEFORE? + EXTERNAL = (1,) + + # Called by modular kernel. + INTERNAL = (2,) + + # Called right before quant_method is executed. + BEFORE_QUANT_METHOD = (3,) + + # Called right after quant_method is executed (possibly with streaming). + AFTER_QUANT_METHOD = (4,) + + def get_layer_from_name(layer_name: str) -> torch.nn.Module: forward_context: ForwardContext = get_forward_context() if layer_name == "from_forward_context": @@ -81,6 +99,7 @@ def _moe_forward( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: +<<<<<<< HEAD layer = get_layer_from_name(_resolve_layer_name(layer_name)) # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -101,6 +120,15 @@ def _moe_forward( router_logits, shared_experts_input, ) +======= + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) +>>>>>>> d11342756e (wip more refactoring) def _moe_forward_fake( @@ -118,6 +146,7 @@ def _moe_forward_shared( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: +<<<<<<< HEAD layer = get_layer_from_name(_resolve_layer_name(layer_name)) # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -177,6 +206,15 @@ def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: ) return reduce_shared_out(shared_out), fused_out +======= + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) +>>>>>>> d11342756e (wip more refactoring) def _moe_forward_shared_fake( @@ -266,6 +304,7 @@ def __init__( # TODO rename these # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. + # TODO: these could be global, i.e. shared by all layers self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() @@ -287,17 +326,69 @@ def __init__( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) - self.allow_shared_experts_stream = ( + self.use_dp_chunking = ( + self.moe_config.moe_parallel_config.use_pplx_kernels + or self.moe_config.moe_parallel_config.use_deepep_ll_kernels + or self.moe_config.moe_parallel_config.use_mori_kernels + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + + # Needed for string -> FusedMoE layer lookup in custom ops. + self.layer_name = layer.layer_name + + self.moe_forward = self._select_forward(layer) + + def _has_external_experts(self) -> bool: + # Disable shared expert overlap if: + # - we are using eplb with non-default backend, because of correctness issues + # - we are using flashinfer with DP, since there nothing to gain + backend = self.moe_config.moe_parallel_config.all2all_backend + return self.shared_experts is not None and not ( + (self.enable_eplb and backend != "allgather_reducescatter") + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) + + def _determine_shared_experts_order( + self, + hidden_states: torch.Tensor, + ) -> tuple[SharedExpertsOrder, bool]: + if self.shared_experts is None: + return SharedExpertsOrder.NONE, False + + if self._has_external_experts(): + return SharedExpertsOrder.EXTERNAL, False + + has_separate_shared_experts = not self.quant_method.mk_owns_shared_expert + + if not has_separate_shared_experts: + return SharedExpertsOrder.INTERNAL, False + + allow_shared_experts_stream = ( current_platform.is_cuda() - and self.has_separate_shared_experts + and has_separate_shared_experts and not self.use_dp_chunking and self.shared_experts_stream is not None + and hidden_states.shape[0] + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) - # Needed for string -> FusedMoE layer lookup in custom ops. - self.layer_name = layer.layer_name + # Check if we need to run shared experts before matrix multiply because + # matrix multiply may modify the hidden_states. + run_shared_experts_before = ( + has_separate_shared_experts and not allow_shared_experts_stream + ) - self.moe_forward = self._select_forward(layer) + if run_shared_experts_before: + return SharedExpertsOrder.BEFORE_QUANT_METHOD, False + else: + return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream + + # XXXXXX used by layer.py and lora/layers/fused_moe.py + def get_shared_experts(self) -> torch.nn.Module | None: + if self._has_external_experts(): + return self.shared_experts + else: + return None def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -312,6 +403,7 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) +<<<<<<< HEAD @property def use_dp_chunking(self) -> bool: return ( @@ -396,35 +488,39 @@ def _maybe_init_dp_chunking(self): device=device, ) - @property - def has_separate_shared_experts(self) -> bool: - return ( - not self.quant_method.mk_owns_shared_expert - and self.shared_experts is not None - ) - - def _apply_shared_experts( + def _maybe_apply_shared_experts( self, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, - allow_streaming: bool = False, + shared_input: torch.Tensor | None, + order: SharedExpertsOrder, ) -> torch.Tensor | None: - shared_output: torch.Tensor | None = None - if self.has_separate_shared_experts: - assert self.shared_experts is not None - - if self.use_shared_experts_stream and allow_streaming: - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self.shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - shared_output = self.shared_experts(hidden_states) - current_stream().wait_stream(self.shared_experts_stream) - else: + experts_order, use_shared_experts_stream = self._determine_shared_experts_order( + hidden_states + ) + if order != experts_order: + return None + + assert self.shared_experts is not None + assert shared_output is None + if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: + hidden_states = self._maybe_setup_shared_experts_stream( + hidden_states, + shared_input, + ) + + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self.shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream shared_output = self.shared_experts(hidden_states) + current_stream().wait_stream(self.shared_experts_stream) + else: + shared_output = self.shared_experts(hidden_states) return shared_output @@ -475,6 +571,18 @@ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Ten return result return hidden_states + # TODO: combine with runner.forward_impl_X + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: + # Reduce shared expert outputs if necessary, since the MLP + # should have been created with reduce_results=False. + if ( + self.reduce_results + and get_tensor_model_parallel_world_size() > 1 + and self.must_reduce_shared_expert_outputs() + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out + def _maybe_reduce_output( self, states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -523,7 +631,7 @@ def _maybe_pad_hidden_states( original_hidden_dim = ( original_hidden_states.shape[-1] if original_hidden_states is not None - else 0 + else hidden_states.shape[-1] ) transformed_hidden_dim = hidden_states.shape[-1] if ( @@ -547,28 +655,22 @@ def _maybe_pad_hidden_states( def _apply_quant_method( self, layer: torch.nn.Module, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, - run_shared_experts_before: bool = True, ) -> tuple[torch.Tensor | None, torch.Tensor]: - shared_output: torch.Tensor | None = None - # Run this before quant_method to avoid inplace issues. - if run_shared_experts_before: - shared_input = shared_input if shared_input is not None else hidden_states - shared_output = self._apply_shared_experts( - shared_input, - False, - ) - else: - hidden_states_clone = self._maybe_setup_shared_experts_stream( - hidden_states, - shared_input, - ) + shared_input = shared_input if shared_input is not None else hidden_states + shared_output = self._maybe_apply_shared_experts( + shared_output, + shared_input, + shared_input, + SharedExpertsOrder.BEFORE_QUANT_METHOD, + ) if self.quant_method.is_monolithic: - result = self.quant_method.apply_monolithic( + fused_out = self.quant_method.apply_monolithic( layer=layer, x=hidden_states, router_logits=router_logits, @@ -579,7 +681,7 @@ def _apply_quant_method( router_logits=router_logits, ) - result = self.quant_method.apply( + fused_out = self.quant_method.apply( layer=layer, x=hidden_states, topk_weights=topk_weights, @@ -587,18 +689,23 @@ def _apply_quant_method( shared_experts_input=shared_input, ) - if isinstance(result, tuple): - assert shared_output is None - shared_output, hidden_states = result + # TODO: try to clean this up + if isinstance(fused_out, tuple): + if fused_out[0] is None: + assert shared_output is not None + hidden_states = fused_out[1] + else: + assert shared_output is None + shared_output, hidden_states = fused_out else: - hidden_states = result + hidden_states = fused_out - if not run_shared_experts_before and self.has_separate_shared_experts: - assert shared_output is None - shared_output = self._apply_shared_experts( - hidden_states_clone, - True, - ) + shared_output = self._maybe_apply_shared_experts( + shared_output, + shared_input, + shared_input, + SharedExpertsOrder.AFTER_QUANT_METHOD, + ) return shared_output, hidden_states @@ -721,10 +828,10 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # For latent MoE: save ORIGINAL hidden_states before transform # (shared_experts need original dimension, routed experts use transformed) - if self.shared_experts is not None: - original_hidden_states = hidden_states - else: - original_hidden_states = None + original_hidden_states = ( + # or shared_experts is not None? + hidden_states if self.routed_input_transform is not None else None + ) # Apply transform for routed experts (e.g., latent projection for latent MoE) hidden_states = self.apply_routed_input_transform(hidden_states) @@ -734,6 +841,7 @@ def forward( hidden_states, ) + # Note: moe_forward will call forward_dispatch fused_output = self.moe_forward( hidden_states, router_logits, @@ -741,8 +849,42 @@ def forward( self._encode_layer_name(), ) + # TODO: figure out how to combine this with maybe_reduce_output + if self.shared_experts is not None: + fused_output = ( + self._maybe_reduce_shared_out(fused_output[0]), + fused_output[1], + ) + return self._maybe_reduce_output(fused_output, og_hidden_dims) + def forward_dispatch( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # TODO(bnell): this can be removed after MK migration is complete. + layer.ensure_moe_quant_config_init() + + router_logits = self._maybe_gate(hidden_states, router_logits) + with self._sequence_parallel_context(): + if self.use_dp_chunking: + return self.forward_impl_chunked( + layer, + hidden_states, + router_logits, + shared_input, + ) + else: + return self.forward_impl( + layer, + hidden_states, + router_logits, + shared_input, + ) + def _slice_and_copy_input( self, out_slice: torch.Tensor, @@ -823,6 +965,7 @@ def forward_impl_chunked( shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, + shared_output=None, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, shared_input=shared_input_chunk, @@ -854,15 +997,11 @@ def forward_impl( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # TODO(bnell): split this into runtime vs. static parts? - self.use_shared_experts_stream = self.allow_shared_experts_stream and ( - hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) - - # Check if we need to run shared experts before matrix multiply because - # matrix multiply may modify the hidden_states. - run_shared_experts_before = ( - self.has_separate_shared_experts and not self.use_shared_experts_stream + shared_output = self._maybe_apply_shared_experts( + None, + hidden_states, + shared_input, + SharedExpertsOrder.EXTERNAL, ) # TODO(bnell): parts of the dispatch/combine steps will go away once @@ -876,10 +1015,10 @@ def forward_impl( shared_output, hidden_states = self._apply_quant_method( layer=layer, + shared_output=shared_output, hidden_states=hidden_states, router_logits=router_logits, shared_input=shared_input, - run_shared_experts_before=run_shared_experts_before, ) return self._maybe_combine( From 2e4ce005b842b58f4d6e0149ca3015f5fd48bb76 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Feb 2026 19:20:36 -0500 Subject: [PATCH 009/108] wip Signed-off-by: Bill Nell --- .../fused_moe/runner/default_moe_runner.py | 157 ++++-------------- 1 file changed, 29 insertions(+), 128 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 8dbf392f430c..51f1afda2400 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -99,28 +99,6 @@ def _moe_forward( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: -<<<<<<< HEAD - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - runner = layer.runner - router_logits = runner._maybe_gate(hidden_states, router_logits) - with runner._sequence_parallel_context(): - if runner.use_dp_chunking: - return runner.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - else: - return runner.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) -======= layer = get_layer_from_name(layer_name) return layer.runner.forward_dispatch( layer, @@ -128,7 +106,6 @@ def _moe_forward( router_logits, shared_experts_input, ) ->>>>>>> d11342756e (wip more refactoring) def _moe_forward_fake( @@ -145,68 +122,7 @@ def _moe_forward_shared( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: -<<<<<<< HEAD - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - runner = layer.runner - router_logits = runner._maybe_gate(hidden_states, router_logits) - with runner._sequence_parallel_context(): - if runner.use_dp_chunking: - return runner.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - else: - return runner.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_shared_new( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: str, ) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(layer_name) - runner = layer.runner - - def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - runner.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and runner.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out - - if not runner.overlap_shared_experts: - shared_out = runner._shared_experts(hidden_states) - # XXXXXXXXXXXXXXXX - shared_experts = None - - router_logits = runner._maybe_gate(hidden_states, router_logits) - with runner._sequence_parallel_context(): - if runner.use_dp_chunking: - shared_out, fused_out = runner.forward_impl_chunked( - layer, hidden_states, router_logits - ) - else: - shared_out = fused_out = runner.forward_impl( - layer, hidden_states, router_logits - ) - - return reduce_shared_out(shared_out), fused_out -======= layer = get_layer_from_name(layer_name) return layer.runner.forward_dispatch( layer, @@ -214,7 +130,6 @@ def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: router_logits, shared_experts_input, ) ->>>>>>> d11342756e (wip more refactoring) def _moe_forward_shared_fake( @@ -330,7 +245,8 @@ def __init__( self.moe_config.moe_parallel_config.use_pplx_kernels or self.moe_config.moe_parallel_config.use_deepep_ll_kernels or self.moe_config.moe_parallel_config.use_mori_kernels - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels + or self.moe_config.moe_parallel_config.use_nixl_ep_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK # Needed for string -> FusedMoE layer lookup in custom ops. @@ -403,32 +319,7 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) -<<<<<<< HEAD - @property - def use_dp_chunking(self) -> bool: - return ( - self.moe_config.moe_parallel_config.use_deepep_ll_kernels - or self.moe_config.moe_parallel_config.use_mori_kernels - or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels - or self.moe_config.moe_parallel_config.use_nixl_ep_kernels - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK - - # TODO(bnell): better name - @property - def overlap_shared_experts( - self, - ) -> bool: - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - # - we are using marlin kernels - backend = self.moe_config.moe_parallel_config.all2all_backend - return self.shared_experts is not None and not ( - (self.enable_eplb and backend != "allgather_reducescatter") - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels - ) - - def _maybe_setup_shared_experts_stream( + def _setup_shared_experts_stream( self, hidden_states: torch.Tensor, shared_input: torch.Tensor | None, @@ -504,7 +395,8 @@ def _maybe_apply_shared_experts( assert self.shared_experts is not None assert shared_output is None if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - hidden_states = self._maybe_setup_shared_experts_stream( + # TODO: fold this in? + hidden_states = self._setup_shared_experts_stream( hidden_states, shared_input, ) @@ -522,6 +414,12 @@ def _maybe_apply_shared_experts( else: shared_output = self.shared_experts(hidden_states) + if order == SharedExpertsOrder.EXTERNAL: + # TODO: figure out how to combine this with maybe_reduce_output? + # or get rid of it completely..... + assert shared_output is not None + shared_output = self._maybe_reduce_shared_out(shared_output) + return shared_output def must_reduce_shared_expert_outputs(self) -> bool: @@ -821,6 +719,11 @@ def _maybe_combine( else: return hidden_states + # + # forward + # - self.moe_forward (_moe_forward or _moe_forward_shared) + # - forward_chunking_dispatch + # - forward_impl or forward_impl_chunked def forward( self, hidden_states: torch.Tensor, @@ -849,13 +752,6 @@ def forward( self._encode_layer_name(), ) - # TODO: figure out how to combine this with maybe_reduce_output - if self.shared_experts is not None: - fused_output = ( - self._maybe_reduce_shared_out(fused_output[0]), - fused_output[1], - ) - return self._maybe_reduce_output(fused_output, og_hidden_dims) def forward_dispatch( @@ -869,10 +765,19 @@ def forward_dispatch( layer.ensure_moe_quant_config_init() router_logits = self._maybe_gate(hidden_states, router_logits) + + shared_output = self._maybe_apply_shared_experts( + None, + hidden_states, + shared_input, + SharedExpertsOrder.EXTERNAL, + ) + with self._sequence_parallel_context(): if self.use_dp_chunking: return self.forward_impl_chunked( layer, + shared_output, hidden_states, router_logits, shared_input, @@ -880,6 +785,7 @@ def forward_dispatch( else: return self.forward_impl( layer, + shared_output, hidden_states, router_logits, shared_input, @@ -908,6 +814,7 @@ def _slice_and_copy_input( def forward_impl_chunked( self, layer: torch.nn.Module, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, @@ -965,7 +872,7 @@ def forward_impl_chunked( shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, - shared_output=None, + shared_output=shared_output, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, shared_input=shared_input_chunk, @@ -993,17 +900,11 @@ def forward_impl_chunked( def forward_impl( self, layer: torch.nn.Module, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - shared_output = self._maybe_apply_shared_experts( - None, - hidden_states, - shared_input, - SharedExpertsOrder.EXTERNAL, - ) - # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain From 0144b8bf9eda7618742b73580e8d9a53bd244efa Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Feb 2026 16:49:41 -0500 Subject: [PATCH 010/108] SharedExperts wip Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 6 +- .../layers/fused_moe/cutlass_moe.py | 2 + .../layers/fused_moe/fused_moe_method_base.py | 4 +- .../fused_moe/fused_moe_modular_method.py | 10 +- vllm/model_executor/layers/fused_moe/layer.py | 65 ++-- .../layers/fused_moe/modular_kernel.py | 58 ++-- .../layers/fused_moe/oracle/fp8.py | 9 +- .../layers/fused_moe/oracle/nvfp4.py | 9 +- .../layers/fused_moe/oracle/unquantized.py | 4 + .../fused_moe/runner/default_moe_runner.py | 301 ++++-------------- .../layers/fused_moe/runner/shared_experts.py | 207 ++++++++++++ .../fused_moe/unquantized_fused_moe_method.py | 10 +- .../layers/quantization/awq_marlin.py | 2 +- .../layers/quantization/bitsandbytes.py | 2 +- .../compressed_tensors_moe.py | 20 +- .../layers/quantization/experts_int8.py | 2 +- .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/gguf.py | 2 +- .../layers/quantization/gptq_marlin.py | 2 +- .../layers/quantization/modelopt.py | 8 +- .../layers/quantization/moe_wna16.py | 2 +- .../layers/quantization/mxfp4.py | 4 +- .../layers/quantization/quark/quark_moe.py | 6 +- 23 files changed, 391 insertions(+), 348 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/runner/shared_experts.py diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 5aa8a56d4115..b6c62b618a09 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,9 +150,11 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), - self.base_layer.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXXXXX + self.base_layer.shared_experts, ) + # TODO: could be incorrect due to monolithic kernel? or add assert it + # is modular? if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.impl.fused_experts, @@ -170,6 +172,7 @@ def wrapper(*args, **kwargs): moe_state_dict["apply_router_weight_on_input"] = kwargs[ "apply_router_weight_on_input" ] + # TODO: global_num_experts/shared_experts_input? result = func(*args, **kwargs) return result @@ -339,6 +342,7 @@ def wrapper(*args, **kwargs): fused_experts = m_fused_moe_fn.impl.fused_experts + # TODO: seems like this could be done with modular kernel subclasses? m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 75ee776646ba..43082b3675a6 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1194,6 +1194,8 @@ def cutlass_moe_w4a8_fp8( quant_config=quant_config, group_size=group_size, ), + shared_experts=None, + inplace=False, ) return fn.apply( diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py index f6a303e7988e..9350a9c7de74 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py @@ -130,7 +130,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: raise NotImplementedError def apply_monolithic( @@ -138,5 +138,5 @@ def apply_monolithic( layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 0065c11f3163..142e180786c6 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -16,6 +16,9 @@ FusedMoEKernel, FusedMoEPrepareAndFinalizeModular, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) logger = init_logger(__name__) @@ -44,7 +47,7 @@ def make( moe_layer: torch.nn.Module, old_quant_method: FusedMoEMethodBase, prepare_finalize: FusedMoEPrepareAndFinalizeModular, - shared_experts: torch.nn.Module | None, + shared_experts: SharedExperts | None, inplace: bool = False, ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( @@ -52,8 +55,7 @@ def make( FusedMoEKernel( prepare_finalize, old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), - shared_experts, - moe_parallel_config=moe_layer.moe_parallel_config, + shared_experts=shared_experts, inplace=inplace, ), ) @@ -89,7 +91,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c6c6dab9dd04..bbf6e2ffedb9 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -42,6 +42,9 @@ from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( DefaultMoERunner, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) @@ -337,7 +340,6 @@ def __init__( super().__init__() self._gate = gate - self._shared_experts = shared_experts self._routed_input_transform = routed_input_transform if params_dtype is None: @@ -565,7 +567,7 @@ def __init__( device=vllm_config.device_config.device, routing_method=self.routing_method_type, # TODO: in_dtype == out_dtype? - disable_inplace=disable_inplace() or self._shared_experts is not None, + disable_inplace=disable_inplace() or self.shared_experts is not None, ) if self.moe_config.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( @@ -632,22 +634,30 @@ def _get_quant_method() -> FusedMoEMethodBase: self.quant_method.create_weights(layer=self, **moe_quant_params) self.base_quant_method = self.quant_method - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - # - we are using marlin kernels - # backend = self.moe_parallel_config.all2all_backend - # self.use_overlapped = ( - # not ( - # (self.enable_eplb and backend != "allgather_reducescatter") - # or self.moe_parallel_config.use_fi_all2allv_kernels - # ) - # and self._shared_experts is not None - # ) - + self._shared_experts = shared_experts + self.shared_experts = self._init_shared_experts() self.runner = self._init_runner() - def _init_runner(self): + def _init_shared_experts(self) -> SharedExperts | None: + if self._shared_experts is None: + return None + + reduce_shared_output = ( + self.reduce_results + # XXXX ordering issue + and self.quant_method.moe_mk is not None + and self.quant_method.moe_mk.output_is_reduced() + ) + + return SharedExperts( + self._shared_experts, + moe_config=self.moe_config, + has_separate_shared_experts=not self.quant_method.mk_owns_shared_expert, + use_dp_chunking=self.use_dp_chunking, # XXXXXXXXXXXXXX + must_reduce_shared_expert_outputs=reduce_shared_output, + ) + + def _init_runner(self) -> DefaultMoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. @@ -657,7 +667,7 @@ def _init_runner(self): router=self.router, routed_input_transform=self._routed_input_transform, gate=self._gate, - shared_experts=self._shared_experts, + shared_experts=self.shared_experts, quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, @@ -700,26 +710,17 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXX + self.runner._get_shared_experts(), # XXXXXXXXXXXXXXXXX inplace=not self.moe_config.disable_inplace, ) ) - # @property - # def shared_experts(self) -> torch.nn.Module | None: - # return self._shared_experts if self.use_overlapped else None - - # TODO(bnell): is this needed? - # @property - # def layer_id(self): - # # Delayed import to avoid circular dependency - # from vllm.model_executor.models.utils import extract_layer_index - - # return extract_layer_index(self.layer_name) + @property + def layer_id(self): + # Delayed import to avoid circular dependency + from vllm.model_executor.models.utils import extract_layer_index - # @property - # def gate(self) -> torch.nn.Module | None: - # return self._gate if self.use_overlapped else None + return extract_layer_index(self.layer_name) @property def tp_size(self): diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index a6b498834017..bbafc682c102 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -20,6 +20,10 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, disable_inplace, @@ -234,6 +238,13 @@ def output_is_reduced(self) -> bool: """ raise NotImplementedError + def supports_async(self) -> bool: + """ + Indicates whether or not this class implements prepare_async and + finalize_async. + """ + return False + # TODO: pass FusedMoEParallelConfig in as ctor parameter? class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): @@ -280,13 +291,6 @@ def prepare( """ raise NotImplementedError - def supports_async(self) -> bool: - """ - Indicates whether or not this class implements prepare_async and - finalize_async. - """ - return False - def prepare_async( self, a1: torch.Tensor, @@ -991,8 +995,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, + shared_experts: SharedExperts | None, inplace: bool = False, ): self.prepare_finalize = prepare_finalize @@ -1000,6 +1003,9 @@ def __init__( self.shared_experts = shared_experts self.moe_parallel_config = moe_parallel_config self.inplace = inplace + + moe_parallel_config = fused_experts.moe_config.moe_parallel_config + self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config self.is_dp_ep = ( moe_parallel_config is not None and moe_parallel_config.dp_size > 1 @@ -1069,6 +1075,17 @@ def _allocate_buffers( return workspace13, workspace2, fused_out + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + ): + if self.shared_experts is not None: + assert shared_experts_input is not None + self.shared_experts.apply( + shared_experts_input, + SharedExpertsOrder.INTERNAL, + ) + def _prepare( self, hidden_states: torch.Tensor, @@ -1241,15 +1258,6 @@ def _finalize( shared_experts_input is the original hidden_states (full dimension) needed by the shared expert MLP. """ - shared_output: torch.Tensor | None = None - - # For latent MoE: shared experts need the original hidden_states - # (full hidden_size), not the latent-projected version used by - # routed experts. - se_hidden_states = ( - shared_experts_input if shared_experts_input is not None else hidden_states - ) - if not self.prepare_finalize.supports_async(): assert not dbo_enabled() @@ -1261,8 +1269,7 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - if self.shared_experts is not None: - shared_output = self.shared_experts(se_hidden_states) + self._maybe_apply_shared_experts(shared_experts_input) else: finalize_ret = self.prepare_finalize.finalize_async( output, @@ -1272,8 +1279,7 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - if self.shared_experts is not None: - shared_output = self.shared_experts(se_hidden_states) + self._maybe_apply_shared_experts(shared_experts_input) # TODO(lucas): refactor this in the alternative schedules followup # currently unpack if we have hook + receiver pair or just @@ -1296,11 +1302,7 @@ def _finalize( receiver() - if self.shared_experts is None: - return output - else: - assert shared_output is not None - return shared_output, output + return output def apply( self, @@ -1314,7 +1316,7 @@ def apply( expert_map: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, shared_experts_input: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """ This function computes a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2, and top-k gating mechanism. diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index a63c02663886..57431b946e24 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -18,6 +18,9 @@ fp8_w8a8_moe_quant_config, fp8_w8a16_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, get_flashinfer_moe_backend, @@ -545,7 +548,7 @@ def make_fp8_moe_kernel( experts_cls: type[mk.FusedMoEExperts], fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - shared_experts: torch.nn.Module | None = None, + shared_experts: SharedExperts | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -583,10 +586,8 @@ def make_fp8_moe_kernel( experts, shared_experts=( shared_experts - if moe_config.moe_parallel_config.use_deepep_ll_kernels - else None + if prepare_finalize.supports_async() else None ), - moe_parallel_config=moe_config.moe_parallel_config, inplace=( not moe_config.disable_inplace and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 8a224cb39e7c..194513981e41 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -18,6 +18,9 @@ nvfp4_moe_quant_config, nvfp4_w4a16_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( prepare_nvfp4_moe_layer_for_fi_or_cutlass, ) @@ -397,7 +400,7 @@ def make_nvfp4_moe_kernel( moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - shared_experts: torch.nn.Module | None = None, + shared_experts: SharedExperts | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -435,10 +438,8 @@ def make_nvfp4_moe_kernel( experts, shared_experts=( shared_experts - if moe_config.moe_parallel_config.use_deepep_ll_kernels - else None + if prepare_finalize.supports_async() else None ), - moe_parallel_config=moe_config.moe_parallel_config, inplace=False, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 9c31da10dd94..926819592c3a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -224,6 +224,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=False, ) @@ -238,6 +239,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.TRITON: @@ -249,6 +251,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.XPU: @@ -260,6 +263,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=not moe_config.disable_inplace, ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 51f1afda2400..2e8ff2a9fd54 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -3,7 +3,6 @@ from collections.abc import Callable from contextlib import nullcontext from typing import TYPE_CHECKING -from enum import IntEnum import torch import torch.nn.functional as F @@ -12,7 +11,6 @@ from vllm.distributed import ( get_ep_group, get_pcp_group, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( @@ -31,6 +29,10 @@ FusedMoERouter, ) from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( @@ -45,23 +47,6 @@ logger = init_logger(__name__) -class SharedExpertsOrder(IntEnum): - # No shared experts. - NONE = (0,) - - # Get rid of this one? combine with BEFORE? - EXTERNAL = (1,) - - # Called by modular kernel. - INTERNAL = (2,) - - # Called right before quant_method is executed. - BEFORE_QUANT_METHOD = (3,) - - # Called right after quant_method is executed (possibly with streaming). - AFTER_QUANT_METHOD = (4,) - - def get_layer_from_name(layer_name: str) -> torch.nn.Module: forward_context: ForwardContext = get_forward_context() if layer_name == "from_forward_context": @@ -154,7 +139,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], + mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -163,7 +148,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward_shared", op_func=_moe_forward_shared, - mutates_args=["hidden_states"], + mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_shared_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -198,7 +183,7 @@ def __init__( router: FusedMoERouter, routed_input_transform: torch.nn.Module | None, gate: torch.nn.Module | None, - shared_experts: torch.nn.Module | None, + shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, @@ -224,23 +209,6 @@ def __init__( self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() - # Allow disabling of the separate shared experts stream for - # debug purposes. - # TODO: Remove this after more extensive testings with TP/DP - # and other execution modes - self.use_shared_experts_stream = False - if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: - logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") - self.shared_experts_stream = None - else: - # TODO(rob): enable shared expert overlap with non-cuda-alike. - # aux_stream() returns None on non-cuda-alike platforms. - self.shared_experts_stream = aux_stream() - if self.shared_experts_stream is not None: - logger.debug_once( - "Enabled separate cuda stream for MoE shared_experts", scope="local" - ) - self.use_dp_chunking = ( self.moe_config.moe_parallel_config.use_pplx_kernels or self.moe_config.moe_parallel_config.use_deepep_ll_kernels @@ -254,57 +222,9 @@ def __init__( self.moe_forward = self._select_forward(layer) - def _has_external_experts(self) -> bool: - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - backend = self.moe_config.moe_parallel_config.all2all_backend - return self.shared_experts is not None and not ( - (self.enable_eplb and backend != "allgather_reducescatter") - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels - ) - - def _determine_shared_experts_order( - self, - hidden_states: torch.Tensor, - ) -> tuple[SharedExpertsOrder, bool]: - if self.shared_experts is None: - return SharedExpertsOrder.NONE, False - - if self._has_external_experts(): - return SharedExpertsOrder.EXTERNAL, False - - has_separate_shared_experts = not self.quant_method.mk_owns_shared_expert - - if not has_separate_shared_experts: - return SharedExpertsOrder.INTERNAL, False - - allow_shared_experts_stream = ( - current_platform.is_cuda() - and has_separate_shared_experts - and not self.use_dp_chunking - and self.shared_experts_stream is not None - and hidden_states.shape[0] - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) - - # Check if we need to run shared experts before matrix multiply because - # matrix multiply may modify the hidden_states. - run_shared_experts_before = ( - has_separate_shared_experts and not allow_shared_experts_stream - ) - - if run_shared_experts_before: - return SharedExpertsOrder.BEFORE_QUANT_METHOD, False - else: - return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream - # XXXXXX used by layer.py and lora/layers/fused_moe.py - def get_shared_experts(self) -> torch.nn.Module | None: - if self._has_external_experts(): - return self.shared_experts - else: - return None + def _get_shared_experts(self) -> SharedExperts | None: + return self.shared_experts def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -319,36 +239,6 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) - def _setup_shared_experts_stream( - self, - hidden_states: torch.Tensor, - shared_input: torch.Tensor | None, - ) -> torch.Tensor | None: - shared_experts_input: torch.Tensor | None = None - if self.use_shared_experts_stream: - assert self.shared_experts_stream is not None - assert self.moe_config.disable_inplace - - shared_experts_input = ( - shared_input if shared_input is not None else hidden_states - ) - - # Record that the shared_experts_input will be used in the - # shared_experts_stream to avoid gc issue from - # deallocation. For more details: - # https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 - # NOTE: We don't need shared_output.record_stream(current_stream()) - # because we synch the streams before using shared_output. - shared_experts_input.record_stream(self.shared_experts_stream) - - # Mark sync start point for the separate shared experts - # stream here since we want to run in parallel with the - # router/gate (next op below) - assert self.shared_experts_stream is not None - self.shared_experts_stream.wait_stream(current_stream()) - - return shared_experts_input - def _maybe_init_dp_chunking(self): if not self.use_dp_chunking: return @@ -379,49 +269,6 @@ def _maybe_init_dp_chunking(self): device=device, ) - def _maybe_apply_shared_experts( - self, - shared_output: torch.Tensor | None, - hidden_states: torch.Tensor, - shared_input: torch.Tensor | None, - order: SharedExpertsOrder, - ) -> torch.Tensor | None: - experts_order, use_shared_experts_stream = self._determine_shared_experts_order( - hidden_states - ) - if order != experts_order: - return None - - assert self.shared_experts is not None - assert shared_output is None - if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - # TODO: fold this in? - hidden_states = self._setup_shared_experts_stream( - hidden_states, - shared_input, - ) - - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self.shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - shared_output = self.shared_experts(hidden_states) - current_stream().wait_stream(self.shared_experts_stream) - else: - shared_output = self.shared_experts(hidden_states) - - if order == SharedExpertsOrder.EXTERNAL: - # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely..... - assert shared_output is not None - shared_output = self._maybe_reduce_shared_out(shared_output) - - return shared_output - def must_reduce_shared_expert_outputs(self) -> bool: """ The shared_experts are typically computed using the RowParallelLinear @@ -449,7 +296,9 @@ def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tens else: return tensor_model_parallel_all_reduce(final_hidden_states) - def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor: + def apply_routed_input_transform( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor | None]: """Apply transform for routed experts (e.g., latent projection). This is called by FusedMoE.forward_native. The original hidden_states @@ -465,21 +314,13 @@ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Ten # ReplicatedLinear returns (output, extra_bias) tuple. # We only need the output tensor; extra_bias is not used here. if isinstance(result, tuple): - return result[0] - return result - return hidden_states - - # TODO: combine with runner.forward_impl_X - def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - self.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out + return result[0], hidden_states + return result, hidden_states + + return ( + hidden_states, + hidden_states if self.shared_experts is not None else None, + ) # XXXXX def _maybe_reduce_output( self, @@ -523,14 +364,10 @@ def _encode_layer_name(self) -> str | ModuleName: def _maybe_pad_hidden_states( self, - original_hidden_states: torch.Tensor | None, + shared_experts_input: torch.Tensor, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, list[int]]: - original_hidden_dim = ( - original_hidden_states.shape[-1] - if original_hidden_states is not None - else hidden_states.shape[-1] - ) + shared_experts_hidden_dim = shared_experts_input.shape[-1] transformed_hidden_dim = hidden_states.shape[-1] if ( not self.quant_method.skip_forward_padding @@ -544,26 +381,31 @@ def _maybe_pad_hidden_states( ) if self.shared_experts is not None: - orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim] + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] else: orig_hidden_dims = [transformed_hidden_dim] return hidden_states, orig_hidden_dims + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + order: SharedExpertsOrder, + ): + if self.shared_experts is not None: + assert shared_experts_input is not None + self.shared_experts.apply(shared_experts_input, order) + def _apply_quant_method( self, layer: torch.nn.Module, - shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: # Run this before quant_method to avoid inplace issues. - shared_input = shared_input if shared_input is not None else hidden_states - shared_output = self._maybe_apply_shared_experts( - shared_output, - shared_input, - shared_input, + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.BEFORE_QUANT_METHOD, ) @@ -584,28 +426,18 @@ def _apply_quant_method( x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, - shared_experts_input=shared_input, + shared_experts_input=shared_experts_input, ) - # TODO: try to clean this up - if isinstance(fused_out, tuple): - if fused_out[0] is None: - assert shared_output is not None - hidden_states = fused_out[1] - else: - assert shared_output is None - shared_output, hidden_states = fused_out - else: - hidden_states = fused_out - - shared_output = self._maybe_apply_shared_experts( - shared_output, - shared_input, - shared_input, + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.AFTER_QUANT_METHOD, ) - return shared_output, hidden_states + return ( + self.shared_experts.output if self.shared_experts is not None else None, + fused_out, + ) def _sequence_parallel_context(self): ctx = get_forward_context() @@ -722,33 +554,28 @@ def _maybe_combine( # # forward # - self.moe_forward (_moe_forward or _moe_forward_shared) - # - forward_chunking_dispatch + # - forward_dispatch # - forward_impl or forward_impl_chunked + # def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # For latent MoE: save ORIGINAL hidden_states before transform - # (shared_experts need original dimension, routed experts use transformed) - original_hidden_states = ( - # or shared_experts is not None? - hidden_states if self.routed_input_transform is not None else None - ) - # Apply transform for routed experts (e.g., latent projection for latent MoE) - hidden_states = self.apply_routed_input_transform(hidden_states) + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( - original_hidden_states, + shared_experts_input, hidden_states, ) - # Note: moe_forward will call forward_dispatch fused_output = self.moe_forward( hidden_states, router_logits, - original_hidden_states, + shared_experts_input, self._encode_layer_name(), ) @@ -759,17 +586,15 @@ def forward_dispatch( layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() router_logits = self._maybe_gate(hidden_states, router_logits) - shared_output = self._maybe_apply_shared_experts( - None, - hidden_states, - shared_input, + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.EXTERNAL, ) @@ -777,18 +602,16 @@ def forward_dispatch( if self.use_dp_chunking: return self.forward_impl_chunked( layer, - shared_output, hidden_states, router_logits, - shared_input, + shared_experts_input, ) else: return self.forward_impl( layer, - shared_output, hidden_states, router_logits, - shared_input, + shared_experts_input, ) def _slice_and_copy_input( @@ -814,10 +637,9 @@ def _slice_and_copy_input( def forward_impl_chunked( self, layer: torch.nn.Module, - shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: final_shared_hidden_states, final_fused_hidden_states = ( self._allocate_dp_chunking_outputs(hidden_states, router_logits) @@ -864,18 +686,17 @@ def forward_impl_chunked( chunk_end, ) - shared_input_chunk = ( - shared_input[chunk_start:chunk_end, :] - if shared_input is not None + shared_experts_input_chunk = ( + shared_experts_input[chunk_start:chunk_end, :] + if shared_experts_input is not None else None ) shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, - shared_output=shared_output, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, - shared_input=shared_input_chunk, + shared_experts_input=shared_experts_input_chunk, ) # Store outputs @@ -900,10 +721,9 @@ def forward_impl_chunked( def forward_impl( self, layer: torch.nn.Module, - shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP @@ -916,10 +736,9 @@ def forward_impl( shared_output, hidden_states = self._apply_quant_method( layer=layer, - shared_output=shared_output, hidden_states=hidden_states, router_logits=router_logits, - shared_input=shared_input, + shared_experts_input=shared_experts_input, ) return self._maybe_combine( diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py new file mode 100644 index 000000000000..ad1bac4dbb03 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from enum import IntEnum + +import torch + +import vllm.envs as envs +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + aux_stream, + current_stream, +) + +logger = init_logger(__name__) + + +class SharedExpertsOrder(IntEnum): + # No shared experts. + NONE = (0,) + + # Get rid of this one? combine with BEFORE? + # Note: this might be important for torch.compile reasons. Can + # get rid of it after _moe_forward is undone. + EXTERNAL = (1,) + + # Called by modular kernel. + INTERNAL = (2,) + + # Called right before quant_method is executed. + BEFORE_QUANT_METHOD = (3,) + + # Called right after quant_method is executed (possibly with streaming). + AFTER_QUANT_METHOD = (4,) + + +# XXXXX add method to prime with shared_experts_input? NO +class SharedExperts: + def __init__( + self, + shared_experts: torch.nn.Module, + moe_config: FusedMoEConfig, + has_separate_shared_experts: bool, # better name + use_dp_chunking: bool, + must_reduce_shared_expert_outputs: bool, # and reduce_results + ): + self._output: torch.Tensor | None = None + self._shared_experts = shared_experts + self._moe_config = moe_config # invariant + self._use_dp_chunking = use_dp_chunking # invariant + + self._has_separate_shared_experts = has_separate_shared_experts # depends on MK + self._must_reduce_shared_expert_outputs = ( + must_reduce_shared_expert_outputs # depends on MK + ) + + # Allow disabling of the separate shared experts stream for + # debug purposes. + # TODO: Remove this after more extensive testings with TP/DP + # and other execution modes + self.use_shared_experts_stream = False + if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: + logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") + self._shared_experts_stream = None + else: + # TODO(rob): enable shared expert overlap with non-cuda-alike. + # aux_stream() returns None on non-cuda-alike platforms. + self._shared_experts_stream = aux_stream() + if self._shared_experts_stream is not None: + logger.debug_once( + "Enabled separate cuda stream for MoE shared_experts", scope="local" + ) + + def _has_external_experts(self) -> bool: + # Disable shared expert overlap if: + # - we are using eplb with non-default backend, because of correctness issues + # - we are using flashinfer with DP, since there nothing to gain + backend = self._moe_config.moe_parallel_config.all2all_backend + return self._shared_experts is not None and not ( + ( + self._moe_config.moe_parallel_config.enable_eplb + and backend != "allgather_reducescatter" + ) + or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) + + def _determine_shared_experts_order( + self, + hidden_states: torch.Tensor, + ) -> tuple[SharedExpertsOrder, bool]: + if self._shared_experts is None: + return SharedExpertsOrder.NONE, False + + if self._has_external_experts(): + return SharedExpertsOrder.EXTERNAL, False + + if ( + not self._has_separate_shared_experts + or not self._moe_config.moe_parallel_config.use_all2all_kernels + ): + return SharedExpertsOrder.INTERNAL, False + + allow_shared_experts_stream = ( + current_platform.is_cuda() + and self._has_separate_shared_experts + and not self._use_dp_chunking + and self._shared_experts_stream is not None + and hidden_states.shape[0] + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + ) + + # Check if we need to run shared experts before matrix multiply because + # matrix multiply may modify the hidden_states. + run_shared_experts_before = ( + self._has_separate_shared_experts and not allow_shared_experts_stream + ) + + if run_shared_experts_before: + return SharedExpertsOrder.BEFORE_QUANT_METHOD, False + else: + return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream + + def _setup_shared_experts_stream( + self, + shared_experts_input: torch.Tensor, + ): + assert self._shared_experts_stream is not None + assert self._moe_config.disable_inplace + + # Record that the clone will be used by shared_experts_stream + # to avoid gc issue from deallocation of hidden_states_clone + # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 + # NOTE: We don't need shared_output.record_stream(current_stream()) + # because we synch the streams before using shared_output. + shared_experts_input.record_stream(self._shared_experts_stream) + + # Mark sync start point for the separate shared experts + # stream here since we want to run in parallel with the + # router/gate (next op below) + assert self._shared_experts_stream is not None + self._shared_experts_stream.wait_stream(current_stream()) + + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: + # Reduce shared expert outputs if necessary, since the MLP + # should have been created with reduce_results=False. + if ( + self._must_reduce_shared_expert_outputs + and get_tensor_model_parallel_world_size() > 1 + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out + + @property + def output(self) -> torch.Tensor | None: + assert self._shared_experts is None == self._output is None + output = self._output + self._output = None + return output + + # add split/join? + + def apply( + self, + shared_experts_input: torch.Tensor, + order: SharedExpertsOrder, + ) -> torch.Tensor | None: + experts_order, use_shared_experts_stream = self._determine_shared_experts_order( + shared_experts_input, + ) + if order != experts_order: + return None + + assert self._shared_experts is not None + assert self._output is None + if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: + # TODO: fold this in? + self._setup_shared_experts_stream(shared_experts_input) + + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self._shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream + self._output = self._shared_experts(shared_experts_input) + current_stream().wait_stream(self._shared_experts_stream) + else: + self._output = self._shared_experts(shared_experts_input) + + if order == SharedExpertsOrder.EXTERNAL: + # TODO: figure out how to combine this with maybe_reduce_output? + # or get rid of it completely..... + assert self._output is not None + self._output = self._maybe_reduce_shared_out(self._output) + + # TODO: do AFTER reduce here? + + return self._output # ? diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index a29d8a7d8dda..e2a5d05320d0 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -93,7 +93,7 @@ def forward_native( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.forward_cuda(layer, x, topk_weights, topk_ids, shared_experts_input) @property @@ -297,7 +297,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.forward( layer=layer, x=x, @@ -322,7 +322,7 @@ def forward_cuda( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.kernel is not None return self.kernel.apply( @@ -343,7 +343,7 @@ def forward_monolithic_cuda( layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: F401 assert self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM @@ -369,7 +369,7 @@ def forward_monolithic_cpu( layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.cpu_fused_moe( layer, x, diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 5b7af3193b03..da717556adbd 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -765,7 +765,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return fused_marlin_moe( x, layer.w13_qweight, diff --git a/vllm/model_executor/layers/quantization/bitsandbytes.py b/vllm/model_executor/layers/quantization/bitsandbytes.py index 716a20090f69..729924663646 100644 --- a/vllm/model_executor/layers/quantization/bitsandbytes.py +++ b/vllm/model_executor/layers/quantization/bitsandbytes.py @@ -483,7 +483,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts # TODO(bnell): Do these need to be called on the hot path? diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 29115fbbc255..8fa25efd1da0 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -349,7 +349,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply( x, @@ -597,7 +597,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -622,7 +622,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply( x, @@ -959,7 +959,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( x, @@ -983,7 +983,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -1123,7 +1123,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( @@ -1607,7 +1607,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.kernel_backend == "Flashinfer" return flashinfer_trtllm_mxint4_moe( x=x, @@ -1634,7 +1634,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.kernel_backend == "Marlin" return fused_marlin_moe( x, @@ -1883,7 +1883,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( @@ -2500,7 +2500,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if layer.enable_eplb: raise NotImplementedError( "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet." diff --git a/vllm/model_executor/layers/quantization/experts_int8.py b/vllm/model_executor/layers/quantization/experts_int8.py index d971f3b5b0d2..301441ff019d 100644 --- a/vllm/model_executor/layers/quantization/experts_int8.py +++ b/vllm/model_executor/layers/quantization/experts_int8.py @@ -141,7 +141,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 5101347cd02a..f967e41a79b3 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -961,7 +961,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -986,7 +986,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/gguf.py b/vllm/model_executor/layers/quantization/gguf.py index 88023349e779..cf72c83d717a 100644 --- a/vllm/model_executor/layers/quantization/gguf.py +++ b/vllm/model_executor/layers/quantization/gguf.py @@ -637,7 +637,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if layer.apply_router_weight_on_input: raise NotImplementedError( "Apply router weight on input is not supported for" diff --git a/vllm/model_executor/layers/quantization/gptq_marlin.py b/vllm/model_executor/layers/quantization/gptq_marlin.py index d7b2a366e1f0..0b43fbd392c9 100644 --- a/vllm/model_executor/layers/quantization/gptq_marlin.py +++ b/vllm/model_executor/layers/quantization/gptq_marlin.py @@ -901,7 +901,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return fused_marlin_moe( x, layer.w13_qweight, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 78644f74d288..a5e1ea52f557 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -932,7 +932,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -957,7 +957,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -1416,7 +1416,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -1441,7 +1441,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index f5c679840432..a327ac17bbc9 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -369,7 +369,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts assert layer.activation == MoEActivation.SILU, ( diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index f992d0f86c4e..16536b85caea 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -1051,7 +1051,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic if layer.enable_eplb: raise NotImplementedError("EPLB is not supported for mxfp4") @@ -1096,7 +1096,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic if layer.enable_eplb: diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 0a5db4e71fdb..fcf03d330d11 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -445,7 +445,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if self.rocm_aiter_moe_enabled: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, @@ -634,7 +634,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, ) @@ -1042,7 +1042,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if not self.emulate: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, From 27ab7694b2249eb36ab96b9ced15bc73b3f9ac3b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Feb 2026 18:47:11 -0500 Subject: [PATCH 011/108] cleanups Signed-off-by: Bill Nell --- .../model_executor/layers/fused_moe/config.py | 10 ++ vllm/model_executor/layers/fused_moe/layer.py | 20 ++-- .../fused_moe/runner/default_moe_runner.py | 98 +++++++++++-------- .../layers/fused_moe/runner/shared_experts.py | 76 ++++++++------ .../model_executor/models/transformers/moe.py | 2 + 5 files changed, 121 insertions(+), 85 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 2500387debe1..cfe50a36cf8e 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -937,6 +937,16 @@ class FusedMoEParallelConfig: all2all_backend: str # all2all backend for MoE communication enable_eplb: bool # whether to enable expert load balancing + @property + def use_dp_chunking(self) -> bool: + return ( + or self.use_deepep_ll_kernels + or self.use_mori_kernels + or self.use_fi_all2allv_kernels + or self.use_fi_nvl_two_sided_kernels + or self.use_nixl_ep_kernels + ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + @property def is_sequence_parallel(self) -> bool: return self.sp_size > 1 diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index bbf6e2ffedb9..c8328ec9b63f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -642,19 +642,16 @@ def _init_shared_experts(self) -> SharedExperts | None: if self._shared_experts is None: return None - reduce_shared_output = ( - self.reduce_results - # XXXX ordering issue - and self.quant_method.moe_mk is not None - and self.quant_method.moe_mk.output_is_reduced() - ) - return SharedExperts( self._shared_experts, moe_config=self.moe_config, - has_separate_shared_experts=not self.quant_method.mk_owns_shared_expert, - use_dp_chunking=self.use_dp_chunking, # XXXXXXXXXXXXXX - must_reduce_shared_expert_outputs=reduce_shared_output, + # Note: For now we must pass quant_method along to SharedExperts so it + # can property determine where the shared experts are supposed to be + # called, i.e. by a MK or by the MoERunner. + # Once the MK can be created upfront, we can just pass in the proper + # flags dervied from the quant_method's MK. + reduce_results=self.reduce_results, + quant_method=self.quant_method, ) def _init_runner(self) -> DefaultMoERunner: @@ -671,7 +668,6 @@ def _init_runner(self) -> DefaultMoERunner: quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, - enable_eplb=self.enable_eplb, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -710,7 +706,7 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.runner._get_shared_experts(), # XXXXXXXXXXXXXXXXX + self.shared_experts, inplace=not self.moe_config.disable_inplace, ) ) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 2e8ff2a9fd54..9a7b32ede979 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -7,7 +7,6 @@ import torch import torch.nn.functional as F -import vllm.envs as envs from vllm.distributed import ( get_ep_group, get_pcp_group, @@ -187,7 +186,6 @@ def __init__( quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, - enable_eplb: bool, ): super().__init__() self.moe_config = moe_config @@ -198,10 +196,10 @@ def __init__( self.quant_method = quant_method self.reduce_results = reduce_results self.enable_dbo = enable_dbo - self.enable_eplb = enable_eplb + self.enable_eplb = moe_config.moe_parallel_config.enable_eplb # Chunked all2all staging tensor - # TODO rename these + # TODO rename these? # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. # TODO: these could be global, i.e. shared by all layers @@ -209,34 +207,31 @@ def __init__( self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() - self.use_dp_chunking = ( - self.moe_config.moe_parallel_config.use_pplx_kernels - or self.moe_config.moe_parallel_config.use_deepep_ll_kernels - or self.moe_config.moe_parallel_config.use_mori_kernels - or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels - or self.moe_config.moe_parallel_config.use_nixl_ep_kernels - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name - self.moe_forward = self._select_forward(layer) + self.moe_forward, self.moe_forward_impl = self._select_forward(layer) - # XXXXXX used by layer.py and lora/layers/fused_moe.py - def _get_shared_experts(self) -> SharedExperts | None: - return self.shared_experts - - def _select_forward(self, layer: torch.nn.Module) -> Callable: + def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: + forward_impl_fn = ( + self.forward_impl_chunked if self.use_dp_chunking else self.forward_impl + ) if current_platform.is_tpu() or current_platform.is_cpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. # Note: CPU doesn't require wrapped forward_impl. - return _moe_forward if self.shared_experts is None else _moe_forward_shared + return ( + _moe_forward if self.shared_experts is None else _moe_forward_shared, + forward_impl_fn, + ) return ( torch.ops.vllm.moe_forward if self.shared_experts is None - else torch.ops.vllm.moe_forward_shared + else torch.ops.vllm.moe_forward_shared, + forward_impl_fn, ) def _maybe_init_dp_chunking(self): @@ -308,6 +303,9 @@ def apply_routed_input_transform( TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be moved inside SharedFusedMoE to all-reduce on the smaller latent dimension. + + Returns (possibly transformed) hidden states and the input for shared + experts (or None if there are no shared experts). """ if self.routed_input_transform is not None: result = self.routed_input_transform(hidden_states) @@ -320,7 +318,7 @@ def apply_routed_input_transform( return ( hidden_states, hidden_states if self.shared_experts is not None else None, - ) # XXXXX + ) def _maybe_reduce_output( self, @@ -364,10 +362,12 @@ def _encode_layer_name(self) -> str | ModuleName: def _maybe_pad_hidden_states( self, - shared_experts_input: torch.Tensor, + shared_experts_input: torch.Tensor | None, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, list[int]]: - shared_experts_hidden_dim = shared_experts_input.shape[-1] + shared_experts_hidden_dim = ( + shared_experts_input.shape[-1] if shared_experts_input is not None else 0 + ) transformed_hidden_dim = hidden_states.shape[-1] if ( not self.quant_method.skip_forward_padding @@ -551,17 +551,39 @@ def _maybe_combine( else: return hidden_states - # - # forward - # - self.moe_forward (_moe_forward or _moe_forward_shared) - # - forward_dispatch - # - forward_impl or forward_impl_chunked - # def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Invoke the fused moe layer. + + Input: + - hidden_states + - router_logits + + Output: + - The new hidden_states. + or + - A tuple of (shared experts output, new hidden_states). + + Calling sequence + - forward + - self.moe_forward (_moe_forward or _moe_forward_shared custom op) + - forward_dispatch + - moe_forward_impl (forward_impl or forward_impl_chunked) + + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due + to the following reasons: + 1. the chunking loop in forward_impl_chunked cannot be compiled by torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward and + _moe_forward_shared must be split. + + If forward_impl_chunked can be implemented via torch.scan we can potentially get + rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into + the 'forward' method. + """ + # Apply transform for routed experts (e.g., latent projection for latent MoE) hidden_states, shared_experts_input = self.apply_routed_input_transform( hidden_states @@ -599,20 +621,12 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - if self.use_dp_chunking: - return self.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - else: - return self.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) + return self.moe_forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) def _slice_and_copy_input( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index ad1bac4dbb03..5081311e57df 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, ) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( aux_stream, @@ -41,31 +44,25 @@ class SharedExpertsOrder(IntEnum): AFTER_QUANT_METHOD = (4,) -# XXXXX add method to prime with shared_experts_input? NO class SharedExperts: def __init__( self, shared_experts: torch.nn.Module, moe_config: FusedMoEConfig, - has_separate_shared_experts: bool, # better name - use_dp_chunking: bool, - must_reduce_shared_expert_outputs: bool, # and reduce_results + quant_method: FusedMoEMethodBase, + reduce_results: bool, ): self._output: torch.Tensor | None = None self._shared_experts = shared_experts - self._moe_config = moe_config # invariant - self._use_dp_chunking = use_dp_chunking # invariant - - self._has_separate_shared_experts = has_separate_shared_experts # depends on MK - self._must_reduce_shared_expert_outputs = ( - must_reduce_shared_expert_outputs # depends on MK - ) + self._moe_config = moe_config + self._quant_method = quant_method + self._reduce_results = reduce_results + self._use_dp_chunking = moe_config.moe_parallel_config.use_dp_chunking # Allow disabling of the separate shared experts stream for # debug purposes. # TODO: Remove this after more extensive testings with TP/DP # and other execution modes - self.use_shared_experts_stream = False if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") self._shared_experts_stream = None @@ -78,6 +75,7 @@ def __init__( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) + @property def _has_external_experts(self) -> bool: # Disable shared expert overlap if: # - we are using eplb with non-default backend, because of correctness issues @@ -91,6 +89,22 @@ def _has_external_experts(self) -> bool: or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels ) + # TODO(bnell): better name + @property + def _has_separate_shared_experts(self) -> bool: + return ( + not self._quant_method.mk_owns_shared_expert + and self._shared_experts is not None + ) + + @property + def _must_reduce_shared_expert_outputs(self) -> bool: + return ( + self._reduce_results + and self._quant_method.moe_mk is not None + and self._quant_method.moe_mk.output_is_reduced() + ) + def _determine_shared_experts_order( self, hidden_states: torch.Tensor, @@ -98,7 +112,7 @@ def _determine_shared_experts_order( if self._shared_experts is None: return SharedExpertsOrder.NONE, False - if self._has_external_experts(): + if self._has_external_experts: return SharedExpertsOrder.EXTERNAL, False if ( @@ -127,10 +141,10 @@ def _determine_shared_experts_order( else: return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream - def _setup_shared_experts_stream( + def _call_with_shared_experts_stream( self, shared_experts_input: torch.Tensor, - ): + ) -> torch.Tensor: assert self._shared_experts_stream is not None assert self._moe_config.disable_inplace @@ -144,9 +158,21 @@ def _setup_shared_experts_stream( # Mark sync start point for the separate shared experts # stream here since we want to run in parallel with the # router/gate (next op below) - assert self._shared_experts_stream is not None self._shared_experts_stream.wait_stream(current_stream()) + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self._shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream + output = self._shared_experts(shared_experts_input) + current_stream().wait_stream(self._shared_experts_stream) + + return output + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: # Reduce shared expert outputs if necessary, since the MLP # should have been created with reduce_results=False. @@ -164,8 +190,6 @@ def output(self) -> torch.Tensor | None: self._output = None return output - # add split/join? - def apply( self, shared_experts_input: torch.Tensor, @@ -174,25 +198,15 @@ def apply( experts_order, use_shared_experts_stream = self._determine_shared_experts_order( shared_experts_input, ) + if order != experts_order: return None assert self._shared_experts is not None assert self._output is None + if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - # TODO: fold this in? - self._setup_shared_experts_stream(shared_experts_input) - - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self._shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - self._output = self._shared_experts(shared_experts_input) - current_stream().wait_stream(self._shared_experts_stream) + self._output = self._call_with_shared_experts_stream(shared_experts_input) else: self._output = self._shared_experts(shared_experts_input) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 5f8352faed50..f65a197abcfc 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -94,6 +94,8 @@ def transformers_moe_forward( self = forward_context.no_compile_layers[layer_name] self._topk_ids = topk_ids # Clone hidden_states because it will be mutated in-place in FusedMoE + # TODO(bnell): figure out a way to avoid calling runner directly. + # it is a hack that the weight are being passed via logits. return self.runner.forward(hidden_states.clone(), topk_weights) From bf8b3f347e366e5d33315b337de861d6cc8637ea Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Feb 2026 18:54:43 -0500 Subject: [PATCH 012/108] fix circular import Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/shared_experts.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 5081311e57df..305697adf699 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -13,8 +13,8 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, ) -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( - FusedMoEMethodBase, +from vllm.model_executor.layers.quantization.base_config import ( + QuantizeMethodBase, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( @@ -49,9 +49,17 @@ def __init__( self, shared_experts: torch.nn.Module, moe_config: FusedMoEConfig, - quant_method: FusedMoEMethodBase, + quant_method: QuantizeMethodBase, reduce_results: bool, ): + from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, + ) + + # quant_method must be a FusedMoEMethodBase but we can't use the type + # due to circular imports. + assert isinstance(quant_method, FusedMoEMethodBase) + self._output: torch.Tensor | None = None self._shared_experts = shared_experts self._moe_config = moe_config From bb541aea61f72418264bf60d0f208fcb6d210d1c Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 16:32:20 -0500 Subject: [PATCH 013/108] fixes Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- .../layers/fused_moe/runner/default_moe_runner.py | 3 +-- vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 4 ++-- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c8328ec9b63f..67d8af288abf 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -741,7 +741,7 @@ def use_ep(self): @property def is_internal_router(self) -> bool: # By default, router/gate is called before FusedMoE forward pass - return self.gate is not None + return self._gate is not None def _maybe_init_expert_routing_tables( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 9a7b32ede979..682e4c0c1224 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -203,12 +203,11 @@ def __init__( # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. # TODO: these could be global, i.e. shared by all layers + self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() - self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking - # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 305697adf699..76a150bd1241 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -120,7 +120,7 @@ def _determine_shared_experts_order( if self._shared_experts is None: return SharedExpertsOrder.NONE, False - if self._has_external_experts: + if self._has_external_experts and not self._use_dp_chunking: return SharedExpertsOrder.EXTERNAL, False if ( @@ -193,7 +193,7 @@ def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: @property def output(self) -> torch.Tensor | None: - assert self._shared_experts is None == self._output is None + assert (self._shared_experts is None) == (self._output is None) output = self._output self._output = None return output From 0eb70c6fb633a9f80eacbf2dd5156336f599ce64 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 17:04:50 -0500 Subject: [PATCH 014/108] renames Signed-off-by: Bill Nell --- tools/ep_kernels/install_python_libraries.sh | 1 + .../fused_moe/runner/default_moe_runner.py | 33 ++++++++++--------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 3372dd10f4dc..c3deb7d6060c 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -103,6 +103,7 @@ pushd "$WORKSPACE" echo "Downloading NVSHMEM ${NVSHMEM_VER} for ${NVSHMEM_SUBDIR} ..." curl -fSL "${NVSHMEM_URL}" -o "${NVSHMEM_FILE}" tar -xf "${NVSHMEM_FILE}" +rm -rf nvshmem mv "${NVSHMEM_FILE%.tar.xz}" nvshmem rm -f "${NVSHMEM_FILE}" rm -rf nvshmem/lib/bin nvshmem/lib/share diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 682e4c0c1224..6dda98a74d63 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -211,12 +211,14 @@ def __init__( # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name - self.moe_forward, self.moe_forward_impl = self._select_forward(layer) + self.forward_entry, self.forward_impl = self._select_forward(layer) def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: + # Select implementation based on presence of DP chunking. forward_impl_fn = ( - self.forward_impl_chunked if self.use_dp_chunking else self.forward_impl + self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl ) + if current_platform.is_tpu() or current_platform.is_cpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. @@ -568,19 +570,20 @@ def forward( Calling sequence - forward - - self.moe_forward (_moe_forward or _moe_forward_shared custom op) + - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - forward_dispatch - - moe_forward_impl (forward_impl or forward_impl_chunked) + - forward_impl (_forward_impl or _forward_impl_chunked) Note: The existence of _moe_forward and _moe_forward_shared custom ops are due to the following reasons: - 1. the chunking loop in forward_impl_chunked cannot be compiled by torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward and - _moe_forward_shared must be split. - - If forward_impl_chunked can be implemented via torch.scan we can potentially get - rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into - the 'forward' method. + 1. the chunking loop in _forward_impl_chunked cannot be compiled by + torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward + and _moe_forward_shared must be split. + + If _forward_impl_chunked can be implemented via torch.scan we can potentially + get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence + into the 'forward' method. """ # Apply transform for routed experts (e.g., latent projection for latent MoE) @@ -593,7 +596,7 @@ def forward( hidden_states, ) - fused_output = self.moe_forward( + fused_output = self.forward_entry( hidden_states, router_logits, shared_experts_input, @@ -620,7 +623,7 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - return self.moe_forward_impl( + return self.forward_impl( layer, hidden_states, router_logits, @@ -647,7 +650,7 @@ def _slice_and_copy_input( out_slice.copy_(orig_slice, non_blocking=True) return out_slice - def forward_impl_chunked( + def _forward_impl_chunked( self, layer: torch.nn.Module, hidden_states: torch.Tensor, @@ -731,7 +734,7 @@ def forward_impl_chunked( assert final_shared_hidden_states is not None return (final_shared_hidden_states, final_fused_hidden_states) - def forward_impl( + def _forward_impl( self, layer: torch.nn.Module, hidden_states: torch.Tensor, From 8351d89eb6fcea3d719da655412c33ac017f2ba2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 17:23:49 -0500 Subject: [PATCH 015/108] add comment Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/default_moe_runner.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 6dda98a74d63..dbb83dd67691 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -77,6 +77,9 @@ def _resolve_layer_name(layer_name: str | ModuleName) -> str: return layer_name.value if isinstance(layer_name, ModuleName) else layer_name +# Note: _moe_forward and _moe_forward_shared should not contain any +# implementation details, They should merely pass along control to +# the runner's 'forward_dispatch' method. def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, From 943e6675c0cfef8f95870033f6174e090c996268 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 17:26:55 -0500 Subject: [PATCH 016/108] more renames Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/shared_experts.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 76a150bd1241..45538daff5f8 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -97,9 +97,8 @@ def _has_external_experts(self) -> bool: or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels ) - # TODO(bnell): better name @property - def _has_separate_shared_experts(self) -> bool: + def _has_mk_owned_shared_experts(self) -> bool: return ( not self._quant_method.mk_owns_shared_expert and self._shared_experts is not None @@ -124,14 +123,14 @@ def _determine_shared_experts_order( return SharedExpertsOrder.EXTERNAL, False if ( - not self._has_separate_shared_experts + not self._has_mk_owned_shared_experts or not self._moe_config.moe_parallel_config.use_all2all_kernels ): return SharedExpertsOrder.INTERNAL, False allow_shared_experts_stream = ( current_platform.is_cuda() - and self._has_separate_shared_experts + and self._has_mk_owned_shared_experts and not self._use_dp_chunking and self._shared_experts_stream is not None and hidden_states.shape[0] @@ -141,7 +140,7 @@ def _determine_shared_experts_order( # Check if we need to run shared experts before matrix multiply because # matrix multiply may modify the hidden_states. run_shared_experts_before = ( - self._has_separate_shared_experts and not allow_shared_experts_stream + self._has_mk_owned_shared_experts and not allow_shared_experts_stream ) if run_shared_experts_before: @@ -202,7 +201,7 @@ def apply( self, shared_experts_input: torch.Tensor, order: SharedExpertsOrder, - ) -> torch.Tensor | None: + ): experts_order, use_shared_experts_stream = self._determine_shared_experts_order( shared_experts_input, ) @@ -220,10 +219,8 @@ def apply( if order == SharedExpertsOrder.EXTERNAL: # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely..... + # or get rid of it completely. assert self._output is not None self._output = self._maybe_reduce_shared_out(self._output) - # TODO: do AFTER reduce here? - - return self._output # ? + # TODO(bnell): potentially do AFTER reduce here insteed of in runner. From 0ae2b4f849f2fb515c1d1354c3b37489a46ec8d7 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Feb 2026 12:55:16 -0500 Subject: [PATCH 017/108] cleanup Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index b6c62b618a09..11332f4cb228 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -172,7 +172,6 @@ def wrapper(*args, **kwargs): moe_state_dict["apply_router_weight_on_input"] = kwargs[ "apply_router_weight_on_input" ] - # TODO: global_num_experts/shared_experts_input? result = func(*args, **kwargs) return result @@ -597,10 +596,6 @@ def forward(self, *args, **kwargs): def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - # @property - # def _shared_experts(self): - # return self.base_layer._shared_experts - @property def quant_method(self): return self.base_layer.quant_method From 657a4ef873834e84c81568ef06493de003cdebea Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Feb 2026 19:38:08 -0500 Subject: [PATCH 018/108] remove memoizing router, not needed yet Signed-off-by: Bill Nell --- .../fused_moe/router/memoizing_router.py | 35 ------------------- 1 file changed, 35 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py deleted file mode 100644 index a55bd2f09d6f..000000000000 --- a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import torch - -from vllm.model_executor.layers.fused_moe.config import ( - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter - - -class MemoizingRouter(FusedMoERouter): - def __init__(self, router: FusedMoERouter): - self.router = router - - def set_capture_fn( - self, - capture_fn: Callable[[torch.Tensor], None] | None, - ) -> None: - self.router.set_capture_fn(capture_fn) - self.results: tuple[torch.Tensor, torch.Tensor] | None = None - - @property - def routing_method_type(self) -> RoutingMethodType: - return self.router.routing_method_type - - def select_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.results is None: - self.results = self.router.select_experts(hidden_states, router_logits) - return self.results From 6b4ef3b3c34eb51c5f3b60850973ebc9f4d4a726 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 18:33:43 -0500 Subject: [PATCH 019/108] fix UBD bug Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 67d8af288abf..772be9e150bb 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -567,7 +567,7 @@ def __init__( device=vllm_config.device_config.device, routing_method=self.routing_method_type, # TODO: in_dtype == out_dtype? - disable_inplace=disable_inplace() or self.shared_experts is not None, + disable_inplace=disable_inplace() or shared_experts is not None, ) if self.moe_config.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( From 125b34b92e293330873a886ed18724ff98269136 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 16:35:53 -0500 Subject: [PATCH 020/108] cleanup merge Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/default_moe_runner.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index dbb83dd67691..91c64a964937 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -37,8 +37,6 @@ from vllm.utils.torch_utils import ( HAS_OPAQUE_TYPE, ModuleName, - aux_stream, - current_stream, direct_register_custom_op, ) from vllm.v1.worker.ubatching import dbo_current_ubatch_id @@ -86,7 +84,7 @@ def _moe_forward( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: - layer = get_layer_from_name(layer_name) + layer = get_layer_from_name(_resolve_layer_name(layer_name)) return layer.runner.forward_dispatch( layer, hidden_states, @@ -110,7 +108,7 @@ def _moe_forward_shared( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(layer_name) + layer = get_layer_from_name(_resolve_layer_name(layer_name)) return layer.runner.forward_dispatch( layer, hidden_states, @@ -141,7 +139,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], # ? + mutates_args=["hidden_states"], # is this still true? fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -150,7 +148,6 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward_shared", op_func=_moe_forward_shared, - mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_shared_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -202,10 +199,8 @@ def __init__( self.enable_eplb = moe_config.moe_parallel_config.enable_eplb # Chunked all2all staging tensor - # TODO rename these? # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. - # TODO: these could be global, i.e. shared by all layers self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None From 78d5ed8647c13c864e2cf32105c1892292ff27c2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 16:46:20 -0500 Subject: [PATCH 021/108] fix merge Signed-off-by: Bill Nell --- tests/kernels/moe/utils.py | 2 -- vllm/model_executor/layers/fused_moe/config.py | 2 +- vllm/model_executor/layers/fused_moe/modular_kernel.py | 8 ++------ 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4b693d8c8a55..2ef4424c2baa 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -603,7 +603,6 @@ def make_shared_experts( def modular_triton_fused_moe( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, - shared_experts: torch.nn.Module | None = None, ) -> FusedMoEKernel: return FusedMoEKernel( maybe_make_prepare_finalize( @@ -613,6 +612,5 @@ def modular_triton_fused_moe( use_monolithic=False, ), TritonExperts(moe_config, quant_config), - shared_experts, inplace=False, ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index cfe50a36cf8e..b3f57e3e376a 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -940,7 +940,7 @@ class FusedMoEParallelConfig: @property def use_dp_chunking(self) -> bool: return ( - or self.use_deepep_ll_kernels + self.use_deepep_ll_kernels or self.use_mori_kernels or self.use_fi_all2allv_kernels or self.use_fi_nvl_two_sided_kernels diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index bbafc682c102..e424d36536b7 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1001,11 +1001,9 @@ def __init__( self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts self.shared_experts = shared_experts - self.moe_parallel_config = moe_parallel_config self.inplace = inplace - moe_parallel_config = fused_experts.moe_config.moe_parallel_config - self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config + self.moe_parallel_config = moe_parallel_config self.is_dp_ep = ( moe_parallel_config is not None and moe_parallel_config.dp_size > 1 @@ -1459,8 +1457,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalize, fused_experts: FusedMoEExperts, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, + shared_experts: SharedExperts | None = None, inplace: bool = False, ): super().__init__() @@ -1475,7 +1472,6 @@ def __init__( prepare_finalize, fused_experts, shared_experts, - moe_parallel_config, inplace, ) From 8a5445a65f19e36b328f6dfd3c48b9c355f096eb Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 17:02:59 -0500 Subject: [PATCH 022/108] fix merge Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 45538daff5f8..d450f21bda34 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -108,8 +108,8 @@ def _has_mk_owned_shared_experts(self) -> bool: def _must_reduce_shared_expert_outputs(self) -> bool: return ( self._reduce_results - and self._quant_method.moe_mk is not None - and self._quant_method.moe_mk.output_is_reduced() + and self._quant_method.moe_kernel is not None + and self._quant_method.moe_kernel.output_is_reduced() ) def _determine_shared_experts_order( From 5ee510f19f6071d5e9a3e54b9ff45db77cafeb85 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 17:15:45 -0500 Subject: [PATCH 023/108] fix typos Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 772be9e150bb..02831d7d3f88 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -649,7 +649,7 @@ def _init_shared_experts(self) -> SharedExperts | None: # can property determine where the shared experts are supposed to be # called, i.e. by a MK or by the MoERunner. # Once the MK can be created upfront, we can just pass in the proper - # flags dervied from the quant_method's MK. + # flags derived from the quant_method's MK. reduce_results=self.reduce_results, quant_method=self.quant_method, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index d450f21bda34..134c500ebeea 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -223,4 +223,4 @@ def apply( assert self._output is not None self._output = self._maybe_reduce_shared_out(self._output) - # TODO(bnell): potentially do AFTER reduce here insteed of in runner. + # TODO(bnell): potentially do AFTER reduce here instead of in runner. From d675882f0666ec9f9271b4fbf0fc79d698be267c Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 17:24:07 +0000 Subject: [PATCH 024/108] fix merge Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/config.py | 1 - vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b3f57e3e376a..83df5e45207b 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -942,7 +942,6 @@ def use_dp_chunking(self) -> bool: return ( self.use_deepep_ll_kernels or self.use_mori_kernels - or self.use_fi_all2allv_kernels or self.use_fi_nvl_two_sided_kernels or self.use_nixl_ep_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 134c500ebeea..05258999ec8c 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -94,7 +94,7 @@ def _has_external_experts(self) -> bool: self._moe_config.moe_parallel_config.enable_eplb and backend != "allgather_reducescatter" ) - or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels + or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels ) @property From 1dc047a73d1211798d224196c7f77ba77e049a82 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 18:44:53 +0000 Subject: [PATCH 025/108] fix format Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/fp8.py | 5 +---- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 57431b946e24..59432c79d853 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -584,10 +584,7 @@ def make_fp8_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=( - shared_experts - if prepare_finalize.supports_async() else None - ), + shared_experts=shared_experts if prepare_finalize.supports_async() else None, inplace=( not moe_config.disable_inplace and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 194513981e41..89cacc8ca01e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -436,10 +436,7 @@ def make_nvfp4_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=( - shared_experts - if prepare_finalize.supports_async() else None - ), + shared_experts=shared_experts if prepare_finalize.supports_async() else None, inplace=False, ) From 111c628aa8c912f2004c80b18f0836ee2e80f04a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 13 Jan 2026 18:44:16 -0500 Subject: [PATCH 026/108] Split of DefaultMoERunner class Signed-off-by: Bill Nell --- .../fused_moe/fused_moe_modular_method.py | 3 + vllm/model_executor/layers/fused_moe/layer.py | 14 +- .../layers/fused_moe/modular_kernel.py | 6 +- .../layers/fused_moe/oracle/fp8.py | 3 + .../fused_moe/router/memoizing_router.py | 35 + .../fused_moe/runner/chunking_moe_runner.py | 219 ++++++ .../fused_moe/runner/default_moe_runner.py | 697 ++---------------- .../fused_moe/runner/moe_runner_base.py | 487 ++++++++++++ .../fused_moe/runner/moe_runner_factory.py | 51 ++ 9 files changed, 879 insertions(+), 636 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 142e180786c6..10db1554f040 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -19,6 +19,9 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) logger = init_logger(__name__) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 02831d7d3f88..e34714c2448a 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -39,8 +39,14 @@ from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) -from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( - DefaultMoERunner, +from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( + MoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import ( + create_moe_runner, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, ) from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -654,11 +660,11 @@ def _init_shared_experts(self) -> SharedExperts | None: quant_method=self.quant_method, ) - def _init_runner(self) -> DefaultMoERunner: + def _init_runner(self) -> MoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. - return DefaultMoERunner( + return create_moe_runner( layer=self, moe_config=self.moe_config, router=self.router, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index e424d36536b7..af805308ab35 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -24,6 +24,10 @@ SharedExperts, SharedExpertsOrder, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, disable_inplace, @@ -995,7 +999,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, - shared_experts: SharedExperts | None, + shared_experts: SharedExperts | None = None, inplace: bool = False, ): self.prepare_finalize = prepare_finalize diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 59432c79d853..236d0a143127 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -21,6 +21,9 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, get_flashinfer_moe_backend, diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py new file mode 100644 index 000000000000..a55bd2f09d6f --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter + + +class MemoizingRouter(FusedMoERouter): + def __init__(self, router: FusedMoERouter): + self.router = router + + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + self.router.set_capture_fn(capture_fn) + self.results: tuple[torch.Tensor, torch.Tensor] | None = None + + @property + def routing_method_type(self) -> RoutingMethodType: + return self.router.routing_method_type + + def select_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.results is None: + self.results = self.router.select_experts(hidden_states, router_logits) + return self.results diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py new file mode 100644 index 000000000000..1130d96ce967 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.forward_context import ( + get_forward_context, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.ubatching import dbo_current_ubatch_id +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + + +class ChunkingMoERunner(MoERunnerBase): + """ + MoE runner wrapper that adds chunked processing to any MoERunnerBase. + + This runner wraps an inner MoERunnerBase and overrides forward_impl to + process large batches by breaking them into smaller chunks. Each chunk + is delegated to the inner runner's forward_impl, making chunking + composable with any runner implementation. + + All MoERunnerBase state (moe_config, router, quant_method, etc.) is + transparently delegated to the inner runner via __getattr__. + ChunkingMoERunner only owns chunking-specific state: the pre-allocated + workspace buffers and the reduce_results override. + + Key behaviors: + - Pre-allocates workspace tensors for CUDA graph compatibility + - Processes chunks via inner.forward_impl per chunk + - Never reduces results (reduce_results always returns False) + """ + + def __init__(self, inner: MoERunnerBase): + # Assert that _maybe_dispatch/_maybe_combine will be nops. + assert inner.moe_config.pcp_size == 1 + + # Skip MoERunnerBase.__init__ — all state is delegated to inner + # via __getattr__. Only chunking-specific state lives here. + self._inner = inner + + # Pre-allocated staging buffers. These need to exist ahead of time + # due to CUDA graph construction needing fixed buffer addresses. + self.batched_hidden_states, self.batched_router_logits = ( + self._init_dp_chunking() + ) + + def __getattr__(self, name): + # Delegate attribute access to the inner runner. This is only + # called when normal lookup (instance __dict__, class MRO) fails, + # so ChunkingMoERunner's own attributes and methods take priority. + return getattr(self._inner, name) + + @property + def reduce_results(self) -> bool: + return False + + def _init_dp_chunking(self) -> list[torch.Tensor]: + states_shape: tuple[int, ...] + logits_shape: tuple[int, ...] + + moe = self.moe_config + + if self.enable_dbo: + states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) + logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) + else: + states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) + logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) + + # Does this need some kind of profiling run check like modular_kernel.py? + return current_workspace_manager().get_simultaneous( + (states_shape, moe.in_dtype), + (logits_shape, moe.router_logits_dtype), + ) + + def _allocate_dp_chunking_outputs( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + # Assert the inputs are of the proper type and shape. + assert self.batched_hidden_states is not None + assert self.batched_router_logits is not None + + assert self.batched_hidden_states.dtype == hidden_states.dtype, ( + f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" + ) + assert self.batched_router_logits.dtype == router_logits.dtype, ( + f"{self.batched_router_logits.dtype} == {router_logits.dtype}" + ) + + # Check size compatibility. + assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) + assert self.batched_router_logits.size(-1) == router_logits.size(-1) + + final_fused_hidden_states = torch.empty_like(hidden_states) + if self.shared_experts is not None: + final_shared_hidden_states = torch.empty_like(hidden_states) + else: + final_shared_hidden_states = None + + return final_shared_hidden_states, final_fused_hidden_states + + def _slice_and_copy_input( + self, + out_slice: torch.Tensor, + orig: torch.Tensor | None, + start: int, + end: int, + ) -> torch.Tensor: + assert orig is not None + slice_size = end - start + orig_slice = orig[start:end, :] + if self.enable_dbo: + assert out_slice.dim() == 3 + batch_buffer_idx = dbo_current_ubatch_id() + out_slice = out_slice[batch_buffer_idx, :] + + assert out_slice.size(0) >= slice_size + out_slice = out_slice[:slice_size, :] + out_slice.copy_(orig_slice, non_blocking=True) + return out_slice + + def forward_impl( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + final_shared_hidden_states, final_fused_hidden_states = ( + self._allocate_dp_chunking_outputs(hidden_states, router_logits) + ) + + ctx = get_forward_context() + # flashinfer_cutlass_kernels can handle: optional DP + TP/EP + max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu + moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens + + # If the input to the MoE is sequence parallel then divide by sp_size + # to find the maximum number of tokens for any individual dispatcher. + if self.moe_config.is_sequence_parallel: + max_tokens_across_dispatchers = cdiv( + max_tokens_across_dispatchers, self.moe_config.sp_size + ) + + num_tokens = hidden_states.size(0) + for chunk_idx, chunk_start_ in enumerate( + range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) + ): + chunk_start = chunk_start_ + chunk_end = min( + chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers + ) + # clamp start and end + chunk_start = min(chunk_start, num_tokens - 1) + chunk_end = min(chunk_end, num_tokens) + chunk_sizes = ctx.dp_metadata.chunked_sizes( + self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx + ) + with chunk_sizes: + hidden_states_chunk = self._slice_and_copy_input( + self.batched_hidden_states, + hidden_states, + chunk_start, + chunk_end, + ) + + router_logits_chunk = self._slice_and_copy_input( + self.batched_router_logits, + router_logits, + chunk_start, + chunk_end, + ) + + shared_experts_input_chunk = ( + shared_experts_input[chunk_start:chunk_end, :] + if shared_experts_input is not None + else None + ) + + # Delegate per-chunk computation to the inner runner. + chunk_result = self._inner.forward_impl( + layer=layer, + hidden_states=hidden_states_chunk, + router_logits=router_logits_chunk, + shared_experts_input=shared_experts_input_chunk, + ) + + # Store outputs + # TODO(bnell): document when chunk_start >= num_tokens + if chunk_start < num_tokens: + if self.shared_experts is not None: + assert isinstance(chunk_result, tuple) + shared_output_chunk, hidden_states_chunk = chunk_result + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + hidden_states_chunk, non_blocking=True + ) + assert shared_output_chunk is not None + assert final_shared_hidden_states is not None + final_shared_hidden_states[chunk_start:chunk_end, :].copy_( + shared_output_chunk, non_blocking=True + ) + else: + assert isinstance(chunk_result, torch.Tensor) + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + chunk_result, non_blocking=True + ) + + if self.shared_experts is None: + return final_fused_hidden_states + else: + assert final_shared_hidden_states is not None + return (final_shared_hidden_states, final_fused_hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 91c64a964937..586ab2834fb5 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -5,17 +5,10 @@ from typing import TYPE_CHECKING import torch -import torch.nn.functional as F from vllm.distributed import ( get_ep_group, get_pcp_group, - tensor_model_parallel_all_reduce, -) -from vllm.forward_context import ( - ForwardContext, - get_forward_context, - is_forward_context_available, ) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( @@ -27,152 +20,39 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, - SharedExpertsOrder, -) -from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import ( - HAS_OPAQUE_TYPE, - ModuleName, - direct_register_custom_op, ) -from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) -def get_layer_from_name(layer_name: str) -> torch.nn.Module: - forward_context: ForwardContext = get_forward_context() - if layer_name == "from_forward_context": - all_moe_layers = forward_context.all_moe_layers - assert all_moe_layers is not None - moe_layer_index = forward_context.moe_layer_index - if moe_layer_index >= len(all_moe_layers): - raise AssertionError( - "We expected the number of MOE layers in `all_moe_layers` " - "to be equal to the number of " - "{vllm.moe_forward, vllm.moe_forward_shared} calls." - ) - layer_name = all_moe_layers[moe_layer_index] - forward_context.moe_layer_index += 1 - return forward_context.no_compile_layers[layer_name] - - -# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; -# on older versions it remains a plain str. -if TYPE_CHECKING: - from typing import TypeAlias - - _layer_name_type: TypeAlias = str | ModuleName -else: - _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str - - -def _resolve_layer_name(layer_name: str | ModuleName) -> str: - return layer_name.value if isinstance(layer_name, ModuleName) else layer_name - - -# Note: _moe_forward and _moe_forward_shared should not contain any -# implementation details, They should merely pass along control to -# the runner's 'forward_dispatch' method. -def _moe_forward( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> torch.Tensor: - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -def _moe_forward_shared( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_shared_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: - # Output shapes: - # - fused_out: same as hidden_states (routed experts use transformed size) - # - shared_out: same as shared_experts_input if provided, else same as - # hidden_states - # (For latent MoE: shared experts use original hidden_size, not latent size) - fused_out = torch.empty_like(hidden_states) - if shared_experts_input is not None: - shared_out = torch.empty_like(shared_experts_input) - else: - shared_out = torch.empty_like(hidden_states) - return shared_out, fused_out - - -direct_register_custom_op( - op_name="moe_forward", - op_func=_moe_forward, - mutates_args=["hidden_states"], # is this still true? - fake_impl=_moe_forward_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -direct_register_custom_op( - op_name="moe_forward_shared", - op_func=_moe_forward_shared, - fake_impl=_moe_forward_shared_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -class DefaultMoERunner(MoERunner): +class DefaultMoERunner(MoERunnerBase): """ - Default implementation of the MoE runner for executing Mixture of Experts layers. + Standard MoE runner implementation for executing Mixture of Experts layers. - This class provides a comprehensive implementation for running MoE computations - with support for: - - Expert routing and token dispatching + This is the primary concrete implementation of MoE execution logic, providing + comprehensive support for standard MoE operations. It handles: + - Expert routing and token dispatching using various routing strategies - Shared experts computation with optional parallel execution using CUDA streams - - Data parallel (DP) chunking for large batch processing - Tensor model parallel and expert parallel operations - - Various quantization methods and custom operators + - Multiple quantization methods and optimized kernel selection - Both monolithic and decomposed expert execution paths + - Integration with various parallel execution modes (TP, EP, DP) + + The runner orchestrates the complete MoE forward pass including routing tokens + to experts, executing expert computations in parallel, and combining results. + It supports advanced features like overlapped execution of shared experts, + optimized kernels for different parallel configurations, and seamless + integration with vLLM's distributed execution framework. - The runner handles the complete MoE forward pass including routing tokens to - experts, executing expert computations, and combining results. It supports - advanced features like overlapped execution of shared experts and optimized - kernels for different parallel execution modes. + This implementation is suitable for most standard MoE use cases. For specialized + scenarios like large batch chunking, alternative runners like ChunkingMoERunner + may be more appropriate. - Eventually, this class will be split up and specialized for different - configurations, e.g. the presence or absence of shared experts, a gate, etc. + Eventually, this class may be split into more specialized implementations + for different configurations (e.g., with/without shared experts, gates, etc.). """ def __init__( @@ -187,307 +67,21 @@ def __init__( reduce_results: bool, enable_dbo: bool, ): - super().__init__() - self.moe_config = moe_config - self.router = router - self.routed_input_transform = routed_input_transform - self.gate = gate - self.shared_experts = shared_experts - self.quant_method = quant_method - self.reduce_results = reduce_results - self.enable_dbo = enable_dbo - self.enable_eplb = moe_config.moe_parallel_config.enable_eplb - - # Chunked all2all staging tensor - # These need to exist ahead of time due to CUDAgraph construction - # needing a fixed buffer address. - self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking - self.batched_hidden_states: torch.Tensor | None = None - self.batched_router_logits: torch.Tensor | None = None - self._maybe_init_dp_chunking() - - # Needed for string -> FusedMoE layer lookup in custom ops. - self.layer_name = layer.layer_name - - self.forward_entry, self.forward_impl = self._select_forward(layer) - - def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: - # Select implementation based on presence of DP chunking. - forward_impl_fn = ( - self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl - ) - - if current_platform.is_tpu() or current_platform.is_cpu(): - # TODO: Once the OOM issue for the TPU backend is resolved, we - # will switch to using the moe_forward custom op. - # Note: CPU doesn't require wrapped forward_impl. - return ( - _moe_forward if self.shared_experts is None else _moe_forward_shared, - forward_impl_fn, - ) - - return ( - torch.ops.vllm.moe_forward - if self.shared_experts is None - else torch.ops.vllm.moe_forward_shared, - forward_impl_fn, - ) - - def _maybe_init_dp_chunking(self): - if not self.use_dp_chunking: - return - - assert self.batched_hidden_states is None - states_shape: tuple[int, ...] - logits_shape: tuple[int, ...] - - moe = self.moe_config - - if self.enable_dbo: - states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) - else: - states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) - - device = torch.accelerator.current_device_index() - self.batched_hidden_states = torch.zeros( - states_shape, - dtype=moe.in_dtype, - device=device, - ) - - self.batched_router_logits = torch.zeros( - logits_shape, - dtype=moe.router_logits_dtype, - device=device, - ) - - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - if self.must_reduce_shared_expert_outputs(): - return final_hidden_states - else: - return tensor_model_parallel_all_reduce(final_hidden_states) - - def apply_routed_input_transform( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Apply transform for routed experts (e.g., latent projection). - - This is called by FusedMoE.forward_native. The original hidden_states - is saved separately so shared experts get [S, hidden_size] while - routed experts get the transformed [S, moe_latent_size]. - - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside SharedFusedMoE to all-reduce on the smaller latent - dimension. - - Returns (possibly transformed) hidden states and the input for shared - experts (or None if there are no shared experts). - """ - if self.routed_input_transform is not None: - result = self.routed_input_transform(hidden_states) - # ReplicatedLinear returns (output, extra_bias) tuple. - # We only need the output tensor; extra_bias is not used here. - if isinstance(result, tuple): - return result[0], hidden_states - return result, hidden_states - - return ( - hidden_states, - hidden_states if self.shared_experts is not None else None, - ) - - def _maybe_reduce_output( - self, - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - trunc_sizes: list[int], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return x[..., :trunc_size] - - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) - - if ( - not self.moe_config.is_sequence_parallel - and not self.use_dp_chunking - and self.reduce_results - and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - ): - func = reduce_and_trunc - else: - func = trunc - - if isinstance(states, tuple): - return tuple( - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] - ) - else: - assert len(trunc_sizes) == 1 - return func(states, trunc_sizes[0]) - - def _encode_layer_name(self) -> str | ModuleName: - if HAS_OPAQUE_TYPE: - return ModuleName(self.layer_name) - # Can be unavailable or None in unittests - if ( - is_forward_context_available() - and get_forward_context().all_moe_layers is not None - ): - return "from_forward_context" - return self.layer_name - - def _maybe_pad_hidden_states( - self, - shared_experts_input: torch.Tensor | None, - hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, list[int]]: - shared_experts_hidden_dim = ( - shared_experts_input.shape[-1] if shared_experts_input is not None else 0 - ) - transformed_hidden_dim = hidden_states.shape[-1] - if ( - not self.quant_method.skip_forward_padding - and self.moe_config.hidden_dim != transformed_hidden_dim - ): - hidden_states = F.pad( - hidden_states, - (0, self.moe_config.hidden_dim - transformed_hidden_dim), - mode="constant", - value=0.0, - ) - - if self.shared_experts is not None: - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] - else: - orig_hidden_dims = [transformed_hidden_dim] - - return hidden_states, orig_hidden_dims - - def _maybe_apply_shared_experts( - self, - shared_experts_input: torch.Tensor | None, - order: SharedExpertsOrder, - ): - if self.shared_experts is not None: - assert shared_experts_input is not None - self.shared_experts.apply(shared_experts_input, order) - - def _apply_quant_method( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - # Run this before quant_method to avoid inplace issues. - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.BEFORE_QUANT_METHOD, - ) - - if self.quant_method.is_monolithic: - fused_out = self.quant_method.apply_monolithic( - layer=layer, - x=hidden_states, - router_logits=router_logits, - ) - else: - topk_weights, topk_ids = self.router.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - ) - - fused_out = self.quant_method.apply( - layer=layer, - x=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_experts_input=shared_experts_input, - ) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.AFTER_QUANT_METHOD, - ) - - return ( - self.shared_experts.output if self.shared_experts is not None else None, - fused_out, - ) - - def _sequence_parallel_context(self): - ctx = get_forward_context() - return ( - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) - if ctx.dp_metadata - else nullcontext() - ) - - def _allocate_dp_chunking_outputs( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - assert self.use_dp_chunking - - # Assert the inputs are of the proper type and shape. - assert self.batched_hidden_states is not None - assert self.batched_router_logits is not None - - assert self.batched_hidden_states.dtype == hidden_states.dtype, ( - f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" - ) - assert self.batched_router_logits.dtype == router_logits.dtype, ( - f"{self.batched_router_logits.dtype} == {router_logits.dtype}" + super().__init__( + layer, + moe_config, + router, + routed_input_transform, + gate, + shared_experts, + quant_method, + reduce_results, + enable_dbo, ) - # Check size compatibility. - assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) - assert self.batched_router_logits.size(-1) == router_logits.size(-1) - - final_fused_hidden_states = torch.empty_like(hidden_states) - if self.shared_experts is not None: - final_shared_hidden_states = torch.empty_like(hidden_states) - else: - final_shared_hidden_states = None - - return final_shared_hidden_states, final_fused_hidden_states - - def _maybe_gate( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - # If router/gate provided, then apply it here. - # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoE via - # separate cuda stream) - if self.gate is not None: - router_logits, _ = self.gate(hidden_states) - return router_logits + @property + def reduce_results(self) -> bool: + return self._reduce_results @property def do_naive_dispatch_combine(self) -> bool: @@ -500,21 +94,45 @@ def _maybe_dispatch( layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # For naive dispatch/combine Dp/Ep, dispatch the hidden states and - # router logits to all experts. - # NOTE: this will be removed once all kernels are migrated into the - # MoEKernel framework. + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + extra_tensor: torch.Tensor | None = None + if self.do_naive_dispatch_combine: - hidden_states, router_logits = get_ep_group().dispatch_router_logits( - hidden_states, + post_quant_allgather = ( + self.moe_config.dp_size > 1 + and self.moe_config.use_ep + and getattr(self.quant_method, "do_post_quant_allgather", False) + ) + + extra_tensors: list[torch.Tensor] | None = None + + if post_quant_allgather: + hidden_states_to_dispatch, extra_tensors = ( + self.quant_method.prepare_dp_allgather_tensor( + layer, hidden_states, router_logits + ) + ) + else: + hidden_states_to_dispatch = hidden_states + + result = get_ep_group().dispatch_router_logits( + hidden_states_to_dispatch, router_logits, self.moe_config.is_sequence_parallel, + extra_tensors=extra_tensors, ) + if len(result) == 3: + hidden_states, router_logits, extra_tensors = result + assert isinstance(extra_tensors, list) and len(extra_tensors) == 1 + extra_tensor = extra_tensors[0] + else: + hidden_states, router_logits = result + # NOTE: Similar with DP, PCP also needs dispatch and combine. For # simplicity, AgRsAll2All was added separately for PCP here. Maybe # we should modify All2AllManager abstraction to better support PCP. + # TODO(bnell): see what we can do here if self.moe_config.pcp_size > 1: hidden_states = get_pcp_group().all_gather( hidden_states, @@ -525,7 +143,7 @@ def _maybe_dispatch( dim=0, ) - return hidden_states, router_logits + return hidden_states, router_logits, extra_tensor def _maybe_combine( self, @@ -542,7 +160,6 @@ def _maybe_combine( hidden_states, dim=0, ) - # need RS for shared_output? if self.shared_experts is not None: assert shared_output is not None @@ -554,195 +171,12 @@ def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Invoke the fused moe layer. - - Input: - - hidden_states - - router_logits - - Output: - - The new hidden_states. - or - - A tuple of (shared experts output, new hidden_states). - - Calling sequence - - forward - - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - - forward_dispatch - - forward_impl (_forward_impl or _forward_impl_chunked) - - Note: The existence of _moe_forward and _moe_forward_shared custom ops are due - to the following reasons: - 1. the chunking loop in _forward_impl_chunked cannot be compiled by - torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward - and _moe_forward_shared must be split. - - If _forward_impl_chunked can be implemented via torch.scan we can potentially - get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence - into the 'forward' method. - """ - - # Apply transform for routed experts (e.g., latent projection for latent MoE) - hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states - ) - - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( - shared_experts_input, - hidden_states, - ) - - fused_output = self.forward_entry( - hidden_states, - router_logits, - shared_experts_input, - self._encode_layer_name(), - ) - - return self._maybe_reduce_output(fused_output, og_hidden_dims) - - def forward_dispatch( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - - router_logits = self._maybe_gate(hidden_states, router_logits) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.EXTERNAL, - ) - - with self._sequence_parallel_context(): - return self.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - def _slice_and_copy_input( - self, - out_slice: torch.Tensor, - orig: torch.Tensor | None, - start: int, - end: int, - ) -> torch.Tensor: - assert orig is not None - slice_size = end - start - orig_slice = orig[start:end, :] - if self.enable_dbo: - assert out_slice.dim() == 3 - batch_buffer_idx = dbo_current_ubatch_id() - out_slice = out_slice[batch_buffer_idx, :] - - assert out_slice.size(0) >= slice_size - out_slice = out_slice[:slice_size, :] - out_slice.copy_(orig_slice, non_blocking=True) - return out_slice - - def _forward_impl_chunked( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - final_shared_hidden_states, final_fused_hidden_states = ( - self._allocate_dp_chunking_outputs(hidden_states, router_logits) - ) - - ctx = get_forward_context() - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens - - # If the input to the MoE is sequence parallel then divide by sp_size - # to find the maximum number of tokens for any individual dispatcher. - if self.moe_config.is_sequence_parallel: - max_tokens_across_dispatchers = cdiv( - max_tokens_across_dispatchers, self.moe_config.sp_size - ) - - num_tokens = hidden_states.size(0) - for chunk_idx, chunk_start_ in enumerate( - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) - ): - chunk_start = chunk_start_ - chunk_end = min( - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers - ) - # clamp start and end - chunk_start = min(chunk_start, num_tokens - 1) - chunk_end = min(chunk_end, num_tokens) - chunk_sizes = ctx.dp_metadata.chunked_sizes( - self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx - ) - with chunk_sizes: - hidden_states_chunk = self._slice_and_copy_input( - self.batched_hidden_states, - hidden_states, - chunk_start, - chunk_end, - ) - - router_logits_chunk = self._slice_and_copy_input( - self.batched_router_logits, - router_logits, - chunk_start, - chunk_end, - ) - - shared_experts_input_chunk = ( - shared_experts_input[chunk_start:chunk_end, :] - if shared_experts_input is not None - else None - ) - - shared_output_chunk, hidden_states_chunk = self._apply_quant_method( - layer=layer, - hidden_states=hidden_states_chunk, - router_logits=router_logits_chunk, - shared_experts_input=shared_experts_input_chunk, - ) - - # Store outputs - # TODO(bnell): document when chunk_start >= num_tokens - if chunk_start < num_tokens: - final_fused_hidden_states[chunk_start:chunk_end, :].copy_( - hidden_states_chunk, non_blocking=True - ) - if self.shared_experts is not None: - assert shared_output_chunk is not None - assert final_shared_hidden_states is not None - final_shared_hidden_states[chunk_start:chunk_end, :].copy_( - shared_output_chunk, non_blocking=True - ) - - if self.shared_experts is None: - return final_fused_hidden_states - else: - assert final_shared_hidden_states is not None - return (final_shared_hidden_states, final_fused_hidden_states) - - def _forward_impl( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain - hidden_states, router_logits = self._maybe_dispatch( + hidden_states, router_logits, extra_tensor = self._maybe_dispatch( layer, hidden_states, router_logits, @@ -751,6 +185,7 @@ def _forward_impl( shared_output, hidden_states = self._apply_quant_method( layer=layer, hidden_states=hidden_states, + extra_tensor=extra_tensor, router_logits=router_logits, shared_experts_input=shared_experts_input, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py new file mode 100644 index 000000000000..076045c923da --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from abc import abstractmethod +from collections.abc import Callable +from contextlib import nullcontext + +import torch +import torch.nn.functional as F + +from vllm.distributed import ( + tensor_model_parallel_all_reduce, +) +from vllm.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + direct_register_custom_op, +) + +logger = init_logger(__name__) + + +def get_layer_from_name(layer_name: str) -> torch.nn.Module: + forward_context: ForwardContext = get_forward_context() + if layer_name == "from_forward_context": + all_moe_layers = forward_context.all_moe_layers + assert all_moe_layers is not None + moe_layer_index = forward_context.moe_layer_index + if moe_layer_index >= len(all_moe_layers): + raise AssertionError( + "We expected the number of MOE layers in `all_moe_layers` " + "to be equal to the number of " + "{vllm.moe_forward, vllm.moe_forward_shared} calls." + ) + layer_name = all_moe_layers[moe_layer_index] + forward_context.moe_layer_index += 1 + return forward_context.no_compile_layers[layer_name] + + +# Note: _moe_forward and _moe_forward_shared should not contain any +# implementation details, They should merely pass along control to +# the runner's 'forward_dispatch' method. +def _moe_forward( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> torch.Tensor: + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +def _moe_forward_shared( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_shared_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + # Output shapes: + # - fused_out: same as hidden_states (routed experts use transformed size) + # - shared_out: same as shared_experts_input if provided, else same as + # hidden_states + # (For latent MoE: shared experts use original hidden_size, not latent size) + fused_out = torch.empty_like(hidden_states) + + if shared_experts_input is not None: + shared_out = torch.empty_like(shared_experts_input) + else: + shared_out = torch.empty_like(hidden_states) + + return shared_out, fused_out + + +direct_register_custom_op( + op_name="moe_forward", + op_func=_moe_forward, + mutates_args=["hidden_states"], # ? + fake_impl=_moe_forward_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +direct_register_custom_op( + op_name="moe_forward_shared", + op_func=_moe_forward_shared, + mutates_args=["hidden_states"], # ? + fake_impl=_moe_forward_shared_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class MoERunnerBase(MoERunner): + """ + Abstract base class providing common functionality for MoE runner implementations. + + This class serves as the foundation for concrete MoE runner implementations by + providing shared state management and common utilities. It handles: + - Common initialization and configuration management + - Shared expert output reduction logic for tensor parallel scenarios + - Base methods for tensor model parallel reductions + - Common properties and utility functions used across different runner types + + Concrete subclasses must implement the abstract methods to define their specific + execution strategies, such as standard execution, chunked processing, or other + specialized approaches. The base class provides the infrastructure while + allowing flexibility in the actual MoE computation implementation. + + Key abstract methods that subclasses must implement: + - reduce_results: Determines whether results should be reduced across ranks + - forward_impl: The core MoE computation logic specific to each runner type + """ + + def __init__( + self, + layer: torch.nn.Module, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_input_transform: torch.nn.Module | None, + gate: torch.nn.Module | None, + shared_experts: SharedExperts | None, + quant_method: FusedMoEMethodBase, + reduce_results: bool, + enable_dbo: bool, + ): + super().__init__() + self.moe_config = moe_config + self.router = router + self.routed_input_transform = routed_input_transform + self.gate = gate + self.shared_experts = shared_experts + self.quant_method = quant_method + self._reduce_results = reduce_results + self.enable_dbo = enable_dbo + self.enable_eplb = moe_config.moe_parallel_config.enable_eplb + + # Needed for string -> FusedMoE layer lookup in custom ops. + self.layer_name = layer.layer_name + + self.forward_entry = self._select_forward(layer) + + def _select_forward(self, layer: torch.nn.Module) -> Callable: + if current_platform.is_tpu() or current_platform.is_cpu(): + # TODO: Once the OOM issue for the TPU backend is resolved, we + # will switch to using the moe_forward custom op. + # Note: CPU doesn't require wrapped forward_impl. + return _moe_forward if self.shared_experts is None else _moe_forward_shared + + return ( + torch.ops.vllm.moe_forward + if self.shared_experts is None + else torch.ops.vllm.moe_forward_shared + ) + + @property + @abstractmethod + def reduce_results(self) -> bool: + raise NotImplementedError + + def must_reduce_shared_expert_outputs(self) -> bool: + """ + The shared_experts are typically computed using the RowParallelLinear + layer. The result of this function is typically used as + the reduce_results argument to the module. + When just tensor-parallel is used, it is not required to reduce + the shared_experts results immediately. Instead we reduce at the + once at the end of the MoE op. (Refer to DeepSeekV2MoE module) + With EP and all2all kernels - this is no longer viable as all + GPU ranks in DP, produce the complete set of hidden_states. + Therefore it is required that we reduce the shared_experts output + early. + """ + return ( + self.quant_method.moe_mk is not None + and self.quant_method.moe_mk.output_is_reduced() + ) + + def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): + """ + Some combine kernels reduce across GPU ranks by default. + """ + if self.must_reduce_shared_expert_outputs(): + return final_hidden_states + else: + return tensor_model_parallel_all_reduce(final_hidden_states) + + def apply_routed_input_transform( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Apply transform for routed experts (e.g., latent projection). + + This is called by FusedMoE.forward_native. The original hidden_states + is saved separately so shared experts get [S, hidden_size] while + routed experts get the transformed [S, moe_latent_size]. + + TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be + moved inside SharedFusedMoE to all-reduce on the smaller latent + dimension. + + Returns (possibly transformed) hidden states and the input for shared + experts (or None if there are no shared experts). + """ + if self.routed_input_transform is not None: + result = self.routed_input_transform(hidden_states) + # ReplicatedLinear returns (output, extra_bias) tuple. + # We only need the output tensor; extra_bias is not used here. + if isinstance(result, tuple): + return result[0], hidden_states + return result, hidden_states + + return ( + hidden_states, + hidden_states if self.shared_experts is not None else None, + ) + + def _maybe_reduce_output( + self, + states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + trunc_sizes: list[int], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: + return x[..., :trunc_size] + + def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: + return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) + + if ( + not self.moe_config.is_sequence_parallel + and self.reduce_results + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + ): + func = reduce_and_trunc + else: + func = trunc + + if isinstance(states, tuple): + return tuple( + [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] + ) + else: + assert len(trunc_sizes) == 1 + return func(states, trunc_sizes[0]) + + def _encode_layer_name(self) -> str: + # Can be unavailable or None in unittests + if ( + is_forward_context_available() + and get_forward_context().all_moe_layers is not None + ): + return "from_forward_context" + return self.layer_name + + def _maybe_pad_hidden_states( + self, + shared_experts_input: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, list[int]]: + shared_experts_hidden_dim = ( + shared_experts_input.shape[-1] if shared_experts_input is not None else 0 + ) + transformed_hidden_dim = hidden_states.shape[-1] + if self.moe_config.hidden_dim != transformed_hidden_dim: + hidden_states = F.pad( + hidden_states, + (0, self.moe_config.hidden_dim - transformed_hidden_dim), + mode="constant", + value=0.0, + ) + + if self.shared_experts is not None: + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] + else: + orig_hidden_dims = [transformed_hidden_dim] + + return hidden_states, orig_hidden_dims + + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + order: SharedExpertsOrder, + ): + if self.shared_experts is not None: + assert shared_experts_input is not None + self.shared_experts.apply(shared_experts_input, order) + + def _apply_quant_method( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + extra_tensor: torch.Tensor | None, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + # Run this before quant_method to avoid inplace issues. + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.BEFORE_QUANT_METHOD, + ) + + # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014). + # Figure out nicer way to do this. + x_arg = hidden_states if extra_tensor is None else (hidden_states, extra_tensor) + + if self.quant_method.is_monolithic: + fused_out = self.quant_method.apply_monolithic( + layer=layer, + x=x_arg, + router_logits=router_logits, + ) + else: + topk_weights, topk_ids = self.router.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + + fused_out = self.quant_method.apply( + layer=layer, + x=x_arg, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts_input=shared_experts_input, + ) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.AFTER_QUANT_METHOD, + ) + + return ( + self.shared_experts.output if self.shared_experts is not None else None, + fused_out, + ) + + def _sequence_parallel_context(self): + ctx = get_forward_context() + return ( + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) + if ctx.dp_metadata + else nullcontext() + ) + + def _maybe_gate( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + # If router/gate provided, then apply it here. + # (Note: This code runs only when "overlapped mode" is on to allow + # parallel execution of shared experts with the FusedMoE via + # separate cuda stream) + if self.gate is not None: + router_logits, _ = self.gate(hidden_states) + return router_logits + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Invoke the fused moe layer. + + Input: + - hidden_states + - router_logits + + Output: + - The new hidden_states. + or + - A tuple of (shared experts output, new hidden_states). + + Calling sequence + - forward + - self.forward_entry (_moe_forward or _moe_forward_shared custom op) + - forward_dispatch + - forward_impl + + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due + to the following reasons: + 1. the chunking loop in _forward_impl_chunked cannot be compiled by + torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward + and _moe_forward_shared must be split. + + If _forward_impl_chunked can be implemented via torch.scan we can potentially + get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence + into the 'forward' method. + """ + + # Apply transform for routed experts (e.g., latent projection for latent MoE) + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) + + hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + shared_experts_input, + hidden_states, + ) + + fused_output = self.forward_entry( + hidden_states, + router_logits, + shared_experts_input, + self._encode_layer_name(), + ) + + return self._maybe_reduce_output(fused_output, og_hidden_dims) + + def forward_dispatch( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # TODO(bnell): this can be removed after MK migration is complete. + layer.ensure_moe_quant_config_init() + + router_logits = self._maybe_gate(hidden_states, router_logits) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.EXTERNAL, + ) + + with self._sequence_parallel_context(): + return self.forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + @abstractmethod + def forward_impl( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py new file mode 100644 index 000000000000..4bda1ca89304 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import ( + ChunkingMoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( + DefaultMoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + + +def create_moe_runner( + layer: torch.nn.Module, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_input_transform: torch.nn.Module | None, + gate: torch.nn.Module | None, + shared_experts: SharedExperts | None, + quant_method: FusedMoEMethodBase, + reduce_results: bool, + enable_dbo: bool, +) -> MoERunner: + runner = DefaultMoERunner( + layer, + moe_config, + router, + routed_input_transform, + gate, + shared_experts, + quant_method, + reduce_results, + enable_dbo, + ) + if moe_config.moe_parallel_config.use_dp_chunking: + return ChunkingMoERunner(runner) + return runner From ced8feca1f662f5fae42a1b6237660f269d1e3e8 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 15:11:39 +0000 Subject: [PATCH 027/108] fix merge Signed-off-by: Bill Nell --- .../fused_moe/fused_moe_modular_method.py | 3 - vllm/model_executor/layers/fused_moe/layer.py | 3 - .../layers/fused_moe/modular_kernel.py | 4 -- .../layers/fused_moe/oracle/fp8.py | 3 - .../fused_moe/runner/chunking_moe_runner.py | 10 +-- .../fused_moe/runner/default_moe_runner.py | 49 +++---------- .../fused_moe/runner/moe_runner_base.py | 69 +++++++++++-------- 7 files changed, 56 insertions(+), 85 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 10db1554f040..142e180786c6 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -19,9 +19,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, ) -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, -) logger = init_logger(__name__) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index e34714c2448a..90e02907e86a 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -48,9 +48,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, ) -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, -) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index af805308ab35..bdf3b7c44721 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -24,10 +24,6 @@ SharedExperts, SharedExpertsOrder, ) -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, - SharedExpertsOrder, -) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, disable_inplace, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 236d0a143127..59432c79d853 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -21,9 +21,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, ) -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, -) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, get_flashinfer_moe_backend, diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py index 1130d96ce967..6626fad557b3 100644 --- a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -19,9 +19,9 @@ class ChunkingMoERunner(MoERunnerBase): """ MoE runner wrapper that adds chunked processing to any MoERunnerBase. - This runner wraps an inner MoERunnerBase and overrides forward_impl to + This runner wraps an inner MoERunnerBase and overrides _forward_impl to process large batches by breaking them into smaller chunks. Each chunk - is delegated to the inner runner's forward_impl, making chunking + is delegated to the inner runner's _forward_impl, making chunking composable with any runner implementation. All MoERunnerBase state (moe_config, router, quant_method, etc.) is @@ -31,7 +31,7 @@ class ChunkingMoERunner(MoERunnerBase): Key behaviors: - Pre-allocates workspace tensors for CUDA graph compatibility - - Processes chunks via inner.forward_impl per chunk + - Processes chunks via inner._forward_impl per chunk - Never reduces results (reduce_results always returns False) """ @@ -126,7 +126,7 @@ def _slice_and_copy_input( out_slice.copy_(orig_slice, non_blocking=True) return out_slice - def forward_impl( + def _forward_impl( self, layer: torch.nn.Module, hidden_states: torch.Tensor, @@ -185,7 +185,7 @@ def forward_impl( ) # Delegate per-chunk computation to the inner runner. - chunk_result = self._inner.forward_impl( + chunk_result = self._inner._forward_impl( layer=layer, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 586ab2834fb5..7b583bda14cc 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable -from contextlib import nullcontext -from typing import TYPE_CHECKING import torch @@ -94,45 +91,21 @@ def _maybe_dispatch( layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - extra_tensor: torch.Tensor | None = None - + ) -> tuple[torch.Tensor, torch.Tensor]: + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and + # router logits to all experts. + # NOTE: this will be removed once all kernels are migrated into the + # MoEKernel framework. if self.do_naive_dispatch_combine: - post_quant_allgather = ( - self.moe_config.dp_size > 1 - and self.moe_config.use_ep - and getattr(self.quant_method, "do_post_quant_allgather", False) - ) - - extra_tensors: list[torch.Tensor] | None = None - - if post_quant_allgather: - hidden_states_to_dispatch, extra_tensors = ( - self.quant_method.prepare_dp_allgather_tensor( - layer, hidden_states, router_logits - ) - ) - else: - hidden_states_to_dispatch = hidden_states - - result = get_ep_group().dispatch_router_logits( - hidden_states_to_dispatch, + hidden_states, router_logits = get_ep_group().dispatch_router_logits( + hidden_states, router_logits, self.moe_config.is_sequence_parallel, - extra_tensors=extra_tensors, ) - if len(result) == 3: - hidden_states, router_logits, extra_tensors = result - assert isinstance(extra_tensors, list) and len(extra_tensors) == 1 - extra_tensor = extra_tensors[0] - else: - hidden_states, router_logits = result - # NOTE: Similar with DP, PCP also needs dispatch and combine. For # simplicity, AgRsAll2All was added separately for PCP here. Maybe # we should modify All2AllManager abstraction to better support PCP. - # TODO(bnell): see what we can do here if self.moe_config.pcp_size > 1: hidden_states = get_pcp_group().all_gather( hidden_states, @@ -143,7 +116,7 @@ def _maybe_dispatch( dim=0, ) - return hidden_states, router_logits, extra_tensor + return hidden_states, router_logits def _maybe_combine( self, @@ -167,8 +140,9 @@ def _maybe_combine( else: return hidden_states - def forward( + def _forward_impl( self, + layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, @@ -176,7 +150,7 @@ def forward( # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain - hidden_states, router_logits, extra_tensor = self._maybe_dispatch( + hidden_states, router_logits = self._maybe_dispatch( layer, hidden_states, router_logits, @@ -185,7 +159,6 @@ def forward( shared_output, hidden_states = self._apply_quant_method( layer=layer, hidden_states=hidden_states, - extra_tensor=extra_tensor, router_logits=router_logits, shared_experts_input=shared_experts_input, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 076045c923da..d06449889c6e 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -3,6 +3,7 @@ from abc import abstractmethod from collections.abc import Callable from contextlib import nullcontext +from typing import TYPE_CHECKING import torch import torch.nn.functional as F @@ -32,6 +33,8 @@ ) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( + HAS_OPAQUE_TYPE, + ModuleName, direct_register_custom_op, ) @@ -55,6 +58,20 @@ def get_layer_from_name(layer_name: str) -> torch.nn.Module: return forward_context.no_compile_layers[layer_name] +# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; +# on older versions it remains a plain str. +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | ModuleName +else: + _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str + + +def _resolve_layer_name(layer_name: str | ModuleName) -> str: + return layer_name.value if isinstance(layer_name, ModuleName) else layer_name + + # Note: _moe_forward and _moe_forward_shared should not contain any # implementation details, They should merely pass along control to # the runner's 'forward_dispatch' method. @@ -62,9 +79,9 @@ def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, - layer_name: str, + layer_name: _layer_name_type, ) -> torch.Tensor: - layer = get_layer_from_name(layer_name) + layer = get_layer_from_name(_resolve_layer_name(layer_name)) return layer.runner.forward_dispatch( layer, hidden_states, @@ -77,7 +94,7 @@ def _moe_forward_fake( hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, - layer_name: str, + layer_name: _layer_name_type, ) -> torch.Tensor: return torch.empty_like(hidden_states) @@ -86,9 +103,9 @@ def _moe_forward_shared( hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, - layer_name: str, + layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(layer_name) + layer = get_layer_from_name(_resolve_layer_name(layer_name)) return layer.runner.forward_dispatch( layer, hidden_states, @@ -101,7 +118,7 @@ def _moe_forward_shared_fake( hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, - layer_name: str, + layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: # Output shapes: # - fused_out: same as hidden_states (routed experts use transformed size) @@ -109,19 +126,17 @@ def _moe_forward_shared_fake( # hidden_states # (For latent MoE: shared experts use original hidden_size, not latent size) fused_out = torch.empty_like(hidden_states) - if shared_experts_input is not None: shared_out = torch.empty_like(shared_experts_input) else: shared_out = torch.empty_like(hidden_states) - return shared_out, fused_out direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], # ? + mutates_args=["hidden_states"], # is this still true? fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -130,7 +145,6 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward_shared", op_func=_moe_forward_shared, - mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_shared_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -154,7 +168,7 @@ class MoERunnerBase(MoERunner): Key abstract methods that subclasses must implement: - reduce_results: Determines whether results should be reduced across ranks - - forward_impl: The core MoE computation logic specific to each runner type + - _forward_impl: The core MoE computation logic specific to each runner type """ def __init__( @@ -189,7 +203,7 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. - # Note: CPU doesn't require wrapped forward_impl. + # Note: CPU doesn't require wrapped _forward_impl. return _moe_forward if self.shared_experts is None else _moe_forward_shared return ( @@ -217,8 +231,8 @@ def must_reduce_shared_expert_outputs(self) -> bool: early. """ return ( - self.quant_method.moe_mk is not None - and self.quant_method.moe_mk.output_is_reduced() + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() ) def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): @@ -287,7 +301,9 @@ def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: assert len(trunc_sizes) == 1 return func(states, trunc_sizes[0]) - def _encode_layer_name(self) -> str: + def _encode_layer_name(self) -> str | ModuleName: + if HAS_OPAQUE_TYPE: + return ModuleName(self.layer_name) # Can be unavailable or None in unittests if ( is_forward_context_available() @@ -333,7 +349,6 @@ def _apply_quant_method( self, layer: torch.nn.Module, hidden_states: torch.Tensor, - extra_tensor: torch.Tensor | None, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: @@ -343,14 +358,10 @@ def _apply_quant_method( SharedExpertsOrder.BEFORE_QUANT_METHOD, ) - # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014). - # Figure out nicer way to do this. - x_arg = hidden_states if extra_tensor is None else (hidden_states, extra_tensor) - if self.quant_method.is_monolithic: fused_out = self.quant_method.apply_monolithic( layer=layer, - x=x_arg, + x=hidden_states, router_logits=router_logits, ) else: @@ -361,7 +372,7 @@ def _apply_quant_method( fused_out = self.quant_method.apply( layer=layer, - x=x_arg, + x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, shared_experts_input=shared_experts_input, @@ -418,18 +429,18 @@ def forward( - forward - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - forward_dispatch - - forward_impl + - _forward_impl Note: The existence of _moe_forward and _moe_forward_shared custom ops are due to the following reasons: - 1. the chunking loop in _forward_impl_chunked cannot be compiled by + 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by torch.compile 2. pytorch cannot handle union types in custom op signatures so _moe_forward and _moe_forward_shared must be split. - If _forward_impl_chunked can be implemented via torch.scan we can potentially - get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence - into the 'forward' method. + If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can + potentially get rid of _moe_forward and _moe_forward_shared and collapse the + whole sequence into the 'forward' method. """ # Apply transform for routed experts (e.g., latent projection for latent MoE) @@ -469,7 +480,7 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - return self.forward_impl( + return self._forward_impl( layer, hidden_states, router_logits, @@ -477,7 +488,7 @@ def forward_dispatch( ) @abstractmethod - def forward_impl( + def _forward_impl( self, layer: torch.nn.Module, hidden_states: torch.Tensor, From 5c8d40d688d13de1c979f85e7b375dd2b3708e27 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 17:27:52 +0000 Subject: [PATCH 028/108] fix merge Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/moe_runner_base.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index d06449889c6e..2a6e5713827b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -321,7 +321,10 @@ def _maybe_pad_hidden_states( shared_experts_input.shape[-1] if shared_experts_input is not None else 0 ) transformed_hidden_dim = hidden_states.shape[-1] - if self.moe_config.hidden_dim != transformed_hidden_dim: + if ( + not self.quant_method.skip_forward_padding + and self.moe_config.hidden_dim != transformed_hidden_dim + ): hidden_states = F.pad( hidden_states, (0, self.moe_config.hidden_dim - transformed_hidden_dim), From 74eeba9d9955408aa661d019a92314582ceef0e1 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 26 Feb 2026 18:00:45 -0500 Subject: [PATCH 029/108] attempt to fix zero experts Signed-off-by: Bill Nell --- .../moe-refactor/LongCat-Flash-Chat-FP8.yaml | 10 ++ .../configs/moe-refactor/config-h100.txt | 1 + vllm/model_executor/layers/fused_moe/layer.py | 43 +++--- .../layers/fused_moe/zero_expert_fused_moe.py | 128 ++++++++++-------- vllm/model_executor/models/longcat_flash.py | 2 +- 5 files changed, 105 insertions(+), 79 deletions(-) create mode 100644 tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml b/tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml new file mode 100644 index 000000000000..ca5c9a00ed37 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml @@ -0,0 +1,10 @@ +model_name: "meituan-longcat/LongCat-Flash-Chat-FP8" +accuracy_threshold: 0.70 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1200 +server_args: >- + --enforce-eager + --max-model-len 4096 + --tensor-parallel-size 8 + --enable-expert-parallel diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt b/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt index 7397fc4e4626..912ba878fae6 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt +++ b/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt @@ -10,3 +10,4 @@ Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml Qwen3-30B-A3B-Fp8-CT-Channel-vllm-cutlass.yaml Qwen3-30B-A3B-BF16-fi-cutlass.yaml Qwen3-30B-A3B-BF16-triton.yaml +LongCat-Flash-Chat-FP8.yaml diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 90e02907e86a..38e663e89a13 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -36,6 +36,9 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( init_aiter_topK_meta_data, ) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) @@ -339,6 +342,7 @@ def __init__( gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, + router: FusedMoERouter | None = None, ): super().__init__() @@ -511,24 +515,27 @@ def __init__( # TODO(bnell): we should not have to create a router if the kernel is # monolithic. - self.router = create_fused_moe_router( - top_k=top_k, - global_num_experts=self.global_num_experts, - eplb_state=self.eplb_state, - renormalize=renormalize, - use_grouped_topk=use_grouped_topk, - num_expert_group=num_expert_group, - topk_group=topk_group, - custom_routing_function=custom_routing_function, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - num_fused_shared_experts=self.num_fused_shared_experts, - enable_eplb=enable_eplb, - # TODO(bnell): once we can construct the MK at init time, we - # can make this a value. - indices_type_getter=lambda: self.quant_method.topk_indices_dtype, - ) + if router is not None: + self.router = router + else: + self.router = create_fused_moe_router( + top_k=top_k, + global_num_experts=self.global_num_experts, + eplb_state=self.eplb_state, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=self.num_fused_shared_experts, + enable_eplb=enable_eplb, + # TODO(bnell): once we can construct the MK at init time, we + # can make this a value. + indices_type_getter=lambda: self.quant_method.topk_indices_dtype, + ) self.routing_method_type: RoutingMethodType = self.router.routing_method_type # Round up hidden size before creating moe_config. diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py index 97d21767f4fc..eaadf0ad6eed 100644 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import contextmanager - import torch -from torch import nn from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) class ZeroExpertFusedMoE(FusedMoE): @@ -25,7 +25,7 @@ def __init__( self, zero_expert_num: int, zero_expert_type: str, - router: nn.Module, + e_score_correction_bias: torch.Tensor | None = None, **kwargs, ): # ZeroExpertFusedMoE manages its own custom_routing_function for memoization @@ -37,19 +37,35 @@ def __init__( "It manages its own for routing memoization." ) - # Automatically slice router's e_score_correction_bias to only include - # real experts (not zero_experts) for the base FusedMoE. - # The full bias will be used temporarily in forward() for routing. - if hasattr(router, "e_score_correction_bias") and "num_experts" in kwargs: + # Slice e_score_correction_bias to only include real experts + # (not zero experts) for the base FusedMoE. The full bias will be + # used temporarily in forward() for routing. + if e_score_correction_bias is not None and "num_experts" in kwargs: num_real_experts = kwargs["num_experts"] - router_bias = router.e_score_correction_bias user_bias = kwargs.get("e_score_correction_bias") - # Use router's bias if: - # 1. User didn't provide bias, or - # 2. User provided full bias (same size as router) - if user_bias is None or user_bias.shape[0] == router_bias.shape[0]: - kwargs["e_score_correction_bias"] = router_bias[:num_real_experts] + if ( + user_bias is None + or user_bias.shape[0] == e_score_correction_bias.shape[0] + ): + kwargs["e_score_correction_bias"] = e_score_correction_bias[ + :num_real_experts + ] + + # Create memoizing routing function BEFORE super().__init__() so it + # gets passed to create_fused_moe_router, which will create a + # CustomRoutingRouter that uses it. The closure captures `self` and + # accesses memoization state at call time (not definition time). + def custom_routing_function(hidden_states, gating_output, topk, renormalize): + """Return memoized `topk_weights` and `topk_ids`.""" + if self._memoized_topk_weights is None or self._memoized_topk_ids is None: + raise RuntimeError( + "ZeroExpertFusedMoE: routing results not memoized. " + "Call select_experts first to compute routing." + ) + return self._memoized_topk_weights, self._memoized_topk_ids + + kwargs["custom_routing_function"] = custom_routing_function # FusedMoE no longer accepts zero_expert_num/zero_expert_type. # We handle zero experts ourselves in forward(). @@ -57,7 +73,8 @@ def __init__( # Store the actual zero_expert_num and zero_expert_type for our own use self._actual_zero_expert_num = zero_expert_num self._actual_zero_expert_type = zero_expert_type - self._router = router # Full router (includes zero experts) + # Full e_score_correction_bias (includes zero experts) + self._full_e_score_correction_bias = e_score_correction_bias # Expose zero_expert_num and zero_expert_type as attributes for # compatibility with quantization methods that check these attributes @@ -68,38 +85,36 @@ def __init__( self._memoized_topk_weights: torch.Tensor | None = None self._memoized_topk_ids: torch.Tensor | None = None - # Create custom_routing_function to reuse memoized routing results - def custom_routing_function(hidden_states, gating_output, topk, renormalize): - """Return memoized `topk_weights` and `topk_ids`.""" - if self._memoized_topk_weights is None or self._memoized_topk_ids is None: - raise RuntimeError( - "ZeroExpertFusedMoE: routing results not memoized. " - "Call select_experts first to compute routing." - ) - return self._memoized_topk_weights, self._memoized_topk_ids - - self.custom_routing_function = custom_routing_function + def select_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Route with full e_score_correction_bias (including zero experts). - @contextmanager - def _temporarily_set_attrs(self, **attrs): - """ - Temporarily set attributes using object.__setattr__ and restore them. - - This bypasses nn.Module.__setattr__ to avoid Dynamo tracing issues. - When PyTorch Dynamo traces the forward pass, it cannot handle - nn.Module.__setattr__ calls (which include parameter registration logic), - resulting in "Unsupported" errors. Using object.__setattr__ directly - sets the attribute without triggering nn.Module's custom __setattr__, - allowing Dynamo to trace the code successfully. + This bypasses self.router (which is a memoizing CustomRoutingRouter) + to perform the actual routing computation with the full bias. """ - originals = {key: getattr(self, key) for key in attrs} - try: - for key, value in attrs.items(): - object.__setattr__(self, key, value) - yield - finally: - for key, value in originals.items(): - object.__setattr__(self, key, value) + if self._full_e_score_correction_bias is not None: + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + e_score_correction_bias=self._full_e_score_correction_bias.data, + topk=self.top_k, + renormalize=self.renormalize, + scoring_func=self.scoring_func, + ) + else: + from vllm.model_executor.layers.fused_moe import fused_topk + + topk_weights, topk_ids = fused_topk( + hidden_states, router_logits, self.top_k, self.renormalize + ) + + if self.routed_scaling_factor != 1.0: + topk_weights = topk_weights * self.routed_scaling_factor + + return topk_weights, topk_ids def _compute_zero_expert_result( self, @@ -138,21 +153,14 @@ def forward( Returns: Combined output from real experts and zero experts """ - # Prepare temporary attribute overrides for routing computation - temp_attrs = { - "custom_routing_function": None, # Disable for first routing - } - if self._router is not None: - temp_attrs["e_score_correction_bias"] = self._router.e_score_correction_bias - - # Compute routing with temporary attributes - # Pass full router_logits (including zero experts) so that zero experts - # can be properly identified in topk_ids - with self._temporarily_set_attrs(**temp_attrs): - topk_weights, topk_ids = self.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, # Full logits (includes zero experts) - ) + # Compute routing with full logits (including zero experts) so that + # zero experts can be properly identified in topk_ids. + # This bypasses self.router (the memoizing CustomRoutingRouter) and + # performs routing directly with the full e_score_correction_bias. + topk_weights, topk_ids = self.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) # Compute zero expert result if needed zero_expert_result = self._compute_zero_expert_result( diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index a9e2c2268ee1..b695e062d276 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -297,7 +297,7 @@ def __init__( self.experts = ZeroExpertFusedMoE( zero_expert_num=config.zero_expert_num, zero_expert_type=config.zero_expert_type, - router=self.router, + e_score_correction_bias=self.router.e_score_correction_bias, num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, From 45a48b058d2d8d62b09140011ad107e043cfec18 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 26 Feb 2026 19:06:00 -0500 Subject: [PATCH 030/108] simplify ZeroExpertFusedMoE and add ZeroExpertRouter Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 114 ++++++++++ .../fused_moe/router/zero_expert_router.py | 115 +++++++++++ .../layers/fused_moe/zero_expert_fused_moe.py | 194 ++++++------------ 3 files changed, 287 insertions(+), 136 deletions(-) create mode 100644 tests/kernels/moe/test_zero_expert_moe.py create mode 100644 vllm/model_executor/layers/fused_moe/router/zero_expert_router.py diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py new file mode 100644 index 000000000000..40ac7d90741d --- /dev/null +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for ZeroExpertFusedMoE. + +Verifies that: +- The ZeroExpertRouter is properly created and used as the layer router. +- A forward pass through ZeroExpertFusedMoE produces correct output. +""" + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) +from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import ( + ZeroExpertFusedMoE, +) +from vllm.v1.worker.workspace import init_workspace_manager + + +@pytest.fixture +def zero_expert_moe(dist_init, default_vllm_config): + """Create a ZeroExpertFusedMoE layer with zero experts.""" + num_experts = 4 + top_k = 2 + hidden_size = 128 + intermediate_size = 256 + zero_expert_num = 1 + + e_score_correction_bias = torch.zeros( + num_experts + zero_expert_num, + dtype=torch.float32, + device="cuda", + ) + + vllm_config = VllmConfig() + vllm_config.compilation_config.static_forward_context = dict() + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + init_workspace_manager(torch.cuda.current_device()) + + layer = ZeroExpertFusedMoE( + zero_expert_num=zero_expert_num, + zero_expert_type="identity", + e_score_correction_bias=e_score_correction_bias, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + params_dtype=torch.bfloat16, + prefix="test_zero_expert_moe", + renormalize=False, + routed_scaling_factor=1.0, + scoring_func="softmax", + ).cuda() + + layer.quant_method.process_weights_after_loading(layer) + + yield layer, vllm_config + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens): + """Verify that ZeroExpertFusedMoE creates a ZeroExpertRouter.""" + layer, _ = zero_expert_moe + assert isinstance(layer.router, ZeroExpertRouter), ( + f"Expected ZeroExpertRouter but got {type(layer.router).__name__}." + ) + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens): + """Verify that custom_routing_function is not set (routing is handled + by ZeroExpertRouter, not a memoizing closure).""" + layer, _ = zero_expert_moe + assert layer.custom_routing_function is None + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): + """Run a forward pass through ZeroExpertFusedMoE and verify output shape.""" + layer, vllm_config = zero_expert_moe + + hidden_size = layer.hidden_size + num_experts = 4 + zero_expert_num = 1 + total_experts = num_experts + zero_expert_num + + hidden_states = torch.randn( + num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda" + ) + router_logits = torch.randn( + num_tokens, total_experts, dtype=torch.float32, device="cuda" + ) + + # Initialize weights to small random values to avoid NaN from + # uninitialized memory. + with torch.no_grad(): + for param in layer.parameters(): + if param.dtype.is_floating_point: + param.normal_(0, 0.01) + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + get_forward_context().all_moe_layers = None + output = layer.forward(hidden_states, router_logits) + + assert output.shape == hidden_states.shape, ( + f"Expected output shape {hidden_states.shape}, got {output.shape}" + ) + assert output.dtype == hidden_states.dtype + assert not torch.isnan(output).any(), "Output contains NaN values" diff --git a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py new file mode 100644 index 000000000000..c87070bc5acf --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import torch + +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, + get_routing_method_type, +) +from vllm.model_executor.layers.fused_moe.fused_moe import ( + zero_experts_compute_triton, +) +from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) + + +class ZeroExpertRouter(BaseRouter): + """Router that handles zero expert computation as part of routing. + + Routes over all experts (real + zero) using full e_score_correction_bias. + Computes zero expert identity contributions as a side effect during routing. + Remaps zero expert IDs to real expert ID 0 (with weight 0) so downstream + MoE computation can ignore them. + """ + + def __init__( + self, + top_k: int, + global_num_experts: int, + eplb_state: EplbLayerState, + e_score_correction_bias: torch.Tensor, + num_logical_experts: int, + zero_expert_type: str, + scoring_func: str = "softmax", + renormalize: bool = False, + routed_scaling_factor: float = 1.0, + enable_eplb: bool = False, + indices_type_getter: Callable[[], torch.dtype | None] | None = None, + ): + super().__init__( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + enable_eplb=enable_eplb, + indices_type_getter=indices_type_getter, + ) + self.e_score_correction_bias = e_score_correction_bias + self.num_logical_experts = num_logical_experts + self.zero_expert_type = zero_expert_type + self.scoring_func = scoring_func + self.renormalize = renormalize + self.routed_scaling_factor = routed_scaling_factor + self._zero_expert_output: torch.Tensor | None = None + + @property + def routing_method_type(self) -> RoutingMethodType: + return get_routing_method_type( + scoring_func=self.scoring_func, + top_k=self.top_k, + renormalize=self.renormalize, + num_expert_group=None, + has_e_score_bias=True, + ) + + def _compute_routing( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + indices_type: torch.dtype | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute routing with full bias, compute zero expert output, + mask zero expert IDs.""" + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + e_score_correction_bias=self.e_score_correction_bias.data, + topk=self.top_k, + renormalize=self.renormalize, + scoring_func=self.scoring_func, + indices_type=indices_type, + ) + + if self.routed_scaling_factor != 1.0: + topk_weights *= self.routed_scaling_factor + + # Compute zero expert output using pre-EPLB topk_ids/weights. + # zero_experts_compute_triton modifies its inputs in-place, so + # pass clones. + self._zero_expert_output = zero_experts_compute_triton( + expert_indices=topk_ids.clone(), + expert_scales=topk_weights.clone(), + num_experts=self.num_logical_experts, + zero_expert_type=self.zero_expert_type, + hidden_states=hidden_states, + ) + + # Mask zero expert entries: remap zero expert IDs to 0 with weight 0 + # so downstream MoE computation ignores them. + zero_mask = topk_ids >= self.num_logical_experts + topk_ids[zero_mask] = 0 + topk_weights[zero_mask] = 0.0 + + return topk_weights, topk_ids + + @property + def zero_expert_output(self) -> torch.Tensor | None: + """Retrieve and clear the zero expert output.""" + output = self._zero_expert_output + self._zero_expert_output = None + return output diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py index eaadf0ad6eed..5530bce9d26e 100644 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py @@ -3,10 +3,9 @@ import torch -from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton from vllm.model_executor.layers.fused_moe.layer import FusedMoE -from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( - fused_topk_bias, +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, ) @@ -16,9 +15,10 @@ class ZeroExpertFusedMoE(FusedMoE): Zero experts perform identity operations (scaled pass-through) instead of full MLP computations. - This class uses memoization to avoid redundant routing computation: - routing is computed once and reused for both zero expert computation - and the main FusedMoE forward pass. + Uses a ZeroExpertRouter as the layer's main router. The router handles + routing over all experts (real + zero) with the full e_score_correction_bias, + computes zero expert contributions as a side effect, and remaps zero expert + IDs so downstream MoE computation only processes real experts. """ def __init__( @@ -28,22 +28,27 @@ def __init__( e_score_correction_bias: torch.Tensor | None = None, **kwargs, ): - # ZeroExpertFusedMoE manages its own custom_routing_function for memoization assert ( "custom_routing_function" not in kwargs or kwargs.get("custom_routing_function") is None ), ( "ZeroExpertFusedMoE does not support external custom_routing_function. " - "It manages its own for routing memoization." + "Routing is handled by ZeroExpertRouter." ) + assert "router" not in kwargs or kwargs.get("router") is None, ( + "ZeroExpertFusedMoE creates its own ZeroExpertRouter. Do not pass a router." + ) + + # Remove custom_routing_function from kwargs if present + kwargs.pop("custom_routing_function", None) + kwargs.pop("router", None) + # Slice e_score_correction_bias to only include real experts - # (not zero experts) for the base FusedMoE. The full bias will be - # used temporarily in forward() for routing. - if e_score_correction_bias is not None and "num_experts" in kwargs: - num_real_experts = kwargs["num_experts"] + # for the base FusedMoE router factory (which we'll replace anyway). + num_real_experts = kwargs["num_experts"] + if e_score_correction_bias is not None: user_bias = kwargs.get("e_score_correction_bias") - if ( user_bias is None or user_bias.shape[0] == e_score_correction_bias.shape[0] @@ -52,99 +57,48 @@ def __init__( :num_real_experts ] - # Create memoizing routing function BEFORE super().__init__() so it - # gets passed to create_fused_moe_router, which will create a - # CustomRoutingRouter that uses it. The closure captures `self` and - # accesses memoization state at call time (not definition time). - def custom_routing_function(hidden_states, gating_output, topk, renormalize): - """Return memoized `topk_weights` and `topk_ids`.""" - if self._memoized_topk_weights is None or self._memoized_topk_ids is None: - raise RuntimeError( - "ZeroExpertFusedMoE: routing results not memoized. " - "Call select_experts first to compute routing." - ) - return self._memoized_topk_weights, self._memoized_topk_ids - - kwargs["custom_routing_function"] = custom_routing_function - - # FusedMoE no longer accepts zero_expert_num/zero_expert_type. - # We handle zero experts ourselves in forward(). super().__init__(**kwargs) - # Store the actual zero_expert_num and zero_expert_type for our own use - self._actual_zero_expert_num = zero_expert_num - self._actual_zero_expert_type = zero_expert_type - # Full e_score_correction_bias (includes zero experts) - self._full_e_score_correction_bias = e_score_correction_bias - - # Expose zero_expert_num and zero_expert_type as attributes for - # compatibility with quantization methods that check these attributes - self.zero_expert_num = 0 - self.zero_expert_type = None - # Memoization state for routing results - self._memoized_topk_weights: torch.Tensor | None = None - self._memoized_topk_ids: torch.Tensor | None = None + # Replace the factory-created router with our ZeroExpertRouter. + # Uses self.eplb_state created by super().__init__() so EPLB state + # is shared between the layer and the router. + self.router = ZeroExpertRouter( + top_k=self.top_k, + global_num_experts=self.global_num_experts, + eplb_state=self.eplb_state, + e_score_correction_bias=e_score_correction_bias, + num_logical_experts=self.logical_num_experts, + zero_expert_type=zero_expert_type, + scoring_func=self.scoring_func, + renormalize=self.renormalize, + routed_scaling_factor=self.routed_scaling_factor, + enable_eplb=self.enable_eplb, + indices_type_getter=lambda: self.quant_method.topk_indices_dtype, + ) - def select_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Route with full e_score_correction_bias (including zero experts). + # Update routing_method_type to match the new router + self.routing_method_type = self.router.routing_method_type - This bypasses self.router (which is a memoizing CustomRoutingRouter) - to perform the actual routing computation with the full bias. - """ - if self._full_e_score_correction_bias is not None: - topk_weights, topk_ids = fused_topk_bias( - hidden_states=hidden_states, - gating_output=router_logits, - e_score_correction_bias=self._full_e_score_correction_bias.data, - topk=self.top_k, - renormalize=self.renormalize, - scoring_func=self.scoring_func, - ) - else: - from vllm.model_executor.layers.fused_moe import fused_topk - - topk_weights, topk_ids = fused_topk( - hidden_states, router_logits, self.top_k, self.renormalize - ) - - if self.routed_scaling_factor != 1.0: - topk_weights = topk_weights * self.routed_scaling_factor - - return topk_weights, topk_ids - - def _compute_zero_expert_result( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - ) -> torch.Tensor | None: - """Compute zero expert results using pre-computed routing.""" - if ( - self._actual_zero_expert_num is None - or self._actual_zero_expert_num <= 0 - or self._actual_zero_expert_type is None - ): - return None - - return zero_experts_compute_triton( - expert_indices=topk_ids.clone(), - expert_scales=topk_weights.clone(), - num_experts=self.logical_num_experts, - zero_expert_type=self._actual_zero_expert_type, - hidden_states=hidden_states, - ) + # Re-init runner with the new router + self.runner = self._init_runner() + + # Expose zero_expert_num=0 and zero_expert_type=None for + # compatibility with quantization methods that check these attributes. + # The actual zero expert handling is done by ZeroExpertRouter. + self.zero_expert_num = 0 + self.zero_expert_type = None def forward( self, hidden_states: torch.Tensor, - router_logits: torch.Tensor, # Full logits including zero experts + router_logits: torch.Tensor, ) -> torch.Tensor: """ - Forward pass with zero expert support and routing memoization. + Forward pass with zero expert support. + + The ZeroExpertRouter handles routing with full logits (including zero + experts), computes zero expert contributions internally, and returns + masked topk_ids suitable for real expert MoE computation. Args: hidden_states: Input hidden states @@ -153,45 +107,13 @@ def forward( Returns: Combined output from real experts and zero experts """ - # Compute routing with full logits (including zero experts) so that - # zero experts can be properly identified in topk_ids. - # This bypasses self.router (the memoizing CustomRoutingRouter) and - # performs routing directly with the full e_score_correction_bias. - topk_weights, topk_ids = self.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - ) - - # Compute zero expert result if needed - zero_expert_result = self._compute_zero_expert_result( - hidden_states=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - ) - - # Memoize routing results for reuse in super().forward() - self._memoized_topk_weights = topk_weights - self._memoized_topk_ids = topk_ids - - # Slice router_logits for real experts only - router_logits_sliced = router_logits[..., : self.logical_num_experts] - - # Compute real expert results (will reuse memoized routing via - # custom_routing_function) - # zero_expert_num is already 0, so FusedMoE won't handle zero experts - fused_out = super().forward( - hidden_states=hidden_states, - router_logits=router_logits_sliced, - ) - - # Combine results - # Both zero_expert_result and fused_out are computed from the same - # hidden_states, so they should be on the same device. - if zero_expert_result is not None: - fused_out = fused_out + zero_expert_result - - # Clear memoization after use - self._memoized_topk_weights = None - self._memoized_topk_ids = None + # The router handles full logits internally: routes over all experts + # (real + zero), computes zero expert output, masks zero expert IDs. + fused_out = super().forward(hidden_states, router_logits) + + # Retrieve zero expert output computed during routing + zero_expert_output = self.router.zero_expert_output + if zero_expert_output is not None: + fused_out = fused_out + zero_expert_output return fused_out From 65be7f7feb3560b672d46e62f9f6d55e3c727b6e Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 26 Feb 2026 19:17:52 -0500 Subject: [PATCH 031/108] add value test Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 143 +++++++++++++++++++++- 1 file changed, 141 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 40ac7d90741d..7713e414ea3a 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -5,6 +5,7 @@ Verifies that: - The ZeroExpertRouter is properly created and used as the layer router. - A forward pass through ZeroExpertFusedMoE produces correct output. +- The output decomposes correctly into real expert + zero expert contributions. """ import pytest @@ -12,6 +13,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe.layer import FusedMoE from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( ZeroExpertRouter, ) @@ -26,8 +28,10 @@ def zero_expert_moe(dist_init, default_vllm_config): """Create a ZeroExpertFusedMoE layer with zero experts.""" num_experts = 4 top_k = 2 - hidden_size = 128 - intermediate_size = 256 + # hidden_size must be >= 256 for the zero expert identity kernel to + # produce output (its BLOCK_SIZE=256 causes grid=0 when hidden_dim<256). + hidden_size = 256 + intermediate_size = 512 zero_expert_num = 1 e_score_correction_bias = torch.zeros( @@ -112,3 +116,138 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): ) assert output.dtype == hidden_states.dtype assert not torch.isnan(output).any(), "Output contains NaN values" + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): + """Validate that ZeroExpertFusedMoE output equals real expert output + plus zero expert contribution. + + The key invariant is: + layer.forward(h, r) == FusedMoE.forward(h, r) + zero_expert_output + + FusedMoE.forward() computes only the real expert MoE output (the + ZeroExpertRouter masks zero expert entries to weight=0), while the + zero expert contribution is computed as a side effect during routing + and added on top by ZeroExpertFusedMoE.forward(). + """ + layer, vllm_config = zero_expert_moe + num_experts = 4 + zero_expert_num = 1 + total_experts = num_experts + zero_expert_num + + hidden_states = torch.randn( + num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + router_logits = torch.randn( + num_tokens, total_experts, dtype=torch.float32, device="cuda" + ) + + with torch.no_grad(): + for param in layer.parameters(): + if param.dtype.is_floating_point: + param.normal_(0, 0.01) + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + get_forward_context().all_moe_layers = None + + # Get the real expert output only (bypasses ZeroExpertFusedMoE.forward, + # calls FusedMoE.forward directly). The ZeroExpertRouter still runs and + # stores zero_expert_output as a side effect. + real_output = FusedMoE.forward(layer, hidden_states, router_logits) + zero_output = layer.router.zero_expert_output + + # Get the full combined output. + full_output = layer.forward(hidden_states, router_logits) + + assert zero_output is not None, "Zero expert output should not be None" + assert not torch.isnan(real_output).any(), "Real expert output has NaN" + assert not torch.isnan(zero_output).any(), "Zero expert output has NaN" + assert not torch.isnan(full_output).any(), "Full output has NaN" + + expected = real_output + zero_output + torch.testing.assert_close( + full_output, + expected, + atol=0, + rtol=0, + msg="ZeroExpertFusedMoE output should equal real expert output " + "plus zero expert contribution", + ) + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens): + """Validate zero expert identity behavior. + + When routing strongly favors the zero expert, its contribution should + be a scaled version of hidden_states (identity operation). We verify + this by manually computing the expected zero expert output from the + routing weights and comparing against what the router produces. + """ + layer, vllm_config = zero_expert_moe + num_experts = 4 + zero_expert_num = 1 + total_experts = num_experts + zero_expert_num + + hidden_states = torch.randn( + num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + # Strongly bias toward the zero expert (index 4). + router_logits = torch.full( + (num_tokens, total_experts), -10.0, dtype=torch.float32, device="cuda" + ) + router_logits[:, num_experts] = 10.0 # zero expert gets high logit + + with torch.no_grad(): + for param in layer.parameters(): + if param.dtype.is_floating_point: + param.normal_(0, 0.01) + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + get_forward_context().all_moe_layers = None + + # Run routing to get topk_weights/topk_ids before masking. + from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, + ) + + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + e_score_correction_bias=layer.router.e_score_correction_bias.data, + topk=layer.top_k, + renormalize=layer.router.renormalize, + scoring_func=layer.router.scoring_func, + ) + + # Manually compute expected zero expert identity output: + # For each token, sum routing weights assigned to zero expert slots, + # then multiply by hidden_states. + zero_mask = topk_ids >= num_experts + zero_weight_per_token = (topk_weights * zero_mask.float()).sum( + dim=-1, keepdim=True + ) + expected_zero_output = (hidden_states.float() * zero_weight_per_token).to( + hidden_states.dtype + ) + + # Run the layer forward to trigger routing and get the actual + # zero expert output from the router. + FusedMoE.forward(layer, hidden_states, router_logits) + actual_zero_output = layer.router.zero_expert_output + + assert actual_zero_output is not None + assert zero_mask.any(), ( + "With high zero expert logit, at least some slots should route " + "to the zero expert" + ) + + torch.testing.assert_close( + actual_zero_output, + expected_zero_output, + atol=1e-3, + rtol=1e-3, + msg="Zero expert identity output should equal " + "hidden_states * sum(zero_expert_weights)", + ) From 3b88ab26c17ef0c3a06a73a122457be46a84beed Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 15:04:12 -0500 Subject: [PATCH 032/108] move ZeroExpertRouter construction into router factory Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 3 ++ .../layers/fused_moe/router/router_factory.py | 42 ++++++++++++++++-- .../layers/fused_moe/zero_expert_fused_moe.py | 44 ++++--------------- 3 files changed, 50 insertions(+), 39 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 38e663e89a13..532386b68d99 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -343,6 +343,7 @@ def __init__( shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, router: FusedMoERouter | None = None, + zero_expert_type: str | None = None, ): super().__init__() @@ -535,6 +536,8 @@ def __init__( # TODO(bnell): once we can construct the MK at init time, we # can make this a value. indices_type_getter=lambda: self.quant_method.topk_indices_dtype, + zero_expert_type=zero_expert_type, + num_logical_experts=self.logical_num_experts, ) self.routing_method_type: RoutingMethodType = self.router.routing_method_type diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index 11027e894bee..42d418d7e537 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import ( RoutingSimulatorRouter, ) +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) EMPTY_EPLB_STATE: EplbLayerState = EplbLayerState() @@ -49,6 +52,9 @@ def create_fused_moe_router( # eplb parameters enable_eplb: bool = False, eplb_state: EplbLayerState = EMPTY_EPLB_STATE, + # zero expert parameters + zero_expert_type: str | None = None, + num_logical_experts: int | None = None, ) -> FusedMoERouter: """ Factory function to create the appropriate FusedMoERouter subclass based on @@ -56,10 +62,11 @@ def create_fused_moe_router( The selection logic follows this priority order: 1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set - 2. GroupedTopKRouter - if use_grouped_topk is True - 3. CustomRoutingRouter - if custom_routing_function is not None - 4. FusedTopKBiasRouter - if e_score_correction_bias is not None - 5. FusedTopKRouter - default fallback + 2. ZeroExpertRouter - if zero_expert_type is not None + 3. GroupedTopKRouter - if use_grouped_topk is True + 4. CustomRoutingRouter - if custom_routing_function is not None + 5. FusedTopKBiasRouter - if e_score_correction_bias is not None + 6. FusedTopKRouter - default fallback Common arguments: top_k: Number of experts to select per token @@ -86,6 +93,12 @@ def create_fused_moe_router( enable_eplb: Whether EPLB is enabled eplb_state: EPLB (Expert Parallelism Load Balancing) state + Zero expert arguments: + zero_expert_type: Type of zero expert (e.g. identity). If not None, + creates a ZeroExpertRouter. + num_logical_experts: Number of real (non-zero) experts. Required when + zero_expert_type is not None. + Returns: An instance of the appropriate FusedMoERouter subclass """ @@ -100,6 +113,27 @@ def create_fused_moe_router( indices_type_getter=indices_type_getter, ) + if zero_expert_type is not None: + assert num_logical_experts is not None, ( + "num_logical_experts is required when zero_expert_type is set" + ) + assert e_score_correction_bias is not None, ( + "e_score_correction_bias is required when zero_expert_type is set" + ) + return ZeroExpertRouter( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + e_score_correction_bias=e_score_correction_bias, + num_logical_experts=num_logical_experts, + zero_expert_type=zero_expert_type, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + enable_eplb=enable_eplb, + indices_type_getter=indices_type_getter, + ) + if use_grouped_topk: assert custom_routing_function is None if num_expert_group is None or topk_group is None: diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py index 5530bce9d26e..4ebfa0d3637a 100644 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py @@ -4,9 +4,6 @@ import torch from vllm.model_executor.layers.fused_moe.layer import FusedMoE -from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( - ZeroExpertRouter, -) class ZeroExpertFusedMoE(FusedMoE): @@ -44,43 +41,20 @@ def __init__( kwargs.pop("custom_routing_function", None) kwargs.pop("router", None) - # Slice e_score_correction_bias to only include real experts - # for the base FusedMoE router factory (which we'll replace anyway). + # Pass the full e_score_correction_bias (real + zero experts) and + # zero_expert_type through to the router factory, which will create + # a ZeroExpertRouter. num_real_experts = kwargs["num_experts"] if e_score_correction_bias is not None: - user_bias = kwargs.get("e_score_correction_bias") - if ( - user_bias is None - or user_bias.shape[0] == e_score_correction_bias.shape[0] - ): - kwargs["e_score_correction_bias"] = e_score_correction_bias[ - :num_real_experts - ] + kwargs["e_score_correction_bias"] = e_score_correction_bias + kwargs["zero_expert_type"] = zero_expert_type super().__init__(**kwargs) - # Replace the factory-created router with our ZeroExpertRouter. - # Uses self.eplb_state created by super().__init__() so EPLB state - # is shared between the layer and the router. - self.router = ZeroExpertRouter( - top_k=self.top_k, - global_num_experts=self.global_num_experts, - eplb_state=self.eplb_state, - e_score_correction_bias=e_score_correction_bias, - num_logical_experts=self.logical_num_experts, - zero_expert_type=zero_expert_type, - scoring_func=self.scoring_func, - renormalize=self.renormalize, - routed_scaling_factor=self.routed_scaling_factor, - enable_eplb=self.enable_eplb, - indices_type_getter=lambda: self.quant_method.topk_indices_dtype, - ) - - # Update routing_method_type to match the new router - self.routing_method_type = self.router.routing_method_type - - # Re-init runner with the new router - self.runner = self._init_runner() + # Fix self.e_score_correction_bias to only cover real experts, + # for compatibility with monolithic kernels that read it directly. + if e_score_correction_bias is not None: + self.e_score_correction_bias = e_score_correction_bias[:num_real_experts] # Expose zero_expert_num=0 and zero_expert_type=None for # compatibility with quantization methods that check these attributes. From 6cd80fc438d5293b5225fda10736191c7e83a1ae Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 15:34:13 -0500 Subject: [PATCH 033/108] move zero expert handling into MoERunnerBase Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 45 +++++-------------- .../fused_moe/runner/moe_runner_base.py | 20 ++++++++- .../layers/fused_moe/zero_expert_fused_moe.py | 30 ------------- 3 files changed, 31 insertions(+), 64 deletions(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 7713e414ea3a..60675e50d56f 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -13,7 +13,6 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.forward_context import get_forward_context, set_forward_context -from vllm.model_executor.layers.fused_moe.layer import FusedMoE from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( ZeroExpertRouter, ) @@ -119,18 +118,9 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): @pytest.mark.parametrize("num_tokens", [1, 32]) -def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): - """Validate that ZeroExpertFusedMoE output equals real expert output - plus zero expert contribution. - - The key invariant is: - layer.forward(h, r) == FusedMoE.forward(h, r) + zero_expert_output - - FusedMoE.forward() computes only the real expert MoE output (the - ZeroExpertRouter masks zero expert entries to weight=0), while the - zero expert contribution is computed as a side effect during routing - and added on top by ZeroExpertFusedMoE.forward(). - """ +def test_zero_expert_moe_output_deterministic(zero_expert_moe, num_tokens): + """Validate that two forward calls with the same input produce + identical output (determinism).""" layer, vllm_config = zero_expert_moe num_experts = 4 zero_expert_num = 1 @@ -151,28 +141,17 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): get_forward_context().all_moe_layers = None - # Get the real expert output only (bypasses ZeroExpertFusedMoE.forward, - # calls FusedMoE.forward directly). The ZeroExpertRouter still runs and - # stores zero_expert_output as a side effect. - real_output = FusedMoE.forward(layer, hidden_states, router_logits) - zero_output = layer.router.zero_expert_output - - # Get the full combined output. - full_output = layer.forward(hidden_states, router_logits) + output1 = layer.forward(hidden_states, router_logits) + output2 = layer.forward(hidden_states, router_logits) - assert zero_output is not None, "Zero expert output should not be None" - assert not torch.isnan(real_output).any(), "Real expert output has NaN" - assert not torch.isnan(zero_output).any(), "Zero expert output has NaN" - assert not torch.isnan(full_output).any(), "Full output has NaN" + assert not torch.isnan(output1).any(), "Output contains NaN" - expected = real_output + zero_output torch.testing.assert_close( - full_output, - expected, + output1, + output2, atol=0, rtol=0, - msg="ZeroExpertFusedMoE output should equal real expert output " - "plus zero expert contribution", + msg="Two forward calls with same input should produce identical output", ) @@ -232,9 +211,9 @@ def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens): hidden_states.dtype ) - # Run the layer forward to trigger routing and get the actual - # zero expert output from the router. - FusedMoE.forward(layer, hidden_states, router_logits) + # Run routing directly to trigger zero expert computation + # without going through the runner (which consumes the output). + layer.router.select_experts(hidden_states, router_logits) actual_zero_output = layer.router.zero_expert_output assert actual_zero_output is not None diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 2a6e5713827b..0c3aef0c4a7b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -26,6 +26,9 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -412,6 +415,19 @@ def _maybe_gate( router_logits, _ = self.gate(hidden_states) return router_logits + def _maybe_add_zero_expert_output( + self, + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if isinstance(self.router, ZeroExpertRouter): + zero_expert_output = self.router.zero_expert_output + assert zero_expert_output is not None + if isinstance(result, tuple): + result = (result[0], result[1] + zero_expert_output) + else: + result = result + zero_expert_output + return result + def forward( self, hidden_states: torch.Tensor, @@ -463,7 +479,9 @@ def forward( self._encode_layer_name(), ) - return self._maybe_reduce_output(fused_output, og_hidden_dims) + result = self._maybe_reduce_output(fused_output, og_hidden_dims) + + return self._maybe_add_zero_expert_output(result) def forward_dispatch( self, diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py index 4ebfa0d3637a..4a6f2a5500c4 100644 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py @@ -61,33 +61,3 @@ def __init__( # The actual zero expert handling is done by ZeroExpertRouter. self.zero_expert_num = 0 self.zero_expert_type = None - - def forward( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - """ - Forward pass with zero expert support. - - The ZeroExpertRouter handles routing with full logits (including zero - experts), computes zero expert contributions internally, and returns - masked topk_ids suitable for real expert MoE computation. - - Args: - hidden_states: Input hidden states - router_logits: Full router logits (including zero experts) - - Returns: - Combined output from real experts and zero experts - """ - # The router handles full logits internally: routes over all experts - # (real + zero), computes zero expert output, masks zero expert IDs. - fused_out = super().forward(hidden_states, router_logits) - - # Retrieve zero expert output computed during routing - zero_expert_output = self.router.zero_expert_output - if zero_expert_output is not None: - fused_out = fused_out + zero_expert_output - - return fused_out From 992672b48982a6b73f36d37a30fb9b405a5e82c0 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 15:43:55 -0500 Subject: [PATCH 034/108] slightly improved test Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 47 ++++++++++++++++++----- 1 file changed, 38 insertions(+), 9 deletions(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 60675e50d56f..1c34cfe32a28 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -118,9 +118,17 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): @pytest.mark.parametrize("num_tokens", [1, 32]) -def test_zero_expert_moe_output_deterministic(zero_expert_moe, num_tokens): - """Validate that two forward calls with the same input produce - identical output (determinism).""" +def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): + """Validate that the layer output equals real expert output plus zero + expert contribution. + + The key invariant is: + layer.forward(h, r) == quant_method.apply(routing) + zero_expert_output + + We compute routing and zero expert output via router.select_experts(), + then compute real expert output via quant_method.apply() directly, and + verify that the layer forward produces their sum. + """ layer, vllm_config = zero_expert_moe num_experts = 4 zero_expert_num = 1 @@ -141,17 +149,38 @@ def test_zero_expert_moe_output_deterministic(zero_expert_moe, num_tokens): with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): get_forward_context().all_moe_layers = None - output1 = layer.forward(hidden_states, router_logits) - output2 = layer.forward(hidden_states, router_logits) + # Compute routing and zero expert output directly (without the + # runner consuming zero_expert_output). + topk_weights, topk_ids = layer.router.select_experts( + hidden_states, router_logits + ) + zero_output = layer.router.zero_expert_output + + # Compute real expert output via quant_method.apply(). + real_output = layer.quant_method.apply( + layer=layer, + x=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts_input=None, + ) + + # Get the combined output from layer.forward(). + full_output = layer.forward(hidden_states, router_logits) - assert not torch.isnan(output1).any(), "Output contains NaN" + assert zero_output is not None, "Zero expert output should not be None" + assert not torch.isnan(real_output).any(), "Real expert output has NaN" + assert not torch.isnan(zero_output).any(), "Zero expert output has NaN" + assert not torch.isnan(full_output).any(), "Full output has NaN" + expected = real_output + zero_output torch.testing.assert_close( - output1, - output2, + full_output, + expected, atol=0, rtol=0, - msg="Two forward calls with same input should produce identical output", + msg="Layer output should equal real expert output " + "plus zero expert contribution", ) From 900fc407804408cb94fb7ea88dbd41d8612a647a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 16:00:34 -0500 Subject: [PATCH 035/108] simplifications Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 1 - vllm/model_executor/layers/fused_moe/layer.py | 47 ++++++++----------- .../layers/fused_moe/zero_expert_fused_moe.py | 7 --- vllm/model_executor/models/longcat_flash.py | 2 - 4 files changed, 20 insertions(+), 37 deletions(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 1c34cfe32a28..811bd35a869c 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -46,7 +46,6 @@ def zero_expert_moe(dist_init, default_vllm_config): init_workspace_manager(torch.cuda.current_device()) layer = ZeroExpertFusedMoE( - zero_expert_num=zero_expert_num, zero_expert_type="identity", e_score_correction_bias=e_score_correction_bias, num_experts=num_experts, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 532386b68d99..cb729656fc6a 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -36,9 +36,6 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( init_aiter_topK_meta_data, ) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( - FusedMoERouter, -) from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) @@ -342,7 +339,6 @@ def __init__( gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, - router: FusedMoERouter | None = None, zero_expert_type: str | None = None, ): super().__init__() @@ -516,29 +512,26 @@ def __init__( # TODO(bnell): we should not have to create a router if the kernel is # monolithic. - if router is not None: - self.router = router - else: - self.router = create_fused_moe_router( - top_k=top_k, - global_num_experts=self.global_num_experts, - eplb_state=self.eplb_state, - renormalize=renormalize, - use_grouped_topk=use_grouped_topk, - num_expert_group=num_expert_group, - topk_group=topk_group, - custom_routing_function=custom_routing_function, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - num_fused_shared_experts=self.num_fused_shared_experts, - enable_eplb=enable_eplb, - # TODO(bnell): once we can construct the MK at init time, we - # can make this a value. - indices_type_getter=lambda: self.quant_method.topk_indices_dtype, - zero_expert_type=zero_expert_type, - num_logical_experts=self.logical_num_experts, - ) + self.router = create_fused_moe_router( + top_k=top_k, + global_num_experts=self.global_num_experts, + eplb_state=self.eplb_state, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=self.num_fused_shared_experts, + enable_eplb=enable_eplb, + # TODO(bnell): once we can construct the MK at init time, we + # can make this a value. + indices_type_getter=lambda: self.quant_method.topk_indices_dtype, + zero_expert_type=zero_expert_type, + num_logical_experts=self.logical_num_experts, + ) self.routing_method_type: RoutingMethodType = self.router.routing_method_type # Round up hidden size before creating moe_config. diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py index 4a6f2a5500c4..59781f75ccff 100644 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py @@ -20,7 +20,6 @@ class ZeroExpertFusedMoE(FusedMoE): def __init__( self, - zero_expert_num: int, zero_expert_type: str, e_score_correction_bias: torch.Tensor | None = None, **kwargs, @@ -55,9 +54,3 @@ def __init__( # for compatibility with monolithic kernels that read it directly. if e_score_correction_bias is not None: self.e_score_correction_bias = e_score_correction_bias[:num_real_experts] - - # Expose zero_expert_num=0 and zero_expert_type=None for - # compatibility with quantization methods that check these attributes. - # The actual zero expert handling is done by ZeroExpertRouter. - self.zero_expert_num = 0 - self.zero_expert_type = None diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index b695e062d276..ddd46a2c4382 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -292,10 +292,8 @@ def __init__( prefix=f"{prefix}.gate", ) - assert config.zero_expert_num is not None assert config.zero_expert_type is not None self.experts = ZeroExpertFusedMoE( - zero_expert_num=config.zero_expert_num, zero_expert_type=config.zero_expert_type, e_score_correction_bias=self.router.e_score_correction_bias, num_experts=num_experts, From 71f8fb871bc0b7d72eab56d7071e8e938bfeb60b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 16:21:48 -0500 Subject: [PATCH 036/108] better test Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 49 +++++++++++++++++------ 1 file changed, 36 insertions(+), 13 deletions(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 811bd35a869c..066506ec8420 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -13,6 +13,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe.layer import FusedMoE from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( ZeroExpertRouter, ) @@ -118,15 +119,16 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): - """Validate that the layer output equals real expert output plus zero - expert contribution. + """Validate that the ZeroExpertFusedMoE output equals a plain FusedMoE + output (real experts only) plus the zero expert contribution. The key invariant is: - layer.forward(h, r) == quant_method.apply(routing) + zero_expert_output + zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real) + + zero_expert_output - We compute routing and zero expert output via router.select_experts(), - then compute real expert output via quant_method.apply() directly, and - verify that the layer forward produces their sum. + We create a plain FusedMoE layer with the same weights and real-expert-only + router logits, compute the zero expert output via the ZeroExpertRouter, and + verify the sum matches the ZeroExpertFusedMoE output. """ layer, vllm_config = zero_expert_moe num_experts = 4 @@ -148,23 +150,44 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): get_forward_context().all_moe_layers = None - # Compute routing and zero expert output directly (without the - # runner consuming zero_expert_output). + # Create a plain FusedMoE layer with the same config but no zero + # experts. Use a separate prefix to avoid collision. + plain_layer = FusedMoE( + num_experts=num_experts, + top_k=layer.top_k, + hidden_size=layer.hidden_size, + intermediate_size=layer.intermediate_size_per_partition, + params_dtype=torch.bfloat16, + prefix="test_zero_expert_moe_plain", + renormalize=False, + scoring_func="softmax", + e_score_correction_bias=layer.e_score_correction_bias, + ).cuda() + + # Share weights from the zero expert layer. + plain_layer.w13_weight.data.copy_(layer.w13_weight.data) + plain_layer.w2_weight.data.copy_(layer.w2_weight.data) + plain_layer.quant_method.process_weights_after_loading(plain_layer) + + # Compute routing via the ZeroExpertRouter. This produces masked + # topk_weights/topk_ids (zero expert entries have weight=0, id=0) + # and stores zero_expert_output as a side effect. topk_weights, topk_ids = layer.router.select_experts( hidden_states, router_logits ) zero_output = layer.router.zero_expert_output - # Compute real expert output via quant_method.apply(). - real_output = layer.quant_method.apply( - layer=layer, + # Compute real expert output using the plain layer with the masked + # routing from the ZeroExpertRouter. + real_output = plain_layer.quant_method.apply( + layer=plain_layer, x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, shared_experts_input=None, ) - # Get the combined output from layer.forward(). + # Get the combined output from the zero expert layer. full_output = layer.forward(hidden_states, router_logits) assert zero_output is not None, "Zero expert output should not be None" @@ -178,7 +201,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): expected, atol=0, rtol=0, - msg="Layer output should equal real expert output " + msg="ZeroExpertFusedMoE output should equal plain FusedMoE output " "plus zero expert contribution", ) From c0e12b5047256a11e1af1163ac0a073009212f7b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 16:33:41 -0500 Subject: [PATCH 037/108] remove ZeroExpertFusedMoE Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 21 +++---- .../layers/fused_moe/__init__.py | 4 -- vllm/model_executor/layers/fused_moe/layer.py | 8 +++ .../layers/fused_moe/zero_expert_fused_moe.py | 56 ------------------- vllm/model_executor/models/longcat_flash.py | 6 +- 5 files changed, 20 insertions(+), 75 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 066506ec8420..92507eabf22f 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for ZeroExpertFusedMoE. +"""Tests for FusedMoE with zero experts. Verifies that: - The ZeroExpertRouter is properly created and used as the layer router. -- A forward pass through ZeroExpertFusedMoE produces correct output. +- A forward pass through FusedMoE with zero experts produces correct output. - The output decomposes correctly into real expert + zero expert contributions. """ @@ -17,15 +17,12 @@ from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( ZeroExpertRouter, ) -from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import ( - ZeroExpertFusedMoE, -) from vllm.v1.worker.workspace import init_workspace_manager @pytest.fixture def zero_expert_moe(dist_init, default_vllm_config): - """Create a ZeroExpertFusedMoE layer with zero experts.""" + """Create a FusedMoE layer with zero experts.""" num_experts = 4 top_k = 2 # hidden_size must be >= 256 for the zero expert identity kernel to @@ -46,7 +43,7 @@ def zero_expert_moe(dist_init, default_vllm_config): with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): init_workspace_manager(torch.cuda.current_device()) - layer = ZeroExpertFusedMoE( + layer = FusedMoE( zero_expert_type="identity", e_score_correction_bias=e_score_correction_bias, num_experts=num_experts, @@ -67,7 +64,7 @@ def zero_expert_moe(dist_init, default_vllm_config): @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens): - """Verify that ZeroExpertFusedMoE creates a ZeroExpertRouter.""" + """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter.""" layer, _ = zero_expert_moe assert isinstance(layer.router, ZeroExpertRouter), ( f"Expected ZeroExpertRouter but got {type(layer.router).__name__}." @@ -84,7 +81,7 @@ def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens): @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): - """Run a forward pass through ZeroExpertFusedMoE and verify output shape.""" + """Run a forward pass through FusedMoE with zero experts and verify output shape.""" layer, vllm_config = zero_expert_moe hidden_size = layer.hidden_size @@ -119,7 +116,7 @@ def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): @pytest.mark.parametrize("num_tokens", [1, 32]) def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): - """Validate that the ZeroExpertFusedMoE output equals a plain FusedMoE + """Validate that the FusedMoE output equals a plain FusedMoE output (real experts only) plus the zero expert contribution. The key invariant is: @@ -128,7 +125,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): We create a plain FusedMoE layer with the same weights and real-expert-only router logits, compute the zero expert output via the ZeroExpertRouter, and - verify the sum matches the ZeroExpertFusedMoE output. + verify the sum matches the FusedMoE output. """ layer, vllm_config = zero_expert_moe num_experts = 4 @@ -201,7 +198,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): expected, atol=0, rtol=0, - msg="ZeroExpertFusedMoE output should equal plain FusedMoE output " + msg="FusedMoE output should equal plain FusedMoE output " "plus zero expert contribution", ) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index f56a2e63bf40..a28fd06930c6 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -33,9 +33,6 @@ from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) -from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import ( - ZeroExpertFusedMoE, -) from vllm.triton_utils import HAS_TRITON _config: dict[str, Any] | None = None @@ -68,7 +65,6 @@ def get_config() -> dict[str, Any] | None: "GateLinear", "RoutingMethodType", "SharedFusedMoE", - "ZeroExpertFusedMoE", "activation_without_mul", "apply_moe_activation", "override_config", diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index cb729656fc6a..172b62e822ed 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -534,6 +534,14 @@ def __init__( ) self.routing_method_type: RoutingMethodType = self.router.routing_method_type + # When using zero experts, slice e_score_correction_bias to cover + # only real experts, for compatibility with monolithic kernels that + # read it directly. + if zero_expert_type is not None and e_score_correction_bias is not None: + self.e_score_correction_bias = e_score_correction_bias[ + : self.logical_num_experts + ] + # Round up hidden size before creating moe_config. # This way moe_config is created with the correct hidden_size from the start. unpadded_hidden_size = hidden_size diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py deleted file mode 100644 index 59781f75ccff..000000000000 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ /dev/null @@ -1,56 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -from vllm.model_executor.layers.fused_moe.layer import FusedMoE - - -class ZeroExpertFusedMoE(FusedMoE): - """ - A FusedMoE operation that also computes the results of zero experts. - Zero experts perform identity operations (scaled pass-through) instead - of full MLP computations. - - Uses a ZeroExpertRouter as the layer's main router. The router handles - routing over all experts (real + zero) with the full e_score_correction_bias, - computes zero expert contributions as a side effect, and remaps zero expert - IDs so downstream MoE computation only processes real experts. - """ - - def __init__( - self, - zero_expert_type: str, - e_score_correction_bias: torch.Tensor | None = None, - **kwargs, - ): - assert ( - "custom_routing_function" not in kwargs - or kwargs.get("custom_routing_function") is None - ), ( - "ZeroExpertFusedMoE does not support external custom_routing_function. " - "Routing is handled by ZeroExpertRouter." - ) - - assert "router" not in kwargs or kwargs.get("router") is None, ( - "ZeroExpertFusedMoE creates its own ZeroExpertRouter. Do not pass a router." - ) - - # Remove custom_routing_function from kwargs if present - kwargs.pop("custom_routing_function", None) - kwargs.pop("router", None) - - # Pass the full e_score_correction_bias (real + zero experts) and - # zero_expert_type through to the router factory, which will create - # a ZeroExpertRouter. - num_real_experts = kwargs["num_experts"] - if e_score_correction_bias is not None: - kwargs["e_score_correction_bias"] = e_score_correction_bias - kwargs["zero_expert_type"] = zero_expert_type - - super().__init__(**kwargs) - - # Fix self.e_score_correction_bias to only cover real experts, - # for compatibility with monolithic kernels that read it directly. - if e_score_correction_bias is not None: - self.e_score_correction_bias = e_score_correction_bias[:num_real_experts] diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index ddd46a2c4382..375b0b69b1f9 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -46,7 +46,7 @@ from vllm.distributed import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoE, ZeroExpertFusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -293,7 +293,7 @@ def __init__( ) assert config.zero_expert_type is not None - self.experts = ZeroExpertFusedMoE( + self.experts = FusedMoE( zero_expert_type=config.zero_expert_type, e_score_correction_bias=self.router.e_score_correction_bias, num_experts=num_experts, @@ -330,7 +330,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states_padded.to(self.router_params_dtype) ) - # ZeroExpertFusedMoE handles routing memoization and zero expert computation + # FusedMoE handles routing memoization and zero expert computation # internally. Pass full router_logits (including zero experts) so that # zero experts can be properly identified in routing. final_hidden_states = self.experts( From 47205a3c2a813e9853aaa77a745ef31bf0574bf2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 2 Mar 2026 08:26:28 -0500 Subject: [PATCH 038/108] Add comment Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 92507eabf22f..24edca00b6af 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -6,6 +6,8 @@ - The ZeroExpertRouter is properly created and used as the layer router. - A forward pass through FusedMoE with zero experts produces correct output. - The output decomposes correctly into real expert + zero expert contributions. + +Note: tests generated with Claude. """ import pytest From a7af12c7e886d0cf77e508544216004facd7e235 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 3 Mar 2026 14:12:29 -0500 Subject: [PATCH 039/108] move shared expert all gather to SharedExperts Signed-off-by: Bill Nell --- .../fused_moe/runner/moe_runner_base.py | 5 ++-- .../layers/fused_moe/runner/shared_experts.py | 30 +++++-------------- 2 files changed, 11 insertions(+), 24 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 0c3aef0c4a7b..1d31d2eb8405 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -297,8 +297,9 @@ def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: func = trunc if isinstance(states, tuple): - return tuple( - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] + return ( + trunc(states[0], trunc_sizes[0]), # shared: already reduced + func(states[1], trunc_sizes[1]), # fused: reduce here ) else: assert len(trunc_sizes) == 1 diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 05258999ec8c..a70c54d36395 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -104,14 +104,6 @@ def _has_mk_owned_shared_experts(self) -> bool: and self._shared_experts is not None ) - @property - def _must_reduce_shared_expert_outputs(self) -> bool: - return ( - self._reduce_results - and self._quant_method.moe_kernel is not None - and self._quant_method.moe_kernel.output_is_reduced() - ) - def _determine_shared_experts_order( self, hidden_states: torch.Tensor, @@ -180,21 +172,23 @@ def _call_with_shared_experts_stream( return output - def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. + def _maybe_reduce_shared_output(self, output: torch.Tensor) -> torch.Tensor: if ( - self._must_reduce_shared_expert_outputs + self._reduce_results + and not self._moe_config.is_sequence_parallel + and not self._use_dp_chunking and get_tensor_model_parallel_world_size() > 1 ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out + output = tensor_model_parallel_all_reduce(output) + return output @property def output(self) -> torch.Tensor | None: assert (self._shared_experts is None) == (self._output is None) output = self._output self._output = None + if output is not None: + output = self._maybe_reduce_shared_output(output) return output def apply( @@ -216,11 +210,3 @@ def apply( self._output = self._call_with_shared_experts_stream(shared_experts_input) else: self._output = self._shared_experts(shared_experts_input) - - if order == SharedExpertsOrder.EXTERNAL: - # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely. - assert self._output is not None - self._output = self._maybe_reduce_shared_out(self._output) - - # TODO(bnell): potentially do AFTER reduce here instead of in runner. From 51da43bfbd9b7a629aa4d2540cc2823eb565dbd8 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 3 Mar 2026 14:22:50 -0500 Subject: [PATCH 040/108] remove must_reduce_shared_expert_outputs external method Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 15 ------- .../layers/fused_moe/runner/moe_runner.py | 4 -- .../fused_moe/runner/moe_runner_base.py | 4 +- vllm/model_executor/models/exaone_moe.py | 40 +++++++++---------- 4 files changed, 22 insertions(+), 41 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 172b62e822ed..05920bfd57cc 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1493,21 +1493,6 @@ def moe_quant_config(self) -> FusedMoEQuantConfig | None: self.ensure_moe_quant_config_init() return self.quant_method.moe_quant_config - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return self.runner.must_reduce_shared_expert_outputs() - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): """ Some combine kernels reduce across GPU ranks by default. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index b298cc2d0c4c..067b3d00c680 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -22,10 +22,6 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError - @abstractmethod - def must_reduce_shared_expert_outputs(self) -> bool: - raise NotImplementedError - @abstractmethod def maybe_all_reduce_tensor_model_parallel( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 1d31d2eb8405..f0a3c0aee415 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -220,7 +220,7 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: def reduce_results(self) -> bool: raise NotImplementedError - def must_reduce_shared_expert_outputs(self) -> bool: + def _must_reduce_shared_expert_outputs(self) -> bool: """ The shared_experts are typically computed using the RowParallelLinear layer. The result of this function is typically used as @@ -242,7 +242,7 @@ def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tens """ Some combine kernels reduce across GPU ranks by default. """ - if self.must_reduce_shared_expert_outputs(): + if self._must_reduce_shared_expert_outputs(): return final_hidden_states else: return tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index d7282edcf4f6..66d05e095ac5 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -31,6 +31,7 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -116,7 +117,22 @@ def __init__( self.physical_expert_start + self.n_local_physical_experts ) - self.experts = FusedMoE( + if getattr(config, "num_shared_experts", 0) > 0: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = ExaoneMoeGatedMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, @@ -135,34 +151,18 @@ def __init__( num_redundant_experts=self.n_redundant_experts, ) - if getattr(config, "num_shared_experts", 0) > 0: - intermediate_size = config.moe_intermediate_size * config.num_shared_experts - self.shared_experts = ExaoneMoeGatedMLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=self.experts.must_reduce_shared_expert_outputs(), - prefix=f"{prefix}.shared_experts", - ) - else: - self.shared_experts = None - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # NOTE: hidden_states can have either 1D or 2D shape. orig_shape = hidden_states.shape hidden_dim = hidden_states.shape[-1] hidden_states = hidden_states.view(-1, hidden_dim) - # router_logits: (num_tokens, n_experts) - router_logits, _ = self.gate(hidden_states) - - final_hidden_states = self.experts( - hidden_states=hidden_states, router_logits=router_logits + shared_output, final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=hidden_states ) if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) + assert shared_output is not None final_hidden_states = final_hidden_states + shared_output if self.tp_size > 1: From 783402152088027bbdade7fb956eaae52699415a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 3 Mar 2026 18:49:56 -0500 Subject: [PATCH 041/108] wip moving experts epilog into MoERunnerBase Signed-off-by: Bill Nell --- .../test_shared_fused_moe_routed_transform.py | 29 ++---- vllm/forward_context.py | 2 +- vllm/lora/layers/fused_moe.py | 3 - vllm/model_executor/layers/fused_moe/layer.py | 22 ++--- .../fused_moe/runner/chunking_moe_runner.py | 7 +- .../fused_moe/runner/default_moe_runner.py | 10 +-- .../layers/fused_moe/runner/moe_runner.py | 9 +- .../fused_moe/runner/moe_runner_base.py | 89 +++++++++---------- .../fused_moe/runner/moe_runner_factory.py | 6 +- .../layers/fused_moe/runner/shared_experts.py | 18 ---- vllm/model_executor/models/afmoe.py | 13 +-- vllm/model_executor/models/aria.py | 7 +- vllm/model_executor/models/bailing_moe.py | 17 +--- vllm/model_executor/models/deepseek_v2.py | 34 +------ vllm/model_executor/models/dots1.py | 13 +-- vllm/model_executor/models/ernie45_moe.py | 11 --- vllm/model_executor/models/ernie45_vl_moe.py | 35 ++------ vllm/model_executor/models/exaone_moe.py | 12 +-- vllm/model_executor/models/glm4_moe.py | 20 +---- vllm/model_executor/models/hunyuan_v1.py | 7 -- vllm/model_executor/models/kimi_linear.py | 46 +++++----- vllm/model_executor/models/lfm2_moe.py | 12 +-- vllm/model_executor/models/llama4.py | 8 +- vllm/model_executor/models/minimax_m2.py | 5 -- vllm/model_executor/models/nemotron_h.py | 30 ++----- vllm/model_executor/models/openpangu.py | 24 +---- vllm/model_executor/models/qwen2_moe.py | 7 -- vllm/model_executor/models/qwen3_moe.py | 10 +-- vllm/model_executor/models/qwen3_next.py | 8 -- vllm/model_executor/models/step3_text.py | 4 - vllm/model_executor/models/step3p5.py | 21 +---- .../model_executor/models/transformers/moe.py | 41 +-------- 32 files changed, 126 insertions(+), 454 deletions(-) diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index b6ef19ddaf3c..ccba92b81782 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -98,7 +98,6 @@ def test_routed_input_transform_inside_vs_outside( top_k=top_k, hidden_size=latent_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=True, params_dtype=dtype, tp_size=1, @@ -116,7 +115,6 @@ def test_routed_input_transform_inside_vs_outside( top_k=top_k, hidden_size=latent_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=True, params_dtype=dtype, tp_size=1, @@ -140,28 +138,19 @@ def test_routed_input_transform_inside_vs_outside( router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype) with set_forward_context(None, vllm_config, num_tokens=num_tokens): - shared_out_A, routed_out_A = moe_with_transform( - hidden_states, router_logits - ) + # Method A: combined output (shared + routed) + combined_A = moe_with_transform(hidden_states, router_logits) + # Method B: manually transform, get routed output, add shared transformed_hidden = routed_transform(hidden_states) - shared_out_B, routed_out_B = moe_without_transform( - transformed_hidden, router_logits - ) + routed_out_B = moe_without_transform(transformed_hidden, router_logits) + shared_out_B = shared_experts(hidden_states) + combined_B = shared_out_B + routed_out_B torch.testing.assert_close( - routed_out_A, - routed_out_B, - atol=1e-3, - rtol=1e-3, - msg="Routed output should match: transform inside vs outside", - ) - - expected_shared_out = shared_experts(hidden_states) - - torch.testing.assert_close( - shared_out_A, - expected_shared_out, + combined_A, + combined_B, atol=1e-3, rtol=1e-3, + msg="Combined output should match: transform inside vs outside", ) diff --git a/vllm/forward_context.py b/vllm/forward_context.py index bf0f9da6eaff..1231263bcfe3 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -213,7 +213,7 @@ class ForwardContext: # For torch.compile cold start times, we need to avoid hard-coding # any strings into the graph. Right now, the vllm.moe_forward - # and vllm.moe_forward_shared custom operators hard-code strings into + # custom operator hard-codes strings into # the graph. # # The workaround is to store a list of the strings that each of those diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 11332f4cb228..fd9023b01fef 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -593,9 +593,6 @@ def set_lora( def forward(self, *args, **kwargs): return self.base_layer.forward(*args, **kwargs) - def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): - return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - @property def quant_method(self): return self.base_layer.quant_method diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 05920bfd57cc..c82e3d16a458 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -295,7 +295,6 @@ class FusedMoE(CustomOp): hidden_size: Input hidden state size of the transformer intermediate_size: Intermediate size of the experts params_dtype: Data type for the parameters. - reduce_results: Whether to all_reduce on the output of the layer renormalize: Whether to renormalize the logits in the fused_moe kernel quant_config: Quantization configure. enable_eplb: Whether to enable expert parallelism load balancer. @@ -311,7 +310,6 @@ def __init__( hidden_size: int, intermediate_size: int, params_dtype: torch.dtype | None = None, - reduce_results: bool = False, renormalize: bool = True, use_grouped_topk: bool = False, num_expert_group: int | None = None, @@ -339,12 +337,16 @@ def __init__( gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + output_scale: float | None = None, zero_expert_type: str | None = None, ): super().__init__() self._gate = gate self._routed_input_transform = routed_input_transform + self._routed_output_transform = routed_output_transform + self._output_scale = output_scale if params_dtype is None: params_dtype = torch.get_default_dtype() @@ -491,7 +493,6 @@ def __init__( assert intermediate_size % self.tp_size == 0 self.intermediate_size_per_partition = intermediate_size // self.tp_size - self.reduce_results = reduce_results self.renormalize = renormalize # TODO(bnell): these attributes are only used by monolithic kernels. @@ -664,7 +665,6 @@ def _init_shared_experts(self) -> SharedExperts | None: # called, i.e. by a MK or by the MoERunner. # Once the MK can be created upfront, we can just pass in the proper # flags derived from the quant_method's MK. - reduce_results=self.reduce_results, quant_method=self.quant_method, ) @@ -680,8 +680,9 @@ def _init_runner(self) -> MoERunner: gate=self._gate, shared_experts=self.shared_experts, quant_method=self.quant_method, - reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, + routed_output_transform=self._routed_output_transform, + output_scale=self._output_scale, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -1493,17 +1494,11 @@ def moe_quant_config(self) -> FusedMoEQuantConfig | None: self.ensure_moe_quant_config_init() return self.quant_method.moe_quant_config - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states) - def forward_native( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.runner.forward( hidden_states, router_logits, @@ -1519,7 +1514,7 @@ def forward_cuda( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.forward_native(hidden_states, router_logits) @classmethod @@ -1576,7 +1571,6 @@ def extra_repr(self) -> str: f"intermediate_size_per_partition={self.intermediate_size_per_partition}, " # noqa: E501 f"tp_size={self.tp_size},\n" f"ep_size={self.ep_size}, " - f"reduce_results={self.reduce_results}, " ) return s diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py index 6626fad557b3..8aee95840cb5 100644 --- a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -27,12 +27,11 @@ class ChunkingMoERunner(MoERunnerBase): All MoERunnerBase state (moe_config, router, quant_method, etc.) is transparently delegated to the inner runner via __getattr__. ChunkingMoERunner only owns chunking-specific state: the pre-allocated - workspace buffers and the reduce_results override. + workspace buffers. Key behaviors: - Pre-allocates workspace tensors for CUDA graph compatibility - Processes chunks via inner._forward_impl per chunk - - Never reduces results (reduce_results always returns False) """ def __init__(self, inner: MoERunnerBase): @@ -55,10 +54,6 @@ def __getattr__(self, name): # so ChunkingMoERunner's own attributes and methods take priority. return getattr(self._inner, name) - @property - def reduce_results(self) -> bool: - return False - def _init_dp_chunking(self) -> list[torch.Tensor]: states_shape: tuple[int, ...] logits_shape: tuple[int, ...] diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 7b583bda14cc..523a5362da51 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -61,8 +61,9 @@ def __init__( gate: torch.nn.Module | None, shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + output_scale: float | None = None, ): super().__init__( layer, @@ -72,14 +73,11 @@ def __init__( gate, shared_experts, quant_method, - reduce_results, enable_dbo, + routed_output_transform=routed_output_transform, + output_scale=output_scale, ) - @property - def reduce_results(self) -> bool: - return self._reduce_results - @property def do_naive_dispatch_combine(self) -> bool: return ( diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 067b3d00c680..43d909ee1458 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -19,12 +19,5 @@ def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - @abstractmethod - def maybe_all_reduce_tensor_model_parallel( - self, - final_hidden_states: torch.Tensor, - ): + ) -> torch.Tensor: raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index f0a3c0aee415..60bdc2280d5b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -54,7 +54,7 @@ def get_layer_from_name(layer_name: str) -> torch.nn.Module: raise AssertionError( "We expected the number of MOE layers in `all_moe_layers` " "to be equal to the number of " - "{vllm.moe_forward, vllm.moe_forward_shared} calls." + "vllm.moe_forward calls." ) layer_name = all_moe_layers[moe_layer_index] forward_context.moe_layer_index += 1 @@ -99,6 +99,10 @@ def _moe_forward_fake( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: + # For latent MoE with reduce_results=True: output has full hidden_size + # (from shared_experts_input), not latent_size (from hidden_states). + if shared_experts_input is not None: + return torch.empty_like(shared_experts_input) return torch.empty_like(hidden_states) @@ -170,7 +174,6 @@ class MoERunnerBase(MoERunner): allowing flexibility in the actual MoE computation implementation. Key abstract methods that subclasses must implement: - - reduce_results: Determines whether results should be reduced across ranks - _forward_impl: The core MoE computation logic specific to each runner type """ @@ -183,8 +186,9 @@ def __init__( gate: torch.nn.Module | None, shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + output_scale: float | None = None, ): super().__init__() self.moe_config = moe_config @@ -193,9 +197,10 @@ def __init__( self.gate = gate self.shared_experts = shared_experts self.quant_method = quant_method - self._reduce_results = reduce_results self.enable_dbo = enable_dbo self.enable_eplb = moe_config.moe_parallel_config.enable_eplb + self.routed_output_transform = routed_output_transform + self.output_scale = output_scale # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name @@ -215,11 +220,6 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) - @property - @abstractmethod - def reduce_results(self) -> bool: - raise NotImplementedError - def _must_reduce_shared_expert_outputs(self) -> bool: """ The shared_experts are typically computed using the RowParallelLinear @@ -257,7 +257,7 @@ def apply_routed_input_transform( routed experts get the transformed [S, moe_latent_size]. TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside SharedFusedMoE to all-reduce on the smaller latent + moved inside FusedMoE to all-reduce on the smaller latent dimension. Returns (possibly transformed) hidden states and the input for shared @@ -276,34 +276,23 @@ def apply_routed_input_transform( hidden_states if self.shared_experts is not None else None, ) - def _maybe_reduce_output( + def _combine_and_reduce( self, - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + states: torch.Tensor, trunc_sizes: list[int], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return x[..., :trunc_size] + ) -> torch.Tensor: + assert len(trunc_sizes) == 1 + result = states[..., : trunc_sizes[0]] - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) + if self.output_scale is not None: + result = result * self.output_scale - if ( - not self.moe_config.is_sequence_parallel - and self.reduce_results - and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + if not self.moe_config.is_sequence_parallel and ( + self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1 ): - func = reduce_and_trunc - else: - func = trunc + result = self.maybe_all_reduce_tensor_model_parallel(result) - if isinstance(states, tuple): - return ( - trunc(states[0], trunc_sizes[0]), # shared: already reduced - func(states[1], trunc_sizes[1]), # fused: reduce here - ) - else: - assert len(trunc_sizes) == 1 - return func(states, trunc_sizes[0]) + return result def _encode_layer_name(self) -> str | ModuleName: if HAS_OPAQUE_TYPE: @@ -336,8 +325,8 @@ def _maybe_pad_hidden_states( value=0.0, ) - if self.shared_experts is not None: - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] + if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: + orig_hidden_dims = [shared_experts_hidden_dim] else: orig_hidden_dims = [transformed_hidden_dim] @@ -418,22 +407,19 @@ def _maybe_gate( def _maybe_add_zero_expert_output( self, - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + result: torch.Tensor, + ) -> torch.Tensor: if isinstance(self.router, ZeroExpertRouter): zero_expert_output = self.router.zero_expert_output assert zero_expert_output is not None - if isinstance(result, tuple): - result = (result[0], result[1] + zero_expert_output) - else: - result = result + zero_expert_output + result = result + zero_expert_output return result def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """Invoke the fused moe layer. Input: @@ -442,8 +428,6 @@ def forward( Output: - The new hidden_states. - or - - A tuple of (shared experts output, new hidden_states). Calling sequence - forward @@ -480,7 +464,7 @@ def forward( self._encode_layer_name(), ) - result = self._maybe_reduce_output(fused_output, og_hidden_dims) + result = self._combine_and_reduce(fused_output, og_hidden_dims) return self._maybe_add_zero_expert_output(result) @@ -490,7 +474,7 @@ def forward_dispatch( hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -502,13 +486,26 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - return self._forward_impl( + result = self._forward_impl( layer, hidden_states, router_logits, shared_experts_input, ) + if isinstance(result, tuple): + shared_output, fused_output = result + # Apply output transform (e.g. latent → full dim) + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + # If combine kernel already reduced fused, reduce shared to match + if self._must_reduce_shared_expert_outputs(): + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = shared_output + fused_output + + return result + @abstractmethod def _forward_impl( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index 4bda1ca89304..e6100814cad8 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -32,8 +32,9 @@ def create_moe_runner( gate: torch.nn.Module | None, shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, - reduce_results: bool, enable_dbo: bool, + routed_output_transform: torch.nn.Module | None = None, + output_scale: float | None = None, ) -> MoERunner: runner = DefaultMoERunner( layer, @@ -43,8 +44,9 @@ def create_moe_runner( gate, shared_experts, quant_method, - reduce_results, enable_dbo, + routed_output_transform=routed_output_transform, + output_scale=output_scale, ) if moe_config.moe_parallel_config.use_dp_chunking: return ChunkingMoERunner(runner) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index a70c54d36395..02a3afcd78e7 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -5,10 +5,6 @@ import torch import vllm.envs as envs -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -50,7 +46,6 @@ def __init__( shared_experts: torch.nn.Module, moe_config: FusedMoEConfig, quant_method: QuantizeMethodBase, - reduce_results: bool, ): from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, @@ -64,7 +59,6 @@ def __init__( self._shared_experts = shared_experts self._moe_config = moe_config self._quant_method = quant_method - self._reduce_results = reduce_results self._use_dp_chunking = moe_config.moe_parallel_config.use_dp_chunking # Allow disabling of the separate shared experts stream for @@ -172,23 +166,11 @@ def _call_with_shared_experts_stream( return output - def _maybe_reduce_shared_output(self, output: torch.Tensor) -> torch.Tensor: - if ( - self._reduce_results - and not self._moe_config.is_sequence_parallel - and not self._use_dp_chunking - and get_tensor_model_parallel_world_size() > 1 - ): - output = tensor_model_parallel_all_reduce(output) - return output - @property def output(self) -> torch.Tensor | None: assert (self._shared_experts is None) == (self._output is None) output = self._output self._output = None - if output is not None: - output = self._maybe_reduce_shared_output(output) return output def apply( diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 22037336411a..5bad52a0c496 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -131,7 +131,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.route_norm if self.score_func == "sigmoid" else False, quant_config=quant_config, use_grouped_topk=True, @@ -152,20 +151,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: router_logits = self.gate(hidden_states.to(dtype=torch.float32)) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_experts is not None: - shared_output, final_hidden_states = fused_moe_out - final_hidden_states = final_hidden_states + shared_output - else: - final_hidden_states = fused_moe_out - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 908581786450..221e991dfc29 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -301,12 +301,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: router_output = torch.nn.functional.linear(hidden_states, self.router_weight) - sparse_expert_output = self.experts(hidden_states, router_output) - - if self.shared_experts is not None: - return sparse_expert_output[0] + sparse_expert_output[1] - else: - return sparse_expert_output + return self.experts(hidden_states, router_output) class AriaTextDecoderLayer(LlamaDecoderLayer): diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 7725dfa2a887..510d605f8046 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -291,7 +291,6 @@ def __init__( top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -301,6 +300,7 @@ def __init__( topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -314,21 +314,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - - if self.shared_experts is not None: - shared_output, final_hidden_states = final_hidden_states - else: - shared_output = None - - final_hidden_states *= self.routed_scaling_factor - - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_size) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index f31e9ac3e840..1a64299789fb 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -312,7 +312,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -320,11 +319,7 @@ def __init__( topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), - # we do scaling outside, set factor to 1.0 to avoid double mul - # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -357,43 +352,20 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - # Fix FP16 overflow - # See DeepseekV2DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - if not self.is_rocm_aiter_moe_enabled: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index 4e393145462a..fc2636db48ab 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -37,7 +37,6 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -120,7 +119,6 @@ def __init__( prefix: str = "", ): super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() self.routed_scaling_factor = config.routed_scaling_factor self.n_shared_experts = config.n_shared_experts @@ -163,7 +161,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -174,6 +171,7 @@ def __init__( # we do scaling outside, set factor to 1.0 to avoid double mul routed_scaling_factor=1.0, e_score_correction_bias=self.gate.e_score_correction_bias, + output_scale=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -182,16 +180,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: router_logits, _ = self.gate(hidden_states) - shared_out, routed_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_experts is not None: - final_hidden_states = (routed_out + shared_out) * self.routed_scaling_factor - else: - final_hidden_states = routed_out * self.routed_scaling_factor - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index f038cfb21f28..c92e230bcd21 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -194,7 +194,6 @@ def __init__( top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -215,16 +214,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states=hidden_states, router_logits=router_logits ) - if self.has_shared_experts: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - else: - final_hidden_states = final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py index 418fdcfa072b..e4b7ac6fb006 100644 --- a/vllm/model_executor/models/ernie45_vl_moe.py +++ b/vllm/model_executor/models/ernie45_vl_moe.py @@ -263,7 +263,6 @@ def __init__( top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size[0], - reduce_results=False, renormalize=True, quant_config=quant_config, e_score_correction_bias=self.e_score_correction_bias[0], @@ -301,7 +300,6 @@ def __init__( top_k=config.moe_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size[1], - reduce_results=False, renormalize=True, quant_config=quant_config, e_score_correction_bias=self.e_score_correction_bias[1], @@ -342,9 +340,6 @@ def forward( visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool() text_token_mask = ~visual_token_mask final_experts_hidden_states = torch.zeros_like(hidden_states) - final_shared_output = ( - torch.zeros_like(hidden_states) if self.has_shared_experts else None - ) text_hidden_states = hidden_states[text_token_mask].reshape( -1, self.hidden_size @@ -356,26 +351,20 @@ def forward( text_router_logits, _ = self.text_experts_gate( text_hidden_states.to(dtype=torch.float32) ) - text_shared_output, text_experts_output = self.text_experts( + text_output = self.text_experts( hidden_states=text_hidden_states, router_logits=text_router_logits ) - final_experts_hidden_states[text_token_mask] = text_experts_output.flatten() - if self.has_shared_experts: - final_shared_output[text_token_mask] = text_shared_output.flatten() + final_experts_hidden_states[text_token_mask] = text_output.flatten() vision_router_logits, _ = self.vision_experts_gate( vision_hidden_states.to(dtype=torch.float32) ) - vision_shared_output, vision_experts_output = self.vision_experts( + vision_output = self.vision_experts( hidden_states=vision_hidden_states, router_logits=vision_router_logits ) - final_experts_hidden_states[visual_token_mask] = ( - vision_experts_output.flatten() - ) - if self.has_shared_experts: - final_shared_output[visual_token_mask] = vision_shared_output.flatten() + final_experts_hidden_states[visual_token_mask] = vision_output.flatten() - final_hidden_states = (final_shared_output, final_experts_hidden_states) + final_hidden_states = final_experts_hidden_states else: # only text modal input text_router_logits, _ = self.text_experts_gate( @@ -386,20 +375,6 @@ def forward( hidden_states=hidden_states, router_logits=text_router_logits ) - if self.has_shared_experts: - # for shared_experts model - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - else: - # for not shared_experts model - final_hidden_states = final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = ( - self.text_experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index 66d05e095ac5..a46cadf007ee 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -137,7 +137,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -157,19 +156,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_dim = hidden_states.shape[-1] hidden_states = hidden_states.view(-1, hidden_dim) - shared_output, final_hidden_states = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index d0e6cb6ada8b..cef934222382 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -184,7 +184,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -192,8 +191,7 @@ def __init__( topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -207,23 +205,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states.to(dtype=torch.float32)) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - - if self.shared_experts is not None: - shared_output, final_hidden_states = fused_moe_out - assert shared_output is not None - final_hidden_states = ( - final_hidden_states * self.routed_scaling_factor + shared_output - ) - else: - final_hidden_states = fused_moe_out * self.routed_scaling_factor - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index a0130402c66f..35d30006a66a 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -39,7 +39,6 @@ get_ep_group, get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention @@ -445,7 +444,6 @@ def __init__( top_k=top_k, hidden_size=config.hidden_size, intermediate_size=intermediate_size, - reduce_results=False, renormalize=top_k > 1, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -464,11 +462,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_mlp is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/kimi_linear.py b/vllm/model_executor/models/kimi_linear.py index 4cd7b63c1472..e586a3ac3469 100644 --- a/vllm/model_executor/models/kimi_linear.py +++ b/vllm/model_executor/models/kimi_linear.py @@ -11,11 +11,10 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.kda import KimiDeltaAttention from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -132,12 +131,25 @@ def __init__( self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) - self.experts = FusedMoE( + if self.num_shared_experts is not None: + intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, num_experts=num_experts, top_k=config.num_experts_per_token, hidden_size=hidden_size, intermediate_size=moe_intermediate_size, - reduce_results=False, renormalize=moe_renormalize, quant_config=quant_config, use_grouped_topk=config.use_grouped_topk, @@ -146,34 +158,16 @@ def __init__( prefix=f"{prefix}.experts", scoring_func=config.moe_router_activation_func, e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, ) - if self.num_shared_experts is not None: - intermediate_size = moe_intermediate_size * self.num_shared_experts - self.shared_experts = KimiMLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=False, - prefix=f"{prefix}.shared_experts", - ) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) - if self.num_shared_experts is not None: - shared_output = self.shared_experts(hidden_states) router_logits, _ = self.gate(hidden_states) - final_hidden_states = ( - self.experts(hidden_states=hidden_states, router_logits=router_logits) - * self.routed_scaling_factor + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits ) - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) @@ -482,7 +476,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if self.config.is_moe: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = FusedMoE.make_expert_params_mapping( + expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index b7ca710eadcb..5d592310e455 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -150,7 +150,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, # needed for softmax score func @@ -161,6 +160,7 @@ def __init__( num_redundant_experts=self.n_redundant_experts, scoring_func="sigmoid", e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -170,16 +170,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - final_hidden_states = ( - self.experts(hidden_states=hidden_states, router_logits=router_logits) - * self.routed_scaling_factor + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits ) - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) - return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index b84b4e2ae512..a1c0ac896052 100644 --- a/vllm/model_executor/models/llama4.py +++ b/vllm/model_executor/models/llama4.py @@ -135,7 +135,6 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): custom_routing_function=Llama4MoE.custom_routing_function, intermediate_size=intermediate_size_moe, apply_router_weight_on_input=True, - reduce_results=False, renormalize=False, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -151,19 +150,14 @@ def forward(self, hidden_states): router_logits, _ = self.router(hidden_states) - shared_out, routed_out = self.experts( + experts_out = self.experts( hidden_states=hidden_states, router_logits=router_logits, ) - experts_out = routed_out + shared_out if self.is_sequence_parallel: experts_out = tensor_model_parallel_all_gather(experts_out, 0) experts_out = experts_out[:num_tokens] - elif self.tp_size > 1: - experts_out = self.experts.maybe_all_reduce_tensor_model_parallel( - experts_out - ) return experts_out diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 2dc0f33cc1c5..1b068c8e4c6c 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -35,7 +35,6 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import FusedMoE @@ -103,7 +102,6 @@ def __init__( e_score_correction_bias=self.e_score_correction_bias, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - reduce_results=False, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -133,9 +131,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - final_hidden_states = final_hidden_states - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 859e34a101db..e9c7798387d3 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -216,7 +216,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=self.moe_hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -231,6 +230,10 @@ def __init__( num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, routed_input_transform=self.fc1_latent_proj, + routed_output_transform=self.fc2_latent_proj + if self.use_latent_moe + else None, + routed_scaling_factor=self.routed_scaling_factor, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -243,38 +246,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - # SharedFusedMoE handles: - # - shared experts (with original hidden_states) - # - routed_input_transform (fc1_latent_proj) for latent MoE - # - multistream parallelism between shared and routed experts - shared_output, final_hidden_states = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - # Fix FP16 overflow - # See DeepseekV2DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - shared_output *= 1.0 / self.routed_scaling_factor - - # TODO: See SharedFusedMoE.apply_routed_input_transform - # for bandwidth optimization - if self.use_latent_moe: - final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states) - - if self.shared_experts is not None: - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 994ae82529ab..789e4260836f 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -206,7 +206,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -214,8 +213,7 @@ def __init__( topk_group=1, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -234,33 +232,15 @@ def forward( router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - if hidden_states.dtype != torch.float16: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index 4b0c756165a5..1234aad77be2 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -167,7 +167,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -184,12 +183,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_expert is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index f2ce070be8b4..f0f69d435379 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -212,7 +212,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -234,22 +233,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - shared_out, fused_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - final_hidden_states = ( - shared_out + fused_out if shared_out is not None else fused_out - ) if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) # return to 1d if input is 1d return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 2f255716547e..989b888acc1f 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -329,7 +329,6 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=getattr(config, "norm_topk_prob", True), quant_config=quant_config, prefix=f"{prefix}.experts", @@ -359,18 +358,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states=hidden_states, router_logits=router_logits ) - if self.shared_expert is not None: - final_hidden_states = final_hidden_states[0] + final_hidden_states[1] - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501 - final_hidden_states - ) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index 18b689166a5f..636a121c590b 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -14,7 +14,6 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul @@ -71,7 +70,6 @@ def __init__( top_k=config.moe_top_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_expert_weight, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -94,8 +92,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index bb4bf14a9632..018f78956029 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -379,7 +379,6 @@ def __init__( top_k=config.moe_top_k, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_expert_weight, quant_config=quant_config, activation=activation, @@ -397,30 +396,16 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.view(-1, hidden_dim) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) + # TODO(bnell): this gate could be moved into the FusedMoE? router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.share_expert is None: - assert shared_output is None - - if self.share_expert is not None: - assert shared_output is not None - final_hidden_states += shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index f65a197abcfc..a654adcc289d 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -204,28 +204,6 @@ def recursive_replace(self): ) assert intermediate_size is not None - # If there are shared experts, the results are - # reduced after mlp.forward() not inside FusedMoE - num_shared_experts = getattr_iter( - text_config, - [ - "n_shared_experts", # DeepSeek, Docs, GLM - "moe_num_shared_experts", # Aria, Ernie - ], - 0, - ) - reduce_results = num_shared_experts == 0 - - def add_all_reduce(mlp: nn.Module): - """Adds an all-reduce to the output of `mlp.forward()`.""" - - class MLPWithAllReduce(mlp.__class__): - def forward(self, *args, **kwargs): - output = super().forward(*args, **kwargs) - return self.experts.maybe_all_reduce_tensor_model_parallel(output) - - mlp.__class__ = MLPWithAllReduce - # Unused kwargs since we use custom_routing_function: # - `scoring_func` and `e_score_correction_bias` only used for grouped # topk routing inside vLLM and are non-trivial to infer @@ -265,7 +243,7 @@ def forward(self, *args, **kwargs): self.num_physical_experts = num_experts + num_redundant_experts self.num_local_physical_experts = self.num_physical_experts // ep_size self.num_routed_experts = num_experts - self.num_shared_experts = num_shared_experts + self.num_shared_experts = 0 self.num_redundant_experts = num_redundant_experts # Recursively fuse MoE layers @@ -289,14 +267,10 @@ def _recursive_replace(module: nn.Module, prefix: str): if "bias" in experts_param_name: has_bias = True break - # Double check there are no shared experts - nonlocal reduce_results - if reduce_results: + # Detect shared experts if config doesn't specify + if self.num_shared_experts == 0: for mlp_param_name, _ in mlp.named_parameters(): if "shared_expert" in mlp_param_name: - reduce_results = False - # If the config does not specify num_shared_experts, but - # the model has shared experts, we assume there is one. self.num_shared_experts = 1 break # Replace experts module with FusedMoE @@ -305,7 +279,7 @@ def _recursive_replace(module: nn.Module, prefix: str): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=reduce_results, + reduce_results=True, renormalize=renormalize, # Hard coded because topk happens in Transformers use_grouped_topk=False, @@ -326,13 +300,6 @@ def _recursive_replace(module: nn.Module, prefix: str): self.moe_layers.append(fused_experts) self.expert_weights.append(fused_experts.get_expert_weights()) self.num_moe_layers += 1 - # If results are not all-reduced in FusedMoE, ensure they - # are all-reduced at the end of mlp.forward() if tensor - # parallel or expert parallel is enabled - if not reduce_results and ( - fused_experts.tp_size > 1 or fused_experts.ep_size > 1 - ): - add_all_reduce(mlp) else: _recursive_replace(child_module, prefix=qual_name) From 3ad7dc596bbd38bc78c15744b93b949f3d021e38 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 3 Mar 2026 18:58:54 -0500 Subject: [PATCH 042/108] cleanups Signed-off-by: Bill Nell --- vllm/forward_context.py | 2 +- .../fused_moe/runner/moe_runner_base.py | 98 ++++++++++++------- 2 files changed, 64 insertions(+), 36 deletions(-) diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 1231263bcfe3..bf0f9da6eaff 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -213,7 +213,7 @@ class ForwardContext: # For torch.compile cold start times, we need to avoid hard-coding # any strings into the graph. Right now, the vllm.moe_forward - # custom operator hard-codes strings into + # and vllm.moe_forward_shared custom operators hard-code strings into # the graph. # # The workaround is to store a list of the strings that each of those diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 60bdc2280d5b..3729e83aafba 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -54,7 +54,7 @@ def get_layer_from_name(layer_name: str) -> torch.nn.Module: raise AssertionError( "We expected the number of MOE layers in `all_moe_layers` " "to be equal to the number of " - "vllm.moe_forward calls." + "{vllm.moe_forward, vllm.moe_forward_shared} calls." ) layer_name = all_moe_layers[moe_layer_index] forward_context.moe_layer_index += 1 @@ -99,13 +99,45 @@ def _moe_forward_fake( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: - # For latent MoE with reduce_results=True: output has full hidden_size - # (from shared_experts_input), not latent_size (from hidden_states). - if shared_experts_input is not None: - return torch.empty_like(shared_experts_input) return torch.empty_like(hidden_states) +def _moe_forward_shared( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_shared_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + # Output shapes: + # - fused_out: same as hidden_states (routed experts use transformed size) + # - shared_out: same as shared_experts_input if provided, else same as + # hidden_states + # (For latent MoE: shared experts use original hidden_size, not latent size) + fused_out = torch.empty_like(hidden_states) + + if shared_experts_input is not None: + shared_out = torch.empty_like(shared_experts_input) + else: + shared_out = torch.empty_like(hidden_states) + + return shared_out, fused_out + + def _moe_forward_shared( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -238,7 +270,9 @@ def _must_reduce_shared_expert_outputs(self) -> bool: and self.quant_method.moe_kernel.output_is_reduced() ) - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): + def _maybe_all_reduce_tensor_model_parallel( + self, final_hidden_states: torch.Tensor + ): """ Some combine kernels reduce across GPU ranks by default. """ @@ -256,10 +290,6 @@ def apply_routed_input_transform( is saved separately so shared experts get [S, hidden_size] while routed experts get the transformed [S, moe_latent_size]. - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside FusedMoE to all-reduce on the smaller latent - dimension. - Returns (possibly transformed) hidden states and the input for shared experts (or None if there are no shared experts). """ @@ -279,10 +309,9 @@ def apply_routed_input_transform( def _combine_and_reduce( self, states: torch.Tensor, - trunc_sizes: list[int], + trunc_size: int, ) -> torch.Tensor: - assert len(trunc_sizes) == 1 - result = states[..., : trunc_sizes[0]] + result = states[..., :trunc_size] if self.output_scale is not None: result = result * self.output_scale @@ -290,7 +319,7 @@ def _combine_and_reduce( if not self.moe_config.is_sequence_parallel and ( self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1 ): - result = self.maybe_all_reduce_tensor_model_parallel(result) + result = self._maybe_all_reduce_tensor_model_parallel(result) return result @@ -309,7 +338,7 @@ def _maybe_pad_hidden_states( self, shared_experts_input: torch.Tensor | None, hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, list[int]]: + ) -> tuple[torch.Tensor, int]: shared_experts_hidden_dim = ( shared_experts_input.shape[-1] if shared_experts_input is not None else 0 ) @@ -326,9 +355,9 @@ def _maybe_pad_hidden_states( ) if self.routed_output_transform is not None and shared_experts_hidden_dim > 0: - orig_hidden_dims = [shared_experts_hidden_dim] + orig_hidden_dims = shared_experts_hidden_dim else: - orig_hidden_dims = [transformed_hidden_dim] + orig_hidden_dims = transformed_hidden_dim return hidden_states, orig_hidden_dims @@ -452,19 +481,31 @@ def forward( hidden_states ) - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + hidden_states, og_hidden_dim = self._maybe_pad_hidden_states( shared_experts_input, hidden_states, ) - fused_output = self.forward_entry( + result = self.forward_entry( hidden_states, router_logits, shared_experts_input, self._encode_layer_name(), ) - result = self._combine_and_reduce(fused_output, og_hidden_dims) + if self.shared_experts is not None: + assert isinstance(result, tuple) + shared_output, fused_output = result + # Apply output transform (e.g. latent → full dim) + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + # If combine kernel already reduced fused, reduce shared to match + if self._must_reduce_shared_expert_outputs(): + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = shared_output + fused_output + + result = self._combine_and_reduce(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) @@ -474,7 +515,7 @@ def forward_dispatch( hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -486,26 +527,13 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - result = self._forward_impl( + return self._forward_impl( layer, hidden_states, router_logits, shared_experts_input, ) - if isinstance(result, tuple): - shared_output, fused_output = result - # Apply output transform (e.g. latent → full dim) - if self.routed_output_transform is not None: - r = self.routed_output_transform(fused_output) - fused_output = r[0] if isinstance(r, tuple) else r - # If combine kernel already reduced fused, reduce shared to match - if self._must_reduce_shared_expert_outputs(): - shared_output = tensor_model_parallel_all_reduce(shared_output) - result = shared_output + fused_output - - return result - @abstractmethod def _forward_impl( self, From eb38e88f9b3518fdda1f2327edc38b898773089e Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 3 Mar 2026 21:59:30 -0500 Subject: [PATCH 043/108] apply_scale_to_output flag Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 7 ++++--- .../layers/fused_moe/runner/default_moe_runner.py | 6 ++++-- .../layers/fused_moe/runner/moe_runner_base.py | 12 +++++++----- .../layers/fused_moe/runner/moe_runner_factory.py | 6 ++++-- vllm/model_executor/models/deepseek_v2.py | 7 ++++++- vllm/model_executor/models/dots1.py | 2 +- vllm/model_executor/models/nemotron_h.py | 2 +- vllm/model_executor/models/openpangu.py | 4 +++- 8 files changed, 30 insertions(+), 16 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c82e3d16a458..732854d0182f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -338,7 +338,7 @@ def __init__( shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, routed_output_transform: torch.nn.Module | None = None, - output_scale: float | None = None, + apply_scale_to_output: bool = False, zero_expert_type: str | None = None, ): super().__init__() @@ -346,7 +346,7 @@ def __init__( self._gate = gate self._routed_input_transform = routed_input_transform self._routed_output_transform = routed_output_transform - self._output_scale = output_scale + self._apply_scale_to_output = apply_scale_to_output if params_dtype is None: params_dtype = torch.get_default_dtype() @@ -682,7 +682,8 @@ def _init_runner(self) -> MoERunner: quant_method=self.quant_method, enable_dbo=self.vllm_config.parallel_config.enable_dbo, routed_output_transform=self._routed_output_transform, - output_scale=self._output_scale, + apply_scale_to_output=self._apply_scale_to_output, + routed_scaling_factor=self.routed_scaling_factor, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 523a5362da51..b6441ae38e5a 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -63,7 +63,8 @@ def __init__( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - output_scale: float | None = None, + apply_scale_to_output: bool = False, + routed_scaling_factor: float = 1.0, ): super().__init__( layer, @@ -75,7 +76,8 @@ def __init__( quant_method, enable_dbo, routed_output_transform=routed_output_transform, - output_scale=output_scale, + apply_scale_to_output=apply_scale_to_output, + routed_scaling_factor=routed_scaling_factor, ) @property diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 3729e83aafba..c9e77b0ac8fe 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -220,7 +220,8 @@ def __init__( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - output_scale: float | None = None, + apply_scale_to_output: bool = False, + routed_scaling_factor: float = 1.0, ): super().__init__() self.moe_config = moe_config @@ -232,7 +233,8 @@ def __init__( self.enable_dbo = enable_dbo self.enable_eplb = moe_config.moe_parallel_config.enable_eplb self.routed_output_transform = routed_output_transform - self.output_scale = output_scale + self.apply_scale_to_output = apply_scale_to_output + self.routed_scaling_factor = routed_scaling_factor # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name @@ -313,9 +315,6 @@ def _combine_and_reduce( ) -> torch.Tensor: result = states[..., :trunc_size] - if self.output_scale is not None: - result = result * self.output_scale - if not self.moe_config.is_sequence_parallel and ( self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1 ): @@ -505,6 +504,9 @@ def forward( shared_output = tensor_model_parallel_all_reduce(shared_output) result = shared_output + fused_output + if self.apply_scale_to_output: + result = result * self.routed_scaling_factor + result = self._combine_and_reduce(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index e6100814cad8..c9a7efc431c3 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -34,7 +34,8 @@ def create_moe_runner( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - output_scale: float | None = None, + apply_scale_to_output: bool = False, + routed_scaling_factor: float = 1.0, ) -> MoERunner: runner = DefaultMoERunner( layer, @@ -46,7 +47,8 @@ def create_moe_runner( quant_method, enable_dbo, routed_output_transform=routed_output_transform, - output_scale=output_scale, + apply_scale_to_output=apply_scale_to_output, + routed_scaling_factor=routed_scaling_factor, ) if moe_config.moe_parallel_config.use_dp_chunking: return ChunkingMoERunner(runner) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 1a64299789fb..95b4525cfa3a 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -319,7 +319,12 @@ def __init__( topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), - routed_scaling_factor=self.routed_scaling_factor, + # we do scaling outside, set factor to 1.0 to avoid double mul + # aiter applies routed_scaling_factor internally + routed_scaling_factor=1.0 + if not self.is_rocm_aiter_moe_enabled + else self.routed_scaling_factor, + apply_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index fc2636db48ab..fa867c5893fb 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -171,7 +171,7 @@ def __init__( # we do scaling outside, set factor to 1.0 to avoid double mul routed_scaling_factor=1.0, e_score_correction_bias=self.gate.e_score_correction_bias, - output_scale=self.routed_scaling_factor, + apply_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e9c7798387d3..1a2abe90be7b 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -233,7 +233,7 @@ def __init__( routed_output_transform=self.fc2_latent_proj if self.use_latent_moe else None, - routed_scaling_factor=self.routed_scaling_factor, + apply_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 789e4260836f..0d7e0f544a21 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -213,11 +213,13 @@ def __init__( topk_group=1, prefix=f"{prefix}.experts", scoring_func="sigmoid", - routed_scaling_factor=self.routed_scaling_factor, + # we do scaling outside, set factor to 1.0 to avoid double mul + routed_scaling_factor=1.0, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, + apply_scale_to_output=True, ) def forward( From 1c17922294b09d113c2f6627d2716e092d05171a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 3 Mar 2026 22:36:32 -0500 Subject: [PATCH 044/108] fix fp16 scaling factor stuff Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 4 +- .../fused_moe/runner/moe_runner_base.py | 50 ++++++++++++++----- vllm/model_executor/models/deepseek_v2.py | 4 +- vllm/model_executor/models/dots1.py | 3 +- vllm/model_executor/models/nemotron_h.py | 1 + vllm/model_executor/models/openpangu.py | 5 +- 6 files changed, 45 insertions(+), 22 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 732854d0182f..100a1bfb79a0 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -523,7 +523,9 @@ def __init__( topk_group=topk_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, + routed_scaling_factor=routed_scaling_factor + if not apply_scale_to_output + else 1.0, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=self.num_fused_shared_experts, enable_eplb=enable_eplb, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index c9e77b0ac8fe..c4ffb4f1db6b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -233,7 +233,9 @@ def __init__( self.enable_dbo = enable_dbo self.enable_eplb = moe_config.moe_parallel_config.enable_eplb self.routed_output_transform = routed_output_transform - self.apply_scale_to_output = apply_scale_to_output + self.apply_scale_to_output = ( + apply_scale_to_output and routed_scaling_factor != 1.0 + ) self.routed_scaling_factor = routed_scaling_factor # Needed for string -> FusedMoE layer lookup in custom ops. @@ -268,7 +270,8 @@ def _must_reduce_shared_expert_outputs(self) -> bool: early. """ return ( - self.quant_method.moe_kernel is not None + self.shared_experts is not None + and self.quant_method.moe_kernel is not None and self.quant_method.moe_kernel.output_is_reduced() ) @@ -308,7 +311,7 @@ def apply_routed_input_transform( hidden_states if self.shared_experts is not None else None, ) - def _combine_and_reduce( + def _maybe_reduce_output( self, states: torch.Tensor, trunc_size: int, @@ -492,22 +495,43 @@ def forward( self._encode_layer_name(), ) + ############################################################## + + # Extract outputs from result if self.shared_experts is not None: assert isinstance(result, tuple) shared_output, fused_output = result - # Apply output transform (e.g. latent → full dim) - if self.routed_output_transform is not None: - r = self.routed_output_transform(fused_output) - fused_output = r[0] if isinstance(r, tuple) else r - # If combine kernel already reduced fused, reduce shared to match - if self._must_reduce_shared_expert_outputs(): - shared_output = tensor_model_parallel_all_reduce(shared_output) - result = shared_output + fused_output + else: + shared_output = None + fused_output = result + + # Apply output transform (e.g. latent → full dim) + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + + # If combine kernel already reduced fused, reduce shared to match + if self._must_reduce_shared_expert_outputs(): + assert shared_output is not None + shared_output = tensor_model_parallel_all_reduce(shared_output) if self.apply_scale_to_output: - result = result * self.routed_scaling_factor + # FP16 overflow protection: instead of (shared + fused) * scale, + # compute shared/scale + fused. The decoder layer compensates + # with matching divisions (see DeepseekV2DecoderLayer). + if fused_output.dtype != torch.float16: + fused_output *= self.routed_scaling_factor + elif shared_output is not None: + shared_output *= 1.0 / self.routed_scaling_factor + + if shared_output is not None: + result = shared_output + fused_output + else: + result = fused_output + + ############################################################## - result = self._combine_and_reduce(result, og_hidden_dim) + result = self._maybe_reduce_output(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 95b4525cfa3a..c0d5849a2921 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -321,9 +321,7 @@ def __init__( scoring_func=getattr(config, "scoring_func", "softmax"), # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, apply_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index fa867c5893fb..f0bbde4bcdbe 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -168,9 +168,8 @@ def __init__( topk_group=config.topk_group, prefix=f"{prefix}.experts", scoring_func=config.scoring_func, - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, apply_scale_to_output=True, ) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 1a2abe90be7b..6359f0798b51 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -233,6 +233,7 @@ def __init__( routed_output_transform=self.fc2_latent_proj if self.use_latent_moe else None, + routed_scaling_factor=self.routed_scaling_factor, apply_scale_to_output=True, ) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 0d7e0f544a21..409370a77089 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -213,13 +213,12 @@ def __init__( topk_group=1, prefix=f"{prefix}.experts", scoring_func="sigmoid", - # we do scaling outside, set factor to 1.0 to avoid double mul - routed_scaling_factor=1.0, + routed_scaling_factor=self.routed_scaling_factor, + apply_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, - apply_scale_to_output=True, ) def forward( From e81109505302f5789591eedaa630b07bb9b86dcb Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 4 Mar 2026 09:48:44 -0500 Subject: [PATCH 045/108] cleanups Signed-off-by: Bill Nell --- .../fused_moe/runner/moe_runner_base.py | 140 +++++++++++------- vllm/model_executor/models/flex_olmo.py | 1 - vllm/model_executor/models/gpt_oss.py | 1 - vllm/model_executor/models/granitemoe.py | 1 - vllm/model_executor/models/grok1.py | 1 - vllm/model_executor/models/interns1_pro.py | 1 - vllm/model_executor/models/jamba.py | 1 - vllm/model_executor/models/longcat_flash.py | 1 - vllm/model_executor/models/mimo_v2_flash.py | 1 - vllm/model_executor/models/minimax_text_01.py | 1 - vllm/model_executor/models/mixtral.py | 1 - vllm/model_executor/models/olmoe.py | 1 - vllm/model_executor/models/phimoe.py | 1 - .../model_executor/models/transformers/moe.py | 1 - 14 files changed, 83 insertions(+), 70 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index c4ffb4f1db6b..068d18d8bb9a 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -189,6 +189,15 @@ def _moe_forward_shared_fake( ) +def _unpack( + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor | None, torch.Tensor]: + if isinstance(result, tuple): + return result + else: + return (None, result) + + class MoERunnerBase(MoERunner): """ Abstract base class providing common functionality for MoE runner implementations. @@ -256,36 +265,6 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) - def _must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.shared_experts is not None - and self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - - def _maybe_all_reduce_tensor_model_parallel( - self, final_hidden_states: torch.Tensor - ): - """ - Some combine kernels reduce across GPU ranks by default. - """ - if self._must_reduce_shared_expert_outputs(): - return final_hidden_states - else: - return tensor_model_parallel_all_reduce(final_hidden_states) - def apply_routed_input_transform( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -311,6 +290,58 @@ def apply_routed_input_transform( hidden_states if self.shared_experts is not None else None, ) + def apply_routed_output_transform( + self, + fused_output: torch.Tensor, + ) -> torch.Tensor: + if self.routed_output_transform is not None: + r = self.routed_output_transform(fused_output) + fused_output = r[0] if isinstance(r, tuple) else r + return fused_output + + def _maybe_apply_output_scale( + self, + shared_output: torch.Tensor | None, + fused_output: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + if self.apply_scale_to_output: + # FP16 overflow protection: instead of (shared + fused) * scale, + # compute shared/scale + fused. The decoder layer compensates + # with matching divisions (see DeepseekV2DecoderLayer). + if fused_output.dtype != torch.float16: + fused_output *= self.routed_scaling_factor + elif shared_output is not None: + shared_output *= 1.0 / self.routed_scaling_factor + return shared_output, fused_output + + def _must_reduce_shared_expert_output(self) -> bool: + """ + The shared_experts are typically computed using the RowParallelLinear + layer. The result of this function is typically used as + the reduce_results argument to the module. + When just tensor-parallel is used, it is not required to reduce + the shared_experts results immediately. Instead we reduce at the + once at the end of the MoE op. (Refer to DeepSeekV2MoE module) + With EP and all2all kernels - this is no longer viable as all + GPU ranks in DP, produce the complete set of hidden_states. + Therefore it is required that we reduce the shared_experts output + early. + """ + return ( + self.shared_experts is not None + and self.quant_method.moe_mk is not None + and self.quant_method.moe_mk.output_is_reduced() + ) + + def _maybe_reduce_shared_expert_output( + self, + shared_output: torch.Tensor | None, + ) -> torch.Tensor | None: + if self._must_reduce_shared_expert_output(): + assert shared_output is not None + shared_output = tensor_model_parallel_all_reduce(shared_output) + return shared_output + def _maybe_reduce_output( self, states: torch.Tensor, @@ -318,10 +349,12 @@ def _maybe_reduce_output( ) -> torch.Tensor: result = states[..., :trunc_size] - if not self.moe_config.is_sequence_parallel and ( - self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1 + if ( + not self.moe_config.is_sequence_parallel + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not self._must_reduce_shared_expert_output() ): - result = self._maybe_all_reduce_tensor_model_parallel(result) + result = tensor_model_parallel_all_reduce(result) return result @@ -495,42 +528,34 @@ def forward( self._encode_layer_name(), ) - ############################################################## + # + # Note: there are two all-reduce points below. They are mutually + # exclusive, controlled by _must_reduce_shared_expert_output(): + # - When True: the combine kernel already reduced fused_output, + # so we reduce shared_output here to match, then skip the + # all-reduce in _maybe_reduce_output. + # - When False: neither output is reduced yet, so we combine + # them first and all-reduce the sum in _maybe_reduce_output. # Extract outputs from result - if self.shared_experts is not None: - assert isinstance(result, tuple) - shared_output, fused_output = result - else: - shared_output = None - fused_output = result + shared_output, fused_output = _unpack(result) # Apply output transform (e.g. latent → full dim) - if self.routed_output_transform is not None: - r = self.routed_output_transform(fused_output) - fused_output = r[0] if isinstance(r, tuple) else r + fused_output = self.apply_routed_output_transform(fused_output) - # If combine kernel already reduced fused, reduce shared to match - if self._must_reduce_shared_expert_outputs(): - assert shared_output is not None - shared_output = tensor_model_parallel_all_reduce(shared_output) + # If combine kernel already reduced fused, reduce shared to match. + # See note above re: the two all-reduce points. + shared_output = self._maybe_reduce_shared_expert_output(shared_output) - if self.apply_scale_to_output: - # FP16 overflow protection: instead of (shared + fused) * scale, - # compute shared/scale + fused. The decoder layer compensates - # with matching divisions (see DeepseekV2DecoderLayer). - if fused_output.dtype != torch.float16: - fused_output *= self.routed_scaling_factor - elif shared_output is not None: - shared_output *= 1.0 / self.routed_scaling_factor + shared_output, fused_output = self._maybe_apply_output_scale( + shared_output, fused_output + ) if shared_output is not None: result = shared_output + fused_output else: result = fused_output - ############################################################## - result = self._maybe_reduce_output(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) @@ -545,6 +570,7 @@ def forward_dispatch( # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() + # TODO: move _maybe_gate and _maybe_apply_shared_experts to forward? router_logits = self._maybe_gate(hidden_states, router_logits) self._maybe_apply_shared_experts( diff --git a/vllm/model_executor/models/flex_olmo.py b/vllm/model_executor/models/flex_olmo.py index a2e2adc2a6bd..882d99c5a298 100644 --- a/vllm/model_executor/models/flex_olmo.py +++ b/vllm/model_executor/models/flex_olmo.py @@ -76,7 +76,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): top_k=hf_config.num_experts_per_tok, hidden_size=hf_config.hidden_size, intermediate_size=hf_config.intermediate_size, - reduce_results=True, renormalize=False, quant_config=None, tp_size=tp_size, diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 482056250a1e..083defc9c2db 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -186,7 +186,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, - reduce_results=True, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index 171b2e0ec5a0..f57a8c942bb4 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -104,7 +104,6 @@ def __init__( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/grok1.py b/vllm/model_executor/models/grok1.py index 0bd6a8f3d606..c9aa3d2068f0 100644 --- a/vllm/model_executor/models/grok1.py +++ b/vllm/model_executor/models/grok1.py @@ -209,7 +209,6 @@ def __init__( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=renormalize, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 1c9f1a7bfc16..968d2418819f 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -176,7 +176,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=True, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index 980bcffb5f9b..b4b3b6873db3 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -90,7 +90,6 @@ def __init__( self.intermediate_size, tp_size=tp_size, params_dtype=params_dtype, - reduce_results=True, renormalize=False, use_grouped_topk=False, quant_config=quant_config, diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index 375b0b69b1f9..945fcb61509b 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -300,7 +300,6 @@ def __init__( top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, params_dtype=params_dtype, renormalize=False, quant_config=quant_config, diff --git a/vllm/model_executor/models/mimo_v2_flash.py b/vllm/model_executor/models/mimo_v2_flash.py index 43475ed690c9..0b466f16601a 100644 --- a/vllm/model_executor/models/mimo_v2_flash.py +++ b/vllm/model_executor/models/mimo_v2_flash.py @@ -162,7 +162,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=True, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/minimax_text_01.py b/vllm/model_executor/models/minimax_text_01.py index 21d74d8b0580..67d7cb2d8bcb 100644 --- a/vllm/model_executor/models/minimax_text_01.py +++ b/vllm/model_executor/models/minimax_text_01.py @@ -162,7 +162,6 @@ def __init__( hidden_size=self.hidden_size, intermediate_size=self.intermediate_size * self.tp_size, params_dtype=self.params_dtype, - reduce_results=True, renormalize=True, quant_config=self.quant_config, tp_size=self.tp_size, diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 376fd7a1709d..c182444f667d 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -132,7 +132,6 @@ def __init__( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/olmoe.py b/vllm/model_executor/models/olmoe.py index f0afe0e997cc..fcde2e41afbb 100644 --- a/vllm/model_executor/models/olmoe.py +++ b/vllm/model_executor/models/olmoe.py @@ -98,7 +98,6 @@ def __init__( top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, renormalize=False, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/phimoe.py b/vllm/model_executor/models/phimoe.py index 0b55b7ec8392..7d6083f202e6 100644 --- a/vllm/model_executor/models/phimoe.py +++ b/vllm/model_executor/models/phimoe.py @@ -281,7 +281,6 @@ def __init__( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=params_dtype, - reduce_results=True, renormalize=False, quant_config=quant_config, tp_size=tp_size, diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index a654adcc289d..faafc5037e29 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -279,7 +279,6 @@ def _recursive_replace(module: nn.Module, prefix: str): top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, - reduce_results=True, renormalize=renormalize, # Hard coded because topk happens in Transformers use_grouped_topk=False, From 2df1794d14cb039cd2245a71b4ea673ad451e921 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 4 Mar 2026 10:02:31 -0500 Subject: [PATCH 046/108] add claude generated comments Signed-off-by: Bill Nell --- .../fused_moe/runner/moe_runner_base.py | 116 +++++++++++++++--- 1 file changed, 99 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 068d18d8bb9a..325951199ba5 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -77,7 +77,7 @@ def _resolve_layer_name(layer_name: str | ModuleName) -> str: # Note: _moe_forward and _moe_forward_shared should not contain any # implementation details, They should merely pass along control to -# the runner's 'forward_dispatch' method. +# the runner's '_forward_dispatch' method. def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -85,7 +85,7 @@ def _moe_forward( layer_name: _layer_name_type, ) -> torch.Tensor: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, @@ -109,7 +109,7 @@ def _moe_forward_shared( layer_name: str, ) -> tuple[torch.Tensor, torch.Tensor]: layer = get_layer_from_name(layer_name) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, @@ -294,6 +294,12 @@ def apply_routed_output_transform( self, fused_output: torch.Tensor, ) -> torch.Tensor: + """Apply transform to routed expert output (e.g., latent to full dim). + + Used by latent MoE models (e.g., NemotronH) where routed experts + operate in a compressed latent space and need projection back to + the full hidden dimension before combining with shared expert output. + """ if self.routed_output_transform is not None: r = self.routed_output_transform(fused_output) fused_output = r[0] if isinstance(r, tuple) else r @@ -304,10 +310,15 @@ def _maybe_apply_output_scale( shared_output: torch.Tensor | None, fused_output: torch.Tensor, ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Apply routed_scaling_factor to the output with FP16 overflow + protection. + + When apply_scale_to_output is True, scales the fused expert output + by routed_scaling_factor. For FP16, avoids overflow by dividing + shared_output by the scale instead (the decoder layer compensates + with matching divisions). + """ if self.apply_scale_to_output: - # FP16 overflow protection: instead of (shared + fused) * scale, - # compute shared/scale + fused. The decoder layer compensates - # with matching divisions (see DeepseekV2DecoderLayer). if fused_output.dtype != torch.float16: fused_output *= self.routed_scaling_factor elif shared_output is not None: @@ -337,6 +348,13 @@ def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, ) -> torch.Tensor | None: + """All-reduce shared expert output when the combine kernel already + reduced fused output. + + This is the "early" all-reduce path. When the combine kernel produces + already-reduced fused output, shared output must be reduced separately + to match. See _must_reduce_shared_expert_output for details. + """ if self._must_reduce_shared_expert_output(): assert shared_output is not None shared_output = tensor_model_parallel_all_reduce(shared_output) @@ -347,6 +365,13 @@ def _maybe_reduce_output( states: torch.Tensor, trunc_size: int, ) -> torch.Tensor: + """Truncate padded dimensions and all-reduce the combined output. + + This is the "late" all-reduce path. When neither fused nor shared + output was individually reduced, the combined sum is all-reduced + here. Skipped when sequence-parallel is active (SP handles its + own reduction) or when the early path already reduced both outputs. + """ result = states[..., :trunc_size] if ( @@ -359,6 +384,13 @@ def _maybe_reduce_output( return result def _encode_layer_name(self) -> str | ModuleName: + """Return the layer name string for custom op layer lookup. + + When torch.compile is active, returns "from_forward_context" so the + custom op resolves the layer via ForwardContext at runtime (avoiding + graph breaks). Falls back to the literal layer name for unit tests + or when ForwardContext is unavailable. + """ if HAS_OPAQUE_TYPE: return ModuleName(self.layer_name) # Can be unavailable or None in unittests @@ -374,6 +406,14 @@ def _maybe_pad_hidden_states( shared_experts_input: torch.Tensor | None, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, int]: + """Pad hidden_states to moe_config.hidden_dim and compute the + original dimension for later truncation. + + For latent MoE, the routed hidden_states may be smaller than + hidden_dim. Padding ensures uniform tensor sizes through the + fused MoE kernel. The returned trunc_size is used by + _maybe_reduce_output to strip the padding from the result. + """ shared_experts_hidden_dim = ( shared_experts_input.shape[-1] if shared_experts_input is not None else 0 ) @@ -401,6 +441,13 @@ def _maybe_apply_shared_experts( shared_experts_input: torch.Tensor | None, order: SharedExpertsOrder, ): + """Trigger shared expert computation at the specified ordering point. + + Shared experts can run at different points relative to routed experts + (EXTERNAL, BEFORE_QUANT_METHOD, AFTER_QUANT_METHOD) depending on the + model's overlap strategy. Only fires if shared experts are configured + and the order matches the shared experts' configured execution point. + """ if self.shared_experts is not None: assert shared_experts_input is not None self.shared_experts.apply(shared_experts_input, order) @@ -412,6 +459,12 @@ def _apply_quant_method( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Run expert routing and the fused MoE kernel via the quant method. + + Orchestrates shared expert execution (before/after), expert selection + via the router, and the actual fused MoE computation. Returns + (shared_expert_output, fused_expert_output). + """ # Run this before quant_method to avoid inplace issues. self._maybe_apply_shared_experts( shared_experts_input, @@ -449,6 +502,13 @@ def _apply_quant_method( ) def _sequence_parallel_context(self): + """Return a context manager for sequence-parallel token + redistribution. + + When sequence parallelism is active, returns a context that handles + local size tracking for proper token scatter/gather. Otherwise + returns a no-op context. + """ ctx = get_forward_context() return ( ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) @@ -461,10 +521,12 @@ def _maybe_gate( hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor: - # If router/gate provided, then apply it here. - # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoE via - # separate cuda stream) + """Apply the gate module to compute router logits if provided. + + Used in overlapped mode where shared experts run in parallel with + routed experts on a separate CUDA stream. The gate is separated + from the router to allow this parallel execution. + """ if self.gate is not None: router_logits, _ = self.gate(hidden_states) return router_logits @@ -473,6 +535,12 @@ def _maybe_add_zero_expert_output( self, result: torch.Tensor, ) -> torch.Tensor: + """Add the zero expert's contribution to the final result. + + When a ZeroExpertRouter is used, it computes a bias-like output + from the "zero expert" that is added to the combined routed+shared + expert output. + """ if isinstance(self.router, ZeroExpertRouter): zero_expert_output = self.router.zero_expert_output assert zero_expert_output is not None @@ -496,24 +564,25 @@ def forward( Calling sequence - forward - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - - forward_dispatch + - _forward_dispatch - _forward_impl Note: The existence of _moe_forward and _moe_forward_shared custom ops are due to the following reasons: 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward - and _moe_forward_shared must be split. + 2. pytorch cannot handle union types in custom op signatures so + _moe_forward and _moe_forward_shared must be split. If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can potentially get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into the 'forward' method. """ - # Apply transform for routed experts (e.g., latent projection for latent MoE) + # Apply transform for routed experts (e.g., latent projection + # for latent MoE) hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states + hidden_states, ) hidden_states, og_hidden_dim = self._maybe_pad_hidden_states( @@ -540,7 +609,7 @@ def forward( # Extract outputs from result shared_output, fused_output = _unpack(result) - # Apply output transform (e.g. latent → full dim) + # Apply output transform (e.g. latent -> full dim) fused_output = self.apply_routed_output_transform(fused_output) # If combine kernel already reduced fused, reduce shared to match. @@ -560,13 +629,19 @@ def forward( return self._maybe_add_zero_expert_output(result) - def forward_dispatch( + def _forward_dispatch( self, layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Entry point called by the custom op to run the MoE computation. + + Handles pre-dispatch setup (gate application, external shared expert + triggering, quant config init) then delegates to _forward_impl within + the sequence-parallel context. + """ # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -594,4 +669,11 @@ def _forward_impl( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Core MoE computation to be implemented by subclasses. + + Performs expert routing, fused MoE kernel execution, and shared + expert computation. Returns a single tensor (fused output only) + or a tuple of (shared_output, fused_output) when shared experts + are present. + """ raise NotImplementedError From 08ac61c52b005c0927afd886310c37da9dffdb31 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 4 Mar 2026 10:45:52 -0500 Subject: [PATCH 047/108] move stuff out ot custom op Signed-off-by: Bill Nell --- .../fused_moe/runner/moe_runner_base.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 325951199ba5..ffa5c534b490 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -78,6 +78,8 @@ def _resolve_layer_name(layer_name: str | ModuleName) -> str: # Note: _moe_forward and _moe_forward_shared should not contain any # implementation details, They should merely pass along control to # the runner's '_forward_dispatch' method. +# These functions should never be called directly since they do not +# include all the functionality of the MoE layer. def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -250,7 +252,7 @@ def __init__( # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name - self.forward_entry = self._select_forward(layer) + self._forward_entry = self._select_forward(layer) def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -563,7 +565,7 @@ def forward( Calling sequence - forward - - self.forward_entry (_moe_forward or _moe_forward_shared custom op) + - self._forward_entry (_moe_forward or _moe_forward_shared custom op) - _forward_dispatch - _forward_impl @@ -590,7 +592,14 @@ def forward( hidden_states, ) - result = self.forward_entry( + router_logits = self._maybe_gate(hidden_states, router_logits) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.EXTERNAL, + ) + + result = self._forward_entry( hidden_states, router_logits, shared_experts_input, @@ -645,14 +654,6 @@ def _forward_dispatch( # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() - # TODO: move _maybe_gate and _maybe_apply_shared_experts to forward? - router_logits = self._maybe_gate(hidden_states, router_logits) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.EXTERNAL, - ) - with self._sequence_parallel_context(): return self._forward_impl( layer, From 9d39ea878183f18376116da75ab55f728753bba0 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 4 Mar 2026 15:18:41 -0500 Subject: [PATCH 048/108] fix transformers/moe.py Signed-off-by: Bill Nell --- vllm/model_executor/models/transformers/moe.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index faafc5037e29..baacb528af0d 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -204,6 +204,15 @@ def recursive_replace(self): ) assert intermediate_size is not None + num_shared_experts = getattr_iter( + text_config, + [ + "n_shared_experts", # DeepSeek, Docs, GLM + "moe_num_shared_experts", # Aria, Ernie + ], + 0, + ) + # Unused kwargs since we use custom_routing_function: # - `scoring_func` and `e_score_correction_bias` only used for grouped # topk routing inside vLLM and are non-trivial to infer @@ -243,7 +252,7 @@ def recursive_replace(self): self.num_physical_experts = num_experts + num_redundant_experts self.num_local_physical_experts = self.num_physical_experts // ep_size self.num_routed_experts = num_experts - self.num_shared_experts = 0 + self.num_shared_experts = num_shared_experts self.num_redundant_experts = num_redundant_experts # Recursively fuse MoE layers @@ -267,7 +276,8 @@ def _recursive_replace(module: nn.Module, prefix: str): if "bias" in experts_param_name: has_bias = True break - # Detect shared experts if config doesn't specify + # If the config does not specify num_shared_experts, but + # the model has shared experts, we assume there is one. if self.num_shared_experts == 0: for mlp_param_name, _ in mlp.named_parameters(): if "shared_expert" in mlp_param_name: From 65b0f325a27d8fc7e733e310084215c41e1547cb Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 15:26:11 +0000 Subject: [PATCH 049/108] tweak op registration Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index ffa5c534b490..13ca37a16a26 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -177,7 +177,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], # is this still true? + mutates_args=["hidden_states"], fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) From 9eb54fdd8db0ed4171dd306c7608a090cf2c91dc Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 16:05:56 +0000 Subject: [PATCH 050/108] fix rebase Signed-off-by: Bill Nell --- .../fused_moe/runner/moe_runner_base.py | 38 +------------------ 1 file changed, 1 insertion(+), 37 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 13ca37a16a26..ea7f1c25a74d 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -104,42 +104,6 @@ def _moe_forward_fake( return torch.empty_like(hidden_states) -def _moe_forward_shared( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: str, -) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(layer_name) - return layer.runner._forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_shared_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: str, -) -> tuple[torch.Tensor, torch.Tensor]: - # Output shapes: - # - fused_out: same as hidden_states (routed experts use transformed size) - # - shared_out: same as shared_experts_input if provided, else same as - # hidden_states - # (For latent MoE: shared experts use original hidden_size, not latent size) - fused_out = torch.empty_like(hidden_states) - - if shared_experts_input is not None: - shared_out = torch.empty_like(shared_experts_input) - else: - shared_out = torch.empty_like(hidden_states) - - return shared_out, fused_out - - def _moe_forward_shared( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -147,7 +111,7 @@ def _moe_forward_shared( layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( + return layer.runner._forward_dispatch( layer, hidden_states, router_logits, From 8e531806677ba1777b78ccd699ad1c825056f210 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 18:02:35 +0000 Subject: [PATCH 051/108] fix merge Signed-off-by: Bill Nell --- .../model_executor/layers/fused_moe/runner/moe_runner_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index ea7f1c25a74d..6f99681a1044 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -306,8 +306,8 @@ def _must_reduce_shared_expert_output(self) -> bool: """ return ( self.shared_experts is not None - and self.quant_method.moe_mk is not None - and self.quant_method.moe_mk.output_is_reduced() + and self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() ) def _maybe_reduce_shared_expert_output( From 3b463bc80b8b203574b59fb3ef0a544d3c1488d2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 19:06:20 +0000 Subject: [PATCH 052/108] fix lint Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 24edca00b6af..d8f900256ec3 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -43,7 +43,7 @@ def zero_expert_moe(dist_init, default_vllm_config): vllm_config.compilation_config.static_forward_context = dict() with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): - init_workspace_manager(torch.cuda.current_device()) + init_workspace_manager(torch.accelerator.current_device_index()) layer = FusedMoE( zero_expert_type="identity", From ec88db3d2aa23924d698e53b3f904a0834a7b809 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 19 Mar 2026 16:44:23 +0000 Subject: [PATCH 053/108] fix gate overlap Signed-off-by: Bill Nell --- .../fused_moe/runner/default_moe_runner.py | 38 +++++++++---------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 15ccd47667a5..a09273fc8049 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -84,7 +84,6 @@ def _moe_forward( # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() runner = layer.runner - router_logits = runner._maybe_gate(hidden_states, router_logits) with runner._sequence_parallel_context(): if runner.use_dp_chunking: return runner.forward_impl_chunked( @@ -121,7 +120,6 @@ def _moe_forward_shared( # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() runner = layer.runner - router_logits = runner._maybe_gate(hidden_states, router_logits) with runner._sequence_parallel_context(): if runner.use_dp_chunking: return runner.forward_impl_chunked( @@ -274,8 +272,7 @@ def _maybe_setup_shared_experts_stream( self, hidden_states: torch.Tensor, shared_input: torch.Tensor | None, - ) -> torch.Tensor | None: - shared_experts_input: torch.Tensor | None = None + ): if self.use_shared_experts_stream: assert self.shared_experts_stream is not None assert self.moe_config.disable_inplace @@ -298,8 +295,6 @@ def _maybe_setup_shared_experts_stream( assert self.shared_experts_stream is not None self.shared_experts_stream.wait_stream(current_stream()) - return shared_experts_input - def _maybe_init_dp_chunking(self): if not self.use_dp_chunking: return @@ -486,20 +481,12 @@ def _apply_quant_method( shared_input: torch.Tensor | None, run_shared_experts_before: bool = True, ) -> tuple[torch.Tensor | None, torch.Tensor]: + shared_input = shared_input if shared_input is not None else hidden_states shared_output: torch.Tensor | None = None # Run this before quant_method to avoid inplace issues. if run_shared_experts_before: - shared_input = shared_input if shared_input is not None else hidden_states - shared_output = self._apply_shared_experts( - shared_input, - False, - ) - else: - hidden_states_clone = self._maybe_setup_shared_experts_stream( - hidden_states, - shared_input, - ) + shared_output = self._apply_shared_experts(shared_input, False) if self.quant_method.is_monolithic: result = self.quant_method.apply_monolithic( @@ -529,10 +516,7 @@ def _apply_quant_method( if not run_shared_experts_before and self.has_separate_shared_experts: assert shared_output is None - shared_output = self._apply_shared_experts( - hidden_states_clone, - True, - ) + shared_output = self._apply_shared_experts(shared_input, True) return shared_output, hidden_states @@ -704,6 +688,10 @@ def forward_impl_chunked( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # Gate overlap not supported when chunking is enabled. Run the + # gate first. + router_logits = self._maybe_gate(hidden_states, router_logits) + final_shared_hidden_states, final_fused_hidden_states = ( self._allocate_dp_chunking_outputs(hidden_states, router_logits) ) @@ -805,6 +793,16 @@ def forward_impl( self.has_separate_shared_experts and not self.use_shared_experts_stream ) + # The shared experts stream must be set up before calling the gate so they + # can be overlapped. + if not run_shared_experts_before: + self._maybe_setup_shared_experts_stream( + hidden_states, + shared_input, + ) + + router_logits = self._maybe_gate(hidden_states, router_logits) + # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain From 76aff0a040ab96f92d144c8aea1006c4a2ec5656 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 4 Feb 2026 10:36:16 -0500 Subject: [PATCH 054/108] wip Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 13 ++-- .../fused_moe/runner/default_moe_runner.py | 64 ++++++++++++++++++- .../layers/fused_moe/shared_fused_moe.py | 41 ++---------- 3 files changed, 74 insertions(+), 44 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 2f704569209c..bf6816fc611e 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -329,6 +329,7 @@ def __init__( is_sequence_parallel=False, expert_mapping: list[tuple[str, str, int, str]] | None = None, n_shared_experts: int | None = None, + shared_experts: torch.nn.Module | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -661,6 +662,7 @@ def _init_runner(self): quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, + enable_eplb=self.enable_eplb, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -708,12 +710,13 @@ def maybe_init_modular_kernel(self) -> None: def shared_experts(self) -> torch.nn.Module | None: return self._shared_experts if self.use_overlapped else None - @property - def layer_id(self): - # Delayed import to avoid circular dependency - from vllm.model_executor.models.utils import extract_layer_index + # TODO(bnell): is this needed? + #@property + #def layer_id(self): + # # Delayed import to avoid circular dependency + # from vllm.model_executor.models.utils import extract_layer_index - return extract_layer_index(self.layer_name) + # return extract_layer_index(self.layer_name) @property def gate(self) -> torch.nn.Module | None: diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index a09273fc8049..9703729d93cf 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -11,6 +11,7 @@ from vllm.distributed import ( get_ep_group, get_pcp_group, + get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( @@ -137,6 +138,42 @@ def _moe_forward_shared( ) +def _moe_forward_shared_new( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: str, +) -> tuple[torch.Tensor, torch.Tensor]: + def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: + # Reduce shared expert outputs if necessary, since the MLP + # should have been created with reduce_results=False. + if ( + self.reduce_results + and get_tensor_model_parallel_world_size() > 1 + and self.must_reduce_shared_expert_outputs() + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out + + if False and not self.overlap_shared_experts: + shared_out = self._shared_experts(hidden_states) + # XXXXXXXXXXXx + shared_experts = None + + router_logits = self._maybe_gate(hidden_states, router_logits) + with self._sequence_parallel_context(): + if self.use_dp_chunking: + shared_out, fused_out = self.forward_impl_chunked( + layer, hidden_states, router_logits + ) + else: + shared_out = fused_out = self.forward_impl( + layer, hidden_states, router_logits + ) + + return reduce_shared_out(shared_out), fused_out + + def _moe_forward_shared_fake( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -207,6 +244,7 @@ def __init__( quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, + enable_eplb: bool, ): super().__init__() self.moe_config = moe_config @@ -217,9 +255,12 @@ def __init__( self.quant_method = quant_method self.reduce_results = reduce_results self.enable_dbo = enable_dbo + self.enable_eplb = enable_eplb # Chunked all2all staging tensor - # TODO(bnell) rename these? + # TODO rename these + # These need to exist ahead of time due to CUDAgraph construction + # needing a fixed buffer address. self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() @@ -268,6 +309,22 @@ def use_dp_chunking(self) -> bool: or self.moe_config.moe_parallel_config.use_nixl_ep_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + # TODO(bnell): better name + @property + def overlap_shared_experts( + self, + shared_experts: torch.nn.Module | None, + ) -> bool: + # Disable shared expert overlap if: + # - we are using eplb with non-default backend, because of correctness issues + # - we are using flashinfer with DP, since there nothing to gain + # - we are using marlin kernels + backend = self.moe_config.moe_parallel_config.all2all_backend + return shared_experts is not None and not ( + (self.enable_eplb and backend != "allgather_reducescatter") + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) + def _maybe_setup_shared_experts_stream( self, hidden_states: torch.Tensor, @@ -476,6 +533,7 @@ def _maybe_pad_hidden_states( def _apply_quant_method( self, layer: torch.nn.Module, + shared_experts: torch.nn.Module | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, @@ -514,7 +572,9 @@ def _apply_quant_method( else: hidden_states = result - if not run_shared_experts_before and self.has_separate_shared_experts: + if not run_shared_experts_before and self.has_separate_shared_experts( + shared_experts + ): assert shared_output is None shared_output = self._apply_shared_experts(shared_input, True) diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index 37336df17561..251c05300331 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -3,10 +3,6 @@ import torch -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) from vllm.model_executor.layers.fused_moe.layer import FusedMoE @@ -23,36 +19,7 @@ def forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - if not self.use_overlapped: - if self._shared_experts is not None: - shared_out = self._shared_experts(hidden_states) - - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - self.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - else: - shared_out = None - - fused_out = super().forward( - hidden_states=hidden_states, - router_logits=router_logits, - ) - else: - shared_out, fused_out = super().forward( - hidden_states=hidden_states, - router_logits=router_logits, - ) - # ensure early TP reduction of shared expert outputs when required - if ( - shared_out is not None - and self.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out, fused_out + return super().forward( + hidden_states=hidden_states, + router_logits=router_logits, + ) From 4fab915f53e026f3ca53267efd953f4c0a2b826c Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 9 Feb 2026 18:27:46 -0500 Subject: [PATCH 055/108] fix Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index bf6816fc611e..c992930c8859 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -329,7 +329,6 @@ def __init__( is_sequence_parallel=False, expert_mapping: list[tuple[str, str, int, str]] | None = None, n_shared_experts: int | None = None, - shared_experts: torch.nn.Module | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -711,8 +710,8 @@ def shared_experts(self) -> torch.nn.Module | None: return self._shared_experts if self.use_overlapped else None # TODO(bnell): is this needed? - #@property - #def layer_id(self): + # @property + # def layer_id(self): # # Delayed import to avoid circular dependency # from vllm.model_executor.models.utils import extract_layer_index From d8a7f910a0f40b645a858bd22eeca85a11b03bb8 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 11 Feb 2026 14:36:30 -0500 Subject: [PATCH 056/108] WIP DOUBLE CHECK THIS Signed-off-by: Bill Nell --- .../fused_moe/router/fused_moe_router.py | 8 +++ .../fused_moe/router/memoizing_router.py | 35 +++++++++++++ .../fused_moe/runner/default_moe_runner.py | 50 +++++++++---------- .../layers/fused_moe/shared_fused_moe.py | 2 +- 4 files changed, 69 insertions(+), 26 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py diff --git a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py index c322a8cd4cd6..d7aed4fdeb2b 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections.abc import Callable import torch @@ -13,6 +14,13 @@ class FusedMoERouter(ABC): method that is used for routing hidden states based on router logits. """ + @abstractmethod + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + raise NotImplementedError + @property @abstractmethod def routing_method_type(self) -> RoutingMethodType: diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py new file mode 100644 index 000000000000..a55bd2f09d6f --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter + + +class MemoizingRouter(FusedMoERouter): + def __init__(self, router: FusedMoERouter): + self.router = router + + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + self.router.set_capture_fn(capture_fn) + self.results: tuple[torch.Tensor, torch.Tensor] | None = None + + @property + def routing_method_type(self) -> RoutingMethodType: + return self.router.routing_method_type + + def select_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.results is None: + self.results = self.router.select_experts(hidden_states, router_logits) + return self.results diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 9703729d93cf..d7f3284ee197 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -144,30 +144,33 @@ def _moe_forward_shared_new( shared_experts_input: torch.Tensor | None, layer_name: str, ) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(layer_name) + runner = layer.runner + def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: # Reduce shared expert outputs if necessary, since the MLP # should have been created with reduce_results=False. if ( - self.reduce_results + runner.reduce_results and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() + and runner.must_reduce_shared_expert_outputs() ): shared_out = tensor_model_parallel_all_reduce(shared_out) return shared_out - if False and not self.overlap_shared_experts: - shared_out = self._shared_experts(hidden_states) - # XXXXXXXXXXXx + if not runner.overlap_shared_experts: + shared_out = runner._shared_experts(hidden_states) + # XXXXXXXXXXXXXXXX shared_experts = None - router_logits = self._maybe_gate(hidden_states, router_logits) - with self._sequence_parallel_context(): - if self.use_dp_chunking: - shared_out, fused_out = self.forward_impl_chunked( + router_logits = runner._maybe_gate(hidden_states, router_logits) + with runner._sequence_parallel_context(): + if runner.use_dp_chunking: + shared_out, fused_out = runner.forward_impl_chunked( layer, hidden_states, router_logits ) else: - shared_out = fused_out = self.forward_impl( + shared_out = fused_out = runner.forward_impl( layer, hidden_states, router_logits ) @@ -282,6 +285,13 @@ def __init__( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) + self.allow_shared_experts_stream = ( + current_platform.is_cuda() + and self.has_separate_shared_experts + and not self.use_dp_chunking + and self.shared_experts_stream is not None + ) + # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name @@ -313,14 +323,13 @@ def use_dp_chunking(self) -> bool: @property def overlap_shared_experts( self, - shared_experts: torch.nn.Module | None, ) -> bool: # Disable shared expert overlap if: # - we are using eplb with non-default backend, because of correctness issues # - we are using flashinfer with DP, since there nothing to gain # - we are using marlin kernels backend = self.moe_config.moe_parallel_config.all2all_backend - return shared_experts is not None and not ( + return self.shared_experts is not None and not ( (self.enable_eplb and backend != "allgather_reducescatter") or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels ) @@ -533,7 +542,6 @@ def _maybe_pad_hidden_states( def _apply_quant_method( self, layer: torch.nn.Module, - shared_experts: torch.nn.Module | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, @@ -572,9 +580,7 @@ def _apply_quant_method( else: hidden_states = result - if not run_shared_experts_before and self.has_separate_shared_experts( - shared_experts - ): + if not run_shared_experts_before and self.has_separate_shared_experts: assert shared_output is None shared_output = self._apply_shared_experts(shared_input, True) @@ -836,15 +842,9 @@ def forward_impl( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - self.use_shared_experts_stream = ( - current_platform.is_cuda() - and self.has_separate_shared_experts - and not self.use_dp_chunking - and self.shared_experts_stream is not None - and ( - hidden_states.shape[0] - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) + # TODO(bnell): split this into runtime vs. static parts? + self.use_shared_experts_stream = self.allow_shared_experts_stream and ( + hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) # Check if we need to run shared experts before matrix multiply because diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index 251c05300331..9aa9fe8e0a9a 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -6,7 +6,7 @@ from vllm.model_executor.layers.fused_moe.layer import FusedMoE -# TODO(bnell): Add shared + fused combo function? e.g. + +# TODO(bnell): Remove this entirely class SharedFusedMoE(FusedMoE): """ A FusedMoE operation that also computes the results of shared experts. From 3dec78f98e40bc7a7d734cf832038ccbe7e8727d Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Feb 2026 19:00:11 -0500 Subject: [PATCH 057/108] wip more refactoring Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 7 +- vllm/model_executor/layers/fused_moe/layer.py | 34 +-- .../fused_moe/runner/default_moe_runner.py | 270 ++++++++++++++---- 3 files changed, 231 insertions(+), 80 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 78876ef7c9b0..5aa8a56d4115 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,6 +150,7 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), + self.base_layer.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXXXXX ) if quant_config.use_mxfp4_w4a16: @@ -592,9 +593,9 @@ def forward(self, *args, **kwargs): def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - @property - def _shared_experts(self): - return self.base_layer._shared_experts + # @property + # def _shared_experts(self): + # return self.base_layer._shared_experts @property def quant_method(self): diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c992930c8859..c6c6dab9dd04 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -636,14 +636,14 @@ def _get_quant_method() -> FusedMoEMethodBase: # - we are using eplb with non-default backend, because of correctness issues # - we are using flashinfer with DP, since there nothing to gain # - we are using marlin kernels - backend = self.moe_parallel_config.all2all_backend - self.use_overlapped = ( - not ( - (self.enable_eplb and backend != "allgather_reducescatter") - or self.moe_parallel_config.use_fi_nvl_two_sided_kernels - ) - and self._shared_experts is not None - ) + # backend = self.moe_parallel_config.all2all_backend + # self.use_overlapped = ( + # not ( + # (self.enable_eplb and backend != "allgather_reducescatter") + # or self.moe_parallel_config.use_fi_all2allv_kernels + # ) + # and self._shared_experts is not None + # ) self.runner = self._init_runner() @@ -656,8 +656,8 @@ def _init_runner(self): moe_config=self.moe_config, router=self.router, routed_input_transform=self._routed_input_transform, - gate=self.gate, - shared_experts=self.shared_experts, + gate=self._gate, + shared_experts=self._shared_experts, quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, @@ -700,14 +700,14 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.shared_experts, + self.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXX inplace=not self.moe_config.disable_inplace, ) ) - @property - def shared_experts(self) -> torch.nn.Module | None: - return self._shared_experts if self.use_overlapped else None + # @property + # def shared_experts(self) -> torch.nn.Module | None: + # return self._shared_experts if self.use_overlapped else None # TODO(bnell): is this needed? # @property @@ -717,9 +717,9 @@ def shared_experts(self) -> torch.nn.Module | None: # return extract_layer_index(self.layer_name) - @property - def gate(self) -> torch.nn.Module | None: - return self._gate if self.use_overlapped else None + # @property + # def gate(self) -> torch.nn.Module | None: + # return self._gate if self.use_overlapped else None @property def tp_size(self): diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index d7f3284ee197..3bfefcb442fe 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -3,6 +3,7 @@ from collections.abc import Callable from contextlib import nullcontext from typing import TYPE_CHECKING +from enum import IntEnum import torch import torch.nn.functional as F @@ -44,6 +45,23 @@ logger = init_logger(__name__) +class SharedExpertsOrder(IntEnum): + # No shared experts. + NONE = (0,) + + # Get rid of this one? combine with BEFORE? + EXTERNAL = (1,) + + # Called by modular kernel. + INTERNAL = (2,) + + # Called right before quant_method is executed. + BEFORE_QUANT_METHOD = (3,) + + # Called right after quant_method is executed (possibly with streaming). + AFTER_QUANT_METHOD = (4,) + + def get_layer_from_name(layer_name: str) -> torch.nn.Module: forward_context: ForwardContext = get_forward_context() if layer_name == "from_forward_context": @@ -81,6 +99,7 @@ def _moe_forward( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: +<<<<<<< HEAD layer = get_layer_from_name(_resolve_layer_name(layer_name)) # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -100,6 +119,15 @@ def _moe_forward( router_logits, shared_experts_input, ) +======= + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) +>>>>>>> d11342756e (wip more refactoring) def _moe_forward_fake( @@ -117,6 +145,7 @@ def _moe_forward_shared( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: +<<<<<<< HEAD layer = get_layer_from_name(_resolve_layer_name(layer_name)) # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() @@ -175,6 +204,15 @@ def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: ) return reduce_shared_out(shared_out), fused_out +======= + layer = get_layer_from_name(layer_name) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) +>>>>>>> d11342756e (wip more refactoring) def _moe_forward_shared_fake( @@ -264,6 +302,7 @@ def __init__( # TODO rename these # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. + # TODO: these could be global, i.e. shared by all layers self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() @@ -285,17 +324,69 @@ def __init__( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) - self.allow_shared_experts_stream = ( + self.use_dp_chunking = ( + self.moe_config.moe_parallel_config.use_pplx_kernels + or self.moe_config.moe_parallel_config.use_deepep_ll_kernels + or self.moe_config.moe_parallel_config.use_mori_kernels + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + + # Needed for string -> FusedMoE layer lookup in custom ops. + self.layer_name = layer.layer_name + + self.moe_forward = self._select_forward(layer) + + def _has_external_experts(self) -> bool: + # Disable shared expert overlap if: + # - we are using eplb with non-default backend, because of correctness issues + # - we are using flashinfer with DP, since there nothing to gain + backend = self.moe_config.moe_parallel_config.all2all_backend + return self.shared_experts is not None and not ( + (self.enable_eplb and backend != "allgather_reducescatter") + or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) + + def _determine_shared_experts_order( + self, + hidden_states: torch.Tensor, + ) -> tuple[SharedExpertsOrder, bool]: + if self.shared_experts is None: + return SharedExpertsOrder.NONE, False + + if self._has_external_experts(): + return SharedExpertsOrder.EXTERNAL, False + + has_separate_shared_experts = not self.quant_method.mk_owns_shared_expert + + if not has_separate_shared_experts: + return SharedExpertsOrder.INTERNAL, False + + allow_shared_experts_stream = ( current_platform.is_cuda() - and self.has_separate_shared_experts + and has_separate_shared_experts and not self.use_dp_chunking and self.shared_experts_stream is not None + and hidden_states.shape[0] + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) - # Needed for string -> FusedMoE layer lookup in custom ops. - self.layer_name = layer.layer_name + # Check if we need to run shared experts before matrix multiply because + # matrix multiply may modify the hidden_states. + run_shared_experts_before = ( + has_separate_shared_experts and not allow_shared_experts_stream + ) - self.moe_forward = self._select_forward(layer) + if run_shared_experts_before: + return SharedExpertsOrder.BEFORE_QUANT_METHOD, False + else: + return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream + + # XXXXXX used by layer.py and lora/layers/fused_moe.py + def get_shared_experts(self) -> torch.nn.Module | None: + if self._has_external_experts(): + return self.shared_experts + else: + return None def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -310,6 +401,7 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) +<<<<<<< HEAD @property def use_dp_chunking(self) -> bool: return ( @@ -391,35 +483,39 @@ def _maybe_init_dp_chunking(self): device=device, ) - @property - def has_separate_shared_experts(self) -> bool: - return ( - not self.quant_method.mk_owns_shared_expert - and self.shared_experts is not None - ) - - def _apply_shared_experts( + def _maybe_apply_shared_experts( self, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, - allow_streaming: bool = False, + shared_input: torch.Tensor | None, + order: SharedExpertsOrder, ) -> torch.Tensor | None: - shared_output: torch.Tensor | None = None - if self.has_separate_shared_experts: - assert self.shared_experts is not None - - if self.use_shared_experts_stream and allow_streaming: - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self.shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - shared_output = self.shared_experts(hidden_states) - current_stream().wait_stream(self.shared_experts_stream) - else: + experts_order, use_shared_experts_stream = self._determine_shared_experts_order( + hidden_states + ) + if order != experts_order: + return None + + assert self.shared_experts is not None + assert shared_output is None + if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: + hidden_states = self._maybe_setup_shared_experts_stream( + hidden_states, + shared_input, + ) + + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self.shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream shared_output = self.shared_experts(hidden_states) + current_stream().wait_stream(self.shared_experts_stream) + else: + shared_output = self.shared_experts(hidden_states) return shared_output @@ -470,6 +566,18 @@ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Ten return result return hidden_states + # TODO: combine with runner.forward_impl_X + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: + # Reduce shared expert outputs if necessary, since the MLP + # should have been created with reduce_results=False. + if ( + self.reduce_results + and get_tensor_model_parallel_world_size() > 1 + and self.must_reduce_shared_expert_outputs() + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out + def _maybe_reduce_output( self, states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -518,7 +626,7 @@ def _maybe_pad_hidden_states( original_hidden_dim = ( original_hidden_states.shape[-1] if original_hidden_states is not None - else 0 + else hidden_states.shape[-1] ) transformed_hidden_dim = hidden_states.shape[-1] if ( @@ -542,20 +650,22 @@ def _maybe_pad_hidden_states( def _apply_quant_method( self, layer: torch.nn.Module, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, - run_shared_experts_before: bool = True, ) -> tuple[torch.Tensor | None, torch.Tensor]: - shared_input = shared_input if shared_input is not None else hidden_states - shared_output: torch.Tensor | None = None - # Run this before quant_method to avoid inplace issues. - if run_shared_experts_before: - shared_output = self._apply_shared_experts(shared_input, False) + shared_input = shared_input if shared_input is not None else hidden_states + shared_output = self._maybe_apply_shared_experts( + shared_output, + shared_input, + shared_input, + SharedExpertsOrder.BEFORE_QUANT_METHOD, + ) if self.quant_method.is_monolithic: - result = self.quant_method.apply_monolithic( + fused_out = self.quant_method.apply_monolithic( layer=layer, x=hidden_states, router_logits=router_logits, @@ -566,7 +676,7 @@ def _apply_quant_method( router_logits=router_logits, ) - result = self.quant_method.apply( + fused_out = self.quant_method.apply( layer=layer, x=hidden_states, topk_weights=topk_weights, @@ -574,15 +684,23 @@ def _apply_quant_method( shared_experts_input=shared_input, ) - if isinstance(result, tuple): - assert shared_output is None - shared_output, hidden_states = result + # TODO: try to clean this up + if isinstance(fused_out, tuple): + if fused_out[0] is None: + assert shared_output is not None + hidden_states = fused_out[1] + else: + assert shared_output is None + shared_output, hidden_states = fused_out else: - hidden_states = result + hidden_states = fused_out - if not run_shared_experts_before and self.has_separate_shared_experts: - assert shared_output is None - shared_output = self._apply_shared_experts(shared_input, True) + shared_output = self._maybe_apply_shared_experts( + shared_output, + shared_input, + shared_input, + SharedExpertsOrder.AFTER_QUANT_METHOD, + ) return shared_output, hidden_states @@ -705,10 +823,10 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # For latent MoE: save ORIGINAL hidden_states before transform # (shared_experts need original dimension, routed experts use transformed) - if self.shared_experts is not None: - original_hidden_states = hidden_states - else: - original_hidden_states = None + original_hidden_states = ( + # or shared_experts is not None? + hidden_states if self.routed_input_transform is not None else None + ) # Apply transform for routed experts (e.g., latent projection for latent MoE) hidden_states = self.apply_routed_input_transform(hidden_states) @@ -718,6 +836,7 @@ def forward( hidden_states, ) + # Note: moe_forward will call forward_dispatch fused_output = self.moe_forward( hidden_states, router_logits, @@ -725,8 +844,42 @@ def forward( self._encode_layer_name(), ) + # TODO: figure out how to combine this with maybe_reduce_output + if self.shared_experts is not None: + fused_output = ( + self._maybe_reduce_shared_out(fused_output[0]), + fused_output[1], + ) + return self._maybe_reduce_output(fused_output, og_hidden_dims) + def forward_dispatch( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # TODO(bnell): this can be removed after MK migration is complete. + layer.ensure_moe_quant_config_init() + + router_logits = self._maybe_gate(hidden_states, router_logits) + with self._sequence_parallel_context(): + if self.use_dp_chunking: + return self.forward_impl_chunked( + layer, + hidden_states, + router_logits, + shared_input, + ) + else: + return self.forward_impl( + layer, + hidden_states, + router_logits, + shared_input, + ) + def _slice_and_copy_input( self, out_slice: torch.Tensor, @@ -811,6 +964,7 @@ def forward_impl_chunked( shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, + shared_output=None, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, shared_input=shared_input_chunk, @@ -842,15 +996,11 @@ def forward_impl( router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # TODO(bnell): split this into runtime vs. static parts? - self.use_shared_experts_stream = self.allow_shared_experts_stream and ( - hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) - - # Check if we need to run shared experts before matrix multiply because - # matrix multiply may modify the hidden_states. - run_shared_experts_before = ( - self.has_separate_shared_experts and not self.use_shared_experts_stream + shared_output = self._maybe_apply_shared_experts( + None, + hidden_states, + shared_input, + SharedExpertsOrder.EXTERNAL, ) # The shared experts stream must be set up before calling the gate so they @@ -874,10 +1024,10 @@ def forward_impl( shared_output, hidden_states = self._apply_quant_method( layer=layer, + shared_output=shared_output, hidden_states=hidden_states, router_logits=router_logits, shared_input=shared_input, - run_shared_experts_before=run_shared_experts_before, ) return self._maybe_combine( From e94b8632fbb467c8dc277da894c71ec180b4952a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Feb 2026 19:20:36 -0500 Subject: [PATCH 058/108] wip Signed-off-by: Bill Nell --- .../fused_moe/runner/default_moe_runner.py | 165 +++--------------- 1 file changed, 29 insertions(+), 136 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 3bfefcb442fe..b0b59bbb0c1e 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -99,27 +99,6 @@ def _moe_forward( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: -<<<<<<< HEAD - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - runner = layer.runner - with runner._sequence_parallel_context(): - if runner.use_dp_chunking: - return runner.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - else: - return runner.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) -======= layer = get_layer_from_name(layer_name) return layer.runner.forward_dispatch( layer, @@ -127,7 +106,6 @@ def _moe_forward( router_logits, shared_experts_input, ) ->>>>>>> d11342756e (wip more refactoring) def _moe_forward_fake( @@ -145,66 +123,6 @@ def _moe_forward_shared( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: -<<<<<<< HEAD - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - runner = layer.runner - with runner._sequence_parallel_context(): - if runner.use_dp_chunking: - return runner.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - else: - return runner.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_shared_new( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: str, -) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(layer_name) - runner = layer.runner - - def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - runner.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and runner.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out - - if not runner.overlap_shared_experts: - shared_out = runner._shared_experts(hidden_states) - # XXXXXXXXXXXXXXXX - shared_experts = None - - router_logits = runner._maybe_gate(hidden_states, router_logits) - with runner._sequence_parallel_context(): - if runner.use_dp_chunking: - shared_out, fused_out = runner.forward_impl_chunked( - layer, hidden_states, router_logits - ) - else: - shared_out = fused_out = runner.forward_impl( - layer, hidden_states, router_logits - ) - - return reduce_shared_out(shared_out), fused_out -======= layer = get_layer_from_name(layer_name) return layer.runner.forward_dispatch( layer, @@ -212,7 +130,6 @@ def reduce_shared_out(shared_out: torch.Tensor) -> torch.Tensor: router_logits, shared_experts_input, ) ->>>>>>> d11342756e (wip more refactoring) def _moe_forward_shared_fake( @@ -328,7 +245,8 @@ def __init__( self.moe_config.moe_parallel_config.use_pplx_kernels or self.moe_config.moe_parallel_config.use_deepep_ll_kernels or self.moe_config.moe_parallel_config.use_mori_kernels - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels + or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels + or self.moe_config.moe_parallel_config.use_nixl_ep_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK # Needed for string -> FusedMoE layer lookup in custom ops. @@ -401,32 +319,7 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) -<<<<<<< HEAD - @property - def use_dp_chunking(self) -> bool: - return ( - self.moe_config.moe_parallel_config.use_deepep_ll_kernels - or self.moe_config.moe_parallel_config.use_mori_kernels - or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels - or self.moe_config.moe_parallel_config.use_nixl_ep_kernels - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK - - # TODO(bnell): better name - @property - def overlap_shared_experts( - self, - ) -> bool: - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - # - we are using marlin kernels - backend = self.moe_config.moe_parallel_config.all2all_backend - return self.shared_experts is not None and not ( - (self.enable_eplb and backend != "allgather_reducescatter") - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels - ) - - def _maybe_setup_shared_experts_stream( + def _setup_shared_experts_stream( self, hidden_states: torch.Tensor, shared_input: torch.Tensor | None, @@ -499,7 +392,8 @@ def _maybe_apply_shared_experts( assert self.shared_experts is not None assert shared_output is None if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - hidden_states = self._maybe_setup_shared_experts_stream( + # TODO: fold this in? + hidden_states = self._setup_shared_experts_stream( hidden_states, shared_input, ) @@ -517,6 +411,12 @@ def _maybe_apply_shared_experts( else: shared_output = self.shared_experts(hidden_states) + if order == SharedExpertsOrder.EXTERNAL: + # TODO: figure out how to combine this with maybe_reduce_output? + # or get rid of it completely..... + assert shared_output is not None + shared_output = self._maybe_reduce_shared_out(shared_output) + return shared_output def must_reduce_shared_expert_outputs(self) -> bool: @@ -816,6 +716,11 @@ def _maybe_combine( else: return hidden_states + # + # forward + # - self.moe_forward (_moe_forward or _moe_forward_shared) + # - forward_chunking_dispatch + # - forward_impl or forward_impl_chunked def forward( self, hidden_states: torch.Tensor, @@ -844,13 +749,6 @@ def forward( self._encode_layer_name(), ) - # TODO: figure out how to combine this with maybe_reduce_output - if self.shared_experts is not None: - fused_output = ( - self._maybe_reduce_shared_out(fused_output[0]), - fused_output[1], - ) - return self._maybe_reduce_output(fused_output, og_hidden_dims) def forward_dispatch( @@ -864,10 +762,19 @@ def forward_dispatch( layer.ensure_moe_quant_config_init() router_logits = self._maybe_gate(hidden_states, router_logits) + + shared_output = self._maybe_apply_shared_experts( + None, + hidden_states, + shared_input, + SharedExpertsOrder.EXTERNAL, + ) + with self._sequence_parallel_context(): if self.use_dp_chunking: return self.forward_impl_chunked( layer, + shared_output, hidden_states, router_logits, shared_input, @@ -875,6 +782,7 @@ def forward_dispatch( else: return self.forward_impl( layer, + shared_output, hidden_states, router_logits, shared_input, @@ -903,6 +811,7 @@ def _slice_and_copy_input( def forward_impl_chunked( self, layer: torch.nn.Module, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, @@ -964,7 +873,7 @@ def forward_impl_chunked( shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, - shared_output=None, + shared_output=shared_output, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, shared_input=shared_input_chunk, @@ -992,27 +901,11 @@ def forward_impl_chunked( def forward_impl( self, layer: torch.nn.Module, + shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - shared_output = self._maybe_apply_shared_experts( - None, - hidden_states, - shared_input, - SharedExpertsOrder.EXTERNAL, - ) - - # The shared experts stream must be set up before calling the gate so they - # can be overlapped. - if not run_shared_experts_before: - self._maybe_setup_shared_experts_stream( - hidden_states, - shared_input, - ) - - router_logits = self._maybe_gate(hidden_states, router_logits) - # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain From 6cc50741324a3da05fd797605f7aa5055be82151 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Feb 2026 16:49:41 -0500 Subject: [PATCH 059/108] SharedExperts wip Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 6 +- .../layers/fused_moe/cutlass_moe.py | 2 + .../layers/fused_moe/fused_moe_method_base.py | 4 +- .../fused_moe/fused_moe_modular_method.py | 10 +- vllm/model_executor/layers/fused_moe/layer.py | 65 ++-- .../layers/fused_moe/modular_kernel.py | 58 ++-- .../layers/fused_moe/oracle/fp8.py | 9 +- .../layers/fused_moe/oracle/nvfp4.py | 9 +- .../layers/fused_moe/oracle/unquantized.py | 4 + .../fused_moe/runner/default_moe_runner.py | 298 ++++-------------- .../layers/fused_moe/runner/shared_experts.py | 207 ++++++++++++ .../fused_moe/unquantized_fused_moe_method.py | 10 +- .../layers/quantization/awq_marlin.py | 2 +- .../layers/quantization/bitsandbytes.py | 2 +- .../compressed_tensors_moe.py | 20 +- .../layers/quantization/experts_int8.py | 2 +- .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/gguf.py | 2 +- .../layers/quantization/gptq_marlin.py | 2 +- .../layers/quantization/modelopt.py | 8 +- .../layers/quantization/moe_wna16.py | 2 +- .../layers/quantization/mxfp4.py | 4 +- .../layers/quantization/quark/quark_moe.py | 6 +- 23 files changed, 391 insertions(+), 345 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/runner/shared_experts.py diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 5aa8a56d4115..b6c62b618a09 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,9 +150,11 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), - self.base_layer.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXXXXX + self.base_layer.shared_experts, ) + # TODO: could be incorrect due to monolithic kernel? or add assert it + # is modular? if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.impl.fused_experts, @@ -170,6 +172,7 @@ def wrapper(*args, **kwargs): moe_state_dict["apply_router_weight_on_input"] = kwargs[ "apply_router_weight_on_input" ] + # TODO: global_num_experts/shared_experts_input? result = func(*args, **kwargs) return result @@ -339,6 +342,7 @@ def wrapper(*args, **kwargs): fused_experts = m_fused_moe_fn.impl.fused_experts + # TODO: seems like this could be done with modular kernel subclasses? m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 75ee776646ba..43082b3675a6 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -1194,6 +1194,8 @@ def cutlass_moe_w4a8_fp8( quant_config=quant_config, group_size=group_size, ), + shared_experts=None, + inplace=False, ) return fn.apply( diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py index f6a303e7988e..9350a9c7de74 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py @@ -130,7 +130,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: raise NotImplementedError def apply_monolithic( @@ -138,5 +138,5 @@ def apply_monolithic( layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 0065c11f3163..142e180786c6 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -16,6 +16,9 @@ FusedMoEKernel, FusedMoEPrepareAndFinalizeModular, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) logger = init_logger(__name__) @@ -44,7 +47,7 @@ def make( moe_layer: torch.nn.Module, old_quant_method: FusedMoEMethodBase, prepare_finalize: FusedMoEPrepareAndFinalizeModular, - shared_experts: torch.nn.Module | None, + shared_experts: SharedExperts | None, inplace: bool = False, ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( @@ -52,8 +55,7 @@ def make( FusedMoEKernel( prepare_finalize, old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), - shared_experts, - moe_parallel_config=moe_layer.moe_parallel_config, + shared_experts=shared_experts, inplace=inplace, ), ) @@ -89,7 +91,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c6c6dab9dd04..bbf6e2ffedb9 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -42,6 +42,9 @@ from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( DefaultMoERunner, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) @@ -337,7 +340,6 @@ def __init__( super().__init__() self._gate = gate - self._shared_experts = shared_experts self._routed_input_transform = routed_input_transform if params_dtype is None: @@ -565,7 +567,7 @@ def __init__( device=vllm_config.device_config.device, routing_method=self.routing_method_type, # TODO: in_dtype == out_dtype? - disable_inplace=disable_inplace() or self._shared_experts is not None, + disable_inplace=disable_inplace() or self.shared_experts is not None, ) if self.moe_config.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( @@ -632,22 +634,30 @@ def _get_quant_method() -> FusedMoEMethodBase: self.quant_method.create_weights(layer=self, **moe_quant_params) self.base_quant_method = self.quant_method - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - # - we are using marlin kernels - # backend = self.moe_parallel_config.all2all_backend - # self.use_overlapped = ( - # not ( - # (self.enable_eplb and backend != "allgather_reducescatter") - # or self.moe_parallel_config.use_fi_all2allv_kernels - # ) - # and self._shared_experts is not None - # ) - + self._shared_experts = shared_experts + self.shared_experts = self._init_shared_experts() self.runner = self._init_runner() - def _init_runner(self): + def _init_shared_experts(self) -> SharedExperts | None: + if self._shared_experts is None: + return None + + reduce_shared_output = ( + self.reduce_results + # XXXX ordering issue + and self.quant_method.moe_mk is not None + and self.quant_method.moe_mk.output_is_reduced() + ) + + return SharedExperts( + self._shared_experts, + moe_config=self.moe_config, + has_separate_shared_experts=not self.quant_method.mk_owns_shared_expert, + use_dp_chunking=self.use_dp_chunking, # XXXXXXXXXXXXXX + must_reduce_shared_expert_outputs=reduce_shared_output, + ) + + def _init_runner(self) -> DefaultMoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. @@ -657,7 +667,7 @@ def _init_runner(self): router=self.router, routed_input_transform=self._routed_input_transform, gate=self._gate, - shared_experts=self._shared_experts, + shared_experts=self.shared_experts, quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, @@ -700,26 +710,17 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.runner.get_shared_experts(), # XXXXXXXXXXXXXXXXX + self.runner._get_shared_experts(), # XXXXXXXXXXXXXXXXX inplace=not self.moe_config.disable_inplace, ) ) - # @property - # def shared_experts(self) -> torch.nn.Module | None: - # return self._shared_experts if self.use_overlapped else None - - # TODO(bnell): is this needed? - # @property - # def layer_id(self): - # # Delayed import to avoid circular dependency - # from vllm.model_executor.models.utils import extract_layer_index - - # return extract_layer_index(self.layer_name) + @property + def layer_id(self): + # Delayed import to avoid circular dependency + from vllm.model_executor.models.utils import extract_layer_index - # @property - # def gate(self) -> torch.nn.Module | None: - # return self._gate if self.use_overlapped else None + return extract_layer_index(self.layer_name) @property def tp_size(self): diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index a6b498834017..bbafc682c102 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -20,6 +20,10 @@ FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, disable_inplace, @@ -234,6 +238,13 @@ def output_is_reduced(self) -> bool: """ raise NotImplementedError + def supports_async(self) -> bool: + """ + Indicates whether or not this class implements prepare_async and + finalize_async. + """ + return False + # TODO: pass FusedMoEParallelConfig in as ctor parameter? class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): @@ -280,13 +291,6 @@ def prepare( """ raise NotImplementedError - def supports_async(self) -> bool: - """ - Indicates whether or not this class implements prepare_async and - finalize_async. - """ - return False - def prepare_async( self, a1: torch.Tensor, @@ -991,8 +995,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, + shared_experts: SharedExperts | None, inplace: bool = False, ): self.prepare_finalize = prepare_finalize @@ -1000,6 +1003,9 @@ def __init__( self.shared_experts = shared_experts self.moe_parallel_config = moe_parallel_config self.inplace = inplace + + moe_parallel_config = fused_experts.moe_config.moe_parallel_config + self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config self.is_dp_ep = ( moe_parallel_config is not None and moe_parallel_config.dp_size > 1 @@ -1069,6 +1075,17 @@ def _allocate_buffers( return workspace13, workspace2, fused_out + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + ): + if self.shared_experts is not None: + assert shared_experts_input is not None + self.shared_experts.apply( + shared_experts_input, + SharedExpertsOrder.INTERNAL, + ) + def _prepare( self, hidden_states: torch.Tensor, @@ -1241,15 +1258,6 @@ def _finalize( shared_experts_input is the original hidden_states (full dimension) needed by the shared expert MLP. """ - shared_output: torch.Tensor | None = None - - # For latent MoE: shared experts need the original hidden_states - # (full hidden_size), not the latent-projected version used by - # routed experts. - se_hidden_states = ( - shared_experts_input if shared_experts_input is not None else hidden_states - ) - if not self.prepare_finalize.supports_async(): assert not dbo_enabled() @@ -1261,8 +1269,7 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - if self.shared_experts is not None: - shared_output = self.shared_experts(se_hidden_states) + self._maybe_apply_shared_experts(shared_experts_input) else: finalize_ret = self.prepare_finalize.finalize_async( output, @@ -1272,8 +1279,7 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - if self.shared_experts is not None: - shared_output = self.shared_experts(se_hidden_states) + self._maybe_apply_shared_experts(shared_experts_input) # TODO(lucas): refactor this in the alternative schedules followup # currently unpack if we have hook + receiver pair or just @@ -1296,11 +1302,7 @@ def _finalize( receiver() - if self.shared_experts is None: - return output - else: - assert shared_output is not None - return shared_output, output + return output def apply( self, @@ -1314,7 +1316,7 @@ def apply( expert_map: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, shared_experts_input: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """ This function computes a Mixture of Experts (MoE) layer using two sets of weights, w1 and w2, and top-k gating mechanism. diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index a63c02663886..57431b946e24 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -18,6 +18,9 @@ fp8_w8a8_moe_quant_config, fp8_w8a16_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, get_flashinfer_moe_backend, @@ -545,7 +548,7 @@ def make_fp8_moe_kernel( experts_cls: type[mk.FusedMoEExperts], fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - shared_experts: torch.nn.Module | None = None, + shared_experts: SharedExperts | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -583,10 +586,8 @@ def make_fp8_moe_kernel( experts, shared_experts=( shared_experts - if moe_config.moe_parallel_config.use_deepep_ll_kernels - else None + if prepare_finalize.supports_async() else None ), - moe_parallel_config=moe_config.moe_parallel_config, inplace=( not moe_config.disable_inplace and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 8a224cb39e7c..194513981e41 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -18,6 +18,9 @@ nvfp4_moe_quant_config, nvfp4_w4a16_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( prepare_nvfp4_moe_layer_for_fi_or_cutlass, ) @@ -397,7 +400,7 @@ def make_nvfp4_moe_kernel( moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - shared_experts: torch.nn.Module | None = None, + shared_experts: SharedExperts | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -435,10 +438,8 @@ def make_nvfp4_moe_kernel( experts, shared_experts=( shared_experts - if moe_config.moe_parallel_config.use_deepep_ll_kernels - else None + if prepare_finalize.supports_async() else None ), - moe_parallel_config=moe_config.moe_parallel_config, inplace=False, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 9c31da10dd94..926819592c3a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -224,6 +224,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=False, ) @@ -238,6 +239,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.TRITON: @@ -249,6 +251,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=not moe_config.disable_inplace, ) elif backend == UnquantizedMoeBackend.XPU: @@ -260,6 +263,7 @@ def make_unquantized_moe_kernel( moe_config=moe_config, quant_config=quant_config, ), + shared_experts=None, inplace=not moe_config.disable_inplace, ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index b0b59bbb0c1e..cc177b946cf1 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -3,7 +3,6 @@ from collections.abc import Callable from contextlib import nullcontext from typing import TYPE_CHECKING -from enum import IntEnum import torch import torch.nn.functional as F @@ -12,7 +11,6 @@ from vllm.distributed import ( get_ep_group, get_pcp_group, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( @@ -31,6 +29,10 @@ FusedMoERouter, ) from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( @@ -45,23 +47,6 @@ logger = init_logger(__name__) -class SharedExpertsOrder(IntEnum): - # No shared experts. - NONE = (0,) - - # Get rid of this one? combine with BEFORE? - EXTERNAL = (1,) - - # Called by modular kernel. - INTERNAL = (2,) - - # Called right before quant_method is executed. - BEFORE_QUANT_METHOD = (3,) - - # Called right after quant_method is executed (possibly with streaming). - AFTER_QUANT_METHOD = (4,) - - def get_layer_from_name(layer_name: str) -> torch.nn.Module: forward_context: ForwardContext = get_forward_context() if layer_name == "from_forward_context": @@ -154,7 +139,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], + mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -163,7 +148,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward_shared", op_func=_moe_forward_shared, - mutates_args=["hidden_states"], + mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_shared_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -198,7 +183,7 @@ def __init__( router: FusedMoERouter, routed_input_transform: torch.nn.Module | None, gate: torch.nn.Module | None, - shared_experts: torch.nn.Module | None, + shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, @@ -224,23 +209,6 @@ def __init__( self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() - # Allow disabling of the separate shared experts stream for - # debug purposes. - # TODO: Remove this after more extensive testings with TP/DP - # and other execution modes - self.use_shared_experts_stream = False - if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: - logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") - self.shared_experts_stream = None - else: - # TODO(rob): enable shared expert overlap with non-cuda-alike. - # aux_stream() returns None on non-cuda-alike platforms. - self.shared_experts_stream = aux_stream() - if self.shared_experts_stream is not None: - logger.debug_once( - "Enabled separate cuda stream for MoE shared_experts", scope="local" - ) - self.use_dp_chunking = ( self.moe_config.moe_parallel_config.use_pplx_kernels or self.moe_config.moe_parallel_config.use_deepep_ll_kernels @@ -254,57 +222,9 @@ def __init__( self.moe_forward = self._select_forward(layer) - def _has_external_experts(self) -> bool: - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - backend = self.moe_config.moe_parallel_config.all2all_backend - return self.shared_experts is not None and not ( - (self.enable_eplb and backend != "allgather_reducescatter") - or self.moe_config.moe_parallel_config.use_fi_all2allv_kernels - ) - - def _determine_shared_experts_order( - self, - hidden_states: torch.Tensor, - ) -> tuple[SharedExpertsOrder, bool]: - if self.shared_experts is None: - return SharedExpertsOrder.NONE, False - - if self._has_external_experts(): - return SharedExpertsOrder.EXTERNAL, False - - has_separate_shared_experts = not self.quant_method.mk_owns_shared_expert - - if not has_separate_shared_experts: - return SharedExpertsOrder.INTERNAL, False - - allow_shared_experts_stream = ( - current_platform.is_cuda() - and has_separate_shared_experts - and not self.use_dp_chunking - and self.shared_experts_stream is not None - and hidden_states.shape[0] - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD - ) - - # Check if we need to run shared experts before matrix multiply because - # matrix multiply may modify the hidden_states. - run_shared_experts_before = ( - has_separate_shared_experts and not allow_shared_experts_stream - ) - - if run_shared_experts_before: - return SharedExpertsOrder.BEFORE_QUANT_METHOD, False - else: - return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream - # XXXXXX used by layer.py and lora/layers/fused_moe.py - def get_shared_experts(self) -> torch.nn.Module | None: - if self._has_external_experts(): - return self.shared_experts - else: - return None + def _get_shared_experts(self) -> SharedExperts | None: + return self.shared_experts def _select_forward(self, layer: torch.nn.Module) -> Callable: if current_platform.is_tpu() or current_platform.is_cpu(): @@ -319,33 +239,6 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) - def _setup_shared_experts_stream( - self, - hidden_states: torch.Tensor, - shared_input: torch.Tensor | None, - ): - if self.use_shared_experts_stream: - assert self.shared_experts_stream is not None - assert self.moe_config.disable_inplace - - shared_experts_input = ( - shared_input if shared_input is not None else hidden_states - ) - - # Record that the shared_experts_input will be used in the - # shared_experts_stream to avoid gc issue from - # deallocation. For more details: - # https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 - # NOTE: We don't need shared_output.record_stream(current_stream()) - # because we synch the streams before using shared_output. - shared_experts_input.record_stream(self.shared_experts_stream) - - # Mark sync start point for the separate shared experts - # stream here since we want to run in parallel with the - # router/gate (next op below) - assert self.shared_experts_stream is not None - self.shared_experts_stream.wait_stream(current_stream()) - def _maybe_init_dp_chunking(self): if not self.use_dp_chunking: return @@ -376,49 +269,6 @@ def _maybe_init_dp_chunking(self): device=device, ) - def _maybe_apply_shared_experts( - self, - shared_output: torch.Tensor | None, - hidden_states: torch.Tensor, - shared_input: torch.Tensor | None, - order: SharedExpertsOrder, - ) -> torch.Tensor | None: - experts_order, use_shared_experts_stream = self._determine_shared_experts_order( - hidden_states - ) - if order != experts_order: - return None - - assert self.shared_experts is not None - assert shared_output is None - if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - # TODO: fold this in? - hidden_states = self._setup_shared_experts_stream( - hidden_states, - shared_input, - ) - - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self.shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - shared_output = self.shared_experts(hidden_states) - current_stream().wait_stream(self.shared_experts_stream) - else: - shared_output = self.shared_experts(hidden_states) - - if order == SharedExpertsOrder.EXTERNAL: - # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely..... - assert shared_output is not None - shared_output = self._maybe_reduce_shared_out(shared_output) - - return shared_output - def must_reduce_shared_expert_outputs(self) -> bool: """ The shared_experts are typically computed using the RowParallelLinear @@ -446,7 +296,9 @@ def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tens else: return tensor_model_parallel_all_reduce(final_hidden_states) - def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor: + def apply_routed_input_transform( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor | None]: """Apply transform for routed experts (e.g., latent projection). This is called by FusedMoE.forward_native. The original hidden_states @@ -462,21 +314,13 @@ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Ten # ReplicatedLinear returns (output, extra_bias) tuple. # We only need the output tensor; extra_bias is not used here. if isinstance(result, tuple): - return result[0] - return result - return hidden_states - - # TODO: combine with runner.forward_impl_X - def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: - # Reduce shared expert outputs if necessary, since the MLP - # should have been created with reduce_results=False. - if ( - self.reduce_results - and get_tensor_model_parallel_world_size() > 1 - and self.must_reduce_shared_expert_outputs() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out + return result[0], hidden_states + return result, hidden_states + + return ( + hidden_states, + hidden_states if self.shared_experts is not None else None, + ) # XXXXX def _maybe_reduce_output( self, @@ -520,14 +364,10 @@ def _encode_layer_name(self) -> str | ModuleName: def _maybe_pad_hidden_states( self, - original_hidden_states: torch.Tensor | None, + shared_experts_input: torch.Tensor, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, list[int]]: - original_hidden_dim = ( - original_hidden_states.shape[-1] - if original_hidden_states is not None - else hidden_states.shape[-1] - ) + shared_experts_hidden_dim = shared_experts_input.shape[-1] transformed_hidden_dim = hidden_states.shape[-1] if ( not self.quant_method.skip_forward_padding @@ -541,26 +381,31 @@ def _maybe_pad_hidden_states( ) if self.shared_experts is not None: - orig_hidden_dims = [original_hidden_dim, transformed_hidden_dim] + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] else: orig_hidden_dims = [transformed_hidden_dim] return hidden_states, orig_hidden_dims + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + order: SharedExpertsOrder, + ): + if self.shared_experts is not None: + assert shared_experts_input is not None + self.shared_experts.apply(shared_experts_input, order) + def _apply_quant_method( self, layer: torch.nn.Module, - shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: # Run this before quant_method to avoid inplace issues. - shared_input = shared_input if shared_input is not None else hidden_states - shared_output = self._maybe_apply_shared_experts( - shared_output, - shared_input, - shared_input, + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.BEFORE_QUANT_METHOD, ) @@ -581,28 +426,18 @@ def _apply_quant_method( x=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, - shared_experts_input=shared_input, + shared_experts_input=shared_experts_input, ) - # TODO: try to clean this up - if isinstance(fused_out, tuple): - if fused_out[0] is None: - assert shared_output is not None - hidden_states = fused_out[1] - else: - assert shared_output is None - shared_output, hidden_states = fused_out - else: - hidden_states = fused_out - - shared_output = self._maybe_apply_shared_experts( - shared_output, - shared_input, - shared_input, + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.AFTER_QUANT_METHOD, ) - return shared_output, hidden_states + return ( + self.shared_experts.output if self.shared_experts is not None else None, + fused_out, + ) def _sequence_parallel_context(self): ctx = get_forward_context() @@ -719,33 +554,28 @@ def _maybe_combine( # # forward # - self.moe_forward (_moe_forward or _moe_forward_shared) - # - forward_chunking_dispatch + # - forward_dispatch # - forward_impl or forward_impl_chunked + # def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # For latent MoE: save ORIGINAL hidden_states before transform - # (shared_experts need original dimension, routed experts use transformed) - original_hidden_states = ( - # or shared_experts is not None? - hidden_states if self.routed_input_transform is not None else None - ) - # Apply transform for routed experts (e.g., latent projection for latent MoE) - hidden_states = self.apply_routed_input_transform(hidden_states) + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( - original_hidden_states, + shared_experts_input, hidden_states, ) - # Note: moe_forward will call forward_dispatch fused_output = self.moe_forward( hidden_states, router_logits, - original_hidden_states, + shared_experts_input, self._encode_layer_name(), ) @@ -756,17 +586,15 @@ def forward_dispatch( layer: torch.nn.Module, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() router_logits = self._maybe_gate(hidden_states, router_logits) - shared_output = self._maybe_apply_shared_experts( - None, - hidden_states, - shared_input, + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.EXTERNAL, ) @@ -774,18 +602,16 @@ def forward_dispatch( if self.use_dp_chunking: return self.forward_impl_chunked( layer, - shared_output, hidden_states, router_logits, - shared_input, + shared_experts_input, ) else: return self.forward_impl( layer, - shared_output, hidden_states, router_logits, - shared_input, + shared_experts_input, ) def _slice_and_copy_input( @@ -811,10 +637,9 @@ def _slice_and_copy_input( def forward_impl_chunked( self, layer: torch.nn.Module, - shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # Gate overlap not supported when chunking is enabled. Run the # gate first. @@ -865,18 +690,17 @@ def forward_impl_chunked( chunk_end, ) - shared_input_chunk = ( - shared_input[chunk_start:chunk_end, :] - if shared_input is not None + shared_experts_input_chunk = ( + shared_experts_input[chunk_start:chunk_end, :] + if shared_experts_input is not None else None ) shared_output_chunk, hidden_states_chunk = self._apply_quant_method( layer=layer, - shared_output=shared_output, hidden_states=hidden_states_chunk, router_logits=router_logits_chunk, - shared_input=shared_input_chunk, + shared_experts_input=shared_experts_input_chunk, ) # Store outputs @@ -901,10 +725,9 @@ def forward_impl_chunked( def forward_impl( self, layer: torch.nn.Module, - shared_output: torch.Tensor | None, hidden_states: torch.Tensor, router_logits: torch.Tensor, - shared_input: torch.Tensor | None, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP @@ -917,10 +740,9 @@ def forward_impl( shared_output, hidden_states = self._apply_quant_method( layer=layer, - shared_output=shared_output, hidden_states=hidden_states, router_logits=router_logits, - shared_input=shared_input, + shared_experts_input=shared_experts_input, ) return self._maybe_combine( diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py new file mode 100644 index 000000000000..ad1bac4dbb03 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from enum import IntEnum + +import torch + +import vllm.envs as envs +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + aux_stream, + current_stream, +) + +logger = init_logger(__name__) + + +class SharedExpertsOrder(IntEnum): + # No shared experts. + NONE = (0,) + + # Get rid of this one? combine with BEFORE? + # Note: this might be important for torch.compile reasons. Can + # get rid of it after _moe_forward is undone. + EXTERNAL = (1,) + + # Called by modular kernel. + INTERNAL = (2,) + + # Called right before quant_method is executed. + BEFORE_QUANT_METHOD = (3,) + + # Called right after quant_method is executed (possibly with streaming). + AFTER_QUANT_METHOD = (4,) + + +# XXXXX add method to prime with shared_experts_input? NO +class SharedExperts: + def __init__( + self, + shared_experts: torch.nn.Module, + moe_config: FusedMoEConfig, + has_separate_shared_experts: bool, # better name + use_dp_chunking: bool, + must_reduce_shared_expert_outputs: bool, # and reduce_results + ): + self._output: torch.Tensor | None = None + self._shared_experts = shared_experts + self._moe_config = moe_config # invariant + self._use_dp_chunking = use_dp_chunking # invariant + + self._has_separate_shared_experts = has_separate_shared_experts # depends on MK + self._must_reduce_shared_expert_outputs = ( + must_reduce_shared_expert_outputs # depends on MK + ) + + # Allow disabling of the separate shared experts stream for + # debug purposes. + # TODO: Remove this after more extensive testings with TP/DP + # and other execution modes + self.use_shared_experts_stream = False + if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: + logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") + self._shared_experts_stream = None + else: + # TODO(rob): enable shared expert overlap with non-cuda-alike. + # aux_stream() returns None on non-cuda-alike platforms. + self._shared_experts_stream = aux_stream() + if self._shared_experts_stream is not None: + logger.debug_once( + "Enabled separate cuda stream for MoE shared_experts", scope="local" + ) + + def _has_external_experts(self) -> bool: + # Disable shared expert overlap if: + # - we are using eplb with non-default backend, because of correctness issues + # - we are using flashinfer with DP, since there nothing to gain + backend = self._moe_config.moe_parallel_config.all2all_backend + return self._shared_experts is not None and not ( + ( + self._moe_config.moe_parallel_config.enable_eplb + and backend != "allgather_reducescatter" + ) + or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels + ) + + def _determine_shared_experts_order( + self, + hidden_states: torch.Tensor, + ) -> tuple[SharedExpertsOrder, bool]: + if self._shared_experts is None: + return SharedExpertsOrder.NONE, False + + if self._has_external_experts(): + return SharedExpertsOrder.EXTERNAL, False + + if ( + not self._has_separate_shared_experts + or not self._moe_config.moe_parallel_config.use_all2all_kernels + ): + return SharedExpertsOrder.INTERNAL, False + + allow_shared_experts_stream = ( + current_platform.is_cuda() + and self._has_separate_shared_experts + and not self._use_dp_chunking + and self._shared_experts_stream is not None + and hidden_states.shape[0] + <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + ) + + # Check if we need to run shared experts before matrix multiply because + # matrix multiply may modify the hidden_states. + run_shared_experts_before = ( + self._has_separate_shared_experts and not allow_shared_experts_stream + ) + + if run_shared_experts_before: + return SharedExpertsOrder.BEFORE_QUANT_METHOD, False + else: + return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream + + def _setup_shared_experts_stream( + self, + shared_experts_input: torch.Tensor, + ): + assert self._shared_experts_stream is not None + assert self._moe_config.disable_inplace + + # Record that the clone will be used by shared_experts_stream + # to avoid gc issue from deallocation of hidden_states_clone + # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 + # NOTE: We don't need shared_output.record_stream(current_stream()) + # because we synch the streams before using shared_output. + shared_experts_input.record_stream(self._shared_experts_stream) + + # Mark sync start point for the separate shared experts + # stream here since we want to run in parallel with the + # router/gate (next op below) + assert self._shared_experts_stream is not None + self._shared_experts_stream.wait_stream(current_stream()) + + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: + # Reduce shared expert outputs if necessary, since the MLP + # should have been created with reduce_results=False. + if ( + self._must_reduce_shared_expert_outputs + and get_tensor_model_parallel_world_size() > 1 + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out + + @property + def output(self) -> torch.Tensor | None: + assert self._shared_experts is None == self._output is None + output = self._output + self._output = None + return output + + # add split/join? + + def apply( + self, + shared_experts_input: torch.Tensor, + order: SharedExpertsOrder, + ) -> torch.Tensor | None: + experts_order, use_shared_experts_stream = self._determine_shared_experts_order( + shared_experts_input, + ) + if order != experts_order: + return None + + assert self._shared_experts is not None + assert self._output is None + if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: + # TODO: fold this in? + self._setup_shared_experts_stream(shared_experts_input) + + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self._shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream + self._output = self._shared_experts(shared_experts_input) + current_stream().wait_stream(self._shared_experts_stream) + else: + self._output = self._shared_experts(shared_experts_input) + + if order == SharedExpertsOrder.EXTERNAL: + # TODO: figure out how to combine this with maybe_reduce_output? + # or get rid of it completely..... + assert self._output is not None + self._output = self._maybe_reduce_shared_out(self._output) + + # TODO: do AFTER reduce here? + + return self._output # ? diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index a29d8a7d8dda..e2a5d05320d0 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -93,7 +93,7 @@ def forward_native( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.forward_cuda(layer, x, topk_weights, topk_ids, shared_experts_input) @property @@ -297,7 +297,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.forward( layer=layer, x=x, @@ -322,7 +322,7 @@ def forward_cuda( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.kernel is not None return self.kernel.apply( @@ -343,7 +343,7 @@ def forward_monolithic_cuda( layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: F401 assert self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM @@ -369,7 +369,7 @@ def forward_monolithic_cpu( layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821 x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return self.cpu_fused_moe( layer, x, diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 5b7af3193b03..da717556adbd 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -765,7 +765,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return fused_marlin_moe( x, layer.w13_qweight, diff --git a/vllm/model_executor/layers/quantization/bitsandbytes.py b/vllm/model_executor/layers/quantization/bitsandbytes.py index 716a20090f69..729924663646 100644 --- a/vllm/model_executor/layers/quantization/bitsandbytes.py +++ b/vllm/model_executor/layers/quantization/bitsandbytes.py @@ -483,7 +483,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts # TODO(bnell): Do these need to be called on the hot path? diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 29115fbbc255..8fa25efd1da0 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -349,7 +349,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply( x, @@ -597,7 +597,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -622,7 +622,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply( x, @@ -959,7 +959,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( x, @@ -983,7 +983,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -1123,7 +1123,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( @@ -1607,7 +1607,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.kernel_backend == "Flashinfer" return flashinfer_trtllm_mxint4_moe( x=x, @@ -1634,7 +1634,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.kernel_backend == "Marlin" return fused_marlin_moe( x, @@ -1883,7 +1883,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( @@ -2500,7 +2500,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if layer.enable_eplb: raise NotImplementedError( "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet." diff --git a/vllm/model_executor/layers/quantization/experts_int8.py b/vllm/model_executor/layers/quantization/experts_int8.py index d971f3b5b0d2..301441ff019d 100644 --- a/vllm/model_executor/layers/quantization/experts_int8.py +++ b/vllm/model_executor/layers/quantization/experts_int8.py @@ -141,7 +141,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 5101347cd02a..f967e41a79b3 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -961,7 +961,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -986,7 +986,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/gguf.py b/vllm/model_executor/layers/quantization/gguf.py index 88023349e779..cf72c83d717a 100644 --- a/vllm/model_executor/layers/quantization/gguf.py +++ b/vllm/model_executor/layers/quantization/gguf.py @@ -637,7 +637,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if layer.apply_router_weight_on_input: raise NotImplementedError( "Apply router weight on input is not supported for" diff --git a/vllm/model_executor/layers/quantization/gptq_marlin.py b/vllm/model_executor/layers/quantization/gptq_marlin.py index d7b2a366e1f0..0b43fbd392c9 100644 --- a/vllm/model_executor/layers/quantization/gptq_marlin.py +++ b/vllm/model_executor/layers/quantization/gptq_marlin.py @@ -901,7 +901,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: return fused_marlin_moe( x, layer.w13_qweight, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 78644f74d288..a5e1ea52f557 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -932,7 +932,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -957,7 +957,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( @@ -1416,7 +1416,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( @@ -1441,7 +1441,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply( diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index f5c679840432..a327ac17bbc9 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -369,7 +369,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts assert layer.activation == MoEActivation.SILU, ( diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index f992d0f86c4e..16536b85caea 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -1051,7 +1051,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert not self.is_monolithic if layer.enable_eplb: raise NotImplementedError("EPLB is not supported for mxfp4") @@ -1096,7 +1096,7 @@ def apply_monolithic( layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: assert self.is_monolithic if layer.enable_eplb: diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 0a5db4e71fdb..fcf03d330d11 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -445,7 +445,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if self.rocm_aiter_moe_enabled: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, @@ -634,7 +634,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, ) @@ -1042,7 +1042,7 @@ def apply( topk_weights: torch.Tensor, topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: if not self.emulate: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, From e8865e6c9a7e2e5dc8ccd1cda73f231d69a72528 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Feb 2026 18:47:11 -0500 Subject: [PATCH 060/108] cleanups Signed-off-by: Bill Nell --- .../model_executor/layers/fused_moe/config.py | 10 ++ vllm/model_executor/layers/fused_moe/layer.py | 20 ++-- .../fused_moe/runner/default_moe_runner.py | 98 +++++++++++-------- .../layers/fused_moe/runner/shared_experts.py | 76 ++++++++------ .../model_executor/models/transformers/moe.py | 2 + 5 files changed, 121 insertions(+), 85 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 2500387debe1..cfe50a36cf8e 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -937,6 +937,16 @@ class FusedMoEParallelConfig: all2all_backend: str # all2all backend for MoE communication enable_eplb: bool # whether to enable expert load balancing + @property + def use_dp_chunking(self) -> bool: + return ( + or self.use_deepep_ll_kernels + or self.use_mori_kernels + or self.use_fi_all2allv_kernels + or self.use_fi_nvl_two_sided_kernels + or self.use_nixl_ep_kernels + ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + @property def is_sequence_parallel(self) -> bool: return self.sp_size > 1 diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index bbf6e2ffedb9..c8328ec9b63f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -642,19 +642,16 @@ def _init_shared_experts(self) -> SharedExperts | None: if self._shared_experts is None: return None - reduce_shared_output = ( - self.reduce_results - # XXXX ordering issue - and self.quant_method.moe_mk is not None - and self.quant_method.moe_mk.output_is_reduced() - ) - return SharedExperts( self._shared_experts, moe_config=self.moe_config, - has_separate_shared_experts=not self.quant_method.mk_owns_shared_expert, - use_dp_chunking=self.use_dp_chunking, # XXXXXXXXXXXXXX - must_reduce_shared_expert_outputs=reduce_shared_output, + # Note: For now we must pass quant_method along to SharedExperts so it + # can property determine where the shared experts are supposed to be + # called, i.e. by a MK or by the MoERunner. + # Once the MK can be created upfront, we can just pass in the proper + # flags dervied from the quant_method's MK. + reduce_results=self.reduce_results, + quant_method=self.quant_method, ) def _init_runner(self) -> DefaultMoERunner: @@ -671,7 +668,6 @@ def _init_runner(self) -> DefaultMoERunner: quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, - enable_eplb=self.enable_eplb, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py @@ -710,7 +706,7 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.runner._get_shared_experts(), # XXXXXXXXXXXXXXXXX + self.shared_experts, inplace=not self.moe_config.disable_inplace, ) ) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index cc177b946cf1..5822fa7bcf26 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -7,7 +7,6 @@ import torch import torch.nn.functional as F -import vllm.envs as envs from vllm.distributed import ( get_ep_group, get_pcp_group, @@ -187,7 +186,6 @@ def __init__( quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, - enable_eplb: bool, ): super().__init__() self.moe_config = moe_config @@ -198,10 +196,10 @@ def __init__( self.quant_method = quant_method self.reduce_results = reduce_results self.enable_dbo = enable_dbo - self.enable_eplb = enable_eplb + self.enable_eplb = moe_config.moe_parallel_config.enable_eplb # Chunked all2all staging tensor - # TODO rename these + # TODO rename these? # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. # TODO: these could be global, i.e. shared by all layers @@ -209,34 +207,31 @@ def __init__( self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() - self.use_dp_chunking = ( - self.moe_config.moe_parallel_config.use_pplx_kernels - or self.moe_config.moe_parallel_config.use_deepep_ll_kernels - or self.moe_config.moe_parallel_config.use_mori_kernels - or self.moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels - or self.moe_config.moe_parallel_config.use_nixl_ep_kernels - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK + self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name - self.moe_forward = self._select_forward(layer) + self.moe_forward, self.moe_forward_impl = self._select_forward(layer) - # XXXXXX used by layer.py and lora/layers/fused_moe.py - def _get_shared_experts(self) -> SharedExperts | None: - return self.shared_experts - - def _select_forward(self, layer: torch.nn.Module) -> Callable: + def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: + forward_impl_fn = ( + self.forward_impl_chunked if self.use_dp_chunking else self.forward_impl + ) if current_platform.is_tpu() or current_platform.is_cpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. # Note: CPU doesn't require wrapped forward_impl. - return _moe_forward if self.shared_experts is None else _moe_forward_shared + return ( + _moe_forward if self.shared_experts is None else _moe_forward_shared, + forward_impl_fn, + ) return ( torch.ops.vllm.moe_forward if self.shared_experts is None - else torch.ops.vllm.moe_forward_shared + else torch.ops.vllm.moe_forward_shared, + forward_impl_fn, ) def _maybe_init_dp_chunking(self): @@ -308,6 +303,9 @@ def apply_routed_input_transform( TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be moved inside SharedFusedMoE to all-reduce on the smaller latent dimension. + + Returns (possibly transformed) hidden states and the input for shared + experts (or None if there are no shared experts). """ if self.routed_input_transform is not None: result = self.routed_input_transform(hidden_states) @@ -320,7 +318,7 @@ def apply_routed_input_transform( return ( hidden_states, hidden_states if self.shared_experts is not None else None, - ) # XXXXX + ) def _maybe_reduce_output( self, @@ -364,10 +362,12 @@ def _encode_layer_name(self) -> str | ModuleName: def _maybe_pad_hidden_states( self, - shared_experts_input: torch.Tensor, + shared_experts_input: torch.Tensor | None, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, list[int]]: - shared_experts_hidden_dim = shared_experts_input.shape[-1] + shared_experts_hidden_dim = ( + shared_experts_input.shape[-1] if shared_experts_input is not None else 0 + ) transformed_hidden_dim = hidden_states.shape[-1] if ( not self.quant_method.skip_forward_padding @@ -551,17 +551,39 @@ def _maybe_combine( else: return hidden_states - # - # forward - # - self.moe_forward (_moe_forward or _moe_forward_shared) - # - forward_dispatch - # - forward_impl or forward_impl_chunked - # def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Invoke the fused moe layer. + + Input: + - hidden_states + - router_logits + + Output: + - The new hidden_states. + or + - A tuple of (shared experts output, new hidden_states). + + Calling sequence + - forward + - self.moe_forward (_moe_forward or _moe_forward_shared custom op) + - forward_dispatch + - moe_forward_impl (forward_impl or forward_impl_chunked) + + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due + to the following reasons: + 1. the chunking loop in forward_impl_chunked cannot be compiled by torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward and + _moe_forward_shared must be split. + + If forward_impl_chunked can be implemented via torch.scan we can potentially get + rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into + the 'forward' method. + """ + # Apply transform for routed experts (e.g., latent projection for latent MoE) hidden_states, shared_experts_input = self.apply_routed_input_transform( hidden_states @@ -599,20 +621,12 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - if self.use_dp_chunking: - return self.forward_impl_chunked( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - else: - return self.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) + return self.moe_forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) def _slice_and_copy_input( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index ad1bac4dbb03..5081311e57df 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, ) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( aux_stream, @@ -41,31 +44,25 @@ class SharedExpertsOrder(IntEnum): AFTER_QUANT_METHOD = (4,) -# XXXXX add method to prime with shared_experts_input? NO class SharedExperts: def __init__( self, shared_experts: torch.nn.Module, moe_config: FusedMoEConfig, - has_separate_shared_experts: bool, # better name - use_dp_chunking: bool, - must_reduce_shared_expert_outputs: bool, # and reduce_results + quant_method: FusedMoEMethodBase, + reduce_results: bool, ): self._output: torch.Tensor | None = None self._shared_experts = shared_experts - self._moe_config = moe_config # invariant - self._use_dp_chunking = use_dp_chunking # invariant - - self._has_separate_shared_experts = has_separate_shared_experts # depends on MK - self._must_reduce_shared_expert_outputs = ( - must_reduce_shared_expert_outputs # depends on MK - ) + self._moe_config = moe_config + self._quant_method = quant_method + self._reduce_results = reduce_results + self._use_dp_chunking = moe_config.moe_parallel_config.use_dp_chunking # Allow disabling of the separate shared experts stream for # debug purposes. # TODO: Remove this after more extensive testings with TP/DP # and other execution modes - self.use_shared_experts_stream = False if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") self._shared_experts_stream = None @@ -78,6 +75,7 @@ def __init__( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) + @property def _has_external_experts(self) -> bool: # Disable shared expert overlap if: # - we are using eplb with non-default backend, because of correctness issues @@ -91,6 +89,22 @@ def _has_external_experts(self) -> bool: or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels ) + # TODO(bnell): better name + @property + def _has_separate_shared_experts(self) -> bool: + return ( + not self._quant_method.mk_owns_shared_expert + and self._shared_experts is not None + ) + + @property + def _must_reduce_shared_expert_outputs(self) -> bool: + return ( + self._reduce_results + and self._quant_method.moe_mk is not None + and self._quant_method.moe_mk.output_is_reduced() + ) + def _determine_shared_experts_order( self, hidden_states: torch.Tensor, @@ -98,7 +112,7 @@ def _determine_shared_experts_order( if self._shared_experts is None: return SharedExpertsOrder.NONE, False - if self._has_external_experts(): + if self._has_external_experts: return SharedExpertsOrder.EXTERNAL, False if ( @@ -127,10 +141,10 @@ def _determine_shared_experts_order( else: return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream - def _setup_shared_experts_stream( + def _call_with_shared_experts_stream( self, shared_experts_input: torch.Tensor, - ): + ) -> torch.Tensor: assert self._shared_experts_stream is not None assert self._moe_config.disable_inplace @@ -144,9 +158,21 @@ def _setup_shared_experts_stream( # Mark sync start point for the separate shared experts # stream here since we want to run in parallel with the # router/gate (next op below) - assert self._shared_experts_stream is not None self._shared_experts_stream.wait_stream(current_stream()) + # Run shared experts in parallel on a separate stream + # NOTE: We start the separate stream here and mark the + # sync end point immediately after it is done. This is + # important to avoid excessive stream allocations by the cuda + # graph replay later. + with torch.cuda.stream(self._shared_experts_stream): + # Note that hidden_states clone() is necessary here to avoid + # conflict with the main stream + output = self._shared_experts(shared_experts_input) + current_stream().wait_stream(self._shared_experts_stream) + + return output + def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: # Reduce shared expert outputs if necessary, since the MLP # should have been created with reduce_results=False. @@ -164,8 +190,6 @@ def output(self) -> torch.Tensor | None: self._output = None return output - # add split/join? - def apply( self, shared_experts_input: torch.Tensor, @@ -174,25 +198,15 @@ def apply( experts_order, use_shared_experts_stream = self._determine_shared_experts_order( shared_experts_input, ) + if order != experts_order: return None assert self._shared_experts is not None assert self._output is None + if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - # TODO: fold this in? - self._setup_shared_experts_stream(shared_experts_input) - - # Run shared experts in parallel on a separate stream - # NOTE: We start the separate stream here and mark the - # sync end point immediately after it is done. This is - # important to avoid excessive stream allocations by the cuda - # graph replay later. - with torch.cuda.stream(self._shared_experts_stream): - # Note that hidden_states clone() is necessary here to avoid - # conflict with the main stream - self._output = self._shared_experts(shared_experts_input) - current_stream().wait_stream(self._shared_experts_stream) + self._output = self._call_with_shared_experts_stream(shared_experts_input) else: self._output = self._shared_experts(shared_experts_input) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 5f8352faed50..f65a197abcfc 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -94,6 +94,8 @@ def transformers_moe_forward( self = forward_context.no_compile_layers[layer_name] self._topk_ids = topk_ids # Clone hidden_states because it will be mutated in-place in FusedMoE + # TODO(bnell): figure out a way to avoid calling runner directly. + # it is a hack that the weight are being passed via logits. return self.runner.forward(hidden_states.clone(), topk_weights) From f83e0f5d13963d2a1a61dfd9b692420eae416b04 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Feb 2026 18:54:43 -0500 Subject: [PATCH 061/108] fix circular import Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/shared_experts.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 5081311e57df..305697adf699 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -13,8 +13,8 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, ) -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( - FusedMoEMethodBase, +from vllm.model_executor.layers.quantization.base_config import ( + QuantizeMethodBase, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( @@ -49,9 +49,17 @@ def __init__( self, shared_experts: torch.nn.Module, moe_config: FusedMoEConfig, - quant_method: FusedMoEMethodBase, + quant_method: QuantizeMethodBase, reduce_results: bool, ): + from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, + ) + + # quant_method must be a FusedMoEMethodBase but we can't use the type + # due to circular imports. + assert isinstance(quant_method, FusedMoEMethodBase) + self._output: torch.Tensor | None = None self._shared_experts = shared_experts self._moe_config = moe_config From 88e80b915d2a9f64bd1792a44ee0239c555c0771 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 16:32:20 -0500 Subject: [PATCH 062/108] fixes Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- .../layers/fused_moe/runner/default_moe_runner.py | 3 +-- vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 4 ++-- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c8328ec9b63f..67d8af288abf 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -741,7 +741,7 @@ def use_ep(self): @property def is_internal_router(self) -> bool: # By default, router/gate is called before FusedMoE forward pass - return self.gate is not None + return self._gate is not None def _maybe_init_expert_routing_tables( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 5822fa7bcf26..ac9086444dcf 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -203,12 +203,11 @@ def __init__( # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. # TODO: these could be global, i.e. shared by all layers + self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None self._maybe_init_dp_chunking() - self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking - # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 305697adf699..76a150bd1241 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -120,7 +120,7 @@ def _determine_shared_experts_order( if self._shared_experts is None: return SharedExpertsOrder.NONE, False - if self._has_external_experts: + if self._has_external_experts and not self._use_dp_chunking: return SharedExpertsOrder.EXTERNAL, False if ( @@ -193,7 +193,7 @@ def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: @property def output(self) -> torch.Tensor | None: - assert self._shared_experts is None == self._output is None + assert (self._shared_experts is None) == (self._output is None) output = self._output self._output = None return output From 781d4eaf70af1f7e9a0cf9cdbf56cfad4e6bb9c8 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 17:04:50 -0500 Subject: [PATCH 063/108] renames Signed-off-by: Bill Nell --- tools/ep_kernels/install_python_libraries.sh | 1 + .../fused_moe/runner/default_moe_runner.py | 33 ++++++++++--------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 3372dd10f4dc..c3deb7d6060c 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -103,6 +103,7 @@ pushd "$WORKSPACE" echo "Downloading NVSHMEM ${NVSHMEM_VER} for ${NVSHMEM_SUBDIR} ..." curl -fSL "${NVSHMEM_URL}" -o "${NVSHMEM_FILE}" tar -xf "${NVSHMEM_FILE}" +rm -rf nvshmem mv "${NVSHMEM_FILE%.tar.xz}" nvshmem rm -f "${NVSHMEM_FILE}" rm -rf nvshmem/lib/bin nvshmem/lib/share diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index ac9086444dcf..2a0117e9c933 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -211,12 +211,14 @@ def __init__( # Needed for string -> FusedMoE layer lookup in custom ops. self.layer_name = layer.layer_name - self.moe_forward, self.moe_forward_impl = self._select_forward(layer) + self.forward_entry, self.forward_impl = self._select_forward(layer) def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: + # Select implementation based on presence of DP chunking. forward_impl_fn = ( - self.forward_impl_chunked if self.use_dp_chunking else self.forward_impl + self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl ) + if current_platform.is_tpu() or current_platform.is_cpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. @@ -568,19 +570,20 @@ def forward( Calling sequence - forward - - self.moe_forward (_moe_forward or _moe_forward_shared custom op) + - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - forward_dispatch - - moe_forward_impl (forward_impl or forward_impl_chunked) + - forward_impl (_forward_impl or _forward_impl_chunked) Note: The existence of _moe_forward and _moe_forward_shared custom ops are due to the following reasons: - 1. the chunking loop in forward_impl_chunked cannot be compiled by torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward and - _moe_forward_shared must be split. - - If forward_impl_chunked can be implemented via torch.scan we can potentially get - rid of _moe_forward and _moe_forward_shared and collapse the whole sequence into - the 'forward' method. + 1. the chunking loop in _forward_impl_chunked cannot be compiled by + torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward + and _moe_forward_shared must be split. + + If _forward_impl_chunked can be implemented via torch.scan we can potentially + get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence + into the 'forward' method. """ # Apply transform for routed experts (e.g., latent projection for latent MoE) @@ -593,7 +596,7 @@ def forward( hidden_states, ) - fused_output = self.moe_forward( + fused_output = self.forward_entry( hidden_states, router_logits, shared_experts_input, @@ -620,7 +623,7 @@ def forward_dispatch( ) with self._sequence_parallel_context(): - return self.moe_forward_impl( + return self.forward_impl( layer, hidden_states, router_logits, @@ -647,7 +650,7 @@ def _slice_and_copy_input( out_slice.copy_(orig_slice, non_blocking=True) return out_slice - def forward_impl_chunked( + def _forward_impl_chunked( self, layer: torch.nn.Module, hidden_states: torch.Tensor, @@ -735,7 +738,7 @@ def forward_impl_chunked( assert final_shared_hidden_states is not None return (final_shared_hidden_states, final_fused_hidden_states) - def forward_impl( + def _forward_impl( self, layer: torch.nn.Module, hidden_states: torch.Tensor, From 369501657ac58a748f1947690b278c0cfc2a455d Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 17:23:49 -0500 Subject: [PATCH 064/108] add comment Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/default_moe_runner.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 2a0117e9c933..64fb330d0487 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -77,6 +77,9 @@ def _resolve_layer_name(layer_name: str | ModuleName) -> str: return layer_name.value if isinstance(layer_name, ModuleName) else layer_name +# Note: _moe_forward and _moe_forward_shared should not contain any +# implementation details, They should merely pass along control to +# the runner's 'forward_dispatch' method. def _moe_forward( hidden_states: torch.Tensor, router_logits: torch.Tensor, From 053f66f7cfa0cd59bada3dae156c341420355e07 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 24 Feb 2026 17:26:55 -0500 Subject: [PATCH 065/108] more renames Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/shared_experts.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 76a150bd1241..45538daff5f8 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -97,9 +97,8 @@ def _has_external_experts(self) -> bool: or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels ) - # TODO(bnell): better name @property - def _has_separate_shared_experts(self) -> bool: + def _has_mk_owned_shared_experts(self) -> bool: return ( not self._quant_method.mk_owns_shared_expert and self._shared_experts is not None @@ -124,14 +123,14 @@ def _determine_shared_experts_order( return SharedExpertsOrder.EXTERNAL, False if ( - not self._has_separate_shared_experts + not self._has_mk_owned_shared_experts or not self._moe_config.moe_parallel_config.use_all2all_kernels ): return SharedExpertsOrder.INTERNAL, False allow_shared_experts_stream = ( current_platform.is_cuda() - and self._has_separate_shared_experts + and self._has_mk_owned_shared_experts and not self._use_dp_chunking and self._shared_experts_stream is not None and hidden_states.shape[0] @@ -141,7 +140,7 @@ def _determine_shared_experts_order( # Check if we need to run shared experts before matrix multiply because # matrix multiply may modify the hidden_states. run_shared_experts_before = ( - self._has_separate_shared_experts and not allow_shared_experts_stream + self._has_mk_owned_shared_experts and not allow_shared_experts_stream ) if run_shared_experts_before: @@ -202,7 +201,7 @@ def apply( self, shared_experts_input: torch.Tensor, order: SharedExpertsOrder, - ) -> torch.Tensor | None: + ): experts_order, use_shared_experts_stream = self._determine_shared_experts_order( shared_experts_input, ) @@ -220,10 +219,8 @@ def apply( if order == SharedExpertsOrder.EXTERNAL: # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely..... + # or get rid of it completely. assert self._output is not None self._output = self._maybe_reduce_shared_out(self._output) - # TODO: do AFTER reduce here? - - return self._output # ? + # TODO(bnell): potentially do AFTER reduce here insteed of in runner. From 708dd2bbe3089b8c08d6fd6cd0c5c748b06faa05 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Feb 2026 12:55:16 -0500 Subject: [PATCH 066/108] cleanup Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index b6c62b618a09..11332f4cb228 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -172,7 +172,6 @@ def wrapper(*args, **kwargs): moe_state_dict["apply_router_weight_on_input"] = kwargs[ "apply_router_weight_on_input" ] - # TODO: global_num_experts/shared_experts_input? result = func(*args, **kwargs) return result @@ -597,10 +596,6 @@ def forward(self, *args, **kwargs): def maybe_all_reduce_tensor_model_parallel(self, *args, **kwargs): return self.base_layer.maybe_all_reduce_tensor_model_parallel(*args, **kwargs) - # @property - # def _shared_experts(self): - # return self.base_layer._shared_experts - @property def quant_method(self): return self.base_layer.quant_method From 5748f7c8391e422a7afd12ecac1dbdd0ca444e32 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Feb 2026 19:38:08 -0500 Subject: [PATCH 067/108] remove memoizing router, not needed yet Signed-off-by: Bill Nell --- .../fused_moe/router/memoizing_router.py | 35 ------------------- 1 file changed, 35 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py deleted file mode 100644 index a55bd2f09d6f..000000000000 --- a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import torch - -from vllm.model_executor.layers.fused_moe.config import ( - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter - - -class MemoizingRouter(FusedMoERouter): - def __init__(self, router: FusedMoERouter): - self.router = router - - def set_capture_fn( - self, - capture_fn: Callable[[torch.Tensor], None] | None, - ) -> None: - self.router.set_capture_fn(capture_fn) - self.results: tuple[torch.Tensor, torch.Tensor] | None = None - - @property - def routing_method_type(self) -> RoutingMethodType: - return self.router.routing_method_type - - def select_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.results is None: - self.results = self.router.select_experts(hidden_states, router_logits) - return self.results From 9123f156bdf00f82946cb5a82406f30030b525d2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 27 Feb 2026 18:33:43 -0500 Subject: [PATCH 068/108] fix UBD bug Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 67d8af288abf..772be9e150bb 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -567,7 +567,7 @@ def __init__( device=vllm_config.device_config.device, routing_method=self.routing_method_type, # TODO: in_dtype == out_dtype? - disable_inplace=disable_inplace() or self.shared_experts is not None, + disable_inplace=disable_inplace() or shared_experts is not None, ) if self.moe_config.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( From 04b430f9986072a02ff1106e07fe68c60d7fdd72 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 16:35:53 -0500 Subject: [PATCH 069/108] cleanup merge Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/default_moe_runner.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 64fb330d0487..c48744174d76 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -37,8 +37,6 @@ from vllm.utils.torch_utils import ( HAS_OPAQUE_TYPE, ModuleName, - aux_stream, - current_stream, direct_register_custom_op, ) from vllm.v1.worker.ubatching import dbo_current_ubatch_id @@ -86,7 +84,7 @@ def _moe_forward( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> torch.Tensor: - layer = get_layer_from_name(layer_name) + layer = get_layer_from_name(_resolve_layer_name(layer_name)) return layer.runner.forward_dispatch( layer, hidden_states, @@ -110,7 +108,7 @@ def _moe_forward_shared( shared_experts_input: torch.Tensor | None, layer_name: _layer_name_type, ) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(layer_name) + layer = get_layer_from_name(_resolve_layer_name(layer_name)) return layer.runner.forward_dispatch( layer, hidden_states, @@ -141,7 +139,7 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward", op_func=_moe_forward, - mutates_args=["hidden_states"], # ? + mutates_args=["hidden_states"], # is this still true? fake_impl=_moe_forward_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -150,7 +148,6 @@ def _moe_forward_shared_fake( direct_register_custom_op( op_name="moe_forward_shared", op_func=_moe_forward_shared, - mutates_args=["hidden_states"], # ? fake_impl=_moe_forward_shared_fake, tags=(torch.Tag.needs_fixed_stride_order,), ) @@ -202,10 +199,8 @@ def __init__( self.enable_eplb = moe_config.moe_parallel_config.enable_eplb # Chunked all2all staging tensor - # TODO rename these? # These need to exist ahead of time due to CUDAgraph construction # needing a fixed buffer address. - # TODO: these could be global, i.e. shared by all layers self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking self.batched_hidden_states: torch.Tensor | None = None self.batched_router_logits: torch.Tensor | None = None From 526db3885a5c9bc4433b36e713c1a8d2ab1abf57 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 16:46:20 -0500 Subject: [PATCH 070/108] fix merge Signed-off-by: Bill Nell --- tests/kernels/moe/utils.py | 2 -- vllm/model_executor/layers/fused_moe/config.py | 2 +- vllm/model_executor/layers/fused_moe/modular_kernel.py | 8 ++------ 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 4b693d8c8a55..2ef4424c2baa 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -603,7 +603,6 @@ def make_shared_experts( def modular_triton_fused_moe( moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, - shared_experts: torch.nn.Module | None = None, ) -> FusedMoEKernel: return FusedMoEKernel( maybe_make_prepare_finalize( @@ -613,6 +612,5 @@ def modular_triton_fused_moe( use_monolithic=False, ), TritonExperts(moe_config, quant_config), - shared_experts, inplace=False, ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index cfe50a36cf8e..b3f57e3e376a 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -940,7 +940,7 @@ class FusedMoEParallelConfig: @property def use_dp_chunking(self) -> bool: return ( - or self.use_deepep_ll_kernels + self.use_deepep_ll_kernels or self.use_mori_kernels or self.use_fi_all2allv_kernels or self.use_fi_nvl_two_sided_kernels diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index bbafc682c102..e424d36536b7 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1001,11 +1001,9 @@ def __init__( self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts self.shared_experts = shared_experts - self.moe_parallel_config = moe_parallel_config self.inplace = inplace - moe_parallel_config = fused_experts.moe_config.moe_parallel_config - self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config + self.moe_parallel_config = moe_parallel_config self.is_dp_ep = ( moe_parallel_config is not None and moe_parallel_config.dp_size > 1 @@ -1459,8 +1457,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalize, fused_experts: FusedMoEExperts, - shared_experts: torch.nn.Module | None = None, - moe_parallel_config: FusedMoEParallelConfig | None = None, + shared_experts: SharedExperts | None = None, inplace: bool = False, ): super().__init__() @@ -1475,7 +1472,6 @@ def __init__( prepare_finalize, fused_experts, shared_experts, - moe_parallel_config, inplace, ) From 67bdab221527c4d1fd4eb9ac70469d9d0cb39992 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 17:02:59 -0500 Subject: [PATCH 071/108] fix merge Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 45538daff5f8..d450f21bda34 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -108,8 +108,8 @@ def _has_mk_owned_shared_experts(self) -> bool: def _must_reduce_shared_expert_outputs(self) -> bool: return ( self._reduce_results - and self._quant_method.moe_mk is not None - and self._quant_method.moe_mk.output_is_reduced() + and self._quant_method.moe_kernel is not None + and self._quant_method.moe_kernel.output_is_reduced() ) def _determine_shared_experts_order( From e9afbe67b72707d32e5ad52b0f819d8fedf116e2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 5 Mar 2026 17:15:45 -0500 Subject: [PATCH 072/108] fix typos Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 772be9e150bb..02831d7d3f88 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -649,7 +649,7 @@ def _init_shared_experts(self) -> SharedExperts | None: # can property determine where the shared experts are supposed to be # called, i.e. by a MK or by the MoERunner. # Once the MK can be created upfront, we can just pass in the proper - # flags dervied from the quant_method's MK. + # flags derived from the quant_method's MK. reduce_results=self.reduce_results, quant_method=self.quant_method, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index d450f21bda34..134c500ebeea 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -223,4 +223,4 @@ def apply( assert self._output is not None self._output = self._maybe_reduce_shared_out(self._output) - # TODO(bnell): potentially do AFTER reduce here insteed of in runner. + # TODO(bnell): potentially do AFTER reduce here instead of in runner. From 453ab3d4ad5c000719acc35b2bee3c57810529fa Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 17:24:07 +0000 Subject: [PATCH 073/108] fix merge Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/config.py | 1 - vllm/model_executor/layers/fused_moe/runner/shared_experts.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b3f57e3e376a..83df5e45207b 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -942,7 +942,6 @@ def use_dp_chunking(self) -> bool: return ( self.use_deepep_ll_kernels or self.use_mori_kernels - or self.use_fi_all2allv_kernels or self.use_fi_nvl_two_sided_kernels or self.use_nixl_ep_kernels ) and envs.VLLM_ENABLE_MOE_DP_CHUNK diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 134c500ebeea..05258999ec8c 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -94,7 +94,7 @@ def _has_external_experts(self) -> bool: self._moe_config.moe_parallel_config.enable_eplb and backend != "allgather_reducescatter" ) - or self._moe_config.moe_parallel_config.use_fi_all2allv_kernels + or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels ) @property From 48acc59f18b12d4d2bffa4a60cbe3d2beecd0d17 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 18 Mar 2026 18:44:53 +0000 Subject: [PATCH 074/108] fix format Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/fp8.py | 5 +---- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 57431b946e24..59432c79d853 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -584,10 +584,7 @@ def make_fp8_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=( - shared_experts - if prepare_finalize.supports_async() else None - ), + shared_experts=shared_experts if prepare_finalize.supports_async() else None, inplace=( not moe_config.disable_inplace and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 194513981e41..89cacc8ca01e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -436,10 +436,7 @@ def make_nvfp4_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=( - shared_experts - if prepare_finalize.supports_async() else None - ), + shared_experts=shared_experts if prepare_finalize.supports_async() else None, inplace=False, ) From c0678440b0c6923cd205271ec9048fbe57caf007 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 19 Mar 2026 19:11:20 +0000 Subject: [PATCH 075/108] fix gate overlap Signed-off-by: Bill Nell --- .../fused_moe/runner/default_moe_runner.py | 17 +++++--- .../layers/fused_moe/runner/shared_experts.py | 42 ++++++++++++------- 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index c48744174d76..0f56da79a053 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -476,17 +476,22 @@ def _allocate_dp_chunking_outputs( return final_shared_hidden_states, final_fused_hidden_states - def _maybe_gate( + def _maybe_overlap_gate_with_shared_experts( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: # If router/gate provided, then apply it here. # (Note: This code runs only when "overlapped mode" is on to allow # parallel execution of shared experts with the FusedMoE via # separate cuda stream) + if self.shared_experts is not None: + self.shared_experts.maybe_setup_shared_experts_stream(shared_experts_input) + if self.gate is not None: router_logits, _ = self.gate(hidden_states) + return router_logits @property @@ -613,7 +618,11 @@ def forward_dispatch( # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() - router_logits = self._maybe_gate(hidden_states, router_logits) + router_logits = self._maybe_overlap_gate_with_shared_experts( + hidden_states, + router_logits, + shared_experts_input, + ) self._maybe_apply_shared_experts( shared_experts_input, @@ -655,10 +664,6 @@ def _forward_impl_chunked( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # Gate overlap not supported when chunking is enabled. Run the - # gate first. - router_logits = self._maybe_gate(hidden_states, router_logits) - final_shared_hidden_states, final_fused_hidden_states = ( self._allocate_dp_chunking_outputs(hidden_states, router_logits) ) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 05258999ec8c..3aea8201e509 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -148,24 +148,38 @@ def _determine_shared_experts_order( else: return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream + def maybe_setup_shared_experts_stream( + self, + shared_experts_input: torch.Tensor, + ): + experts_order, use_shared_experts_stream = self._determine_shared_experts_order( + shared_experts_input, + ) + + if ( + experts_order == SharedExpertsOrder.AFTER_QUANT_METHOD + and use_shared_experts_stream + ): + assert self._shared_experts_stream is not None + assert self._moe_config.disable_inplace + + # Record that the clone will be used by shared_experts_stream + # to avoid gc issue from deallocation of hidden_states_clone + # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 + # NOTE: We don't need shared_output.record_stream(current_stream()) + # because we synch the streams before using shared_output. + shared_experts_input.record_stream(self._shared_experts_stream) + + # Mark sync start point for the separate shared experts + # stream here since we want to run in parallel with the + # router/gate (next op below) + self._shared_experts_stream.wait_stream(current_stream()) + def _call_with_shared_experts_stream( self, shared_experts_input: torch.Tensor, ) -> torch.Tensor: - assert self._shared_experts_stream is not None - assert self._moe_config.disable_inplace - - # Record that the clone will be used by shared_experts_stream - # to avoid gc issue from deallocation of hidden_states_clone - # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 - # NOTE: We don't need shared_output.record_stream(current_stream()) - # because we synch the streams before using shared_output. - shared_experts_input.record_stream(self._shared_experts_stream) - - # Mark sync start point for the separate shared experts - # stream here since we want to run in parallel with the - # router/gate (next op below) - self._shared_experts_stream.wait_stream(current_stream()) + # TODO: assert that maybe_setup_shared_experts_stream has been called. # Run shared experts in parallel on a separate stream # NOTE: We start the separate stream here and mark the From bc8297856810510b9419f2a5360437da34b64330 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 19 Mar 2026 20:47:24 +0000 Subject: [PATCH 076/108] renames, revert lora changes Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 4 -- .../layers/fused_moe/runner/shared_experts.py | 39 +++++++++---------- 2 files changed, 18 insertions(+), 25 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 11332f4cb228..c4950cd75d07 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,11 +150,8 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), - self.base_layer.shared_experts, ) - # TODO: could be incorrect due to monolithic kernel? or add assert it - # is modular? if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.impl.fused_experts, @@ -341,7 +338,6 @@ def wrapper(*args, **kwargs): fused_experts = m_fused_moe_fn.impl.fused_experts - # TODO: seems like this could be done with modular kernel subclasses? m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 3aea8201e509..b8c70ac24fae 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -47,7 +47,7 @@ class SharedExpertsOrder(IntEnum): class SharedExperts: def __init__( self, - shared_experts: torch.nn.Module, + layer: torch.nn.Module, moe_config: FusedMoEConfig, quant_method: QuantizeMethodBase, reduce_results: bool, @@ -61,7 +61,7 @@ def __init__( assert isinstance(quant_method, FusedMoEMethodBase) self._output: torch.Tensor | None = None - self._shared_experts = shared_experts + self._layer = layer self._moe_config = moe_config self._quant_method = quant_method self._reduce_results = reduce_results @@ -73,12 +73,12 @@ def __init__( # and other execution modes if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local") - self._shared_experts_stream = None + self._stream = None else: # TODO(rob): enable shared expert overlap with non-cuda-alike. # aux_stream() returns None on non-cuda-alike platforms. - self._shared_experts_stream = aux_stream() - if self._shared_experts_stream is not None: + self._stream = aux_stream() + if self._stream is not None: logger.debug_once( "Enabled separate cuda stream for MoE shared_experts", scope="local" ) @@ -89,7 +89,7 @@ def _has_external_experts(self) -> bool: # - we are using eplb with non-default backend, because of correctness issues # - we are using flashinfer with DP, since there nothing to gain backend = self._moe_config.moe_parallel_config.all2all_backend - return self._shared_experts is not None and not ( + return self._layer is not None and not ( ( self._moe_config.moe_parallel_config.enable_eplb and backend != "allgather_reducescatter" @@ -99,10 +99,7 @@ def _has_external_experts(self) -> bool: @property def _has_mk_owned_shared_experts(self) -> bool: - return ( - not self._quant_method.mk_owns_shared_expert - and self._shared_experts is not None - ) + return not self._quant_method.mk_owns_shared_expert and self._layer is not None @property def _must_reduce_shared_expert_outputs(self) -> bool: @@ -116,7 +113,7 @@ def _determine_shared_experts_order( self, hidden_states: torch.Tensor, ) -> tuple[SharedExpertsOrder, bool]: - if self._shared_experts is None: + if self._layer is None: return SharedExpertsOrder.NONE, False if self._has_external_experts and not self._use_dp_chunking: @@ -132,7 +129,7 @@ def _determine_shared_experts_order( current_platform.is_cuda() and self._has_mk_owned_shared_experts and not self._use_dp_chunking - and self._shared_experts_stream is not None + and self._stream is not None and hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) @@ -160,7 +157,7 @@ def maybe_setup_shared_experts_stream( experts_order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream ): - assert self._shared_experts_stream is not None + assert self._stream is not None assert self._moe_config.disable_inplace # Record that the clone will be used by shared_experts_stream @@ -168,12 +165,12 @@ def maybe_setup_shared_experts_stream( # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 # NOTE: We don't need shared_output.record_stream(current_stream()) # because we synch the streams before using shared_output. - shared_experts_input.record_stream(self._shared_experts_stream) + shared_experts_input.record_stream(self._stream) # Mark sync start point for the separate shared experts # stream here since we want to run in parallel with the # router/gate (next op below) - self._shared_experts_stream.wait_stream(current_stream()) + self._stream.wait_stream(current_stream()) def _call_with_shared_experts_stream( self, @@ -186,11 +183,11 @@ def _call_with_shared_experts_stream( # sync end point immediately after it is done. This is # important to avoid excessive stream allocations by the cuda # graph replay later. - with torch.cuda.stream(self._shared_experts_stream): + with torch.cuda.stream(self._stream): # Note that hidden_states clone() is necessary here to avoid # conflict with the main stream - output = self._shared_experts(shared_experts_input) - current_stream().wait_stream(self._shared_experts_stream) + output = self._layer(shared_experts_input) + current_stream().wait_stream(self._stream) return output @@ -206,7 +203,7 @@ def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: @property def output(self) -> torch.Tensor | None: - assert (self._shared_experts is None) == (self._output is None) + assert (self._layer is None) == (self._output is None) output = self._output self._output = None return output @@ -223,13 +220,13 @@ def apply( if order != experts_order: return None - assert self._shared_experts is not None + assert self._layer is not None assert self._output is None if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: self._output = self._call_with_shared_experts_stream(shared_experts_input) else: - self._output = self._shared_experts(shared_experts_input) + self._output = self._layer(shared_experts_input) if order == SharedExpertsOrder.EXTERNAL: # TODO: figure out how to combine this with maybe_reduce_output? From 24da9fb59895e4bb99b27b0fa2d7a11896590b36 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 19 Mar 2026 21:55:38 +0000 Subject: [PATCH 077/108] fix lint Signed-off-by: Bill Nell --- tests/kernels/moe/test_zero_expert_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py index 24edca00b6af..d8f900256ec3 100644 --- a/tests/kernels/moe/test_zero_expert_moe.py +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -43,7 +43,7 @@ def zero_expert_moe(dist_init, default_vllm_config): vllm_config.compilation_config.static_forward_context = dict() with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): - init_workspace_manager(torch.cuda.current_device()) + init_workspace_manager(torch.accelerator.current_device_index()) layer = FusedMoE( zero_expert_type="identity", From 3dc9d4f6a24b65526acba8ec554e7cfdb4380834 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 20 Mar 2026 17:34:42 +0000 Subject: [PATCH 078/108] review comments + cleanup Signed-off-by: Bill Nell --- .../layers/fused_moe/fused_moe_method_base.py | 4 +- vllm/model_executor/layers/fused_moe/layer.py | 45 +++++++------------ .../layers/fused_moe/modular_kernel.py | 17 +++++-- .../layers/fused_moe/oracle/fp8.py | 2 +- .../layers/fused_moe/oracle/nvfp4.py | 2 +- .../fused_moe/runner/default_moe_runner.py | 22 +++++++-- .../layers/fused_moe/runner/moe_runner.py | 4 ++ .../layers/fused_moe/runner/shared_experts.py | 32 +++---------- 8 files changed, 62 insertions(+), 66 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py index 9350a9c7de74..8f40a63fa211 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_method_base.py @@ -39,9 +39,7 @@ def supports_internal_mk(self) -> bool: def mk_owns_shared_expert(self) -> bool: # NOTE(rob): temporary attribute to indicate support for # completed migration to the new internal MK interface. - return ( - self.moe_kernel is not None and self.moe_kernel.shared_experts is not None - ) + return self.moe_kernel is not None and self.moe_kernel.owns_shared_experts @abstractmethod def create_weights( diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 02831d7d3f88..4ed84a17fa8c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -42,9 +42,6 @@ from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( DefaultMoERunner, ) -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, -) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) @@ -339,7 +336,6 @@ def __init__( ): super().__init__() - self._gate = gate self._routed_input_transform = routed_input_transform if params_dtype is None: @@ -632,29 +628,17 @@ def _get_quant_method() -> FusedMoEMethodBase: moe_quant_params["intermediate_size_full"] = intermediate_size self.quant_method.create_weights(layer=self, **moe_quant_params) - self.base_quant_method = self.quant_method - self._shared_experts = shared_experts - self.shared_experts = self._init_shared_experts() - self.runner = self._init_runner() + # TODO(bnell): Why is this needed? Can probably be removed. + self.base_quant_method = self.quant_method - def _init_shared_experts(self) -> SharedExperts | None: - if self._shared_experts is None: - return None + self.runner = self._init_runner(gate, shared_experts) - return SharedExperts( - self._shared_experts, - moe_config=self.moe_config, - # Note: For now we must pass quant_method along to SharedExperts so it - # can property determine where the shared experts are supposed to be - # called, i.e. by a MK or by the MoERunner. - # Once the MK can be created upfront, we can just pass in the proper - # flags derived from the quant_method's MK. - reduce_results=self.reduce_results, - quant_method=self.quant_method, - ) - - def _init_runner(self) -> DefaultMoERunner: + def _init_runner( + self, + gate: torch.nn.Module | None, + shared_experts: torch.nn.Module | None, + ) -> DefaultMoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. @@ -663,8 +647,8 @@ def _init_runner(self) -> DefaultMoERunner: moe_config=self.moe_config, router=self.router, routed_input_transform=self._routed_input_transform, - gate=self._gate, - shared_experts=self.shared_experts, + gate=gate, + shared_experts=shared_experts, quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, @@ -678,7 +662,10 @@ def _replace_quant_method(self, mk: FusedMoEMethodBase): # We need to force reconstruction of runner because we're swapping out # the quant_method with a FusedMoEModularMethod. This logic can go # away once the FusedMoEModularMethod is eliminated. - self.runner = self._init_runner() + self.runner = self._init_runner( + self.runner.gate, + self.runner.shared_experts, + ) # Note: maybe_init_modular_kernel should only be called by # prepare_communication_buffer_for_model. @@ -706,7 +693,7 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.shared_experts, + self.runner.shared_experts, inplace=not self.moe_config.disable_inplace, ) ) @@ -741,7 +728,7 @@ def use_ep(self): @property def is_internal_router(self) -> bool: # By default, router/gate is called before FusedMoE forward pass - return self._gate is not None + return self.runner.is_internal_router() def _maybe_init_expert_routing_tables( self, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index e424d36536b7..d6e40b69eff2 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1000,7 +1000,10 @@ def __init__( ): self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts - self.shared_experts = shared_experts + # Only accept shared experts if they can be run w/async. + self.shared_experts = ( + shared_experts if prepare_finalize.supports_async() else None + ) self.inplace = inplace moe_parallel_config = fused_experts.moe_config.moe_parallel_config self.moe_parallel_config = moe_parallel_config @@ -1277,7 +1280,8 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - self._maybe_apply_shared_experts(shared_experts_input) + # TODO: remove + # self._maybe_apply_shared_experts(shared_experts_input) # TODO(lucas): refactor this in the alternative schedules followup # currently unpack if we have hook + receiver pair or just @@ -1461,7 +1465,6 @@ def __init__( inplace: bool = False, ): super().__init__() - self.shared_experts = shared_experts # NOTE: check if we can remove # Initialize the implementation (monolithic or modular). self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl @@ -1478,7 +1481,6 @@ def __init__( elif isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): - assert shared_experts is None assert not inplace self.impl = FusedMoEKernelMonolithicImpl( prepare_finalize, @@ -1494,6 +1496,13 @@ def __init__( self._post_init_setup() + @property + def owns_shared_experts(self) -> bool: + if isinstance(self.impl, FusedMoEKernelModularImpl): + return self.impl.shared_experts is not None + else: + return False + @property def is_monolithic(self) -> bool: return isinstance(self.impl, FusedMoEKernelMonolithicImpl) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 59432c79d853..3d9a499027cd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -584,7 +584,7 @@ def make_fp8_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=shared_experts if prepare_finalize.supports_async() else None, + shared_experts=shared_experts, inplace=( not moe_config.disable_inplace and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 89cacc8ca01e..940423694673 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -436,7 +436,7 @@ def make_nvfp4_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=shared_experts if prepare_finalize.supports_async() else None, + shared_experts=shared_experts, inplace=False, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 0f56da79a053..db7378f2e1b7 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -182,7 +182,7 @@ def __init__( router: FusedMoERouter, routed_input_transform: torch.nn.Module | None, gate: torch.nn.Module | None, - shared_experts: SharedExperts | None, + shared_experts: torch.nn.Module | None, quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, @@ -192,7 +192,21 @@ def __init__( self.router = router self.routed_input_transform = routed_input_transform self.gate = gate - self.shared_experts = shared_experts + + self.shared_experts: SharedExperts | None = None + if shared_experts is not None: + self.shared_experts = SharedExperts( + shared_experts, + moe_config=moe_config, + # Note: For now we must pass quant_method along to SharedExperts so it + # can property determine where the shared experts are supposed to be + # called, i.e. by a MK or by the MoERunner. + # Once the MK can be created upfront, we can just pass in the proper + # flags derived from the quant_method's MK. + reduce_results=reduce_results, + quant_method=quant_method, + ) + self.quant_method = quant_method self.reduce_results = reduce_results self.enable_dbo = enable_dbo @@ -233,6 +247,9 @@ def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: forward_impl_fn, ) + def is_internal_router(self) -> bool: + return self.gate is not None + def _maybe_init_dp_chunking(self): if not self.use_dp_chunking: return @@ -547,7 +564,6 @@ def _maybe_combine( hidden_states, dim=0, ) - # need RS for shared_output? if self.shared_experts is not None: assert shared_output is not None diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index b298cc2d0c4c..720e997cda36 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -32,3 +32,7 @@ def maybe_all_reduce_tensor_model_parallel( final_hidden_states: torch.Tensor, ): raise NotImplementedError + + @abstractmethod + def is_internal_router(self) -> bool: + raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index b8c70ac24fae..4d36a700d8f9 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -89,7 +89,7 @@ def _has_external_experts(self) -> bool: # - we are using eplb with non-default backend, because of correctness issues # - we are using flashinfer with DP, since there nothing to gain backend = self._moe_config.moe_parallel_config.all2all_backend - return self._layer is not None and not ( + return not ( ( self._moe_config.moe_parallel_config.enable_eplb and backend != "allgather_reducescatter" @@ -97,10 +97,6 @@ def _has_external_experts(self) -> bool: or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels ) - @property - def _has_mk_owned_shared_experts(self) -> bool: - return not self._quant_method.mk_owns_shared_expert and self._layer is not None - @property def _must_reduce_shared_expert_outputs(self) -> bool: return ( @@ -113,37 +109,24 @@ def _determine_shared_experts_order( self, hidden_states: torch.Tensor, ) -> tuple[SharedExpertsOrder, bool]: - if self._layer is None: - return SharedExpertsOrder.NONE, False - if self._has_external_experts and not self._use_dp_chunking: return SharedExpertsOrder.EXTERNAL, False - if ( - not self._has_mk_owned_shared_experts - or not self._moe_config.moe_parallel_config.use_all2all_kernels - ): + if self._quant_method.mk_owns_shared_expert: return SharedExpertsOrder.INTERNAL, False allow_shared_experts_stream = ( current_platform.is_cuda() - and self._has_mk_owned_shared_experts and not self._use_dp_chunking and self._stream is not None and hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) - # Check if we need to run shared experts before matrix multiply because - # matrix multiply may modify the hidden_states. - run_shared_experts_before = ( - self._has_mk_owned_shared_experts and not allow_shared_experts_stream - ) - - if run_shared_experts_before: - return SharedExpertsOrder.BEFORE_QUANT_METHOD, False + if allow_shared_experts_stream: + return SharedExpertsOrder.AFTER_QUANT_METHOD, True else: - return SharedExpertsOrder.AFTER_QUANT_METHOD, allow_shared_experts_stream + return SharedExpertsOrder.BEFORE_QUANT_METHOD, False def maybe_setup_shared_experts_stream( self, @@ -202,8 +185,8 @@ def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: return shared_out @property - def output(self) -> torch.Tensor | None: - assert (self._layer is None) == (self._output is None) + def output(self) -> torch.Tensor: + assert self._output is not None output = self._output self._output = None return output @@ -220,7 +203,6 @@ def apply( if order != experts_order: return None - assert self._layer is not None assert self._output is None if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: From 12bda3d877e68d9fde63aafdbcf0432f0bf4a1d6 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 20 Mar 2026 17:39:43 +0000 Subject: [PATCH 079/108] remove _must_reduce_shared_expert_outputs Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/shared_experts.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 4d36a700d8f9..3aafb0eb925b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -97,14 +97,6 @@ def _has_external_experts(self) -> bool: or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels ) - @property - def _must_reduce_shared_expert_outputs(self) -> bool: - return ( - self._reduce_results - and self._quant_method.moe_kernel is not None - and self._quant_method.moe_kernel.output_is_reduced() - ) - def _determine_shared_experts_order( self, hidden_states: torch.Tensor, @@ -178,7 +170,9 @@ def _maybe_reduce_shared_out(self, shared_out: torch.Tensor) -> torch.Tensor: # Reduce shared expert outputs if necessary, since the MLP # should have been created with reduce_results=False. if ( - self._must_reduce_shared_expert_outputs + self._reduce_results + and self._quant_method.moe_kernel is not None + and self._quant_method.moe_kernel.output_is_reduced() and get_tensor_model_parallel_world_size() > 1 ): shared_out = tensor_model_parallel_all_reduce(shared_out) From 8aaddea26bf3f3cbc0a483ac5df48ead75f5c69f Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Mar 2026 17:38:19 +0000 Subject: [PATCH 080/108] undo some changes + add Rob's changes Signed-off-by: Bill Nell --- .../fused_moe/deepep_ll_prepare_finalize.py | 2 +- vllm/model_executor/layers/fused_moe/layer.py | 58 ++++++++++++++----- .../layers/fused_moe/modular_kernel.py | 8 +-- .../fused_moe/runner/default_moe_runner.py | 25 ++------ .../layers/fused_moe/runner/shared_experts.py | 47 +++++++-------- 5 files changed, 75 insertions(+), 65 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py index a22b89415364..5bd3f627ab1f 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py @@ -306,7 +306,7 @@ def prepare_async( **(dict(use_nvfp4=True) if use_nvfp4 else dict()), **( dict(x_global_scale=qc_a1_gscale_or_scale) - if qc_a1_gscale_or_scale is not None + if qc_a1_gscale_or_scale is not None and nvfp4_dispatch else dict() ), async_finish=False, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 4ed84a17fa8c..456643c8c39e 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -42,6 +42,9 @@ from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( DefaultMoERunner, ) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) @@ -632,23 +635,43 @@ def _get_quant_method() -> FusedMoEMethodBase: # TODO(bnell): Why is this needed? Can probably be removed. self.base_quant_method = self.quant_method - self.runner = self._init_runner(gate, shared_experts) + # Note: for now, the layer must keep _gate and _shared_experts. + # This is because a number of locations swap out the quant_method + # which requires re-initializing the SharedExperts and DefaultMoERunner. + # Once we've figured out alternatives to swapping out the quant_method, + # we can move ownership of _gate and _shared_experts into the runner. + self._gate = gate + self._shared_experts = shared_experts + self.runner = self._init_runner() + + def _init_shared_experts(self) -> SharedExperts | None: + if self._shared_experts is None: + return None - def _init_runner( - self, - gate: torch.nn.Module | None, - shared_experts: torch.nn.Module | None, - ) -> DefaultMoERunner: + return SharedExperts( + self._shared_experts, + moe_config=self.moe_config, + # Note: For now we must pass quant_method along to SharedExperts so it + # can property determine where the shared experts are supposed to be + # called, i.e. by a MK or by the MoERunner. + # Once the MK can be created upfront, we can just pass in the proper + # flags derived from the quant_method's MK. + reduce_results=self.reduce_results, + quant_method=self.quant_method, + ) + + def _init_runner(self) -> DefaultMoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. + self.shared_experts = self._init_shared_experts() return DefaultMoERunner( layer=self, moe_config=self.moe_config, router=self.router, routed_input_transform=self._routed_input_transform, - gate=gate, - shared_experts=shared_experts, + gate=self._gate, + shared_experts=self.shared_experts, quant_method=self.quant_method, reduce_results=self.reduce_results, enable_dbo=self.vllm_config.parallel_config.enable_dbo, @@ -662,10 +685,7 @@ def _replace_quant_method(self, mk: FusedMoEMethodBase): # We need to force reconstruction of runner because we're swapping out # the quant_method with a FusedMoEModularMethod. This logic can go # away once the FusedMoEModularMethod is eliminated. - self.runner = self._init_runner( - self.runner.gate, - self.runner.shared_experts, - ) + self.runner = self._init_runner() # Note: maybe_init_modular_kernel should only be called by # prepare_communication_buffer_for_model. @@ -693,7 +713,7 @@ def maybe_init_modular_kernel(self) -> None: self, self.base_quant_method, prepare_finalize, - self.runner.shared_experts, + self.shared_experts, inplace=not self.moe_config.disable_inplace, ) ) @@ -1422,7 +1442,12 @@ def _maybe_make_contiguous( assert all( weight.is_contiguous() for name, weight in weights - if not (name.startswith("_shared_experts.") or name.startswith("_gate.")) + if not ( + name.startswith("_shared_experts.") + or name.startswith("_gate.") + or name.startswith("_routed_input_transform.") + or name.startswith("_routed_output_transform.") + ) and name not in NON_EXPERT_WEIGHTS ) @@ -1432,8 +1457,11 @@ def _maybe_make_contiguous( if name not in NON_EXPERT_WEIGHTS and weight.shape != torch.Size([]) and not name.startswith("_shared_experts.") - # exclude parameters from non-expert submodules (e.g. gate/shared) + # exclude parameters from non-expert submodules, + # e.g. gate/shared/transforms. and not name.startswith("_gate.") + and not name.startswith("_routed_input_transform.") + and not name.startswith("_routed_output_transform.") ] def set_eplb_state( diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d6e40b69eff2..870d24a3c479 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1001,6 +1001,8 @@ def __init__( self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts # Only accept shared experts if they can be run w/async. + # The MoERunner/SharedExperts class will coordinate with the MK to ensure + # that the SharedExperts are executed only once. self.shared_experts = ( shared_experts if prepare_finalize.supports_async() else None ) @@ -1084,7 +1086,7 @@ def _maybe_apply_shared_experts( assert shared_experts_input is not None self.shared_experts.apply( shared_experts_input, - SharedExpertsOrder.INTERNAL, + SharedExpertsOrder.MK_INTERNAL_OVERLAPPED, ) def _prepare( @@ -1270,7 +1272,6 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - self._maybe_apply_shared_experts(shared_experts_input) else: finalize_ret = self.prepare_finalize.finalize_async( output, @@ -1280,8 +1281,7 @@ def _finalize( apply_router_weight_on_input, self.fused_experts.finalize_weight_and_reduce_impl(), ) - # TODO: remove - # self._maybe_apply_shared_experts(shared_experts_input) + self._maybe_apply_shared_experts(shared_experts_input) # TODO(lucas): refactor this in the alternative schedules followup # currently unpack if we have hook + receiver pair or just diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index db7378f2e1b7..41a5c69e60d1 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -182,7 +182,7 @@ def __init__( router: FusedMoERouter, routed_input_transform: torch.nn.Module | None, gate: torch.nn.Module | None, - shared_experts: torch.nn.Module | None, + shared_experts: SharedExperts | None, quant_method: FusedMoEMethodBase, reduce_results: bool, enable_dbo: bool, @@ -192,21 +192,7 @@ def __init__( self.router = router self.routed_input_transform = routed_input_transform self.gate = gate - - self.shared_experts: SharedExperts | None = None - if shared_experts is not None: - self.shared_experts = SharedExperts( - shared_experts, - moe_config=moe_config, - # Note: For now we must pass quant_method along to SharedExperts so it - # can property determine where the shared experts are supposed to be - # called, i.e. by a MK or by the MoERunner. - # Once the MK can be created upfront, we can just pass in the proper - # flags derived from the quant_method's MK. - reduce_results=reduce_results, - quant_method=quant_method, - ) - + self.shared_experts = shared_experts self.quant_method = quant_method self.reduce_results = reduce_results self.enable_dbo = enable_dbo @@ -420,9 +406,10 @@ def _apply_quant_method( shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: # Run this before quant_method to avoid inplace issues. + # TODO(bnell): probably not needed anymore since inplace is + # disabled when shared experts are present. self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.BEFORE_QUANT_METHOD, + shared_experts_input, SharedExpertsOrder.NO_OVERLAP ) if self.quant_method.is_monolithic: @@ -447,7 +434,7 @@ def _apply_quant_method( self._maybe_apply_shared_experts( shared_experts_input, - SharedExpertsOrder.AFTER_QUANT_METHOD, + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, ) return ( diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 3aafb0eb925b..c4333bcf9f19 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -34,14 +34,14 @@ class SharedExpertsOrder(IntEnum): # get rid of it after _moe_forward is undone. EXTERNAL = (1,) - # Called by modular kernel. - INTERNAL = (2,) + # No overlap - defensively called before MK. + NO_OVERLAP = (2,) - # Called right before quant_method is executed. - BEFORE_QUANT_METHOD = (3,) + # Overlapped with dispatch/combine in DP/EP - called by the MK. + MK_INTERNAL_OVERLAPPED = (3,) - # Called right after quant_method is executed (possibly with streaming). - AFTER_QUANT_METHOD = (4,) + # Overlapped with the gate, router, experts in aux stream. + MULTI_STREAM_OVERLAPPED = (4,) class SharedExperts: @@ -100,14 +100,14 @@ def _has_external_experts(self) -> bool: def _determine_shared_experts_order( self, hidden_states: torch.Tensor, - ) -> tuple[SharedExpertsOrder, bool]: + ) -> SharedExpertsOrder: if self._has_external_experts and not self._use_dp_chunking: - return SharedExpertsOrder.EXTERNAL, False + return SharedExpertsOrder.EXTERNAL if self._quant_method.mk_owns_shared_expert: - return SharedExpertsOrder.INTERNAL, False + return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED - allow_shared_experts_stream = ( + should_run_shared_in_aux_stream = ( current_platform.is_cuda() and not self._use_dp_chunking and self._stream is not None @@ -115,23 +115,18 @@ def _determine_shared_experts_order( <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD ) - if allow_shared_experts_stream: - return SharedExpertsOrder.AFTER_QUANT_METHOD, True + if should_run_shared_in_aux_stream: + return SharedExpertsOrder.MULTI_STREAM_OVERLAPPED else: - return SharedExpertsOrder.BEFORE_QUANT_METHOD, False + return SharedExpertsOrder.NO_OVERLAP def maybe_setup_shared_experts_stream( self, shared_experts_input: torch.Tensor, ): - experts_order, use_shared_experts_stream = self._determine_shared_experts_order( - shared_experts_input, - ) + experts_order = self._determine_shared_experts_order(shared_experts_input) - if ( - experts_order == SharedExpertsOrder.AFTER_QUANT_METHOD - and use_shared_experts_stream - ): + if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: assert self._stream is not None assert self._moe_config.disable_inplace @@ -147,7 +142,7 @@ def maybe_setup_shared_experts_stream( # router/gate (next op below) self._stream.wait_stream(current_stream()) - def _call_with_shared_experts_stream( + def _run_in_aux_stream( self, shared_experts_input: torch.Tensor, ) -> torch.Tensor: @@ -190,17 +185,15 @@ def apply( shared_experts_input: torch.Tensor, order: SharedExpertsOrder, ): - experts_order, use_shared_experts_stream = self._determine_shared_experts_order( - shared_experts_input, - ) + experts_order = self._determine_shared_experts_order(shared_experts_input) if order != experts_order: return None assert self._output is None - if order == SharedExpertsOrder.AFTER_QUANT_METHOD and use_shared_experts_stream: - self._output = self._call_with_shared_experts_stream(shared_experts_input) + if order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: + self._output = self._run_in_aux_stream(shared_experts_input) else: self._output = self._layer(shared_experts_input) @@ -210,4 +203,6 @@ def apply( assert self._output is not None self._output = self._maybe_reduce_shared_out(self._output) + assert self._output is not None + # TODO(bnell): potentially do AFTER reduce here instead of in runner. From 392f311bb697518d1e15bd9f0f4ffcd1e83358e1 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Mar 2026 18:59:52 +0000 Subject: [PATCH 081/108] hacky fix for unquantized method Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 30 ++++++++++++++++--- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c312ca3bd5a2..aacdc19dad90 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -614,13 +614,35 @@ def _get_quant_method() -> FusedMoEMethodBase: # we can move ownership of _gate and _shared_experts into the runner. self._gate = gate self._shared_experts = shared_experts + self.shared_experts = None self.runner = self._init_runner() - def _init_shared_experts(self) -> SharedExperts | None: + def _init_shared_experts(self): # -> SharedExperts | None: if self._shared_experts is None: - return None + return + + # Note: If the SharedExperts already exist, we reinitialize + # them in place. This is because the MK might be holding a + # reference to the same SharedExperts object. If we create a + # new instance, the MK will still be holding onto the old one, + # including the old quant_method. This is a workaround for + # UnquantizedFusedMoEMethod's handling of MK initialization + # which should be fixed by #36732. + if self.shared_experts is not None: + self.shared_experts.__init__( + self._shared_experts, + moe_config=self.moe_config, + # Note: For now we must pass quant_method along to SharedExperts so it + # can property determine where the shared experts are supposed to be + # called, i.e. by a MK or by the MoERunner. + # Once the MK can be created upfront, we can just pass in the proper + # flags derived from the quant_method's MK. + reduce_results=self.reduce_results, + quant_method=self.quant_method, + ) + return - return SharedExperts( + self.shared_experts = SharedExperts( self._shared_experts, moe_config=self.moe_config, # Note: For now we must pass quant_method along to SharedExperts so it @@ -636,7 +658,7 @@ def _init_runner(self) -> DefaultMoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. - self.shared_experts = self._init_shared_experts() + self._init_shared_experts() return DefaultMoERunner( layer=self, moe_config=self.moe_config, From 7d5adbec163461b7e2323860f609c08f40ce3aa9 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 23 Mar 2026 22:02:18 +0000 Subject: [PATCH 082/108] fix lint Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 4 ++-- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index aacdc19dad90..32802c2239de 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -614,7 +614,7 @@ def _get_quant_method() -> FusedMoEMethodBase: # we can move ownership of _gate and _shared_experts into the runner. self._gate = gate self._shared_experts = shared_experts - self.shared_experts = None + self.shared_experts: SharedExperts | None = None self.runner = self._init_runner() def _init_shared_experts(self): # -> SharedExperts | None: @@ -629,7 +629,7 @@ def _init_shared_experts(self): # -> SharedExperts | None: # UnquantizedFusedMoEMethod's handling of MK initialization # which should be fixed by #36732. if self.shared_experts is not None: - self.shared_experts.__init__( + self.shared_experts.__init__( # type: ignore self._shared_experts, moe_config=self.moe_config, # Note: For now we must pass quant_method along to SharedExperts so it diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 9d1c8e27b7f5..77a53a6c3e5f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -859,7 +859,6 @@ def make_mxfp4_moe_kernel( if moe_config.moe_parallel_config.use_deepep_ll_kernels else None ), - moe_parallel_config=moe_config.moe_parallel_config, inplace=( not moe_config.disable_inplace and mxfp4_backend not in TRTLLM_BACKENDS ), From f3451658df968639419c3445d6be4ab9ad4fbadf Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 12 Feb 2026 18:03:08 -0500 Subject: [PATCH 083/108] fix lint Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 4 + vllm/model_executor/layers/fused_moe/layer.py | 14 +- .../layers/fused_moe/modular_kernel.py | 2 +- .../fused_moe/router/memoizing_router.py | 35 + .../fused_moe/runner/chunking_moe_runner.py | 219 ++++++ .../fused_moe/runner/default_moe_runner.py | 665 +----------------- .../fused_moe/runner/moe_runner_base.py | 501 +++++++++++++ .../fused_moe/runner/moe_runner_factory.py | 51 ++ 8 files changed, 853 insertions(+), 638 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index c4950cd75d07..11332f4cb228 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,8 +150,11 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), + self.base_layer.shared_experts, ) + # TODO: could be incorrect due to monolithic kernel? or add assert it + # is modular? if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.impl.fused_experts, @@ -338,6 +341,7 @@ def wrapper(*args, **kwargs): fused_experts = m_fused_moe_fn.impl.fused_experts + # TODO: seems like this could be done with modular kernel subclasses? m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 32802c2239de..5a83d5138168 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -39,8 +39,14 @@ from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) -from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( - DefaultMoERunner, +from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( + MoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import ( + create_moe_runner, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, ) from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -654,12 +660,12 @@ def _init_shared_experts(self): # -> SharedExperts | None: quant_method=self.quant_method, ) - def _init_runner(self) -> DefaultMoERunner: + def _init_runner(self) -> MoERunner: # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. self._init_shared_experts() - return DefaultMoERunner( + return create_moe_runner( layer=self, moe_config=self.moe_config, router=self.router, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 870d24a3c479..f51c7fa53b43 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -995,7 +995,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, - shared_experts: SharedExperts | None, + shared_experts: SharedExperts | None = None, inplace: bool = False, ): self.prepare_finalize = prepare_finalize diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py new file mode 100644 index 000000000000..a55bd2f09d6f --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter + + +class MemoizingRouter(FusedMoERouter): + def __init__(self, router: FusedMoERouter): + self.router = router + + def set_capture_fn( + self, + capture_fn: Callable[[torch.Tensor], None] | None, + ) -> None: + self.router.set_capture_fn(capture_fn) + self.results: tuple[torch.Tensor, torch.Tensor] | None = None + + @property + def routing_method_type(self) -> RoutingMethodType: + return self.router.routing_method_type + + def select_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.results is None: + self.results = self.router.select_experts(hidden_states, router_logits) + return self.results diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py new file mode 100644 index 000000000000..6626fad557b3 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.forward_context import ( + get_forward_context, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.ubatching import dbo_current_ubatch_id +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + + +class ChunkingMoERunner(MoERunnerBase): + """ + MoE runner wrapper that adds chunked processing to any MoERunnerBase. + + This runner wraps an inner MoERunnerBase and overrides _forward_impl to + process large batches by breaking them into smaller chunks. Each chunk + is delegated to the inner runner's _forward_impl, making chunking + composable with any runner implementation. + + All MoERunnerBase state (moe_config, router, quant_method, etc.) is + transparently delegated to the inner runner via __getattr__. + ChunkingMoERunner only owns chunking-specific state: the pre-allocated + workspace buffers and the reduce_results override. + + Key behaviors: + - Pre-allocates workspace tensors for CUDA graph compatibility + - Processes chunks via inner._forward_impl per chunk + - Never reduces results (reduce_results always returns False) + """ + + def __init__(self, inner: MoERunnerBase): + # Assert that _maybe_dispatch/_maybe_combine will be nops. + assert inner.moe_config.pcp_size == 1 + + # Skip MoERunnerBase.__init__ — all state is delegated to inner + # via __getattr__. Only chunking-specific state lives here. + self._inner = inner + + # Pre-allocated staging buffers. These need to exist ahead of time + # due to CUDA graph construction needing fixed buffer addresses. + self.batched_hidden_states, self.batched_router_logits = ( + self._init_dp_chunking() + ) + + def __getattr__(self, name): + # Delegate attribute access to the inner runner. This is only + # called when normal lookup (instance __dict__, class MRO) fails, + # so ChunkingMoERunner's own attributes and methods take priority. + return getattr(self._inner, name) + + @property + def reduce_results(self) -> bool: + return False + + def _init_dp_chunking(self) -> list[torch.Tensor]: + states_shape: tuple[int, ...] + logits_shape: tuple[int, ...] + + moe = self.moe_config + + if self.enable_dbo: + states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) + logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) + else: + states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) + logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) + + # Does this need some kind of profiling run check like modular_kernel.py? + return current_workspace_manager().get_simultaneous( + (states_shape, moe.in_dtype), + (logits_shape, moe.router_logits_dtype), + ) + + def _allocate_dp_chunking_outputs( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + # Assert the inputs are of the proper type and shape. + assert self.batched_hidden_states is not None + assert self.batched_router_logits is not None + + assert self.batched_hidden_states.dtype == hidden_states.dtype, ( + f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" + ) + assert self.batched_router_logits.dtype == router_logits.dtype, ( + f"{self.batched_router_logits.dtype} == {router_logits.dtype}" + ) + + # Check size compatibility. + assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) + assert self.batched_router_logits.size(-1) == router_logits.size(-1) + + final_fused_hidden_states = torch.empty_like(hidden_states) + if self.shared_experts is not None: + final_shared_hidden_states = torch.empty_like(hidden_states) + else: + final_shared_hidden_states = None + + return final_shared_hidden_states, final_fused_hidden_states + + def _slice_and_copy_input( + self, + out_slice: torch.Tensor, + orig: torch.Tensor | None, + start: int, + end: int, + ) -> torch.Tensor: + assert orig is not None + slice_size = end - start + orig_slice = orig[start:end, :] + if self.enable_dbo: + assert out_slice.dim() == 3 + batch_buffer_idx = dbo_current_ubatch_id() + out_slice = out_slice[batch_buffer_idx, :] + + assert out_slice.size(0) >= slice_size + out_slice = out_slice[:slice_size, :] + out_slice.copy_(orig_slice, non_blocking=True) + return out_slice + + def _forward_impl( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + final_shared_hidden_states, final_fused_hidden_states = ( + self._allocate_dp_chunking_outputs(hidden_states, router_logits) + ) + + ctx = get_forward_context() + # flashinfer_cutlass_kernels can handle: optional DP + TP/EP + max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu + moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens + + # If the input to the MoE is sequence parallel then divide by sp_size + # to find the maximum number of tokens for any individual dispatcher. + if self.moe_config.is_sequence_parallel: + max_tokens_across_dispatchers = cdiv( + max_tokens_across_dispatchers, self.moe_config.sp_size + ) + + num_tokens = hidden_states.size(0) + for chunk_idx, chunk_start_ in enumerate( + range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) + ): + chunk_start = chunk_start_ + chunk_end = min( + chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers + ) + # clamp start and end + chunk_start = min(chunk_start, num_tokens - 1) + chunk_end = min(chunk_end, num_tokens) + chunk_sizes = ctx.dp_metadata.chunked_sizes( + self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx + ) + with chunk_sizes: + hidden_states_chunk = self._slice_and_copy_input( + self.batched_hidden_states, + hidden_states, + chunk_start, + chunk_end, + ) + + router_logits_chunk = self._slice_and_copy_input( + self.batched_router_logits, + router_logits, + chunk_start, + chunk_end, + ) + + shared_experts_input_chunk = ( + shared_experts_input[chunk_start:chunk_end, :] + if shared_experts_input is not None + else None + ) + + # Delegate per-chunk computation to the inner runner. + chunk_result = self._inner._forward_impl( + layer=layer, + hidden_states=hidden_states_chunk, + router_logits=router_logits_chunk, + shared_experts_input=shared_experts_input_chunk, + ) + + # Store outputs + # TODO(bnell): document when chunk_start >= num_tokens + if chunk_start < num_tokens: + if self.shared_experts is not None: + assert isinstance(chunk_result, tuple) + shared_output_chunk, hidden_states_chunk = chunk_result + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + hidden_states_chunk, non_blocking=True + ) + assert shared_output_chunk is not None + assert final_shared_hidden_states is not None + final_shared_hidden_states[chunk_start:chunk_end, :].copy_( + shared_output_chunk, non_blocking=True + ) + else: + assert isinstance(chunk_result, torch.Tensor) + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + chunk_result, non_blocking=True + ) + + if self.shared_experts is None: + return final_fused_hidden_states + else: + assert final_shared_hidden_states is not None + return (final_shared_hidden_states, final_fused_hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 41a5c69e60d1..07d63cab2220 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -1,21 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable -from contextlib import nullcontext -from typing import TYPE_CHECKING import torch -import torch.nn.functional as F from vllm.distributed import ( get_ep_group, get_pcp_group, - tensor_model_parallel_all_reduce, -) -from vllm.forward_context import ( - ForwardContext, - get_forward_context, - is_forward_context_available, ) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( @@ -27,152 +17,39 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, - SharedExpertsOrder, -) -from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import ( - HAS_OPAQUE_TYPE, - ModuleName, - direct_register_custom_op, ) -from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) -def get_layer_from_name(layer_name: str) -> torch.nn.Module: - forward_context: ForwardContext = get_forward_context() - if layer_name == "from_forward_context": - all_moe_layers = forward_context.all_moe_layers - assert all_moe_layers is not None - moe_layer_index = forward_context.moe_layer_index - if moe_layer_index >= len(all_moe_layers): - raise AssertionError( - "We expected the number of MOE layers in `all_moe_layers` " - "to be equal to the number of " - "{vllm.moe_forward, vllm.moe_forward_shared} calls." - ) - layer_name = all_moe_layers[moe_layer_index] - forward_context.moe_layer_index += 1 - return forward_context.no_compile_layers[layer_name] - - -# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; -# on older versions it remains a plain str. -if TYPE_CHECKING: - from typing import TypeAlias - - _layer_name_type: TypeAlias = str | ModuleName -else: - _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str - - -def _resolve_layer_name(layer_name: str | ModuleName) -> str: - return layer_name.value if isinstance(layer_name, ModuleName) else layer_name - - -# Note: _moe_forward and _moe_forward_shared should not contain any -# implementation details, They should merely pass along control to -# the runner's 'forward_dispatch' method. -def _moe_forward( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> torch.Tensor: - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -def _moe_forward_shared( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_shared_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: - # Output shapes: - # - fused_out: same as hidden_states (routed experts use transformed size) - # - shared_out: same as shared_experts_input if provided, else same as - # hidden_states - # (For latent MoE: shared experts use original hidden_size, not latent size) - fused_out = torch.empty_like(hidden_states) - if shared_experts_input is not None: - shared_out = torch.empty_like(shared_experts_input) - else: - shared_out = torch.empty_like(hidden_states) - return shared_out, fused_out - - -direct_register_custom_op( - op_name="moe_forward", - op_func=_moe_forward, - mutates_args=["hidden_states"], # is this still true? - fake_impl=_moe_forward_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -direct_register_custom_op( - op_name="moe_forward_shared", - op_func=_moe_forward_shared, - fake_impl=_moe_forward_shared_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -class DefaultMoERunner(MoERunner): +class DefaultMoERunner(MoERunnerBase): """ - Default implementation of the MoE runner for executing Mixture of Experts layers. + Standard MoE runner implementation for executing Mixture of Experts layers. - This class provides a comprehensive implementation for running MoE computations - with support for: - - Expert routing and token dispatching + This is the primary concrete implementation of MoE execution logic, providing + comprehensive support for standard MoE operations. It handles: + - Expert routing and token dispatching using various routing strategies - Shared experts computation with optional parallel execution using CUDA streams - - Data parallel (DP) chunking for large batch processing - Tensor model parallel and expert parallel operations - - Various quantization methods and custom operators + - Multiple quantization methods and optimized kernel selection - Both monolithic and decomposed expert execution paths + - Integration with various parallel execution modes (TP, EP, DP) - The runner handles the complete MoE forward pass including routing tokens to - experts, executing expert computations, and combining results. It supports - advanced features like overlapped execution of shared experts and optimized - kernels for different parallel execution modes. + The runner orchestrates the complete MoE forward pass including routing tokens + to experts, executing expert computations in parallel, and combining results. + It supports advanced features like overlapped execution of shared experts, + optimized kernels for different parallel configurations, and seamless + integration with vLLM's distributed execution framework. - Eventually, this class will be split up and specialized for different - configurations, e.g. the presence or absence of shared experts, a gate, etc. + This implementation is suitable for most standard MoE use cases. For specialized + scenarios like large batch chunking, alternative runners like ChunkingMoERunner + may be more appropriate. + + Eventually, this class may be split into more specialized implementations + for different configurations (e.g., with/without shared experts, gates, etc.). """ def __init__( @@ -187,316 +64,24 @@ def __init__( reduce_results: bool, enable_dbo: bool, ): - super().__init__() - self.moe_config = moe_config - self.router = router - self.routed_input_transform = routed_input_transform - self.gate = gate - self.shared_experts = shared_experts - self.quant_method = quant_method - self.reduce_results = reduce_results - self.enable_dbo = enable_dbo - self.enable_eplb = moe_config.moe_parallel_config.enable_eplb - - # Chunked all2all staging tensor - # These need to exist ahead of time due to CUDAgraph construction - # needing a fixed buffer address. - self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking - self.batched_hidden_states: torch.Tensor | None = None - self.batched_router_logits: torch.Tensor | None = None - self._maybe_init_dp_chunking() - - # Needed for string -> FusedMoE layer lookup in custom ops. - self.layer_name = layer.layer_name - - self.forward_entry, self.forward_impl = self._select_forward(layer) - - def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: - # Select implementation based on presence of DP chunking. - forward_impl_fn = ( - self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl - ) - - if current_platform.is_tpu() or current_platform.is_cpu(): - # TODO: Once the OOM issue for the TPU backend is resolved, we - # will switch to using the moe_forward custom op. - # Note: CPU doesn't require wrapped forward_impl. - return ( - _moe_forward if self.shared_experts is None else _moe_forward_shared, - forward_impl_fn, - ) - - return ( - torch.ops.vllm.moe_forward - if self.shared_experts is None - else torch.ops.vllm.moe_forward_shared, - forward_impl_fn, + super().__init__( + layer, + moe_config, + router, + routed_input_transform, + gate, + shared_experts, + quant_method, + reduce_results, + enable_dbo, ) def is_internal_router(self) -> bool: return self.gate is not None - def _maybe_init_dp_chunking(self): - if not self.use_dp_chunking: - return - - assert self.batched_hidden_states is None - states_shape: tuple[int, ...] - logits_shape: tuple[int, ...] - - moe = self.moe_config - - if self.enable_dbo: - states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) - else: - states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) - - device = torch.accelerator.current_device_index() - self.batched_hidden_states = torch.zeros( - states_shape, - dtype=moe.in_dtype, - device=device, - ) - - self.batched_router_logits = torch.zeros( - logits_shape, - dtype=moe.router_logits_dtype, - device=device, - ) - - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - if self.must_reduce_shared_expert_outputs(): - return final_hidden_states - else: - return tensor_model_parallel_all_reduce(final_hidden_states) - - def apply_routed_input_transform( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Apply transform for routed experts (e.g., latent projection). - - This is called by FusedMoE.forward_native. The original hidden_states - is saved separately so shared experts get [S, hidden_size] while - routed experts get the transformed [S, moe_latent_size]. - - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside SharedFusedMoE to all-reduce on the smaller latent - dimension. - - Returns (possibly transformed) hidden states and the input for shared - experts (or None if there are no shared experts). - """ - if self.routed_input_transform is not None: - result = self.routed_input_transform(hidden_states) - # ReplicatedLinear returns (output, extra_bias) tuple. - # We only need the output tensor; extra_bias is not used here. - if isinstance(result, tuple): - return result[0], hidden_states - return result, hidden_states - - return ( - hidden_states, - hidden_states if self.shared_experts is not None else None, - ) - - def _maybe_reduce_output( - self, - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - trunc_sizes: list[int], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return x[..., :trunc_size] - - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) - - if ( - not self.moe_config.is_sequence_parallel - and not self.use_dp_chunking - and self.reduce_results - and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - ): - func = reduce_and_trunc - else: - func = trunc - - if isinstance(states, tuple): - return tuple( - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] - ) - else: - assert len(trunc_sizes) == 1 - return func(states, trunc_sizes[0]) - - def _encode_layer_name(self) -> str | ModuleName: - if HAS_OPAQUE_TYPE: - return ModuleName(self.layer_name) - # Can be unavailable or None in unittests - if ( - is_forward_context_available() - and get_forward_context().all_moe_layers is not None - ): - return "from_forward_context" - return self.layer_name - - def _maybe_pad_hidden_states( - self, - shared_experts_input: torch.Tensor | None, - hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, list[int]]: - shared_experts_hidden_dim = ( - shared_experts_input.shape[-1] if shared_experts_input is not None else 0 - ) - transformed_hidden_dim = hidden_states.shape[-1] - if ( - not self.quant_method.skip_forward_padding - and self.moe_config.hidden_dim != transformed_hidden_dim - ): - hidden_states = F.pad( - hidden_states, - (0, self.moe_config.hidden_dim - transformed_hidden_dim), - mode="constant", - value=0.0, - ) - - if self.shared_experts is not None: - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] - else: - orig_hidden_dims = [transformed_hidden_dim] - - return hidden_states, orig_hidden_dims - - def _maybe_apply_shared_experts( - self, - shared_experts_input: torch.Tensor | None, - order: SharedExpertsOrder, - ): - if self.shared_experts is not None: - assert shared_experts_input is not None - self.shared_experts.apply(shared_experts_input, order) - - def _apply_quant_method( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - # Run this before quant_method to avoid inplace issues. - # TODO(bnell): probably not needed anymore since inplace is - # disabled when shared experts are present. - self._maybe_apply_shared_experts( - shared_experts_input, SharedExpertsOrder.NO_OVERLAP - ) - - if self.quant_method.is_monolithic: - fused_out = self.quant_method.apply_monolithic( - layer=layer, - x=hidden_states, - router_logits=router_logits, - ) - else: - topk_weights, topk_ids = self.router.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - ) - - fused_out = self.quant_method.apply( - layer=layer, - x=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_experts_input=shared_experts_input, - ) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, - ) - - return ( - self.shared_experts.output if self.shared_experts is not None else None, - fused_out, - ) - - def _sequence_parallel_context(self): - ctx = get_forward_context() - return ( - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) - if ctx.dp_metadata - else nullcontext() - ) - - def _allocate_dp_chunking_outputs( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - assert self.use_dp_chunking - - # Assert the inputs are of the proper type and shape. - assert self.batched_hidden_states is not None - assert self.batched_router_logits is not None - - assert self.batched_hidden_states.dtype == hidden_states.dtype, ( - f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" - ) - assert self.batched_router_logits.dtype == router_logits.dtype, ( - f"{self.batched_router_logits.dtype} == {router_logits.dtype}" - ) - - # Check size compatibility. - assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) - assert self.batched_router_logits.size(-1) == router_logits.size(-1) - - final_fused_hidden_states = torch.empty_like(hidden_states) - if self.shared_experts is not None: - final_shared_hidden_states = torch.empty_like(hidden_states) - else: - final_shared_hidden_states = None - - return final_shared_hidden_states, final_fused_hidden_states - - def _maybe_overlap_gate_with_shared_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - # If router/gate provided, then apply it here. - # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoE via - # separate cuda stream) - if self.shared_experts is not None: - self.shared_experts.maybe_setup_shared_experts_stream(shared_experts_input) - - if self.gate is not None: - router_logits, _ = self.gate(hidden_states) - - return router_logits + @property + def reduce_results(self) -> bool: + return self._reduce_results @property def do_naive_dispatch_combine(self) -> bool: @@ -558,192 +143,6 @@ def _maybe_combine( else: return hidden_states - def forward( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Invoke the fused moe layer. - - Input: - - hidden_states - - router_logits - - Output: - - The new hidden_states. - or - - A tuple of (shared experts output, new hidden_states). - - Calling sequence - - forward - - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - - forward_dispatch - - forward_impl (_forward_impl or _forward_impl_chunked) - - Note: The existence of _moe_forward and _moe_forward_shared custom ops are due - to the following reasons: - 1. the chunking loop in _forward_impl_chunked cannot be compiled by - torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward - and _moe_forward_shared must be split. - - If _forward_impl_chunked can be implemented via torch.scan we can potentially - get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence - into the 'forward' method. - """ - - # Apply transform for routed experts (e.g., latent projection for latent MoE) - hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states - ) - - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( - shared_experts_input, - hidden_states, - ) - - fused_output = self.forward_entry( - hidden_states, - router_logits, - shared_experts_input, - self._encode_layer_name(), - ) - - return self._maybe_reduce_output(fused_output, og_hidden_dims) - - def forward_dispatch( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - - router_logits = self._maybe_overlap_gate_with_shared_experts( - hidden_states, - router_logits, - shared_experts_input, - ) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.EXTERNAL, - ) - - with self._sequence_parallel_context(): - return self.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - def _slice_and_copy_input( - self, - out_slice: torch.Tensor, - orig: torch.Tensor | None, - start: int, - end: int, - ) -> torch.Tensor: - assert orig is not None - slice_size = end - start - orig_slice = orig[start:end, :] - if self.enable_dbo: - assert out_slice.dim() == 3 - batch_buffer_idx = dbo_current_ubatch_id() - out_slice = out_slice[batch_buffer_idx, :] - - assert out_slice.size(0) >= slice_size - out_slice = out_slice[:slice_size, :] - out_slice.copy_(orig_slice, non_blocking=True) - return out_slice - - def _forward_impl_chunked( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - final_shared_hidden_states, final_fused_hidden_states = ( - self._allocate_dp_chunking_outputs(hidden_states, router_logits) - ) - - ctx = get_forward_context() - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens - - # If the input to the MoE is sequence parallel then divide by sp_size - # to find the maximum number of tokens for any individual dispatcher. - if self.moe_config.is_sequence_parallel: - max_tokens_across_dispatchers = cdiv( - max_tokens_across_dispatchers, self.moe_config.sp_size - ) - - num_tokens = hidden_states.size(0) - for chunk_idx, chunk_start_ in enumerate( - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) - ): - chunk_start = chunk_start_ - chunk_end = min( - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers - ) - # clamp start and end - chunk_start = min(chunk_start, num_tokens - 1) - chunk_end = min(chunk_end, num_tokens) - chunk_sizes = ctx.dp_metadata.chunked_sizes( - self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx - ) - with chunk_sizes: - hidden_states_chunk = self._slice_and_copy_input( - self.batched_hidden_states, - hidden_states, - chunk_start, - chunk_end, - ) - - router_logits_chunk = self._slice_and_copy_input( - self.batched_router_logits, - router_logits, - chunk_start, - chunk_end, - ) - - shared_experts_input_chunk = ( - shared_experts_input[chunk_start:chunk_end, :] - if shared_experts_input is not None - else None - ) - - shared_output_chunk, hidden_states_chunk = self._apply_quant_method( - layer=layer, - hidden_states=hidden_states_chunk, - router_logits=router_logits_chunk, - shared_experts_input=shared_experts_input_chunk, - ) - - # Store outputs - # TODO(bnell): document when chunk_start >= num_tokens - if chunk_start < num_tokens: - final_fused_hidden_states[chunk_start:chunk_end, :].copy_( - hidden_states_chunk, non_blocking=True - ) - if self.shared_experts is not None: - assert shared_output_chunk is not None - assert final_shared_hidden_states is not None - final_shared_hidden_states[chunk_start:chunk_end, :].copy_( - shared_output_chunk, non_blocking=True - ) - - if self.shared_experts is None: - return final_fused_hidden_states - else: - assert final_shared_hidden_states is not None - return (final_shared_hidden_states, final_fused_hidden_states) - def _forward_impl( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py new file mode 100644 index 000000000000..2a6e5713827b --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from abc import abstractmethod +from collections.abc import Callable +from contextlib import nullcontext +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +from vllm.distributed import ( + tensor_model_parallel_all_reduce, +) +from vllm.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + HAS_OPAQUE_TYPE, + ModuleName, + direct_register_custom_op, +) + +logger = init_logger(__name__) + + +def get_layer_from_name(layer_name: str) -> torch.nn.Module: + forward_context: ForwardContext = get_forward_context() + if layer_name == "from_forward_context": + all_moe_layers = forward_context.all_moe_layers + assert all_moe_layers is not None + moe_layer_index = forward_context.moe_layer_index + if moe_layer_index >= len(all_moe_layers): + raise AssertionError( + "We expected the number of MOE layers in `all_moe_layers` " + "to be equal to the number of " + "{vllm.moe_forward, vllm.moe_forward_shared} calls." + ) + layer_name = all_moe_layers[moe_layer_index] + forward_context.moe_layer_index += 1 + return forward_context.no_compile_layers[layer_name] + + +# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; +# on older versions it remains a plain str. +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | ModuleName +else: + _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str + + +def _resolve_layer_name(layer_name: str | ModuleName) -> str: + return layer_name.value if isinstance(layer_name, ModuleName) else layer_name + + +# Note: _moe_forward and _moe_forward_shared should not contain any +# implementation details, They should merely pass along control to +# the runner's 'forward_dispatch' method. +def _moe_forward( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> torch.Tensor: + layer = get_layer_from_name(_resolve_layer_name(layer_name)) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +def _moe_forward_shared( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(_resolve_layer_name(layer_name)) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_shared_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> tuple[torch.Tensor, torch.Tensor]: + # Output shapes: + # - fused_out: same as hidden_states (routed experts use transformed size) + # - shared_out: same as shared_experts_input if provided, else same as + # hidden_states + # (For latent MoE: shared experts use original hidden_size, not latent size) + fused_out = torch.empty_like(hidden_states) + if shared_experts_input is not None: + shared_out = torch.empty_like(shared_experts_input) + else: + shared_out = torch.empty_like(hidden_states) + return shared_out, fused_out + + +direct_register_custom_op( + op_name="moe_forward", + op_func=_moe_forward, + mutates_args=["hidden_states"], # is this still true? + fake_impl=_moe_forward_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +direct_register_custom_op( + op_name="moe_forward_shared", + op_func=_moe_forward_shared, + fake_impl=_moe_forward_shared_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class MoERunnerBase(MoERunner): + """ + Abstract base class providing common functionality for MoE runner implementations. + + This class serves as the foundation for concrete MoE runner implementations by + providing shared state management and common utilities. It handles: + - Common initialization and configuration management + - Shared expert output reduction logic for tensor parallel scenarios + - Base methods for tensor model parallel reductions + - Common properties and utility functions used across different runner types + + Concrete subclasses must implement the abstract methods to define their specific + execution strategies, such as standard execution, chunked processing, or other + specialized approaches. The base class provides the infrastructure while + allowing flexibility in the actual MoE computation implementation. + + Key abstract methods that subclasses must implement: + - reduce_results: Determines whether results should be reduced across ranks + - _forward_impl: The core MoE computation logic specific to each runner type + """ + + def __init__( + self, + layer: torch.nn.Module, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_input_transform: torch.nn.Module | None, + gate: torch.nn.Module | None, + shared_experts: SharedExperts | None, + quant_method: FusedMoEMethodBase, + reduce_results: bool, + enable_dbo: bool, + ): + super().__init__() + self.moe_config = moe_config + self.router = router + self.routed_input_transform = routed_input_transform + self.gate = gate + self.shared_experts = shared_experts + self.quant_method = quant_method + self._reduce_results = reduce_results + self.enable_dbo = enable_dbo + self.enable_eplb = moe_config.moe_parallel_config.enable_eplb + + # Needed for string -> FusedMoE layer lookup in custom ops. + self.layer_name = layer.layer_name + + self.forward_entry = self._select_forward(layer) + + def _select_forward(self, layer: torch.nn.Module) -> Callable: + if current_platform.is_tpu() or current_platform.is_cpu(): + # TODO: Once the OOM issue for the TPU backend is resolved, we + # will switch to using the moe_forward custom op. + # Note: CPU doesn't require wrapped _forward_impl. + return _moe_forward if self.shared_experts is None else _moe_forward_shared + + return ( + torch.ops.vllm.moe_forward + if self.shared_experts is None + else torch.ops.vllm.moe_forward_shared + ) + + @property + @abstractmethod + def reduce_results(self) -> bool: + raise NotImplementedError + + def must_reduce_shared_expert_outputs(self) -> bool: + """ + The shared_experts are typically computed using the RowParallelLinear + layer. The result of this function is typically used as + the reduce_results argument to the module. + When just tensor-parallel is used, it is not required to reduce + the shared_experts results immediately. Instead we reduce at the + once at the end of the MoE op. (Refer to DeepSeekV2MoE module) + With EP and all2all kernels - this is no longer viable as all + GPU ranks in DP, produce the complete set of hidden_states. + Therefore it is required that we reduce the shared_experts output + early. + """ + return ( + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() + ) + + def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): + """ + Some combine kernels reduce across GPU ranks by default. + """ + if self.must_reduce_shared_expert_outputs(): + return final_hidden_states + else: + return tensor_model_parallel_all_reduce(final_hidden_states) + + def apply_routed_input_transform( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Apply transform for routed experts (e.g., latent projection). + + This is called by FusedMoE.forward_native. The original hidden_states + is saved separately so shared experts get [S, hidden_size] while + routed experts get the transformed [S, moe_latent_size]. + + TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be + moved inside SharedFusedMoE to all-reduce on the smaller latent + dimension. + + Returns (possibly transformed) hidden states and the input for shared + experts (or None if there are no shared experts). + """ + if self.routed_input_transform is not None: + result = self.routed_input_transform(hidden_states) + # ReplicatedLinear returns (output, extra_bias) tuple. + # We only need the output tensor; extra_bias is not used here. + if isinstance(result, tuple): + return result[0], hidden_states + return result, hidden_states + + return ( + hidden_states, + hidden_states if self.shared_experts is not None else None, + ) + + def _maybe_reduce_output( + self, + states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + trunc_sizes: list[int], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: + return x[..., :trunc_size] + + def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: + return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) + + if ( + not self.moe_config.is_sequence_parallel + and self.reduce_results + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + ): + func = reduce_and_trunc + else: + func = trunc + + if isinstance(states, tuple): + return tuple( + [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] + ) + else: + assert len(trunc_sizes) == 1 + return func(states, trunc_sizes[0]) + + def _encode_layer_name(self) -> str | ModuleName: + if HAS_OPAQUE_TYPE: + return ModuleName(self.layer_name) + # Can be unavailable or None in unittests + if ( + is_forward_context_available() + and get_forward_context().all_moe_layers is not None + ): + return "from_forward_context" + return self.layer_name + + def _maybe_pad_hidden_states( + self, + shared_experts_input: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, list[int]]: + shared_experts_hidden_dim = ( + shared_experts_input.shape[-1] if shared_experts_input is not None else 0 + ) + transformed_hidden_dim = hidden_states.shape[-1] + if ( + not self.quant_method.skip_forward_padding + and self.moe_config.hidden_dim != transformed_hidden_dim + ): + hidden_states = F.pad( + hidden_states, + (0, self.moe_config.hidden_dim - transformed_hidden_dim), + mode="constant", + value=0.0, + ) + + if self.shared_experts is not None: + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] + else: + orig_hidden_dims = [transformed_hidden_dim] + + return hidden_states, orig_hidden_dims + + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + order: SharedExpertsOrder, + ): + if self.shared_experts is not None: + assert shared_experts_input is not None + self.shared_experts.apply(shared_experts_input, order) + + def _apply_quant_method( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + # Run this before quant_method to avoid inplace issues. + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.BEFORE_QUANT_METHOD, + ) + + if self.quant_method.is_monolithic: + fused_out = self.quant_method.apply_monolithic( + layer=layer, + x=hidden_states, + router_logits=router_logits, + ) + else: + topk_weights, topk_ids = self.router.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + + fused_out = self.quant_method.apply( + layer=layer, + x=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts_input=shared_experts_input, + ) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.AFTER_QUANT_METHOD, + ) + + return ( + self.shared_experts.output if self.shared_experts is not None else None, + fused_out, + ) + + def _sequence_parallel_context(self): + ctx = get_forward_context() + return ( + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) + if ctx.dp_metadata + else nullcontext() + ) + + def _maybe_gate( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + # If router/gate provided, then apply it here. + # (Note: This code runs only when "overlapped mode" is on to allow + # parallel execution of shared experts with the FusedMoE via + # separate cuda stream) + if self.gate is not None: + router_logits, _ = self.gate(hidden_states) + return router_logits + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Invoke the fused moe layer. + + Input: + - hidden_states + - router_logits + + Output: + - The new hidden_states. + or + - A tuple of (shared experts output, new hidden_states). + + Calling sequence + - forward + - self.forward_entry (_moe_forward or _moe_forward_shared custom op) + - forward_dispatch + - _forward_impl + + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due + to the following reasons: + 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by + torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward + and _moe_forward_shared must be split. + + If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can + potentially get rid of _moe_forward and _moe_forward_shared and collapse the + whole sequence into the 'forward' method. + """ + + # Apply transform for routed experts (e.g., latent projection for latent MoE) + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) + + hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + shared_experts_input, + hidden_states, + ) + + fused_output = self.forward_entry( + hidden_states, + router_logits, + shared_experts_input, + self._encode_layer_name(), + ) + + return self._maybe_reduce_output(fused_output, og_hidden_dims) + + def forward_dispatch( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # TODO(bnell): this can be removed after MK migration is complete. + layer.ensure_moe_quant_config_init() + + router_logits = self._maybe_gate(hidden_states, router_logits) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.EXTERNAL, + ) + + with self._sequence_parallel_context(): + return self._forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + @abstractmethod + def _forward_impl( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py new file mode 100644 index 000000000000..4bda1ca89304 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import ( + ChunkingMoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( + DefaultMoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + + +def create_moe_runner( + layer: torch.nn.Module, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_input_transform: torch.nn.Module | None, + gate: torch.nn.Module | None, + shared_experts: SharedExperts | None, + quant_method: FusedMoEMethodBase, + reduce_results: bool, + enable_dbo: bool, +) -> MoERunner: + runner = DefaultMoERunner( + layer, + moe_config, + router, + routed_input_transform, + gate, + shared_experts, + quant_method, + reduce_results, + enable_dbo, + ) + if moe_config.moe_parallel_config.use_dp_chunking: + return ChunkingMoERunner(runner) + return runner From bdefdf5216d6dccce8341e8c4498d57d0c3800a0 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Mar 2026 03:07:43 +0000 Subject: [PATCH 084/108] fix merge Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/default_moe_runner.py | 3 --- .../layers/fused_moe/runner/moe_runner_base.py | 7 +++++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 07d63cab2220..7b583bda14cc 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -76,9 +76,6 @@ def __init__( enable_dbo, ) - def is_internal_router(self) -> bool: - return self.gate is not None - @property def reduce_results(self) -> bool: return self._reduce_results diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 2a6e5713827b..c62eadbf263c 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -212,6 +212,9 @@ def _select_forward(self, layer: torch.nn.Module) -> Callable: else torch.ops.vllm.moe_forward_shared ) + def is_internal_router(self) -> bool: + return self.gate is not None + @property @abstractmethod def reduce_results(self) -> bool: @@ -358,7 +361,7 @@ def _apply_quant_method( # Run this before quant_method to avoid inplace issues. self._maybe_apply_shared_experts( shared_experts_input, - SharedExpertsOrder.BEFORE_QUANT_METHOD, + SharedExpertsOrder.NO_OVERLAP, ) if self.quant_method.is_monolithic: @@ -383,7 +386,7 @@ def _apply_quant_method( self._maybe_apply_shared_experts( shared_experts_input, - SharedExpertsOrder.AFTER_QUANT_METHOD, + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, ) return ( From 377acc859eeb00b6487b5acd906575f0e2019ff7 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Mar 2026 03:24:35 +0000 Subject: [PATCH 085/108] fix merge Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/moe_runner_base.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index c62eadbf263c..078085d23f2e 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -402,17 +402,22 @@ def _sequence_parallel_context(self): else nullcontext() ) - def _maybe_gate( + def _maybe_overlap_gate_with_shared_experts( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: # If router/gate provided, then apply it here. # (Note: This code runs only when "overlapped mode" is on to allow # parallel execution of shared experts with the FusedMoE via # separate cuda stream) + if self.shared_experts is not None: + self.shared_experts.maybe_setup_shared_experts_stream(shared_experts_input) + if self.gate is not None: router_logits, _ = self.gate(hidden_states) + return router_logits def forward( @@ -478,7 +483,11 @@ def forward_dispatch( # TODO(bnell): this can be removed after MK migration is complete. layer.ensure_moe_quant_config_init() - router_logits = self._maybe_gate(hidden_states, router_logits) + router_logits = self._maybe_overlap_gate_with_shared_experts( + hidden_states, + router_logits, + shared_experts_input, + ) self._maybe_apply_shared_experts( shared_experts_input, From 14e58dc8cdf9955c4bdcd42f6e491215698e7fb5 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 25 Mar 2026 19:46:26 +0000 Subject: [PATCH 086/108] don't pass shared_experts to MK in lora code Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 11332f4cb228..93bb4ad759c2 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -150,7 +150,6 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), - self.base_layer.shared_experts, ) # TODO: could be incorrect due to monolithic kernel? or add assert it From 7b86f43016f2b491cb40bc1202654c2100958764 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 1 Apr 2026 15:08:40 +0000 Subject: [PATCH 087/108] remove cruft Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 3 -- .../fused_moe/router/memoizing_router.py | 35 ------------------- 2 files changed, 38 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index c5dfe5e371d8..01efe3e47310 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -155,8 +155,6 @@ def _inject_lora_into_fused_moe(self): ), ) - # TODO: could be incorrect due to monolithic kernel? or add assert it - # is modular? if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.impl.fused_experts, @@ -343,7 +341,6 @@ def wrapper(*args, **kwargs): fused_experts = m_fused_moe_fn.impl.fused_experts - # TODO: seems like this could be done with modular kernel subclasses? m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py deleted file mode 100644 index a55bd2f09d6f..000000000000 --- a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import torch - -from vllm.model_executor.layers.fused_moe.config import ( - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter - - -class MemoizingRouter(FusedMoERouter): - def __init__(self, router: FusedMoERouter): - self.router = router - - def set_capture_fn( - self, - capture_fn: Callable[[torch.Tensor], None] | None, - ) -> None: - self.router.set_capture_fn(capture_fn) - self.results: tuple[torch.Tensor, torch.Tensor] | None = None - - @property - def routing_method_type(self) -> RoutingMethodType: - return self.router.routing_method_type - - def select_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.results is None: - self.results = self.router.select_experts(hidden_states, router_logits) - return self.results From fd7a32480a89665d3c53aa09f1e9e6b19ae56df4 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 1 Apr 2026 16:46:04 +0000 Subject: [PATCH 088/108] review comments Signed-off-by: Bill Nell --- .../layers/fused_moe/modular_kernel.py | 2 +- .../fused_moe/runner/chunking_moe_runner.py | 6 ++-- .../fused_moe/runner/default_moe_runner.py | 36 ------------------- .../fused_moe/runner/moe_runner_base.py | 3 -- 4 files changed, 4 insertions(+), 43 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 4d43d7511315..f2e6e2560e70 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1007,7 +1007,7 @@ def __init__( self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, - shared_experts: SharedExperts | None = None, + shared_experts: SharedExperts | None, inplace: bool = False, ): self.prepare_finalize = prepare_finalize diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py index 6626fad557b3..1eaa6437f9bc 100644 --- a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -6,14 +6,11 @@ from vllm.forward_context import ( get_forward_context, ) -from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase from vllm.utils.math_utils import cdiv from vllm.v1.worker.ubatching import dbo_current_ubatch_id from vllm.v1.worker.workspace import current_workspace_manager -logger = init_logger(__name__) - class ChunkingMoERunner(MoERunnerBase): """ @@ -55,6 +52,9 @@ def __getattr__(self, name): # so ChunkingMoERunner's own attributes and methods take priority. return getattr(self._inner, name) + # Reducing results when chunking is handled by the MK finalize operations + # when DP chunking is enabled.. + # This will be removed by #35949 @property def reduce_results(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 8320bf20f92a..85c6563c0849 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -7,20 +7,8 @@ get_ep_group, get_pcp_group, ) -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, -) -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( - FusedMoEMethodBase, -) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( - FusedMoERouter, -) from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase -logger = init_logger(__name__) - class DefaultMoERunner(MoERunnerBase): """ @@ -49,30 +37,6 @@ class DefaultMoERunner(MoERunnerBase): for different configurations (e.g., with/without shared experts, gates, etc.). """ - def __init__( - self, - layer_name: str, - moe_config: FusedMoEConfig, - router: FusedMoERouter, - routed_input_transform: torch.nn.Module | None, - gate: torch.nn.Module | None, - shared_experts: torch.nn.Module | None, - quant_method: FusedMoEMethodBase, - reduce_results: bool, - enable_dbo: bool, - ): - super().__init__( - layer_name, - moe_config, - router, - routed_input_transform, - gate, - shared_experts, - quant_method, - reduce_results, - enable_dbo, - ) - @property def reduce_results(self) -> bool: return self._reduce_results diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index aa0c8dc064de..775c6df5ecb9 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -16,7 +16,6 @@ get_forward_context, is_forward_context_available, ) -from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, ) @@ -38,8 +37,6 @@ direct_register_custom_op, ) -logger = init_logger(__name__) - def get_layer_from_name(layer_name: str) -> torch.nn.Module: forward_context: ForwardContext = get_forward_context() From ea79ff7a306837251856dad7ba33c08ef54cc3a5 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 1 Apr 2026 16:50:20 +0000 Subject: [PATCH 089/108] fix lint Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/moe_runner.py | 17 +++++++++++++++++ .../layers/fused_moe/runner/moe_runner_base.py | 8 ++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 720e997cda36..9ffbf3108f8e 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -4,6 +4,13 @@ import torch +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + class MoERunner(ABC): """ @@ -36,3 +43,13 @@ def maybe_all_reduce_tensor_model_parallel( @abstractmethod def is_internal_router(self) -> bool: raise NotImplementedError + + @property + @abstractmethod + def shared_experts(self) -> SharedExperts | None: + raise NotImplementedError + + # TODO(bnell): temporary hack, do not call this method. + @abstractmethod + def _replace_quant_method(self, quant_method: FusedMoEMethodBase): + raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 775c6df5ecb9..c26bb42bb2e9 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -189,9 +189,9 @@ def __init__( self._reduce_results = reduce_results self.enable_dbo = enable_dbo - self.shared_experts: SharedExperts | None = None + self._shared_experts: SharedExperts | None = None if shared_experts is not None: - self.shared_experts = SharedExperts( + self._shared_experts = SharedExperts( shared_experts, moe_config=moe_config, # Note: For now we must pass quant_method along to SharedExperts so it @@ -222,6 +222,10 @@ def _select_forward(self) -> Callable: else torch.ops.vllm.moe_forward_shared ) + @property + def shared_experts(self) -> SharedExperts | None: + return self._shared_experts + # TODO(bnell): temporary hack, do not call this method. def _replace_quant_method(self, quant_method: FusedMoEMethodBase): if self.shared_experts is not None: From b6ed07b5fa5ef3982895a8e549d3c29259f87764 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 1 Apr 2026 17:42:31 +0000 Subject: [PATCH 090/108] remove EXTERNAL SharedExperts order Signed-off-by: Bill Nell --- .../fused_moe/runner/moe_runner_base.py | 5 ----- .../layers/fused_moe/runner/shared_experts.py | 22 +++---------------- 2 files changed, 3 insertions(+), 24 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index c26bb42bb2e9..766d6b255c95 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -508,11 +508,6 @@ def forward_dispatch( if self.gate is not None: router_logits, _ = self.gate(hidden_states) - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.EXTERNAL, - ) - with self._sequence_parallel_context(): return self._forward_impl( layer, diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 6d2189cb49b3..3cc4a2d8c5d9 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -32,19 +32,14 @@ class SharedExpertsOrder(IntEnum): # No shared experts. NONE = (0,) - # Get rid of this one? combine with BEFORE? - # Note: this might be important for torch.compile reasons. Can - # get rid of it after _moe_forward is undone. - EXTERNAL = (1,) - # No overlap - defensively called before MK. - NO_OVERLAP = (2,) + NO_OVERLAP = (1,) # Overlapped with dispatch/combine in DP/EP - called by the MK. - MK_INTERNAL_OVERLAPPED = (3,) + MK_INTERNAL_OVERLAPPED = (2,) # Overlapped with the gate, router, experts in aux stream. - MULTI_STREAM_OVERLAPPED = (4,) + MULTI_STREAM_OVERLAPPED = (3,) class SharedExperts: @@ -110,9 +105,6 @@ def _determine_shared_experts_order( self, hidden_states: torch.Tensor, ) -> SharedExpertsOrder: - if self._has_external_experts and not self._use_dp_chunking: - return SharedExpertsOrder.EXTERNAL - if self._quant_method.mk_owns_shared_expert: return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED @@ -205,12 +197,4 @@ def apply( else: self._output[self._output_idx] = self._layer(shared_experts_input) - if order == SharedExpertsOrder.EXTERNAL: - # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely. - assert self._output[self._output_idx] is not None - self._output[self._output_idx] = self._maybe_reduce_shared_out( - self._output[self._output_idx] - ) - assert self._output[self._output_idx] is not None From dd1f23a7d172c658db41618fc844204f50cd76d3 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 1 Apr 2026 17:59:11 +0000 Subject: [PATCH 091/108] make sure some methods are handled properly on ChunkingMoERunner Signed-off-by: Bill Nell --- .../fused_moe/runner/chunking_moe_runner.py | 18 +++++++++++++++ .../fused_moe/runner/moe_runner_base.py | 22 +++++++++---------- 2 files changed, 29 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py index 1eaa6437f9bc..f7146a15a8df 100644 --- a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -6,7 +6,13 @@ from vllm.forward_context import ( get_forward_context, ) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) from vllm.utils.math_utils import cdiv from vllm.v1.worker.ubatching import dbo_current_ubatch_id from vllm.v1.worker.workspace import current_workspace_manager @@ -52,6 +58,18 @@ def __getattr__(self, name): # so ChunkingMoERunner's own attributes and methods take priority. return getattr(self._inner, name) + @property + def shared_experts(self) -> SharedExperts | None: + return self._inner.shared_experts + + # TODO(bnell): temporary hack, do not call this method. + def _replace_quant_method(self, quant_method: FusedMoEMethodBase): + self._inner._replace_quant_method(quant_method) + self.quant_method = quant_method + + def is_internal_router(self) -> bool: + return self._inner.gate is not None + # Reducing results when chunking is handled by the MK finalize operations # when DP chunking is enabled.. # This will be removed by #35949 diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 766d6b255c95..d8788d47d187 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -214,11 +214,11 @@ def _select_forward(self) -> Callable: # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. # Note: CPU doesn't require wrapped _forward_impl. - return _moe_forward if self.shared_experts is None else _moe_forward_shared + return _moe_forward if self._shared_experts is None else _moe_forward_shared return ( torch.ops.vllm.moe_forward - if self.shared_experts is None + if self._shared_experts is None else torch.ops.vllm.moe_forward_shared ) @@ -228,8 +228,8 @@ def shared_experts(self) -> SharedExperts | None: # TODO(bnell): temporary hack, do not call this method. def _replace_quant_method(self, quant_method: FusedMoEMethodBase): - if self.shared_experts is not None: - self.shared_experts._quant_method = quant_method + if self._shared_experts is not None: + self._shared_experts._quant_method = quant_method self.quant_method = quant_method def is_internal_router(self) -> bool: @@ -293,7 +293,7 @@ def apply_routed_input_transform( return ( hidden_states, - hidden_states if self.shared_experts is not None else None, + hidden_states if self._shared_experts is not None else None, ) def _maybe_reduce_output( @@ -355,7 +355,7 @@ def _maybe_pad_hidden_states( value=0.0, ) - if self.shared_experts is not None: + if self._shared_experts is not None: orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] else: orig_hidden_dims = [transformed_hidden_dim] @@ -367,9 +367,9 @@ def _maybe_apply_shared_experts( shared_experts_input: torch.Tensor | None, order: SharedExpertsOrder, ): - if self.shared_experts is not None: + if self._shared_experts is not None: assert shared_experts_input is not None - self.shared_experts.apply(shared_experts_input, order) + self._shared_experts.apply(shared_experts_input, order) def _apply_quant_method( self, @@ -413,7 +413,7 @@ def _apply_quant_method( ) return ( - self.shared_experts.output if self.shared_experts is not None else None, + self._shared_experts.output if self._shared_experts is not None else None, fused_out, ) @@ -433,8 +433,8 @@ def _maybe_sync_shared_experts_stream( # (Note: This code runs only when "overlapped mode" is on to allow # parallel execution of shared experts with the FusedMoE via # separate cuda stream) - if self.shared_experts is not None: - self.shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) + if self._shared_experts is not None: + self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) def forward( self, From acebc422b3fe595dbb9f810c75f07810a5c40bb2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 2 Apr 2026 21:52:19 +0000 Subject: [PATCH 092/108] fixes Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/chunking_moe_runner.py | 10 ++++++++-- .../layers/fused_moe/runner/moe_runner_base.py | 3 +++ vllm/model_executor/models/nemotron_h.py | 1 + 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py index f7146a15a8df..a8c75486d711 100644 --- a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -100,6 +100,7 @@ def _allocate_dp_chunking_outputs( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, ) -> tuple[torch.Tensor | None, torch.Tensor]: # Assert the inputs are of the proper type and shape. assert self.batched_hidden_states is not None @@ -118,7 +119,10 @@ def _allocate_dp_chunking_outputs( final_fused_hidden_states = torch.empty_like(hidden_states) if self.shared_experts is not None: - final_shared_hidden_states = torch.empty_like(hidden_states) + if shared_experts_input is not None: + final_shared_hidden_states = torch.empty_like(shared_experts_input) + else: + final_shared_hidden_states = torch.empty_like(hidden_states) else: final_shared_hidden_states = None @@ -152,7 +156,9 @@ def _forward_impl( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: final_shared_hidden_states, final_fused_hidden_states = ( - self._allocate_dp_chunking_outputs(hidden_states, router_logits) + self._allocate_dp_chunking_outputs( + hidden_states, router_logits, shared_experts_input + ) ) ctx = get_forward_context() diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index d8788d47d187..f750c388999d 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -508,6 +508,9 @@ def forward_dispatch( if self.gate is not None: router_logits, _ = self.gate(hidden_states) + assert hidden_states.dtype == self.moe_config.in_dtype + assert router_logits.dtype == self.moe_config.router_logits_dtype + with self._sequence_parallel_context(): return self._forward_impl( layer, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 4ec794eccf72..8abbc808cce1 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -231,6 +231,7 @@ def __init__( num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, routed_input_transform=self.fc1_latent_proj, + router_logits_dtype=self.gate.out_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: From ed0ff6ef3bf73a91ba523d71f6359f477289ca78 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 3 Apr 2026 01:46:15 +0000 Subject: [PATCH 093/108] remove assert Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index f750c388999d..d8788d47d187 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -508,9 +508,6 @@ def forward_dispatch( if self.gate is not None: router_logits, _ = self.gate(hidden_states) - assert hidden_states.dtype == self.moe_config.in_dtype - assert router_logits.dtype == self.moe_config.router_logits_dtype - with self._sequence_parallel_context(): return self._forward_impl( layer, From 30da43cf40aea46ca7cc33fd0711d75300c4efa1 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 3 Apr 2026 23:06:54 +0000 Subject: [PATCH 094/108] remove memoizing router Signed-off-by: Bill Nell --- .../fused_moe/router/memoizing_router.py | 35 ------------------- 1 file changed, 35 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/router/memoizing_router.py diff --git a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py b/vllm/model_executor/layers/fused_moe/router/memoizing_router.py deleted file mode 100644 index a55bd2f09d6f..000000000000 --- a/vllm/model_executor/layers/fused_moe/router/memoizing_router.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import torch - -from vllm.model_executor.layers.fused_moe.config import ( - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import FusedMoERouter - - -class MemoizingRouter(FusedMoERouter): - def __init__(self, router: FusedMoERouter): - self.router = router - - def set_capture_fn( - self, - capture_fn: Callable[[torch.Tensor], None] | None, - ) -> None: - self.router.set_capture_fn(capture_fn) - self.results: tuple[torch.Tensor, torch.Tensor] | None = None - - @property - def routing_method_type(self) -> RoutingMethodType: - return self.router.routing_method_type - - def select_experts( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.results is None: - self.results = self.router.select_experts(hidden_states, router_logits) - return self.results From 574e5c124d0ec235d23c5bbb8a959ff025f3d3c2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 3 Apr 2026 23:43:13 +0000 Subject: [PATCH 095/108] fix merge + remove new reduce_results Signed-off-by: Bill Nell --- .../test_shared_fused_moe_routed_transform.py | 5 --- vllm/model_executor/layers/fused_moe/layer.py | 2 +- .../fused_moe/runner/moe_runner_base.py | 1 - vllm/model_executor/models/AXK1.py | 35 +++---------------- .../models/bailing_moe_linear.py | 1 - vllm/model_executor/models/dbrx.py | 1 - vllm/model_executor/models/gemma4.py | 1 - vllm/model_executor/models/sarvam.py | 16 +-------- 8 files changed, 7 insertions(+), 55 deletions(-) diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index d861ceaf8dd3..83d40341a79f 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -211,11 +211,6 @@ def test_routed_input_transform_inside_vs_outside( hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype) router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype) - # Clone inputs so any in-place modification by Method A - # cannot affect Method B's computation. - hidden_states_A = hidden_states.clone() - router_logits_A = router_logits.clone() - with set_forward_context(None, vllm_config, num_tokens=num_tokens): # Method A: combined output (shared + routed) combined_A = moe_with_transform(hidden_states, router_logits) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 575585ba125a..144f95d2211e 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -427,7 +427,7 @@ def __init__( ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s." assert intermediate_size % self.tp_size == 0 - self.intermediate_size_per_partition = intermediate_size // self.tp_size + intermediate_size_per_partition = intermediate_size // self.tp_size self.renormalize = renormalize # TODO(bnell): these attributes are only used by monolithic kernels. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 8c4c6dd9d093..11a6cc7af22d 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -218,7 +218,6 @@ def __init__( # called, i.e. by a MK or by the MoERunner. # Once the MK can be created upfront, we can just pass in the proper # flags derived from the quant_method's MK. - reduce_results=reduce_results, quant_method=quant_method, enable_dbo=enable_dbo, ) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index f5ed4400fb65..3d655926e1de 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -100,7 +100,7 @@ def __init__( self.tp_size = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() - self.routed_scaling_factor = config.routed_scaling_factor + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.ep_group = get_ep_group().device_group self.ep_rank = get_ep_group().rank_in_group @@ -170,7 +170,6 @@ def __init__( top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=config.norm_topk_prob, quant_config=quant_config, use_grouped_topk=True, @@ -180,9 +179,8 @@ def __init__( scoring_func=config.scoring_func, # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally - routed_scaling_factor=1.0 - if not self.is_rocm_aiter_moe_enabled - else self.routed_scaling_factor, + routed_scaling_factor=self.routed_scaling_factor, + apply_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, @@ -204,43 +202,20 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = sequence_parallel_chunk(hidden_states) if self.experts.is_internal_router: - # In this case, the gate/router runs inside the FusedMoE class - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states ) else: - # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - fused_moe_out = self.experts( + final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) - shared_output, final_hidden_states = fused_moe_out - if self.shared_experts is None: - assert shared_output is None - - # Fix FP16 overflow - # See AXK1DecoderLayer for more details. - if hidden_states.dtype != torch.float16: - if not self.is_rocm_aiter_moe_enabled: - final_hidden_states *= self.routed_scaling_factor - elif self.shared_experts is not None: - assert shared_output is not None - shared_output *= 1.0 / self.routed_scaling_factor - - if self.shared_experts is not None: - assert shared_output is not None - final_hidden_states += shared_output - if self.is_sequence_parallel: final_hidden_states = tensor_model_parallel_all_gather( final_hidden_states, 0 ) final_hidden_states = final_hidden_states[:num_tokens] - elif self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index ecc5d63ced75..079021126939 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -358,7 +358,6 @@ def __init__( top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py index ca6e6a49a98a..a72f4e487164 100644 --- a/vllm/model_executor/models/dbrx.py +++ b/vllm/model_executor/models/dbrx.py @@ -85,7 +85,6 @@ def __init__( hidden_size=config.d_model, intermediate_size=config.ffn_config.ffn_hidden_size, params_dtype=params_dtype, - reduce_results=True, renormalize=True, quant_config=quant_config, tp_size=get_tensor_model_parallel_world_size(), diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index edb533134995..33553a04ab56 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -233,7 +233,6 @@ def routing_function( "moe_intermediate_size", getattr(config, "expert_intermediate_size", None), ), - reduce_results=True, renormalize=True, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index fa5ec44d7e72..3656fc921b25 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -341,7 +341,6 @@ def __init__( top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -370,20 +369,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: router_logits=router_logits, ) - if self.shared_experts is not None: - shared_output, expert_output = final_hidden - else: - shared_output, expert_output = None, final_hidden - - if shared_output is not None: - expert_output = expert_output + shared_output - - if self.tp_size > 1: - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( - expert_output - ) - - return expert_output.view(num_tokens, hidden_dim) + return final_hidden.view(num_tokens, hidden_dim) class SarvamMLABlock(nn.Module): From 5f86ac411f7e2818dbb9eeea3a8b2934f021c172 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 3 Apr 2026 23:44:25 +0000 Subject: [PATCH 096/108] fix merge Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/moe_runner_base.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 11a6cc7af22d..020a9d384b17 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -570,13 +570,6 @@ def forward( hidden_states, ) - router_logits = self._maybe_gate(hidden_states, router_logits) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.EXTERNAL, - ) - result = self._forward_entry( hidden_states, router_logits, From 0405c75146a1dfdcd47c9b548ff00bee04393957 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 6 Apr 2026 16:50:15 +0000 Subject: [PATCH 097/108] revert bogus changes Signed-off-by: Bill Nell --- vllm/lora/layers/fused_moe.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 46beb7948e79..01efe3e47310 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -153,11 +153,8 @@ def _inject_lora_into_fused_moe(self): self.base_layer.quant_method.select_gemm_impl( prepare_finalize, self.base_layer ), - self.base_layer.shared_experts, ) - # TODO: could be incorrect due to monolithic kernel? or add assert it - # is modular? if quant_config.use_mxfp4_w4a16: assert isinstance( m_fused_moe_fn.impl.fused_experts, @@ -344,7 +341,6 @@ def wrapper(*args, **kwargs): fused_experts = m_fused_moe_fn.impl.fused_experts - # TODO: seems like this could be done with modular kernel subclasses? m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply) fused_experts.activation = act_decorator( self.base_layer, fused_experts.activation From dd4a254868f9aa6e949208660d013f028817f502 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 14 Apr 2026 21:42:11 +0000 Subject: [PATCH 098/108] fix up additional models Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 12 ++++-------- vllm/model_executor/models/aria.py | 1 - .../models/bailing_moe_linear.py | 18 ++---------------- vllm/model_executor/models/param2moe.py | 16 +--------------- 4 files changed, 7 insertions(+), 40 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 07a52264a672..86eeab991233 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -278,10 +278,6 @@ def __init__( ): super().__init__() - self._routed_input_transform = routed_input_transform - self._routed_output_transform = routed_output_transform - self._apply_scale_to_output = apply_scale_to_output - if params_dtype is None: params_dtype = torch.get_default_dtype() self.params_dtype = params_dtype @@ -581,14 +577,14 @@ def _get_quant_method() -> FusedMoEMethodBase: layer_name=self.layer_name, moe_config=self.moe_config, router=self.router, - routed_input_transform=self._routed_input_transform, gate=gate, shared_experts=shared_experts, quant_method=self.quant_method, enable_dbo=self.vllm_config.parallel_config.enable_dbo, - routed_output_transform=self._routed_output_transform, - apply_scale_to_output=self._apply_scale_to_output, - routed_scaling_factor=self.routed_scaling_factor, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + apply_scale_to_output=apply_scale_to_output, + routed_scaling_factor=routed_scaling_factor, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 85cfff05d06b..7a079c565402 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -283,7 +283,6 @@ def __init__( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, quant_config=quant_config, - reduce_results=True, prefix=f"{prefix}.experts", ) diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 079021126939..a62642b5684f 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -367,6 +367,8 @@ def __init__( topk_group=self.topk_group, use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, + routed_scaling_factor=self.routed_scaling_factor, + apply_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -382,22 +384,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states=hidden_states, router_logits=router_logits ) - # Handle tuple return from SharedFusedMoE - if self.shared_experts is not None: - shared_output, final_hidden_states = final_hidden_states - else: - shared_output = None - - final_hidden_states *= self.routed_scaling_factor - - if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output - - if self.tp_size > 1: - final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - return final_hidden_states.view(num_tokens, hidden_size) diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 925f94e06005..fddd1a8f1733 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -359,7 +359,6 @@ def __init__( top_k=self.top_k, hidden_size=self.hidden_size, intermediate_size=config.moe_intermediate_size, - reduce_results=False, renormalize=self.norm_expert_prob, quant_config=quant_config, prefix=f"{prefix}.experts", @@ -388,24 +387,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: self.gate.weight.float(), ).to(hidden_states.dtype) - final_hidden = self.experts( + expert_output = self.experts( hidden_states=hidden_states, router_logits=router_logits, ) - if self.shared_experts is not None: - shared_output, expert_output = final_hidden - else: - shared_output, expert_output = None, final_hidden - - if shared_output is not None: - expert_output = expert_output + shared_output - - if self.tp_size > 1: - expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( - expert_output - ) - return expert_output.view(num_tokens, hidden_dim) From 574e61f459b5c63b350829d65446f09cb4b5da30 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 14 Apr 2026 22:01:42 +0000 Subject: [PATCH 099/108] tweaks Signed-off-by: Bill Nell --- .../model_executor/layers/fused_moe/runner/moe_runner_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 2a188d258181..0826f1f611e7 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -309,7 +309,7 @@ def _maybe_apply_output_scale( protection. When apply_scale_to_output is True, scales the fused expert output - by routed_scaling_factor. For FP16, avoids overflow by dividing + by routed_scaling_factor. For FP16, avoid overflow by dividing shared_output by the scale instead (the decoder layer compensates with matching divisions). """ @@ -562,7 +562,7 @@ def forward( # Apply transform for routed experts (e.g., latent projection # for latent MoE) hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states, + hidden_states ) hidden_states, og_hidden_dim = self._maybe_pad_hidden_states( From fe9ec8b4b2c523698bb91f2696f81726ebc5e985 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 14 Apr 2026 22:03:03 +0000 Subject: [PATCH 100/108] remove tests Signed-off-by: Bill Nell --- .../configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml | 10 ---------- tests/evals/gsm8k/configs/moe-refactor/config-h100.txt | 1 - 2 files changed, 11 deletions(-) delete mode 100644 tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml b/tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml deleted file mode 100644 index ca5c9a00ed37..000000000000 --- a/tests/evals/gsm8k/configs/moe-refactor/LongCat-Flash-Chat-FP8.yaml +++ /dev/null @@ -1,10 +0,0 @@ -model_name: "meituan-longcat/LongCat-Flash-Chat-FP8" -accuracy_threshold: 0.70 -num_questions: 1319 -num_fewshot: 5 -startup_max_wait_seconds: 1200 -server_args: >- - --enforce-eager - --max-model-len 4096 - --tensor-parallel-size 8 - --enable-expert-parallel diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt b/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt index 912ba878fae6..7397fc4e4626 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt +++ b/tests/evals/gsm8k/configs/moe-refactor/config-h100.txt @@ -10,4 +10,3 @@ Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml Qwen3-30B-A3B-Fp8-CT-Channel-vllm-cutlass.yaml Qwen3-30B-A3B-BF16-fi-cutlass.yaml Qwen3-30B-A3B-BF16-triton.yaml -LongCat-Flash-Chat-FP8.yaml From 31c04917f7b7fec6b386b5a047269e1694d26170 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 14 Apr 2026 22:27:56 +0000 Subject: [PATCH 101/108] fix layer test Signed-off-by: Bill Nell --- tests/kernels/moe/test_moe_layer.py | 93 ++++++----------------------- 1 file changed, 17 insertions(+), 76 deletions(-) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index aa2948b8e989..6f944e24d04b 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -236,7 +236,6 @@ class MoETestConfig: use_gate: bool use_routed_input_transform: bool enable_eplb: bool = False - reduce_results: bool = False backend: str | None = None ep_size: int = 1 dp_size: int = 1 @@ -295,7 +294,6 @@ def generate_valid_test_configs( use_shared_experts, use_gate, use_routed_input_transform, - reduce_results, ) in product( SHAPE_COMBOS, NUM_EXPERTS, @@ -304,7 +302,6 @@ def generate_valid_test_configs( [False, True], # shared [False, True], # gate [False, True], # routed input exform - [False, True], # reduce results ): config = MoETestConfig( shape[0], # m @@ -318,7 +315,6 @@ def generate_valid_test_configs( use_gate, use_routed_input_transform, enable_eplb, - reduce_results, backend, ep_size, dp_size, @@ -395,18 +391,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: and config.backend.startswith("flashinfer_nvlink") and not current_platform.has_device_capability(90) ): - return False, "flashinfer_nvlink needs an H100+ GPUs" - - # reduce_results incompatibilities - if config.reduce_results and config.use_shared_experts: - return False, "reduce_results=True is not compatible with shared_experts=True" - - if config.reduce_results and config.quantization is not None: - return ( - False, - "reduce_results=True only tested with unquantized data types in " - "order to limit number of tests run", - ) + return False, "flashinfer_nvlink needs H100+ GPUs" # Backend-specific checks if config.backend is not None: @@ -448,10 +433,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS: return False, f"EPLB not supported with {config.backend}." - world_size = config.tp_size * config.dp_size - if config.reduce_results and world_size == 1: - return False, "reduce_results=True only makes sense for multi-GPU tests" - if ( config.backend is not None and config.backend.startswith("flashinfer_nvlink") @@ -838,7 +819,6 @@ def make_fused_moe_layer( tp_size: int, ep_size: int, dp_size: int, - reduce_results: bool, w1: torch.Tensor, w2: torch.Tensor, top_k: int, @@ -866,7 +846,7 @@ def make_fused_moe_layer( routed_input_transform: torch.nn.Module | None = None, routed_output_transform: torch.nn.Module | None = None, pcp_size: int | None = 1, -) -> tuple[Callable, FusedMoE]: +) -> FusedMoE: quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts) kwargs = dict() @@ -879,8 +859,10 @@ def make_fused_moe_layer( # Add gate and routed_input_transform if provided if gate is not None: kwargs["gate"] = gate + if routed_input_transform is not None: kwargs["routed_input_transform"] = routed_input_transform + kwargs["routed_output_transform"] = routed_output_transform layer = builder( num_experts=global_num_experts, @@ -888,7 +870,6 @@ def make_fused_moe_layer( hidden_size=hidden_size, intermediate_size=intermediate_size, params_dtype=in_dtype, - reduce_results=reduce_results, renormalize=renormalize, use_grouped_topk=use_grouped_topk, num_expert_group=num_expert_group, @@ -928,36 +909,7 @@ def make_fused_moe_layer( layer.quant_method.process_weights_after_loading(layer) - def _moe( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - if shared_experts is None: - final_shared_states = None - final_hidden_states = layer(hidden_states, router_logits) - else: - final_shared_states, final_hidden_states = layer( - hidden_states, router_logits - ) - - # Apply routed output transform if provided - # (e.g., latent space -> original space) - if routed_output_transform is not None: - final_hidden_states = routed_output_transform(final_hidden_states) - - if shared_experts is not None: - assert not reduce_results - assert final_shared_states is not None - final_hidden_states += final_shared_states - - if not reduce_results and layer.tp_size > 1: - final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel( - final_hidden_states - ) - - return final_hidden_states - - return _moe, layer + return layer def make_fake_moe_layer( @@ -991,7 +943,6 @@ def make_fake_moe_layer( tp_size: int = 1, dp_size: int = 1, ep_size: int = 1, - reduce_results: bool = False, ) -> Callable: activation = MoEActivation.from_str(activation) @@ -1093,7 +1044,7 @@ def _moe( def _test_body_regular( - moe_fn: Callable, + moe_layer: Callable, hidden_states: torch.Tensor, router_logits: torch.Tensor, vllm_config: VllmConfig, @@ -1110,13 +1061,12 @@ def _test_body_regular( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output = moe_fn(hidden_states, router_logits) + output = moe_layer(hidden_states, router_logits) return baseline_output, output def _test_body_eplb( - moe_fn: Callable, moe_layer: FusedMoE, hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -1137,7 +1087,6 @@ def _test_body_eplb( n: int, top_k: int, shared_experts, - reduce_results: bool, gate: torch.nn.Module | None, routed_input_transform: torch.nn.Module | None, routed_output_transform: torch.nn.Module | None, @@ -1153,7 +1102,7 @@ def _test_body_eplb( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output_before = moe_fn(hidden_states, router_logits) + output_before = moe_layer(hidden_states, router_logits) # Create a fresh FusedMoE layer with enable_eplb=True # Delete the original layer's registration so the constructor can @@ -1166,7 +1115,7 @@ def _test_body_eplb( # When using routed_input_transform, experts operate in latent space hidden_size_for_layer = k // 2 if routed_input_transform is not None else k - moe_fn, moe_layer = make_fused_moe_layer( + eplb_moe_layer = make_fused_moe_layer( quantization=quantization, use_ep=use_ep, hidden_size=hidden_size_for_layer, @@ -1175,7 +1124,6 @@ def _test_body_eplb( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, w1=w1, w2=w2, top_k=top_k, @@ -1188,14 +1136,14 @@ def _test_body_eplb( ) # Necessary? - if moe_layer._expert_map is not None: - moe_layer._expert_map = moe_layer._expert_map.to(device) + if eplb_moe_layer._expert_map is not None: + eplb_moe_layer._expert_map = eplb_moe_layer._expert_map.to(device) # All ranks must generate the same permutation initial_indices = torch.arange(num_experts, dtype=torch.long) shuffled_indices = initial_indices[torch.randperm(num_experts)] - expert_weights = [list(moe_layer.get_expert_weights())] + expert_weights = [list(eplb_moe_layer.get_expert_weights())] communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), @@ -1219,7 +1167,7 @@ def _test_body_eplb( num_experts, dtype=torch.int32, device=device ) - moe_layer.set_eplb_state( + eplb_moe_layer.set_eplb_state( moe_layer_idx=0, expert_load_view=torch.zeros( (1, num_experts), @@ -1236,7 +1184,7 @@ def _test_body_eplb( ), ) - moe_layer.eplb_state.should_record_tensor = torch.ones( + eplb_moe_layer.eplb_state.should_record_tensor = torch.ones( (), dtype=torch.bool, device=device ) @@ -1247,7 +1195,7 @@ def _test_body_eplb( num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ): - output_after = moe_fn(hidden_states, router_logits) + output_after = eplb_moe_layer(hidden_states, router_logits) return output_before, output_after @@ -1266,7 +1214,6 @@ def _run_one_config( num_experts: int, top_k: int, quantization: str | None, - reduce_results: bool, backend: str | None, test_body_fn: Callable, use_shared_experts: bool, @@ -1333,7 +1280,6 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, ) baseline_output = baseline_layer(hidden_states, router_logits) @@ -1361,7 +1307,7 @@ def _run_one_config( hidden_size_for_layer = k // 2 if routed_input_transform is not None else k # Create initial MoE layer - moe_fn, moe_layer = make_fused_moe_layer( + moe_layer = make_fused_moe_layer( quantization=quantization, use_ep=use_ep, hidden_size=hidden_size_for_layer, @@ -1370,7 +1316,6 @@ def _run_one_config( tp_size=tp_size, ep_size=ep_size, dp_size=dp_size, - reduce_results=reduce_results, w1=w1, w2=w2, top_k=top_k, @@ -1394,7 +1339,6 @@ def _run_one_config( # Call the test body function with all necessary context expected, actual = test_body_fn( - moe_fn=moe_fn, moe_layer=moe_layer, hidden_states=hidden_states, router_logits=router_logits, @@ -1415,7 +1359,6 @@ def _run_one_config( m=m, top_k=top_k, shared_experts=shared_experts, - reduce_results=reduce_results, gate=gate, routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, @@ -1512,7 +1455,6 @@ def test_moe_layer_no_parallel( test_config.num_experts, test_config.top_k, test_config.quantization, - test_config.reduce_results, test_config.backend, _test_body_regular, use_shared_experts=test_config.use_shared_experts, @@ -1570,7 +1512,6 @@ def _parallel_worker( test_config.num_experts, test_config.top_k, test_config.quantization, - test_config.reduce_results, test_config.backend, functools.partial( _test_body_config, test_config=test_config, cpu_group=cpu_group @@ -1589,7 +1530,7 @@ def _parallel_worker( failed = failed + 1 if verbosity > 0: traceback.print_exc() - print(f"\n{str(ex)}\nFAILED {ex.__class__}") + print(f"\n{str(ex)}\nFAILED") else: print("F", end="") finally: From 2a250665559885f540d25c6f3fd0021cccb59e17 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 16 Apr 2026 19:17:49 +0000 Subject: [PATCH 102/108] review comments Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 6 +- .../fused_moe/runner/moe_runner_base.py | 58 ++++++++----------- .../fused_moe/runner/moe_runner_factory.py | 4 +- vllm/model_executor/models/AXK1.py | 2 +- .../models/bailing_moe_linear.py | 2 +- vllm/model_executor/models/deepseek_v2.py | 2 +- vllm/model_executor/models/dots1.py | 2 +- vllm/model_executor/models/glm4_moe.py | 1 + vllm/model_executor/models/nemotron_h.py | 2 +- vllm/model_executor/models/openpangu.py | 2 +- 10 files changed, 35 insertions(+), 46 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 86eeab991233..d99e5c94cadb 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -273,7 +273,7 @@ def __init__( shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, routed_output_transform: torch.nn.Module | None = None, - apply_scale_to_output: bool = False, + apply_routed_scale_to_fused_output: bool = False, zero_expert_type: str | None = None, ): super().__init__() @@ -454,7 +454,7 @@ def __init__( custom_routing_function=custom_routing_function, scoring_func=scoring_func, routed_scaling_factor=routed_scaling_factor - if not apply_scale_to_output + if not apply_routed_scale_to_fused_output else 1.0, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=self.num_fused_shared_experts, @@ -583,7 +583,7 @@ def _get_quant_method() -> FusedMoEMethodBase: enable_dbo=self.vllm_config.parallel_config.enable_dbo, routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, - apply_scale_to_output=apply_scale_to_output, + apply_routed_scale_to_fused_output=apply_routed_scale_to_fused_output, routed_scaling_factor=routed_scaling_factor, ) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 0826f1f611e7..4c5f79bb6ee2 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -199,7 +199,7 @@ def __init__( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - apply_scale_to_output: bool = False, + apply_routed_scale_to_fused_output: bool = False, # Flip sense and change name? routed_scaling_factor: float = 1.0, ): super().__init__() @@ -207,13 +207,17 @@ def __init__( self.router = router self.routed_input_transform = routed_input_transform self.routed_output_transform = routed_output_transform - self.apply_scale_to_output = ( - apply_scale_to_output and routed_scaling_factor != 1.0 + self.apply_routed_scale_to_fused_output = ( + apply_routed_scale_to_fused_output and routed_scaling_factor != 1.0 ) self.routed_scaling_factor = routed_scaling_factor self.gate = gate self.quant_method = quant_method self.enable_dbo = enable_dbo + self._fused_output_is_reduced = ( + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() + ) self._shared_experts: SharedExperts | None = None if shared_experts is not None: @@ -300,7 +304,7 @@ def apply_routed_output_transform( fused_output = r[0] if isinstance(r, tuple) else r return fused_output - def _maybe_apply_output_scale( + def _maybe_apply_routed_scale_to_fused_output( self, shared_output: torch.Tensor | None, fused_output: torch.Tensor, @@ -308,37 +312,18 @@ def _maybe_apply_output_scale( """Apply routed_scaling_factor to the output with FP16 overflow protection. - When apply_scale_to_output is True, scales the fused expert output + When apply_routed_scale_to_fused_output is True, scales the fused expert output by routed_scaling_factor. For FP16, avoid overflow by dividing shared_output by the scale instead (the decoder layer compensates with matching divisions). """ - if self.apply_scale_to_output: + if self.apply_routed_scale_to_fused_output: if fused_output.dtype != torch.float16: fused_output *= self.routed_scaling_factor elif shared_output is not None: shared_output *= 1.0 / self.routed_scaling_factor return shared_output, fused_output - def _must_reduce_shared_expert_output(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.shared_experts is not None - and self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, @@ -348,14 +333,14 @@ def _maybe_reduce_shared_expert_output( This is the "early" all-reduce path. When the combine kernel produces already-reduced fused output, shared output must be reduced separately - to match. See _must_reduce_shared_expert_output for details. + to match. """ - if self._must_reduce_shared_expert_output(): + if self._fused_output_is_reduced: assert shared_output is not None shared_output = tensor_model_parallel_all_reduce(shared_output) return shared_output - def _maybe_reduce_output( + def _maybe_reduce_final_output( self, states: torch.Tensor, trunc_size: int, @@ -369,10 +354,13 @@ def _maybe_reduce_output( """ result = states[..., :trunc_size] + # We don't need to reduce the final output if: + # - We are not running with TP or DP + # - The MK already reduced the fused output itself. if ( not self.moe_config.is_sequence_parallel and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - and not self._must_reduce_shared_expert_output() + and not self._fused_output_is_reduced ): result = tensor_model_parallel_all_reduce(result) @@ -400,7 +388,7 @@ def _maybe_pad_hidden_states( For latent MoE, the routed hidden_states may be smaller than hidden_dim. Padding ensures uniform tensor sizes through the fused MoE kernel. The returned trunc_size is used by - _maybe_reduce_output to strip the padding from the result. + _maybe_reduce_final_output to strip the padding from the result. """ shared_experts_hidden_dim = ( shared_experts_input.shape[-1] if shared_experts_input is not None else 0 @@ -579,12 +567,12 @@ def forward( # # Note: there are two all-reduce points below. They are mutually - # exclusive, controlled by _must_reduce_shared_expert_output(): + # exclusive, controlled by _fused_output_is_reduced # - When True: the combine kernel already reduced fused_output, # so we reduce shared_output here to match, then skip the - # all-reduce in _maybe_reduce_output. + # all-reduce in _maybe_reduce_final_output. # - When False: neither output is reduced yet, so we combine - # them first and all-reduce the sum in _maybe_reduce_output. + # them first and all-reduce the sum in _maybe_reduce_final_output. # Extract outputs from result shared_output, fused_output = _unpack(result) @@ -596,7 +584,7 @@ def forward( # See note above re: the two all-reduce points. shared_output = self._maybe_reduce_shared_expert_output(shared_output) - shared_output, fused_output = self._maybe_apply_output_scale( + shared_output, fused_output = self._maybe_apply_routed_scale_to_fused_output( shared_output, fused_output ) @@ -605,7 +593,7 @@ def forward( else: result = fused_output - result = self._maybe_reduce_output(result, og_hidden_dim) + result = self._maybe_reduce_final_output(result, og_hidden_dim) return self._maybe_add_zero_expert_output(result) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index 06d2a2c13b20..a454eddf28b6 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -31,7 +31,7 @@ def create_moe_runner( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - apply_scale_to_output: bool = False, + apply_routed_scale_to_fused_output: bool = False, routed_scaling_factor: float = 1.0, ) -> MoERunner: return DefaultMoERunner( @@ -44,6 +44,6 @@ def create_moe_runner( quant_method, enable_dbo, routed_output_transform=routed_output_transform, - apply_scale_to_output=apply_scale_to_output, + apply_routed_scale_to_fused_output=apply_routed_scale_to_fused_output, routed_scaling_factor=routed_scaling_factor, ) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 3d655926e1de..21f54ac4438f 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -180,7 +180,7 @@ def __init__( # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, - apply_scale_to_output=not self.is_rocm_aiter_moe_enabled, + apply_routed_scale_to_fused_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index a62642b5684f..afd42d9da8c8 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -368,7 +368,7 @@ def __init__( use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, routed_scaling_factor=self.routed_scaling_factor, - apply_scale_to_output=True, + apply_routed_scale_to_fused_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 9451d04be324..ebe36b8c28ec 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -324,7 +324,7 @@ def __init__( # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, - apply_scale_to_output=not self.is_rocm_aiter_moe_enabled, + apply_routed_scale_to_fused_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index f0bbde4bcdbe..e0bd9ff150a8 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -170,7 +170,7 @@ def __init__( scoring_func=config.scoring_func, e_score_correction_bias=self.gate.e_score_correction_bias, routed_scaling_factor=self.routed_scaling_factor, - apply_scale_to_output=True, + apply_routed_scale_to_fused_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index cef934222382..530f471580b5 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -192,6 +192,7 @@ def __init__( prefix=f"{prefix}.experts", scoring_func="sigmoid", routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scale_to_fused_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index d44e4c69d574..75f959ee43fa 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -234,7 +234,7 @@ def __init__( if self.use_latent_moe else None, routed_scaling_factor=self.routed_scaling_factor, - apply_scale_to_output=True, + apply_routed_scale_to_fused_output=True, router_logits_dtype=self.gate.out_dtype, ) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 409370a77089..8e586c3139fc 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -214,7 +214,7 @@ def __init__( prefix=f"{prefix}.experts", scoring_func="sigmoid", routed_scaling_factor=self.routed_scaling_factor, - apply_scale_to_output=True, + apply_routed_scale_to_fused_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, From 73ecb6bde76276b96a764f4e9229b4980cf38e9f Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 16 Apr 2026 22:09:03 +0000 Subject: [PATCH 103/108] fix SharedFusedMoE return type Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/shared_fused_moe.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index ed243e992636..11c5130be18a 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -23,7 +23,4 @@ def forward( hidden_states=hidden_states, router_logits=router_logits, ) - if self.shared_experts is None: - return None, result - else: - return result + return None, result From f0fa198852466b0a55131b0587e6e55c1d95cd0d Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 16 Apr 2026 22:13:58 +0000 Subject: [PATCH 104/108] try again Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/shared_fused_moe.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py index 11c5130be18a..9cfcb1baa9bb 100644 --- a/vllm/model_executor/layers/fused_moe/shared_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/shared_fused_moe.py @@ -18,9 +18,8 @@ def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - result = super().forward( + ) -> torch.Tensor: + return super().forward( hidden_states=hidden_states, router_logits=router_logits, ) - return None, result From b83a93b6dec4387a1f00073d137aefeb03d2e3bd Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 17 Apr 2026 03:07:55 +0000 Subject: [PATCH 105/108] renames, add comments, simplify scaling factor handling Signed-off-by: Bill Nell --- .../test_vllm_fusion_pattern_matcher_pass.py | 8 +++---- .../layers/fused_moe/fused_marlin_moe.py | 2 +- vllm/model_executor/layers/fused_moe/layer.py | 24 +++++++++++++++---- .../fused_moe/runner/moe_runner_base.py | 19 ++++++--------- .../fused_moe/runner/moe_runner_factory.py | 2 -- vllm/model_executor/models/AXK1.py | 2 +- .../models/bailing_moe_linear.py | 2 +- vllm/model_executor/models/deepseek_v2.py | 3 +-- vllm/model_executor/models/dots1.py | 2 +- vllm/model_executor/models/glm4_moe.py | 2 +- vllm/model_executor/models/nemotron_h.py | 2 +- vllm/model_executor/models/openpangu.py | 2 +- 12 files changed, 38 insertions(+), 32 deletions(-) diff --git a/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py b/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py index ee2eef009fc2..381c215a9ae1 100644 --- a/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py +++ b/tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py @@ -6,14 +6,13 @@ import vllm.config from tests.compile.backend import TestBackend -from vllm.platforms import current_platform from vllm.compilation.passes.vllm_inductor_pass import ( VllmFusionPatternMatcherPass, VllmPatternMatcherPass, VllmPatternReplacement, ) from vllm.config import CompilationConfig, CompilationMode, VllmConfig - +from vllm.platforms import current_platform class ReluToAbsPattern(VllmPatternReplacement): @@ -58,7 +57,6 @@ def get_inputs(self) -> list[torch.Tensor]: return [self.empty_fp32(4)] - class ReluFusionPass(VllmFusionPatternMatcherPass): def __init__(self, config: VllmConfig) -> None: super().__init__(config, "test_relu_fusion") @@ -72,13 +70,13 @@ def __init__(self, config: VllmConfig) -> None: self.register(ExpToSqrtPattern()) - @pytest.fixture def vllm_config(): return VllmConfig( compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE), ) + @pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA") def test_register_tracks_patterns(vllm_config): """register() appends each VllmPatternReplacement to _pattern_replacements.""" @@ -96,7 +94,7 @@ def test_uuid_stable(vllm_config): with vllm.config.set_current_vllm_config(vllm_config): p1 = ReluFusionPass(vllm_config) p2 = ReluFusionPass(vllm_config) - p3= TwoPatternFusionPass(vllm_config) + p3 = TwoPatternFusionPass(vllm_config) assert p1.uuid() == p2.uuid() assert p1.uuid() != p3.uuid() diff --git a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py index 6c916cf3cb66..daae5b6bd16c 100644 --- a/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_marlin_moe.py @@ -716,7 +716,7 @@ def apply( ): assert self.w1_scale is not None assert self.w2_scale is not None - return fused_marlin_moe( + fused_marlin_moe( hidden_states=hidden_states, w1=w1, w2=w2, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index d99e5c94cadb..0df0c9036f04 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -234,6 +234,14 @@ class FusedMoE(PluggableLayer): quant_config: Quantization configure. enable_eplb: Whether to enable expert parallelism load balancer. router_logits_dtype: Data type for router logits buffers. + routed_scaling_factor: A scaling factor that is applied to the topk_weights + by the router or the output of the layer depending + on the value of `apply_routed_scale_to_output` + apply_routed_scale_to_output: Determine whether or not `routed_scaling_factor` + is applied to the topk_weights or to the experts + output. It is applied to the experts output + instead of the topk_weights when this feature is + not supported by the router (or the experts). """ # --8<-- [end:fused_moe] @@ -273,7 +281,7 @@ def __init__( shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, routed_output_transform: torch.nn.Module | None = None, - apply_routed_scale_to_fused_output: bool = False, + apply_routed_scale_to_output: bool = False, zero_expert_type: str | None = None, ): super().__init__() @@ -453,8 +461,12 @@ def __init__( topk_group=topk_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, + # When apply_routed_scale_to_output is True, we set + # the scaling factor to 1.0 so it ends up being a nop. + # Applying the scale will be handled by the runner + # in this case. routed_scaling_factor=routed_scaling_factor - if not apply_routed_scale_to_fused_output + if not apply_routed_scale_to_output else 1.0, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=self.num_fused_shared_experts, @@ -583,8 +595,12 @@ def _get_quant_method() -> FusedMoEMethodBase: enable_dbo=self.vllm_config.parallel_config.enable_dbo, routed_input_transform=routed_input_transform, routed_output_transform=routed_output_transform, - apply_routed_scale_to_fused_output=apply_routed_scale_to_fused_output, - routed_scaling_factor=routed_scaling_factor, + # When apply_routed_scale_to_output is True, we allow + # the scaling factor to be passed to the runner, otherwise + # we pass 1.0 so it ends up being a nop. + routed_scaling_factor=routed_scaling_factor + if apply_routed_scale_to_output + else 1.0, ) # TODO(bnell): This method is provided as a hook so vllm/lora/layers/fused_moe.py diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 4c5f79bb6ee2..0204a6539fee 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -199,7 +199,6 @@ def __init__( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - apply_routed_scale_to_fused_output: bool = False, # Flip sense and change name? routed_scaling_factor: float = 1.0, ): super().__init__() @@ -207,9 +206,6 @@ def __init__( self.router = router self.routed_input_transform = routed_input_transform self.routed_output_transform = routed_output_transform - self.apply_routed_scale_to_fused_output = ( - apply_routed_scale_to_fused_output and routed_scaling_factor != 1.0 - ) self.routed_scaling_factor = routed_scaling_factor self.gate = gate self.quant_method = quant_method @@ -304,7 +300,7 @@ def apply_routed_output_transform( fused_output = r[0] if isinstance(r, tuple) else r return fused_output - def _maybe_apply_routed_scale_to_fused_output( + def _maybe_apply_routed_scale_to_output( self, shared_output: torch.Tensor | None, fused_output: torch.Tensor, @@ -312,12 +308,11 @@ def _maybe_apply_routed_scale_to_fused_output( """Apply routed_scaling_factor to the output with FP16 overflow protection. - When apply_routed_scale_to_fused_output is True, scales the fused expert output - by routed_scaling_factor. For FP16, avoid overflow by dividing - shared_output by the scale instead (the decoder layer compensates - with matching divisions). + Scale the fused expert output by routed_scaling_factor. For FP16, + avoid overflow by dividing shared_output by the scale instead + (the decoder layer compensates with matching divisions). """ - if self.apply_routed_scale_to_fused_output: + if self.routed_scaling_factor != 1.0: if fused_output.dtype != torch.float16: fused_output *= self.routed_scaling_factor elif shared_output is not None: @@ -352,7 +347,7 @@ def _maybe_reduce_final_output( here. Skipped when sequence-parallel is active (SP handles its own reduction) or when the early path already reduced both outputs. """ - result = states[..., :trunc_size] + result = states[..., :trunc_size].contiguous() # We don't need to reduce the final output if: # - We are not running with TP or DP @@ -584,7 +579,7 @@ def forward( # See note above re: the two all-reduce points. shared_output = self._maybe_reduce_shared_expert_output(shared_output) - shared_output, fused_output = self._maybe_apply_routed_scale_to_fused_output( + shared_output, fused_output = self._maybe_apply_routed_scale_to_output( shared_output, fused_output ) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index a454eddf28b6..feb4614d8372 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -31,7 +31,6 @@ def create_moe_runner( quant_method: FusedMoEMethodBase, enable_dbo: bool, routed_output_transform: torch.nn.Module | None = None, - apply_routed_scale_to_fused_output: bool = False, routed_scaling_factor: float = 1.0, ) -> MoERunner: return DefaultMoERunner( @@ -44,6 +43,5 @@ def create_moe_runner( quant_method, enable_dbo, routed_output_transform=routed_output_transform, - apply_routed_scale_to_fused_output=apply_routed_scale_to_fused_output, routed_scaling_factor=routed_scaling_factor, ) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 21f54ac4438f..d42fbed42ae2 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -180,7 +180,7 @@ def __init__( # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=not self.is_rocm_aiter_moe_enabled, + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index afd42d9da8c8..a63ad83f45bb 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -368,7 +368,7 @@ def __init__( use_grouped_topk=self.use_grouped_topk, router_logits_dtype=self.router_dtype, routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=True, + apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index dd8f7b7d3dc2..1b01caded94c 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -325,10 +325,9 @@ def __init__( topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), - # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=not self.is_rocm_aiter_moe_enabled, + apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/dots1.py b/vllm/model_executor/models/dots1.py index e0bd9ff150a8..c176b7365689 100644 --- a/vllm/model_executor/models/dots1.py +++ b/vllm/model_executor/models/dots1.py @@ -170,7 +170,7 @@ def __init__( scoring_func=config.scoring_func, e_score_correction_bias=self.gate.e_score_correction_bias, routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=True, + apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 530f471580b5..671e868da0ad 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -192,7 +192,7 @@ def __init__( prefix=f"{prefix}.experts", scoring_func="sigmoid", routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=True, + apply_routed_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 75f959ee43fa..2813f4a278a3 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -234,7 +234,7 @@ def __init__( if self.use_latent_moe else None, routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=True, + apply_routed_scale_to_output=True, router_logits_dtype=self.gate.out_dtype, ) diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 8e586c3139fc..7de84da51935 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -214,7 +214,7 @@ def __init__( prefix=f"{prefix}.experts", scoring_func="sigmoid", routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_fused_output=True, + apply_routed_scale_to_output=True, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, From 0c1f88c6a345e32d00b1d6db20850face25b7e12 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 17 Apr 2026 17:57:03 +0000 Subject: [PATCH 106/108] fix bug in routed scale initialization. move output transform to proper place Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/layer.py | 17 +++++++++-------- .../layers/fused_moe/modular_kernel.py | 2 +- .../layers/fused_moe/runner/moe_runner_base.py | 6 +++--- vllm/model_executor/models/nemotron_h.py | 4 +--- 4 files changed, 14 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 0df0c9036f04..bf10bc9d5c4c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -442,7 +442,14 @@ def __init__( self.topk_group = topk_group self.custom_routing_function = custom_routing_function self.scoring_func = scoring_func - self.routed_scaling_factor = routed_scaling_factor + # When apply_routed_scale_to_output is True, we set the scaling factor + # to 1.0 so it ends up being a nop. Applying the scale will be handled + # by the runner in this case. + # The member variable must be set in the same way as the router since + # some quantization methods can access it. + self.routed_scaling_factor = ( + routed_scaling_factor if not apply_routed_scale_to_output else 1.0 + ) self.e_score_correction_bias = e_score_correction_bias # TODO(bnell): end attributes @@ -461,13 +468,7 @@ def __init__( topk_group=topk_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, - # When apply_routed_scale_to_output is True, we set - # the scaling factor to 1.0 so it ends up being a nop. - # Applying the scale will be handled by the runner - # in this case. - routed_scaling_factor=routed_scaling_factor - if not apply_routed_scale_to_output - else 1.0, + routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=e_score_correction_bias, num_fused_shared_experts=self.num_fused_shared_experts, enable_eplb=enable_eplb, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f2e6e2560e70..376fcdf5b65e 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1261,7 +1261,7 @@ def _finalize( topk_ids: torch.Tensor, apply_router_weight_on_input: bool, shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor: """ The _finalize method is a wrapper around self.prepare_finalize.finalize that handles DBO, async and shared expert overlap. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 0204a6539fee..32d348070620 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -572,9 +572,6 @@ def forward( # Extract outputs from result shared_output, fused_output = _unpack(result) - # Apply output transform (e.g. latent -> full dim) - fused_output = self.apply_routed_output_transform(fused_output) - # If combine kernel already reduced fused, reduce shared to match. # See note above re: the two all-reduce points. shared_output = self._maybe_reduce_shared_expert_output(shared_output) @@ -583,6 +580,9 @@ def forward( shared_output, fused_output ) + # Apply output transform (e.g. latent -> full dim) + fused_output = self.apply_routed_output_transform(fused_output) + if shared_output is not None: result = shared_output + fused_output else: diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 2813f4a278a3..fa068639648e 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -230,9 +230,7 @@ def __init__( num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, routed_input_transform=self.fc1_latent_proj, - routed_output_transform=self.fc2_latent_proj - if self.use_latent_moe - else None, + routed_output_transform=self.fc2_latent_proj, routed_scaling_factor=self.routed_scaling_factor, apply_routed_scale_to_output=True, router_logits_dtype=self.gate.out_dtype, From 4d9179154eb4a76a1f6e09a111428faa8ad396da Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 17 Apr 2026 19:01:11 +0000 Subject: [PATCH 107/108] move contiguous back to AR call Signed-off-by: Bill Nell --- .../model_executor/layers/fused_moe/runner/moe_runner_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 32d348070620..0eac93c00c99 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -347,7 +347,7 @@ def _maybe_reduce_final_output( here. Skipped when sequence-parallel is active (SP handles its own reduction) or when the early path already reduced both outputs. """ - result = states[..., :trunc_size].contiguous() + result = states[..., :trunc_size] # We don't need to reduce the final output if: # - We are not running with TP or DP @@ -357,7 +357,7 @@ def _maybe_reduce_final_output( and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not self._fused_output_is_reduced ): - result = tensor_model_parallel_all_reduce(result) + result = tensor_model_parallel_all_reduce(result.contiguous()) return result From 2ed708826d494b4a5f8a387e6c171cdf6eaebe9d Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Sun, 19 Apr 2026 18:30:33 +0000 Subject: [PATCH 108/108] move trunc after reduce Signed-off-by: Bill Nell --- .../layers/fused_moe/runner/moe_runner_base.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 0eac93c00c99..2d2d0bb4ebe6 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -347,8 +347,6 @@ def _maybe_reduce_final_output( here. Skipped when sequence-parallel is active (SP handles its own reduction) or when the early path already reduced both outputs. """ - result = states[..., :trunc_size] - # We don't need to reduce the final output if: # - We are not running with TP or DP # - The MK already reduced the fused output itself. @@ -357,9 +355,9 @@ def _maybe_reduce_final_output( and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not self._fused_output_is_reduced ): - result = tensor_model_parallel_all_reduce(result.contiguous()) + states = tensor_model_parallel_all_reduce(states) - return result + return states[..., :trunc_size] def _encode_layer_name(self) -> str | LayerName: if _USE_LAYERNAME: