What changes when prefix saves and restores go to TransferQueue in batches
Explains NVIDIA-NeMo/RL #4508 (stacked on the prefix-recovery PRs, real base #4612). All permalinks pin the PR head 5498a0a; "before" links pin the base 2548c67. The review comments are in a pending review; their links work once it is published.
Before
At a checkpoint, each unfinished model call saved its tokens with its own TransferQueue write. After a restore, each resumed request read its own saved rows, on its own thread.
This PR
Groups the saves into a few large writes (up to 256 rows each). It also merges the reads of requests that resume at the same time into one shared read. Two new config keys set the batch limits.
Status
Proposed: the speed-up is not measured yet. The PR has no timing numbers. The review asks for two, and for a fix to one shared-read bug (below).
Four pages on this PR. This round-1 page describes an older head (5498a0a); see round 2 for the current one. This one: how the batched path differs from the old one.
Background. When RL saves a checkpoint in the middle of generation, every unfinished model call has its tokens so far written to TransferQueue (TQ), one row per call. After a restart, the call resumes from that row instead of starting over. Dotted words have a short explanation: hover or tap them.
Before this PR the save loop called sink.stage_generation_prefix once per call.
Now _checkpoint_generation_cut_batched collects rows and sends each group with one stage_generation_prefix_batch call.
Each call stays locked until the whole group's write is back, so no row is marked saved until every row in its group checks out.
On restore, the new PrefixReadBatcher waits 2 ms, merges the keys of every request that asked, and runs one read.
The two paths, on the same small example
Where the two paths meet, and what still differs
They share the per-call code. The old body that freezes one call's tokens, writes them, and checks the reply is now the generator _generation_cut_transaction. It hands its row out (L1148) and waits for the write result. The batched loop drives many of these at once. Reads meet in _fetch_generation_cut_rows, which both the old one-request read and the new shared read call.
What still differs
Before
This PR
Writes per save
one per call
one per group of ≤ 256 rows / ≤ 4,194,304 tokens
How long a call stays locked
its own write
the whole group's write. A call that errors meanwhile blocks the server's event loop until the write is back (open on #4506)
So what. The write side is a fair shape, and its rollback and deadline handling checked out in review. Whether the PR is worth its 2,802 lines depends on two numbers it does not report yet: save time before and after with hundreds of unfinished calls, and restore time with many requests resuming at once. The shared read also needs one fix, so that a bad row fails only the request that asked for it.