From 8e8e278dccad51b9d73a0b717a521946aa85e25e Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 14 Jul 2026 14:17:47 -0400 Subject: [PATCH] Run async scheduling primers as standalone steps Signed-off-by: Lawrence McAfee --- .../core/inference/engines/dynamic_engine.py | 9 +- .../text_generation_controller.py | 69 +++++++--- .../test_dynamic_engine_async_sched.py | 45 +++++++ .../test_text_generation_controller.py | 121 ++++++++++++------ 4 files changed, 186 insertions(+), 58 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index c0372127e98..7ce580eb94e 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -44,6 +44,7 @@ ) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + DynamicBatchControllerStepResult, TextGenerationController, ) from megatron.core.inference.utils import Counter, InferenceMode, await_process_call @@ -1971,7 +1972,13 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]: if will_log_this_step: self.step_start_event.record() - result = await self.controller.async_generate_output_tokens_dynamic_batch() + while True: + controller_result: DynamicBatchControllerStepResult = ( + await self.controller.async_generate_output_tokens_dynamic_batch() + ) + if not controller_result.primer_only: + result = controller_result.output + break if will_log_this_step: self.step_end_event.record() self.step_end_event.synchronize() diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index d265ca5fa23..6b2a8b1b494 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -110,6 +110,19 @@ def clear(self) -> None: self.ready_event = None +@dataclass(frozen=True) +class DynamicBatchControllerStepResult: + """Result of one dynamic-batching controller step. + + Attributes: + output: Sampled-step output, or ``None`` when no output was produced. + primer_only: Whether the step launched only an async-scheduling primer. + """ + + output: Optional[Dict] = None + primer_only: bool = False + + @dataclass class _AsyncScheduleResolveResult: """State produced by async scheduling request resolution.""" @@ -2073,11 +2086,13 @@ def _run_async_sched_resolve( compaction_done_event=compaction_done_event, ) - async def _run_async_sched_step(self, *, overlap: bool) -> Optional[Dict]: + async def _run_async_sched_step(self, *, overlap: bool) -> DynamicBatchControllerStepResult: """Run one decode-only step using the async scheduling path. - The first decode step launches and completes a forward primer so logits - exist. Steady-state overlap follows this schedule:: + A controller step launches at most one model forward. When no pending + logits exist, this method launches only a forward primer and returns. + The engine immediately invokes the controller again to consume that + primer. Steady-state overlap follows this schedule:: CPU: prepare request state N+1 compute stream: forward N -> sample N -> copy input N+1 @@ -2097,8 +2112,8 @@ async def _run_async_sched_step(self, *, overlap: bool) -> Optional[Dict]: for current-step GPU work. Returns: - Optional[Dict]: Step result for sampled and finished requests, or - `None` when no requests are active. + DynamicBatchControllerStepResult: A primer-only indication or the + sampled step output. The output is `None` when no work is active. """ context = self.inference_wrapped_model.inference_context @@ -2109,7 +2124,7 @@ async def _run_async_sched_step(self, *, overlap: bool) -> Optional[Dict]: active_request_count = context.total_request_count - context.paused_request_count if context.active_token_count == 0 and active_request_count == 0: self._async_sched_logits.clear() - return None + return DynamicBatchControllerStepResult() # ------------------------------------------------------------------------- # Primer @@ -2117,15 +2132,23 @@ async def _run_async_sched_step(self, *, overlap: bool) -> Optional[Dict]: # Launch the forward primer if no existing logits state can be reused. primer_launched, primer_bookkeeping_done_event = self._run_async_sched_forward_primer() + current_logits_ready_event = self._async_sched_logits.ready_event + + if primer_launched: + # Overlap only waits until preparation may safely reuse the pinned + # CPU bookkeeping source. Serial mode completes the primer forward. + if overlap: + self._synchronize_async_sched_event(primer_bookkeeping_done_event) + else: + self._synchronize_async_sched_event(current_logits_ready_event) + return DynamicBatchControllerStepResult(primer_only=True) + with torch.inference_mode(): - current_logits_ready_event = self._async_sched_logits.ready_event cuda_graph_request_count = self._async_sched_logits.cuda_graph_request_count - # Serial mode waits for logits; overlap only waits for a new primer's H2D source read. + # Serial mode waits for the pending logits before preparing the next step. if not overlap: self._synchronize_async_sched_event(current_logits_ready_event) - elif primer_launched: - self._synchronize_async_sched_event(primer_bookkeeping_done_event) # ------------------------------------------------------------------------- # Prepare @@ -2214,7 +2237,7 @@ async def _run_async_sched_step(self, *, overlap: bool) -> Optional[Dict]: # Yield only after resolution is complete and forward N+1 is already submitted. await asyncio.sleep(0) - return result + return DynamicBatchControllerStepResult(output=result) # ------------------------------------------------------------------------- # End async scheduling methods @@ -2385,7 +2408,7 @@ async def _run_legacy_step(self, skip_bookkeeping: Optional[bool] = False) -> Op async def async_generate_output_tokens_dynamic_batch( self, skip_bookkeeping: Optional[bool] = False - ) -> Optional[Dict]: + ) -> DynamicBatchControllerStepResult: """Forward step the model and update the inference context. Args: @@ -2393,14 +2416,15 @@ async def async_generate_output_tokens_dynamic_batch( on the legacy path. Returns: - Optional[Dict]: Step result for sampled and finished requests, or - `None` when no requests are active. + DynamicBatchControllerStepResult: One controller-step result. """ context = self.inference_wrapped_model.inference_context mode = context.config.async_sched_mode if mode == AsyncScheduleMode.LEGACY or context.num_prefill_requests != 0: - return await self._run_legacy_step(skip_bookkeeping) + return DynamicBatchControllerStepResult( + output=await self._run_legacy_step(skip_bookkeeping) + ) if mode == AsyncScheduleMode.SERIAL: assert not skip_bookkeeping, "Async scheduling requires request bookkeeping." return await self._run_async_sched_step(overlap=False) @@ -2413,9 +2437,20 @@ async def async_generate_output_tokens_dynamic_batch( def generate_output_tokens_dynamic_batch( self, loop: Optional[asyncio.AbstractEventLoop] = None ) -> Optional[Dict]: - """Synchronous wrapper for `self.async_generate_output_tokens_dynamic_batch.""" + """Synchronously run dynamic batching through any primer-only calls. + + Args: + loop (Optional[asyncio.AbstractEventLoop]): Event loop used to run + the asynchronous controller. + + Returns: + Optional[Dict]: Step output, or `None` when no work is active. + """ loop = get_asyncio_loop(loop) - return loop.run_until_complete(self.async_generate_output_tokens_dynamic_batch()) + while True: + result = loop.run_until_complete(self.async_generate_output_tokens_dynamic_batch()) + if not result.primer_only: + return result.output def _update_top_n_logprobs_dict( self, 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 abab280497f..d0836aaa40a 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 @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import asyncio from types import SimpleNamespace from unittest import mock @@ -7,7 +8,11 @@ from megatron.core.inference.config import AsyncScheduleMode from megatron.core.inference.engines import DynamicInferenceEngine +from megatron.core.inference.engines.dynamic_engine import EngineState from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + DynamicBatchControllerStepResult, +) def _make_engine(async_sched_mode=AsyncScheduleMode.SERIAL, **overrides): @@ -102,3 +107,43 @@ def test_add_request_runs_async_sched_request_validation(): engine._add_request(request) engine._validate_async_sched_support_for_request.assert_called_once_with(request) + + +def test_async_forward_reenters_controller_after_primer_without_rescheduling(): + """The engine immediately consumes a primer without crossing its step boundary.""" + engine = DynamicInferenceEngine.__new__(DynamicInferenceEngine) + engine.state = EngineState.RUNNING + engine.logging_step_interval = 0 + engine.metrics_writer = None + engine.schedule_waiting_requests = mock.Mock() + engine.context = SimpleNamespace( + step_count=4, + prefix_cache_lru_clock=7, + active_token_count=2, + is_decode_only=mock.Mock(return_value=True), + ) + expected_output = {"sample": "tokens"} + engine.controller = SimpleNamespace( + async_generate_output_tokens_dynamic_batch=mock.AsyncMock( + side_effect=[ + DynamicBatchControllerStepResult(primer_only=True), + DynamicBatchControllerStepResult(output=expected_output), + ] + ) + ) + + with ( + mock.patch("megatron.core.inference.engines.dynamic_engine.nvtx_range_push") as range_push, + mock.patch("megatron.core.inference.engines.dynamic_engine.nvtx_range_pop") as range_pop, + ): + output, context_state, step_time = asyncio.run(engine.async_forward()) + + assert output is expected_output + assert context_state == {"active_token_count": 2, "step_count": 4, "kv_stats": None} + assert step_time == 0.0 + assert engine.context.step_count == 5 + assert engine.context.prefix_cache_lru_clock == 8 + engine.schedule_waiting_requests.assert_called_once_with() + assert engine.controller.async_generate_output_tokens_dynamic_batch.await_count == 2 + range_push.assert_called_once_with("Decode") + range_pop.assert_called_once_with("Decode") 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 f0939fa0ad9..9339e1690fd 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 @@ -34,6 +34,7 @@ from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.text_generation_controllers.text_generation_controller import ( AsyncScheduleLogitsState, + DynamicBatchControllerStepResult, TextGenerationController, ) from megatron.core.inference.utils import InferenceMode @@ -496,7 +497,7 @@ def test_run_async_sched_forward_primer(is_valid): controller._run_async_sched_forward.assert_called_once_with(input_ids, position_ids) -def test_async_sched_step_returns_none_without_active_requests(): +def test_async_sched_step_returns_empty_result_without_active_requests(): context = _make_async_sched_context(total_request_count=0) context.active_token_count = 0 controller = _make_async_sched_controller(context) @@ -507,13 +508,52 @@ def test_async_sched_step_returns_none_without_active_requests(): result = asyncio.run(controller._run_async_sched_step(overlap=False)) - assert result is None + assert result == DynamicBatchControllerStepResult() assert not controller._async_sched_logits.is_valid assert controller._async_sched_logits.cuda_graph_request_count is None assert controller._async_sched_logits.ready_event is None controller._validate_async_sched_support_for_step.assert_called_once_with() +@pytest.mark.parametrize( + "overlap, expected_wait", [(False, "forward"), (True, "primer_bookkeeping")] +) +def test_async_sched_primer_step_returns_before_prepare(overlap, expected_wait): + """A primer-only call launches no other async-scheduling phase.""" + context = _make_async_sched_context(total_request_count=3) + controller = _make_async_sched_controller(context) + controller._async_sched_logits = AsyncScheduleLogitsState() + controller._validate_async_sched_support_for_step = mock.Mock() + + def run_primer(): + controller._async_sched_logits.set_pending(7, "forward") + return True, "primer_bookkeeping" + + controller._run_async_sched_forward_primer = mock.Mock(side_effect=run_primer) + controller._synchronize_async_sched_event = mock.Mock() + controller._run_async_sched_prepare = mock.Mock() + controller._run_async_sched_sample = mock.Mock() + controller._copy_async_sched_sample_to_cpu = mock.Mock() + context.copy_async_sched_sample_to_forward = mock.Mock() + controller._run_async_sched_publish_bookkeeping = mock.Mock() + controller._run_async_sched_forward = mock.Mock() + controller._run_async_sched_resolve = mock.Mock() + + result = asyncio.run(controller._run_async_sched_step(overlap=overlap)) + + assert result == DynamicBatchControllerStepResult(primer_only=True) + controller._synchronize_async_sched_event.assert_called_once_with(expected_wait) + controller._run_async_sched_prepare.assert_not_called() + controller._run_async_sched_sample.assert_not_called() + controller._copy_async_sched_sample_to_cpu.assert_not_called() + context.copy_async_sched_sample_to_forward.assert_not_called() + controller._run_async_sched_publish_bookkeeping.assert_not_called() + controller._run_async_sched_forward.assert_not_called() + controller._run_async_sched_resolve.assert_not_called() + assert context.async_sched_step_count == 0 + assert context.async_sched_compaction_step_count == 0 + + def test_run_async_sched_sample_reuses_gpu_buffer(): context = _make_async_sched_context(total_request_count=3) controller = _make_async_sched_controller(context) @@ -622,11 +662,10 @@ def compact_logits(survivor_idxs): @pytest.mark.parametrize( - "overlap, has_valid_logits, expected_call_order", + "overlap, expected_call_order", [ ( False, - True, [ "wait:current", "prepare", @@ -645,26 +684,7 @@ def compact_logits(survivor_idxs): ), ( True, - True, - [ - "prepare", - "sample", - "copy_input", - "copy_sample", - "publish", - "forward", - "wait:sample", - "wait:bookkeeping", - "resolve", - "yield", - ], - ), - ( - True, - False, [ - "primer", - "wait:primer_bookkeeping", "prepare", "sample", "copy_input", @@ -679,7 +699,7 @@ def compact_logits(survivor_idxs): ), ], ) -def test_async_sched_step_order(overlap, has_valid_logits, expected_call_order): +def test_async_sched_step_order(overlap, expected_call_order): sample_tokens = torch.tensor([1, 2, 3], dtype=torch.int64) sampled_tokens_cpu = sample_tokens.clone() input_ids = torch.tensor([[101, 102, 103]]) @@ -687,20 +707,12 @@ def test_async_sched_step_order(overlap, has_valid_logits, expected_call_order): context = _make_async_sched_context(total_request_count=3) controller = _make_async_sched_controller(context) controller._async_sched_logits = AsyncScheduleLogitsState( - is_valid=has_valid_logits, - cuda_graph_request_count=7 if has_valid_logits else None, - ready_event="current" if has_valid_logits else None, + is_valid=True, cuda_graph_request_count=7, ready_event="current" ) controller._validate_async_sched_support_for_step = mock.Mock() call_order = [] - def run_primer(): - if not has_valid_logits: - call_order.append("primer") - controller._async_sched_logits.set_pending(7, "current") - return not has_valid_logits, "primer_bookkeeping" if not has_valid_logits else None - - controller._run_async_sched_forward_primer = mock.Mock(side_effect=run_primer) + controller._run_async_sched_forward_primer = mock.Mock(return_value=(False, None)) controller._synchronize_async_sched_event = mock.Mock( side_effect=lambda event: call_order.append(f"wait:{event}") ) @@ -742,8 +754,9 @@ async def yield_to_event_loop(_delay): ): result = asyncio.run(controller._run_async_sched_step(overlap=overlap)) - assert result["sample"].tolist() == sample_tokens.tolist() - assert result["cuda_graph_request_count"] == 7 + assert not result.primer_only + assert result.output["sample"].tolist() == sample_tokens.tolist() + assert result.output["cuda_graph_request_count"] == 7 assert context.async_sched_step_count == 1 assert context.async_sched_compaction_step_count == 1 assert call_order == expected_call_order @@ -784,8 +797,10 @@ def run_forward(*_args): controller._run_async_sched_forward = mock.Mock(side_effect=run_forward) - result = asyncio.run(controller._run_async_sched_step(overlap=False)) + step_result = asyncio.run(controller._run_async_sched_step(overlap=False)) + result = step_result.output + assert not step_result.primer_only assert torch.equal(result["sample"], sampled_tokens) assert result["finished_request_ids"].tolist() == expected_finished_ids context.copy_async_sched_sample_to_forward.assert_called_once() @@ -826,8 +841,10 @@ async def run_step(): ) return await controller._run_async_sched_step(overlap=True) - result = asyncio.run(run_step()) + step_result = asyncio.run(run_step()) + result = step_result.output + assert not step_result.primer_only assert result["sample"].tolist() == [1] assert observed == [(1, False)] @@ -851,12 +868,36 @@ def test_async_generate_output_tokens_dynamic_batch_routes( controller = _make_async_sched_controller(context) controller._run_legacy_step = mock.AsyncMock(return_value="legacy") controller._run_async_sched_step = mock.AsyncMock( - side_effect=lambda *, overlap: "overlap" if overlap else "async" + side_effect=lambda *, overlap: DynamicBatchControllerStepResult( + output="overlap" if overlap else "async" + ) ) result = asyncio.run(controller.async_generate_output_tokens_dynamic_batch(skip_bookkeeping)) - assert result == expected_result + assert not result.primer_only + assert result.output == expected_result + + +def test_generate_output_tokens_dynamic_batch_consumes_primer_result(): + """The synchronous API hides primer-only controller calls from its caller.""" + controller = _make_async_sched_controller() + expected_output = {"sample": torch.tensor([1])} + controller.async_generate_output_tokens_dynamic_batch = mock.AsyncMock( + side_effect=[ + DynamicBatchControllerStepResult(primer_only=True), + DynamicBatchControllerStepResult(output=expected_output), + ] + ) + loop = asyncio.new_event_loop() + + try: + result = controller.generate_output_tokens_dynamic_batch(loop=loop) + finally: + loop.close() + + assert result is expected_output + assert controller.async_generate_output_tokens_dynamic_batch.await_count == 2 @pytest.mark.parametrize(