Enable MnnvlThroughput in Trtllm MoE.

Signed-off-by: Bo Li <[email protected]>

Author: bobboli

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -6,6 +6,7 @@
 from torch import nn
 
 from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe
+from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll
 from tensorrt_llm._utils import get_sm_version
 
 from ...custom_ops.trtllm_gen_custom_ops import \
@@ -112,11 +113,26 @@ def __init__(
         self.alltoall_workspace = None
         self.alltoall_prepare_workspace = None
         if self.enable_alltoall:
-            MnnvlMemory.initialize()
-            self.alltoall_workspace = MnnvlMoe.get_moe_workspaces(
-                model_config.mapping)
-            self.alltoall_prepare_workspace = MnnvlMoe.get_moe_prepare_workspace(
-                model_config.mapping)
+            if self.moe_alltoall_backend == "MnnvlLatency":
+                MnnvlMemory.initialize()
+                self.alltoall_workspace = MnnvlMoe.get_moe_workspaces(
+                    model_config.mapping)
+                self.alltoall_prepare_workspace = MnnvlMoe.get_moe_prepare_workspace(
+                    model_config.mapping)
+            elif self.moe_alltoall_backend == "MnnvlThroughput":
+                workspace_mb = int(
+                    os.environ.get("TRTLLM_MOE_A2A_WORKSPACE_MB", "512"))
+                self.moe_a2a = MoeAlltoAll(
+                    mapping=self.mapping,
+                    max_num_tokens_per_rank=model_config.max_num_tokens,
+                    top_k=self.routing_method.experts_per_token,
+                    num_experts=self.num_experts,
+                    workspace_size_per_rank=workspace_mb * 1024 * 1024,
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported moe alltoall backend: {self.moe_alltoall_backend}"
+                )
 
         self._weights_created = False
         if not model_config.skip_create_weights_in_init:
@@ -131,6 +147,12 @@ def enable_alltoall(self):
                 and os.environ.get("TRTLLM_MOE_DISABLE_ALLTOALLV", "0") != "1"
                 and MnnvlMemory.supports_mnnvl())
 
+    @cached_property
+    def moe_alltoall_backend(self):
+        # "MnnvlLatency" (default) or "MnnvlThroughput"
+        return os.environ.get("TRTLLM_MOE_ALLTOALL_BACKEND",
+                              "MnnvlLatency").strip().lower()
+
     def _check_configs(self):
         assert self.has_deepseek_fp8_block_scales \
             or self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_nvfp4_fp8 \
