Extend dynamic inference asynchronous scheduling support

[lmcafee-nvidia]

Summary

Extend dynamic inference asynchronous scheduling across the remaining request and output features needed by the primary inference path. The controller keeps one pending forward while CPU request work is reordered around GPU sampling and forward execution.

Architecture

• Route asynchronous iterations through phase orders matched to the current work:
  • no-overlap: sample or verify the previous output, resolve lifecycle changes, prepare and admit requests, then launch the next forward
  • overlap: prepare the next decode step, sample the previous output, launch the next forward, then resolve the previous output
  • MTP overlap: verify drafts, rewind rejected state, prepare, launch the next forward, then resolve
• Keep the first admitted request as a primer-only forward and preserve one base-model forward per controller iteration.
• Keep expert-parallel ranks in a consistent MTP-before-base-forward order, including ranks with dummy work.
• Add single-bank Mamba handling by deferring admission until pending state is safe to reuse.
• Support chunked prefill and KV/Mamba prefix caching within asynchronous scheduling.
• Use the existing sampling implementations for greedy and non-greedy sampling with Torch and FlashInfer backends.
• Preserve sampled log probabilities, top-n log probabilities, and stop-word completion behavior across request compaction.
• Expose exactly two scheduling modes: legacy and async.

Routing replay remains explicitly unsupported by this PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[sidsingh-nvidia]

The name _run_async_sched_prefill_step is misleading (controller L2437). It can fire for both prefills and decodes

1. admit a new request - actual prefill;
2. transition an in-flight prefill → decode (num_prefill_requests > 0 is set from the previous prefill step) — produces decode output;
3. cold start (not _async_sched_logits.is_valid) - launches a primer (can be a prefill)

The real common thread is that all three flush the in-flight forward and issue the next one after resolve/admit.

Recommendation.

1. Rename the schedule to reflect what it is. I recommend using async_sched_flush since the one invariant in this schedule is that we flush out all GPU work from the past turn by sampling and resolving it.
2. Suggest breaking down _run_async_sched_prefill_step into two methods, which can be called from the engine.

• async_sched_flush_and_prepare() — L2459–2497 (sample/MTP → resolve → prepare/commit); returns the resolved output.
• async_sched_post_flush_forward() — L2506–2526 (issue exactly one base forward: real / primer / dummy).

Engine flow becomes async_sched_flush_and_prepare() → admit() → async_sched_post_flush_forward(). This removes the current callback inversion where schedule_waiting_requests is passed into the controller and called back out.
3. Unify the dummy forward. Route both dummy cases through a single controller dummy-forward primitive called from the engine, with an arg (e.g. run_mtp_dummy: bool) to skip the MTP dummy:

• flush schedule, drained-to-zero → run_mtp_dummy=False (the drain already issued the real MTP collectives this iteration, so only the base half is left);
• idle rank (local_pending == 0) → run_mtp_dummy=True (no drain ran, so both MTP and base must be faked).

I have left a separate comment about the "5 counter restore hack", which we can avoid with a simple fix.

[sidsingh-nvidia]

/claude strict-review

[claude]

Strict review passed — no significant issues found. LGTM

Reviewed for implementation correctness, distributed correctness (TP/PP/SP/CP/EP), performance, and backward compatibility. Key paths verified:

• MTP sequence-length termination: resolved_sequence_lengths = get_active_sequence_lengths() + 1 is correct in the prefill and MTP-decode steps — _rewind_kv_cache folds accepted speculative tokens into kv_length_offsets before the +1, matching the documented legacy ordering. Prefill rewind is a no-op and +1 accounts for the newly sampled base token.
• Accepted-token gating: accepted_tokens gated on num_decode_requests > 0 is consistent with prepare_next_forward_pass only populating decode rows; the None path threads cleanly through the CPU-copy/resolve chain.
• Grouped-token layout: base/draft writes validate MTP shape, transpose drafts correctly, and zero padding; dtype-safe.
• resolve_requests: hole-filling survivor order matches legacy update_requests; token-row moves correctly skipped for the prefill→decode transition (prepare rebuilds them) with active_token_count = 0.
• EP alignment: primer issues dummy MTP-before-base forward to match steady-state collective ordering; dummy-step ordering covered by new tests.

Compliance checks: no new direct parallel_state.get_*_group() reads in megatron/core; changed-signature APIs have only internal/test callers (low external-compat risk); record_cg_wait probe path is non-mutating; no newly-introduced unused variables or arguments.

[lmcafee-nvidia]

@sidsingh-nvidia Thanks for the detailed review.

The naming concern is addressed by renaming the schedules according to their actual execution property: no_overlap and overlap. The no-overlap path also handles lifecycle fallbacks beyond prefill, so I avoided flush; it consumes the pending output but launches the successor forward without globally flushing GPU work.

I'd prefer to keep the phase orchestration in the controller. Sampling/MTP, pending-logit state, request resolution, events, and forward submission constitute one controller step. Splitting that step around engine admission would expose partially completed controller state across calls and make the ordering contract harder to preserve. The synchronous admission callback keeps engine-owned queue mutation at one explicit point without moving controller state into the engine.

I agree that the two async dummy cases can share one controller primitive. I'll consolidate those using a run_mtp_dummy argument while preserving the existing phase ordering.

The counter workaround mentioned at the end has also been removed in 7164358b4: counter preservation is now explicit and independent of prefix-cache preservation.

[lmcafee-nvidia]

/ok to test 5b236ba

[lmcafee-nvidia]

/ok to test 2c957eb

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30374569144