kvcache-ai / kvcache-ai/Mooncake
[RFC]: Add a VeRL Store-backed Weight-Update Backend to `mooncake-rl`
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
## Purpose
Add a Store-backed weight-update path from VeRL trainers to rollout workers, with four properties:
- **No partial updates.** Each policy update is an immutable generation, announced only when complete.
- **Restart and scale-out without republish.** The active generation stays in the Store; rollout workers can restart or scale out and load it directly.
- **Minimal Store traffic on both sides.** The trainer puts exactly **one** copy of each tensor, allocated preferentially in the publisher's node-local segment (near-zero publish network traffic). The whole rollout cluster reads each published byte from the Store exactly **once**.
- **No NCCL.** Bytes fetched from the Store are redistributed among rollout workers over Transfer Engine one-sided writes, following the point-to-point approach of fabric-lib (arXiv:2510.27656). The data plane uses only Mooncake Store and Transfer Engine.
## Summary
Trainer ranks collectively reconstruct each tensor; a deterministic, byte-balanced assignment picks one publisher, which stores it as a single-replica Store object in its node-local segment. After all writes succeed, the coordinator atomically announces the generation.
On the rollout side, all workers' demands are merged into **disjoint** source byte ranges. Each range has exactly one **reader** that fetches it via `get_into_ranges()`, then pushes it to the other consumers over Transfer Engine one-sided writes. Two consumption modes share this path:
- **Compatibility mode**: per-tensor leaders read complete tensors and push them to peers; every worker feeds VeRL's existing loader.
- **Parameter-direct mode**: readers fetch only the bytes that land in live parameters and redistribution writes them directly into parameter buffers.
## Foundations
| Foundation | Status | Role |
| --- | --- | --- |
| PR #1717 `get_into_ranges()`, PR #3000 scatter batching | Merged | Ranged reads from Store objects. |
| PR #3187 resource and weight manifests | Merged | Placement/runtime-binding contracts, reused **unchanged** — no new types or fields. |
| PR #3441 N-D logical transfer planner | Open | Source/target overlap planning; parameter-direct mode depends on it once stable. |
| Transfer Engine one-sided write, buffer registration | Merged | Carries all reader-to-consumer redistribution. |
This RFC stays narrower than RFC #3111: VeRL reconstructs one full tensor at a time, adapters describe target views, and a Store lowering maps plans to keys, ranged reads, and Transfer Engine writes. VeRL orchestration, rollout quiescence, routing, and registration lifetime live here; contracts live in #3187; generic planning lives in #3441.
## Design
### Publish: one copy, prefer-local
For each tensor, all trainer ranks reconstruct it collectively; a deterministic assignment picks one publisher, which converts it to the rollout dtype and calls `put_from()` with `ReplicateConfig{replica_num = 1, preferred_segments = [local segment], with_hard_pin = true}`. The put degenerates to an intra-node copy, so publication is near-zero network traffic. The semantics are *prefer* local, never *require* local; correctness does not depend on where an object lands.
The assignment balances **bytes** across trainer nodes, so the published model spreads roughly evenly and rollout reads later fan out across all trainer NICs instead of hot-spotting one node.
On a write failure, trainers finish the iterator to preserve collective ordering, clean up the failed generation, and keep the previous one active. Objects use deterministic model/generation/fragment keys and raw contiguous payloads; the announcement carries the source placement manifest and digest, and may carry per-object segment locations as locality hints (correctness never depends on them).
### Consume: read each byte once
**The single-pull invariant.** For each generation, every byte of every source object is read from the Store exactly once, cluster-wide, in both modes. Demands from all workers are merged per source object — take the union of demanded intervals, partition it into disjoint ranges — so overlapping demands deduplicate into one read plus multiple pushes, and bytes nobody needs are never read.
**Reader assignment.** Each range gets one reader; other ranks needing it are consumers. The assignment balances each reader's NIC load (`bytes read + bytes pushed cross-node`; intra-node pushes ride Transfer Engine local transports off-NIC), prefers a reader that is itself a consumer, and splits wide-fan-out objects across readers. It is a pure function of `(source digest, target digest, topology)`: computed once, cached, and reused every generation until the digests change.
**Compatibility mode.** Leaders are assigned per tensor, byte-balanced, so no single worker carries the push load. A leader pushes each tensor hierarchically: one Transfer Engine write per remote node into a node delegate's staging buffer, then intra-node fan-out over local transports — `num_nodes − 1` NIC sends per tensor instead of `num_workers − 1`. The `(N − 1)`-fold delivery volume is inherent here, because VeRL's loader wants every full tensor on every worker.
**Parameter-direct mode.** The redundancy disappears because the demand itself shrinks: each worker declares its exact shards through `WeightPlacementPart`, the planner computes source/target overlap, and each worker receives only bytes that land in its parameter buffers — nothing is received and discarded. Total delivery equals the bytes the target placement actually stores (its lower bound): with 4 rollout instances × TP8, Store reads stay at 1× model and delivery is ~4× (one per DP replica) instead of ~31×.
Direct loading runs in three steps. **Prepare** validates the routing table (full coverage, no overlap), objects, generation/lease fences, and registrations, touching nothing. **Apply** executes reads and pushes and verifies every returned byte count; consumers stay passive since all writes are one-sided. **Post-apply** invalidates rollout caches and advances the generation only after all-rollout consensus. Rollout drains in-flight requests before prepare; prepare failure resumes the old generation (and may re-route around a failed rank); failure after apply begins pauses rollout and rebuilds replicas — no in-place retry.
### Pipelined execution
Plans are cut into byte-bounded batches flowing through read → redistribute → scatter. A range's pushes depend only on its own read, so while batch `k` is being pushed, batch `k+1`'s read is already in flight; the only bubble is the first read. On a reader's NIC the two stages use opposite directions (reads inbound, pushes outbound), so overlapping them saturates the full-duplex link rather than contending. A staging-memory watermark bounds in-flight batches and must admit at least two (double buffering).
### Registration
One manager handles both domains — Store-client registrations (ranged-read destinations) and Transfer Engine registrations (remote-write destinations) — keyed by epoch, instance, device, base, and size. Registrations are provisional until every rank finishes prepare, then promoted to a persistent cache and reused across generations while client, device, address, and size are unchanged; reconnect invalidates the previous epoch; teardown unregisters exact bases.
### Module boundaries
- **mooncake-reshard**: #3187 contracts (unchanged) and, once landed, #3441 planning.
- **mooncake-rl / Store adapter**: generations, publication, reader assignment and routing, lowering, Transfer Engine redistribution, registration, cleanup.
- **VeRL/vLLM/SGLang adapters**: drain/resume, runtime bindings, cache invalidation, model-specific mappings.
Model-specific semantics stay outside Mooncake core; no second generic planner and no manifest changes.
## Trade-offs
- **Single copy.** If a trainer node holding part of a generation dies, that generation is unreadable; consumers detect it at prepare and stay on the previous one. `replica_num` remains a knob for deployments that want durability at the cost of publish traffic.
- **Reader failure.** Before apply, prepare fails and the assignment can re-route; after apply, existing pause-and-rebuild semantics govern.
- **Fan-out.** A range with `k` cross-node consumers costs `k × size` outbound on its reader; the assignment balances read-plus-push bytes and splits large objects to keep NICs even.
## Implementation plan
1. **Compatibility mode and generation lifecycle** — single-replica prefer-local publication, version lifecycle, per-tensor leaders with hierarchical Transfer Engine push into VeRL's existing loader. The single-pull invariant holds from this phase.
2. **Parameter-direct mode** — lower planner output to ranged reads plus push records writing into rollout parameters; update-scoped registration; all-prepare-before-apply; rebuild on partial apply.
3. **Persistent registration** — promote Store and Transfer Engine registrations to a persistent cache reused across generations; invalidate on reconnect or allocation change.
4. **Static schedule caching and pipelining** — cache the routing table by digests and topology; run the read → redistribute → scatter pipeline under the staging watermark.
5. **Maintenance and availability** — locality-aware reader assignment from segment hints; optional multi-replica publication.
Contributor guide
Assessment
This issue has not been assessed yet.