feat: add Mooncake RDMA transport for rollout data

[zxpdemonio]

Summary

Add Mooncake structured object store as an alternative transport for rollout data between rollout manager and trainer, activated by --rollout-data-transport=mooncake. This bypasses Ray's plasma object store and transfers rollout tensors directly over RDMA, reducing cross-node data transfer latency by ~3x for typical rollout payloads.

Performance

Benchmark: slime rollout dict , cross-machine (eRDMA, A10 cluster):

Size	Ray PUT	Ray GET	Ray GET BW	MoonCake PUT	MoonCake PUT BW	MoonCake GET	MoonCake GET BW	GET Speedup
128MB	61ms	65ms	1.9 GB/s	30ms	3.97 GB/s	18ms	6.5 GB/s	3.6x
512MB	207ms	211ms	2.4 GB/s	71ms	6.72 GB/s	59ms	8.0 GB/s	3.6x
1GB	296ms	432ms	2.3 GB/s	130ms	7.35 GB/s	114ms	8.3 GB/s	3.8x
4GB	340ms	1561ms	2.6 GB/s	477ms	7.99 GB/s	430ms	8.9 GB/s	3.6x
16GB	708ms	5802ms	2.8 GB/s	1903ms	8.02 GB/s	1701ms	9.0 GB/s	3.4x

~3x faster on the trainer-side get path, which is on the critical path of the training loop.

Why faster than Ray

Ray object store serializes the entire dict with pickle/msgpack, then transfers the serialized blob over gRPC/plasma. On the receiver side it deserializes the full payload.

Mooncake takes a different approach:

• Zero-copy RDMA transfer: Tensor buffers are transferred directly via RDMA without serialization overhead
• Schema-driven typed_ragged codec: Variable-length tensor fields (tokens, loss_masks, log_probs, etc.) are encoded as flat buffer + offset array, enabling scatter-gather RDMA reads without per-element pickling
• BufferPool-based memory registration: Pre-registered RDMA memory avoids per-transfer registration cost
• Chunked parallel transfer: Large payloads are split into chunks and transferred with configurable inflight concurrency

The net effect: Ray pays serialize + network + deserialize; Mooncake pays only network (RDMA) + lightweight metadata decode.

Design

• Thin integration: 1 new file (slime/utils/data_transfer.py, 84 lines) + minimal dispatch in existing code (+103/−3 total)
• Schema-driven encoding: Field schemas derived from existing _ROLLOUT_DATA_TENSOR_DTYPES (statically cached), passed to put_legacy_dict for typed_ragged codec on
  variable-length tensor fields
• Same Box wrapper pattern as nixl transport — transparent to training loop
• Cleanup: Explicit cleanup_mooncake_rollout_refs() after each training step to free remote buffers

Usage

python train.py \                                                                                                                                                                  
  --rollout-data-transport mooncake \                                                                                                                                              
  # Mooncake store config via env vars:                                                                                                                                            
  # MC_LOCAL_HOSTNAME, MC_METADATA_SERVER, MC_PROTOCOL, MC_DEVICE, MC_SEGMENT_SIZE, MC_BUFFER_SIZE

Testing

• ruff check slime/utils/data_transfer.py slime/utils/arguments.py slime/utils/data.py slime/ray/rollout.py train.py train_async.py
• git diff HEAD --check
• Cross-machine benchmark (4.59 GB, eRDMA A10 cluster, 70→73): Ray ~1800ms vs Mooncake ~600ms remote get

No benchmark scripts or benchmark outputs are included in this PR.

Checklist

• Mooncake rollout data transport wiring (--rollout-data-transport=mooncake)
• Schema-driven typed_ragged encoding via _ROLLOUT_DATA_TENSOR_DTYPES
• Driver-side post-training cleanup lifecycle (cleanup_mooncake_rollout_refs)
• Lazy store initialization with env-var / args configuration
• Format/lint checks (pre-commit pass)

[Bo-Vincent]

I re-ran the comparison using a PR-aligned bandwidth accounting boundary.

The key point is that NIXL raw ray.put(..., _tensor_transport="nixl") is not directly comparable to Mooncake PUT. In the Mooncake PR benchmark, PUT measures put_legacy_dict(data) after the rollout dict has already been constructed. Therefore, for NIXL, the closest producer-side equivalent is:

NIXL aligned PUT = pack_rollout_to_cuda(data) + ray.put(packed, _tensor_transport="nixl")

For this reason, I report NIXL pack and pure PUT separately below, and calculate NIXL PUT bandwidth with pack + pure PUT as the denominator. pure PUT shows the raw transport-stage cost, while NIXL PUT BW reflects the PR-aligned producer-side cost for this structured rollout dict workload.

I also added unpack/reconstruct timing on the NIXL reader side. So the NIXL GET bandwidth below uses pure GET + unpack as the denominator, which is closer to Mooncake get_legacy_dict(ref) in terms of start/end boundary. I still report NIXL pure GET separately to show the raw packed-payload transport cost.

The main observation is that NIXL's pure packed-payload transfer is fast, but reconstructing the legacy rollout dict is expensive. In other words, the large gap between NIXL pure GET BW and aligned NIXL GET BW mostly comes from unpack/reconstruct cost rather than the raw NIXL data path.

