[PyTorch] [torch.compile] Prepare the custom-op framework and LayerNorm modules for compile wiring

[pggPL]

Description

Behavior-neutral groundwork for registering LayerNormLinear and LayerNormMLP as torch.compile custom ops (framework from NVIDIA#3053). Part 2 of 3: sits on #42 (the module restructuring) and is the base of the wiring PR (#43).

Type of change

• Code refactoring

Changes

dynamo/custom_op.py

• A backward args dataclass may declare GRAD_OUTPUT_FIELDS, naming one grad field per user output. The LN modules have two differentiable outputs (out and the optionally returned norm output), so a single grad_output slot is not enough.
• Dict[str, <simple>] annotations are bundle-simple (needed for activation_params).

module/_common.py: sp_out_leading / sp_inp_leading / fake_workspace_valid shared by Linear and the LN modules (moved out of linear.py, linear.py updated to import them).

layernorm_linear.py / layernorm_mlp.py

• GRAD_OUTPUT_FIELDS on the bwd args; fp8_output carried in the fwd args (the compile gate needs it).
• inp_shape is rederived from grad_output in backward instead of being stored on the bwd args (SymInt dims are not hashable in the value bundle under dynamic=True).
• LayerNormMLP: the recipe object no longer rides on the backward args; the properties the backward needs (float8_block_scaling, custom, dbias-dact fusion availability) are bools computed in Module.forward. The activation tables are split into per-activation (act, dact) pairs plus the fused dbias kernels, so the backward looks them up without a recipe object.
• The returned-norm-output grad is tolerated as None in backward.

Testing (RTX Ada, single GPU)

Same suites as #42, unchanged results: 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; existing test_torch_compile.py (Linear / DPA) unchanged; pylint 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