[PyTorch] torch.compile for an OperationFuser group holding one operation

[pggPL]

Description

First step towards torch.compile(fullgraph=True) support for transformer_engine.pytorch.ops.

The approach: keep the pipeline-level _OperationFuserAutogradFunction and let Dynamo trace it as a higher-order op, with each fusible operation calling its own custom op inside. That keeps the forward and backward fusion layouts independent (so the backward-only fusions survive), keeps OperationContext inside the traced scope, and bounds op registration to one entry per op class.

This PR makes that work for an OperationFuser group holding one operation. It converts no real operation: the fuser's path is exercised by a test-only operation, so it does not depend on which operations happen to declare their compute halves. Converting the operations themselves is a follow-up.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• register_custom_op now defines an operation's forward and backward as two independent custom ops and hands back both, leaving autograd to the caller. That is what lets a pipeline-level autograd.Function decide how the two are wired. The variant that wires autograd itself keeps the old behaviour under register_custom_op_with_autograd. Both are built on a single-op primitive (_register_op returning a _RegisteredOp), so the pair is two calls and a forward-only op is possible later. The autograd-free forward returns a plain (output, aux) tuple, where aux holds only fresh tensors produced inside the op.
• BasicOperation gains the plumbing an operation needs to opt in: declare two argument containers and implement four compute classmethods, and __init_subclass__ registers the custom ops while op_forward / op_backward are written once in the base. setup_context(ctx, args, aux) decides what to save for backward, so inputs and parameters (which a custom op may not return) are saved from args. compile_unsupported_reason lets an operation say why it cannot be compiled; it sits on the operation rather than on the args because in ops/ the compile boundary is the fuser group.
• OperationFuser runs a supported group through its operations' custom ops under compile. An operation may declare read-only forward kwargs (fwd_kwarg_names); they are resolved into its args container like any other config. Only tensor kwargs are accepted on the compiled path: a value kwarg becomes a symbolic scalar on its second value, which the opaque value bundle cannot carry, so it sends the group to eager deterministically instead. A 0-d tensor is the way to pass a scalar.

Side effects that had to go

All of them reached outside the higher-order op's scope:

• OperationContext objects are created in the forward, but the backward is a separate subgraph, so writing to them there mutates an enclosing scope. The backward copies them into its own scope.
• requires_grad_ on an output: AOTAutograd's functionalization drops it anyway, and autograd marks the outputs of an apply() itself.
• _do_not_clear on inputs and outputs.

They are gated on being traced, not on using the custom ops. Under fullgraph=True there is no leaving the graph, so an unsupported operation does not "fall back": the pipeline is traced either way and only the choice of implementation changes. The gate reports its reason through warn_compile_eager_fallback, which is safe to call from the traced region.

Testing

test_torch_compile.py gains two test-only operations and five tests: a single-operation group compiles and matches eager (output, input gradient, parameter gradient); a backward-only fusion sends the group to the eager implementations; setup_context saves the parameter so a mutation between forward and backward is detected; an operation without the compute halves still runs, on its eager implementation; and a quantized tensor forward kwarg reaches the op through its custom op, changes between calls without a recompilation, while a value kwarg is gated onto eager.

test_torch_compile.py: 123 passed, 46 skipped, 1 xpassed. test_fusible_ops.py unchanged. RTX Ada.

Gated out and untouched: multi-operation groups, fused operations, operations with extra tensor inputs/outputs, grouped operations, userbuffers, delayed scaling, FP8 block scaling.

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