[macandro96]

What does this PR do ?

Batches the TransferQueue I/O behind generation-prefix recovery: a checkpoint cut stages all of a worker's active prefixes in a few bounded PUTs instead of one PUT per call, and resumed requests share coalesced GETs for their restored prefixes.

On 8,192 rows across 32 TQ clients, a cut's writes finish 23× faster at 1k-token prefixes (10.1 s → 0.44 s) and 2.7× faster at 16k (10.2 s → 3.75 s). A mass resume of 8,192 restores across 128 workers finishes 6–13× faster, with p50 restore latency down from 6.9 s to 1.3 s. See Benchmarks.

Stacked on #4612 (router-replay prefix recovery). The diff shows only this PR's commits on top of the prefix-recovery stack.

Issues

How it works

flowchart LR
    subgraph Cut["Checkpoint cut (per vLLM DP leader)"]
        I[Gym inventory<br/>active calls] --> F[Freeze each call<br/>lifecycle lock held]
        F --> B[Group rows into batches<br/>≤ batch_size rows, ≤ max_tokens]
        B --> P[One TQ PUT per batch]
        P --> V[Validate every ack] --> S[Seal all calls]
    end
    subgraph Restore["Resume after restore"]
        R1[request A] --> Q[PrefixReadBatcher]
        R2[request B] --> Q
        Q --> G[Coalesced TQ GET<br/>≤ batch_size keys, deduplicated,<br/>chain order kept]
    end

Loading

• Writes (generation_cut_capture.py, tq_token_sink.py): each active call becomes a generator transaction that freezes its buffer under its lifecycle lock, yields its prefix record, and seals only after every acknowledgement in the batch validates. Every cut takes this path; generation_prefix_batch_size: 1 writes each call in its own PUT. stage_generation_prefix_batch groups rows by field dtype/shape, writes one put_many per group, and returns the written keys; on a failed put it discards every key of the batch, since earlier groups may already be written, and raises. A row larger than the token limit is written alone and never split.

• Failure handling: as in the base stack, failures are per call: a call that cannot be cut is acknowledged as durable_failure and counted under the base's failure reasons (cut_error, or control_deadline past the deadline), its frozen cut rolls back, and it regenerates after a restore, while the other calls keep their cuts. A failed batched PUT is retried one row per PUT, so one bad row fails only its own call. A failed acknowledgement or seal fails only its call, whose transaction clears its own row. Rows written but never handed to a call, which only an interruption can leave, are cleared too, so no unsealed key is ever advertised to Gym.

• Cut deadline: the base stack bounds a cut by Gym's control timeout. Past the deadline no new TQ write starts. Rows frozen for the next batch roll back, and they and every call not yet reached are acknowledged as durable_failure; they regenerate after restore.

• Reads (prefix_read_batcher.py, vllm_worker_async.py): PrefixReadBatcher is an event-loop-local coalescer for restore-time GETs. Concurrent requests share GETs of at most generation_prefix_batch_size distinct keys, with keys deduplicated and each request's chain order preserved. Cancelling one request does not cancel the shared GET for others. If a multi-key GET fails for any reason, each key is re-read alone once, so a bad row fails only the requests that asked for it, and an outage costs N+1 GETs. The rows go to the base stack's _resolve_usable_generation_cut, which regenerates the call if its prefix still cannot be used.

• Config: two new knobs under rollout_recovery. Their only defaults are the RolloutRecoveryConfig fields; setup_single_controller passes both to GenerationInterface.setup_token_capture and on to the vLLM workers. Megatron generation accepts and ignores them.

  rollout_recovery:
    generation_prefix_batch_size: 256         # rows per cut-time PUT and keys per restore GET
    generation_prefix_batch_max_tokens: 4194304  # token rows per cut-time PUT

  Both are runtime knobs, excluded from the rollout-checkpoint bootstrap fingerprint, so changing them does not invalidate an existing checkpoint.

Benchmarks

One 128-CPU node, TQ simple backend, synthetic rows (no routes or media). The tools call the production sink and PrefixReadBatcher directly; they are kept off this PR on branch amahishi/ckpt-v2-prefix-io-benchmarks, with the exact commands in its commit message.

Cut-time writes

8,192 rows from 32 process-isolated TQ clients sharing one deployment. Batch size 1 is one PUT per call, as before this PR.

prefix tokens	rows/PUT	PUT calls	wall (s)	rows/s	mean row wait (ms)
1,024	1	8,192	10.08	813	36.8
1,024	256	32	0.44	18,471	188.4
16,384	1	8,192	10.20	803	26.1
16,384	256	32	3.75	2,184	561.1

A row waits longer for its own PUT, but a cut returns its receipt only after every row is written, so wall time is what the cut deadline sees.

Restore reads

"per-request" is the old path: each restore reads its keys in its own asyncio.to_thread GET. "coalesced" is PrefixReadBatcher (256 keys/GET, 2 ms coalescing window). Each mode ran twice, in alternating order; the table shows the slower run.

scenario	mode	GETs	keys/GET	wall (s)	request p50 (ms)	request p99 (ms)
mass resume: 128 workers × 64 restores, 16k tokens	per-request	8,192	1.0	12.57	6,877	9,731
	coalesced	128	64.0	1.80	1,334	1,742
