[Rockdu]

What

With --use-output-store, the rollout process no longer reads a training rollout's decoded output. The reward actors read it from the store.

• Reward pools: every AsyncRewardActorPool wraps its actor class with _StoredOutputRewardActor. The wrapper reads the stored outputs it is given, then calls the scorer's score_batch. The reward actor classes keep their scoring code.
• PickScore: a stored video reaches the actor whole, and the actor picks --pickscore-num-frames of its frames. They are the frames the rollout process picks for an inline output.
• After the reward, release_generated_outputs deletes the samples' output fields from the store and sets generated_output to None.
• W&B media: on a rollout that logs media, the first --wandb-log-num-images scored samples are read back as tensors before the delete. This is the only read left in the rollout process. A rollout that logs no media reads nothing there.
• Reads through registered memory: RegisteredTensorReader reads stored outputs into one buffer that stays registered with the store. It registers the buffer when first needed and again only to grow. OutputStore.read_many registers fresh memory on every call.
• Removed: load_generated_outputs and its store_read dashboard stage from feat(rollout): read the generated output from the output store right after generation #376.
• Default segment: --output-store-global-segment-size goes from 16gb to 72gb, so a rollout's outputs fit while they wait for their reward. The krea2 async recipe drops its 64gb override.
• Default transport: --output-store-protocol goes from tcp to rdma. A node without an RDMA device passes tcp.

Why

#376 reads every output in the rollout manager right after generation. Every output then passes through that one process twice: store to manager, then manager to reward actor over Ray. The manager also holds a whole rollout's outputs until the rollout is converted (about 17 GB for LTX-2.3).

Now an output goes engine, store, reward actor, and the rollout process holds only the media samples.

Contract changes

• In the rollout process, sample.generated_output is a StoredTensor until the reward and None after it, except the kept media samples. Hooks that run after scoring see that.
• A --custom-rm-path function cannot read the tensor itself. It scores through a reward pool: a built-in *_rm, --custom-api-rm-config, or its own AsyncRewardActorPool(args=args, ...). weighted_mixture_rm already does.
• --group-rm rewards get refs the same way; the outputs are released after the group reward.
• A custom generate function cannot read the output.
• AsyncRewardActorPool takes args.
• A rollout logs at most --wandb-log-num-images samples: a group dropped by a filter can use up a slot.

Store capacity

The slot is released after generation, so outputs wait in the store until they are scored, with no cap.

• Measured on LTX-2.3 (64 videos a rollout, 269 MB each, one PickScore actor): generation finishes about 2.3 min into the rollout and scoring takes 7.5 min, so the store peaks at 16.0 GB, nearly the whole rollout's outputs.
• Default flip: every store-on run now lends 72gb of host memory per training node; under RDMA it is pinned. The old 16gb was just too small for LTX-2.3: 51 engine puts failed on a full store and were regenerated, with metrics still bitwise equal.
• Sizing rule (in the --output-store-global-segment-size help): the segments should hold one rollout's outputs plus two rollouts of latents.
• When the store is full: the engine's put fails and the request is retried, which regenerates it. A reward that stays slower than generation with a full store fails the rollout after the retries run out.

Transport and the read buffer

Measured on one H200 node with a 1.23 GB object, the size of one MiniMax-H3 output, in a segment lent by another process:

	put	read_many (registers per call)	read into registered memory
tcp	3.35 s	0.35 s	0.36 to 0.64 s
rdma	0.17 s	0.23 s, of which 0.18 s is the registration	0.03 to 0.05 s

In the H3 e2e, where the reward actor reads one 1.33 GB output per sample:

	per read	rollout 0 / 1
tcp, read_many	0.33 to 0.44 s	3:39 / 3:37
rdma, read_many	0.2 s when memory is free; 1 to 10 s otherwise; 21.6, 27.9 and 39.9 s seen	3:51 / 3:59
rdma, registered buffer	0.04 s transfer + 0.58 s copy	3:39 / 3:38

• Where read_many loses its time: timing its three steps in the H3 e2e shows the transfer always takes 0.03 to 0.06 s and the registration takes the rest. Each slow registration comes with direct reclaim and compaction in the kernel (for the 21.6 s read: 608 allocation stalls, 1218 compaction stalls, 77k pages migrated); fast ones have none.
• Cause: read_many allocates fresh memory per call, and an RDMA registration pins it at once. On a NUMA node whose free memory is all page cache (2 GB free of 1 TB on the test node), the kernel has to reclaim and move pages inside the registration.
• With the buffer registered once, that cost is paid at most once per reader. The rest of a read is copying the tensor out of the buffer, single-threaded in a Ray actor.
• On tcp, one whole-output read outlasted Mooncake's 10 s lease (LEASE_EXPIRED), which fails the rollout. It was seen once and not reproduced.

Not in this PR

