Extend dynamic inference asynchronous scheduling support

examples/inference/advanced/gpt_dynamic_inference.py

@@ -128,7 +128,10 @@ def _process_step_result(result):
         """Process a single engine step result, updating bookkeeping state."""
         nonlocal total_output_tokens, num_requests_finished
 
-        is_decode_only = engine.is_decode_only
+        decode_only = engine.decode_only
+        is_decode_only = (
+            decode_only.launched if decode_only.launched is not None else decode_only.consumed
+        )
 
         # Record cuda_graph_request_count.
         cuda_graph_request_count = result["cuda_graph_request_count"]
@@ -228,7 +231,7 @@ def _process_step_result(result):
             add_times.append(get_curr_time(do_broadcast=False) - add_start)
 
             # Step inference engine (i.e., generate a token for each active request).
-            # Before step, we haven't done the scheduling, so we cannot know the is_decode_only
+            # The engine reports the consumed and launched decode-only states after scheduling.
             try:
                 result = engine.step_modern()
             except EngineSuspendedError as e:

megatron/core/inference/config.py

@@ -139,11 +139,8 @@ class AsyncScheduleMode(str, Enum):
     LEGACY = "legacy"
     """Resolve requests before preparing the next forward pass."""
 
-    SERIAL = "serial"
-    """Prepare and forward speculatively before resolving the sampled requests."""
-
-    OVERLAP = "overlap"
-    """Overlap async scheduling prepare/sample and forward/resolve phases."""
+    ASYNC = "async"
+    """Overlap asynchronous scheduling phases by reordering them to prepare-before-resolve."""
 
 
 @dataclass

megatron/training/arguments.py

@@ -2012,13 +2012,11 @@ def _add_inference_args(parser):
                             'is requested but the package is not installed.')
     group.add_argument('--inference-dynamic-batching-async-sched-mode',
                        type=str, default='legacy',
-                       choices=['legacy', 'serial', 'overlap'],
+                       choices=['legacy', 'async'],
                        help='Async scheduling mode for dynamic batching. '
                             '"legacy" (default) preserves the existing resolve-before-prepare '
-                            'path. "serial" speculatively prepares and forwards decode-only '
-                            'steps before resolving finished requests. "overlap" uses the same '
-                            'async scheduling path while overlapping prepare/sample and '
-                            'forward/resolve phases.')
+                            'path. "async" overlaps asynchronous scheduling phases by reordering '
+                            'them to prepare-before-resolve.')
     group.add_argument('--inference-dynamic-batching-logprobs-mode',
                        type=str, default='raw_logprobs',
                        choices=['raw_logprobs', 'processed_logprobs'],

megatron/training/config/inference_config.py

@@ -136,11 +136,10 @@ class InferenceSetupConfig:
     """Which sampling kernels to use during inference. Falls back to "torch" with a warning if
     "flashinfer" is requested but the package is not installed."""
 
-    inference_dynamic_batching_async_sched_mode: Literal["legacy", "serial", "overlap"] = "legacy"
+    inference_dynamic_batching_async_sched_mode: Literal["legacy", "async"] = "legacy"
     """Async scheduling mode for dynamic batching. "legacy" (default) preserves the
-    existing resolve-before-prepare path. "serial" speculatively prepares and forwards decode-only
-    steps before resolving finished requests. "overlap" uses the same async scheduling path while
-    overlapping prepare/sample and forward/resolve phases."""
+    existing resolve-before-prepare path. "async" overlaps asynchronous scheduling phases by
+    reordering them to prepare-before-resolve."""
 
     inference_dynamic_batching_logprobs_mode: Literal["raw_logprobs", "processed_logprobs"] = (
         "raw_logprobs"