Benchmark: slime rollout dict, cross-machine RDMA, 2-node H20 cluster, 1000-sample rollout payload tiled/subsampled to target sizes.

Size	Mooncake PUT	Mooncake PUT BW	Mooncake GET	Mooncake GET BW	Mooncake E2E	NIXL pack	NIXL pure PUT	NIXL PUT BW	NIXL pure GET	NIXL pure GET BW	NIXL unpack	NIXL aligned GET	NIXL GET BW	NIXL E2E
128MB	50.5ms	2.47 GiB/s	8.0ms	15.61 GiB/s	58.5ms	183.8ms	5.5ms	0.66 GiB/s	27.5ms	4.54 GiB/s	71.8ms	99.2ms	1.26 GiB/s	288.6ms
512MB	40.5ms	12.33 GiB/s	21.7ms	23.02 GiB/s	62.2ms	181.9ms	5.7ms	2.66 GiB/s	56.0ms	8.92 GiB/s	293.2ms	352.1ms	1.42 GiB/s	536.8ms
1GB	79.5ms	13.20 GiB/s	42.1ms	24.93 GiB/s	121.6ms	416.7ms	5.8ms	2.48 GiB/s	101.3ms	10.36 GiB/s	611.6ms	711.7ms	1.47 GiB/s	1135.4ms
4GB	280.5ms	14.97 GiB/s	181.9ms	23.09 GiB/s	462.4ms	1847.4ms	5.6ms	2.27 GiB/s	142.6ms	29.44 GiB/s	2413.9ms	2558.1ms	1.64 GiB/s	4409.5ms
16GB	1002.5ms	16.54 GiB/s	749.5ms	22.13 GiB/s	1752.0ms	6945.0ms	5.7ms	2.39 GiB/s	457.4ms	36.26 GiB/s	9673.9ms	10130.5ms	1.64 GiB/s	17082.0ms

My interpretation is that NIXL has a very strong raw transport path once the payload is already packed: the pure PUT cost is only around 5.5ms in this benchmark. However, for this PR's target workload, namely structured/ragged rollout dict transfer, the pack step is part of the producer-side cost. Under this PR-aligned PUT bandwidth accounting, Mooncake has higher PUT BW across all tested sizes.

On the GET side, after including unpack/reconstruct back to the legacy rollout dict shape, Mooncake also has higher GET BW across all tested sizes. NIXL pure GET is still useful as a raw packed-payload transport reference, but it should not be treated as a complete legacy-dict GET replacement.

[Bo-Vincent]

Additional hardware note for the benchmark above:

The cross-machine RDMA benchmark was run on a 2-node H20 cluster with BlueField-3 E-series SuperNICs.

RDMA NIC model:
MT41692 - 900-9D3B4-00CV-EA0, BlueField-3 E-series SuperNIC, 200GbE/NDR200 IB, Dual-port QSFP112, PCIe Gen5.0 x16.

RDMA bandwidth:
Each node has 4 RDMA bonds (mlx5_bond_0 to mlx5_bond_3). Each bond is reported as 400Gbps by ethtool, backed by 2 x 200Gbps physical links. The benchmark used a single RDMA bond (mlx5_bond_0), so the effective link bandwidth for this test was 400Gbps.

Also note that the original performance numbers in this PR were measured on an eRDMA setup, while the later numbers I reported were measured on this RDMA setup. Therefore, the later RDMA results are expected to be better than the original eRDMA results.

[zhuzilin]

Great! Do we need to upgrade mooncake for this PR?

[zxpdemonio]

Great! Do we need to upgrade mooncake for this PR?

@zhuzilin Yes. This PR should be used with the latest Mooncake main. The required structured-object transfer APIs have already landed there.

Once Mooncake cuts a release containing the current main APIs, we can pin/document that release version here.

[stmatengss]

The mooncake package has been upgraded to version 0.3.12.post1 on the sglang side, fully supporting all features in this PR.

[zxpdemonio]

@lilei199908 This PR has been rebased onto the latest main. I also added English and Chinese user guides covering prerequisites, complete two-node setup, TCP/RDMA configuration, sync/async entrypoints, the configuration reference, and troubleshooting.\n\nThe documented path was validated with the same Mooncake wheel on both nodes: both TCP and RDMA structured-object transfer passed, including GET result release and object cleanup, and a two-node Qwen3-4B job completed several rollout/training iterations through the Mooncake RDMA backend. Could you please take another look when convenient?

[zxpdemonio]

@lilei199908 CI update after the full run completed: the original failing test_qwen3_30B_A3B.py passed, and the transient Moonlight model-loading failure also passed. The branch is already based on the latest main (behind_by: 0), so rebasing would be a no-op. The run was ultimately blocked by two infrastructure or baseline failures: (1) test_full_disk_weight_update.py stayed in Execute on h20-node0 for nearly six hours and was cancelled, although the same official test normally passes in about five minutes on h20-node1; (2) test_qwen3.6_35B_A3B_pd_mooncake.py failed in torch_memory_saver.pause() with cudaError invalid argument. Current run: https://github.com/THUDM/slime/actions/runs/31824820879. Could you please reset or inspect h20-node0 and rerun the failed jobs?