Who owns Mooncake memory — and so who every checkpoint save has to wait on

Explainer for NVIDIA-NeMo/RL #4465 (faster Mooncake CPU-RDMA checkpoints). All permalinks pin the PR head 7847eeb.

Before
Every Mooncake process held data in its own RAM, including the vLLM serving workers. A save had to call all of them, and the vLLM ones answered late because they were busy serving rollouts.
This PR
vLLM workers stop holding data. An opt-in mode moves all data into small CPU-only "storage unit" actors, so no GPU process is called during a save.
Status
Works as described for the SingleController path. Open review items on a few combinations — see the last section.

Background. The data plane moves rollout rows between processes. With the mooncake_cpu backend, every process is a Mooncake client (it puts and gets rows over RDMA), and most are also an owner: a block of their RAM, the segment, holds rows. A checkpoint save writes each owner's rows to disk, so the save must call every owner and waits for the slowest. Dotted words have a short explanation: hover or tap them.

A process takes part in a save only if its segment is bigger than 0 (tq_mooncake_checkpoint.py:1390). So "who is on the save path" is just "who has a segment". The PR changes that in two steps. vLLM workers now attach with segment_size=0 (vllm_worker_async.py:559). And when mooncake_cpu.storage_unit_segment_size > 0, every client gets 0 (transfer_queue.py:859) and new MooncakeStorageUnit actors own all memory (mooncake_storage_unit.py:42). That is the same shape the simple backend already had with TransferQueue's SimpleStorageUnit.

The same 2-node job, three ways

owns a segment — the save calls it owns a segment, busy with rollouts client only (segment 0) — not called Before this PR every process owns 64 GiB train node T0 T1 T2 ctrl 3 trainer workers + the controller inference node V0 V1 2 vLLM workers (token capture on) save calls 6 owners waits behind rollout traffic on V0, V1 86 nodes, timing/train/checkpointing: mean 43.0 s · worst 63.9 s After — default storage_unit_segment_size: 0 train node T0 T1 T2 ctrl inference node V0 V1 segment_size=0 · :559 save calls 4 owners vLLM rows land in trainer/ctrl RAM non-owners dropped once · :1424 86 nodes (changes 1–3, earlier branch): mean 14.6 s · worst 24.4 s After — storage units on segment > 0, placement: inference, 5 units train node T0 T1 T2 ctrl inference node V0 V1 U0 U1 U2 U3 U4 vLLM puts prefer a unit on its node MooncakeStorageUnit (1 CPU each) · :42 save calls 5 CPU units only no GPU process is on the save path 86 nodes (changes 1–4, units on vLLM): mean 10.9 s · worst 13.5 s Numbers are from the PR description. The 86-node rows ran on an earlier branch of this code; this head was tested on 12 nodes. Trainer puts with units on go over RDMA to any unit. With placement "inference", trainer nodes hold no data plane memory at all.

Where the change is, in code

Where it lands: RL only. No TransferQueue or Mooncake change. Three lines decide who owns a segment; the rest is wiring: start_storage_units in SingleController setup, and the units joining the save list in single_controller.py:414.

@@ vllm_worker_async.py L550 → L559 · setup_token_capture
-        dp_client = build_data_plane_client(dp_cfg, bootstrap=False)
+        # Own no segment: a save would otherwise wait on this actor, whose event
+        # loop is busy serving rollouts. Capture writes land in other segments.
+        dp_client = build_data_plane_client(dp_cfg, bootstrap=False, segment_size=0)

@@ transfer_queue.py L812 → L859-L861 · _init_tq (the config every client attaches with)
-                    "global_segment_size": int(mooncake_cfg.global_segment_size),
+                    "global_segment_size": 0
+                    if mooncake_cfg.storage_unit_segment_size
+                    else int(mooncake_cfg.global_segment_size),

@@ tq_mooncake_checkpoint.py L1423-L1425 · configure_checkpoint_workers
         manager._checkpoint_workers = list(workers)
+        _, owners = _live_participants(manager)
+        manager._checkpoint_workers = list(owners.values())

The save code itself did not change. Both layouts run the same save: call every process with a segment (:1388-1391). The only difference is which processes have one. The simple backend gets the same node picker through simple.storage_unit_placement, which pins TransferQueue's own storage units to chosen nodes (transfer_queue.py:740).

What the review asks for

The design holds; a few combinations need a guard. These are in the review comments (the review is still pending, so its links work only for its author until it is published):