feat(cuda-graphs): capture fixed-capacity MoE execution

[yaoyu-33]

What does this PR do ?

Adds a fixed-capacity moe CUDA Graph scope on top of #2917. The captured boundary starts after routing and includes token dispatch, local Transformer Engine GroupedLinear experts, and token combine. Attention and MoE scopes can be enabled together with cuda_graph_modules: [attn, moe] under FSDP2.

This is part of #1027. It is built on #2917, whose scoped partial-capture implementation credits @hemildesai as co-author.

Changelog

• Add cuda_graph_modules: [moe] and cuda_graph_moe_capacity_factor configuration.
• Support EP=1 with local drop-and-pad permutation, fixed TE expert splits, and local unpermutation, including multi-rank FSDP2 jobs.
• Support EP>1 through HybridEP with Megatron-style per-expert drop-and-pad routing.
• Set the static HybridEP receive size to source_capacity * ep_size * num_local_experts, avoiding tokens_per_expert.sum() and its host synchronization.
• Pass attention and TE expert parameters as explicit graph inputs and guard replay with the existing tensor/control contract checks.
• Unshard the owning FSDP2 transformer block while capturing attention or MoE, then replay with the live materialized parameters inside the block pre-forward lifecycle.
• Keep BF16, activation checkpointing, dynamic shapes, and unsupported dispatcher combinations fail-closed.

Design notes

For N source tokens, top-k K, E global experts, and capacity factor F, source capacity is:

ceil(N * K * F / E)

Every source rank selects exactly that many token slots for every global expert. Real assignments are selected by routing probability; unused slots have zero routing probability. Therefore every local expert receives exactly source_capacity * ep_size rows and TE receives constant m_splits.

This follows the Megatron Core moe_pad_expert_input_to_capacity / HybridEP path. Legacy DeepEP remains rejected: its current dispatch path waits on dynamic communication metadata on the CPU, and padding only after dispatch cannot make capture safe. A DeepEP v2 static-buffer integration would be separate work.

CUDA experiments

Combined-scope attribution

Tiny Qwen3-MoE, 2 layers, hidden size 512, 8 experts, top-2, BF16, sequence length 512, local/global batch size 2/4, FSDP2 over 2 H100s. Timing is 8 measured forward/backward/Adam steps after 4 warmup steps.

Scope	Iteration	Latency vs eager	Throughput vs eager	Peak allocated	Graph stats
eager	27.072 ms	-	-	0.35 GB	-
attn	24.882 ms	-8.1%	+8.8%	0.53 GB	captured=2, replayed=22, fallback=0
moe	23.702 ms	-12.4%	+14.2%	0.40 GB	captured=2, replayed=22, fallback=0
attn, moe	20.920 ms	-22.7%	+29.4%	0.60 GB	captured=4, replayed=44, fallback=0

All four modes report the same 12-step loss sequence. This tiny-model run is useful for attribution, but launch overhead is overrepresented.

Qwen3-30B-A3B low-load sanity check

Qwen/Qwen3-30B-A3B, 48 layers, BF16, sequence length 1024, local/global batch size 1/8, FSDP2 + EP8 HybridEP over 8 H100s, real learned router, 25 forward/backward/Adam steps with 5 warmup steps. The order was eager 1, graph 1, graph 2, eager 2. Setup, TE compilation, and graph capture are excluded from steady-state iteration timing.

Mode	Run 1	Run 2	Mean	Mean MFU	Peak allocated
eager	457.622 ms	455.914 ms	456.768 ms	4.39%	32.56 GB
attn, moe	304.391 ms	300.639 ms	302.515 ms	6.63%	53.53 GB

Mean latency decreases by 33.8%; equivalent training throughput increases by 51.0%. Each graph run captured all 96 entries (attention + MoE for 48 layers), replayed 2,304 calls, and had zero fallback. This is a deliberately underloaded launch-overhead result (4.39% eager MFU), not a production-throughput claim.

The memory cost is material: +20.97 GB/GPU in this configuration. Capture is also a one-time startup cost and is not included in the steady-state speedup.

