[RFC] RL Training Weight Offload and Restore Paths for Mooncake Store

# [RFC]: RL Training Weight Offload and Restore Paths for Mooncake Store

## Changes proposed

This RFC proposes a staged design for using Mooncake Store as a high-performance offload/restore backend for GPU-resident model weights during RL training.

The primary scenario is an iterative RL training loop:

1. The trainer or rollout worker owns updated model weights in GPU HBM.
2. The worker offloads those GPU tensors into Mooncake Store when the weights need to leave HBM.
3. Mooncake keeps the weight object in Store segment memory, with hard-pin/upsert semantics provided by the weight-storage work.
4. A later trainer or rollout step restores the Store object directly into GPU tensors.
5. In the common production deployment, the Python process is a DummyClient and the resource-owning process is a standalone RealClient on the same node.

The goal is to make this same-node GPU <-> Store memory path bandwidth-optimal while preserving the existing cross-node RDMA behavior.

This RFC is a follow-up to Mooncake issue #1621, "[Call for Contribution] Enhance Model Weight Storage for Mooncake Store". That RFC identifies model weights as a first-class Mooncake Store workload and calls out the semantic/API work needed for that workload, especially hard pin support, native upsert, and documentation/examples for RL and model-management systems.

This proposal focuses on the next layer: the high-performance tensor data path for weight offload and restore. It is also complementary to the SGLang/LMSYS P2P weight-transfer work, which uses Mooncake TransferEngine for cluster-wide weight distribution. In short:

- Issue #1621 defines the Store semantics needed by model weights.
- The P2P weight-transfer path optimizes cross-node weight distribution.
- This RFC optimizes Store-backed same-node GPU <-> CPU memory movement for RL weight offload/restore.

Together they form a complete in-memory weight lifecycle: fast distribution across workers, first-class Store object semantics, and fast local offload/restore when weights must be retained outside HBM.

## Motivation

Mooncake Store already has the right high-level building blocks:

- RealClient-owned Store segments for distributed memory storage.
- DummyClient + RealClient deployment mode for applications that should not own Store segments directly.
- Tensor APIs for Python callers.
- Pinned Store segment memory support for faster GPU <-> CPU DMA.
- TransferEngine/RDMA for cross-node transfer.

However, the current weight offload/restore path can still pay avoidable copies or force data through staging buffers in the same-node DummyClient case. For large model weights, these copies dominate latency. The important distinction is that the same-node path and cross-node path should not be optimized in the same way:

- Same-node GPU <-> Store segment should use CUDA copy / CUDA IPC handoff and avoid client staging buffers.
- Cross-node memory transfer should continue to use TransferEngine/RDMA and registered buffers.

## Non-goals

- This RFC does not replace issue #1621. It builds on the weight-storage semantics proposed there.
- This RFC does not replace the P2P weight-transfer path based on TransferEngine GPU/RDMA distribution.
- This RFC does not require changing the Store object consistency model.
- This RFC does not make all local buffers pinned. Pinned memory must remain quota-controlled and limited to Store/offload-owned memory, not arbitrary temporary buffers.

## Existing and in-review building blocks

### Already split or merged

- Model-weight storage semantics proposed in issue #1621:
- Hard pin support for objects that must not be evicted unexpectedly.
- Native upsert support for frequently updated weight objects.
- Documentation/examples for RL and model-management scenarios.
- Safe GPU tensor staging direction fix: write paths must copy caller GPU buffers with device-to-host semantics when staging is required.
- Store segment pinned memory quota: only Store segment memory is pinned, under user-configured quota. This is intentionally different from pinning arbitrary client local buffers.

### Currently under review / current branch

- Tensor no-staging path for same-node memory writes.
- DummyClient tensor write support for CUDA IPC handoff:
- Dummy process exports the caller source CUDA tensor as a CUDA IPC handle.
- RealClient opens the handle and directly reads from the caller GPU allocation.
- Metadata still uses DummyClient shared memory because it is small and already follows the existing DummyClient control/data handoff model.

### Planned follow-up work

- Get-side no-staging for same-node tensor reads.
- Extend the CUDA IPC handoff design to cover both write and read tensor APIs cleanly.
- Direct tensor reads into GPU memory:
- Read a Store tensor object directly into a caller-provided CUDA tensor.
- In DummyClient mode, export the destination CUDA tensor by CUDA IPC so the RealClient can materialize data into GPU memory without routing the payload through DummyClient SHM.

## Proposed architecture

### RL same-node weight offload/restore fast path

