feat(data_plane): add Mooncake GDR support by zyzhou5 · Pull Request #3501 · NVIDIA-NeMo/RL

zyzhou5 · 2026-08-05T18:14:04Z

What does this PR do?

Adds the minimal NeMo-RL configuration plumbing needed to enable TransferQueue v0.1.9's Mooncake GPUDirect RDMA (GDR) path. TransferQueue v0.1.9 is already pinned on current main by #3423.

adds optional data_plane.use_gdr; omitting it defaults to false, so existing recipes are unchanged;
passes protocol: rdma and use_gdr: true to Mooncake when GDR is enabled;
leaves RDMA device selection local to each Mooncake client unless explicitly configured;
requires CUDA to be initialized before a GPU worker attaches its GDR-enabled TQ client, avoiding TQ's CPU fallback;
retains NeMo-RL's existing transfer/decode paths and TQ's default staging-buffer size; and
documents how users select TCP, CPU RDMA, or GPU RDMA.

Persistent queued Mooncake objects remain in CPU RAM. GDR changes the CUDA client's transfer path by using a registered GPU staging buffer; it does not turn the persistent queue into GPU storage.

User configuration

data_plane:
  enabled: true
  impl: transfer_queue
  backend: mooncake_cpu
  use_gdr: true

Relationship to #2935

This PR does not stack on or require #2935. The TQ dependency update is already on main via #3423. The runtime-environment cleanup proposed by #2935 is separate and is not included here.

Relationship to #3837

To use GDR on the put side instead of falling back to CPU RDMA, we need #3837 to get merged since TQ must be initialized inside the generation worker after CUDA is initialized so it selects GDR for PUTs.

E2E test

Backend	W&B run
Simple	View run
CPU RDMA	View run
GDR	View run

copy-pr-bot · 2026-08-05T18:14:07Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

zyzhou5 · 2026-08-07T22:49:23Z

/ok to test d2edefb

zyzhou5 · 2026-08-17T21:16:43Z

/ok to test 31f9117

zyzhou5 · 2026-08-26T17:44:18Z

/ok to test 5bdc8e7

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

zyzhou5 · 2026-09-08T18:30:26Z

Self-review with `/review-pr-team`

Run with the repo's multi-agent review skill — RL-codebase, bug-finder, test and design reviewers, plus an adversarial devil's-advocate pass that cut roughly two thirds of the raw findings. Everything was re-verified against the pinned TransferQueue source (c516143). Five inline comments on this PR.

Note on the first comment. It was written before the rebase, so it shows as outdated. That one is already fixed: the branch is now rebased onto cbbb57e7a, and _gdr_requested no longer rides on _promote_1d, which main deleted in #2957 — it keys off the backend directly.

What checks out. All three claims the new README prose makes about TransferQueue are accurate — the GDR staging buffer really is serialized behind a single lock, split_by_bytes really does group aggregate fetches, and the CPU PUT path really does omit the chunk metadata an oversized GDR GET needs. gdr_staging_buffer_mb: PositiveInt is a good catch: TQ treats 0 as "silently fall back to CPU RDMA even with use_gdr=True", and this turns that into a config-time error instead. The CUDA guard in TQWorkerMixin is on the right seam — all four non-mixin client-construction sites are genuinely CPU-only, so test_cpu_only_client_may_attach_with_gdr_config is a correct spec rather than a gap. The two exemplar YAMLs are in sync, and the new test file's pytest marks and CI shard are right.

Still outstanding: CI has never run on this code. The old head 480ccc707 had 5 check-runs, all metadata; 5bdc8e7ec had 83 including the full L0 matrix and lint, all green. The last /ok to test names 5bdc8e7, so the fail-closed commit has never been validated — and the rebased head needs a fresh /ok to test.

Method caveat. Nothing was executed: the review host has no uv, pytest, or importable torch/pydantic. Every finding is static reading of this branch plus pinned TransferQueue source, and anything that would need a run to confirm is labelled as such inline.

Several candidate findings were investigated and deliberately dropped rather than posted: three docstring asks (no ruff rule requires a Raises: section, and D417 is ignored here), a suspected missing pydantic import (it is imported, and pre-dates this PR), a YAML comment-alignment nit (lint was green on the commit containing it), a hardening ask against a hypothetical future TransferQueue rename (all attribute names are correct at the current pin), and several items with no constructible failing case.

Generated by Claude Code

zyzhou5 · 2026-09-08T18:33:01Z