For a controlled replay parity check, the same 30B configuration was run with lr=0 for 6 steps. Eager and graph produced the same reported loss on every step: 10.3457, 10.3424, 10.3292, 10.3191, 10.3709, 10.3198. The graph run captured 96 entries, replayed 480 calls, and had zero fallback. With nonzero learning rate, repeated distributed HybridEP/BF16 runs are not bitwise deterministic, so the controlled no-update run is the forward/replay parity check.

Existing-recipe scaling check

The repository H100 recipe uses local batch 4, sequence length 4096, TE projections/norm/RoPE, DeepEP, and activation checkpointing. It cannot be used as a pure graph-on/graph-off comparison yet: whole-attention graph currently requires torch projections, torch RMSNorm, and unfused RoPE; full MoE graph requires HybridEP and rejects activation checkpointing.

Using the closest currently supported backend on 8 H100-80GB GPUs (linear: torch, rms_norm: torch, rope_fusion: false, TE attention/experts, HybridEP, FSDP2 + EP8), the doubled-batch result is:

Scope	Local batch x sequence	Iteration	Latency vs eager	Throughput vs eager	Peak allocated	Graph stats
eager	2 x 1024	481.675 ms	-	-	39.00 GB	-
attn	2 x 1024	441.681 ms	-8.3%	+9.1%	45.50 GB	captured=48, replayed=672, fallback=0
moe	2 x 1024	405.294 ms	-15.9%	+18.9%	54.00 GB	captured=48, replayed=672, fallback=0
attn, moe	2 x 1024	OOM during capture	-	-	75.43 GiB process usage	failed at layer 45 MoE

At local batch 2 x sequence 2048, eager runs at 539.834 ms (15.80% MFU, 53.61 GB), but combined capture OOMs earlier at layer 27 attention with 75.37 GiB process usage. At 2 x 1024, combined capture reaches layer 45 before OOM. This identifies cumulative graph-pool memory across 96 per-layer entries as the immediate scaling blocker. Both 48-entry single scopes fit and replay without fallback.

Validation

• uv run ruff format .
• uv run ruff check --fix .
• uv run pytest -q tests/unit_tests/moe tests/unit_tests/recipes/test_partial_cuda_graphs.py
  • 539 passed, 54 skipped.
• Slurm job 5579217: 2-H100 FSDP2 eager/attention/MoE/combined attribution.
• Slurm job 5579222: 8-H100 Qwen3-30B-A3B symmetric repeated E2E benchmark.
• Slurm job 5579279: 8-H100 Qwen3-30B-A3B controlled lr=0 replay parity.
• Slurm jobs 5579426 and 5579467: recipe-derived combined-capture memory scaling at 2 x 2048 and 2 x 1024.
• Slurm job 5579620: recipe-derived attention-only and MoE-only attribution at 2 x 1024.

Known limitations

• Fixed input shape and explicit opt-in only.
• BF16 TE experts only; FP8 is rejected.
• Activation checkpointing and pipeline parallelism are rejected for the full MoE scope.
• Legacy DeepEP is not graph-safe and is rejected; EP>1 currently requires HybridEP.
• Whole-attention graph currently requires torch projections, torch RMSNorm, and unfused RoPE, while the existing H100 recipe uses the TE variants.
• The Qwen3-30B run validates FSDP2 for the transformer blocks and EP8 distribution for experts. Additional FSDP sharding of each rank's local expert weights (ep_shard > 1) is not validated.
• CUDA Graph pools increase persistent GPU memory; the 30B combined run used 53.53 GB/GPU versus 32.56 GB/GPU eager.
• Combined attention + MoE capture does not fit at local batch 2 even with sequence length 1024 on H100-80GB; graph-pool sharing/reuse is required before matching existing recipe operating points.

Before your PR is "Ready for review"

• Read and followed contributor guidelines.
• Added unit coverage for config, capacity/drop-and-pad metadata, static HybridEP sizing, combined scopes, and FSDP2 graph lifecycle.
• Add user documentation after the API and supported matrix are agreed.

Additional Information

• Based on perf(moe): add scoped partial CUDA graphs #2917
• Related to (partial) Cudagraph + MoE #1027

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