Run async scheduling primers as standalone controller steps

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

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,23 +2124,31 @@ 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
         # -------------------------------------------------------------------------
         # 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,22 +2408,23 @@ 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:
             skip_bookkeeping (Optional[bool]): If true, skip context bookkeeping
                 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,

tests/unit_tests/inference/engines/test_dynamic_engine_async_sched.py

@@ -1,13 +1,18 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import asyncio
 from types import SimpleNamespace
 from unittest import mock
 
 import pytest
 
 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")