RFC: Optional Mooncake DataProto Transfer Backend for ROLL

[zxpdemonio]

Background

ROLL uses DataProto as the core data container between rollout, reward, and training stages. A DataProto object can contain tensor data, non-tensor data, and metadata:

DataProto                                                                                                               
├── batch: TensorDict | None                                                                                            
├── non_tensor_batch: dict[str, np.ndarray]                                                                             
└── meta_info: dict[str, Any]                                                                                           
                                                                                              

For large-scale RL and VLM workloads, rollout payloads can become large and may include both tensor-heavy data and multimodal / object-heavy fields. Efficiently transferring this structured data across workers and nodes is important for distributed training throughput.

This RFC proposes adding Mooncake as an optional transfer backend for ROLL DataProto.

Mooncake is not required for all deployments. Existing transfer behavior remains available, and users can enable Mooncake only when they want RDMA-backed large-payload transfer.

Goals

The Mooncake backend should provide:

• Optional integration through ROLL's transfer backend abstraction.
• Compatibility with the complete ROLL DataProto object model.
• Support for tensor, non-tensor, metadata, and multimodal payloads.
• Efficient cross-node transfer for large rollout data.
• Minimal changes to ROLL pipeline logic.
• Safe and maintainable memory ownership boundaries.
• A configuration-driven way to enable or disable Mooncake.

Non-Goals

This RFC does not propose:

• Making Mooncake a mandatory dependency for all ROLL users.
• Replacing all existing transfer mechanisms.
• Adding Mooncake-specific logic into pipeline implementations.
• Changing the public DataProto data model.
• Adding benchmark numbers into this RFC.
• Implementing direct write into Mooncake local buffers in the current scope.

Benchmark results should be reported separately in the PR log.

Optional Backend Design

Mooncake should be integrated as one backend implementation behind ROLL's existing transfer backend interface.

The normal ROLL flow remains:

DataProto.to_remote(...)
-> transfer_backend.put(...)
-> remote reference
-> transfer_backend.get(...)
-> reconstructed DataProto

When Mooncake is enabled, transfer_backend.put() stores the DataProto through Mooncake and returns a lightweight remote
reference. The receiver later materializes the object through transfer_backend.get().

When Mooncake is not enabled, ROLL continues to use the existing transfer behavior.

This keeps Mooncake optional and avoids coupling pipeline code to one specific transport.

Configuration

Mooncake can be enabled through configuration:

transfer_backend:
backend_name: Mooncake
backend_config:
protocol: rdma
client_scope: node

A deployment can choose whether to enable Mooncake based on cluster environment, RDMA availability, payload size, and
workload characteristics.

For environments that do not have Mooncake or RDMA configured, users can keep the default backend.

DataProto Compatibility

Mooncake supports structured transfer of the full ROLL DataProto object model.

A ROLL DataProto may contain:

DataProto
├── batch: TensorDict | None
├── non_tensor_batch: dict[str, np.ndarray]
└── meta_info: dict[str, Any]

Mooncake preserves all three components during remote transfer.

Tensor Batch Support

Mooncake supports DataProto.batch, including tensor fields commonly used by rollout, reward, and training stages, such
as:

• token ids
• attention masks
• position ids
• log probabilities
• rewards
• advantages
• masks
• other tensor or TensorDict-style fields

This allows tensor-heavy rollout data to be transferred efficiently while preserving the existing ROLL worker
interfaces.

Non-Tensor Batch Support

Mooncake supports DataProto.non_tensor_batch, including:

• NumPy arrays
• object arrays
• text fields
• per-sample metadata
• variable-length fields
• multimodal input structures
• reward-side auxiliary fields

This is important because ROLL workloads, especially VLM and RLVR workloads, may carry data that cannot be represented
cleanly as fixed-shape tensors.

Meta Info Support

Mooncake supports DataProto.meta_info, including auxiliary metadata used by ROLL pipeline execution, scheduling,
post-processing, and downstream workers.

This allows metadata to remain attached to the transferred DataProto instead of being handled through a separate side
channel.

Multimodal DataProto Support

For VLM workloads, a single DataProto may contain both tensor data and multimodal non-tensor fields, for example:

