[GTP] Preserve independent expert prefetch chains under CUDA graphs

[fanshiqing]

Description:

Extend #6057’s separate grouped-expert FC1/FC2 prefetch chains to full-iteration and local CUDA graphs.

Previously, captured expert weights shared the general graphed chain: current FC1 prefetched current FC2 instead of next-block FC1, reducing the available communication-overlap window. Additionally, recording gather completion when consumed could make current FC2 wait for unrelated later FC1 communication on the shared all-gather stream.

This change:

• Preserves independent captured FC1/FC2 chains and their buffer ownership.
• Uses consistent captured-chain filtering for communication drains and stream joins.
• Records grouped-gather completion immediately after issue, including synchronous fallbacks without a Work handle.
• Replays gradient-ready hooks on each parameter’s actual reduce-scatter stream.

How is the prefetch chain(s) changed under full-iter CG model

BEFORE: one logical captured chain

  [Dense] -> [Grouped FC1] -> [Grouped FC2] -> [Dense] -> ...
                 all classified as GTP_graphed

AFTER: three logical captured chains

  Dense chain:  [Dense 0] -> [Dense 1] -> [Dense 2] -> ...

  FC1 chain:    [FC1 L0]  -> [FC1 L1]  -> [FC1 L2]  -> ...

  FC2 chain:    [FC2 L0]  -> [FC2 L1]  -> [FC2 L2]  -> ...

Why this fix is critical for GTP+full-iter CG's performance

• Earlier prefetch: The merged chain prefetches the next weight in execution order. Separate chains let FC1/FC2 prefetch their counterparts in the next MoE layer, giving communication more time to overlap computation.
• Fewer unnecessary waits: The completion-event fix lets a GEMM start once its own weights arrive, without waiting for unrelated gathers queued afterward.
• CG makes this overlap more useful: Full-iteration CG removes much of the CPU launch overhead. Better prefetching then helps keep the GPU compute stream busy instead of waiting for weights.

Validation:

• 89 targeted tests passed per rank on four GB300 GPUs: 80 basic GTP, 7 grouped-CG, and 2 partial-CG tests.
• Covers BF16 full/local graph replay, eager/captured completion ordering, synchronous gather fallbacks, and native-MXFP8 SwiGLU replay across six Adam updates.
• Fixed test isolation so captures do not join stale streams from earlier tests.

Perf Results on DSV4-pro-half-proxy: 35% e2e perf gain

Both runs use 256 NVIDIA GB300 GPUs, MXFP8, Adam, and the DSV4-Pro 30-layer proxy + MTP1, with TP1/PP1/CP1/EP64/GTP64/EGTP4, full-iteration CUDA graphs, and profiling enabled.

Run	GBS	MBS	Sequence length	Tokens/s/GPU	Model TFLOP/s/GPU	Step time (ms)	Peak allocated / sampled device memory (GiB)
Eager baseline	512	2	4096	4,222.68	718.93	1,940.00	204.39 / 248.54
full-iter CG, w/o this MR	512	2	4096	3,799.28	632.70	2,156.20	208.65 / 256.83
full-iter CG with this MR	512	2	4096	5,146.16	857.00	1,591.87	210.30 / 257.81

• Throughput improves by 35.45% compared to CG case, , 19% compared to eager mode.
• Both runs use the same image, training arguments and profiling settings. The source difference consists of separate grouped-expert FC1/FC2 prefetch chains under CUDA graphs and recording gather completion before later prefetches.

Related: #6057

[copy-pr-bot]

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.

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

[fanshiqing]

/review

[fanshiqing]

/ok to test 6914806

[fanshiqing]

/ok to test 8b41aaa