[PyTorch] Compile Linear and Bias operations with forward fusion

[pggPL]

Description

Enable torch.compile for fusible BasicLinear and Bias operations, including multi-operation pipelines and Linear + Bias forward fusion. For example, te.ops.Sequential(te.ops.BasicLinear(...), te.ops.Bias(...)) can retain the GEMM bias epilogue in the compiled forward pass while dispatching the two backward operations separately through custom ops.

This extends the operation compilation interface introduced in #3496. Eager and compiled execution share argument packing, compute methods, and context setup, so their numerical behavior stays aligned.

Type of change

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

Changes

• Register real and fake kernels for BasicLinear, Bias, and the supported Linear + Bias forward fusion. Evaluate forward and backward compilation support independently, allowing supported passes to compile when another pass uses an eager-only fusion.
• Support column- and row-parallel BasicLinear, with and without sequence parallelism, including pipelines with Bias. Fake kernels account for the shape changes from distributed collectives.
• Preserve quantized tensors needed by FP8 dequantized backward and cover both high-precision and dequantized backward overrides.
• Pass in_custom_op directly to compute methods. Keep eager storage release and main_grad bookkeeping inside BasicLinear.backward_compute, without duplicating dispatch or introducing postprocessing hooks. The flag is bound during custom-op registration rather than added to operator schemas.
• Add numerical and graph checks for standalone, fused, and unfused operations, inference and gradient combinations, parameter/input updates, and graph reuse. Numerical checks compare compiled outputs and requested gradients against native PyTorch FP64 references using exactly representable inputs, then compare subsequent random iterations against eager execution.
• Add distributed eager/compile coverage at world sizes 1 and 2 for column/row TP, sequence parallelism, optional bias, and BF16/FP8 Current Scaling. Strengthen existing eager main_grad checks.

Compilation still excludes accumulate_into_main_grad, CPU activation offload, and unsupported quantizers or extra tensor input/output plumbing. TP support here applies to BasicLinear and its pipelines with Bias; standalone AllReduce/ReduceScatter operations, including those used by row-parallel te.ops.Linear, still need their own custom ops.

Operation API change: in_custom_op

The FusibleOperation compute interface adds a keyword-only execution-context argument to both real compute methods:

@classmethod
def forward_compute(cls, args, *, in_custom_op: bool = False):
    ...

@classmethod
def backward_compute(cls, args, *, in_custom_op: bool = False):
    ...

Direct eager calls keep the default False. Custom-op registration binds in_custom_op=True with functools.partial, so the real kernel receives it whenever execution goes through the registered custom op. It identifies the execution path, including runtime kernel execution, rather than whether Python is currently being traced. The flag is not added to argument dataclasses or the PyTorch operator schema; fake-kernel signatures and compute return values remain unchanged.

Operation subclasses using this registration interface must accept the new keyword in their real compute methods, even if they do not need to branch on it. A fusion delegating to another operation's compute method forwards the flag, as Linear + Bias forward fusion does when calling BasicLinear.forward_compute.

This lets shared compute methods retain eager-only side effects while respecting the custom op's non-mutating schema. In BasicLinear.backward_compute, input storage release (clear_tensor_data) and dummy weight-gradient bookkeeping on the original parameter run only when in_custom_op=False. Compilation with accumulate_into_main_grad remains unsupported; the flag preserves its existing eager behavior.

Checklist

• I have read and followed the contributing guidelines (full L0 validation remains pending)
• The functionality described above 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
• Targeted new and existing unit tests pass with my changes