[Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op

[vthumbe1503]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

[greptile-apps]

Greptile Summary

The PR adds named internal extra-tensor channels to PyTorch fusible operations, including forward routing, backward gradient accumulation, validation, documentation, and tests. The stale-routing fix now makes channel bindings immutable after fuser construction, but its lock also applies to temporary fusers used for standalone operation calls.

• Adds producer-to-consumer channel resolution with fan-out support.
• Keeps unbound extra tensors in the public Sequential interface.
• Routes and accumulates internal channel gradients during backward.
• Adds channel validation, mutation locking, tests, and usage documentation.

Confidence Score: 4/5

The PR should not merge until standalone BasicOperation execution stops permanently preventing later channel configuration.

A direct BasicOperation call creates and discards a temporary fuser, but construction irreversibly locks the operation, causing later public channel setters to fail even though no retained routing snapshot exists.

Files Needing Attention: transformer_engine/pytorch/ops/fuser.py and transformer_engine/pytorch/ops/op.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fuser.py	Adds internal channel routing and gradient propagation, but permanently locks operations even when the fuser is temporary.
transformer_engine/pytorch/ops/op.py	Adds the public channel-binding API and irreversible mutation guard used by both persistent and transient fusers.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Initializes the conditional extra-input count before BasicOperation sizes its channel metadata.
tests/pytorch/test_fusible_ops.py	Adds broad channel routing and gradient coverage but does not cover configuring an operation after a standalone call.
docs/examples/op_fuser/op_fuser.rst	Documents channel topology, public versus internal slots, fan-out, fuser boundaries, and fused-operation responsibilities.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A["BasicOperation.forward()"] --> B["Temporary OperationFuser"]
  B --> C["Capture channel routing"]
  C --> D["Permanently lock BasicOperation"]
  D --> E["Run operation and discard fuser"]
  E --> F["Later channel setter"]
  F --> G["RuntimeError"]

Loading

_{Reviews (2): Last reviewed commit: "address review comment" | Re-trigger Greptile}

[greptile-apps]

Transient fusers permanently lock channels

When a BasicOperation is executed directly before its channels are configured, BasicOperation.forward creates and discards a temporary OperationFuser, but this constructor permanently locks the operation. A later set_extra_input_channel or set_extra_output_channel call therefore raises RuntimeError even though no retained fuser depends on the captured routing.

Knowledge Base Used: PyTorch Fusible-Operation Framework

[ptrendx]

Additional review comments from Codex:

  1. [High] Internal channel outputs lose the signal that their gradient is required.
     transformer_engine/pytorch/ops/fuser.py:170 only calls requires_grad_ for public outputs at line 194. A fresh tensor created by an internal producer inside
     torch.autograd.Function.forward therefore arrives at its consumer with requires_grad=False. Existing operations such as transformer_engine/pytorch/ops/basic/
     swiglu.py:468 and ScaledSReLU use that flag to decide whether to compute the extra-input gradient. They consequently return None, silently dropping the
     gradient to a differentiable producer such as the documented router-probability dispatch. The new tests do not expose this because MakeExtraOutput returns the
     original input tensor.

  2. [Medium] Channel routing violates the declared iterable output contract.
     transformer_engine/pytorch/ops/fuser.py:140 applies len() and indexing to a producer’s extra outputs, but transformer_engine/pytorch/ops/op.py:92 permits any
     Iterable[Iterable[Tensor]]. A custom Dispatch returning a generator works under the old flattening logic but now fails when a later channel consumer executes.

  3. [Medium] A standalone operation call permanently prevents later channel configuration.
     transformer_engine/pytorch/ops/op.py:598 constructs a temporary OperationFuser, while transformer_engine/pytorch/ops/fuser.py:509 permanently locks every
     attached operation. Calling an operation once through its normal forward, then placing it into a channel-connected Sequential, makes either setter raise even
     though the temporary fuser no longer exists.

  4. [Medium] The tests do not exercise two major routing branches.
     tests/pytorch/test_fusible_ops.py:460 registers only a forward fusion, so backward remains unfused and never exercises the same-fusion skip at
     transformer_engine/pytorch/ops/fuser.py:323. The multi-output test at tests/pytorch/test_fusible_ops.py:624 only checks duplicate-name rejection; its custom
     operation never runs. Thus mixed bound/unbound slot ordering, filtered autograd returns, and the modified two-input GroupedLinear(scale_bias=True) behavior
     remain unproved.

  ## Suggested repairs

  - For finding 1: Preserve the gradient-requirement flag on every extra output before classifying it as public or internal, or carry equivalent explicit per-slot
    metadata. Add a producer that creates a fresh tensor and verify gradient propagation through ScaledSwiGLU or ScaledSReLU.

  - For finding 2: Materialize and validate each operation’s extra outputs as a tuple immediately after fuser_forward; store that tuple for later consumers and
    lifetime tracking.

  - For finding 3: Make transient fusers created by BasicOperation.forward non-locking, while persistent Sequential fusers retain immutable routing. Add a call-
    then-bind regression test.

  - For finding 4: Add a joint/backward fused residual operation and assert fusion selection plus input/parameter/channel gradients. Add a successful multi-input/
    multi-output routing test and a GroupedLinear(scale_bias=True) channel case.

I would like you to also take a look at the tests - multiple of them duplicate each other (e.g. test_channel_fan_out_accumulates_grad is a stronger duplicate of test_internal_extra_tensor_channel_fanout).

[ptrendx]

There should be some explanation what a channel is before we say "no channel assigned".
It could be as simple as saying that the extra inputs/outputs may optionally specify
a channel (and then proceed with the rest of the text).

[ptrendx]

I would actually keep both of these examples here.

[ptrendx]

This limitation that there has to be at least one consumer in the named channel seems
arbitrary to me. If we do not strictly need this behavior then we shouldn't have that
as it would introduce friction when somebody needs to refactor the code using those
named channels by splitting the sequential - now they also need to remove the channel
names. In fact, I would expect people to generally want to name their extra outputs and
inputs even if they would not be reused inside the sequential. That could also enable
us to accept and return the dictionary rather than a list (which would make it less
fragile).