• text tokens
• image-related inputs
• multimodal processor outputs
• sample-level metadata
• reward-related information

Mooncake transfers these fields as one logical structured object, so the receiver can reconstruct the expected DataProto
layout.

Partial DataProto Support

Different ROLL stages may produce or consume different subsets of the DataProto structure. Mooncake supports these
combinations, including:

• tensor-only DataProto
• non-tensor-only DataProto
• metadata-only DataProto
• full tensor + non-tensor + metadata DataProto

This allows Mooncake to support both current tensor-heavy paths and future multimodal / object-heavy paths.

Remote Reference Support

Mooncake returns a lightweight remote reference after put() instead of eagerly materializing the full object everywhere.

The receiver can materialize the DataProto through the existing backend interface when needed:

transfer_backend.get(remote_reference)
-> DataProto

This keeps the integration compatible with ROLL's existing DataProto.to_remote() flow.

Compatibility Summary

DataProto Component	Supported	Notes
batch	Yes	Tensor / TensorDict-style rollout and training data
non_tensor_batch	Yes	NumPy arrays, object arrays, multimodal fields, variable-length metadata
meta_info	Yes	Pipeline and worker metadata
Tensor-only payloads	Yes	Supports current tensor-heavy transfer paths
Non-tensor-heavy payloads	Yes	Important for VLM / multimodal / object-heavy workloads
Mixed multimodal payloads	Yes	Preserves one logical DataProto object
Remote references	Yes	Compatible with ROLL to_remote / get flow

Why Mooncake

Mooncake provides several advantages for large-scale ROLL workloads.

1. Efficient Large Payload Transfer

Mooncake is designed for high-throughput object transfer and can use RDMA for cross-node movement of large payloads.

This is a good fit for rollout data movement in distributed RL training.

2. Structured Data Transfer

Mooncake can transfer the complete DataProto structure instead of treating it as a single opaque object.

This improves compatibility with ROLL's data model and provides room for future optimization.

3. Better Fit for Multimodal Workloads

Multimodal RL workloads often mix tensors, metadata, and object-style fields.

Mooncake's structured DataProto support is suitable for this mixed payload structure and can support both tensor-heavy
and non-tensor-heavy workloads.

4. Optional Deployment

Mooncake can be enabled only in environments where it is useful.

For example, deployments with large rollout payloads, multi-node training, and RDMA support can enable Mooncake, while
smaller or single-node deployments can keep the existing backend.

5. Node-Scoped Client Support

Mooncake can use node-scoped client resources. This is useful when transport resources should be managed once per node
instead of repeatedly initialized by each worker.

This can reduce initialization overhead and better match RDMA resource management.

6. Safe Buffer Management

The current implementation uses BufferPool-based staging to keep memory ownership explicit and safe.

This provides a conservative production-ready default while still supporting efficient transfer.

Integration Scope

The proposed integration should keep ROLL-side changes minimal.

Expected ROLL-side scope:

• Add Mooncake as a registered transfer backend.
• Allow selecting Mooncake through config.
• Preserve the existing DataProto.to_remote() interface.
• Preserve existing pipeline and worker interfaces.
• Avoid adding Mooncake-specific branches inside pipeline logic.

Expected Mooncake-side scope:

• Provide structured DataProto put / get support.
• Preserve tensor, non-tensor, and metadata fields.
• Return portable remote references.
• Support cleanup of remote objects.
• Support safe buffer management.

Future Optimization Direction

A future optimization is to allow producers or serializers to write directly into Mooncake BufferPool local buffers.

Conceptually:

producer / serializer
-> Mooncake BufferPool local buffer
-> transfer

This could reduce one staging copy and move the path closer to zero-copy transfer.

However, this would require more invasive changes to serialization, ownership, and buffer lifetime management.
Therefore, it is intentionally deferred and is not part of the current integration.

Rollout Plan

1. Add Mooncake as an optional transfer backend.
2. Enable configuration-based backend selection.
3. Support full DataProto transfer:

• batch
• non_tensor_batch
• meta_info

4. Validate tensor-heavy and multimodal / non-tensor payload compatibility.
5. Report final optimized benchmark results in the PR log.
6. Keep future zero-copy BufferPool write support as follow-up work.