fix(nemo-gym): refresh final-token route across turns

[zpqiu]

What changed

• Refresh the previous assistant turn final routed-expert row from the next turn prompt prefill.
• Extend the existing router-replay slicing test with a two-turn regression that replaces the prior dummy route while preserving normal token slicing.

Why

The final generated token in a request is not fed back through the model, so router replay pads its route with a valid dummy row. On the next model request that token is part of the prompt and receives a real prefill route, but NeMo Gym previously sliced only the unseen prompt suffix and left the historical dummy route in the training trajectory. Replaying that dummy route can change hidden states and inflate rollout-versus-training logprob error.

The fix updates only the routed-expert metadata for the last already-seen token. It does not modify seen_token_ids or change token slicing.

Concrete two-turn failure

Assume one MoE layer with top-k=2.

The first request has:

prompt tokens:     [1, 2]
generated tokens: [3]

routes:
token 1 -> [10, 11]  (real prompt-prefill route)
token 2 -> [20, 21]  (real prompt-prefill route)
token 3 -> [ 0,  1]  (dummy padding)

Token 3 is sampled from the logits produced while processing token 2. Because the request ends immediately afterward, token 3 is never itself processed by the model in this request, so no real route exists yet.

After a tool result adds tokens [4, 5], the second request sends the full history:

prompt tokens:     [1, 2, 3, 4, 5]
generated tokens: [6, 7]

second-request prefill route:
token 3 -> [30, 31]  (now real)

At this point, seen_token_ids == [1, 2, 3]. The existing suffix slicing correctly starts at index 3 and adds only the new prompt tokens [4, 5], but that also means token 3 is not appended again. Before this fix, its already-stored route was never updated:

before: training replay keeps token 3 -> [0, 1]   (wrong dummy route)
after:  previous assistant token 3 -> [30, 31]   (real prefill route)

This is not a tokenization or sequence-offset change. The token history remains [1, 2, 3, 4, 5, 6, 7]; only the route metadata for the previous turn final token is refreshed from:

routed_experts[len(seen_token_ids) - 1]

The stale dummy route is harmful even though it belongs to a historical token: R3 forces training replay through those dummy experts, changing that token's hidden state and therefore the logprobs of later tokens that attend to it.

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

[zpqiu]

/ok to test e3ba8c2

[zpqiu]

/ok to test 09a2e40

[zpqiu]

/ok to test 09a2e40

[zpqiu]

/ok to test 09a2e40