[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 have a custom op. 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

What an operation implements

A BasicOperation opts in by declaring two argument containers and implementing the methods below. op_forward / op_backward are then inherited from the base class; the operation does not write any custom-op plumbing itself.

Declarations

• fwd_args_type, bwd_args_type: dataclasses. Their fields define the custom-op schema (tensors, quantized tensors, quantizers, plain values), exactly as LinearFwdArgs / LinearBwdArgs do for te.Linear.

Traced part (runs in the Dynamo-traced region, may read self, module state and ctx)

• resolve_fwd_args(input_, *, requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer, <kwargs>) -> fwd_args_type: gathers everything the forward needs into the container. Forward kwargs the operation accepts are declared as additional keyword-only parameters with defaults; fwd_kwarg_names is derived from this signature.
• resolve_bwd_args(ctx, grad_output) -> bwd_args_type: rebuilds the backward's inputs from the saved context.
• setup_context(ctx, args, aux) (optional): decides what to save for backward. The default saves aux. An operation whose backward needs an input or a parameter saves it here from args, because a custom op may not return one of its own inputs.

Custom op implementation (classmethods; read nothing but args, mutate nothing)

• forward_impl(args) -> (output, aux): the forward. aux is a tuple of fresh tensors produced inside the op that the backward needs (norm statistics, a quantized copy of the input). Inputs and parameters are not aux.
• forward_fake(args) -> (output, aux): the same over TensorSpec, allocation-free.
• backward_impl(args) -> (grad_input, *grad_params): the backward. grad_input may be None, meaning "grad_output, unchanged".
• backward_fake(args): its TensorSpec twin.

The *_impl methods are called directly on the eager path and are the bodies of the registered custom ops on the compiled path, so there is one implementation of the math. Because they are custom-op bodies, they may not mutate their arguments: no clear_tensor_data on inputs or saved tensors, no in-place writes to tensors from an enclosing scope.

Derived by the base class

• compile_ops: the registered (forward_fn, backward_fn) pair, one registration per class, under the lower-cased class name.
• fwd_kwarg_names: the extra keyword-only parameters of resolve_fwd_args.
• compile_unsupported_reason(): None, or why the operation must run its eager implementation. The base class rejects operations without a custom op and quantizers that are not value-opaque (delayed scaling).

Changes

• register_custom_op now defines an operation's forward and backward as two independent custom ops and returns both, leaving autograd to the caller. 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.
• BasicOperation gains the contract above; __init_subclass__ registers the custom op once per class.
• OperationFuser decides once per group whether to run its operations' custom ops. The group runs its eager implementations when: it holds more than one operation, either pass has a fusion, an operation has extra tensor inputs/outputs, a kwarg is undeclared or not a tensor, or an operation reports compile_unsupported_reason. Only tensor kwargs are accepted because a value kwarg becomes a symbolic scalar on its second value, which the opaque value bundle cannot carry; 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 group of two operations is gated onto the eager implementations;
• a backward-only fusion is gated onto the eager implementations;
• an operation without a custom op still runs, on its eager implementation, under fullgraph=True;
• a quantized tensor forward kwarg reaches the op through its custom op and changes between calls without a recompilation, while a value kwarg is gated onto eager.

test_torch_compile.py + test_fusible_ops.py (without grouped / userbuffers cases): 1616 passed, 1039 skipped, 1 xpassed. 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