@@ -298,45 +320,89 @@ def forward_impl(
             else:
                 token_final_scales = token_final_scales.to(torch.float32)
 
-            assert self.alltoall_prepare_workspace is not None, "alltoall_prepare_workspace should be initialized"
-            alltoall_info, _ = MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather(
-                token_selected_experts,
-                None,
-                self.alltoall_prepare_workspace,
-                max_num_token,
-                self.ep_rank,
-                self.ep_size,
-                self.num_experts,
-                self.num_slots,
-                top_k,
-            )
+            if self.moe_alltoall_backend == "MnnvlLatency":
+                assert self.alltoall_prepare_workspace is not None, "alltoall_prepare_workspace should be initialized"
+                alltoall_info, _ = MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather(
+                    token_selected_experts,
+                    None,
+                    self.alltoall_prepare_workspace,
+                    max_num_token,
+                    self.ep_rank,
+                    self.ep_size,
+                    self.num_experts,
+                    self.num_slots,
+                    top_k,
+                )
 
-            if x_sf is not None:
-                x_sf = x_sf.view(x_row, ceil_div(x_col,
-                                                 self.scaling_vector_size))
+                if x_sf is not None:
+                    x_sf = x_sf.view(x_row,
+                                     ceil_div(x_col, self.scaling_vector_size))
 
-            x, x_sf, token_selected_experts, token_final_scales = MnnvlMoe.mnnvl_moe_alltoallv(
-                [x, x_sf, token_selected_experts, token_final_scales],
-                alltoall_info,
-                self.alltoall_workspace,
-                self.ep_rank,
-                self.ep_size,
-            )
+                x, x_sf, token_selected_experts, token_final_scales = MnnvlMoe.mnnvl_moe_alltoallv(
+                    [x, x_sf, token_selected_experts, token_final_scales],
+                    alltoall_info,
+                    self.alltoall_workspace,
+                    self.ep_rank,
+                    self.ep_size,
+                )
 
-            torch.ops.trtllm.memset_expert_ids(
-                token_selected_experts,
-                alltoall_info.recv_rank_count_cumsum,
-                max_num_token,
-                top_k,
-                self.num_slots,
-                self.ep_size,
-            )
+                torch.ops.trtllm.memset_expert_ids(
+                    token_selected_experts,
+                    alltoall_info.recv_rank_count_cumsum,
+                    max_num_token,
+                    top_k,
+                    self.num_slots,
+                    self.ep_size,
+                )
 
-            if x_sf is not None:
-                x_sf = x_sf.flatten()
+                if x_sf is not None:
+                    x_sf = x_sf.flatten()
+
+                if token_final_scales is not None:
+                    token_final_scales = token_final_scales.to(torch.bfloat16)
+            elif self.moe_alltoall_backend == "MnnvlThroughput":
+                if x_sf is not None:
+                    x_sf = x_sf.view(x_row,
+                                     ceil_div(x_col, self.scaling_vector_size))
+
+                payloads = []
+                payloads.append(x)
+                if x_sf is not None:
+                    payloads.append(x_sf)
+                    expert_id_payload_index = 2
+                else:
+                    expert_id_payload_index = 1
+                payloads.append(token_selected_experts)
+                payloads.append(token_final_scales)
+
+                recv_buffers = self.moe_a2a.dispatch(
+                    token_selected_experts,
+                    payloads,
+                    invalid_token_expert_id=
+                    -1,  # Note Cutlass MoE uses num_experts as invalid token expert id
+                    expert_id_payload_index=expert_id_payload_index,
+                )
 
-            if token_final_scales is not None:
-                token_final_scales = token_final_scales.to(torch.bfloat16)
+                if x_sf is not None:
+                    x_recv, x_sf_recv, token_selected_experts_recv, token_final_scales_recv = recv_buffers
+                    x_sf = x_sf_recv.view(-1, x_sf_recv.shape[-1])
+                else:
+                    x_recv, token_selected_experts_recv, token_final_scales_recv = recv_buffers
+                x = x_recv.view(-1, x_recv.shape[-1])
+                token_selected_experts = token_selected_experts_recv.view(
+                    -1, token_selected_experts_recv.shape[-1])
+                token_final_scales = token_final_scales_recv.view(
+                    -1, token_final_scales_recv.shape[-1])
+
+                if x_sf is not None:
+                    x_sf = x_sf.flatten()
+
+                if token_final_scales is not None:
+                    token_final_scales = token_final_scales.to(torch.bfloat16)
+            else:
+                raise ValueError(
+                    f"Unsupported moe alltoall backend: {self.moe_alltoall_backend}"
+                )
 
         elif run_post_quant_allgather:
             if x_sf is not None:
@@ -600,16 +666,28 @@ def forward_impl(
             )
 
         # Combine results if using alltoall
-        if self.enable_alltoall and alltoall_info is not None:
-            final_hidden_states = MnnvlMoe.mnnvl_moe_alltoallv_combine(
-                final_hidden_states,
-                alltoall_info,
-                self.alltoall_workspace,
-                ep_rank=self.ep_rank,
-                ep_size=self.ep_size,
-                top_k=top_k,
-                token_count=token_count,
-            )
+        if self.enable_alltoall:
+            if self.moe_alltoall_backend == "MnnvlLatency":
+                if alltoall_info is not None:
+                    final_hidden_states = MnnvlMoe.mnnvl_moe_alltoallv_combine(
+                        final_hidden_states,
+                        alltoall_info,
+                        self.alltoall_workspace,
+                        ep_rank=self.ep_rank,
+                        ep_size=self.ep_size,
+                        top_k=top_k,
+                        token_count=token_count,
+                    )
+            elif self.moe_alltoall_backend == "MnnvlThroughput":
+                hidden = final_hidden_states.shape[-1]
+                payload = final_hidden_states.view(
+                    self.ep_size, self.moe_a2a.max_num_tokens_per_rank, hidden)
+                final_hidden_states = self.moe_a2a.combine(
+                    payload, payload_in_workspace=False)
+            else:
+                raise ValueError(
+                    f"Unsupported moe alltoall backend: {self.moe_alltoall_backend}"
+                )
 
         final_hidden_states = self.reducescatter_or_allreduce(
             final_hidden_states,