Support async scheduling for MoE and expert parallelism

megatron/core/inference/engines/dynamic_engine.py

@@ -1010,10 +1010,6 @@ def _validate_async_sched_support_for_config(self) -> None:
             raise ValueError("Async scheduling does not support prefix caching.")
         if not self.materialize_only_last_token_logits:
             raise ValueError("Async scheduling requires materialize_only_last_token_logits=True.")
-        if model_config.expert_model_parallel_size > 1:
-            raise ValueError("Async scheduling does not support expert parallelism.")
-        if model_config.num_moe_experts is not None:
-            raise ValueError("Async scheduling does not support MoE models.")
         if model_config.moe_enable_routing_replay:
             raise ValueError("Async scheduling does not support routing replay.")
 
@@ -1945,6 +1941,102 @@ def schedule_chunked_prefill(self):
             else:
                 self.waiting_request_ids.extendleft(reversed(pending_request_ids))
 
+    async def _ep_establish_consensus(
+        self, local_work: int, signal_consensus: bool
+    ) -> tuple[int, bool]:
+        """EP all-reduce to share work counts and pause consensus.
+
+        All-reduces two integers at once:
+        - local_work: actual pending request count (always >= 0).
+        - consensus flag: -1 if this rank wants to pause, 0 otherwise.
+
+        Using max for both:
+        - max(work) > 0 means at least one EP peer has real work.
+        - max(consensus) == -1 means ALL peers signaled -1 (all PAUSING).
+          Any RUNNING peer contributes 0, pulling the max to 0.
+
+        Args:
+            local_work: Pending request count for this rank.
+            signal_consensus: True if this rank is ready to pause.
+
+        Returns:
+            (global_work, all_pausing): max work across EP, and whether
+                all peers signaled consensus.
+        """
+        nvtx_range_push("_ep_establish_consensus")
+
+        consensus_val = -1 if signal_consensus else 0
+
+        # Signals can be received asynchronously on EP ranks.
+        # We do not want a rank to pause prematurely if its peers have yet to receive the signal.
+        # So this is an *attempt* to process the signal. This rank has received the signal
+        # and passes -1 to the all-reduce. If any other rank in the EP group has not received
+        # the signal yet, it will pass a zero value to the all-reduce, hence the global consensus
+        # will be zero and we will defer processing the signal.
+        # When all ranks receive the signal, global consensus will be -1 and we can process.
+
+        if self.ep_world_size > 1:
+            # Note that it is important to use a non-blocking asyncio-friendly all-reduce here.
+            # The user may have other tasks running in the event loop that need to be serviced.
+            # Do not using a torch.distributed blocking all-reduce here using nccl/gloo.
+            # We have tried that and it blocks the event loop in megatron-rl.
+            global_work, global_consensus = (
+                await self.expert_parallel_zmq_communicator.all_reduce_max(
+                    local_work, consensus_val, async_op=(not self.use_synchronous_zmq_collectives)
+                )
+            )
+        else:
+            global_work, global_consensus = local_work, consensus_val
+
+        nvtx_range_pop("_ep_establish_consensus")
+        return global_work, global_consensus == -1
+
+    async def _advance_ep_consensus(
+        self, local_pending: int, signal_consensus: bool
+    ) -> tuple[int, bool]:
+        """Advance the cached EP consensus at its configured cadence.
+
+        Args:
+            local_pending (int): Local active and waiting request count.
+            signal_consensus (bool): Whether this rank is ready to pause.
+
+        Returns:
+            tuple[int, bool]: Cached global work and all-ranks-pausing state.
+        """
+        global_work_from_last_consensus, _ = self._last_ep_consensus
+        # Recheck immediately while idle and periodically while work remains.
+        if (
+            global_work_from_last_consensus == 0
+            or self._ep_consensus_loop_counter % self.ep_consensus_interval == 0
+        ):
+            self._last_ep_consensus = await self._ep_establish_consensus(
+                local_pending, signal_consensus=signal_consensus
+            )
+
+        self._ep_consensus_loop_counter += 1
+        return self._last_ep_consensus
+
+    async def _advance_async_sched_primer(self) -> None:
+        """Advance the existing EP cadence before consuming primer logits.
+
+        Primer consumption is atomic with respect to local request admission,
+        but it still occupies another EP forward slot. This method mirrors the
+        coordinator's existing consensus cadence without adding a primer-specific
+        collective or waiting for the primer forward.
+        """
+        if not self.use_coordinator or self.ep_world_size <= 1:
+            return
+
+        if self.disable_ep_consensus:
+            await asyncio.sleep(0)
+            return
+
+        local_pending = self.context.get_active_request_count() + len(self.waiting_request_ids)
+        # Keep the primer and its consumption atomic. A pending pause is
+        # reconsidered by the normal coordinator loop after this slot.
+        global_work, _ = await self._advance_ep_consensus(local_pending, signal_consensus=False)
+        assert global_work > 0, "An async-scheduling primer requires active EP work."
+
     async def async_forward(self) -> Tuple[Dict, Dict, float]:
         """Uses `asyncio` for continuous generation.
         Sleeps when no requests are available, until new requests have been added.
@@ -2008,6 +2100,7 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
             if not controller_result.primer_only:
                 result = controller_result.output
                 break
+            await self._advance_async_sched_primer()
 
         if admission_deferred:
             # Admit against the resolved batch, then leave its mixed forward pending.
@@ -2016,6 +2109,8 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
             assert (
                 primer_result.primer_only or primer_result.output is None
             ), "Async admission may only launch a forward primer."
+            if primer_result.primer_only:
+                await self._advance_async_sched_primer()
         if will_log_this_step:
             self.step_end_event.record()
             self.step_end_event.synchronize()
@@ -2606,55 +2701,6 @@ async def run_engine(self, *, loop: Optional[asyncio.AbstractEventLoop] = None):
         except asyncio.CancelledError:
             pass
 
-    async def _ep_establish_consensus(
-        self, local_work: int, signal_consensus: bool
-    ) -> tuple[int, bool]:
-        """EP all-reduce to share work counts and pause consensus.
-
-        All-reduces two integers at once:
-        - local_work: actual pending request count (always >= 0).
-        - consensus flag: -1 if this rank wants to pause, 0 otherwise.
-
-        Using max for both:
-        - max(work) > 0 means at least one EP peer has real work.
-        - max(consensus) == -1 means ALL peers signaled -1 (all PAUSING).
-          Any RUNNING peer contributes 0, pulling the max to 0.
-
-        Args:
-            local_work: Pending request count for this rank.
-            signal_consensus: True if this rank is ready to pause.
-        Returns:
-            (global_work, all_pausing): max work across EP, and whether
-            all peers signaled consensus.
-        """
-        nvtx_range_push("_ep_establish_consensus")
-
-        consensus_val = -1 if signal_consensus else 0
-
-        # Signals can be received asynchronously on EP ranks.
-        # We do not want a rank to pause prematurely if its peers have yet to receive the signal.
-        # So this is an *attempt* to process the signal. This rank has received the signal
-        # and passes -1 to the all-reduce. If any other rank in the EP group has not received
-        # the signal yet, it will pass a zero value to the all-reduce, hence the global consensus
-        # will be zero and we will defer processing the signal.
-        # When all ranks receive the signal, global consensus will be -1 and we can process.
-
-        if self.ep_world_size > 1:
-            # Note that it is important to use a non-blocking asyncio-friendly all-reduce here.
-            # The user may have other tasks running in the event loop that need to be serviced.
-            # Do not using a torch.distributed blocking all-reduce here using nccl/gloo.
-            # We have tried that and it blocks the event loop in megatron-rl.
-            global_work, global_consensus = (
-                await self.expert_parallel_zmq_communicator.all_reduce_max(
-                    local_work, consensus_val, async_op=(not self.use_synchronous_zmq_collectives)
-                )
-            )
-        else:
-            global_work, global_consensus = local_work, consensus_val
-
-        nvtx_range_pop("_ep_establish_consensus")
-        return global_work, global_consensus == -1
-
     async def _world_barrier(self):
         """World-wide ZMQ all-reduce barrier for global rank consensus.
 
