[JF-D]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Generalizes fixed-address persistent storage for GTP backward operations captured by local CUDA graphs. It also enables NCCL symmetric-memory one-shot reduce-scatter for GTP + local CUDA graphs.

Problem

Local CUDA graph replay requires every graph-owned communication buffer to retain a stable address. The existing specialized wgrad-ring mechanism has several limitations:

• it only manages padded wgrad reduce-scatter inputs;
• it requires an explicitly configured ring size;
• it does not cover asynchronous reduce-scatter outputs or FP32-accumulation scratch;
• it cannot naturally allocate graph-owned wgrad inputs from the exact registered symmetric-memory domain required by dense GTP or expert GTP one-shot reduce-scatter;
• repeated parameter usage, including MTP, can require multiple distinct buffers within one graph.

Consequently, adding another graph-owned GTP tensor requires another specialized lifetime mechanism, while symmetric one-shot RS cannot safely reuse ordinary graph storage.

Solution

This PR replaces the specialized wgrad ring with an automatically sized persistent-buffer plan:

1. Eager backward discovers each CUDA graph runner’s aligned storage requirements, grouped by allocation domain and dtype.
2. Before capture, the plan allocates a bounded number of alternating fixed-address arenas outside the CUDA graph memory pool.
3. Capture bump-allocates stable tensor views from the runner’s assigned arena.
4. Replay-time CUDA events prevent an arena generation from being reused while earlier graph work is still in flight.
5. Capture validates that it consumes exactly the capacity discovered during eager execution.

Symmetric reduce-scatter inputs are allocated from the registered memory pool for their exact GTP or expert-GTP process group. Dense and expert GTP therefore cannot alias storage registered for a different communicator, allowing NCCL one-shot RS to work with local CUDA graphs.

Persistent storage coverage

The plan now manages:

• padded wgrad reduce-scatter inputs;
• graph-owned asynchronous reduce-scatter outputs;
• FP32 reduce-scatter all-to-all scratch;
• FP32 accumulation output scratch;
• repeated GTP parameter usage, including MTP.

Ungraphed execution continues to use the existing eager buffer pools and cache paths.

The centralized discovery, allocation, and replay-fencing API can also support additional graph-owned GTP storage in the future. This PR does not add CUDA graph offloading support.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Convergence

1/4 Nemotron-4 Proxy, 8 nodes, blue: eager, red: partialcg, no mtp.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[JF-D]

/claude strict-review

[JF-D]

@claude strict-review

[JF-D]

/claude review

[fanshiqing]

/claude strict-review

[JF-D]

This PR will be split into 3 PRs. Two split PRs before this one: #6862, #6867. This PR is a stack version and will be updated to support oneshotRS after previous PRs are merged.