```mermaid
flowchart LR
subgraph App["Training / Rollout process"]
GPUTensor["GPU weight tensor
trainer / rollout"]
Dummy["Mooncake DummyClient"]
end

subgraph RealProc["Standalone RealClient process on same node"]
Real["RealClient"]
Segment["Store segment memory
optionally pinned"]
StoreObj["Mooncake Store object"]
end

GPUTensor -->|CUDA IPC handle| Dummy
Dummy -->|RPC: tensor metadata in SHM
CUDA IPC handle for payload| Real
Real -->|open CUDA IPC mapping| GPUTensor
Real -->|same-node CUDA memcpy
no client staging buffer| Segment
Segment --> StoreObj

StoreObj --> Segment
Segment -->|same-node CUDA memcpy H2D
direct get into CUDA tensor
no client staging buffer| Real
Real -->|open CUDA IPC destination| GPUTensor
```

Key properties:

- The DummyClient does not need a Store segment.
- The DummyClient does not pass RealClient Store segment pointers back to Python.
- Large tensor payloads are handed off by CUDA IPC, not copied through DummyClient shared memory.
- Read-side tensor APIs should support direct materialization into caller-provided CUDA tensors (`get_tensor_into(cuda_tensor)` / batch variants), not only returning CPU-backed buffers that the application later copies to GPU.
- Metadata remains in shared memory because it is small and already fits the existing DummyClient model.
- The RealClient decides after allocation whether the chosen Store segment is local. If it is local and local memcpy is enabled, it uses the same-node fast path. The allocator must not change default placement semantics just to force locality.
- If the allocated segment is remote, the request remains on the normal TransferEngine/RDMA path.

## API and implementation plan

### Phase 1: Same-node tensor write fast path

Scope:

- Tensor put/upsert paths.
- DummyClient + RealClient mode.
- RealClient mode where the caller can provide a CUDA tensor directly.

Implementation requirements:

- Use the existing allocator result to determine whether the destination Store segment is local.
- For local segments and `MC_STORE_MEMCPY` enabled, copy directly from CUDA IPC mapped source tensor to Store segment memory.
- Do not add a user-visible switch for this fast path. It is an implementation detail of same-node placement and local memcpy capability.
- Do not fallback to staging after the system has already determined a local fast path is possible, except for explicit invalid-input failures.
- Cover `put_tensor`, `batch_put_tensor`, `upsert_tensor`, and batch upsert variants.

### Phase 2: Same-node tensor read fast path

Scope:

- Tensor get-into APIs where the destination is a CUDA tensor.
- Direct reads from Mooncake Store tensor objects into GPU memory.
- DummyClient + RealClient mode.

Implementation requirements:

- Add or extend tensor get-into APIs so callers can pass preallocated CUDA tensors, for example `get_tensor_into(key, cuda_tensor)` and `batch_get_tensor_into(keys, cuda_tensors)`.
- Export destination tensor by CUDA IPC in the DummyClient process.
- RealClient opens the destination CUDA IPC handle and writes directly from local Store segment memory into the GPU tensor when the selected source segment is local.
- The API should validate dtype, shape, size, and contiguity against the stored tensor metadata before starting the data transfer.
- The fast path should avoid a CPU staging buffer on memory hits. A Store memory object should be readable into GPU memory as Store segment -> CUDA tensor.
- Preserve the normal RDMA path for remote memory replicas.

## Relationship to issue #1621 and existing weight-transfer work

Issue #1621 frames the model-weight storage problem for Mooncake Store. It points out that KVCache-oriented assumptions are not always ideal for weights: weights should often be protected from eviction, and RL workflows need frequent update semantics. The proposed answer there is composable Store features rather than a fixed "RL mode": hard pin, upsert, and clear documentation/examples.

This RFC assumes those semantics and asks what the optimal data path should be once weights are stored as first-class Mooncake Store objects.

The existing P2P weight-transfer work demonstrates that Mooncake TransferEngine can move large model weights across many GPUs efficiently using RDMA/zero-copy transport. In the published SGLang/LMSYS write-up, the P2P path uses a source-side CPU engine replica and Mooncake TransferEngine RDMA to reduce 1T-parameter Kimi-K2 weight update time from 53s to 7.2s. This RFC does not compete with that path. It fills a different layer:

| Layer | Primary problem | Best path |
| --- | --- | --- |
| Store object semantics | Keep/update model weights correctly as Store objects | #1621 hard pin + upsert + docs |
| Cluster-wide weight broadcast / update | Move weights from producer GPUs to consumer GPUs across nodes | TransferEngine GPU/RDMA weight-transfer path |
| Same-node Store write/read | Move GPU tensors into/out of local Store memory | CUDA IPC handoff + local CUDA memcpy + pinned Store segment |
| Direct tensor get to GPU | Materialize Store tensor objects directly into CUDA tensors | `get_tensor_into(cuda_tensor)` / batch variants with CUDA IPC destination |

