feat(data_plane): add Mooncake GDR support

examples/configs/grpo_math_1B.yaml

@@ -565,6 +565,8 @@ data_plane:
     local_buffer_size:    4294967296 # 4 GiB/process
     reuse_registered_buffers: true   # reuse RDMA-registered buffers
     staging_buffer_size:   268435456 # 256 MiB/pool slot; bigger transfers bypass the pool
+    use_gdr: false                    # GPU-memory RDMA staging in CUDA clients
+    gdr_staging_buffer_mb: 1024       # persistent MiB per active GDR client
   # observability:                    # NotRequired
   #   enabled: false
 

nemo_rl/data_plane/README.md

@@ -437,6 +437,8 @@ data_plane:
     local_buffer_size:    4294967296   # 4 GiB/process
     reuse_registered_buffers: true     # reuse RDMA-registered buffers
     staging_buffer_size:   268435456   # 256 MiB/pool slot; bigger transfers bypass the pool
+    use_gdr: false                      # GPU-memory RDMA staging in CUDA clients
+    gdr_staging_buffer_mb: 1024         # persistent MiB per active GDR client
   # observability:                     # NotRequired
   #   enabled: false
 ```
@@ -445,13 +447,29 @@ These keys used to sit directly under `data_plane:`. That spelling is not
 rejected — it is simply never read. A config still using it silently gets
 this backend's defaults instead of its own values: an inherited config
 supplies the nested block, so a surviving flat key always loses the merge,
-with no warning either way.
+with no warning either way. Earlier revisions of PR #3501 also used flat
+`use_gdr` and `gdr_staging_buffer_mb` keys; those two spellings are rejected
+with a migration error because silently disabling GDR would invalidate a run.
 
 Backend choice:
 - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests
   and small runs.
-- **`mooncake_cpu`** — Mooncake transfer engine; higher throughput at
-  scale. Required for multi-node clusters with large bulk volume.
+- **`mooncake_cpu`** — Mooncake's RDMA-only transfer engine. By default,
+  tensors transfer through registered CPU staging. Set
+  `mooncake_cpu.use_gdr: true` to let CUDA-initialized clients use
+  TransferQueue's GDR staging path. CPU-only clients, such as a
+  SingleController producer, continue to use CPU RDMA. GDR changes the
+  client-side tensor transfer and staging path; queued objects still reside in
+  Mooncake-managed host-memory segments.
+
+The CPU host staging pool's `staging_buffer_size` is independent of the GDR
+buffer. `gdr_staging_buffer_mb` is the persistent GPU staging capacity per
+active CUDA client and defaults to 1024 MiB. Transfers through that per-client
+buffer are serialized. Aggregate fetches may exceed the buffer and are split
+into groups. In the mixed CPU-producer/GDR-receiver flow used by
+SingleController, however, each individual tensor must currently fit because
+the CPU PUT path does not create the chunk metadata required by an oversized
+GDR GET.
 
 Capacity rule of thumb (any backend):
 

nemo_rl/data_plane/adapters/transfer_queue.py

@@ -27,6 +27,7 @@
 import importlib
 import ipaddress
 import json
+import logging
 import os
 import resource
 import socket
@@ -55,6 +56,8 @@
     data_plane_supports_checkpointing,
 )
 
+LOGGER = logging.getLogger(__name__)
+
 # ──────────────────────────────────────────────────────────────────────────
 # Backend init — lifted from rl-arena/arena/backends.py.
 # ──────────────────────────────────────────────────────────────────────────
@@ -684,6 +687,8 @@ def _init_tq(cfg: DataPlaneConfig) -> None:
                     "metadata_server": f"{local_ip}:50050",
                     "master_server_address": f"{local_ip}:50051",
                     **_mooncake_transport_config(),
+                    "use_gdr": bool(mooncake_cfg.use_gdr),
+                    "gdr_staging_buffer_mb": int(mooncake_cfg.gdr_staging_buffer_mb),
                 },
             },
         }
@@ -817,6 +822,13 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None:
 
         self._backend = cfg["backend"]
         self._supports_checkpointing = data_plane_supports_checkpointing(cfg)
+        # GDR is a mooncake_cpu-only transport knob, so key it off the backend
+        # directly rather than off any incidental per-backend flag.
+        self._gdr_requested = self._backend == "mooncake_cpu" and bool(
+            backend_config(cfg).use_gdr
+        )
+        self._gdr_put_confirmed = False
+
         if bootstrap:
             _init_tq(cfg)
         else:
@@ -1031,6 +1043,28 @@ def put_samples(
             wire_fields = detached_fields
             field_names = [str(key) for key in detached_fields.keys()]
 
+        confirm_gdr_put = bool(
+            self._gdr_requested
+            and not self._gdr_put_confirmed
+            and torch.cuda.is_initialized()
+            and wire_fields is not None
+            and any(
+                isinstance(wire_fields.get(key), torch.Tensor)
+                for key in wire_fields.keys()
+            )
+        )
+        gdr_staging = None
+        if confirm_gdr_put:
+            tq_client = tq.get_client()
+            storage_manager = getattr(tq_client, "storage_manager", None)
+            storage_client = getattr(storage_manager, "storage_client", None)
+            gdr_staging = getattr(storage_client, "_gdr_staging", None)
+            if not getattr(storage_client, "use_gdr", False) or gdr_staging is None:
+                raise RuntimeError(
+                    "GDR was requested for a CUDA-initialized TransferQueue "
+                    "client, but TransferQueue selected CPU RDMA for tensor PUTs"
+                )
+
         self._mark_data_operation_started()
         # TQ's wire vocabulary is `keys=` — translation point.
         tq.kv_batch_put(
@@ -1039,6 +1073,16 @@ def put_samples(
             fields=wire_fields,
             tags=user_tags,
         )
+        if confirm_gdr_put:
+            if not getattr(gdr_staging, "_initialized", False):
+                raise RuntimeError(
+                    "TransferQueue tensor PUT completed without initializing "
+                    "the requested GDR staging buffer"
+                )
+            LOGGER.info(
+                "TransferQueue GDR tensor PUT active (partition=%s)", partition_id
+            )
+            self._gdr_put_confirmed = True
 
         return KVBatchMeta(
             partition_id=partition_id,

nemo_rl/data_plane/interfaces.py

@@ -40,7 +40,7 @@
 from pathlib import Path
 from typing import Annotated, Any, Callable, Literal, NotRequired, Sequence, TypedDict
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PositiveInt
 from tensordict import TensorDict
 
 DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2
@@ -80,6 +80,11 @@ class MooncakeCpuConfig(BaseModel, extra="allow"):
     admitted and never shrink — so raise it only when a per-key payload (one
     sample of one field) genuinely exceeds it, not for headroom.
 
+    ``use_gdr`` lets CUDA-initialized clients transfer through TransferQueue's
+    persistent GPU staging buffer. ``gdr_staging_buffer_mb`` is the positive
+    HBM capacity of that buffer per active GDR client. CPU-only clients keep
+    using the registered host-buffer path.
+
     Every RDMA rail on the host is offered to mooncake (see ``rdma_devices``).
     That is only safe with ``MC_ENABLE_DEST_DEVICE_AFFINITY=1``, which pins each
     transfer's peer rail to the local one by name; on a rail-isolated RoCE
@@ -91,6 +96,8 @@ class MooncakeCpuConfig(BaseModel, extra="allow"):
     local_buffer_size: int = 4294967296  # 4 GiB per client process
     reuse_registered_buffers: bool = True
     staging_buffer_size: int = 268435456  # 256 MiB per pool slot
+    use_gdr: bool = False
+    gdr_staging_buffer_mb: PositiveInt = 1024
 
 
 class DataPlaneConfig(TypedDict):
@@ -117,7 +124,9 @@ class DataPlaneConfig(TypedDict):
     ``local_buffer_size`` used to sit at this level. A config still using that
     spelling is not rejected — the flat key is simply never read, and
     :func:`backend_config` resolves the nested block (or its defaults) as if it
-    were absent. See there.
+    were absent. Earlier revisions of PR #3501 also put ``use_gdr`` and
+    ``gdr_staging_buffer_mb`` here; those two spellings are rejected explicitly
+    so a stale GDR recipe cannot silently run with GDR disabled. See there.
     """
 
     enabled: bool
@@ -130,6 +139,14 @@ class DataPlaneConfig(TypedDict):
     ack_timeout_ms: NotRequired[int]
     observability: NotRequired["ObservabilityConfig"]
 
+    # Declared only so Pydantic preserves the first GDR integration's flat
+    # spellings until backend_config can raise the migration error below.
+    use_gdr: NotRequired[bool]
+    gdr_staging_buffer_mb: NotRequired[int]
+
+
+_FIRST_GDR_FLAT_KEYS = ("use_gdr", "gdr_staging_buffer_mb")
+
 
 _CHECKPOINTABLE_BACKENDS: frozenset[str] = frozenset({"simple"})
 
@@ -158,9 +175,16 @@ def backend_config(cfg: DataPlaneConfig) -> Any:
     (block already coerced to a model) or as a plain dict from a test.
 
     Sizing is read only from the nested block. A config still using the
-    pre-nesting flat spelling gets this backend's defaults, not its own values.
+    pre-nesting flat sizing spelling gets this backend's defaults, not its own
+    values. The first GDR integration's flat spelling raises instead because
+    silently disabling a requested transport would invalidate a benchmark.
     """
     backend = cfg["backend"]
+    stale_gdr = [key for key in _FIRST_GDR_FLAT_KEYS if key in cfg]
+    if stale_gdr:
+        keys = ",".join(stale_gdr)
+        raise ValueError(f"data_plane.{{{keys}}} moved under data_plane.mooncake_cpu")
+
     nested = cfg.get(backend) or {}
     if isinstance(nested, BaseModel):
         nested = nested.model_dump(exclude_unset=True)

nemo_rl/data_plane/worker_mixin.py

@@ -36,10 +36,9 @@
 import numpy as np
 import torch
 
-FetchPolicy = Literal["auto", "independent", "leader_broadcast"]
-
 from nemo_rl.data.llm_message_utils import attach_message_log_view
 from nemo_rl.data.multimodal_utils import PackedTensor
+from nemo_rl.data_plane.interfaces import backend_config
 from nemo_rl.data_plane.schema import (
     ELEM_COUNTS_PER_GB,
     GLOBAL_FORWARD_PAD_SEQLEN,
@@ -64,6 +63,8 @@
         DataPlaneRuntimeConfig,
     )
 
+FetchPolicy = Literal["auto", "independent", "leader_broadcast"]
+
 
 def _broadcast_batched_data_dict(
     data: Optional[BatchedDataDict[Any]],
@@ -327,6 +328,15 @@ def setup_data_plane(self, cfg: DataPlaneRuntimeConfig) -> None:
         self._route_fallback_counts = Counter()
         from nemo_rl.data_plane import build_data_plane_client
 
+        if (
+            cfg["backend"] == "mooncake_cpu"
+            and backend_config(cfg).use_gdr
+            and not torch.cuda.is_initialized()
+        ):
+            raise RuntimeError(
+                "CUDA must be initialized before attaching TransferQueue with GDR"
+            )
+
         # bootstrap=False — the driver already created the named
         # controller actor; this process attaches as a client.
         self._dp_client = build_data_plane_client(cfg, bootstrap=False)

tests/unit/data_plane/README.md

@@ -107,6 +107,12 @@ Generated audit of every test function under `tests/unit/data_plane/` with a one
 - `test_local_node_ip_returns_empty_on_exception` — DNS exception → returns empty string (no crash).
 - `test_mc_tcp_bind_address_overwrites_existing` — TQDataPlaneClient `__init__` uses direct assignment (not `setdefault`).
 
+## `test_mooncake_gdr.py` (3 tests)
+
+- `test_init_tq_forwards_nested_gdr_config_and_keeps_rdma` — Nested GDR settings reach TransferQueue without changing #2935's RDMA/all-rail transport.
+- `test_cpu_only_client_may_attach_with_gdr_config` — A CPU-only producer may attach because GDR eligibility is client-local.
+- `test_gdr_receiver_requires_cuda_initialized` — A policy receiver must initialize CUDA before attaching its GDR client.
+
 ## `test_message_log_decompose.py` (11 tests)
 
 - `test_decompose_message_log_basic_shapes` — Basic shapes of decompose output.

tests/unit/data_plane/test_backend_config.py

@@ -61,12 +61,19 @@ def test_checkpointing_capability_defaults_to_unsupported(
 def test_nested_block_is_used() -> None:
     cfg = _cfg(
         "mooncake_cpu",
-        mooncake_cpu={"global_segment_size": 111, "reuse_registered_buffers": False},
+        mooncake_cpu={
+            "global_segment_size": 111,
+            "reuse_registered_buffers": False,
+            "use_gdr": True,
+            "gdr_staging_buffer_mb": 512,
+        },
     )
     resolved = backend_config(cfg)
     assert isinstance(resolved, MooncakeCpuConfig)
     assert resolved.global_segment_size == 111
     assert resolved.reuse_registered_buffers is False
+    assert resolved.use_gdr is True
+    assert resolved.gdr_staging_buffer_mb == 512
 
 
 def test_absent_block_falls_back_to_model_defaults() -> None:
@@ -81,6 +88,23 @@ def test_absent_block_falls_back_to_model_defaults() -> None:
     assert resolved.local_buffer_size == 4294967296  # 4 GiB per client process
     # The opt-out flag defaults on, so omitting it must not disable the pool.
     assert resolved.reuse_registered_buffers is True
+    assert resolved.use_gdr is False
+    assert resolved.gdr_staging_buffer_mb == 1024
+
+
+@pytest.mark.parametrize("value", [0, -1])
+def test_gdr_staging_size_rejects_non_positive_values(value: int) -> None:
+    with pytest.raises(pydantic.ValidationError):
+        backend_config(
+            _cfg("mooncake_cpu", mooncake_cpu={"gdr_staging_buffer_mb": value})
+        )
+
+
+@pytest.mark.parametrize("backend", ["simple", "mooncake_cpu"])
+@pytest.mark.parametrize("key", ["use_gdr", "gdr_staging_buffer_mb"])
+def test_flat_gdr_key_from_first_integration_raises(backend: str, key: str) -> None:
+    with pytest.raises(ValueError, match="moved under data_plane.mooncake_cpu"):
+        backend_config(_cfg(backend, **{key: True}))
 
 
 def test_accepts_an_already_coerced_model() -> None:
@@ -133,3 +157,19 @@ def test_schema_validates_without_any_backend_block() -> None:
     validate, otherwise MasterConfig fails before training starts.
     """
     TypeAdapter(DataPlaneConfig).validate_python(_cfg("simple"))
+
+
+@pytest.mark.parametrize(
+    ("key", "value"),
+    [("use_gdr", True), ("gdr_staging_buffer_mb", 512)],
+)
+def test_pydantic_preserves_flat_gdr_keys_for_migration_error(
+    key: str, value: object
+) -> None:
+    validated = TypeAdapter(DataPlaneConfig).validate_python(
+        _cfg("mooncake_cpu", **{key: value})
+    )
+
+    assert key in validated
+    with pytest.raises(ValueError, match="moved under data_plane.mooncake_cpu"):
+        backend_config(validated)