mass resume, 2k tokens	per-request	8,192	1.0	12.02	6,719	9,240
	coalesced	128	64.0	0.86	535	823
mass resume, 4 cut chunks per restore	per-request	8,192	4.0	30.02	17,253	27,205
	coalesced	155	211.4	2.33	1,465	2,203
128 workers, a restore every 5 ms per worker	per-request	8,192	1.0	12.38	6,625	9,235
	coalesced	300	27.3	2.19	1,735	2,093
128 workers, a restore every 20 ms per worker	per-request	8,192	1.0	11.89	6,465	8,016
	coalesced	488	16.8	2.41	1,181	1,663
light: 8 workers, a restore every 10 ms	per-request	512	1.0	2.54	363	596
	coalesced	96	5.3	0.74	98	161
light: 8 workers, a restore every 50 ms	per-request	512	1.0	4.52	17.3	24.2
	coalesced	512	1.0	3.19	27.4	34.2

• Mass resume is the production pattern: after a restart, every cut call on a worker resumes at once. Coalescing cuts the GET count by 64–98× and wall time by 6–13×.
• Light load is the worst case. At one restore per 50 ms per worker there is nothing to coalesce. GET time is the same in both modes (25.0 vs 24.9 ms in the second repeat), so a restore pays the 2 ms coalescing window plus noise. By one restore per 10 ms per worker, coalescing is already 3.7× faster.
• Head-of-line blocking: the batcher issues one GET at a time per worker. A GET that hangs now stalls every restore on that worker, where before it stalled only its own request. None of the runs above triggered this.
• Caveats: all clients ran on one node, so the per-request numbers include client-side contention that spreading workers over nodes would reduce; the TQ storage units are shared either way. The read tool reads through TQTokenSource.fetch rather than fetch_for_finalization. A restore's staging_chain reads through ChainPrefixCache are not coalesced in either mode.

Changes

Area	Files
Batched cut writes + deadline	nemo_rl/models/generation/generation_cut_capture.py, nemo_rl/data_plane/tq_token_sink.py
Coalesced restore reads	nemo_rl/models/generation/prefix_read_batcher.py (new), nemo_rl/models/generation/vllm/vllm_worker_async.py
Config plumbing	single_controller_utils/{config,setup,rollout_checkpoint}.py, generation/{interfaces,vllm/vllm_generation,megatron/megatron_generation}.py, example SC configs
Tests	test_prefix_read_batcher.py (new); test_vllm_token_capture_hosting.py, test_megatron_token_capture_hosting.py, test_interfaces.py, test_tq_token_sink.py, test_captured_media.py, test_setup.py, test_rollout_checkpoint.py

Lines: +1758 / −307 (product +767 / −270, tests +991 / −37).

Testing

On this branch (6c72a0f19), with Gym at the pinned 5911d628:

• Unit, Gym lane (--nemo-gym-only):
  • test_vllm_token_capture_hosting.py + test_megatron_token_capture_hosting.py + test_prefix_read_batcher.py + test_tq_token_sink.py + test_captured_media.py: 407 passed.
  • Batched-path coverage includes 1000-prefix batching with a cached receipt, ragged and oversized rows under the token limit, a failed batch PUT retried row by row, a call whose rows always fail leaving the other calls' cuts intact, a failed acknowledgement failing only its call, an interruption clearing rows no call received, a later batch's failure keeping earlier cuts, and the cut deadline passing both between batches and while rows are pending. The retry, interruption-cleanup and deadline tests fail if that code is removed.
  • Read-side coverage: a bad row fails only its own caller, a failed batch re-reads each key alone exactly once, and a restore uses the chunks the batcher fetched instead of reading them again.
• Unit, default lane: test_interfaces.py + tests/unit/single_controller/test_setup.py + test_rollout_checkpoint.py: 381 passed.
• Functional: tests/functional/grpo_async_gym_single_controller_prefix_recovery.sh passes locally on perf(checkpoint): batch obsolete prefix cleanup in background #4655's head (fd86c61e7), which is this branch plus background prefix cleanup, run with that cleanup enabled: 8/8 rows match the uninterrupted run, no restored row whose phase-1 prefix differed from the reference, 8 restored prefixes (148–229 tokens), 99.09% bit-identical generation logprobs over 9,312 tokens (max |diff| 0.18), invalid_row_rate and capture_poisoned_rollouts both 0, and no failed cut, fallback regeneration or failed clear in the log.

Additional Information

• Pending: timing/rollout_checkpoint/total_save_seconds before vs after, from a run with hundreds of active calls per DP leader. The microbenchmarks above isolate TQ I/O; that run will show the end-to-end save time the training loop loses.
• During phase 2 of the functional run the TQ controller logged one non-fatal dtype mismatch: existing=torch.int64, incoming=torch.float32 for partition rollout_staging right after the native TQ checkpoint restore. The batched writer groups rows by field dtype, and all checks passed; I have not yet confirmed whether the same line appears on the base stack.
• Pre-existing on main, not from this stack: test_single_controller_actor.py::test_opd_advantage_stage_pools_raw_gap_separately_under_tropd fails with AttributeError: ... '_train_data_dump' (reproduces on main at 19b245a8d).