The combined design gives a complete lifecycle:

1. Use #1621 semantics so model weights have explicit hard-pin/upsert behavior when needed.
2. Use the TransferEngine weight-transfer path to distribute fresh weights across the cluster.
3. Use Mooncake Store tensor APIs to persist/offload local weight tensors without extra same-node staging copies.
4. Restore Mooncake Store tensor objects directly into GPU tensors for the next RL training or rollout step.

## Correctness and safety considerations

- CUDA IPC lifetime: the exporting process must keep the tensor allocation alive until the RealClient finishes the operation. The API must document that the input/output tensor cannot be freed or resized concurrently.
- DummyClient SHM lifetime: metadata buffers remain under the existing DummyClient shared-memory mapping and RPC lifetime rules.
- Store segment lifetime: local fast path must operate only on allocated Store segment ranges and must not expose raw Store segment pointers to Python.
- Pinned memory quota: pinned memory is limited to Store/offload-owned regions. Exceeding quota should fall back to pageable Store memory, not to opportunistic per-operation pinning.
- Same-node detection: locality should be derived from the actual allocated segment or selected replica, not from a user-specified preference that changes placement semantics.
- Cross-node behavior: remote replicas continue to use TransferEngine/RDMA and must not require CUDA IPC.

## Performance expectations

Expected benefits are concentrated in the same-node GPU <-> CPU memory path:

- GPU -> Store memory put: avoid DummyClient payload copy and client staging buffer; use pinned Store segment when available.
- Store memory -> GPU get: avoid client staging buffer on same-node reads; use CUDA H2D from Store-owned memory.
- Direct tensor get to GPU: avoid returning CPU-backed tensors followed by a framework-level `.to("cuda")`; materialize directly into the caller-provided CUDA tensor.

Cross-node RDMA is expected to see less benefit from same-node no-staging changes because the network transfer remains the dominant path. Cross-node optimization should be evaluated separately and should reuse TransferEngine primitives rather than adding a parallel data path.

## Open questions

1. Should direct GPU tensor reads be exposed as a dedicated CUDA API or, preferably, as a tensor-aware overload/extension of the existing `get_tensor_into` and `batch_get_tensor_into` APIs?
2. How should the Store expose capability information to upper layers: CUDA IPC available and pinned Store segment quota available?
3. Should weight objects use hard pin / soft pin / group IDs by default, or should this remain application-controlled through `ReplicateConfig`?
4. What benchmark matrix should be required before merging each phase?

## Proposed benchmark matrix

- Same-node `put_tensor(cuda)` and `batch_put_tensor(cuda)`:
- pageable Store segment vs pinned Store segment
- staging vs no-staging
- DummyClient + RealClient mode
- Same-node direct tensor reads into GPU memory:
- `get_tensor_into(key, cuda_tensor)`
- `batch_get_tensor_into(keys, cuda_tensors)`
- tensor object restore into preallocated CUDA tensors
- pageable Store segment vs pinned Store segment
- staging vs no-staging
- DummyClient + RealClient mode
- Cross-node RDMA:
- confirm no regression on existing RDMA put/get paths

## Rollout plan

1. Merge correctness fixes for tensor GPU copy direction and CUDA runtime linkage.
2. Merge Store segment pinned memory quota support.
3. Merge same-node tensor write no-staging support.
4. Merge same-node tensor read no-staging support.

Each phase should include focused tests, scoped documentation, and performance data for the exact path it changes.

## References

- Mooncake issue #1621, "[Call for Contribution] Enhance Model Weight Storage for Mooncake Store": https://github.com/kvcache-ai/Mooncake/issues/1621
- Mooncake README update: https://github.com/kvcache-ai/Mooncake#-updates
- SGLang/LMSYS P2P weight transfer write-up: https://www.lmsys.org/blog/2026-04-29-p2p-update/
- Mooncake governance RFC requirement: https://github.com/kvcache-ai/Mooncake/blob/main/docs/source/community/governance.md
- Mooncake Store design: https://github.com/kvcache-ai/Mooncake/blob/main/docs/source/design/mooncake-store.md

Review the existing tensor put/upsert and get APIs, including put_tensor, batch_put_tensor, upsert_tensor, and the proposed get_tensor_into and batch_get_tensor_into entry points. Trace the DummyClient/RealClient handoff, allocator locality decision, and TransferEngine/RDMA fallback. Done means same-node write and read paths avoid staging while validation, lifetime, quota, and remote-path requirements remain intact.

Written by the indexing model from the issue text.