[MoE Refactor][13/N] Convert FI to Use PFNoEP (vllm-project#31533)

Signed-off-by: Robert Shaw <[email protected]>
Signed-off-by: Robert Shaw <[email protected]>
Signed-off-by: Robert Shaw <[email protected]>
Co-authored-by: Robert Shaw <[email protected]>
Co-authored-by: Wentao Ye <[email protected]>

Author: 3 people

vllm/model_executor/layers/fused_moe/all2all_utils.py

@@ -15,6 +15,9 @@
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
     FusedMoEPrepareAndFinalize,
 )
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+    build_flashinfer_fp8_cutlass_moe_prepare_finalize,
+)
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_deep_ep, has_pplx
 
@@ -77,10 +80,17 @@ def maybe_make_prepare_finalize(
 
     prepare_finalize: FusedMoEPrepareAndFinalize | None = None
 
-    # TODO: could allow this now
-    assert not moe.use_flashinfer_cutlass_kernels, "Must be created in modelopt.py"
+    if moe.use_flashinfer_cutlass_kernels:
+        assert quant_config is not None
+        use_deepseek_fp8_block_scale = (
+            quant_config is not None and quant_config.is_block_quantized
+        )
+        prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
+            moe=moe,
+            use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
+        )
 
-    if moe.use_pplx_kernels:
+    elif moe.use_pplx_kernels:
         assert quant_config is not None
 
         hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

@@ -10,6 +10,9 @@
 )
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
@@ -349,14 +352,23 @@ def create_flashinfer_prepare_finalize(
     use_nvfp4: bool = False,
     enable_alltoallv: bool = False,
     use_deepseek_fp8_block_scale: bool = False,
-) -> FlashInferCutlassMoEPrepareAndFinalize:
+) -> FlashInferCutlassMoEPrepareAndFinalize | MoEPrepareAndFinalizeNoEP:
     """Factory function to create the appropriate FlashInfer implementation."""
+    # TODO(rob): migrate non-DP cases to MoEPrepareAndFinalizeNoEP
+    # once we complete the FP8 refactor.
     if use_nvfp4:
         if enable_alltoallv:
             return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)
         else:
             return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)
-    # FP8 path currently supported via AllGather; optionally enable block-scale
-    return FlashInferAllGatherMoEPrepareAndFinalize(
-        use_dp=use_dp, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale
-    )
+
+    # FP8 DP path currently supported via AllGather.
+    if use_dp:
+        return FlashInferAllGatherMoEPrepareAndFinalize(
+            use_dp=True,
+            use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
+        )
+    else:
+        # NOTE(rob): CUTLASS FP8 block quant executes the input
+        # quantzation and grouped gemm in a single kernel.
+        return MoEPrepareAndFinalizeNoEP(defer_input_quant=use_deepseek_fp8_block_scale)

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

@@ -49,7 +49,6 @@ def make(
                 prepare_finalize,
                 old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),
                 shared_experts,
-                getattr(moe_layer, "shared_experts_stream", None),
                 moe_parallel_config=moe_layer.moe_parallel_config,
             ),
         )

vllm/model_executor/layers/fused_moe/layer.py

