[pggPL]

Description

Add full-graph torch.compile support for LayerNormLinear and LayerNormMLP, including backward, dynamic batch dimensions and normalization outputs.

Targets main. This branch includes the NVFP4 packed-buffer, TensorSpec and Bias correctness fixes from NVIDIA/TransformerEngine#3628.

Changes

• Describe forward outputs, saved tensors, aliases and backward gradients for both LayerNorm modules; retain eager fallbacks for unsupported configurations.

• Share output-gradient setup and quantizer validation with Linear. Adapt the newly inherited fused-attention backward container to the same gradient setup contract, preserving its existing first-output gradient assignment.

• Validate both MLP GEMMs, including gated activations and sequence-parallel dimensions.

• Support compiled FP8 weight caching in both LayerNorm modules, following Linear: replace the cache on the first microbatch and reuse it on subsequent microbatches, including both MLP weights. GPU-controlled cache updates and fused wgrad accumulation retain eager fallbacks.

• Share the cache regression across Linear, LayerNormLinear and LayerNormMLP, checking independent eager outputs/gradients, cache refresh/reuse, optimizer updates, immediate/deferred backward and CUDA Graph replay.

• Describe LayerNorm/RMSNorm output, mean and reciprocal standard deviation through dynamo.fake.fake_norm, used by both LayerNorm modules. Share the fused-quantization layout predicate with activation output metadata, including the RHT-fusion disable flag and output device; remove duplicate layout helpers. Preserve MXFP8/cuDNN normalization layouts and gathered LayerNormMLP output shapes.

• Parameterize module compile tests across recipes, shapes and modes, with LayerNorm-specific outputs, gradients and eager fallbacks. Retain the non-matrix Linear regressions and expanded cold CUDA Graph recipe coverage inherited from [PyTorch] Fix NVFP4 metadata in compiled Linear operations NVIDIA/TransformerEngine#3628.

• Preserve the native RMSNorm MXFP8 2D scale layout in fake normalization metadata; cover metadata reconstruction and add CustomRecipe MXFP8 2D cases to the compiled normalization scale-layout matrix.

Validation

2026-10-23 RMSNorm follow-up: the new metadata regression detects two failures on the old helper. The corrected selection passes 113 tests, with 83 hardware skips on Ada. The final helper also passes all 32 metadata cases after the lint-only condition refactor. Native MXFP8 2D compiled forward/backward cases are added for supported hardware but remain unverified on Ada. Existing native libraries reused.

2026-10-23, workstation RTX 5880 Ada, PyTorch 2.15.0a0+875d815502.nvinternal.main; existing native libraries reused, no fresh native build:

• After merging NVIDIA/main dd02f3f7: fake-versus-real outputs, LayerNorm variants and inherited Linear cache selection: 85 passed, 15 skipped.
• Final cache and eager-fallback selection: 21 passed, 36 hardware skips. The cache tests cover all three modules, five minibatches with optimizer updates, three microbatches each, both compile modes and immediate/deferred backward. Both MLP caches are checked. Float8CurrentScaling passes; Float8BlockScaling, MXFP8 and NVFP4 cache cases require supported hardware and remain unverified for this change.
• Black 24.4.2, changed-module pylint, license, AST and whitespace checks pass. Full Python lint retains the pre-existing fused_mla_q_uproj.py:19 C0411 import-order diagnostic.

Historical validation before this merge/cache extension:

B100/SM100, PyTorch 2.15.0a0+875d815502.nvinternal.main, CUDA 13.4; existing October 7 native libraries reused, no fresh native build.

• Focused LayerNorm compile selection: 180 passed, 4 expected skips for unsupported LayerNormMLP backward overrides. Includes full-graph forward/backward, non-matrix inputs, normalization output gradients, variants and eager fallbacks.
• The selection includes 60 normalization scale-layout cases: both modules, LayerNorm/RMSNorm, MXFP8 and NVFP4, cuDNN settings, shapes, returned norm output, and NVFP4 RHT fusion enabled/disabled.
• The inherited cold CUDA Graph regression is identical to the version validated in [PyTorch] Fix NVFP4 metadata in compiled Linear operations NVIDIA/TransformerEngine#3628: 48 passed with an independent eager model and no implicit recipe construction for unquantized inputs.
• Pre-commit, license, syntax and whitespace checks passed. Full PyTorch Python lint retains an unchanged C0411 import-order diagnostic in attention/fused_mla_q_uproj.py in the GPU container; no unrelated source cleanup is included. Pylint on all changed production files passed.

No complete compile-suite rerun or distributed TP/SP/Userbuffers validation. The gathered-output test emulates the collective and does not establish multi-GPU correctness.