fix(moe): exclude padding tokens from dropless HybridEP routing

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(

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])