From 54109c2bd1ad6571cacf11c20e0a7706ab849942 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Fri, 21 Aug 2026 19:14:54 +0000 Subject: [PATCH] add UnfinalizedMoEOutput proto Signed-off-by: yewentao256 --- .../layers/fused_moe/runner/moe_runner.py | 15 +++++++++++---- .../fused_moe/unquantized_fused_moe_method.py | 3 ++- .../compressed_tensors_moe_w4a4_nvfp4.py | 3 ++- .../layers/quantization/modelopt.py | 3 ++- .../layers/quantization/online/moe_base.py | 3 ++- .../layers/quantization/quark/quark_moe.py | 3 ++- 6 files changed, 21 insertions(+), 9 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 55d50e44ab82..66ad630e680b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -825,8 +825,13 @@ def _maybe_combine( | tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput] ): if self.do_naive_dispatch_combine: + if isinstance(hidden_states, UnfinalizedMoEOutput): + raise RuntimeError( + "Naive expert-parallel combine cannot consume a deferred " + "MoE output." + ) hidden_states = get_ep_group().combine( - cast(torch.Tensor, hidden_states), + hidden_states, self.moe_config.is_sequence_parallel, ) @@ -834,9 +839,11 @@ def _maybe_combine( self.moe_config.pcp_size > 1 and not self.moe_config.moe_parallel_config.use_all2all_kernels ): - hidden_states = get_pcp_group().reduce_scatter( - cast(torch.Tensor, hidden_states), dim=0 - ) + if isinstance(hidden_states, UnfinalizedMoEOutput): + raise RuntimeError( + "PCP reduce-scatter cannot consume a deferred MoE output." + ) + hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0) if self.shared_experts is not None: assert shared_output is not None 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 6ed1730737e9..31404c93a6fb 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 @@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, convert_to_unquantized_kernel_format, @@ -299,7 +300,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 92728691d2de..cf98e52bb61a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -14,6 +14,7 @@ FusedMoEConfig, FusedMoEQuantConfig, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( NvFp4MoeBackend, convert_to_nvfp4_moe_kernel_format, @@ -278,7 +279,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 2adf21b46f56..241a7c8adaf3 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -24,6 +24,7 @@ RoutedExperts, SharedExperts, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, convert_to_fp8_moe_kernel_format, @@ -1624,7 +1625,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 8f47768a64a1..866c88b3d649 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -10,6 +10,7 @@ RoutedExperts, SharedExperts, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.model_loader.reload.layerwise import ( initialize_online_processing, ) @@ -139,7 +140,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 752d2e423ad2..1000b427051e 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -26,6 +26,7 @@ mxfp4_w4a16_moe_quant_config, ocp_mx_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, @@ -1444,7 +1445,7 @@ def apply_monolithic( x: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | UnfinalizedMoEOutput: assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic(