[TRTLLM-8957][feat] create communication related classes

tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+MoE Communication Strategies Module
+
+This module provides various communication strategies for expert parallelism in MoE models.
+
+Available Communication Methods:
+- AllGatherReduceScatter: Default fallback method, always available
+- MnnvlLatency: MNNVL-optimized communication for latency
+- MNNVLThroughput: MNNVL-optimized communication for throughput
+- DeepEP: Deep Expert Parallelism with support for large batches
+- DeepEPLowLatency: Deep Expert Parallelism optimized for low latency
+
+Factory:
+- CommunicationFactory: Automatically selects the best communication method
+"""
+
+from .allgather_reducescatter import AllGatherReduceScatter
+from .base import Communication
+from .communication_factory import CommunicationFactory
+from .deep_ep import DeepEP
+from .deep_ep_low_latency import DeepEPLowLatency
+from .mnnvl_latency import MnnvlLatency
+from .mnnvl_throughput import MNNVLThroughput
+
+__all__ = [
+    # Base classes and types
+    "Communication",
+    # Communication strategies
+    "AllGatherReduceScatter",
+    "MnnvlLatency",
+    "MNNVLThroughput",
+    "DeepEP",
+    "DeepEPLowLatency",
+    # Factory
+    "CommunicationFactory",
+]

tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+AllGather + ReduceScatter Communication Strategy
+
+This module implements the AllGather + ReduceScatter communication method for MoE.
+This is the default fallback strategy that always works.
+
+AllGather ALWAYS supports post-quant dispatch (quantize → allgather)
+"""
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from tensorrt_llm._torch.distributed import allgather, reducescatter
+from tensorrt_llm.mapping import Mapping
+
+from .base import Communication
+
+
+class AllGatherReduceScatter(Communication):
+    def __init__(
+        self,
+        mapping: Mapping,
+    ):
+        super().__init__(mapping)
+
+        # Initialize dispatch state
+        self._dispatch_state = {}
+
+    def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool:
+        """
+        Check if AllGather is feasible for the given workload at runtime.
+
+        AllGather is always available as fallback, so this always returns True.
+        """
+        return True
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_sf: Optional[torch.Tensor],
+        token_selected_slots: torch.Tensor,
+        token_final_scales: Optional[torch.Tensor],
+        all_rank_num_tokens: List[int],
+        use_dp_padding: Optional[bool] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]:
+        """
+        AllGather dispatch (always post-quant dispatch)
+        """
+        sizes = None if use_dp_padding else all_rank_num_tokens
+
+        hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = allgather(
+            [hidden_states, hidden_states_sf, token_selected_slots, token_final_scales],
+            self.mapping,
+            dim=0,
+            sizes=sizes,
+        )
+
+        # Store sizes for combine
+        self._dispatch_state["sizes"] = sizes
+
+        return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales
+
+    def combine(
+        self,
+        final_hidden_states: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        ReduceScatter combine phase
+        """
+        outputs = reducescatter(
+            final_hidden_states, self.mapping, dim=0, sizes=self._dispatch_state.get("sizes")
+        )
+        return outputs

tensorrt_llm/_torch/modules/fused_moe/communication/base.py

@@ -0,0 +1,162 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Base Classes for MoE Communication Strategies
+
+This module defines the abstract base class and common types for MoE communication methods.
+
+Key Design: Communication dispatch can happen BEFORE or AFTER quantization
+- Pre-quant dispatch: dispatch → quantize → allgather (DeepEP pre-quant)
+- Post-quant dispatch: quantize → allgather → dispatch (MNNVL, DeepEP post-quant)
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Optional, Tuple
+
+import torch
+
+from tensorrt_llm.mapping import Mapping
+
+
+class Communication(ABC):
+    """
+    Abstract base class for MoE communication methods
+
+    Key Design: Supports both pre-quant and post-quant dispatch
+    - Pre-quant: dispatch() called BEFORE quantization
+    - Post-quant: dispatch() called AFTER quantization
+
+    The communication method declares which mode(s) it supports via supports_post_quant_dispatch()
+    """
+
+    def __init__(
+        self,
+        mapping: Mapping,
+    ):
+        self.mapping = mapping
+        self.ep_size = mapping.moe_ep_size
+        self.ep_rank = mapping.moe_ep_rank
+
+    @abstractmethod
+    def is_workload_feasible(
+        self,
+        all_rank_num_tokens: List[int],
+        num_chunks: int,
+    ) -> bool:
+        """
+        Check if this communication strategy is feasible for the given workload at runtime.
+
+        This method performs runtime checks based on workload characteristics such as
+        token counts, number of chunks, and other runtime parameters.
+        """
+        raise NotImplementedError
+
+    def supports_post_quant_dispatch(self) -> bool:
+        """
+        Check if this strategy supports post-quantization dispatch
+
+        Returns:
+            True: Dispatch should happen AFTER quantization
+            False: Dispatch should happen BEFORE quantization
+
+        Default: True for most strategies (post-quant is more common)
+        """
+        return True
+
+    def prepare_dispatch(
+        self,
+        token_selected_slots: torch.Tensor,  # [local_num_tokens, top_k]
+        all_rank_num_tokens: List[int],  # [ep_size]
+        local_statistic_tensor: Optional[torch.Tensor] = None,  # [num_experts], for EPLB
+    ) -> Optional[torch.Tensor]:
+        """
+        Prepare dispatch metadata and information (BEFORE quantization)
+
+        This method is called before quantization to:
+        1. Gather EPLB statistics (if needed, e.g., MNNVL)
+        2. Prepare communication metadata (stored internally for later use in dispatch)
+
+        Args:
+            token_selected_slots: Selected expert slots [local_num_tokens, top_k]
+            all_rank_num_tokens: Token counts per rank [ep_size]
+            local_statistic_tensor: Local EPLB statistics [num_experts] (optional)
+
+        Returns:
+            gathered_stats: Gathered EPLB statistics across all ranks (optional)
+
+        Side effects:
+            May store internal state for use in dispatch() (e.g., MNNVL stores alltoall_info)
+        """
+        # Default: No preparation needed (AllGather, DeepEP, DeepEPLowLatency)
+        return None
+
+    @abstractmethod
+    def dispatch(
+        self,
+        # Core data
+        hidden_states: torch.Tensor,  # [local_num_tokens, hidden_size]
+        hidden_states_sf: Optional[torch.Tensor],  # [local_num_tokens, sf_size]
+        token_selected_slots: torch.Tensor,  # [local_num_tokens, top_k]
+        token_final_scales: Optional[torch.Tensor],  # [local_num_tokens, top_k]
+        # Metadata
+        all_rank_num_tokens: List[int],  # [ep_size]
+        # Optional parameters for flexibility
+        use_dp_padding: Optional[bool] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Dispatch phase: scatter/send data to different ranks
+
+        This method may read internal state from prepare_dispatch() and stores
+        dispatch metadata in self._dispatch_state for later use in combine().
+
+        Args:
+            hidden_states: Input tensor [local_num_tokens, hidden_size]
+            hidden_states_sf: Input scaling factor [local_num_tokens, sf_size]
+            token_selected_slots: Selected expert slots [local_num_tokens, top_k]
+            token_final_scales: Router weights [local_num_tokens, top_k]
+            all_rank_num_tokens: Token counts per rank [ep_size]
+            use_dp_padding: Whether to use DP padding (optional)
+            **kwargs: Strategy-specific arguments
+
+        Returns:
+            Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales)
+
+        Side effects:
+            May read from internal state set by prepare_dispatch() (e.g., MNNVL reads alltoall_info)
+            Stores dispatch state in self._dispatch_state for combine()
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def combine(
+        self,
+        final_hidden_states: torch.Tensor,  # MoE computation output
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Combine phase: gather/receive data from different ranks
+
+        This method reads dispatch metadata from self._dispatch_state that was set by dispatch().
+
+        Args:
+            final_hidden_states: Output from MoE computation
+            **kwargs: Strategy-specific arguments
+
+        Returns:
+            Combined output tensor [local_num_tokens, hidden_size]
+        """
+        raise NotImplementedError