Enable async scheduling by default for dynamic inference (#6851)

Signed-off-by: Lawrence McAfee <lmcafee@nvidia.com>

Author: lmcafee-nvidia

examples/inference/README.md

@@ -111,7 +111,7 @@ Inference can collect traces two ways.
 
 | Path | Enable with | Captures | CUDA graphs |
 |------|-------------|----------|------------|
-| **Sink** | `--moe-enable-routing-replay` | top-K indices only | on |
+| **Sink** | `--moe-enable-routing-replay` + legacy scheduling | top-K indices only | on |
 | **Hook** | no replay + `--cuda-graph-impl none` | indices **+ hidden states + router weights** | must be off |
 
 Only the hook path captures the hidden states and router weights that
@@ -134,10 +134,14 @@ Forward hooks do not fire during CUDA graph replay. MoE cudagraphs must be disab
 
 **Inference — sink** (routing indices only, graphs on):
 
+Routing replay requires legacy scheduling. Because async scheduling is enabled by
+default, explicitly select legacy mode when enabling the routing replay sink.
+
 ```bash
 --moe-routing-trace-path /path/to/trace_dir
 --moe-routing-trace-max-inference-steps 200
 --moe-enable-routing-replay
+--inference-dynamic-batching-async-sched-mode legacy
 ```
 
 **Inference — hook** (adds hidden states + weights for predictability):

examples/inference/advanced/gpt_dynamic_inference.py

@@ -506,7 +506,7 @@ def escape_str(s):
         p_count = len(p_times)
         d_count = len(d_times)
 
-        p_mean = p_total / p_count
+        p_mean = p_total / p_count if p_count != 0 else 0.0
         d_mean = d_total / d_count if d_count != 0 else 0.0
 
         # Commented out for now as the step/add/output times are not calculated correctly.

megatron/core/inference/config.py

@@ -556,8 +556,9 @@ class InferenceConfig:
     enabled), then all DP ranks share the same sampling / generation seed.
     """
 
-    async_sched_mode: AsyncScheduleMode = AsyncScheduleMode.LEGACY
-    """Mode used to schedule dynamic batching inference work."""
+    async_sched_mode: AsyncScheduleMode = AsyncScheduleMode.ASYNC
+    """Mode used to schedule dynamic batching inference work. Defaults to async scheduling; use
+    ``AsyncScheduleMode.LEGACY`` to disable it."""
 
     logprobs_mode: Literal['raw_logprobs', 'processed_logprobs'] = 'raw_logprobs'
     """Whether returned log-probs are modified by the sampling parameters or not."""

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -3405,11 +3405,15 @@ async def async_generate_output_tokens_dynamic_batch(
         self._validate_async_sched_support_for_step(run_async_overlap)
 
         active_request_count = context.total_request_count - context.paused_request_count
-        if context.active_token_count == 0 and active_request_count == 0 and run_async_overlap:
+        if context.active_token_count == 0 and active_request_count == 0:
+            # A lifecycle reset such as RECOMPUTE suspend/resume can remove every
+            # request represented by a pending forward. Discard those stale logits
+            # before the no-overlap path admits and primes the restored requests.
             self._async_sched_logits.clear()
-            return DynamicBatchControllerStepResult(
-                decode_only=DecodeOnly(consumed=None, launched=None)
-            )
+            if run_async_overlap:
+                return DynamicBatchControllerStepResult(
+                    decode_only=DecodeOnly(consumed=None, launched=None)
+                )
 
         if not run_async_overlap or not self._async_sched_logits.is_valid:
             return await self._run_async_sched_step_no_overlap(

megatron/training/arguments.py

@@ -2301,12 +2301,12 @@ def _add_inference_args(parser):
                        help='Use the same inference sampling seed on every data-parallel rank. '
                             '--deterministic-mode also uses the same seed on every DP rank.')
     group.add_argument('--inference-dynamic-batching-async-sched-mode',
-                       type=str, default='legacy',
-                       choices=['legacy', 'async'],
+                       type=str, default='async',
+                       choices=['async', 'legacy'],
                        help='Async scheduling mode for dynamic batching. '
-                            '"legacy" (default) preserves the existing resolve-before-prepare '
-                            'path. "async" overlaps asynchronous scheduling phases by reordering '
-                            'them to prepare-before-resolve.')
+                            '"async" (default) overlaps asynchronous scheduling phases by '
+                            'reordering them to prepare-before-resolve. Select "legacy" to '
+                            'disable async scheduling and use the resolve-before-prepare path.')
     group.add_argument('--inference-dynamic-batching-logprobs-mode',
                        type=str, default='raw_logprobs',
                        choices=['raw_logprobs', 'processed_logprobs'],

megatron/training/config/inference_config.py

@@ -142,10 +142,10 @@ class InferenceSetupConfig:
     generation seed. Disable with --use-same-sampling-seed-across-dp-ranks. Also forced off when
     --deterministic-mode is enabled."""
 
