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

docs/examples/op_fuser/op_fuser.rst

@@ -151,6 +151,139 @@ arguments and the extra outputs will be returned.
    the block has been split into two sections, each with one branching
    operation.
 
+Extra tensor channels
+"""""""""""""""""""""
+
+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.
+
+With a channel, the residual block above can be expressed using one
+``Sequential``:
+
+.. code-block:: python
+
+    import torch
+    import transformer_engine.pytorch as te
+
+    make_residual = te.ops.MakeExtraOutput()
+    add_residual = te.ops.AddExtraInput()
+    make_residual.set_extra_output_channel(0, "residual")
+    add_residual.set_extra_input_channel(0, "residual")
+
+    block = te.ops.Sequential(
+        te.ops.LayerNorm(4096),
+        make_residual,
+        te.ops.Linear(4096, 28672),
+        te.ops.SwiGLU(),
+        te.ops.Linear(14336, 4096),
+        add_residual,
+    )
+
+    # The residual is routed internally and is also returned to the caller.
+    x = torch.randn(16384, 4096, device="cuda")
+    y, residual = block(x)
+
+Channels are also useful for mixture-of-experts blocks. The following
+example assumes custom ``Dispatch`` and ``Combine`` basic operations.
+``Dispatch`` has one public extra input containing router probabilities
+and three extra outputs: split sizes, token probabilities, and a
+routing map. ``Combine`` consumes the routing map.
+
+.. code-block:: python
+
+    import transformer_engine.pytorch as te
+    from my_ops import Dispatch, Combine
+
+    num_experts = 8
+    hidden_size = 4096
+    ffn_size = 14336
+
+    dispatch = Dispatch(num_experts)
+    fc1 = te.ops.GroupedLinear(
+        num_experts, hidden_size, 2 * ffn_size, bias=False
+    )
+    activation = te.ops.ScaledSwiGLU()
+    fc2 = te.ops.GroupedLinear(
+        num_experts, ffn_size, hidden_size, bias=False
+    )
+    combine = Combine(num_experts)
+
+    # Dispatch extra outputs:
+    #   0: split sizes, 1: token probabilities, 2: routing map
+    dispatch.set_extra_output_channel(0, "m_splits")
+    dispatch.set_extra_output_channel(1, "probs")
+    dispatch.set_extra_output_channel(2, "routing_map")
+
+    fc1.set_extra_input_channel(0, "m_splits")
+    activation.set_extra_input_channel(0, "probs")
+    fc2.set_extra_input_channel(0, "m_splits")
+    combine.set_extra_input_channel(0, "routing_map")
+
+    moe = te.ops.Sequential(dispatch, fc1, activation, fc2, combine)
+
+    # Dispatch's extra input has no channel, so the caller passes router_probs.
+    # Channels supply all later extra inputs internally, while Dispatch's
+    # extra outputs are still returned in their original order.
+    y, m_splits, probs, routing_map = moe(x, router_probs)
+
+Channels cannot connect operations in different ``OperationFuser``
+instances. In particular, an ordinary PyTorch module inside a
+``Sequential`` splits the fusible operations on either side into
+separate fusers. The following channel connection is therefore not
+supported:
+
+.. code-block:: python
+
+    make_residual = te.ops.MakeExtraOutput()
+    add_residual = te.ops.AddExtraInput()
+    make_residual.set_extra_output_channel(0, "residual")
+    add_residual.set_extra_input_channel(0, "residual")
+
+    block = te.ops.Sequential(
+        make_residual,
+        torch.nn.Identity(),  # Splits the operations into separate fusers.
+        add_residual,
+    )
+
+Use the public extra output and extra input interfaces, as in the
+two-``Sequential`` example above, when the producer and consumer cannot
+be placed in the same ``OperationFuser``.
+
+The following conditions apply to extra tensor channels:
+
+- A producer must appear before all of its consumers. Backward edges
+  and cycles are not supported.
+- An output channel name has at most one producer, but its output may
+  fan out to multiple consumers.
+- A named output does not require a consumer. It is still returned as
+  a public extra output.
+- A channel is scoped to one ``OperationFuser``. In a ``Sequential``,
+  ordinary PyTorch modules split adjacent fusible operations into
+  separate fusers, and channels cannot cross that boundary.
+- The caller passes extra inputs that are not connected to an earlier
+  producer in the same fuser. Channel-connected extra input slots do
+  not appear in the ``Sequential`` arguments.
+- The caller receives every extra output in the original basic-operation
+  and slot order. This includes channel-bound outputs that are also
+  consumed internally. Gradients supplied for a returned output are
+  combined with gradients from its internal channel consumers.
+- Channel bindings are captured when an ``OperationFuser`` (or the
+  fusers inside a ``Sequential``) is first constructed. Changing
+  ``set_extra_input_channel`` / ``set_extra_output_channel`` afterward
+  requires constructing a new ``OperationFuser`` or ``Sequential``.
+
+Channel-connected basic operations may still be replaced by registered
+``FusedOperation`` implementations. If a fused operation contains both
+the producer and consumer of a channel, its ``fuser_forward`` and
+``fuser_backward`` implementations are responsible for routing the
+tensor and its gradient between those basic operations.
+
 Developer guide
 ---------------