@@ -2722,23 +2768,9 @@ async def run_engine_with_coordinator(
                             # delivery, request scheduling) run between steps.
                             await asyncio.sleep(0)
                         continue
-                    global_work_from_last_consensus, _ = self._last_ep_consensus
-                    if (
-                        global_work_from_last_consensus == 0
-                        or self._ep_consensus_loop_counter % self.ep_consensus_interval == 0
-                    ):
-                        # selectively enter ep_establish_consensus if
-                        # 1. there is no global work -> engine is idle. At any step in the future
-                        #    one of the ranks can receive work. So we should be eagerly checking for that
-                        # 2. it has been 20 steps since we last established consensus, and that consensus
-                        #    had some work.
-                        # In the worst case, this delays pausing by 20 steps which is around
-                        # 200-400 milliseconds.
-                        self._last_ep_consensus = await self._ep_establish_consensus(
-                            local_pending, signal_consensus=(self.state == EngineState.PAUSING)
-                        )
-                    global_work, all_pausing = self._last_ep_consensus
-                    self._ep_consensus_loop_counter += 1
+                    global_work, all_pausing = await self._advance_ep_consensus(
+                        local_pending, signal_consensus=(self.state == EngineState.PAUSING)
+                    )
 
                     if all_pausing:
                         # All EP peers are PAUSING: pause immediately.

tests/unit_tests/inference/engines/test_dynamic_engine.py

@@ -53,6 +53,7 @@
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
 from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope
+from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_fa_min_version, is_te_min_version
 from tests.unit_tests.test_utilities import Utils, clear_nvte_env_vars
@@ -246,6 +247,7 @@ def _build_requests(cls, test_config: DynamicEngineTestConfig) -> List[DynamicIn
                 termination_id=(
                     -1 if test_config.use_fixed_output_lengths else test_config.vocab_size - 1
                 ),
+                top_k=(1 if test_config.async_sched_mode != AsyncScheduleMode.LEGACY else 0),
                 return_log_probs=test_config.return_log_probs,
                 skip_prompt_log_probs=test_config.skip_prompt_log_probs,
             )
@@ -5008,6 +5010,7 @@ class TestDynamicInferenceEngineParallel(DynamicInferenceEngineTestBase):
     """
 
     def teardown_method(self, method):
+        NVLSAllGatherVDispatcher._delete_buffers()
         delete_cuda_graphs()
         Utils.destroy_model_parallel()
 
@@ -5081,6 +5084,31 @@ def test_parallel_inference(
             transformer_impl=transformer_impl,
         )
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
+    @pytest.mark.parametrize(
+        "async_sched_mode", [AsyncScheduleMode.SERIAL, AsyncScheduleMode.OVERLAP]
+    )
+    @pytest.mark.parametrize("inference_moe_token_dispatcher_type", ["nccl", "nvls"])
+    @torch.inference_mode()
+    def test_async_sched_moe_expert_parallel(
+        self, async_sched_mode, inference_moe_token_dispatcher_type
+    ):
+        """Run async scheduling end to end through an EP=2 MoE model."""
+        if Utils.world_size < 2:
+            pytest.skip("Test requires at least 2 GPUs")
+
+        self._run_test(
+            model_provider="gpt",
+            expert_model_parallel_size=2,
+            async_sched_mode=async_sched_mode,
+            inference_moe_token_dispatcher_type=inference_moe_token_dispatcher_type,
+            num_gap_steps=0,
+            num_tokens_to_generate=4,
+        )
+
     @pytest.mark.internal
     @pytest.mark.skipif(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"

tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py

@@ -57,8 +57,9 @@ def _make_engine(async_sched_mode=AsyncScheduleMode.SERIAL, **overrides):
         ({"context_is_hybrid_model": True}, False),
         ({"context_enable_prefix_caching": True}, True),
         ({"materialize_only_last_token_logits": False}, True),
-        ({"model_config_expert_model_parallel_size": 2}, True),
-        ({"model_config_num_moe_experts": 4}, True),
+        ({"model_config_expert_model_parallel_size": 2}, False),
+        ({"model_config_num_moe_experts": 4}, False),
+        ({"model_config_expert_model_parallel_size": 2, "model_config_num_moe_experts": 4}, False),
         ({"model_config_moe_enable_routing_replay": True}, True),
     ],
 )
@@ -120,6 +121,7 @@ def test_async_forward_reenters_controller_after_primer_without_rescheduling():
     engine.logging_step_interval = 0
     engine.metrics_writer = None
     engine.schedule_waiting_requests = mock.Mock(return_value=False)
+    engine._advance_async_sched_primer = mock.AsyncMock()
     engine.context = SimpleNamespace(
         step_count=4,
         prefix_cache_lru_clock=7,
@@ -151,6 +153,7 @@ def test_async_forward_reenters_controller_after_primer_without_rescheduling():
     engine.controller.async_generate_output_tokens_dynamic_batch.assert_has_awaits(
         [mock.call(drain_pending_forward=False), mock.call(drain_pending_forward=False)]
     )
+    engine._advance_async_sched_primer.assert_awaited_once_with()
     range_push.assert_called_once_with("Decode")
     range_pop.assert_called_once_with("Decode")
 
@@ -196,6 +199,7 @@ def test_async_forward_drains_then_admits_and_primes():
     engine.logging_step_interval = 0
     engine.metrics_writer = None
     engine.schedule_waiting_requests = mock.Mock(side_effect=[True, False])
+    engine._advance_async_sched_primer = mock.AsyncMock()
     engine.context = SimpleNamespace(
         step_count=4,
         prefix_cache_lru_clock=7,
@@ -219,3 +223,71 @@ def test_async_forward_drains_then_admits_and_primes():
     engine.controller.async_generate_output_tokens_dynamic_batch.assert_has_awaits(
         [mock.call(drain_pending_forward=True), mock.call()]
     )
+    engine._advance_async_sched_primer.assert_awaited_once_with()
+
+
+@pytest.mark.parametrize("consensus_due", [False, True])
+def test_advance_async_sched_primer_uses_existing_ep_consensus_cadence(consensus_due):
+    """Primer continuation adds no EP collective unless normal consensus is already due."""
+    engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine)
+    engine.use_coordinator = True
+    engine.ep_world_size = 2
+    engine.disable_ep_consensus = False
+    engine.ep_consensus_interval = 20
+    engine._ep_consensus_loop_counter = 20 if consensus_due else 1
+    engine._last_ep_consensus = (3, False)
+    engine.waiting_request_ids = deque([20])
+    engine.context = SimpleNamespace(get_active_request_count=mock.Mock(return_value=2))
+    engine._ep_establish_consensus = mock.AsyncMock(return_value=(3, False))
+
+    asyncio.run(engine._advance_async_sched_primer())
+
+    assert engine._ep_consensus_loop_counter == (21 if consensus_due else 2)
+    if consensus_due:
+        engine._ep_establish_consensus.assert_awaited_once_with(3, signal_consensus=False)
+    else:
+        engine._ep_establish_consensus.assert_not_awaited()
+
+
+def test_advance_ep_consensus_forwards_pause_signal():
+    """The shared cadence helper updates and returns a due pause consensus."""
+    engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine)
+    engine.ep_consensus_interval = 20
+    engine._ep_consensus_loop_counter = 20
+    engine._last_ep_consensus = (3, False)
+    engine._ep_establish_consensus = mock.AsyncMock(return_value=(4, True))
+
+    result = asyncio.run(engine._advance_ep_consensus(2, signal_consensus=True))
+
+    assert result == (4, True)
+    assert engine._last_ep_consensus == (4, True)
+    assert engine._ep_consensus_loop_counter == 21
+    engine._ep_establish_consensus.assert_awaited_once_with(2, signal_consensus=True)
+
+
+@pytest.mark.parametrize(
+    "use_coordinator, ep_world_size, disable_ep_consensus, expected_yield",
+    [(False, 2, False, False), (True, 1, False, False), (True, 2, True, True)],
+)
+def test_advance_async_sched_primer_skips_unneeded_ep_cadence(
+    use_coordinator, ep_world_size, disable_ep_consensus, expected_yield
+):
+    """Primer continuation changes cadence only for coordinated EP execution."""
+    engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine)
+    engine.use_coordinator = use_coordinator
+    engine.ep_world_size = ep_world_size
+    engine.disable_ep_consensus = disable_ep_consensus
+    engine._ep_consensus_loop_counter = 3
+    engine._ep_establish_consensus = mock.AsyncMock()
+
+    with mock.patch(
+        "megatron.core.inference.engines.dynamic_engine.asyncio.sleep", new=mock.AsyncMock()
+    ) as sleep:
+        asyncio.run(engine._advance_async_sched_primer())
+
+    assert engine._ep_consensus_loop_counter == 3
+    engine._ep_establish_consensus.assert_not_awaited()
+    if expected_yield:
+        sleep.assert_awaited_once_with(0)
+    else:
+        sleep.assert_not_awaited()

tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py

@@ -289,6 +289,21 @@ def test_validate_async_sched_support_for_step_ignores_immutable_restrictions():
     controller._validate_async_sched_support_for_step()
 
 
+@pytest.mark.parametrize("expert_model_parallel_size, num_moe_experts", [(1, 4), (2, 4)])
+def test_validate_async_sched_support_for_step_accepts_moe(
+    expert_model_parallel_size, num_moe_experts
+):
+    context = _make_async_sched_context(total_request_count=2)
+    model_config = SimpleNamespace(
+        params_dtype=torch.float32,
+        expert_model_parallel_size=expert_model_parallel_size,
+        num_moe_experts=num_moe_experts,
+        moe_enable_routing_replay=False,
+    )
+
+    _make_async_sched_controller(context, model_config)._validate_async_sched_support_for_step()
+
+
 @pytest.mark.parametrize("unsupported_case", ["paused_request", "chunked_prefill"])
 def test_validate_async_sched_support_for_step_errors(unsupported_case):
     context = _make_async_sched_context(total_request_count=2)