🤖 Generated with Claude Code

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

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 079933e (PR #4508 from amahishi/ckpt-v2-prefix-batched-io)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 7eb9a01 (PR #4508 from amahishi/ckpt-v2-prefix-batched-io)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: a82e2a2 (PR #4508 from amahishi/ckpt-v2-prefix-batched-io)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[terrykong]

Thanks for this. The write-side batching has a good shape: pulling out _encode_record so single-row and batch writes share one encoder, and the freeze → yield → seal generator, are both the right moves. These hold up: partial-PUT rollback, ack-before-seal, the deadline release, the oversized row written alone, lock order, dtype grouping, and the fingerprint exclusion (a mutation test confirms the test catches it). All the PR's CPU tests pass.

Before merge

• Rebase: GitHub shows conflicts with the base branch.
• The description says "Stacked on #4507", but the base is #4612's branch (amahishi/r3-prefix-recovery). Please update it.
• This is a perf PR with no numbers. Please add timing/rollout_checkpoint/total_save_seconds before vs after to the description, from a run with hundreds of active calls per DP leader. It is already logged, and dispatch is blocked during the save, so it is the time the run loses.
• Please close #4389 as superseded. It is the older copy of this PR, and its tools/benchmark_tq_prefix_writes.py is byte-for-byte the same.

Size. This PR is +2,802 lines, on a feature that is already over 20k lines across the stack. With the inline asks it drops to about +1,000 (≈ +500 product):

• one cut path (comment at L1399): ≈ −75 product, −30 to −40 test
• no copied config defaults or limits class (comment at L647): ≈ −30
• benchmark tools out (comment on the writes tool): −1,235
• no per-row results in the batch write (comment at L741): ≈ −20 to −30
• read side to its own PR (comment on PrefixReadBatcher): ≈ −440

Smaller and optional: put could call put_many (−8), and a shared "rolled back" assert helper (first copy) plus pytest-asyncio in the batcher tests would save ≈ 23.

Open #4506 threads that also apply here

• r4262413829 (abort on the event loop waits on lifecycle_lock): please fix it in #4506; this PR then needs no change. Until then, batching widens the stall from one call to up to 256 calls held across one put.
• r4262414074 / r4262414095 (call on self, not through the class, for the test fake): this PR adds 3 more class-qualified calls (generation_cut_capture.py L577, vllm_worker_async.py L903, L964) and grows the fake's name list from 21 to 24 (test L649-L652).
• r4262413947 (simplify the frozen/sealed buffer): this PR adds one more rollback (L1096) and widens the rollback except to BaseException (L1188).

Question, no change needed. Cut batches count tokens only (L1340). With MoE router replay, each token also carries layers × top-k int16 routes: about 3.2 GB at 48 layers / top-8 for a full 4M-token batch (arithmetic only; each chunk holds only the tokens since the last cut). What is the largest single put you have seen with routes on?

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4508-batched-prefix-io.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4508-size-and-reuse.html

Generated by Claude Code

[macandro96]

Thanks for the thorough review. Status of the summary items:

Before merge

• Rebase: done. perf(checkpoint): batch generation prefix TQ I/O #4508 is now 4 commits on feat(checkpoint): recover router replay across prefixes #4612's head (def6e59f0b), +1,557 / −264.
• Description: now says it's stacked on feat(checkpoint): recover router replay across prefixes #4612, and documents the single cut path, row-only read limits and the new failure handling.
• Numbers: the description now has TQ-level benchmarks for both sides. Cut writes are 23× faster at 1k-token prefixes and 2.7× at 16k. A mass resume's restore reads are 6–13× faster. At light load, where there's nothing to coalesce, a restore pays the 2 ms coalescing window. timing/rollout_checkpoint/total_save_seconds before vs after, from a run with hundreds of active calls per DP leader, is still pending; I'll add it to the description when it's done.
• perf(checkpoint): batch generation-prefix TQ writes #4389: closed as superseded.

Size: the inline asks are in, except moving the read side to its own PR. I kept the batcher here because the mass-resume numbers back it, and I noted the one-GET-at-a-time head-of-line trade-off in the description. The benchmark tools now live on amahishi/ckpt-v2-prefix-io-benchmarks.

#4506 threads: all three are fixed in #4506, which this stack includes:

• The event-loop abort no longer waits on lifecycle_lock (blocking=False).
• The test worker is a real worker built with object.__new__. The last two class-qualified _fetch_generation_cut_rows calls in this PR now go through self.
• Request state is one token list plus staged_token_count. This PR's transaction uses the same seal_frozen_cut / rollback_frozen_cut.

Route bytes: I haven't measured the largest put with routes on yet. Your bound is right: the cut limit counts token rows only, and a full 4M-token batch would carry about 3.2 GB of routes at 48 layers / top-8. The realistic worst case is a worker's first cut, since each call's first chunk includes its prompt. I'll measure it on an R3 run and report back.

[macandro96]

/ok to test def6e59