[MCore] Enable dense routing maps for Flex dispatch

megatron/core/extensions/transformer_engine.py

@@ -3673,10 +3673,20 @@ def te_general_gemm(
         fused_topk_with_score_function,
     )
 
+    try:
+        _fused_topk_sig = inspect.signature(fused_topk_with_score_function)
+        fused_topk_with_score_function_supports_topk_indices = (
+            "topk_indices" in _fused_topk_sig.parameters
+        )
+        del _fused_topk_sig
+    except (TypeError, ValueError):
+        fused_topk_with_score_function_supports_topk_indices = False
+
 else:
     fused_topk_with_score_function = None
     fused_compute_score_for_moe_aux_loss = None
     fused_moe_aux_loss = None
+    fused_topk_with_score_function_supports_topk_indices = False
 
 
 def set_save_original_input(module):

megatron/core/transformer/moe/fused_a2a.py

@@ -4,7 +4,7 @@
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
 import inspect
-from typing import Optional
+from typing import Callable, Optional
 
 from megatron.core.utils import internal_api
 
@@ -268,12 +268,37 @@ def set_deepep_num_sms(num_sms):
     set_deepep_num_sms = None
 
 
+def _has_parameter(function: Callable, parameter: str) -> bool:
+    """Return whether a callable exposes a named parameter."""
+    try:
+        return parameter in inspect.signature(function).parameters
+    except (TypeError, ValueError):
+        return False
+
+
 try:
     from deep_ep import HybridEPBuffer
 
     HAVE_HYBRIDEP = True
+    HAVE_HYBRIDEP_EXPLICIT_DENSE_ROUTING = _has_parameter(
+        HybridEPBuffer.dispatch_with_permute, "dense_routing"
+    )
+    try:
+        import hybrid_ep_cpp
+
+        HAVE_HYBRIDEP_INFERRED_DENSE_ROUTING = hasattr(
+            hybrid_ep_cpp.HybridEpConfigInstance(), "topk"
+        )
+    except (ImportError, AttributeError, TypeError, ValueError):
+        HAVE_HYBRIDEP_INFERRED_DENSE_ROUTING = False
+    HAVE_HYBRIDEP_DENSE_ROUTING = (
+        HAVE_HYBRIDEP_EXPLICIT_DENSE_ROUTING or HAVE_HYBRIDEP_INFERRED_DENSE_ROUTING
+    )
 except ImportError:
     HAVE_HYBRIDEP = False
+    HAVE_HYBRIDEP_EXPLICIT_DENSE_ROUTING = False
+    HAVE_HYBRIDEP_INFERRED_DENSE_ROUTING = False
+    HAVE_HYBRIDEP_DENSE_ROUTING = False
 
 _hybrid_ep_buffer = None
 
@@ -376,16 +401,16 @@ def forward(
         num_permuted_tokens=None,
         pad_multiple=None,
         num_sms_preprocessing_api=108,
+        topk_idx=None,
+        num_of_experts=None,
     ):
         '''
         Forward pass of fused dispatch of the HybridEP backend
         '''
         if fused or num_blocks_permute is not None or num_blocks_unpermute is not None:
-            import inspect
             import warnings
 