• The trainer still reads latents with read_many. Its reads are about 150 MB, and train time was the same on both transports (305 s against 304 s).
• Reading only the frames PickScore scores. It would cut a read from 269 MB to 14 MB on LTX-2.3 and shrink the actor's buffer, but rollout time does not move: reads overlap generation.
• Returning views of the buffer instead of copies. It would save the 0.58 s copy per 1.33 GB read and needs a release step.
• CLI reference rows for the output store flags. The page has none today, for any of them.
• A warning when a rollout's stored bytes will not fit the segment, and a metric for retried rollout requests.

Validation

• tests/fast/rollout/test_generated_output_store.py, 13 test functions, numbered in its header diagram:
  • The store read equals the inline parse; the reader registers once across reads and again only for a larger read; 200 reads from 8 threads each get their own tensor.
  • The wrapper reads refs and passes inline (eval) tensors through.
  • Pools wrap their actors only with the store on; the PickScore pool passes its frame count to the actor.
  • pickscore_rm hands a stored video over unsliced, and the actor scores a whole video as the frames picked inline, picking before the uint8 quantisation.
  • Every call site (microgroup, custom generate, group reward) hands the reward refs with the slot released, reads nothing in the rollout process, and releases the outputs.
  • Release deletes only output fields; concurrent microgroups keep at most N for media; a failed media read or delete does not fail the rollout.
• Mutations (no release, no release after a group reward, media counter not decremented, media filter letting refs through, PickScore actor without its frame count, registering on every read, no lock) each fail at least one test.
• tests/fast/rollout + tests/fast/utils + tests/fast/dashboard: 421 passed on the current head, and on each of its three commits.
• CI, on test(e2e): run every RL e2e with the output store on, generated output included (CI only, do not merge) #377 (this PR's three commits plus ci: attach the GPU job container to an ipvlan network on the node's RoCE bond, so Mooncake's rdma transport can connect #384's CI network change and two CI-only commits): all nine e2e tests pass against their standards. The eight RL tests run with the store on and every Mooncake client on the rdma transport; no put failed and no lease expired.
• This PR needs ci: attach the GPU job container to an ipvlan network on the node's RoCE bond, so Mooncake's rdma transport can connect #384 first. The CI node's RDMA is RoCE, which a job container on the runner's bridge network cannot connect over, so with rdma as the default every e2e test that uses the store fails without ci: attach the GPU job container to an ipvlan network on the node's RoCE bond, so Mooncake's rdma transport can connect #384.
• End to end on 4x H200 with sglang-miles head, store on over rdma, compared bitwise with the e2e standards. These ran on the head before the final renames and comment edits, which change no behaviour:
  • SD3.5 NFT and H3 t2va GRPO: pass.
  • LTX-2.3: bitwise identical to unmodified main on the same machine (the LTX standard does not reproduce there).
  • A weighted_mixture_rm run (SD3.5 OCR + PickScore through --custom-rm-path) succeeds.
  • No put failed and no lease expired in any of the four.
• Train/rollout log-prob mean abs diff: 1e-6 to 4e-5.

Files

• miles/utils/output_store.py: RegisteredTensorReader.
• miles/rollout/rm_hub/core.py: _StoredOutputRewardActor; the pool takes args.
• miles/rollout/rm_hub/{pickscore,hps,hpsv3,ocr,api}.py: pass args; the PickScore actor takes its frame count, and pickscore_rm hands a stored video over unsliced.
• miles/rollout/sglang_diffusion_rollout.py: release_generated_outputs after both reward sites; the media counter; load_generated_outputs removed.
• miles/utils/diffusion_rollout_response.py: read_generated_outputs reads through the reader.
• miles/utils/tracking_utils.py, miles/ray/rollout.py: logs_rollout_media; media logging takes the samples that hold a tensor.
• miles/dashboard/events.py: the store_read kinds removed.
• miles/utils/arguments.py: the --use-output-store, segment size and protocol help; the 72gb and rdma defaults.
• scripts/run_diffusion_grpo_krea2_async.py: no segment override.
• miles/utils/types.py, docs/user-guide/customization.md, docs/advanced/streaming-reward.md: comment and docs for the contract above.
• Tests: test_generated_output_store.py rewritten; four test fakes get the fields the pool, the actor and the rollout now read; the default assertion in test_output_store_args.py.

Checklist

• pre-commit passes on the changed files
• Added/updated tests for new behaviour
• pytest is green (tests/fast/rollout, tests/fast/utils, tests/fast/dashboard)
• Two defaults changed (--output-store-global-segment-size, --output-store-protocol)
• GPU e2e: CI on test(e2e): run every RL e2e with the output store on, generated output included (CI only, do not merge) #377 (nine e2e tests, the eight RL ones with the store on over rdma), and SD3.5, H3, LTX-2.3 and a custom-reward run on a devbox