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

[pggPL]

Description

Restructure LayerNormLinear and LayerNormMLP the same way Linear was restructured in NVIDIA#2967, so that the follow-up PR registering them as torch.compile custom ops (on top of the framework merged in NVIDIA#3053) is a pure wiring change.

No functional change intended. Part 1 of 3 (then #44 groundwork, #43 compile wiring).

Type of change

• Code refactoring

Changes

• LayerNormLinearFwdArgs / LayerNormLinearBwdArgs and LayerNormMLPFwdArgs / LayerNormMLPBwdArgs dataclasses replace the positional non_tensor_args tuple and the ad-hoc ctx.<attr> plumbing.
• The autograd Functions become thin wrappers around module-level _layernorm_{linear,mlp}_forward_impl / _setup_ctx / _backward_impl that take one dataclass argument (the shape the custom-op framework expects). Saved tensors that alias forward inputs or outputs are deduplicated by name (saved_tensor_aliases) and rebuilt in setup_ctx.
• LayerNormMLP keeps activation checkpointing: the forward saves only its inputs, and _layernorm_mlp_recompute re-runs the forward impl in backward with recompute_for_bwd set, filling the backward args from the recomputed state.
• Weight workspaces move from positional autograd inputs into the args bag and are returned as extra outputs (as in Linear). Backward-only recipe knobs (dgrad/wgrad_use_split_accumulator, backward_override, debug UB flags, delayed wgrad_store) are resolved in Module.forward.
• check_fp8_reduce_and_update moves to module/_common.py and is shared with Linear.

Testing (RTX Ada, single GPU)

• tests/pytorch/test_sanity.py -k "layernorm_linear or layernorm_mlp": 7344 passed
• tests/pytorch/test_numerics.py -k "layernorm_linear or layernorm_mlp": 738 passed
• tests/pytorch/test_backward_override.py: 1316 passed
• tests/pytorch/test_cuda_graphs.py -k "layernorm or mlp": 105 passed
• test_gpt_full_activation_recompute (checkpoint recompute path): same pass/fail pattern as main (known Ada flakiness)
• pylint on the touched files: 10/10

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