Read the routes once per cut, not on every token: about a quarter of the source can go

Size review of NVIDIA-NeMo/RL #4612. Line counts are from git diff 0bf5789 2548c67 (the PR against its real base, #4583). Permalinks pin the head 2548c679. Review comments are in a pending review; their links work once it is published.

Before
The prefix-recovery feature is a stack of open PRs, over 20,000 added lines. This PR adds about 1,080 of them: ~610 source, ~470 test.
This PR
Copies new route rows out of vLLM on every generated token of every request, so that a checkpoint, which is rare, has them ready.
Status
Proposed — not implemented. Each cut below is a review comment; none is applied. The diff on this page has not been run.

Three pages on this PR — this one: how big it is, and what can go.

vLLM 0.29 keeps each running request's routes in a private list, routed_experts_chunks, and only hands them out when the request ends. That list is only ever appended to. Rows already in it never change. So the rows a cut needs are simply the first N rows of that list at the moment of the cut. The PR instead copies every new piece as it arrives, carries it through the HTTP layer, and keeps it in buffers that a cut swaps, saves, or rolls back.

Two ways to get the same rows

This PR: on every token of every request, once route capture is on next_output copy the new pieces of vLLM's private route list +1 row [1, L, K] per token 3 extra fields on vLLM's output object read back in the HTTP layer count, pieces, error _observe_routed_experts checks a piece counter, skips old rows, keeps a list counters + lists per call freeze / seal / rollback swap the route lists next to the token lists at a cut: join the pieces runs ~every token · logs an ERROR on every token if a vLLM upgrade renames a private field Proposed: only when a cut is taken freeze_for_checkpoint freeze the tokens, as today tokens [12, 13] → N = 3 rows (prompt 10, 11 + 12; 13 has no row yet) One hop to vLLM's loop (timeout) cat(routed_experts_chunks)[:N] → [1, 2, 3] shape [3, L, K] finished already? use its final routes Save, as today rows [1, 2, 3, ph] next to the tokens resume join and final merge stay exactly as in the PR runs once per cut · nothing added to the token path · route fields on each call: 8 → 3

Same tiny example as page 1: prompt 10, 11, tokens 12, 13 at the first cut. L = MoE layers, K = experts per token. Both ways give the rows [1, 2, 3]. Two reviewers checked the vLLM source: the scheduler never re-sends rows for a position, and a finished request's list is dropped in the same step as its final output, which is why a cut that races the end must use the final routes.

What can go

What the PR writesReplace withsource lines
Per-token route copies, the 3 extra output fields, the piece counter, route buffers in freeze/seal/rollback, the unused public snapshot_new_routed_expert_chunks

review comment — vllm_worker_async.py:391 review comment — vllm_worker_async.py:449

One read at cut time (above). Removes ~190, adds ~40.~150
A third copy of the "which int type is this route string" lookup (:1442)

review comment — vllm_worker_async.py:1442

Let decode_routed_experts keep the type written in the string~7
The placeholder row, built twice (:231, :622). The two must match byte for byte or every resume fails its check.

review comment — generation_cut_capture.py:231

One helper used in both places~6
Checks that cannot fire: two "changed model dimensions" checks, one length check, and a fetch branch no source reaches (:558)

review comment — generation_cut_capture.py:183 review comment — generation_cut_capture.py:558

Delete~14
Tests: the two large router-replay tests repeat ~17-line fake vLLM outputs and inventories by hand

review comment — test_vllm_token_capture_hosting.py:1376

A small _observe helper, the file's existing _single_generation_cut_inventory, dataclasses.replace~125 test

Not worth changing: the resume join and the final merge. A reviewer tried to reuse the existing route assembler, execute_route_plan, there. The saved pieces have a different row layout, so it would need a new mode in shared training code.

So what. About 160 of ~610 source lines (~26%) and ~125 test lines can go. That leaves roughly 450 source and 300 test lines, including the ~45 lines of tests the review asks for. The per-token copy is also the part that breaks quietly on a vLLM upgrade. Asking now matters because #4508 will be rebased on top of these files.