How MInf router replay gets its expert routes into training, next to the vLLM path

Explainer for lauradang/RL #2, a PR into NVIDIA-NeMo/RL #4129. All permalinks pin the PR head b311ce9.

Before
Router replay worked only with vLLM generation. The SingleController setup rejected router replay with Megatron generation (setup.py:1196, before).
This PR
Reshapes the expert routes from Megatron Inference (MInf) into the same per-call rows vLLM already writes, so the rest of the pipeline reads them unchanged.
Status
Proposed — open PR into #4129, under review.

Two pages on this PR, one question each — this one: how the new MInf path compares with the vLLM path.

Router replay makes the trainer use the same experts, per token and per MoE layer, that the generation engine picked, so the two do not disagree on routing. MInf records one route row per token except the last: shape [T-1, L, K] for T tokens, L MoE layers and top-K experts (inference_request.py:917). But each model call in a multi-turn rollout stores only its new tokens: rows prev_len to T, where prev_len is the length the earlier turns already stored. So the PR adds one all -1 row and cuts off the first prev_len rows (_delta_align_minf_routing_indices). vLLM also gives [T-1, L, K]; its NeMo-RL worker already adds a final row (filled with experts 0
K-1) and cuts rows the same way. The two paths differ only in where that row is added and what goes in it.

Two engines, one staging row

Call 2 of the example below: 21 tokens, the first 7 already stored by call 1, 2 MoE layers, top-8 experts MInf engine routing_indices [20, 2, 8] int16 no row for the last token inference_request.py:917 NeMo-RL: MInf staging step add a -1 row → [21, 2, 8] keep rows [7:] → [14, 2, 8] missing or wrong-length routes: call fails tq_token_sink.py:451-459 _stage_admitted :694 Gym commit complete_call(extras=
) routes passed in by hand: Gym's Megatron reader returns no extras tq_token_sink.py:707 Gym megatron.py:82 vLLM engine routed_experts [20, 2, 8] no row for the last token vllm/utils.py:261 NeMo-RL: vLLM worker add a 0
7 row → [21, 2, 8] keep rows [7:] → [14, 2, 8] wrong-length routes: warn, drop, go on vllm/utils.py:310-322 vllm_worker_async.py:634 Gym commit complete_call_from_response Gym's vLLM reader pulls the trimmed routes out of the reply vllm_worker_async.py:674 Gym vllm.py:67 Shared from here TQTokenSink.stage rows must equal the new token count: 14 tq_token_sink.py:305 execute_route_plan joins both calls into one [21, 2, 8] row route_assembly.py:110 trainer -1 row → its own router router_replay.py:365 MInf only (new in this PR) vLLM only (already there) one shared code path for both

The same 2-turn rollout through both

Call 1: a 3-token prompt, 4 generated tokens, 7 in all. Call 2 continues it: those 7, then 5 new user tokens (a 12-token prompt), then 9 generated tokens, 21 in all. The model has 2 MoE layers and picks the top 8 experts.

# call 2 — prev_len = 7 (call 1's whole length), T = 21, so 14 new tokens MInf routing_indices [20, 2, 8] int16 # one row per token except the last add one all -1 row → [21, 2, 8] # tq_token_sink.py:451-458 keep rows [7:] → [14, 2, 8] # :459 vLLM routed_experts [20, 2, 8] # also one row per token except the last, utils.py:261 add one 0
7 row → [21, 2, 8] # utils.py:310-322 keep rows [7:] → [14, 2, 8] # vllm_worker_async.py:634 both TQTokenSink.stage: 14 rows == 14 new tokens → staged # tq_token_sink.py:305
The full rollout the trainer sees: one column per token. Which route does each token get? call 1: prev_len 0, stores 7 rows call 2: prev_len 7, stores 14 rows token 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 MInf -1 -1 vLLM 0
7 0
7 prompt / user token generated token the engine's real route -1: trainer's router picks experts 0
7 forced Column 6 (dashed) is call 1's last generated token. In call 1 no engine had routed it yet, so each engine put its filler there. In call 2 both engines do route it, since it is part of call 2's prompt. But that row sits before prev_len = 7, so [7:] drops it. Column 20 is the rollout's last token. Its route only shapes its own output, which training does not score. documented in token-capture-ledger.md:136

What still differs after they meet

MInf (this PR)vLLM (already there)
Filler for a token with no route all -1: the trainer's own router picks (route_assembly.py:45) experts 0
K-1, forced (utils.py:310)
Routes missing or the wrong length the call fails, so the whole rollout is thrown out (tq_token_sink.py:694, :662, nemo_gym.py:972) wrong length: routes dropped with a warning; the rollout is kept and those rows become -1 (vllm_worker_async.py:636, route_assembly.py:140)
token_capture.defer_routed_experts_to_policy must stay false; setup rejects true (setup.py:1205) either value works

Why keep both? vLLM stays the usual generation engine. MInf lets a run generate with Megatron itself, and until this PR such a run could not use router replay at all. In exchange, the MInf path is stricter — one bad route costs the whole rollout — and it has no deferred mode yet.

What this means for you. Router replay now works with generation.backend: megatron on the SingleController + Gym token-capture path, and once the routes are staged the two engines share every line of code. Before you switch, keep defer_routed_experts_to_policy: false, and expect missing routes to show up as rejected rollouts rather than warnings. The token at each turn boundary (column 6) gets no replayed route on either engine; MInf lets the trainer's router choose there, while vLLM forces experts 0
7.