@@ -356,14 +356,14 @@ def __init__(
         # TODO: Remove this after more extensive testings with TP/DP
         # and other execution modes
         if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:
-            logger.info_once("Disabling MoE shared_experts cuda 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.info_once(
+                logger.debug_once(
                     "Enabled separate cuda stream for MoE shared_experts", scope="local"
                 )
 

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -21,7 +21,6 @@
     count_expert_num_tokens,
     disable_inplace,
 )
-from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.v1.worker.ubatching import (
     dbo_enabled,
@@ -682,14 +681,12 @@ def __init__(
         prepare_finalize: FusedMoEPrepareAndFinalize,
         fused_experts: FusedMoEPermuteExpertsUnpermute,
         shared_experts: torch.nn.Module | None = None,
-        shared_experts_stream: torch.cuda.Stream | None = None,
         moe_parallel_config: FusedMoEParallelConfig | None = None,
     ):
         super().__init__()
         self.prepare_finalize = prepare_finalize
         self.fused_experts = fused_experts
         self.shared_experts = shared_experts
-        self.shared_experts_stream = shared_experts_stream
 
         # prefer an explicit FusedMoEParallelConfig when available (from
         # FusedMoE layers / tests).
@@ -904,34 +901,6 @@ def _slice_expert_tokens_metadata(
             expert_num_tokens_cpu=c_expert_num_tokens_cpu,
         )
 
-    def _maybe_setup_shared_experts_stream(
-        self, hidden_states: torch.Tensor
-    ) -> tuple[bool, torch.Tensor | None]:
-        # decide whether to run shared experts on a separate CUDA stream to
-        # overlap with the main fused MoE kernel.
-        use_shared_experts_stream = (
-            self.shared_experts is not None
-            and self.shared_experts_stream is not None
-            and hidden_states.is_cuda
-            and (
-                hidden_states.shape[0]
-                <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
-            )
-        )
-
-        hidden_states_clone: torch.Tensor | None = None
-        if use_shared_experts_stream and self.shared_experts_stream is not None:
-            # TODO: Optimize this (complicated)
-            # Note: this clone adds overhead but is required
-            # for correctness with multiple CUDA streams and CUDA graph capture.
-            hidden_states_clone = hidden_states.clone()
-            # record that the clone will be used by the separate stream so its
-            # lifetime is correctly tracked.
-            hidden_states_clone.record_stream(self.shared_experts_stream)
-            self.shared_experts_stream.wait_stream(torch.cuda.current_stream())
-
-        return use_shared_experts_stream, hidden_states_clone
-
     def _prepare(
         self,
         hidden_states: torch.Tensor,
@@ -1119,30 +1088,12 @@ def _finalize(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
         apply_router_weight_on_input: bool,
-        hidden_states_clone: torch.Tensor | None = None,
-        use_shared_experts_stream: bool = False,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """
         The _finalize method is a wrapper around self.prepare_finalize.finalize
         that handles DBO, async and shared expert overlap.
         """
-
-        def maybe_run_shared_experts() -> torch.Tensor | None:
-            if self.shared_experts is None:
-                return None
-
-            if (
-                not use_shared_experts_stream
-                or self.shared_experts_stream is not None
-                and (not hidden_states.is_cuda or not torch.cuda.is_available())
-            ):
-                # fall back to running on the current stream
-                return self.shared_experts(hidden_states)
-
-            assert hidden_states_clone is not None
-            # launch shared experts on the dedicated stream.
-            with torch.cuda.stream(self.shared_experts_stream):
-                return self.shared_experts(hidden_states_clone)
+        shared_output: torch.Tensor | None = None
 
         if not self.prepare_finalize.supports_async():
             assert not dbo_enabled()
@@ -1155,7 +1106,8 @@ def maybe_run_shared_experts() -> torch.Tensor | None:
                 apply_router_weight_on_input,
                 self.fused_experts.finalize_weight_and_reduce_impl(),
             )
-            shared_output = maybe_run_shared_experts()
+            if self.shared_experts is not None:
+                shared_output = self.shared_experts(hidden_states)
         else:
             finalize_ret = self.prepare_finalize.finalize_async(
                 output,
@@ -1165,8 +1117,8 @@ def maybe_run_shared_experts() -> torch.Tensor | None:
                 apply_router_weight_on_input,
                 self.fused_experts.finalize_weight_and_reduce_impl(),
             )
-
-            shared_output = maybe_run_shared_experts()
+            if self.shared_experts is not None:
+                shared_output = self.shared_experts(hidden_states)
 
             # TODO(lucas): refactor this in the alternative schedules followup
             # currently unpack if we have hook + receiver pair or just
@@ -1189,28 +1141,12 @@ def maybe_run_shared_experts() -> torch.Tensor | None:
 
             receiver()
 
-        self._wait_for_shared_experts_stream(hidden_states, use_shared_experts_stream)
-
         if self.shared_experts is None:
             return output
         else:
             assert shared_output is not None
             return shared_output, output
 
-    def _wait_for_shared_experts_stream(
-        self, hidden_states: torch.Tensor, use_shared_experts_stream: bool
-    ) -> None:
-        # ensure that any work enqueued on the shared_experts_stream is
-        # completed before the shared_output tensor is consumed
-        if (
-            self.shared_experts is not None
-            and use_shared_experts_stream
-            and self.shared_experts_stream is not None
-            and hidden_states.is_cuda
-            and current_platform.is_cuda()
-        ):
-            torch.cuda.current_stream().wait_stream(self.shared_experts_stream)
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -1257,10 +1193,6 @@ def forward(
         else:
             output = torch.zeros_like(hidden_states)
 
-        use_shared_experts_stream, hidden_states_clone = (
-            self._maybe_setup_shared_experts_stream(hidden_states)
-        )
-
         local_num_experts = w1.size(0)
         if global_num_experts == -1:
             global_num_experts = local_num_experts
@@ -1297,6 +1229,4 @@ def forward(
             topk_weights,
             topk_ids,
             apply_router_weight_on_input,
-            hidden_states_clone=hidden_states_clone,
-            use_shared_experts_stream=use_shared_experts_stream,
         )