[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

The PR adds named extra-tensor channels so later fusible operations can consume outputs from earlier operations in the same fuser while preserving public extra outputs and their gradients.

• Adds channel binding and version tracking to basic operations.
• Resolves internal channel routing and fan-out in OperationFuser forward and backward passes.
• Adds validation, documentation, and broad tests covering routing, gradients, fusion, and invalid configurations.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; stale routing is rejected explicitly, transient and discarded fusers no longer leave channel locks, and structural Sequential changes rebuild routing from current operation state.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fuser.py	Implements channel discovery, internal forward routing, gradient fan-out accumulation, public-slot mapping, and stale-routing detection.
transformer_engine/pytorch/ops/op.py	Adds validated channel setters and version tracking to BasicOperation and strengthens the extra-output sequence contract.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Initializes dynamic extra-input arity before BasicOperation allocates its channel arrays.
tests/pytorch/test_fusible_ops.py	Adds coverage for channel routing, fusion ownership, fan-out gradients, mixed public/internal slots, invalid bindings, and stale configuration handling.
docs/examples/op_fuser/op_fuser.rst	Documents channel semantics, constraints, public output behavior, fuser boundaries, and fused-operation responsibilities.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Caller["Caller inputs"] --> Fuser["OperationFuser"]
  Fuser --> Producer["Producer basic op"]
  Producer -->|"main output"| Consumer["Later basic op"]
  Producer -->|"named extra-output channel"| Consumer
  Producer -->|"public extra output"| Caller
  Consumer --> Result["Sequential result"]

Loading

_{Reviews (17): Last reviewed commit: "fix lint" | 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.

[vthumbe1503]

/te-ci L1 pytorch

[timmoon10]

Is this working correctly? I recall having trouble setting attrs in torch.nn.Module before calling __init__ since it does special handling for params and modules: https://github.com/pytorch/pytorch/blob/99ee33c9719e0ceeb9cbda00e4cdd12a76c9ae45/torch/nn/modules/module.py#L1973

[timmoon10]

Nit: The "Forward op" comment originally described the entire code block with op.fuser_forward, but it's become meaningless now that it's so far away.

Suggested change

	# Forward op. Resolve internal channel inputs from outputs of
	# earlier basic ops. When a fusion contains both producer and
	# consumer, leave the consumer slot unset so the fused op can
	# wire the channel itself
	for idx in basic_op_idxs:
	for input_idx, source in enumerate(fuser._basic_op_extra_input_sources[idx]):
	if source is None:
	continue
	producer_idx, output_idx = source
	if producer_idx in basic_op_idxs:
	# fused op will wire the channel itself internally
	continue
	producer_outputs = extra_outputs[producer_idx]
	basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx]
	# Resolve internal channel inputs from outputs of
	# earlier basic ops. When a fusion contains both producer and
	# consumer, leave the consumer slot unset so the fused op can
	# wire the channel itself
	for idx in basic_op_idxs:
	for input_idx, source in enumerate(fuser._basic_op_extra_input_sources[idx]):
	if source is None:
	continue
	producer_idx, output_idx = source
	if producer_idx in basic_op_idxs:
	# fused op will wire the channel itself internally
	continue
	producer_outputs = extra_outputs[producer_idx]
	basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx]
	# Prepare args for op forward

[timmoon10]

I too look forward to the day when complex-valued operations become mainstream in DL. FFT and eigendecomposition are cool. Until then, we can avoid some CPU overhead:

Suggested change

	and (y.is_floating_point() or y.is_complex())
	and y.is_floating_point()

[timmoon10]

This is fine, but isn't basic_op_extra_output_is_internal redundant? If the channel is unused, channel_grads.get(channel) will return None, and we handle that case gracefully.

[timmoon10]

We should error out if the channel is missing a producer. Handling this edge case adds complexity. It also makes the contract less intuitive without providing any benefit.

[timmoon10]

We shouldn't reach this:

Suggested change

	model = te_ops.Sequential(producer, consumer)
	y, route = model(x, extra)
	torch.testing.assert_close(y, x + extra)
	torch.testing.assert_close(route, x)

[timmoon10]

We should remove this test if we require that input channels are valid.

[timmoon10]

This is fine, but somewhat redundant with test_mixed_channel_outputs_are_public.

[timmoon10]

It would be nice if Sequential could handle channels across OperationFusers, but the implementation would be quite hairy and not worth it for the current effort.

[timmoon10]

If you read pendantically, then this technically describes our current behavior where we ignore unmatched input channels. However, this is subtle and non-obvious. Better to error out if the input channel is invalid.

Suggested change

	Extra inputs and Extra outputs may optionally specify a channel. Assigning
	the same channel name to an extra output and one or more later extra
	inputs routes the tensor internally within the same
	``OperationFuser``. An extra input connected to an earlier producer is
	removed from the public ``Sequential`` arguments because the channel
	supplies it.
	Extra outputs remain in the public ``Sequential`` return value,
	including outputs that are also consumed through a channel.
	Branching operations can also route their extra inputs and outputs within the same ``Sequential`` via named channels. Extra output tensors with a specified channel can be consumed by other operations, in addition to being returned from the ``Sequential``. Extra input tensors with a specified channel are accessed internally instead of being provided as arguments to ``Seqential``.