feat(rollout): replay MInf routing indices

[lauradang]

Summary

• enable router replay with Megatron generation on the Gym token-capture path
• join MInf's native [T-1, L, K] route payload with Gym's CaptureAdmission in the serving worker
• validate and delta-align routes before synchronously staging the canonical call row in TransferQueue
• return only ng_commit_coords with the HTTP completion, then commit a token-free CallRecord to Gym's lineage ledger
• verify and linearize staged calls in the CPU reassembler before publishing [B, S, L, K] routes to the canonical training batch

Builds on #4129 and NVIDIA/Megatron-LM#7015 (both merged). Reviewed on the fork as lauradang#2.

Flow

flowchart LR
    subgraph Gym["Gym model server"]
        Admit["resolve parent<br/>create CaptureAdmission"]
        Ledger["CaptureLedger<br/>token-free CallRecord"]
    end

    subgraph Serve["MInf serving worker"]
        Forward["prepare admitted request<br/>MoE forward records top-k IDs"]
        Native["OffloadedRequestPayload<br/>tokens + logprobs<br/>routes: [T-1, L, K]"]
        Stage["TQMegatronTokenStager<br/>validate T-1 and (L, K)<br/>append -1<br/>slice from prev_len"]
    end

    subgraph TQ["TransferQueue"]
        Staging["canonical staging row<br/>key: rollout_id/model_call_id<br/>routes: [delta_len, L, K]"]
        Training["canonical training batch<br/>routed_experts: [B, S, L, K]"]
    end

    Reply["HTTP response<br/>content + ng_commit_coords<br/>no token or route payload"]

    subgraph Finalizer["NeMo-RL CPU reassembler"]
        Verify["fetch staged rows<br/>verify digests and lineage<br/>linearize call chain"]
        Assemble["assemble call route spans<br/>pad the rollout group"]
    end

    Trainer["Megatron trainer<br/>force expert IDs<br/>compute current router scores"]

    Admit -->|"request + offload_params.ng_capture"| Forward
    Forward -. "preparer reads staging_chain prefix" .- Staging
    Forward --> Native --> Stage
    Admit -. "prev_len + lineage" .-> Stage
    Stage -->|"synchronous durable write"| Staging
    Stage -->|"staged, or capture_failed coords"| Reply
    Reply -->|"validate coordinates; strip transport fields"| Ledger
    Ledger -->|"RolloutReceipt manifest"| Verify
    Staging --> Verify --> Assemble --> Training --> Trainer

Loading

The conversion is owned by NeMo-RL through an adapter installed into the MInf
engine, but it runs in the serving worker. Gym forwards the lineage admission,
including prev_len, in offload_params, so routes are already canonical and
delta-aligned when they enter TransferQueue. Before admission, the prompt
preparer in the same worker splices in the prior turns' tokens from the staged
rows named by the admission's staging_chain. Neither the route payload nor
token arrays return in the HTTP response. If staging fails, the reply still
carries ng_commit_coords, marked capture_failed, so Gym records the call as
poisoned.

Test plan

• Ruff check and format check pass for all touched Python files
• On macOS against this branch merged with main: tests/unit/single_controller/test_setup.py (141 passed), tests/unit/algorithms/test_grpo.py (210 passed), tests/unit/data_plane/test_rollout_reassembler.py (14 passed; the MInf stager test needs megatron.core, which only the Linux CI environment has)
• tests/unit/models/megatron/test_router_replay.py, tests/unit/data_plane/test_tq_token_sink.py, and tests/unit/models/generation/test_megatron_token_capture_hosting.py need Megatron/Linux and are covered in CI

🤖 Generated with Claude Code

[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.

[terrykong]

Thanks, Laura. This round addresses the fork review (lauradang#2): 17 of its 21 threads are fixed, including the colocated MInf + KL > 0 stall, and 3 more no longer apply. The last one is the vLLM arange(K) question for @zyzhou5, re-asked inline. The Gym capture API calls are correct against Gym source, and keeping the MInf route alignment separate from vLLM's is the right call. All changed tests pass on a CPU venv and sit in CI lanes that run them.

CI: no unit or functional CI has run yet (copy-pr-bot has not vetted the PR). ruff check and ruff format pass on all changed files.

Most important, in order:

1. Colocated MInf leaves the policy routers recording: the reference pass can crash when prev-logprobs are skipped.
2. Hybrid layer count change: vLLM with deferred routes quietly turns router replay off for hybrid MoE models.
3. async_sched_mode check: an unset key passes, and MCore defaults to async.

Not raised on purpose: router replay on the PPO and distillation entry points fails late for both vLLM and MInf, because only GRPO and SingleController call configure_vllm_for_router_replay (grpo.py, setup.py). That is not new in this PR; it may be worth a tracking issue.

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4307-minf-vs-vllm-routes.html

Generated by Claude Code

[terrykong]

No action needed from @lauradang in this PR — tracked in #4355, assigned to @zyzhou5.

TL;DR — vLLM fills the last token's route row with experts 0…K-1, and on the SingleController path that row becomes a prompt token of the next turn, so the trainer replays made-up experts there.

@zyzhou5, bringing this over from the fork review (lauradang#2). This line says MInf and vLLM fill the last row differently. MInf uses -1, which means "let the trainer's router pick".

• vLLM builds that row from arange(K) in pad_and_align_routed_expert_indices.
• The SingleController vLLM stager keeps the row through experts[prev_len:] (vllm_worker_async.py).
• In a multi-turn rollout, the parent call's last sampled token is a prompt token of the next turn, so the trainer replays experts 0…K-1 for it instead of using its own router.

Follow-up

Tracked in #4355, which covers the whole fix: pad_and_align_routed_expert_indices writes R3_MISSING_ROUTE_SENTINEL (-1) into the last real token's row, batch-padding rows keep arange(K), plus a unit test. Once that lands, this doc line no longer needs to say MInf "deliberately differs from vLLM".