Naive dispatch combine POC

vllm/distributed/device_communicators/all2all.py

@@ -62,10 +62,11 @@ def naive_multicast(
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if extra_tensors is not None:
             raise NotImplementedError(
                 "extra_tensors is not supported for NaiveAll2AllManager"
@@ -78,11 +79,14 @@ def dispatch(
         hidden_states = self.naive_multicast(
             hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel
         )
-        router_logits = self.naive_multicast(
-            router_logits, cu_tokens_across_sp_cpu, is_sequence_parallel
+        topk_weights = self.naive_multicast(
+            topk_weights, cu_tokens_across_sp_cpu, is_sequence_parallel
+        )
+        topk_ids = self.naive_multicast(
+            topk_ids, cu_tokens_across_sp_cpu, is_sequence_parallel
         )
 
-        return hidden_states, router_logits
+        return hidden_states, topk_weights, topk_ids
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
@@ -117,12 +121,13 @@ def __init__(self, cpu_group):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> (
-        tuple[torch.Tensor, torch.Tensor]
-        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         """
         Gather hidden_states and router_logits from all dp ranks.
@@ -134,7 +139,7 @@ def dispatch(
         dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
 
-        tensors_to_gather = [hidden_states, router_logits]
+        tensors_to_gather = [hidden_states, topk_weights, topk_ids]
         if extra_tensors is not None:
             tensors_to_gather.extend(extra_tensors)
 
@@ -144,9 +149,14 @@ def dispatch(
             sizes=sizes,
         )
 
-        if extra_tensors is not None:
-            return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
-        return gathered_tensors[0], gathered_tensors[1]
+        hidden_states = gathered_tensors[0]
+        topk_weights = gathered_tensors[1]
+        topk_ids = gathered_tensors[2]
+
+        if extra_tensors is None:
+            return hidden_states, topk_weights, topk_ids
+
+        return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
@@ -219,10 +229,11 @@ def get_handle(self, kwargs):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         raise NotImplementedError
 
     def combine(
@@ -267,10 +278,11 @@ def get_handle(self, kwargs):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         raise NotImplementedError
 
     def combine(

vllm/distributed/device_communicators/base_device_communicator.py

@@ -67,7 +67,8 @@ def get_handle(self, kwargs):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> Any:
@@ -283,15 +284,16 @@ def prepare_communication_buffer_for_model(self, model: torch.nn.Module) -> None
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Dispatch the hidden states and router logits to the appropriate device.
         This is a no-op in the base class.
         """
-        return hidden_states, router_logits
+        return hidden_states, topk_weights, topk_ids
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

vllm/distributed/device_communicators/cuda_communicator.py

@@ -321,17 +321,19 @@ def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
     def dispatch(  # type: ignore[override]
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> (
-        tuple[torch.Tensor, torch.Tensor]
-        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch(
             hidden_states,
-            router_logits,
+            topk_weights,
+            topk_ids,
             is_sequence_parallel,
             extra_tensors,  # type: ignore[call-arg]
         )

vllm/distributed/device_communicators/xpu_communicator.py

@@ -76,14 +76,16 @@ def broadcast(self, input_: torch.Tensor, src: int = 0) -> None:
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch(
             hidden_states,
-            router_logits,
+            topk_weights,
+            topk_ids,
             is_sequence_parallel,
             extra_tensors,  # type: ignore[call-arg]
         )

vllm/distributed/parallel_state.py

@@ -1005,22 +1005,24 @@ def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> (
-        tuple[torch.Tensor, torch.Tensor]
-        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         if self.device_communicator is not None:
             return self.device_communicator.dispatch(  # type: ignore[call-arg]
                 hidden_states,
-                router_logits,
+                topk_weights,
+                topk_ids,
                 is_sequence_parallel,
                 extra_tensors,
             )
         else:
-            return hidden_states, router_logits
+            return hidden_states, topk_weights, topk_ids
 
     def combine(
         self, hidden_states, is_sequence_parallel: bool = False

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.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNaiveEP,
+)
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_deep_ep, has_pplx
 
@@ -170,4 +173,7 @@ def maybe_make_prepare_finalize(
             local_expert_global_ids=local_expert_global_ids,
         )
 
+    elif moe.use_naive_kernels:
+        prepare_finalize = MoEPrepareAndFinalizeNaiveEP()
+
     return prepare_finalize

vllm/model_executor/layers/fused_moe/config.py

@@ -855,6 +855,12 @@ def use_deepep_ht_kernels(self):
     def use_deepep_ll_kernels(self):
         return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
 
+    @property
+    def use_naive_kernels(self):
+        return self.use_all2all_kernels and (
+            self.all2all_backend in ["allgather_reducescatter", "naive"]
+        )
+
     @staticmethod
     def flatten_tp_across_dp_and_pcp(
         tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1076,6 +1082,10 @@ def use_deepep_ht_kernels(self):
     def use_deepep_ll_kernels(self):
         return self.moe_parallel_config.use_deepep_ll_kernels
 
+    @property
+    def use_naive_kernels(self):
+        return self.moe_parallel_config.use_naive_kernels
+
     @property
     def use_flashinfer_cutlass_kernels(self):
         """

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

@@ -11,6 +11,7 @@
 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 (
+    MoEPrepareAndFinalizeNaiveEP,
     MoEPrepareAndFinalizeNoEP,
 )
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -363,13 +364,11 @@
         else:
             return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)
 
-    # FP8 DP path currently supported via AllGather.
+    # NOTE(rob): CUTLASS FP8 block quant executes the input
+    # quantzation and grouped gemm in a single kernel.
     if use_dp:
-        return FlashInferAllGatherMoEPrepareAndFinalize(
-            use_dp=True,
-            use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
+        return MoEPrepareAndFinalizeNaiveEP(
+            defer_input_quant=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/prepare_finalize.py

@@ -4,6 +4,7 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.distributed import get_ep_group
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceContiguous,
@@ -12,6 +13,115 @@
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 
 
+class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize):
+    def __init__(self, defer_input_quant: bool = False) -> None:
+        super().__init__()
+        self.defer_input_quant = defer_input_quant
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return None
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return None
+
+    def num_dispatchers(self) -> int:
+        return 1
+
+    def output_is_reduced(self) -> bool:
+        return False
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+    ) -> mk.PrepareResultType:
+        if apply_router_weight_on_input:
+            topk = topk_ids.size(1)
+            assert topk == 1, (
+                "apply_router_weight_on_input is only implemented for topk=1"
+            )
+            # Note: do not use inplace for shared experts overlap
+            a1 = a1 * topk_weights.to(a1.dtype)
+
+        # Defer input quantization to the MoE kernel.
+        if self.defer_input_quant:
+            a1q = a1
+            a1q_scale = None
+        else:
+            a1q, a1q_scale = moe_kernel_quantize_input(
+                a1,
+                quant_config.a1_scale,
+                quant_config.quant_dtype,
+                quant_config.per_act_token_quant,
+                quant_config.block_shape,
+            )
+        # TODO - this is just for deepgemm?
+        expert_tokens_meta = None
+
+        from vllm.platforms import current_platform
+
+        # The torch ops do not support fp8, so use an int8 view.
+        # Since dispatch does not do a reduce, this is safe to do.
+        use_int8_view = a1q.dtype == current_platform.fp8_dtype()
+        if use_int8_view:
+            a1q = a1q.view(torch.int8)
+
+        # Skip gathering scales if we have static quantization
+        # (the scale is a scalar, replicated on all ranks) or
+        # if quantization is deferred.
+        skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0
+        scales = None if skip_gather_scales else [a1q_scale]
+
+        res = get_ep_group().dispatch(
+            a1q,
+            topk_weights,
+            topk_ids,
+            is_sequence_parallel=False,  # TODO: support SP
+            extra_tensors=scales,
+        )
+        if skip_gather_scales:
+            a1q, topk_weights, topk_ids = res
+        else:
+            a1q, topk_weights, topk_ids, scales = res
+            a1q_scale = res[0]
+
+        if use_int8_view:
+            a1q = a1q.view(current_platform.fp8_dtype())
+
+        return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
+            weight_and_reduce_impl = TopKWeightAndReduceContiguous()
+
+        out = weight_and_reduce_impl.apply(
+            output=None,
+            fused_expert_output=fused_expert_output,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )
+
+        output.copy_(get_ep_group().combine(out, is_sequence_parallel=False))
+
+
 class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
     def __init__(self, defer_input_quant: bool = False) -> None:
         super().__init__()