/ok to test 1376417

The check ran after `tq.kv_batch_put`, by which point TQ has already written the rows into Mooncake and called `notify_data_update`, so raising could not prevent the fallback it was meant to catch. It only denied the caller its `KVBatchMeta` — and at the one call site where it was reachable, `column_io.write_columns` via `TQWorkerMixin._write_back`, there is no rollback, so the committed rows were left orphaned. It was also unreachable as a genuine failure. The pre-PUT check already tests TQ's own routing condition (`use_gdr and _gdr_staging is not None`), and `_put_tensors_gdr` then calls `lazy_init()` unconditionally, which sets `_initialized` or raises. The only way to trip the post-PUT check was an upstream rename of a TransferQueue private attribute, which would have aborted a healthy GDR run with a message describing a fallback that never happened. The pre-PUT check is kept: it runs before anything is written, so it still fails closed on a real CPU-RDMA fallback, which is what this commit's predecessor set out to do. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

ZhiyuLi-Nvidia · 2026-09-08T18:51:46Z

+- **`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


Do you think we should change the name?

Option one

moocnake

cpu

gdr

Option 2

mooncake_cpu

mooncake_gdr

This sounds reasonable to me. But maybe we should have a separate renaming PR since it touches a lot of files.

created an issue to track this #4049

ZhiyuLi-Nvidia · 2026-09-08T19:00:49Z

@zyzhou5 awesome PR!

ZhiyuLi-Nvidia · 2026-09-08T19:03:14Z

@zyzhou5 I remember you mentioned we need to some changes in both put and get side to avoid unnecessary H2D or D2H.

I think the best we can do is
vllm/put(gpu tensor) ===GDR===> TQ cpu storage ====GDR===> policy/get (gpu tensor)

Have we achieved that?

`test_gdr_tensor_put_is_confirmed_once_and_never_falls_back` was added without updating the section, which claims to audit every test function under `tests/unit/data_plane/`. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

…fer_mb The knob reserves its full size in HBM per active GDR client for the life of the run, and it has a hard lower bound: below the largest per-sample, per-field payload TransferQueue splits the value into chunk sub-keys, and a GDR read of a value written by a CPU-RDMA client then fails on an assertion inside TransferQueue that carries no message. Sizing also has a knee. An aggregate fetch larger than the buffer is split into groups that transfer one at a time, so lowering it multiplies the group count and costs throughput on both the CPU and GDR paths, while raising it above the default buys little and costs HBM linearly. GDR's gain shows up only once the buffer is large enough that group count is not the bottleneck. None of this was stated anywhere a user would look. The exemplar config is where the value actually gets set, so the guidance goes there; the transfer mechanics stay in nemo_rl/data_plane/README.md. Also aligns the two GDR comments to the block's existing column. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

zyzhou5 · 2026-09-08T21:01:56Z

@zyzhou5 I remember you mentioned we need to some changes in both put and get side to avoid unnecessary H2D or D2H.

I think the best we can do is vllm/put(gpu tensor) ===GDR===> TQ cpu storage ====GDR===> policy/get (gpu tensor)

Have we achieved that?

not yet... the PUT operation starts from CPU so there is an extra H2D. The PUT inference-side fix would need a bigger change so should be placed in another PR. Created an issue to track it #4051.

The flat `data_plane.use_gdr` spelling never shipped — `git log -S"use_gdr"` at this branch's base returns no commits, and it appears nowhere on main — so no merged config could ever trip this check. Making it fire also required declaring `use_gdr` and `gdr_staging_buffer_mb` as `NotRequired` on `DataPlaneConfig`, which advertises them to any reader as legal top-level keys: the opposite of the intent. It also left two policies for one class of mistake three lines apart, since the pre-nesting flat sizing keys are silently ignored. `backend_config` now has one rule: only the nested block is read. Removes the two declarations, `_FIRST_GDR_FLAT_KEYS`, the raise, the two docstring passages, the README sentence, and the two tests that pinned it. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

`setup_data_plane` takes `DataPlaneRuntimeConfig`, which is `DataPlaneConfig | LocalDataPlaneConfig`. The guard subscripted `cfg` directly, so a `LocalDataPlaneConfig` — a BaseModel, not a mapping — would raise `TypeError` before ever reaching `build_data_plane_client`. The guard was written when `setup_data_plane` still took the narrower `DataPlaneConfig`; main widened the parameter, and the two changes never touched the same lines, so the rebase merged cleanly and broke the semantics. `pyrefly` caught it: Argument `LocalDataPlaneConfig | TypedDict[DataPlaneConfig]` is not assignable to parameter `cfg` with type `TypedDict[DataPlaneConfig]` Narrow with `isinstance` first, matching how `factory.py` handles the same union. The local plane has no TQ and no mooncake, so it has no GDR to order against a CUDA context. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

zyzhou5 · 2026-09-08T21:30:04Z

/ok to test 9896a67

zyzhou5 requested review from a team as code owners August 5, 2026 18:14

zyzhou5 marked this pull request as draft August 5, 2026 18:14

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 5, 2026

zyzhou5 force-pushed the zezhou/tq-mooncake-gdr branch from 8fae9a1 to 6e39a2b Compare August 7, 2026 01:57

aoshen02 mentioned this pull request Aug 7, 2026

feat(vllm): support KV-transfer connectors with external store reset on refit #3530

Open

zyzhou5 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 7, 2026

zyzhou5 marked this pull request as ready for review August 7, 2026 22:27

zyzhou5 force-pushed the zezhou/tq-mooncake-gdr branch from 6e39a2b to d2edefb Compare August 7, 2026 22:37

zyzhou5 force-pushed the zezhou/tq-mooncake-gdr branch from d2edefb to 31f9117 Compare August 17, 2026 21:13

zyzhou5 requested a review from a team as a code owner August 17, 2026 21:13

github-actions Bot added the CI Relating to CI label Aug 17, 2026

zyzhou5 force-pushed the zezhou/tq-mooncake-gdr branch from 31f9117 to 5bdc8e7 Compare August 26, 2026 17:42

github-actions Bot removed the CI Relating to CI label Aug 26, 2026

zyzhou5 commented Sep 8, 2026

View reviewed changes

Comment thread nemo_rl/data_plane/adapters/transfer_queue.py Outdated

Comment thread nemo_rl/data_plane/adapters/transfer_queue.py

Comment thread tests/unit/data_plane/README.md Outdated

Comment thread nemo_rl/data_plane/interfaces.py

Comment thread nemo_rl/data_plane/interfaces.py Outdated

zyzhou5 added 2 commits September 8, 2026 17:58

feat(data-plane): add Mooncake GDR support

0f7b614

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

fix(data-plane): fail closed on GDR PUT fallback

1376417

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

zyzhou5 force-pushed the zezhou/tq-mooncake-gdr branch from 480ccc7 to 1376417 Compare September 8, 2026 18:26

ZhiyuLi-Nvidia reviewed Sep 8, 2026

View reviewed changes

Comment thread nemo_rl/data_plane/README.md

ZhiyuLi-Nvidia reviewed Sep 8, 2026

View reviewed changes

Comment thread nemo_rl/data_plane/README.md

docs(data-plane): list the fourth GDR test in the unit-test audit

d3b727a

`test_gdr_tensor_put_is_confirmed_once_and_never_falls_back` was added without updating the section, which claims to audit every test function under `tests/unit/data_plane/`. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

zyzhou5 mentioned this pull request Sep 8, 2026

TQ: reconsider the mooncake_cpu backend name now that GDR exists #4049

Open

zyzhou5 mentioned this pull request Sep 8, 2026

TQ: GDR PUTs stage CPU-resident tensors through the GPU buffer #4051

Open

zyzhou5 added 2 commits September 8, 2026 21:11

Uh oh!

Conversation

zyzhou5 commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

User configuration

Relationship to #2935

Relationship to #3837

E2E test

Uh oh!

copy-pr-bot Bot commented Aug 5, 2026

Uh oh!

zyzhou5 commented Aug 7, 2026

Uh oh!

zyzhou5 commented Aug 17, 2026

Uh oh!

zyzhou5 commented Aug 26, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zyzhou5 commented Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Self-review with /review-pr-team

Uh oh!

zyzhou5 commented Sep 8, 2026

Uh oh!

ZhiyuLi-Nvidia Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

zyzhou5 Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

zyzhou5 Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia commented Sep 8, 2026

Uh oh!

ZhiyuLi-Nvidia commented Sep 8, 2026

Uh oh!

zyzhou5 commented Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zyzhou5 commented Sep 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zyzhou5 commented Aug 5, 2026 •

edited

Loading

zyzhou5 commented Sep 8, 2026 •

edited

Loading

Self-review with `/review-pr-team`

zyzhou5 commented Sep 8, 2026 •

edited

Loading