From 5b7ddd4006c07386700c83dedc0dd49afa9727b6 Mon Sep 17 00:00:00 2001 From: Seonjin Na Date: Fri, 26 Jun 2026 14:32:28 -0700 Subject: [PATCH 1/3] Handle HybridEP packed padding masks Signed-off-by: Seonjin Na --- megatron/core/transformer/moe/moe_layer.py | 19 ++++++++++++++++--- megatron/core/transformer/moe/router.py | 17 +++++++++++++---- 2 files changed, 29 insertions(+), 7 deletions(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index deebd3472ea..b5826f7a336 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -611,9 +611,9 @@ def forward( Args: hidden_states (torch.Tensor): The input tensor shape [seq_length, bsz, hidden_size]. - 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 tokens. + Shape [seq_length, bsz]. True for padding tokens, + False for valid tokens. Defaults to None. Returns: A tuple containing the output tensor and the MLP bias, if any. """ @@ -637,6 +637,19 @@ def forward( # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states if padding_mask is not None: padding_mask = padding_mask.transpose(0, 1).bool() + if ( + self.config.sequence_parallel + and padding_mask.shape[0] != hidden_states.shape[0] + ): + padding_mask = tensor_parallel.scatter_to_sequence_parallel_region( + padding_mask, + group=self.tp_group, + ) + if padding_mask.shape[:2] != hidden_states.shape[:2]: + raise RuntimeError( + f"padding_mask shape {padding_mask.shape} must match hidden_states " + f"shape {hidden_states.shape[:2]} before MoE routing" + ) # MoE forward: route -> dispatch -> compute -> combine def custom_forward(hidden_states, intermediate_tensors=None, padding_mask=None): diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 03317b65f1c..ba1a4bc1074 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -609,7 +609,7 @@ def _apply_expert_bias( if self.enable_expert_bias and torch.is_grad_enabled(): with torch.no_grad(): if padding_mask is not None: - routing_map = routing_map & (~padding_mask) + routing_map = routing_map & (~padding_mask).unsqueeze(-1) self.local_tokens_per_expert += routing_map.sum(dim=0) def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): @@ -617,9 +617,9 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N 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 tokens. + Shape [seq_length, bsz]. True for padding tokens, + False for valid tokens. Defaults to None. Returns: probs (torch.Tensor): The probabilities of token to experts assignment. @@ -653,6 +653,15 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N router_replay=self.router_replay, ) + if ( + padding_mask is not None + and self.config.moe_token_dispatcher_type == "flex" + and self.config.moe_flex_dispatcher_backend == "hybridep" + ): + valid_token_mask = (~padding_mask).unsqueeze(-1) + routing_map = routing_map & valid_token_mask + probs = probs * valid_token_mask.to(dtype=probs.dtype) + # 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( From 4d04e7625c5e84f984a9f01aef58cb006b0aa7ac Mon Sep 17 00:00:00 2001 From: Seonjin Na Date: Fri, 26 Jun 2026 15:32:04 -0700 Subject: [PATCH 2/3] Test HybridEP packed padding router masks Signed-off-by: Seonjin Na --- megatron/core/transformer/moe/moe_layer.py | 8 ++--- .../transformer/moe/test_routers.py | 36 +++++++++++++++++++ 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index b5826f7a336..ee6e2d92db6 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -637,13 +637,9 @@ def forward( # Transpose from [bsz, seq_length] to [seq_length, bsz] to align with hidden_states if padding_mask is not None: padding_mask = padding_mask.transpose(0, 1).bool() - if ( - self.config.sequence_parallel - and padding_mask.shape[0] != hidden_states.shape[0] - ): + if self.config.sequence_parallel and padding_mask.shape[0] != hidden_states.shape[0]: padding_mask = tensor_parallel.scatter_to_sequence_parallel_region( - padding_mask, - group=self.tp_group, + padding_mask, group=self.tp_group ) if padding_mask.shape[:2] != hidden_states.shape[:2]: raise RuntimeError( diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index 9f33dd01920..5413b8f0f72 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -175,6 +175,42 @@ def test_router_with_padding_mask(self): # 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(), reason="CUDA not available") + def test_hybridep_router_masks_padding_tokens_from_dispatch(self): + """Test that HybridEP dispatch metadata excludes padding tokens.""" + self.router = self.router.cuda() + self.router.config.moe_token_dispatcher_type = "flex" + self.router.config.moe_flex_dispatcher_backend = "hybridep" + seq_len = 32 + batch_size = 2 + hidden_size = self.router.config.hidden_size + + hidden_states = torch.randn((seq_len, batch_size, hidden_size)).cuda().bfloat16() + + padding_mask = torch.zeros((seq_len, batch_size), dtype=torch.bool, device='cuda') + padding_mask[seq_len // 2 :, :] = True + + with torch.no_grad(): + probs, routing_map = self.router(hidden_states, padding_mask=padding_mask) + probs_without_mask, routing_map_without_mask = self.router( + hidden_states[: seq_len // 2, :, :] + ) + + probs = probs.reshape(seq_len, batch_size, -1) + routing_map = routing_map.reshape(seq_len, batch_size, -1) + + assert torch.count_nonzero(probs[seq_len // 2 :, :, :]) == 0 + assert not routing_map[seq_len // 2 :, :, :].any() + assert (routing_map[: seq_len // 2, :, :].sum(dim=-1) == self.router.topk).all() + assert torch.equal( + probs[: seq_len // 2, :, :].reshape(-1, probs.shape[-1]), probs_without_mask + ) + assert torch.equal( + routing_map[: seq_len // 2, :, :].reshape(-1, routing_map.shape[-1]), + routing_map_without_mask, + ) + @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_router_dtype(self): From 278cc9128c233a38ea9fa8ac7cf9de22e434efa6 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 22 Aug 2026 01:06:08 -0700 Subject: [PATCH 3/3] Exclude padding tokens from HybridEP routing Signed-off-by: seonjinn --- megatron/core/transformer/moe/router.py | 14 +++++ .../transformer/moe/test_routers.py | 51 ++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index c3f99f7f999..a2dada0c026 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -789,6 +789,20 @@ def routing(self, logits: torch.Tensor, padding_mask: Optional[torch.Tensor] = N router_replay=self.router_replay, ) + # Dropless HybridEP consumes the sparse routing map directly, so exclude padding + # rows before dispatch. Other dispatchers retain their existing fixed-route + # assumptions until they support sparse routing maps end to end. + 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 + routing_map = routing_map & valid_tokens + # 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/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index 6a2ab7908aa..8fc63194c2c 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -135,9 +135,15 @@ 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" seq_len = 32 batch_size = 2 hidden_size = self.router.config.hidden_size @@ -177,9 +183,50 @@ 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(), 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])