diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 51b63983007..3b7f81909eb 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -3752,10 +3752,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): diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index c302387df9e..5f3cba1bb2f 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/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,7 +463,6 @@ 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, @@ -432,6 +470,7 @@ def forward( 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 diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 91981b5afcc..2dcd6c9cf7c 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/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( diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 00b7dce63d1..be7a378ec7f 100644 --- a/megatron/core/transformer/moe/router.py +++ b/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,9 +774,27 @@ 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 @@ -759,7 +808,8 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N 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( diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index ecb45cf7960..7075fd5880c 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -21,6 +21,7 @@ ) from megatron.core.transformer.enums import CudaGraphModule from megatron.core.transformer.moe.fused_a2a import ( + HAVE_HYBRIDEP_DENSE_ROUTING, HYBRIDEP_TOKEN_ALIGNMENT, alloc_ep_symm_buffer, ensure_nccl_ep_bootstrapped, @@ -60,6 +61,8 @@ logger = logging.getLogger(__name__) +_HYBRIDEP_INT16_EXPERT_LIMIT = 1 << 15 + class MoETokenDispatcher: """ @@ -964,9 +967,8 @@ class _DispatchManager(ABC): """ A manager class to handle dispatch and combine processes for MoE models. - DispatcherManager handles token dispatching according to the routing_map of format - [num_local_tokens, world_size, num_instances]. The routing_map is a 3D tensor where each - element indicates whether a token should be sent to a specific rank. + DispatcherManager handles token dispatching from either a bool routing map of shape + [num_local_tokens, world_size, num_instances] or dense top-k expert indices. num_instances is the maximum number of tokens instances dispatched into a target rank, it can be the number of local experts, or the size of sub_group. @@ -1019,6 +1021,7 @@ def __init__( num_local_experts: int, num_experts: int, config: TransformerConfig, + router_topk: Optional[int] = None, ): """ Initialize the HybridEP dispatcher. @@ -1029,11 +1032,13 @@ def __init__( num_local_experts (int): The number of local experts. num_experts (int): The total number of experts in the group. config (TransformerConfig): The configuration for the transformer model. + router_topk (int, optional): The top-k width after expert-TP expansion. """ self.group = group self.num_local_experts = num_local_experts self.num_experts = num_experts self.config = config + self.router_topk = router_topk if router_topk is not None else config.moe_router_topk self.permute_fusion = config.moe_permute_fusion self.capacity_factor = config.moe_expert_capacity_factor # Drop and pad the input to capacity. @@ -1083,18 +1088,62 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT self._padded_num_tokens = padded_num_tokens - routing_map = routing_map.reshape(num_tokens, self.num_experts) probs = probs.reshape(num_tokens, self.num_experts) + provided_topk_idx = None + + if routing_map.dtype == torch.bool: + routing_map = routing_map.reshape(num_tokens, self.num_experts) + if padded_num_tokens > num_tokens: + pad_rows = padded_num_tokens - num_tokens + routing_map = torch.cat( + [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0 + ) + self.routing_map = routing_map + else: + if not HAVE_HYBRIDEP_DENSE_ROUTING: + raise RuntimeError( + "HybridEP dense routing map was provided, but the installed HybridEPBuffer " + "does not support dense topk_idx metadata. Use a newer HybridEP backend or " + "disable dense routing." + ) + self.routing_map = None + provided_topk_idx = routing_map.reshape(num_tokens, self.router_topk).contiguous() + if padded_num_tokens > num_tokens: + pad_rows = padded_num_tokens - num_tokens + provided_topk_idx = torch.cat( + [ + provided_topk_idx, + provided_topk_idx.new_full((pad_rows, self.router_topk), -1), + ], + dim=0, + ) + if padded_num_tokens > num_tokens: pad_rows = padded_num_tokens - num_tokens - routing_map = torch.cat( - [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0 - ) probs = torch.cat([probs, probs.new_zeros((pad_rows, self.num_experts))], dim=0) - self.routing_map = routing_map self.token_probs = probs + if provided_topk_idx is not None: + if self.num_experts > _HYBRIDEP_INT16_EXPERT_LIMIT: + raise RuntimeError( + "HybridEP dense routing requires int16 expert ids, but the expert-TP-expanded " + f"expert count is {self.num_experts}; the maximum is " + f"{_HYBRIDEP_INT16_EXPERT_LIMIT}." + ) + self.topk_idx = provided_topk_idx.to(torch.int16) + elif ( + HAVE_HYBRIDEP_DENSE_ROUTING + and self.config.moe_hybridep_routing_map_mode == "indices" + and self.num_experts <= _HYBRIDEP_INT16_EXPERT_LIMIT + ): + _, self.topk_idx = torch.topk(self.token_probs, self.router_topk, dim=-1) + self.topk_idx = self.topk_idx.to(torch.int16) + invalid_routes = ~self.routing_map.gather(1, self.topk_idx.long()) + self.topk_idx = self.topk_idx.masked_fill(invalid_routes, -1) + else: + self.topk_idx = None + if self.moe_expert_rank_capacity_factor is not None: pad_multiple = get_align_size_for_quantization(self.config) # Static upper bound on permuted tokens passed to HybridEP (dropless EP rank @@ -1162,6 +1211,8 @@ def dispatch( pad_multiple=self.pad_multiple, fused=self.config.moe_permute_fusion_into_hybridep, num_sms_preprocessing_api=self.config.moe_hybridep_num_sms_preprocessing, + topk_idx=self.topk_idx, + num_of_experts=self.num_experts, ) ) if self.moe_expert_rank_capacity_factor is not None: @@ -1294,10 +1345,16 @@ def __init__( def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): num_tokens = routing_map.shape[0] - routing_map = routing_map.reshape(num_tokens, self.num_experts) probs = probs.reshape(num_tokens, self.num_experts) - # Convert the format of routing map from multihot to indices. - self.token_probs, self.token_indices = torch.topk(probs, self.router_topk, dim=-1) + if routing_map.dtype == torch.bool: + routing_map = routing_map.reshape(num_tokens, self.num_experts) + # Convert the format of routing map from multihot to indices. + self.token_probs, self.token_indices = torch.topk(probs, self.router_topk, dim=-1) + else: + self.token_indices = routing_map.reshape(num_tokens, self.router_topk).contiguous() + if self.token_indices.dtype != torch.int64: + self.token_indices = self.token_indices.to(torch.int64) + self.token_probs = probs.gather(1, self.token_indices) # Mask the indices of dropped tokens with -1 if self.capacity_factor is not None: mask = self.token_probs == 0 @@ -1610,8 +1667,15 @@ def __init__( def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor): num_tokens = routing_map.shape[0] probs = probs.reshape(num_tokens, self.num_experts) - # Convert the multihot routing map to (topk weights, topk indices), like DeepEP. - self.token_probs, self.token_indices = torch.topk(probs, self.router_topk, dim=-1) + if routing_map.dtype == torch.bool: + # Convert the multihot routing map to (topk weights, topk indices). + self.token_probs, self.token_indices = torch.topk(probs, self.router_topk, dim=-1) + else: + # Consume TE's direct top-k output without reconstructing it from a sparse map. + self.token_indices = routing_map.reshape(num_tokens, self.router_topk).contiguous() + if self.token_indices.dtype != torch.int64: + self.token_indices = self.token_indices.to(torch.int64) + self.token_probs = probs.gather(1, self.token_indices) self.num_local_tokens = num_tokens def _ensure_bootstrap(self): @@ -1848,8 +1912,13 @@ def __init__( num_local_experts=self.num_local_experts, num_experts=self.tp_size * self.config.num_moe_experts, config=self.config, + router_topk=self.tp_size * self.config.moe_router_topk, ) - self.cudagraph_attrs = ['_comm_manager.token_probs', '_comm_manager.routing_map'] + self.cudagraph_attrs = [ + '_comm_manager.token_probs', + '_comm_manager.routing_map', + '_comm_manager.topk_idx', + ] elif self.config.moe_flex_dispatcher_backend == "ncclep": assert self.tp_size * self.ep_size > 1, "NCCL EP dispatcher requires TPxEP > 1" self._comm_manager = _NCCLEPManager( @@ -1896,20 +1965,33 @@ def _initialize_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor) - This design decouples the communication group from underlying model parallelism groups, such that the communication strategy of tokens can be agnostic of TP size and EP size. - This function expands the routing_map from shape [num_local_tokens, num_experts] to - [num_local_tokens, world_size, num_local_experts]. Each element in the routing_map - indicates whether a token should be sent to a specific rank. Specifically, the - routing_map is replicated across TP group since each TP ranks in a TP group should - receive the same tokens. + Bool routing maps are expanded from [num_local_tokens, num_experts] to + [num_local_tokens, world_size, num_local_experts]. Dense top-k indices are expanded + from [num_local_tokens, topk] to [num_local_tokens, topk * expert_tp_size]. """ num_local_tokens = routing_map.shape[0] world_size = self.tp_size * self.ep_size - # Organize routing map and probs to [num_local_tokens, world_size, num_local_experts] - routing_map = ( - routing_map.reshape(num_local_tokens, self.ep_size, 1, self.num_local_experts) - .expand(-1, -1, self.tp_size, -1) - .reshape(num_local_tokens, world_size, self.num_local_experts) - ).contiguous() + if routing_map.dtype == torch.bool: + routing_map = ( + routing_map.reshape(num_local_tokens, self.ep_size, 1, self.num_local_experts) + .expand(-1, -1, self.tp_size, -1) + .reshape(num_local_tokens, world_size, self.num_local_experts) + ).contiguous() + else: + topk_indices = routing_map.long() + invalid_routes = topk_indices < 0 + expert_parallel_idx = topk_indices // self.num_local_experts + local_expert_idx = topk_indices % self.num_local_experts + tensor_parallel_idx = torch.arange( + self.tp_size, device=routing_map.device, dtype=topk_indices.dtype + ).view(1, 1, self.tp_size) + expanded_indices = ( + expert_parallel_idx.unsqueeze(-1) * self.tp_size + tensor_parallel_idx + ) * self.num_local_experts + local_expert_idx.unsqueeze(-1) + expanded_indices = expanded_indices.masked_fill(invalid_routes.unsqueeze(-1), -1) + routing_map = ( + expanded_indices.reshape(num_local_tokens, -1).to(routing_map.dtype).contiguous() + ) probs = ( probs.reshape(num_local_tokens, self.ep_size, 1, self.num_local_experts) .expand(-1, -1, self.tp_size, -1) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index f3fda5c2c59..0751b4d4451 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1024,6 +1024,12 @@ class TransformerConfig(ModelParallelConfig): moe_hybridep_num_sms_preprocessing: int = 108 """Number of SMs to use for HybridEP preprocessing (metadata scan kernel).""" + moe_hybridep_routing_map_mode: Literal['indices', 'bool'] = 'indices' + """Routing-map format for HybridEP. ``indices`` requests int16 top-k indices and is the + default, while ``bool`` forces the bool token-to-expert map. Index routing remains gated on + Transformer Engine and HybridEP support and the int16 expert limit; unsupported configurations + fall back to the bool routing-map path.""" + moe_ncclep_zero_copy: bool = False """For the 'ncclep' flex dispatcher: use the NCCL symmetric-memory zero-copy IO path (ep_bootstrap zero_copy + symm-mem-backed receive/combine buffers) instead of the default HBM @@ -1919,6 +1925,9 @@ def __post_init__(self): "moe_pad_expert_input_to_capacity" ) + if self.moe_hybridep_routing_map_mode not in ("indices", "bool"): + raise ValueError("moe_hybridep_routing_map_mode must be one of 'indices' or 'bool'.") + if self.moe_flex_dispatcher_backend == "ncclep": if self.moe_token_dispatcher_type != "flex": raise ValueError( diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index e0fe107302c..61031b1c747 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -2319,6 +2319,7 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('moe_router_score_function', force=True) _set_arg('moe_router_enable_expert_bias', force=True) _set_arg('moe_router_topk_scaling_factor', force=True) + _set_arg('moe_hybridep_routing_map_mode', force=False) # Mamba args. _set_arg('mamba_state_dim', force=True) diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index d5c625ca8ca..4e98ef78f90 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -76,6 +76,35 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() +@pytest.mark.parametrize( + ("runtime_mode", "checkpoint_args", "expected_mode"), + [ + ("indices", SimpleNamespace(moe_hybridep_routing_map_mode="bool"), "indices"), + (None, SimpleNamespace(moe_hybridep_routing_map_mode="bool"), "bool"), + (None, SimpleNamespace(), None), + ], +) +def test_load_args_preserves_runtime_hybridep_routing_map_mode( + runtime_mode, checkpoint_args, expected_mode +): + args = SimpleNamespace( + load="checkpoint", + iteration=0, + moe_hybridep_routing_map_mode=runtime_mode, + use_tokenizer_model_from_checkpoint_args=False, + use_mp_args_from_checkpoint_args=False, + ) + state_dict = {"args": checkpoint_args, "iteration": 12} + + with mock.patch( + "megatron.training.checkpointing._load_base_checkpoint", + return_value=(state_dict, "checkpoint", False, CheckpointType.LEGACY), + ): + restored_args, _ = load_args_from_checkpoint(args) + + assert restored_args.moe_hybridep_routing_map_mode == expected_mode + + def test_maybe_save_dataloader_state_uses_explicit_process_groups(tmp_path): """Dataloader checkpoints use the supplied module groups and canonical model-parallel path.""" groups = { diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index d6c8fcbd919..6a6a9915c23 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -1,12 +1,14 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. - import dataclasses +from types import SimpleNamespace from typing import cast import pytest import torch +import megatron.core.transformer.moe.moe_utils as moe_utils +import megatron.core.transformer.moe.router as router_module from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules from megatron.core.transformer.moe.moe_utils import ( @@ -14,7 +16,7 @@ router_gating_linear, topk_routing_with_score_function, ) -from megatron.core.transformer.moe.router import Router +from megatron.core.transformer.moe.router import Router, TopKRouter from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.initialize import _set_random_seed @@ -30,6 +32,70 @@ except Exception: # pragma: no cover - defensive HAVE_ROUTER_FUSION = False +HAVE_DENSE_ROUTER_FUSION = ( + HAVE_ROUTER_FUSION and moe_utils.fused_topk_with_score_function_supports_topk_indices +) + + +@pytest.mark.parametrize( + "backend,routing_map_mode,num_experts,capacity_factor,expected_dtype", + [ + ("deepep", "bool", 8, None, torch.int64), + ("ncclep", "bool", 8, None, torch.int64), + ("hybridep", "bool", 8, None, None), + ("hybridep", "indices", 1 << 15, None, torch.int16), + ("hybridep", "indices", (1 << 15) + 1, None, None), + ("hybridep", "indices", 8, 1.0, None), + ], +) +def test_dense_route_indices_dtype( + monkeypatch, backend, routing_map_mode, num_experts, capacity_factor, expected_dtype +): + monkeypatch.setattr(router_module, "fused_topk_with_score_function_supports_topk_indices", True) + monkeypatch.setattr(router_module, "HAVE_HYBRIDEP_DENSE_ROUTING", True) + router = SimpleNamespace( + config=SimpleNamespace( + moe_router_fusion=True, + moe_token_dispatcher_type="flex", + moe_expert_capacity_factor=capacity_factor, + moe_flex_dispatcher_backend=backend, + moe_hybridep_routing_map_mode=routing_map_mode, + num_moe_experts=num_experts, + ), + expt_tp_group=SimpleNamespace(size=lambda: 1), + ) + + assert TopKRouter._dense_route_indices_dtype(router) == expected_dtype + + +@pytest.mark.parametrize("supports_topk_indices", [False, True]) +def test_fused_router_only_forwards_supported_topk_indices(monkeypatch, supports_topk_indices): + received_kwargs = {} + + def fake_fused_router(**kwargs): + received_kwargs.update(kwargs) + return torch.zeros_like(kwargs["logits"]), kwargs.get( + "topk_indices", torch.zeros_like(kwargs["logits"], dtype=torch.bool) + ) + + monkeypatch.setattr(moe_utils, "HAVE_TE", True) + monkeypatch.setattr(moe_utils, "fused_topk_with_score_function", fake_fused_router) + monkeypatch.setattr( + moe_utils, "fused_topk_with_score_function_supports_topk_indices", supports_topk_indices + ) + logits = torch.randn(4, 8) + topk_indices = torch.empty(4, 2, dtype=torch.int64) + + topk_routing_with_score_function(logits, 2, fused=True, topk_indices=topk_indices) + + assert ("topk_indices" in received_kwargs) is supports_topk_indices + if supports_topk_indices: + assert received_kwargs["topk_indices"] is topk_indices + + received_kwargs.clear() + topk_routing_with_score_function(logits, 2, fused=True) + assert "topk_indices" not in received_kwargs + class TestTop2Router: def setup_method(self, method): @@ -150,9 +216,16 @@ def test_aux_loss(self): @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_router_with_padding_mask(self): - """Test that padding mask correctly excludes padding tokens from routing.""" + @pytest.mark.parametrize("router_fusion", [False, True]) + def test_router_with_padding_mask(self, router_fusion): + """Test that HybridEP excludes padding tokens from routing.""" + if router_fusion and not HAVE_ROUTER_FUSION: + pytest.skip("TE fused router ops not available") self.router = self.router.cuda() + self.router.config.moe_router_fusion = router_fusion + self.router.config.moe_token_dispatcher_type = "flex" + self.router.config.moe_flex_dispatcher_backend = "hybridep" + self.router.config.moe_hybridep_routing_map_mode = "bool" seq_len = 32 batch_size = 2 hidden_size = self.router.config.hidden_size @@ -192,9 +265,77 @@ def test_router_with_padding_mask(self): self.router.config.num_moe_experts, ) + padding_rows = padding_mask.reshape(-1) + assert torch.count_nonzero(probs_with_mask[padding_rows]) == 0 + assert not routing_map_with_mask[padding_rows].any() + # Verify that probs for valid tokens are similar assert torch.equal(probs_valid_part, probs_without_mask) + @pytest.mark.internal + @pytest.mark.skipif( + not torch.cuda.is_available() or not HAVE_DENSE_ROUTER_FUSION, + reason="TE dense fused router output is not available", + ) + def test_hybridep_dense_routing_masks_padding(self, monkeypatch): + monkeypatch.setattr(router_module, "HAVE_HYBRIDEP_DENSE_ROUTING", True) + self.router = self.router.cuda() + self.router.config.moe_router_fusion = True + self.router.config.moe_token_dispatcher_type = "flex" + self.router.config.moe_flex_dispatcher_backend = "hybridep" + self.router.config.moe_hybridep_routing_map_mode = "indices" + hidden_states = torch.randn( + (8, 2, self.router.config.hidden_size), device="cuda", dtype=torch.bfloat16 + ) + padding_mask = torch.zeros((8, 2), dtype=torch.bool, device="cuda") + padding_mask[4:, :] = True + + with torch.no_grad(): + probs, routing_map = self.router(hidden_states, padding_mask=padding_mask) + + padding_rows = padding_mask.reshape(-1) + assert routing_map.dtype == torch.int16 + assert routing_map.shape == (16, self.router.config.moe_router_topk) + assert torch.all(routing_map[padding_rows] == -1) + assert torch.count_nonzero(probs[padding_rows]) == 0 + + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize( + "dispatcher,backend,capacity_factor,rank_capacity_factor", + [ + ("allgather", "deepep", None, None), + ("alltoall", "deepep", None, None), + ("flex", "deepep", None, None), + ("flex", "deepepv2", None, None), + ("flex", "hybridep", 1.0, None), + ("flex", "hybridep", None, 1.0), + ], + ) + def test_padding_mask_preserves_routes_outside_dropless_hybridep( + self, dispatcher, backend, capacity_factor, rank_capacity_factor + ): + """Only dropless HybridEP may consume a sparse route map.""" + self.router = self.router.cuda() + self.router.config.moe_token_dispatcher_type = dispatcher + self.router.config.moe_flex_dispatcher_backend = backend + self.router.config.moe_expert_capacity_factor = capacity_factor + self.router.config.moe_expert_rank_capacity_factor = rank_capacity_factor + hidden_states = torch.randn( + (16, 2, self.router.config.hidden_size), device="cuda", dtype=torch.bfloat16 + ) + padding_mask = torch.zeros((16, 2), dtype=torch.bool, device="cuda") + padding_mask[8:, :] = True + + with torch.no_grad(): + probs_with_mask, routing_map_with_mask = self.router( + hidden_states, padding_mask=padding_mask + ) + probs_without_mask, routing_map_without_mask = self.router(hidden_states) + + torch.testing.assert_close(probs_with_mask, probs_without_mask) + assert torch.equal(routing_map_with_mask, routing_map_without_mask) + @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("with_padding_mask", [False, True]) @@ -535,6 +676,57 @@ def test_router_forward_aux_free(self): # Print some debug info print("Updated bias after first forward pass:", updated_bias) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("deterministic", [False, True]) + def test_dense_expert_bias_token_counts(self, deterministic): + self.router = self.router.cuda() + self.router.local_tokens_per_expert.zero_() + topk_indices = torch.tensor( + [[0, 3], [1, 4], [0, 7], [2, 5]], device="cuda", dtype=torch.int16 + ) + padding_mask = torch.tensor([False, True, False, False], device="cuda") + + previous_deterministic = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(deterministic) + try: + self.router._apply_expert_bias(topk_indices, padding_mask=padding_mask) + finally: + torch.use_deterministic_algorithms(previous_deterministic) + + expected = torch.tensor([2, 0, 1, 1, 0, 1, 0, 1], device="cuda", dtype=torch.float32) + torch.testing.assert_close(self.router.local_tokens_per_expert, expected) + + @pytest.mark.internal + @pytest.mark.skipif( + not torch.cuda.is_available() or not HAVE_DENSE_ROUTER_FUSION, + reason="TE dense fused router output is not available", + ) + def test_fused_dense_routing_with_expert_bias(self): + self.router = self.router.cuda() + self.router.config.moe_router_fusion = True + self.router.config.moe_token_dispatcher_type = "flex" + self.router.config.moe_flex_dispatcher_backend = "deepep" + self.router.local_tokens_per_expert.zero_() + self.router.expert_bias.copy_( + torch.arange(self.router.config.num_moe_experts, device="cuda", dtype=torch.float32) + ) + hidden_states = torch.randn( + (4, 2, self.router.config.hidden_size), device="cuda" + ).bfloat16() + padding_mask = torch.tensor( + [[False, True], [False, False], [True, False], [False, False]], device="cuda" + ) + + _, topk_indices = self.router(hidden_states, padding_mask=padding_mask) + + assert topk_indices.dtype == torch.int64 + assert topk_indices.shape == (8, self.router.config.moe_router_topk) + expected = torch.bincount( + topk_indices[~padding_mask.reshape(-1)].reshape(-1), + minlength=self.router.config.num_moe_experts, + ).to(torch.float32) + torch.testing.assert_close(self.router.local_tokens_per_expert, expected) + @pytest.mark.internal @pytest.mark.skipif( not torch.cuda.is_available() or not HAVE_ROUTER_FUSION, diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher.py b/tests/unit_tests/transformer/moe/test_token_dispatcher.py index 547e9488798..c0caa97f687 100644 --- a/tests/unit_tests/transformer/moe/test_token_dispatcher.py +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher.py @@ -1,21 +1,36 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. import dataclasses +from types import SimpleNamespace import pytest import torch import torch.nn.functional as F +import megatron.core.transformer.moe.fused_a2a as fused_a2a +import megatron.core.transformer.moe.token_dispatcher as token_dispatcher from megatron.core import config, parallel_state +from megatron.core.extensions.transformer_engine import ( + fused_topk_with_score_function_supports_topk_indices, +) from megatron.core.fp8_utils import get_fp8_context from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_submodules, get_gpt_layer_with_transformer_engine_spec, ) -from megatron.core.transformer.moe.fused_a2a import HYBRIDEP_TOKEN_ALIGNMENT, reset_hybrid_ep_buffer +from megatron.core.transformer.moe.fused_a2a import ( + HAVE_HYBRIDEP_DENSE_ROUTING, + HYBRIDEP_TOKEN_ALIGNMENT, + reset_hybrid_ep_buffer, +) from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules from megatron.core.transformer.moe.moe_utils import get_capacity -from megatron.core.transformer.moe.token_dispatcher import _HybridEPManager +from megatron.core.transformer.moe.token_dispatcher import ( + MoEFlexTokenDispatcher, + _DeepepManager, + _HybridEPManager, + _NCCLEPManager, +) from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module @@ -99,7 +114,9 @@ def __init__( sequence_parallel=tp_size > 1, add_bias_linear=kwargs.get("add_bias_linear", False), moe_permute_fusion=kwargs.get("moe_permute_fusion", False), + moe_router_fusion=kwargs.get("moe_router_fusion", False), moe_flex_dispatcher_backend=kwargs.get("moe_flex_dispatcher_backend", None), + moe_hybridep_routing_map_mode=kwargs.get("moe_hybridep_routing_map_mode", "bool"), moe_expert_rank_capacity_factor=kwargs.get("moe_expert_rank_capacity_factor", None), moe_ncclep_zero_copy=kwargs.get("moe_ncclep_zero_copy", False), moe_dispatch_fwd_dtype=kwargs.get("moe_dispatch_fwd_dtype", 'bf16'), @@ -560,6 +577,7 @@ def test_hybridep_pad_uneven_dispatch_inputs_metadata(monkeypatch): manager.group = object() manager.num_local_experts = 2 manager.num_experts = 4 + manager.router_topk = 2 manager.config = TransformerConfig( num_layers=1, hidden_size=16, @@ -598,6 +616,112 @@ def fake_all_reduce(tensor, op=None, group=None): assert not manager.token_probs[local_num_tokens:].any() +def test_hybridep_sparse_fallback_marks_empty_routes_invalid(monkeypatch): + monkeypatch.setattr(token_dispatcher, "HAVE_HYBRIDEP_DENSE_ROUTING", True) + manager = object.__new__(_HybridEPManager) + manager.config = SimpleNamespace( + moe_hybridep_pad_uneven_dispatch_inputs=False, moe_hybridep_routing_map_mode="indices" + ) + manager.group = object() + manager.num_experts = 2 + manager.router_topk = 1 + manager.moe_expert_rank_capacity_factor = None + manager.drop_and_pad = False + + routing_map = torch.tensor([[True, False], [False, False]]) + probs = torch.tensor([[1.0, 0.0], [0.0, 0.0]]) + + manager.setup_metadata(routing_map, probs) + + assert torch.equal(manager.topk_idx, torch.tensor([[0], [-1]], dtype=torch.int16)) + + +def test_hybridep_dense_input_requires_backend_support(monkeypatch): + monkeypatch.setattr(token_dispatcher, "HAVE_HYBRIDEP_DENSE_ROUTING", False) + manager = object.__new__(_HybridEPManager) + manager.config = SimpleNamespace( + moe_hybridep_pad_uneven_dispatch_inputs=False, moe_hybridep_routing_map_mode="indices" + ) + manager.group = object() + manager.num_experts = 4 + manager.router_topk = 2 + manager.moe_expert_rank_capacity_factor = None + manager.drop_and_pad = False + + with pytest.raises(RuntimeError, match="does not support dense topk_idx metadata"): + manager.setup_metadata(torch.tensor([[0, 2]], dtype=torch.int16), torch.ones(1, 4)) + + +def test_flex_dense_metadata_preserves_invalid_routes(): + dispatcher = object.__new__(MoEFlexTokenDispatcher) + dispatcher.tp_size = 2 + dispatcher.ep_size = 2 + dispatcher.num_local_experts = 2 + routing_map = torch.tensor([[0, -1], [3, 1]], dtype=torch.int16) + probs = torch.ones((2, 4)) + + expanded_routes, _ = dispatcher._initialize_metadata(routing_map, probs) + + expected = torch.tensor([[0, 2, -1, -1], [5, 7, 1, 3]], dtype=torch.int16) + assert torch.equal(expanded_routes, expected) + + +@pytest.mark.parametrize("manager_cls", [_DeepepManager, _NCCLEPManager]) +def test_dense_required_manager_accepts_dense_indices(monkeypatch, manager_cls): + manager = object.__new__(manager_cls) + manager.num_experts = 4 + manager.router_topk = 2 + if isinstance(manager, _DeepepManager): + manager.capacity_factor = None + dense_indices = torch.tensor([[0, 2], [3, 1]], dtype=torch.int16) + probs = torch.tensor([[0.6, 0.0, 0.4, 0.0], [0.0, 0.3, 0.0, 0.7]]) + + monkeypatch.setattr( + torch, "topk", lambda *args, **kwargs: pytest.fail("dense routing must not call torch.topk") + ) + manager.setup_metadata(dense_indices, probs) + + assert manager.token_indices.dtype == torch.int64 + assert torch.equal(manager.token_indices, dense_indices.long()) + torch.testing.assert_close(manager.token_probs, torch.tensor([[0.6, 0.4], [0.7, 0.3]])) + + +@pytest.mark.parametrize("explicit_dense_routing", [False, True]) +def test_hybridep_dispatch_supports_inferred_and_explicit_dense_apis( + monkeypatch, explicit_dense_routing +): + class FakeHybridEPBuffer: + def __init__(self): + self.kwargs = None + + def dispatch_with_permute(self, **kwargs): + self.kwargs = kwargs + return ( + kwargs["hidden"], + kwargs["probs"], + None, + torch.ones(2, dtype=torch.int32), + ("handle",), + ) + + fake_buffer = FakeHybridEPBuffer() + monkeypatch.setattr(fused_a2a, "_hybrid_ep_buffer", fake_buffer) + monkeypatch.setattr(fused_a2a, "HAVE_HYBRIDEP_DENSE_ROUTING", True) + monkeypatch.setattr(fused_a2a, "HAVE_HYBRIDEP_EXPLICIT_DENSE_ROUTING", explicit_dense_routing) + hidden = torch.randn(2, 4) + probs = torch.randn(2, 4) + topk_idx = torch.tensor([[0, 1], [2, 3]], dtype=torch.int16) + + fused_a2a.HybridEPDispatch.forward( + SimpleNamespace(), hidden, None, probs, object(), 2, topk_idx=topk_idx, num_of_experts=4 + ) + + assert fake_buffer.kwargs["topk_idx"] is topk_idx + assert fake_buffer.kwargs["num_of_experts"] == 4 + assert "routing_map" not in fake_buffer.kwargs + assert ("dense_routing" in fake_buffer.kwargs) is explicit_dense_routing + + @pytest.mark.skipif( not is_deep_ep_available() and not is_hybrid_ep_available(), reason="Deep EP and Hybrid EP are not available", @@ -610,6 +734,60 @@ def teardown_method(self, method): reset_hybrid_ep_buffer() Utils.destroy_model_parallel() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.internal + @pytest.mark.parametrize("routing_map_mode", ["bool", "indices"]) + def test_hybridep_routing_map_modes_forward_backward(self, routing_map_mode): + if not is_hybrid_ep_available(): + pytest.skip("Hybrid EP is not available") + if routing_map_mode == "indices" and ( + not fused_topk_with_score_function_supports_topk_indices + or not HAVE_HYBRIDEP_DENSE_ROUTING + ): + pytest.skip("Dense TE/HybridEP routing APIs are not available") + + container = MoEModelTestContainer( + tp_size=1, + ep_size=8, + pp_size=1, + num_moe_experts=8, + moe_router_topk=2, + moe_router_load_balancing_type="aux_loss", + moe_token_dispatcher_type="flex", + moe_router_fusion=True, + moe_flex_dispatcher_backend="hybridep", + moe_hybridep_routing_map_mode=routing_map_mode, + hidden_size=1024, + test_dtype=torch.bfloat16, + ) + container.dispatcher_dropless_test() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.internal + @pytest.mark.parametrize("backend", ["deepep", "ncclep"]) + def test_dense_required_backend_forward_backward(self, backend): + if backend == "deepep" and not is_deep_ep_available(): + pytest.skip("Deep EP is not available") + if backend == "ncclep" and not is_nccl_ep_available(): + pytest.skip("NCCL EP is not available") + if not fused_topk_with_score_function_supports_topk_indices: + pytest.skip("Dense TE routing output is not available") + + container = MoEModelTestContainer( + tp_size=1, + ep_size=8, + pp_size=1, + num_moe_experts=8, + moe_router_topk=2, + moe_router_load_balancing_type="aux_loss", + moe_token_dispatcher_type="flex", + moe_router_fusion=True, + moe_flex_dispatcher_backend=backend, + hidden_size=1024, + test_dtype=torch.bfloat16, + ) + container.dispatcher_dropless_test() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.internal @pytest.mark.parametrize("tp_size,ep_size", [(1, 8), (8, 1), (4, 2)])