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

[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

This PR adds named extra-tensor channels so an earlier fusible operation’s extra output can feed later operations within the same OperationFuser, while retaining every extra output in the public result.

• Adds channel configuration, routing snapshots, and stale-configuration detection.
• Routes internal tensors and accumulates gradients across channel consumers.
• Updates compatible basic and fused operation contracts, including dynamic GroupedLinear input sizing.
• Adds forward, backward, fan-out, fusion, validation, and lifecycle tests plus user documentation.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported stale-routing and retained-lock lifecycle issues are addressed by fail-fast version checks, disposable direct-operation fusers, and fresh routing construction after Sequential structural changes.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fuser.py	Implements channel resolution, external-slot arity, internal forward routing, gradient fan-out accumulation, output validation, and stale-routing rejection.
transformer_engine/pytorch/ops/op.py	Adds validated channel setters and mutation versioning to BasicOperation.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Determines scale-bias extra-input arity before base initialization so channel slot storage has the correct size.
tests/pytorch/test_fusible_ops.py	Adds broad coverage for channel routing, gradients, fusion ownership, fan-out, mixed public inputs, validation, and post-construction mutation.
docs/examples/op_fuser/op_fuser.rst	Documents channel semantics, scope, ordering, public outputs, reconstruction requirements, and fused-operation responsibilities.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Caller["Caller inputs"] --> Fuser["OperationFuser"]
  Producer["Earlier BasicOperation"] -->|"named extra output"| Channel["Internal channel"]
  Channel -->|"one or more consumers"| Consumer["Later BasicOperation(s)"]
  Producer -->|"extra output remains public"| CallerOutput["Caller outputs"]
  Consumer --> MainOutput["Main output"]
  CallerGrad["Public output gradient"] --> ProducerGrad["Producer extra-output gradient"]
  ConsumerGrad["Consumer gradients"] --> ProducerGrad

Loading

_{Reviews (15): Last reviewed commit: "minor doc" | Re-trigger Greptile}

[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]

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).

[vthumbe1503]

Originally I wanted to restrict the channel naming as a way to just do internal routing of tensors to reduce possibility of errors and the friction was kind of intentional. But I see your point of making it more seamless for user in future to construct a big a sequential op. If they want to refactor a code from 1 to 2 below

1. single sequential having internal routing
2. Two sequentials with one sequential passing extra output as extra input to another sequential

This can indeed be a problem since there might be some use-cases just supporting 2 but not 1.

And so I have removed that restriction. However, supporting extra_input and extra_output as dictionaries would be a problem from backwards compatibility perspective. Also, I want to restrict the scope of this PR. And allowing for dict based extra_input and extra output can be a seperate PR.

I have one extra requirement from named extra input channel added currently. If two different extra inputs share the same channel name, and is not internally connected to extra output of a previous op. Caller/User should still provide the extra_input two times.

op1.set_extra_input_channel(0, "common_name") # has 1 extra input
op2.set_extra_input_channel(0, "common_name") # has 1 extra input
model = te.Sequential(op1,op2)
y = model(input, extra_input, extra_input) 
# y = model(input, extra_input) --> wrong

This is done so that user's code doesnt have to change while naming an input channel vs not naming it. Also as you can see introducing extra_input dict is also going to make this tricky from backward compatibility perspective.