How a model call with an image is saved mid-generation and resumed, next to a text-only call

NVIDIA-NeMo/RL #4583, stacked on #4506 (text prefix recovery). All permalinks pin the PR head 0bf5789.

Before
A checkpoint could save a half-finished text generation and resume it later. Setup refused to turn this on when the run captured images or video.
This PR
Removes that refusal for Omni image and video. The first saved row for a call also stores the call's image metadata and pixels; on resume the image is processed again and checked against them.
Status
New feature. The review found no bug. It asks for a smaller design (see “size” page) and for tests of the reject paths.

Two pages on this PR, one question each. This one: what changes when the cut call has an image.

Background. In a long rollout, a checkpoint can land while vLLM is still generating a model call. Instead of throwing those tokens away, the parent PR saves them as prefix rows, one per checkpoint. After a restart, Gym sends the same request again with the list of saved rows, and vLLM continues from the saved tokens. This PR makes that work when the prompt has an image. Dotted words have a short explanation: hover or tap them.

Both kinds of call share the same save and resume code. The only extra work for an image call is at three points: the first saved row stores the image's place in the prompt and its pixels (generation_cut_capture.py:717-726); resume reads those pixels back in a second read (fetch_generation_cut); and resume processes the image again and requires the same prompt tokens and the same pixels (vllm_worker_async.py:1753-1784, :1149-1169). Any mismatch rejects the restore, and the call is generated again from the start of the turn.

One call, cut twice, then resumed

text call image call (18 = image placeholder token) same in both lanes only the image call 1 · cut #1 writes row 0 tokens [10,12,11, 41,42] extras: none tokens [10,18,18,11, 41,42] + media_spans [{1..3}] + imgs f32[3,2,3] generation_cut_capture.py:717-726 2 · cut #2 writes row 1 tokens [43] extras: none tokens [43] extras: none (no pixels again) 3 · restore read rows from TQ 1 read: rows 0,1 → tokens check digest (same code) read 1: rows 0,1 → tokens + media_spans read 2: imgs f32[3,2,3] tq_token_sink.py:914 fetch_generation_cut generation_cut_capture.py:321 4 · check before continuing nothing more: the saved tokens are the prompt, no re-render process image again → [10,18,18,11] must equal saved prompt tokens new pixels must equal saved imgs vllm_worker_async.py:1149-1169 :1753-1777 both lanes meet here: engine prompt = saved tokens vllm_worker_async.py:1751 5 · decode the rest prefill [10,12,11,41,42,43] → [44,45] prefill [10,18,18,11,41,42,43] → [44,45] 6 · final row for training all tokens, then prefix rows cleared all tokens + media_spans + imgs f32[3,2,3] written again vllm_worker_async.py:1361-1365

Token ids and the pixel shape come from the PR's unit test test_generation_prefix_fetch_restores_and_verifies_media; the generated ids 41–45 are illustrative.

Where the two paths meet, and what still differs

They share one function for the hard part: _rebuild_generation_cut_snapshot joins the rows and checks their digest for both. The digest now also covers the image metadata (:383). After the extra checks in step 4, both paths use the saved tokens as the engine prompt (vllm_worker_async.py:1751). Only two things still differ after that point: vLLM needs the original image again to prefill, because the saved pixels are not in the form vLLM takes as input, and the final row carries the pixels.

So what. The design reuses the text path well. Only three steps are new, and every mismatch fails closed. The cost is in step 1: each saved image call writes its pixels to TQ one extra time at every checkpoint, under the checkpoint's time limit. Those saved pixels are only ever compared in step 4, never used, and the final row writes them again. The second page shows how a hash would replace them.

🔵 review comment — vllm_worker_async.py:1169 review comment — vllm_worker_async.py:1184 review comment — tq_token_sink.py:927