diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py new file mode 100644 index 000000000000..ece7131ecbc3 --- /dev/null +++ b/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", +] diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py b/tensorrt_llm/_torch/modules/fused_moe/communication/allgather_reducescatter.py new file mode 100644 index 000000000000..9e175853d525 --- /dev/null +++ b/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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py b/tensorrt_llm/_torch/modules/fused_moe/communication/base.py new file mode 100644 index 000000000000..bfacf9f54db9 --- /dev/null +++ b/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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py new file mode 100644 index 000000000000..586b1cadd4de --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -0,0 +1,287 @@ +# 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. + +""" +Communication Method Factory for MoE + +Factory for creating and selecting the best communication method based on +hardware support and configuration. +""" + +import os +from typing import Optional + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._utils import local_mpi_size + +from .allgather_reducescatter import AllGatherReduceScatter +from .base import Communication +from .deep_ep import DeepEP +from .deep_ep_low_latency import DeepEPLowLatency +from .mnnvl_latency import MnnvlLatency +from .mnnvl_throughput import MNNVLThroughput + + +def is_high_throughput() -> bool: + """ + Check if high throughput mode is enabled + """ + return True + + +def is_deepep_feasible(num_ranks: int) -> bool: + """ + Check if DeepEP is feasible for the given number of ranks + + DeepEP supports two modes: + 1. Intranode: Single node with 2, 4, or 8 ranks + 2. Internode: 2, 4, 8, or 16 nodes with 8 ranks per node + """ + NUM_INTRANODE_SUPPORTED_RANKS = {2, 4, 8} + REQUIRED_LOCAL_MPI_SIZE = 8 + NUM_INTERNODE_SUPPORTED_RDMA_RANKS = {2, 4, 8, 16} + mpi_size = local_mpi_size() + + # Intranode cases + if num_ranks == mpi_size and num_ranks in NUM_INTRANODE_SUPPORTED_RANKS: + return True + + # Internode cases + if mpi_size != REQUIRED_LOCAL_MPI_SIZE: + return False + num_rdma_nodes = num_ranks // mpi_size + return num_rdma_nodes in NUM_INTERNODE_SUPPORTED_RDMA_RANKS + + +class CommunicationFactory: + """ + Factory for creating MoE communication methods + + Selects the best communication method based on: + - Hardware support (MNNVL, DeepEP) + - Configuration settings + - Workload characteristics + """ + + @staticmethod + def create_strategy( + model_config: ModelConfig, + num_experts: int, + num_slots: int, + top_k: int, + expert_size_per_partition: int, + payload_in_workspace: bool = False, + alltoall_result_do_sum: bool = False, + ) -> Optional[Communication]: + """ + Create the best communication method for the given configuration + + Selection priority: + 1. Force method (if specified via TRTLLM_FORCE_ALLTOALL_METHOD env) + 2. MNNVL (if hardware supports) + - Selects latency or throughput backend based on TRTLLM_MOE_ALLTOALL_BACKEND env + - Default: "mnnvllatency", alternative: "mnnvlthroughput" + 3. DeepEP / DeepEPLowLatency (if enabled and hardware supports) + 4. AllGather + ReduceScatter (fallback, always works) + + Args: + model_config: Model configuration containing mapping, quant_config, max_num_tokens, etc. + num_experts: Total number of experts + num_slots: Total number of expert slots + top_k: Number of experts per token + expert_size_per_partition: Number of experts per partition (required for DeepEP) + payload_in_workspace: If True, final_hidden_states is already in workspace (for MNNVLThroughput) + alltoall_result_do_sum: If True, sum the alltoall results (for MnnvlLatency) + + Returns: + The selected communication method, or None if attention does not use DP + + Note: + Most parameters are extracted from model_config. Only MoE-specific parameters + (num_experts, num_slots, top_k, expert_size_per_partition) need to be provided separately. + """ + # Extract parameters from model_config + mapping = model_config.mapping + hidden_size = model_config.pretrained_config.hidden_size + weight_dtype = model_config.torch_dtype + quant_config = model_config.quant_config + max_num_tokens = model_config.max_num_tokens + moe_max_num_tokens = model_config.moe_max_num_tokens + use_cuda_graph = model_config.use_cuda_graph + use_low_precision_combine = model_config.use_low_precision_moe_combine + + # If attention does not use data parallelism (either uses TP or single card), no MoE communication is needed + if (not mapping.enable_attention_dp) or mapping.dp_size == 1: + return None + + # If no attention DP, or if MoE TP is enabled, use AllGather + ReduceScatter + # AlltoAll cannot support MoE TP + if mapping.moe_tp_size != 1: + return AllGatherReduceScatter(mapping) + + # Check if forced method is specified via environment variable + force_method = os.environ.get("TRTLLM_FORCE_ALLTOALL_METHOD") + + if force_method is not None: + # Validate platform support for forced method + method_upper = force_method.upper() + if method_upper in ["MNNVLLATENCY", "MNNVLTHROUGHPUT"]: + if not MnnvlLatency.is_platform_supported(): + raise RuntimeError( + f"Forced method '{force_method}' is not supported on this platform. " + "MNNVLLATENCY and MNNVLTHROUGHPUT require compatible hardware." + ) + elif method_upper in ["DEEPEP", "DEEPEPLOWLATENCY"]: + if not DeepEP.is_platform_supported(mapping): + raise RuntimeError( + f"Forced method '{force_method}' is not supported on this platform. " + "DeepEP requires compatible hardware and TRTLLM_CAN_USE_DEEP_EP=1." + ) + + return CommunicationFactory._create_forced_method( + force_method, + model_config, + num_experts, + num_slots, + top_k, + expert_size_per_partition, + payload_in_workspace, + alltoall_result_do_sum, + ) + + # Try MNNVL first (highest priority) + if MnnvlLatency.is_platform_supported(): + if is_high_throughput(): + # Currently, MNNVLThroughput shows better performance at all scenarios + return MNNVLThroughput( + mapping, + num_experts, + top_k, + max_num_tokens_per_rank=max_num_tokens, + payload_in_workspace=payload_in_workspace, + ) + else: + return MnnvlLatency( + mapping, + num_experts, + num_slots, + top_k, + use_low_precision_combine, + alltoall_result_do_sum=alltoall_result_do_sum, + ) + + # Try DeepEP + if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") == "1": + if weight_dtype == torch.bfloat16: + if DeepEP.is_platform_supported(mapping) and is_deepep_feasible( + mapping.moe_ep_size + ): + return DeepEP( + mapping, + num_slots, + hidden_size, + weight_dtype, + quant_config, + expert_size_per_partition, + use_cuda_graph, + ) + else: + # Use DeepEP Low Latency as fallback (when not feasible or not supported) + return DeepEPLowLatency( + mapping, + num_slots, + hidden_size, + weight_dtype, + quant_config, + expert_size_per_partition, + max_num_tokens, + use_low_precision_combine, + moe_max_num_tokens, + ) + + # Fallback to AllGather + ReduceScatter + return AllGatherReduceScatter(mapping) + + @staticmethod + def _create_forced_method( + method: str, + model_config: ModelConfig, + num_experts: int, + num_slots: int, + top_k: int, + expert_size_per_partition: int, + payload_in_workspace: bool, + alltoall_result_do_sum: bool, + ) -> Communication: + """Create a specific method (for debugging/testing)""" + # Extract parameters from model_config + mapping = model_config.mapping + hidden_size = model_config.pretrained_config.hidden_size + weight_dtype = model_config.torch_dtype + quant_config = model_config.quant_config + max_num_tokens = model_config.max_num_tokens + moe_max_num_tokens = model_config.moe_max_num_tokens + use_cuda_graph = model_config.use_cuda_graph + use_low_precision_combine = model_config.use_low_precision_moe_combine + + method = method.upper() + + if method == "MNNVLLATENCY": + return MnnvlLatency( + mapping, + num_experts, + num_slots, + top_k, + use_low_precision_combine, + alltoall_result_do_sum=alltoall_result_do_sum, + ) + elif method == "MNNVLTHROUGHPUT": + # MNNVLThroughput requires max_num_tokens_per_rank + # max_num_tokens is per-rank value (as passed from callers like cutlass) + return MNNVLThroughput( + mapping, + num_experts, + top_k, + max_num_tokens_per_rank=max_num_tokens, + payload_in_workspace=payload_in_workspace, + ) + elif method == "DEEPEP": + return DeepEP( + mapping, + num_slots, + hidden_size, + weight_dtype, + quant_config, + expert_size_per_partition, + use_cuda_graph, + ) + elif method == "DEEPEPLOWLATENCY": + return DeepEPLowLatency( + mapping, + num_slots, + hidden_size, + weight_dtype, + quant_config, + expert_size_per_partition, + max_num_tokens, + use_low_precision_combine, + moe_max_num_tokens, + ) + elif method == "ALLGATHER": + return AllGatherReduceScatter(mapping) + else: + raise ValueError(f"Unknown communication method: {method}") diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py new file mode 100644 index 000000000000..e188d4479ac5 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py @@ -0,0 +1,238 @@ +# 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. + +""" +DeepEP Communication Strategy + +This module implements the DeepEP (Deep Expert Parallelism) communication method for MoE. +DeepEP supports both pre-quant and post-quant dispatch modes. +""" + +import os +from typing import List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + +from .base import Communication + + +class DeepEP(Communication): + """ + DeepEP strategy supporting both pre-quant and post-quant dispatch + + """ + + def __init__( + self, + mapping: Mapping, + num_slots: int, + hidden_size: int, + weight_dtype: torch.dtype, + quant_config: QuantConfig, + expert_size_per_partition: int = 0, + use_cuda_graph: bool = False, + ): + super().__init__(mapping) + + # Store needed parameters + self.num_slots = num_slots + self.hidden_size = hidden_size + self.weight_dtype = weight_dtype + self.quant_config = quant_config + + self.expert_size_per_partition = expert_size_per_partition + self.use_cuda_graph = use_cuda_graph + self.enable_postquant_alltoall = ( + os.environ.get("TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1") == "1" + ) + + # Initialize DeepEP buffer + self.deep_ep_buffer = buffer_pool.get_buffer(mapping) + self.deep_ep_buffer.reserve(hidden_size, weight_dtype) + + @staticmethod + def is_platform_supported(mapping: Mapping) -> bool: + """ + Check if DeepEP is supported on the current platform + """ + if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") != "1": + return False + return deep_ep_installed + + def supports_post_quant_dispatch(self) -> bool: + """ + DeepEP supports post-quant dispatch only for nvfp4 + """ + has_nvfp4 = self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4() + + return self.enable_postquant_alltoall and has_nvfp4 + + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: + """ + Check if DeepEP 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 weight dtype compatibility. + """ + if num_chunks > 1: + return False + if self.weight_dtype != torch.bfloat16: + return False + return self.is_platform_supported(self.mapping) + + 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]]: + """ + DeepEP dispatch + """ + all_rank_max_num_tokens = max(all_rank_num_tokens) + + if not self.supports_post_quant_dispatch(): + # Pre-quant dispatch (unquantized data) + ( + hidden_states, + recv_topk_idx, + token_final_scales, + num_recv_tokens_per_expert_list, + deep_ep_handle, + ) = self.deep_ep_buffer.dispatch( + hidden_states, + token_selected_slots, + token_final_scales, + self.num_slots, + self.expert_size_per_partition * self.ep_rank, + all_rank_max_num_tokens, + self.ep_size, + self.use_cuda_graph, + ) + + padded, hidden_states, _, token_selected_slots, token_final_scales = ( + self._pad_empty_recv_tensors(hidden_states, None, recv_topk_idx, token_final_scales) + ) + + # Store dispatch state for combine + self._dispatch_state = { + "deep_ep_handle": deep_ep_handle, + "padded": padded, + } + + else: + # Post-quant dispatch (quantized data, nvfp4 only) + + if hidden_states_sf is not None: + # Adapter between `hidden_states_sf` and DeepEP + # TODO: remove the adapter by adding dtype support to DeepEP + sf_dtype = hidden_states_sf.dtype + hidden_states_sf = hidden_states_sf.view(torch.float32) + + ( + (hidden_states, hidden_states_sf), + recv_topk_idx, + token_final_scales, + num_recv_tokens_per_expert_list, + deep_ep_handle, + ) = self.deep_ep_buffer.dispatch( + (hidden_states, hidden_states_sf), + token_selected_slots, + token_final_scales, + self.num_slots, + self.expert_size_per_partition * self.ep_rank, + all_rank_max_num_tokens, + self.ep_size, + self.use_cuda_graph, + ) + + padded, hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = ( + self._pad_empty_recv_tensors( + hidden_states, hidden_states_sf, recv_topk_idx, token_final_scales + ) + ) + + if hidden_states_sf is not None: + hidden_states_sf = hidden_states_sf.view(sf_dtype) + + # Store dispatch state for combine + self._dispatch_state = { + "deep_ep_handle": deep_ep_handle, + "padded": padded, + } + + return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales + + def combine( + self, + final_hidden_states: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """ + DeepEP combine - reads from self._dispatch_state + """ + deep_ep_handle = self._dispatch_state["deep_ep_handle"] + padded = self._dispatch_state["padded"] + + final_hidden_states = self._unpad_tensors(padded, final_hidden_states) + final_hidden_states = self.deep_ep_buffer.combine(final_hidden_states, deep_ep_handle) + + return final_hidden_states + + def _pad_empty_recv_tensors( + self, + x: torch.Tensor, + x_sf: Optional[torch.Tensor], + recv_topk_idx: torch.Tensor, + token_final_scales: torch.Tensor, + ) -> Tuple[bool, torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: + """ + Pad empty recv tensors to avoid zero-size tensor issues + """ + if x.shape[0] == 0: + padded = True + x = torch.zeros((1, x.shape[1]), dtype=x.dtype, device=x.device) + if x_sf is not None: + x_sf = torch.zeros((1, x_sf.shape[1]), dtype=x_sf.dtype, device=x_sf.device) + recv_topk_idx = torch.full( + (1, recv_topk_idx.shape[1]), + self.num_slots, + dtype=recv_topk_idx.dtype, + device=recv_topk_idx.device, + ) + token_final_scales = torch.ones( + (1, token_final_scales.shape[1]), + dtype=token_final_scales.dtype, + device=token_final_scales.device, + ) + else: + padded = False + return padded, x, x_sf, recv_topk_idx, token_final_scales + + def _unpad_tensors(self, padded: bool, final_hidden_states: torch.Tensor) -> torch.Tensor: + """ + Unpad tensors if they were padded in dispatch + """ + if padded: + final_hidden_states = final_hidden_states[:0] + return final_hidden_states diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py new file mode 100644 index 000000000000..9f25956467f9 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py @@ -0,0 +1,360 @@ +# 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. + +""" +DeepEP Low Latency Communication Strategy + +This module implements the DeepEP Low Latency communication method for MoE. +DeepEP Low Latency is optimized for small token counts with minimal communication overhead. +""" + +import os +from typing import List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import buffer_pool, deep_ep_installed +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + +from .base import Communication + + +class DeepEPLowLatency(Communication): + """ + DeepEP Low Latency strategy supporting both pre-quant and post-quant + """ + + def __init__( + self, + mapping: Mapping, + num_slots: int, + hidden_size: int, + weight_dtype: torch.dtype, + quant_config: QuantConfig, + expert_size_per_partition: int = 0, + max_num_tokens: int = 1024, + use_low_precision_combine: bool = False, + moe_max_num_tokens: Optional[int] = None, + ): + super().__init__(mapping) + + # Store needed parameters + self.num_slots = num_slots + self.hidden_size = hidden_size + self.weight_dtype = weight_dtype + self.quant_config = quant_config + self.moe_max_num_tokens = moe_max_num_tokens + + self.expert_size_per_partition = expert_size_per_partition + self.use_low_precision_combine = use_low_precision_combine + # Read from environment variable, same as wideEP + self.enable_postquant_alltoall = ( + os.environ.get("TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1") == "1" + ) + + # Calculate deep_ep_max_num_tokens + assert moe_max_num_tokens is not None + default_limit = min(max_num_tokens, moe_max_num_tokens) + self.deep_ep_max_num_tokens = int( + os.environ.get("TRTLLM_DEEP_EP_TOKEN_LIMIT", str(default_limit)) + ) + + # Set nvshmem queue pair depth larger than the number of on-flight WRs + # (ref: https://github.com/deepseek-ai/DeepEP/issues/427) + os.environ["NVSHMEM_QP_DEPTH"] = str(2 * (self.deep_ep_max_num_tokens + 1)) + + self.deep_ep_buffer = buffer_pool.get_low_latency_buffer(mapping) + self.deep_ep_buffer.reserve(self.deep_ep_max_num_tokens, hidden_size, num_slots) + + @staticmethod + def is_platform_supported(mapping: Mapping) -> bool: + """ + Check if DeepEP Low Latency is supported on the current platform + """ + if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") != "1": + return False + if not deep_ep_installed: + return False + return True + + def supports_post_quant_dispatch(self) -> bool: + """ + DeepEP Low Latency supports post-quant for: fp8_qdq, nvfp4, w4afp8 + """ + if not self.enable_postquant_alltoall: + return False + + return self._has_nvfp4() or self._has_fp8_qdq() or self._has_w4afp8() + + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: + """ + Check if DeepEP Low Latency 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 weight dtype compatibility. + """ + if num_chunks > 1: + return False + all_rank_max_num_tokens = max(all_rank_num_tokens) + if all_rank_max_num_tokens > self.deep_ep_max_num_tokens: + return False + if self.weight_dtype != torch.bfloat16: + return False + return self.is_platform_supported(self.mapping) + + 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, + pre_quant_scale: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """ + DeepEP Low Latency dispatch + """ + all_rank_max_num_tokens = max(all_rank_num_tokens) + + assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens + + deep_ep_topk_idx = token_selected_slots + deep_ep_topk_weights = token_final_scales + + if not self.supports_post_quant_dispatch(): + # Pre-quant dispatch (unquantized data) + hidden_states, recv_expert_count, deep_ep_handle = ( + self.deep_ep_buffer.low_latency_dispatch( + hidden_states, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots + ) + ) + + hidden_states, _, token_selected_slots, token_final_scales = ( + self._modify_output_to_adapt_fused_moe( + hidden_states, None, recv_expert_count, token_final_scales.dtype + ) + ) + + # Store dispatch state for combine + self._dispatch_state = { + "deep_ep_handle": deep_ep_handle, + "deep_ep_topk_idx": deep_ep_topk_idx, + "deep_ep_topk_weights": deep_ep_topk_weights, + "recv_expert_count": recv_expert_count, + } + + else: + # Post-quant dispatch (quantized data) + if self._has_fp8_qdq(): + assert hidden_states.dtype == torch.float8_e4m3fn and hidden_states_sf is None, ( + "hidden_states should be torch.float8_e4m3fn and hidden_states_sf should be None " + "in fp8 postquant alltoall" + ) + + hidden_states = hidden_states.view(torch.bfloat16) + hidden_states, recv_expert_count, deep_ep_handle = ( + self.deep_ep_buffer.low_latency_dispatch( + hidden_states, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots + ) + ) + hidden_states = hidden_states.view(torch.float8_e4m3fn) + + elif self._has_nvfp4(): + token_num = hidden_states.shape[0] + # For nvfp4, hidden_states.shape[1] is the quantized dimension (hidden_size // 2) + # We need to calculate the original hidden_size + # note: we use uint8 to store 2 fp4 values + hidden_size = hidden_states.shape[1] * 2 + + # Pre-dispatch assertions + assert ( + hidden_states.dtype == torch.uint8 + and hidden_states_sf is not None + and hidden_states_sf.dtype == torch.uint8 + ) + assert hidden_size % 32 == 0, ( + "HiddenSize should be divisible by 32 in nvfp4 postquant alltoall" + ) + assert ( + hidden_states_sf.shape[0] == token_num + and hidden_states_sf.shape[1] == hidden_size // 16 + ) + assert ( + hidden_states.shape[0] == token_num + and hidden_states.shape[1] == hidden_size // 2 + ) + + hidden_states, hidden_states_sf, recv_expert_count, deep_ep_handle = ( + self.deep_ep_buffer.low_latency_dispatch_fp4( + hidden_states, + hidden_states_sf, + deep_ep_topk_idx, + all_rank_max_num_tokens, + self.num_slots, + ) + ) + + # Post-dispatch assertions + assert hidden_states.dtype == torch.uint8 and hidden_states_sf.dtype == torch.uint8 + assert hidden_states.dim() == 3 and hidden_states_sf.dim() == 3 + assert ( + hidden_states.shape[2] == hidden_size // 2 + and hidden_states_sf.shape[2] == hidden_size // 16 + ) + + elif self._has_w4afp8(): + assert pre_quant_scale is not None, "W4AFP8 requires pre_quant_scale" + assert ( + pre_quant_scale.shape == (1, hidden_states.shape[1]) + and pre_quant_scale.dtype == hidden_states.dtype + ) + + hidden_states = ( + (hidden_states * pre_quant_scale).to(torch.float8_e4m3fn).view(torch.bfloat16) + ) + hidden_states, recv_expert_count, deep_ep_handle = ( + self.deep_ep_buffer.low_latency_dispatch( + hidden_states, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots + ) + ) + hidden_states = hidden_states.view(torch.float8_e4m3fn) + + else: + raise ValueError("Unsupported quantization mode for post-quant DeepEPLowLatency") + + hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = ( + self._modify_output_to_adapt_fused_moe( + hidden_states, hidden_states_sf, recv_expert_count, token_final_scales.dtype + ) + ) + + # Store dispatch state for combine + self._dispatch_state = { + "deep_ep_handle": deep_ep_handle, + "deep_ep_topk_idx": deep_ep_topk_idx, + "deep_ep_topk_weights": deep_ep_topk_weights, + "recv_expert_count": recv_expert_count, + } + + return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales + + def combine( + self, + final_hidden_states: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """ + DeepEP Low Latency combine - reads from self._dispatch_state + """ + deep_ep_handle = self._dispatch_state["deep_ep_handle"] + deep_ep_topk_idx = self._dispatch_state["deep_ep_topk_idx"] + deep_ep_topk_weights = self._dispatch_state["deep_ep_topk_weights"] + recv_expert_count = self._dispatch_state["recv_expert_count"] + + all_rank_max_num_tokens = kwargs.get("all_rank_max_num_tokens") + assert all_rank_max_num_tokens is not None, ( + "all_rank_max_num_tokens must be provided in kwargs" + ) + num_tokens_per_expert = self.mapping.moe_ep_size * all_rank_max_num_tokens + + final_hidden_states = final_hidden_states.view( + self.expert_size_per_partition, num_tokens_per_expert, self.hidden_size + ) + + if self.use_low_precision_combine: + if self._has_nvfp4(): + precision = "nvfp4" + global_scales = torch.ops.trtllm.calculate_nvfp4_global_scale( + final_hidden_states, recv_expert_count + ) + else: + precision = "fp8" + global_scales = None + + final_hidden_states = self.deep_ep_buffer.low_latency_combine_low_precision( + precision, + final_hidden_states, + global_scales, + deep_ep_topk_idx, + deep_ep_topk_weights, + deep_ep_handle, + ) + else: + final_hidden_states = self.deep_ep_buffer.low_latency_combine( + final_hidden_states, deep_ep_topk_idx, deep_ep_topk_weights, deep_ep_handle + ) + + return final_hidden_states + + def _has_nvfp4(self) -> bool: + """Check if NVFP4 quantization is enabled""" + return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4() + + def _has_fp8_qdq(self) -> bool: + """Check if FP8 QDQ quantization is enabled""" + return self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp8_qdq() + + def _has_w4afp8(self) -> bool: + """Check if W4AFP8 quantization is enabled""" + return ( + self.quant_config is not None + and self.quant_config.quant_mode.is_int4_weight_only_per_group() + ) + + def _modify_output_to_adapt_fused_moe( + self, + hidden_states: torch.Tensor, + hidden_states_sf: Optional[torch.Tensor], + recv_expert_count: torch.Tensor, + final_scales_dtype: torch.dtype, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: + """ + Adapter for DeepEP output to match fused_moe interface + + hidden_states shape: [#local experts, EP size * all_rank_max_num_tokens, hidden_size] + recv_expert_count shape: [#local experts] + + TODO: remove the adapter by changing `torch.ops.trtllm.fused_moe` API + """ + mask = torch.arange( + hidden_states.shape[1], dtype=torch.int32, device=hidden_states.device + ).expand(hidden_states.shape[0], hidden_states.shape[1]) < recv_expert_count.unsqueeze(1) + + token_selected_slots = torch.where( + mask, + torch.arange( + hidden_states.shape[0] * self.mapping.moe_ep_rank, + hidden_states.shape[0] * (self.mapping.moe_ep_rank + 1), + dtype=torch.int32, + device=hidden_states.device, + ).unsqueeze(1), + self.num_slots, + ) + + hidden_states = hidden_states.reshape( + hidden_states.shape[0] * hidden_states.shape[1], hidden_states.shape[2] + ) + if hidden_states_sf is not None: + hidden_states_sf = hidden_states_sf.reshape( + hidden_states_sf.shape[0] * hidden_states_sf.shape[1], hidden_states_sf.shape[2] + ) + + token_selected_slots = token_selected_slots.view(hidden_states.shape[0], 1) + token_final_scales = torch.ones_like(token_selected_slots, dtype=final_scales_dtype) + + return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/mnnvl_latency.py b/tensorrt_llm/_torch/modules/fused_moe/communication/mnnvl_latency.py new file mode 100644 index 000000000000..82a162beaf49 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/mnnvl_latency.py @@ -0,0 +1,199 @@ +# 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. + +""" +MNNVL AllToAll Communication Strategy + +This module implements the MNNVL AllToAll communication method for MoE. +MNNVL is an optimized communication strategy for NVIDIA GPU clusters. + +MNNVL supports post-quant dispatch for all quantization modes + +""" + +import os +from typing import List, Optional, Tuple + +import torch + +from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe +from tensorrt_llm.mapping import Mapping + +from .base import Communication + + +class MnnvlLatency(Communication): + """ + MNNVL AllToAll strategy for latency scenarios + """ + + def __init__( + self, + mapping: Mapping, + num_experts: int, + num_slots: int, + top_k: int = 1, + use_low_precision_combine: bool = False, + alltoall_result_do_sum: bool = False, + ): + super().__init__(mapping) + + # Store needed parameters + self.num_experts = num_experts + self.num_slots = num_slots + self.top_k = top_k + + self.use_low_precision_combine = use_low_precision_combine + self.alltoall_result_do_sum = alltoall_result_do_sum + # Read from environment variable, same as wideEP + self.enable_postquant_alltoall = ( + os.environ.get("TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1") == "1" + ) + + # Initialize MNNVL workspaces + MnnvlMemory.initialize() + self.alltoall_workspace = MnnvlMoe.get_moe_workspaces(mapping) + self.alltoall_prepare_workspace = MnnvlMoe.get_moe_prepare_workspace(mapping) + + # Initialize dispatch state + self._dispatch_state = {} + + @staticmethod + def is_platform_supported() -> bool: + """ + Check if MNNVL is supported on current hardware + """ + return MnnvlMemory.supports_mnnvl() + + def supports_post_quant_dispatch(self) -> bool: + """ + MNNVL supports post-quant for all modes + """ + return self.enable_postquant_alltoall + + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: + """ + Check if MNNVL 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. + """ + return self.is_platform_supported() + + def prepare_dispatch( + self, + token_selected_slots: torch.Tensor, + all_rank_num_tokens: List[int], + local_statistic_tensor: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """ + MNNVL prepare dispatch: gather EPLB statistics and prepare alltoall_info + """ + all_rank_max_num_tokens = max(all_rank_num_tokens) + top_k = token_selected_slots.shape[1] + + # Call MNNVL prepare to get alltoall_info and gather EPLB statistics + alltoall_info, gathered_local_statistic_tensor = ( + MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather( + token_selected_slots, + local_statistic_tensor, + self.alltoall_prepare_workspace, + all_rank_max_num_tokens, + self.ep_rank, + self.ep_size, + self.num_experts, + self.num_slots, + top_k, + ) + ) + + # Store alltoall_info in dispatch_state for use in dispatch() + self._dispatch_state["alltoall_info"] = alltoall_info + + return gathered_local_statistic_tensor + + 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]]: + """ + MNNVL dispatch (post-quant, uses alltoall_info from prepare_dispatch) + """ + # Read alltoall_info from dispatch_state (set by prepare_dispatch) + alltoall_info = self._dispatch_state.get("alltoall_info") + if alltoall_info is None: + raise ValueError("MNNVL dispatch requires prepare_dispatch() to be called first") + + all_rank_max_num_tokens = max(all_rank_num_tokens) + original_token_count = hidden_states.shape[0] # Store for combine + top_k = token_selected_slots.shape[1] + + # Dispatch quantized data using AllToAll + hidden_states, hidden_states_sf, token_selected_slots, token_final_scales = ( + MnnvlMoe.mnnvl_moe_alltoallv( + [hidden_states, hidden_states_sf, token_selected_slots, token_final_scales], + alltoall_info, + self.alltoall_workspace, + self.ep_rank, + self.ep_size, + ) + ) + + # Set expert IDs after alltoall + torch.ops.trtllm.memset_expert_ids( + token_selected_slots, + alltoall_info.recv_rank_count_cumsum, + all_rank_max_num_tokens, + top_k, + self.num_slots, + self.ep_size, + ) + + # Store original_token_count for combine (alltoall_info already stored in prepare_dispatch) + self._dispatch_state["original_token_count"] = original_token_count + + return hidden_states, hidden_states_sf, token_selected_slots, token_final_scales + + def combine( + self, + final_hidden_states: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """ + MNNVL combine - reads from self._dispatch_state + + """ + if isinstance(final_hidden_states, list): + final_hidden_states = final_hidden_states[0] + + final_hidden_states = MnnvlMoe.mnnvl_moe_alltoallv_combine( + final_hidden_states, + self._dispatch_state["alltoall_info"], + self.alltoall_workspace, + ep_rank=self.ep_rank, + ep_size=self.ep_size, + top_k=self.top_k, + token_count=self._dispatch_state["original_token_count"], + use_low_precision_combine=self.use_low_precision_combine, + do_reduce=self.alltoall_result_do_sum, + ) + + return final_hidden_states diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/mnnvl_throughput.py b/tensorrt_llm/_torch/modules/fused_moe/communication/mnnvl_throughput.py new file mode 100644 index 000000000000..d21b0041acac --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/mnnvl_throughput.py @@ -0,0 +1,383 @@ +# 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. + +""" +MNNVL AllToAll Throughput Communication Strategy + +This module implements the MNNVL AllToAll throughput communication method for MoE. +MNNVL Throughput uses Python-based AllToAll operations for high throughput scenarios. + +MNNVL Throughput supports post-quant dispatch +""" + +import os +from typing import List, Optional, Tuple + +import torch + +from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm.bindings import internal as _tllm_internal +from tensorrt_llm.logger import logger as tllm_logger +from tensorrt_llm.mapping import Mapping + +from .base import Communication + + +class MNNVLThroughput(Communication): + """ + MNNVL AllToAll strategy for throughput scenarios + + This class uses Python-based AllToAll operations for high throughput scenarios. + It manages workspace allocation and synchronization for cross-GPU communication. + """ + + # Constants from C++ (must match moeAlltoAllKernels.h) + MAX_RANKS = 64 + MAX_TOP_K = 8 + MAX_PAYLOADS = 8 + + # Single shared workspace/memory across the process + _WORKSPACE: dict | None = None + + # MetaInfo indices - initialized from C++ constants + FLAG_VAL_OFFSET_INDEX = None + LOCAL_TOKEN_COUNTER_OFFSET_INDEX = None + SEND_COUNTERS_OFFSET_INDEX = None + RECV_COUNTERS_OFFSET_INDEX = None + DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX = None + COMBINE_COMPLETION_FLAGS_OFFSET_INDEX = None + PAYLOAD_DATA_OFFSET_INDEX = None + + @classmethod + def _init_constants(cls): + """Initialize constants from C++ if not already done.""" + if cls.FLAG_VAL_OFFSET_INDEX is None: + thop = _tllm_internal.thop + cls.FLAG_VAL_OFFSET_INDEX = int(thop.MOE_A2A_FLAG_VAL_OFFSET_INDEX) + cls.LOCAL_TOKEN_COUNTER_OFFSET_INDEX = int( + thop.MOE_A2A_LOCAL_TOKEN_COUNTER_OFFSET_INDEX + ) + cls.SEND_COUNTERS_OFFSET_INDEX = int(thop.MOE_A2A_SEND_COUNTERS_OFFSET_INDEX) + cls.RECV_COUNTERS_OFFSET_INDEX = int(thop.MOE_A2A_RECV_COUNTERS_OFFSET_INDEX) + cls.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX = int( + thop.MOE_A2A_DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX + ) + cls.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX = int( + thop.MOE_A2A_COMBINE_COMPLETION_FLAGS_OFFSET_INDEX + ) + cls.PAYLOAD_DATA_OFFSET_INDEX = int(thop.MOE_A2A_PAYLOAD_DATA_OFFSET_INDEX) + + def __init__( + self, + mapping: Mapping, + num_experts: int, + top_k: int = 1, + max_num_tokens_per_rank: Optional[int] = None, + payload_in_workspace: bool = False, + ): + """ + Initialize MNNVLThroughput with workspace allocation. + + Args: + mapping: TensorRT-LLM Mapping object containing rank information + num_experts: Total number of experts + top_k: Number of experts per token + max_num_tokens_per_rank: Maximum number of tokens per rank (for workspace allocation) + payload_in_workspace: If True, final_hidden_states is already in workspace + """ + super().__init__(mapping) + + # Store needed parameters + self.num_experts = num_experts + self.top_k = top_k + + self.max_num_tokens_per_rank = max_num_tokens_per_rank + self.payload_in_workspace = payload_in_workspace + + # Initialize constants from C++ + self._init_constants() + + # Get workspace size from environment variable (default 512MB) + workspace_mb = int(os.environ.get("TRTLLM_MOE_A2A_WORKSPACE_MB", "512")) + self.workspace_size_per_rank = workspace_mb * 1024 * 1024 + # Initialize or reuse workspace + MnnvlMemory.initialize() + + if self._WORKSPACE is None: + tllm_logger.info( + f"MoE AlltoAll: Allocating workspace with size {self.workspace_size_per_rank} bytes. " + f"ep_rank: {self.ep_rank}, ep_size: {self.ep_size}, " + f"max_num_tokens_per_rank: {self.max_num_tokens_per_rank}" + ) + mnnvl_mem = MnnvlMemory(mapping, self.workspace_size_per_rank) + workspace = mnnvl_mem.as_torch_strided_tensor(torch.uint8) + metainfo = torch.ops.trtllm.moe_a2a_initialize( + workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens_per_rank, + ) + MNNVLThroughput._WORKSPACE = { + "workspace_size_per_rank": self.workspace_size_per_rank, + "max_num_tokens_per_rank": self.max_num_tokens_per_rank, + "ep_rank": self.ep_rank, + "ep_size": self.ep_size, + "mnnvl_mem": mnnvl_mem, + "workspace": workspace, + "metainfo": metainfo, + } + else: + assert self._WORKSPACE["workspace_size_per_rank"] == self.workspace_size_per_rank, ( + "reuse workspace with different workspace_size_per_rank" + ) + assert self._WORKSPACE["max_num_tokens_per_rank"] == self.max_num_tokens_per_rank, ( + "reuse workspace with different max_num_tokens_per_rank" + ) + assert self._WORKSPACE["ep_rank"] == self.ep_rank, ( + "reuse workspace with different ep_rank" + ) + assert self._WORKSPACE["ep_size"] == self.ep_size, ( + "reuse workspace with different ep_size" + ) + + self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] + self.workspace = self._WORKSPACE["workspace"] + self.moe_a2a_metainfo = self._WORKSPACE["metainfo"] + self.max_num_tokens_per_rank = self._WORKSPACE["max_num_tokens_per_rank"] + + # Initialize dispatch state + self._dispatch_state = {} + + # Internal state + self._state: str = "idle" # idle | dispatched + + # Invalid token expert ID (default to num_experts) + self.invalid_token_expert_id: int = self.num_experts + + @staticmethod + def is_platform_supported() -> bool: + """ + Check if MNNVL is supported on current hardware + """ + return MnnvlMemory.supports_mnnvl() + + def supports_post_quant_dispatch(self) -> bool: + """ + MNNVL Throughput supports post-quant dispatch + """ + return True + + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: + """ + Check if MNNVL Throughput 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. + """ + return self.is_platform_supported() + + 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]]: + """ + Dispatch phase: scatter/send data to different ranks + + 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 (unused) + + Returns: + Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) + Each tensor has shape [ep_size, max_tokens_per_rank, ...] + """ + if self._state == "dispatched": + raise RuntimeError("dispatch called twice without an intervening combine") + + # Build payloads list - token_selected_slots is always first + payloads = [] + payloads.append(token_selected_slots) + payloads.append(hidden_states) + if hidden_states_sf is not None: + payloads.append(hidden_states_sf) + if token_final_scales is not None: + payloads.append(token_final_scales) + + # Call AllToAll dispatch + ( + recv_buffers, + send_counters, + recv_counters, + topk_target_ranks, + topk_send_indices, + combine_payload_offset, + ) = torch.ops.trtllm.moe_a2a_dispatch( + token_selected_slots, + payloads, + self.workspace, + self.max_num_tokens_per_rank, + self.ep_rank, + self.ep_size, + self.top_k, + self.num_experts, + ) + + self._state = "dispatched" + + # Store all dispatch state for combine (no class variables) + self._dispatch_state["topk_target_ranks"] = topk_target_ranks + self._dispatch_state["topk_send_indices"] = topk_send_indices + self._dispatch_state["send_counters"] = send_counters + self._dispatch_state["recv_counters"] = recv_counters + self._dispatch_state["combine_payload_offset"] = int(combine_payload_offset) + + # Sanitize expert IDs for invalid tokens if needed + # token_selected_slots is always at index 0 in recv_buffers + recv_token_selected_slots = recv_buffers[0] + torch.ops.trtllm.moe_a2a_sanitize_expert_ids( + recv_token_selected_slots, + recv_counters, + int(self.invalid_token_expert_id), + ) + + # Extract results from recv_buffers + # Payload order: [token_selected_slots, hidden_states, hidden_states_sf (optional), + # token_final_scales (optional)] + token_selected_slots_recv = recv_buffers[0] + hidden_states_recv = recv_buffers[1] + if hidden_states_sf is not None: + hidden_states_sf_recv = recv_buffers[2] + token_final_scales_recv = recv_buffers[3] if token_final_scales is not None else None + else: + hidden_states_sf_recv = None + token_final_scales_recv = recv_buffers[2] if token_final_scales is not None else None + + return ( + hidden_states_recv, + hidden_states_sf_recv, + token_selected_slots_recv, + token_final_scales_recv, + ) + + def combine( + self, + final_hidden_states: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + """ + Combine phase: gather/receive data from different ranks + + Args: + final_hidden_states: Output from MoE computation + Shape: [ep_size, max_tokens_per_rank, hidden_size] or + [ep_size * max_tokens_per_rank, hidden_size] (will be reshaped) + + Returns: + Combined output tensor [local_num_tokens, hidden_size] + + """ + if self._state != "dispatched": + raise RuntimeError("combine called before a successful dispatch") + + # Read dispatch state + topk_target_ranks = self._dispatch_state.get("topk_target_ranks") + topk_send_indices = self._dispatch_state.get("topk_send_indices") + recv_counters = self._dispatch_state.get("recv_counters") + combine_payload_offset = self._dispatch_state.get("combine_payload_offset") + + if topk_target_ranks is None or topk_send_indices is None or recv_counters is None: + raise RuntimeError("combine called but dispatch state is missing") + + # Reshape if needed (handle case where input is flattened) + if final_hidden_states.dim() == 2: + # Flattened: [ep_size * max_tokens_per_rank, hidden_size] + # Reshape to: [ep_size, max_tokens_per_rank, hidden_size] + hidden_size = final_hidden_states.shape[-1] + final_hidden_states = final_hidden_states.view( + self.ep_size, self.max_num_tokens_per_rank, hidden_size + ) + elif final_hidden_states.dim() == 3: + # Already shaped: [ep_size, max_tokens_per_rank, hidden_size] + pass + else: + raise ValueError( + f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" + ) + + # Call AllToAll combine + output = torch.ops.trtllm.moe_a2a_combine( + topk_target_ranks, + topk_send_indices, + recv_counters, + final_hidden_states, + self.workspace, + self.max_num_tokens_per_rank, + self.ep_rank, + self.ep_size, + self.top_k, + int(combine_payload_offset), + bool(self.payload_in_workspace), + ) + + # Reset state for next round + self._state = "idle" + self._dispatch_state.clear() + + return output + + def get_combine_payload_tensor_in_workspace( + self, hidden_size: int, dtype: torch.dtype + ) -> torch.Tensor: + """ + Return the combine payload tensor in the workspace, which could be used + as the output of MoE kernel to avoid extra copy. + See "payload_in_workspace" in combine method. + + Args: + hidden_size: Hidden dimension size + dtype: Data type + + Returns: + Tensor view into workspace [ep_size, max_tokens_per_rank, hidden_size] + """ + if self._state != "dispatched": + raise RuntimeError( + "get_combine_payload_tensor_in_workspace called before a successful dispatch" + ) + + combine_payload_offset = self._dispatch_state.get("combine_payload_offset") + if combine_payload_offset is None: + raise RuntimeError("combine_payload_offset not found in dispatch state") + + return torch.ops.trtllm.moe_a2a_get_combine_payload_tensor( + self.workspace, + int(self.ep_rank), + int(self.ep_size), + int(self.max_num_tokens_per_rank), + int(combine_payload_offset), + dtype, + int(hidden_size), + )