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

[zxpdemonio]

[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 [Call for Contribution] Enhance Model Weight Storage for Mooncake Store #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 [Call for Contribution] Enhance Model Weight Storage for Mooncake Store #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 [Call for Contribution] Enhance Model Weight Storage for Mooncake Store #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

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

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

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

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

Loading

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 [Call for Contribution] Enhance Model Weight Storage for Mooncake Store #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 [Call for Contribution] Enhance Model Weight Storage for Mooncake Store #1621, "[Call for Contribution] Enhance Model Weight Storage for Mooncake Store": [Call for Contribution] Enhance Model Weight Storage for Mooncake Store #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

[github-actions]

Thanks for opening this issue, @zxpdemonio!

Field	Value
Issue	#3229
GitHub user ID	16206312
Reporter	@zxpdemonio

A maintainer will triage this when possible. To help us respond faster, please include:

• Mooncake version or commit SHA
• Environment (OS, CUDA/driver, RDMA stack if relevant)
• Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.