What changes when prefix saves and restores go to TransferQueue in batches (current head)

NVIDIA-NeMo/RL #4508, part 9 of 11 in the partial-checkpointing stack, on top of #4612. Permalinks pin the PR head 6c72a0f; "before" links pin the base f47c6b4.

Before
At a checkpoint, each unfinished model call saved its tokens with its own TransferQueue write. After a restart, each resumed request read its own saved rows on its own thread.
This PR
Groups the saves into a few large writes, and merges the reads of requests that resume together into one shared read. Two config keys set the batch limits.
Status
Open, round 2 of review. All 7 round-1 comments are fixed. The microbenchmarks show a 23× faster save write at 1k tokens and a 6–13× faster mass resume. The end-to-end save time is still pending, and round 2 asks for size cuts (C8–C12).

Four pages on this PR. The round-1 pages describe an older head. This one: how the batched path differs from the old one, today.

Background. When RL saves a checkpoint in the middle of generation, each unfinished model call writes its tokens so far to TransferQueue (TQ), one row per call. After a restart, Gym sends the call back with the keys of its saved rows, and the call continues from them instead of starting over. Dotted words have a short explanation: hover or tap them.

Before, the save loop wrote one call at a time (_checkpoint_active_generation_cut → stage_generation_prefix). Now _checkpoint_generation_cut_batched freezes calls one by one, keeps each locked, and writes up to 256 rows in one stage_generation_prefix_batch call. On restore, the new PrefixReadBatcher waits 2 ms, merges the keys of every request that asked, and runs one get.

The two paths, on one small example

Example: 3 unfinished calls (c0, c1, c2) and a batch limit of 2 rows. On restore, c1 has 2 saved rows (one from an earlier checkpoint), so the 3 requests need 4 keys.

At a checkpoint (save) After a restart (restore) Before This PR c0 · 40 tok c1 · 700 tok c2 · 25 tok put [c0] put [c1] put [c2] 3 puts stage_generation_prefix each call is locked only for its own write c0 · 40 tok c1 · 700 tok c2 · 25 tok put [c0, c1] put [c2] 2 puts stage_generation_prefix_batch dashed = calls held locked until their group's put returns r0 · 1 key r1 · 2 keys r2 · 1 key get [k0] get [k1, k2] get [k3] 3 gets _resolve_usable_generation_cut each request reads its own rows on its own thread r0 · 1 key r1 · 2 keys r2 · 1 key wait 2 ms get [k0,k1,k2,k3] 1 get PrefixReadBatcher rows go back to each request, then into _resolve_usable_generation_cut(…, chunks) if the shared get fails, each request reads its own keys again

Where the two paths meet, and what still differs

Both write paths build a row with _encode_record, which is the old single-row body moved, not new code. The per-call freeze, ack and seal logic is the old body too, now the generator _generation_cut_transaction, so a batch can keep several calls frozen at once. Both restore paths end in _resolve_usable_generation_cut, which checks the prefix digest and lets the call regenerate if anything is wrong. The new path only passes in rows it already read.

BeforeThis PR
Writes per saveone per callone per group of ≤ 256 rows and ≤ 4,194,304 tokens
A failed write fails…its own callits own call: a failed group put is retried one row at a time
How long a call stays lockedits own writeits group's whole write. A call that finishes meanwhile waits for its lock on an event-loop worker thread (asked in the review body)
Reads at restoreone per request, side by sideone shared get at a time per worker
A failed shared readn/are-read key by key inside the batcher, then again per request. Round 2 asks to keep only the second (C10)

C8 — generation_cut_capture.py:1299 C10 — prefix_read_batcher.py:95

So what. The design reuses the old code well: the row encoding and the per-call logic are moved, not rewritten, and the restore still goes through the same check-or-regenerate step. What remains is size. The write side keeps a two-step handshake that only made sense for the old all-or-nothing batch, and the read side retries a failed read twice. The cuts are on the what can shrink page.