[PyTorch] [torch.compile] torch.compile support for LayerNormLinear and LayerNormMLP

[pggPL]

Description

Register LayerNormLinear and LayerNormMLP as torch.compile custom ops on top of the framework merged in NVIDIA#3053, so both trace under fullgraph=True with the FP8 / MXFP8 / NVFP4 recipes. Part 3 of 3: sits on #44 (framework + module groundwork), which sits on #42 (module restructuring). This PR is only the compile wiring itself.

Type of change

• New feature (non-breaking change which adds functionality)

Changes

layernorm_linear.py / layernorm_mlp.py (mirrors Linear)

• compile_unsupported_reason gate on the fwd args.
• Data-free _*_forward_fake / _*_backward_fake on TensorSpec, mirroring the impls' saved-tensor layout and alias dedup.
• register_custom_op(...), @no_torch_dynamo eager wrapper (_*_eager), compiled dispatch + _compile_eager_fallback_reason / _forward_eager_fallback in Module.forward; @no_torch_dynamo comes off Module.forward.
• Eager fallbacks (graph break under fullgraph=True): debug, quantized input, fsdp_group, differentiable fp8_output, CPU offload, delayed wgrad, FP8 weight caching (is_first_microbatch), fuse_wgrad_accumulation, non-opaque quantizers; MLP additionally checkpoint=True and gemm_gelu_fusion under FP8.

Tests (tests/pytorch/test_torch_compile.py): recipes x compile modes (default, reduce-overhead) for both modules, norm / activation / return_layernorm_output / zero_centered_gamma variants, FP8 primary weight, no-bias and frozen-weight slots, dynamic shapes, eager fallbacks.

Testing (RTX Ada, single GPU)

• tests/pytorch/test_torch_compile.py: 160 passed, 46 skipped, 1 xpassed (43 new LN tests)
• Eager regression on this tree: test_sanity.py -k "layernorm_linear or layernorm_mlp" 7344 passed, test_numerics.py -k "layernorm_linear or layernorm_mlp" 738 passed, test_backward_override.py 1316 passed, test_cuda_graphs.py -k "layernorm or mlp" 105 passed
• pylint on the touched files: 10/10

Not covered on a single GPU: TP/SP and Userbuffers under compile.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes