feat(mdp): run per-layer CUDA graphs on the MDP decoder

[BestJuly]

Summary

Unlocks per-layer ("partial") CUDA graphs on the MDP decoder. Per-layer
graphs were not structurally incompatible with MDP's replay-iterator
mechanism — P4 runs the unmodified decoder schedule and the bridge only
touches embedding leaves — but capture had never actually reached a working
replay before this branch.

Built on top of --thd-static-packing / --mdp-greedy-packing (the
preceding PR): per-layer graph capture needs a batch-independent shape
contract, which is exactly what that PR provides.

One commit, 3c8e72cd5.

What shipped

• Shared thd_shapes_are_static(config) predicate replacing four conflated
  call sites (transformer/module.py, transformer_config.py x2,
  training/arguments.py).
• MDP's CUDA-graph gate rewritten: full_iteration graphs stay rejected,
  per-layer graphs are permitted under thd_static_packing
  (megatron/core/mdp/config.py:_validate_cuda_graph_options).
• _discover_layers fixed to forward through model-wrapper properties and to
  fail loudly (not silently) when every chunk's decoder is unreachable
  (examples/multimodal_dev/models/base.py,
  megatron/core/transformer/cuda_graphs.py).
• pad_between_seqs derived from the collator's row alignment instead of
  forced True — the change that actually makes attn graphs replay (see
  below).

The bug: v1 silently captured 0 layers

The original attempt at this (an earlier, now-superseded branch) never
reached replay because _discover_layers only unwrapped .module to find
the decoder, while under MDP the decoder lives at
self.language_model — a different wrapper path. Capture therefore silently
found 0 graphable layers on every rank, and CUDA graphs were a no-op.
_discover_layers now forwards through the actual wrapper chain, and a new
loud failure (test_unreachable_decoder_fails_loudly) makes a future
zero-layer regression a hard error instead of a silent no-op. On this branch:
INFO:...cuda_graphs:Rank 0: 4 graphable layers. — all 4 decoder layers, on
every rank.

The other bug: pad_between_seqs=True was the real blocker, not a second-order cost

_reconstruct_packed_seq_params_from_kwargs hardcoded
pad_between_seqs=True, reasoning that the value can't be a graph input and
inferring it from CUDA tensors would synchronize during capture. With that,
attn graphs OOM'd at 32 GiB (T=8192): forcing True makes FlashAttention
ineligible, and on this configuration (head_dim 256, TP=CP=1) cuDNN fused
attention declines too, so TE fell through to
UnfusedDotProductAttention's O(T²) torch.baddbmm. The value doesn't have
to be guessed: under --thd-static-packing it's a function of the
collator's row alignment alone (thd_collate_row_alignment), which is
batch-independent by construction — exactly what a graph needs. At
TP=CP=1 the alignment is 1, no sample is ever padded, cu_seqlens and
cu_seqlens_padded coincide, and the honest answer is False. With it
derived, FlashAttention is eligible and attn graphs replay. The conservative
True is retained for every configuration that doesn't set
thd_static_packing.

moe_router fusion: originally blocked, now confirmed working

At review time, --cuda-graph-modules moe_router with
--moe-router-fusion (default on) failed capture inside
fused_moe_aux_loss_fwd with cudaErrorStreamCaptureInvalidated — confirmed
unrelated to MDP (a plain pretrain_gpt.py decoder-only THD run reproduced
the identical failure), and attributed to a TE-side dynamic-shape gap in the
fused aux-loss kernel under stream capture.

Retested 2026-08-21 on an image with TE 2.18.0+27486e0 (up from
2.16.0.dev0+cf64ef56 at review time), attributed to
Megatron-LM PR #4359
adding dynamic-shape support to the fused router: moe_router now captures
cleanly with the real fused aux_loss router (no
--moe-router-load-balancing-type none workaround needed), and so does the
full attn moe_router moe_preprocess scope — MDP enabled throughout
(--mdp-enable).

Performance (4x GB300, MDP enabled, lognormal lengths, budget 8192, cap 16, TP=PP=CP=1, DP=4, 4 decoder layers / 8 experts top-2, 20 iterations, mean over iterations 6-20)

	mean ms/iter	TFLOP/s/GPU	vs eager (static pad)
eager, no static pad	217.6	290.4	—
eager, static pad (baseline)	214.5	294.6	—
graph attn	193.9	325.5	-9.6% / +10.5%
graph moe_router (fused)	192.6	329.2	-10.2% / +11.8%
graph attn + moe_router	172.5	367.3	-19.6% / +24.7%
graph attn + moe_router + moe_preprocess (max scope)	171.4	369.4	-20.1% / +25.4%

moe_router alone reproduces the eager baseline's loss/grad-norm bit for
bit; every graph combination stays within ~3.5e-4 relative loss of eager over
20 iterations. moe_preprocess only adds ~0.6% beyond attn+moe_router
alone on this proxy (its own op is cheap here); a production model may see a
larger contribution. mlp graphs are inapplicable to this proxy (every
layer is MoE — --num-experts 8, no --moe-layer-freq).

Recommendation: enable the maximal scope,
--cuda-graph-modules attn moe_router moe_preprocess, together with
--thd-static-packing, for MDP at TP=CP=1.

Strict review

strict_review_v2.md (base 5200c89a3, head reviewed = this PR's tip
3c8e72cd5, i.e. reviewed at the exact commit in this PR): 0 CRITICAL,
1 IMPORTANT (an all-chunks-zero guard in _discover_layers that would have
turned legitimate zero-graphable-layer models into hard failures — fixed to
key off lookup failures instead), 2 SUGGESTION (a row-alignment rule
restated in a third place; an unreachable-today pad_to_multiple desync
risk, now an explicit assert) — all addressed before this commit.

Tests

4x GB300, TE 2.18.0+27486e0:

torchrun --nproc-per-node 1 -m pytest -q tests/unit_tests/transformer/test_thd_static_shape_predicate.py -> 7 passed
torchrun --nproc-per-node 1 -m pytest -q examples/multimodal_dev/tests/test_mdp_cuda_graph.py             -> 4 passed
torchrun --nproc-per-node 4 -m pytest -q tests/unit_tests/mdp/                                            -> 162 passed

Not covered

Production Qwen3.5-VL shape — this proxy is launch-bound, so the
static-pad tax is unmeasurable here (eager_static vs eager_nopad: +1.4%
step time) and will show up more at production layer weight, shrinking the
graph win shown above. CP>1, PP>1, SP+PP>1 (_needs_full_local_padding_mask
starts applying to MDP and is unverified), and seq_aux_loss (only the
default aux_loss router was retested with the newer TE).

No push to NVIDIA/Megatron-LM and no upstream PR is created by this update.