From 28563504dd72ce02716b786d3907e0222ae4d862 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 10 Jul 2026 15:50:41 -0400 Subject: [PATCH 1/2] Support async scheduling for MoE and expert parallelism Signed-off-by: Lawrence McAfee --- .../core/inference/engines/dynamic_engine.py | 137 +++++++++++------- .../inference/engines/test_dynamic_engine.py | 28 ++++ .../test_dynamic_engine_async_sched.py | 58 +++++++- .../test_text_generation_controller.py | 15 ++ 4 files changed, 183 insertions(+), 55 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index ac45c1525b9..e94b3172560 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/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,87 @@ 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_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 + + 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 + ): + 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. + self._last_ep_consensus = await self._ep_establish_consensus( + local_pending, signal_consensus=False + ) + + global_work, _ = self._last_ep_consensus + self._ep_consensus_loop_counter += 1 + 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 +2085,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 +2094,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 +2686,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. diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 2513bd78b6c..1a5310c1ce4 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/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" diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py b/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py index b34b0e1714d..b7d82519d51 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py +++ b/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") @@ -219,3 +222,54 @@ 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()] ) + + +@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() + + +@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() diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index 5074767a1bd..3d6ecf5f0eb 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/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) From 7ee19802d0005ba1e6300c88dfb2d98ce1cd480b Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 13 Jul 2026 17:26:21 -0400 Subject: [PATCH 2/2] Share expert consensus cadence Signed-off-by: Lawrence McAfee --- .../core/inference/engines/dynamic_engine.py | 63 ++++++++++--------- .../test_dynamic_engine_async_sched.py | 18 ++++++ 2 files changed, 50 insertions(+), 31 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index e94b3172560..be2588db0bd 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1991,6 +1991,31 @@ async def _ep_establish_consensus( 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. @@ -2006,20 +2031,10 @@ async def _advance_async_sched_primer(self) -> None: await asyncio.sleep(0) return - 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 - ): - 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. - self._last_ep_consensus = await self._ep_establish_consensus( - local_pending, signal_consensus=False - ) - - global_work, _ = self._last_ep_consensus - self._ep_consensus_loop_counter += 1 + 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]: @@ -2753,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. diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py b/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py index b7d82519d51..70cdd4b3081 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py @@ -199,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, @@ -222,6 +223,7 @@ 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]) @@ -247,6 +249,22 @@ def test_advance_async_sched_primer_uses_existing_ep_consensus_cadence(consensus 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)],