[JF-D]

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

What does this PR do?

Problem

Local CUDA graphs need fixed-address storage for GTP wgrad reduce-scatter inputs whose asynchronous lifetime can extend beyond the graph that produced them.

The existing per-parameter wgrad ring has several limitations:

• its size is configured manually through graph_wgrad_ring_size;
• allocation depends on parameter and scheduling-domain preclassification;
• a graph with more same-key writers than available slots fails during capture;
• the specialized ring lifecycle cannot be reused by other graph-owned temporary buffers.

Solution

Replace the wgrad ring with a runner-scoped persistent-buffer plan:

1. During eager backward, each local CUDA graph runner records its aligned storage requirements by allocator domain and dtype.
2. Before capture, GraphPersistentBufferPlan allocates one arena generation when one local graph module is configured and two generations when multiple modules are configured. One generation remains correct because reuse is event-fenced; two preserves the expected cross-graph overlap. Larger counts are not used because they increase persistent memory.
3. During capture, requests bump-allocate fixed-address views from the runner's assigned arena. Capture validates that it consumes exactly the capacity discovered eagerly.
4. During replay, completion events prevent a generation from being overwritten while an earlier graph using that storage is still in flight.
5. Symmetric-memory reduce-scatter inputs are allocated from the exact registered GTP or expert-GTP process-group pool.

The GTP wgrad reduce-scatter input path now uses this generic API, and the manual graph_wgrad_ring_size configuration and ring-specific API are removed.

Scope

This PR intentionally contains only the persistent-buffer allocation and lifetime refactor. It does not change:

• repeated-MTP gradient-finalization accounting;
• reduce-scatter output ownership;
• FP32 reduce-scatter scratch or accumulation buffers;
• activation offloading;
• GTP communication scheduling.

Those behaviors can be added or fixed independently on top of this foundation.

Validation

• Repository autoformatter checks: Black, isort, Pylint, and Ruff passed.
• Focused 4-GPU tests: 12 passed, covering persistent allocation/reuse, symmetric-memory ownership, and one-/multi-module arena selection.

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

[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 review

[JF-D]

/ok to test 87b89e9

[claude]

Light review — the refactor reads well and the new unit tests cover the interesting cases (distinct slices, symmetric domains, empty runner, reordered replay, capacity validation). Three inline notes plus two things that are better discussed at top level.

1. Discovery state stays globally armed after the last runner's backward

_CudagraphRecordNode.backward arms _ACTIVE_DISCOVERY_STATE for a runner and relies on the next record node to switch it. That means after the final backward runner records, the global remains pointed at that runner's state until _create_cudagraphs() clears it — and create_cudagraphs() runs after finalize_model_grads. Any GTP wgrad request issued in that window (e.g. a GRAPHED-chain deferred cascade, or a second microbatch's backward on a runner whose bwd_graph_recorded flag is already set for other runners) gets recorded into the last runner's capacity.

The failure is loud rather than silent — validate_complete() raises consumed N/M at capture — but the message points at the wrong runner and gives no hint about the extra eager request. Consider disarming at the end of the discovery backward rather than at the start of the next one, e.g. in _CudagraphGlobalRecord.record_bwd_graph returning, or asserting in _create_cudagraphs that the global is already None:

# in _create_cudagraphs, before the plan is built
assert (
    gtp_cuda_graphs._ACTIVE_DISCOVERY_STATE is None
), "persistent-buffer discovery was still armed after backward; a post-backward GTP wgrad request would inflate a runner's capacity"
set_graph_persistent_buffer_discovery(None)

2. _get_gtp_persistent_buffer_max_inflight counts module kinds, not concurrent runners

The generation count is derived from len(config.cuda_graph_modules), but the number of backward runners whose persistent buffers can be live concurrently is a function of how many local graphs overlap at replay, not how many distinct module kinds are configured. With cuda_graph_modules=["attn"] on a 32-layer model there are 32 backward runners, all mapped to generation 0, so every backward serializes on one _ready_events[0] — correct, but it silently removes the cross-graph RS overlap that §3.6 of the doc describes as the point of this machinery.

That may well be the intended trade-off, but the docstring says "one arena generation is enough for one configured local-CG module" which reads as a correctness statement rather than a deliberate overlap sacrifice. Worth either clarifying the comment or basing the count on something closer to the actual overlap depth (the runner-stream pool size is 3, for instance).

3. Missing regression coverage for the whole-layer / cuda_graph_modules=[] case

test_gtp_persistent_buffer_max_inflight_follows_local_graph_scope covers the helper in isolation, and test_gtp_partial_cg.py runs with ["attn"] and ["moe_router"]. Neither exercises max_inflight=1 end to end through create_cudagraphs() + multi-replay, which is exactly the new single-generation path that all-runners-share-one-arena depends on. Since the ring's two-slot default previously guaranteed at least one generation of separation, a parametrization of the existing partial-CG worker with the whole-layer scope (or an explicit max_inflight=1 monkeypatch) would guard the new serialization fence numerically rather than structurally.

[JF-D]

/claude strict-review

[JF-D]

/claude strict-review