A resumed call whose saved output was already finished crashes when router replay is on

NVIDIA-NeMo/RL #4612 · permalinks pin head 2548c679. The review comment is in a pending review; its link works once the review is published.

Before
Setup refused to start a run with both prefix recovery and router replay, so this case could not happen.
This PR
Removes that check and saves routes with each prefix. But one resume path never asks vLLM for anything, so it has no routes to return.
Status
Open — not fixed in this PR. Raised in the review. Found by two reviewers separately and reproduced with the real route-attach function.

Three pages on this PR — this one: the one resume path that has no routes.

Background. Router replay saves which experts each token went to in an MoE model while vLLM generates. Training then forces the same experts. Prefix recovery saves the tokens a call has generated so far at each checkpoint (a "cut"). After a restart, the call resumes from them instead of starting over. Dotted words have a short explanation: hover or tap them.

A cut can land after vLLM has produced a call's last token but before the call's final row is saved. It can also land when the saved tokens already hit the output limit. On resume, _classify_restored_prefix_terminal sees this and skips decoding. _build_restored_prefix_terminal_output then builds a fake vLLM output with no routed_experts. With router replay on, the reply always goes through the route-attach step, and that step raises when there are no routes. The saved routes were fetched and checked, then never used.

What happens

Example: prompt [10, 11]. The call generated 12, 13, 14, and 14 was the last token its limit allows. Then the cut was taken. 1. Saved at the cut tokens [12, 13, 14] routes [1, 2, 3, 4, ph] finish: "length" 2. Resume: already finished? _classify_restored_prefix_terminal → yes. No decoding. vLLM is never asked for routes 3. Fake vLLM output _build_restored_prefix_terminal_output token_ids=[] finish_reason="length" routed_experts = None 4. Route-attach step (router replay on) attach_routed_experts_to_chat_… RuntimeError: … the generation output did not include routed_experts. 5. The call fails HTTP 500, capture aborted The saved prefix is lost; a retry regenerates the call from scratch. saved routes: fetched and checked, then never used where it breaks ph = the placeholder row for the last position, which never gets a route of its own

Side effect: the generated_len == 0 case in the new merge, replace_len = restored.shape[0] if generated_len == 0 else …, is written for exactly this call, but the call never gets there. No test covers a finished resume with router replay on. A GPU run with router replay and a forced restart would have shown it.

The fix — proposed, not in the PR

Nothing below is implemented, and it has not been run.

Where the fix lands: RL only, one file, vllm_worker_async.py. No Gym change: the saved routes are already fetched and checked on resume. They only need to be handed to the fake output.

@@ vllm_worker_async.py L26, L71-L72 · imports
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from nemo_rl.models.generation.vllm.utils import (
+    R3_MISSING_ROUTE_SENTINEL,
     attach_routed_experts_to_chat_response_choices,

@@ vllm_worker_async.py L92-L99 · _RestoredPrefixTerminal
     finish_reason: Literal["stop", "length"]
     stop_reason: str | int | None = None
+    # Saved routes for every prompt position, so the route-attach step has rows
+    # even though no decoding runs. -1 where no route was ever computed.
+    routed_experts: torch.Tensor | None = None

@@ vllm_worker_async.py L182-L190 · _build_restored_prefix_terminal_output
                 finish_reason=terminal.finish_reason,
                 stop_reason=terminal.stop_reason,
+                routed_experts=terminal.routed_experts,
             )

@@ vllm_worker_async.py L2035 · preprocess_chat
                 if restored_prefix_terminal is not None:
+                    if generation_cut_routed_experts is not None:
+                        saved = generation_cut_routed_experts
+                        # Rows before prev_len were never routed by this call.
+                        unrouted = torch.full(
+                            (admission.prev_len, *saved.shape[1:]),
+                            R3_MISSING_ROUTE_SENTINEL,
+                            dtype=saved.dtype,
+                        )
+                        # Drop the saved placeholder row; route-attach adds its own.
+                        restored_prefix_terminal = replace(
+                            restored_prefix_terminal,
+                            routed_experts=torch.cat((unrouted, saved[:-1])),
+                        )
                     request.min_tokens = 0

@@ vllm_worker_async.py L1459 · _merge_generation_cut_routed_experts
-        replace_len = restored.shape[0] if generated_len == 0 else restored.shape[0] - 1
+        # The saved last row is a placeholder; the row from this call wins.
+        replace_len = restored.shape[0] - 1

On the example: the fake output now carries [1, 2, 3, 4]. The route-attach step adds a -1 for the last position (same as every normal reply), giving [1, 2, 3, 4, -1]. That is exactly what the call would have saved if no restart had happened. The merge then has no special case left. A simpler fallback is to refuse the cut, so the call regenerates, when routes are on and the call is already finished. That works today only because the last token and its finish signal are recorded under one lock.

🔴 review comment — vllm_worker_async.py:1459

So what. This PR removes the setup check that kept these two features apart. Any MoE run with both on can now fail a resumed call this way. Gym uses a saved prefix only once (a used cut is not handed back), so nothing loops: the episode fails, and the prefix the checkpoint saved is wasted. The fix is about 15 lines plus one test, and it removes a special case from the merge.