-    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. "async" overlaps asynchronous scheduling phases by
-    reordering them to prepare-before-resolve."""
+    inference_dynamic_batching_async_sched_mode: Literal["async", "legacy"] = "async"
+    """Async scheduling mode for dynamic batching. "async" (default) overlaps asynchronous
+    scheduling phases by reordering them to prepare-before-resolve. Select "legacy" to disable
+    async scheduling and use the resolve-before-prepare path."""
 
     inference_dynamic_batching_logprobs_mode: Literal["raw_logprobs", "processed_logprobs"] = (
         "raw_logprobs"

tests/functional_tests/dynamic_inference_functional_tests.md

@@ -121,7 +121,7 @@ CLI flags below are verified to exist in `megatron/training/arguments.py` and/or
 | Family | Features |
 |---|---|
 | Hybrid (Mamba+Attn) | `--mamba-inference-conv-states-dtype`, `--mamba-inference-ssm-states-dtype`, mamba chunk size |
-| MoE | `--moe-enable-routing-replay` (router replay), `--moe-grouped-gemm`, `--moe-token-dispatcher-type` |
+| MoE | `--moe-enable-routing-replay` (router replay; requires `--inference-dynamic-batching-async-sched-mode legacy`), `--moe-grouped-gemm`, `--moe-token-dispatcher-type` |
 
 ---
 

tests/functional_tests/test_cases/gpt/gpt_dynamic_inference_tp1_pp1_583m_async_sched/model_config.yaml

@@ -39,9 +39,8 @@ MODEL_ARGS:
   --seq-length: 1024
   --temperature: 1.0
   --top_k: 1
-  # Async scheduling only supports greedy sampling (top_k=1, top_p=0.0) and does
-  # not support log probabilities, stop words, chunked prefill, or prefix
-  # caching (see dynamic_engine._validate_async_sched_support_for_request).
+  # Keep the mode explicit so this named test always exercises async scheduling,
+  # independent of the default.
   --inference-dynamic-batching-buffer-size-gb: 20
   --inference-dynamic-batching-async-sched-mode: async
   --dist-ckpt-strictness: log_unexpected

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_cudagraph_zmq/model_config.yaml

@@ -83,6 +83,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-num-cuda-graphs: -1 
   --inference-dynamic-batching-max-requests: 16
   --inference-logging-step-interval: 1
+  --inference-dynamic-batching-async-sched-mode: legacy
   --moe-enable-routing-replay: true
 METRICS:
   - "generated_tokens"

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq/model_config.yaml

@@ -80,6 +80,7 @@ MODEL_ARGS:
   --inference-dynamic-batching-buffer-size-gb: 20
   --inference-dynamic-batching-max-requests: 16
   --inference-logging-step-interval: 1
+  --inference-dynamic-batching-async-sched-mode: legacy
   --moe-enable-routing-replay: true
 METRICS:
   - "generated_tokens"