-            sig = inspect.signature(HybridEPBuffer.dispatch_with_permute)
-            if 'fuse_permute_dispatch' not in sig.parameters:
+            if not _has_parameter(HybridEPBuffer.dispatch_with_permute, 'fuse_permute_dispatch'):
                 warnings.warn(
                     "Current DeepEP version does not support fused permute dispatch or "
                     "num_blocks_permute/num_blocks_unpermute. Falling back to unfused "
@@ -415,7 +440,21 @@ def forward(
         # If we provide the num_permuted_tokens, we do not need to use sync to
         # wait for the data in pinned memory ready
         non_blocking = num_permuted_tokens is not None
-        # Process the dispatch
+        use_dense = topk_idx is not None and HAVE_HYBRIDEP_DENSE_ROUTING
+        if use_dense:
+            assert num_of_experts is not None, "num_of_experts is required for dense routing"
+            dense_kwargs = {"dense_routing": True} if HAVE_HYBRIDEP_EXPLICIT_DENSE_ROUTING else {}
+            dispatch_kwargs = {
+                "topk_idx": topk_idx,
+                "num_of_experts": num_of_experts,
+                **dense_kwargs,
+            }
+        else:
+            assert (
+                routing_map is not None
+            ), "routing_map is required when dense HybridEP routing is unavailable"
+            dispatch_kwargs = {"routing_map": routing_map}
+
         (
             dispatched_hidden,
             dispatched_probs,
@@ -424,14 +463,14 @@ def forward(
             handle,
         ) = _hybrid_ep_buffer.dispatch_with_permute(
             hidden=x,
-            routing_map=routing_map,
             probs=probs,
             scaling_factor=None,
             num_of_experts_per_rank=num_local_experts,
             pad_multiple=pad_multiple,
             num_permuted_tokens=num_permuted_tokens,
             non_blocking=non_blocking,
             **({"fuse_permute_dispatch": fused} if fused else {}),
+            **dispatch_kwargs,
         )
 
         ctx.handle = handle
@@ -472,6 +511,8 @@ def backward(ctx, grad_x, grad_probs, grad_scaling_factor, grad_tokens_per_exper
             None,
             None,
             None,
+            None,
+            None,
         )
 
 
@@ -532,6 +573,8 @@ def hybrid_ep_dispatch(
         num_permuted_tokens=None,
         pad_multiple=None,
         num_sms_preprocessing_api=108,
+        topk_idx=None,
+        num_of_experts=None,
     ):
         '''
         Perform fused dispatch for "permute + dispatch a2a + permute" using the
@@ -565,6 +608,10 @@ def hybrid_ep_dispatch(
                 is performed.
             num_sms_preprocessing_api (int):
                 Number of SMs used by the preprocessing (metadata scan) kernel.
+            topk_idx (torch.Tensor, optional):
+                Dense top-k expert indices with shape [num_tokens, topk].
+            num_of_experts (int, optional):
+                Total number of experts. Required when topk_idx is provided.
         '''
         return HybridEPDispatch.apply(
             x,
@@ -580,6 +627,8 @@ def hybrid_ep_dispatch(
             num_permuted_tokens,
             pad_multiple,
             num_sms_preprocessing_api,
+            topk_idx,
+            num_of_experts,
         )
 
     @internal_api

megatron/core/transformer/moe/moe_utils.py

@@ -42,6 +42,7 @@
         fused_sort_chunks_by_index,
         fused_sort_chunks_by_index_with_probs,
         fused_topk_with_score_function,
+        fused_topk_with_score_function_supports_topk_indices,
         fused_unpermute,
         te_general_gemm,
     )
@@ -55,9 +56,10 @@
         fused_sort_chunks_by_index,
         fused_sort_chunks_by_index_with_probs,
         fused_topk_with_score_function,
+        fused_topk_with_score_function_supports_topk_indices,
         fused_unpermute,
         te_general_gemm,
-    ) = (None, None, None, None, None, None, None, None, None, None)
+    ) = (None, None, None, None, None, None, None, None, False, None, None)
 
 
 def switch_load_balancing_loss_func(
@@ -776,6 +778,7 @@ def topk_routing_with_score_function(
     router_replay: Optional['RouterReplay'] = None,
     dense_output: bool = False,
     precomputed_indices: Optional[torch.Tensor] = None,
+    topk_indices: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Compute the routing probabilities and map for top-k selection with score function.
 
@@ -804,6 +807,8 @@ def topk_routing_with_score_function(
                                        selected by the caller. When given, the score function's
                                        own top-k is bypassed and probs are computed at these
                                        indices (e.g. for quantile balancing). Defaults to None.
+        topk_indices (torch.Tensor, optional): Optional dense top-k index output buffer with shape
+                                               [num_tokens, topk]. Only used by the fused TE path.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]:
@@ -813,7 +818,8 @@ def topk_routing_with_score_function(
                   entries correspond to the top-k selected experts per token.
                 - routing_map (torch.Tensor): Shape [num_tokens, num_experts]. Boolean mask where
                   True indicates the token is routed to that expert (i.e. the expert was in the
-                  token's top-k selection).
+                  token's top-k selection). When topk_indices is provided, this is instead that
+                  [num_tokens, topk] dense index buffer.
             When dense_output=True:
                 - probs (torch.Tensor): Shape [num_tokens, topk]. The normalized routing
                   probabilities for each token's top-k selected experts.
@@ -844,6 +850,11 @@ def topk_routing_with_score_function(
             scaling_factor=scaling_factor,
             score_function=score_function,
             expert_bias=expert_bias,
+            **(
+                {"topk_indices": topk_indices}
+                if fused_topk_with_score_function_supports_topk_indices and topk_indices is not None
+                else {}
+            ),
         )
 
     def _compute_topk(

megatron/core/transformer/moe/router.py

@@ -11,6 +11,7 @@
     is_batch_invariant_mode_enabled,
 )
 from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe.fused_a2a import HAVE_HYBRIDEP_DENSE_ROUTING
 from megatron.core.transformer.moe.moe_logging import get_moe_metrics_tracker
 from megatron.core.transformer.moe.moe_utils import (
     MoEAuxLossAutoScaler,
@@ -19,6 +20,7 @@
     apply_random_logits,
     apply_router_token_dropping,
     compute_routing_scores_for_aux_loss,
+    fused_topk_with_score_function_supports_topk_indices,
     get_tokens_per_expert_and_token_count,
     qb_dual_update,
     router_gating_linear,
@@ -30,6 +32,8 @@
 from megatron.core.transformer.moe.router_replay import RouterReplay
 from megatron.core.transformer.transformer_config import TransformerConfig
 
+_HYBRIDEP_INT16_EXPERT_LIMIT = 1 << 15
+
 
 class Router(ABC, MegatronModule):
     """Base Router class"""
@@ -55,6 +59,7 @@ def __init__(
         self.layer_number = None
         self.is_mtp_layer = is_mtp_layer
         self.tp_group = pg_collection.tp
+        self.expt_tp_group = pg_collection.expt_tp
         self.cp_group = pg_collection.cp
         self.tp_cp_group = pg_collection.tp_cp
         self.tp_dp_cp_group = pg_collection.tp_dp_cp
@@ -411,6 +416,32 @@ def is_aux_loss_enabled(self) -> bool:
                 return True
         return False
 
+    def _dense_route_indices_dtype(self) -> Optional[torch.dtype]:
+        """Return the route-index dtype for Flex backends that consume dense top-k indices."""
+        if not self.config.moe_router_fusion:
+            return None
+        if self.config.moe_token_dispatcher_type != "flex":
+            return None
+        if self.config.moe_expert_capacity_factor is not None:
+            return None
+        if not fused_topk_with_score_function_supports_topk_indices:
+            return None
+
+        backend = self.config.moe_flex_dispatcher_backend
+        if backend in ("deepep", "ncclep"):
+            return torch.int64
+        if backend != "hybridep":
+            return None
+        if self.config.moe_hybridep_routing_map_mode != "indices":
+            return None
+        if not HAVE_HYBRIDEP_DENSE_ROUTING:
+            return None
+
+        num_experts = self.expt_tp_group.size() * self.config.num_moe_experts
+        if num_experts <= _HYBRIDEP_INT16_EXPERT_LIMIT:
+            return torch.int16
+        return None
+
     def _apply_aux_loss(
         self,
         probs: torch.Tensor,
@@ -743,23 +774,42 @@ def _apply_expert_bias(
         """
         if self.enable_expert_bias and torch.is_grad_enabled():
             with torch.no_grad():
+                use_dense_indices = routing_map.dtype != torch.bool
                 if padding_mask is not None:
-                    routing_map = routing_map & (~padding_mask).unsqueeze(-1)
-                self.local_tokens_per_expert += routing_map.sum(dim=0)
+                    flat_mask = padding_mask.reshape(-1)
+                    assert (
+                        flat_mask.shape[0] == routing_map.shape[0]
+                    ), f"padding_mask flat {flat_mask.shape} vs routing_map {routing_map.shape}"
+                    if use_dense_indices:
+                        routing_map = routing_map[~flat_mask]
+                    else:
+                        routing_map = routing_map & (~flat_mask).unsqueeze(-1)
+                if use_dense_indices:
+                    expert_indices = routing_map.reshape(-1).to(torch.long)
+                    token_counts = torch.ones_like(
+                        expert_indices, dtype=self.local_tokens_per_expert.dtype
+                    )
+                    if torch.are_deterministic_algorithms_enabled():
+                        self.local_tokens_per_expert.index_add_(0, expert_indices, token_counts)
+                    else:
+                        self.local_tokens_per_expert.scatter_add_(0, expert_indices, token_counts)
+                else:
+                    self.local_tokens_per_expert += routing_map.sum(dim=0)
 
     def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None):
         """Top-k routing function
 
         Args:
             logits (torch.Tensor): Logits tensor after gating.
-            padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens.
-                                                   Shape [seq_length, bsz]. True for valid tokens,
-                                                   False for padding tokens. Defaults to None.
+            padding_mask (torch.Tensor, optional): Boolean mask indicating padding positions.
+                                                   Shape [seq_length, bsz]. True = padding,
+                                                   False = valid. Defaults to None.
 
         Returns:
             probs (torch.Tensor): The probabilities of token to experts assignment.
             routing_map (torch.Tensor): The mapping of token to experts assignment,
-                with shape [num_tokens, num_experts].
+                with shape [num_tokens, num_experts], or dense top-k indices with shape
+                [num_tokens, topk] for supported Flex backends.
         """
         seq_length, bsz = logits.shape[:2]
         logits = logits.view(-1, self.config.num_moe_experts)
@@ -780,6 +830,14 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
             ), "Quantile balancing routing does not support padding masks yet."
             probs, routing_map = self.quantile_balancing(logits)
         else:
+            topk_indices_dtype = self._dense_route_indices_dtype()
+            topk_indices = (
+                torch.empty(
+                    (logits.shape[0], self.topk), dtype=topk_indices_dtype, device=logits.device
+                )
+                if topk_indices_dtype is not None
+                else None
+            )
             probs, routing_map = topk_routing_with_score_function(
                 logits,
                 self.topk,
@@ -791,8 +849,25 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N
                 expert_bias=self.expert_bias,
                 fused=self.config.moe_router_fusion,
                 router_replay=self.router_replay,
+                topk_indices=topk_indices,
             )
 
+        # Dropless HybridEP consumes routing metadata directly, so exclude padding rows before
+        # dispatch. Other dispatchers retain their existing fixed-route assumptions.
+        use_dropless_hybridep = (
+            self.config.moe_token_dispatcher_type == "flex"
+            and self.config.moe_flex_dispatcher_backend == "hybridep"
+            and self.config.moe_expert_capacity_factor is None
+            and self.config.moe_expert_rank_capacity_factor is None
+        )
+        if padding_mask is not None and use_dropless_hybridep:
+            valid_tokens = (~padding_mask).unsqueeze(-1)
+            probs = probs * valid_tokens
+            if routing_map.dtype == torch.bool:
+                routing_map = routing_map & valid_tokens
+            else:
+                routing_map = routing_map.masked_fill(padding_mask.unsqueeze(-1), -1)
+
         # Apply token dropping to probs and routing_map.
         if self.config.moe_expert_capacity_factor is not None:
             probs, routing_map = apply_router_token_dropping(