From b1574f0a9791d8142ea58a51eab04d5f2028d28e Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 28 Jul 2026 23:05:14 +0000 Subject: [PATCH 01/83] produce/consume extra output Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 42 ++++++ transformer_engine/pytorch/ops/fuser.py | 174 ++++++++++++++++++++---- transformer_engine/pytorch/ops/op.py | 38 ++++++ 3 files changed, 224 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..fc262b30cf 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -436,6 +436,48 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x3, x3_orig + x2 + b) torch.testing.assert_close(x4, x4_orig + x3) + def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: + """An internal extra output can feed multiple later consumers.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + + # Main path: x -> x + route -> x + route + route. + torch.testing.assert_close(y, 3 * x) + y.sum().backward() + # The channel fan-out contributes two independent gradient paths. + torch.testing.assert_close(x.grad, torch.full_like(x, 3)) + + # Internal slots are unavailable before forward, so grad discovery + # must tolerate them when no public input requires gradients. + x_no_grad = x.detach() + torch.testing.assert_close(model(x_no_grad), 3 * x_no_grad) + + def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: + """Unbound slots remain public when other slots use internal channels.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + + torch.testing.assert_close(y, 2 * x + extra) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + torch.testing.assert_close(extra.grad, torch.ones_like(extra)) + class TestFuser: """Tests for operation fusion infrastructure""" diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..33742137db 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,12 +102,16 @@ def forward( for tensor in (input_,) + params_and_extra_inputs: tensor._do_not_clear = True - # Unflatten list of parameters and extra tensor inputs - extra_inputs = params_and_extra_inputs[-fuser.num_extra_inputs :] - basic_op_extra_inputs = [] - for op in fuser._basic_ops: - xs, extra_inputs = _split_tuple(extra_inputs, op.num_extra_inputs) - basic_op_extra_inputs.append(xs) + # Place public extra inputs into their basic-op slots. Slots bound to + # internal channels are filled lazily as their producers execute. + extra_inputs = params_and_extra_inputs[len(fuser._flat_basic_op_params) :] + basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_inputs for op in fuser._basic_ops + ] + for tensor, (op_idx, input_idx) in zip( + extra_inputs, fuser._external_extra_input_slots + ): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops x = input_ @@ -118,8 +122,31 @@ def forward( for idx in basic_op_idxs: basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward - # Forward op - extra_inputs = [basic_op_extra_inputs[idx] for idx in basic_op_idxs] + # Forward op. Resolve internal channel inputs from outputs of + # earlier basic ops. A fusion may consume an earlier channel, but + # may not contain both its producer and consumer. + 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: + raise RuntimeError( + "An operation fusion contains both producer and consumer " + f"of extra tensor channel " + f"{fuser._basic_op_extra_output_channels[producer_idx][output_idx]!r}" + ) + producer_outputs = extra_outputs[producer_idx] + if producer_outputs is None: + raise RuntimeError( + f"Extra tensor channel producer op {producer_idx} has not run" + ) + basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] + op_extra_inputs = [ + tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs + ] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None prev_op_grad_output_quantizer = None @@ -134,18 +161,21 @@ def forward( x, fused_op_extra_outputs = op.fuser_forward( [basic_op_ctxs[idx] for idx in basic_op_idxs], x, - basic_op_extra_inputs=extra_inputs, + basic_op_extra_inputs=op_extra_inputs, prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): for y in ys: - if set_output_requires_grad: + if set_output_requires_grad and ( + y.is_floating_point() or y.is_complex() + ): y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs[idx] = ys - # Flatten list of extra outputs + # Validate extra outputs and flatten only public slots. Outputs bound + # to channels stay internal to the fuser. extra_outputs_flat = [] for idx, ys in enumerate(extra_outputs): ys = list(ys) @@ -156,7 +186,9 @@ def forward( "{num_extra_outputs} extra inputs, " f"but got {len(ys)}" ) - extra_outputs_flat.extend(ys) + for output_idx, y in enumerate(ys): + if fuser._basic_op_extra_output_channels[idx][output_idx] is None: + extra_outputs_flat.append(y) # Save context for backward pass if func_ctx is not None: @@ -188,10 +220,12 @@ def forward( func_ctx.basic_op_num_params = fuser._basic_op_num_params func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) + func_ctx.fuser = fuser func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward - for tensor in [x] + extra_outputs_flat: + all_extra_outputs = [y for ys in extra_outputs for y in ys] + for tensor in [x] + all_extra_outputs: tensor._do_not_clear = True if set_output_requires_grad: @@ -224,21 +258,28 @@ def backward( ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None - # Unflatten list of extra tensor output grads + fuser = func_ctx.fuser + + # Place public extra-output grads into their basic-op slots. Internal + # output grads are accumulated from channel consumers during backward. if len(grad_extra_outputs) != func_ctx.num_extra_outputs: raise ValueError( f"Expected grads for {func_ctx.num_extra_outputs} extra tensor outputs, " f"but got {len(grad_extra_outputs)}" ) - basic_op_grad_extra_outputs = [] - for op in basic_ops: - dys, grad_extra_outputs = _split_tuple(grad_extra_outputs, op.num_extra_outputs) - basic_op_grad_extra_outputs.append(dys) + basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_outputs for op in basic_ops + ] + for grad, (op_idx, output_idx) in zip( + grad_extra_outputs, fuser._external_extra_output_slots + ): + basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops dx = grad_output grad_params = [None for _ in range(len(basic_ops))] grad_extra_inputs = [None for _ in range(len(basic_ops))] + channel_grads: dict[str, torch.Tensor] = {} for op, basic_op_idxs in reversed(backward_ops): # Stop if no more gradients are required @@ -246,18 +287,39 @@ def backward( dx = None break - # Backward op - grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] + # Backward op. Supply gradients accumulated from every consumer of + # each internal channel. + for idx in basic_op_idxs: + for output_idx, channel in enumerate( + fuser._basic_op_extra_output_channels[idx] + ): + if channel is not None: + basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get( + channel + ) + op_grad_extra_outputs = [ + tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs + ] dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( [basic_op_ctxs[idx] for idx in basic_op_idxs], dx, - basic_op_grad_extra_outputs=grad_extra_outputs, + basic_op_grad_extra_outputs=op_grad_extra_outputs, ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs + for input_idx, grad in enumerate(dxs): + source = fuser._basic_op_extra_input_sources[idx][input_idx] + if source is None or grad is None: + continue + producer_idx, output_idx = source + channel = fuser._basic_op_extra_output_channels[producer_idx][output_idx] + previous_grad = channel_grads.get(channel) + channel_grads[channel] = ( + grad if previous_grad is None else previous_grad + grad + ) # Flatten list of parameter gradients grad_params_flat = [] @@ -288,7 +350,9 @@ def backward( f"for {num_extra_inputs} extra inputs, " f"but got {len(dxs)}" ) - grad_extra_inputs_flat.extend(dxs) + for input_idx, grad in enumerate(dxs): + if fuser._basic_op_extra_input_sources[idx][input_idx] is None: + grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -342,7 +406,52 @@ def __init__( # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) - self.num_extra_inputs: int = sum(self._basic_op_num_extra_inputs) + self._basic_op_extra_input_sources: list[list[Optional[tuple[int, int]]]] = [ + [None] * op.num_extra_inputs for op in basic_ops + ] + self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ + list(op._extra_output_channels) for op in basic_ops + ] + self._external_extra_input_slots: list[tuple[int, int]] = [] + self._external_extra_output_slots: list[tuple[int, int]] = [] + + # Resolve named channels in pipeline order. Channels deliberately only + # connect an output to later inputs, which keeps execution acyclic. + channel_producers: dict[str, tuple[int, int]] = {} + consumed_channels: set[str] = set() + for op_idx, op in enumerate(basic_ops): + for input_idx, channel in enumerate(op._extra_input_channels): + if channel is None: + self._external_extra_input_slots.append((op_idx, input_idx)) + continue + if channel not in channel_producers: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no earlier producer" + ) + self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[ + channel + ] + consumed_channels.add(channel) + for output_idx, channel in enumerate(op._extra_output_channels): + if channel is None: + self._external_extra_output_slots.append((op_idx, output_idx)) + continue + if channel in channel_producers: + producer_idx, _ = channel_producers[channel] + raise ValueError( + f"Extra tensor channel {channel!r} has multiple producers " + f"(ops {producer_idx} and {op_idx})" + ) + channel_producers[channel] = (op_idx, output_idx) + + unused_channels = channel_producers.keys() - consumed_channels + if unused_channels: + channels = ", ".join(repr(channel) for channel in sorted(unused_channels)) + raise ValueError(f"Extra tensor channels have no consumers: {channels}") + + self.num_extra_inputs = len(self._external_extra_input_slots) + self.num_extra_outputs = len(self._external_extra_output_slots) # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] @@ -432,7 +541,9 @@ def maybe_fuse_ops( first_op_requiring_backward = self._num_basic_ops for op_idx in range(self._num_basic_ops): op_inputs = itertools.chain(self._basic_op_params[op_idx], extra_inputs[op_idx]) - if any(tensor.requires_grad for tensor in op_inputs): + if any( + tensor is not None and tensor.requires_grad for tensor in op_inputs + ): first_op_requiring_backward = op_idx break @@ -517,12 +628,15 @@ def __call__( if basic_op_kwargs is None: basic_op_kwargs = [{}] * self._num_basic_ops - # Unflatten list of extra tensor inputs - extra_inputs_copy = list(extra_inputs) - basic_op_extra_inputs = [] - for op in self._basic_ops: - xs, extra_inputs_copy = _split_tuple(extra_inputs_copy, op.num_extra_inputs) - basic_op_extra_inputs.append(xs) + # Place public extra inputs into their basic-op slots. Internal slots + # are not available until forward executes their producers. + basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_inputs for op in self._basic_ops + ] + for tensor, (op_idx, input_idx) in zip( + extra_inputs, self._external_extra_input_slots + ): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state recipe = None diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..3159ae0f0d 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -187,10 +187,48 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): def __init__(self) -> None: super().__init__() + # Optional names for extra-tensor channels internal to an OperationFuser. + # Unbound slots remain public inputs/outputs, preserving the original API. + self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs + self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs + # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Bind an extra input slot to an internal fuser channel. + + A bound slot receives the matching extra output from an earlier + operation in the same fuser instead of consuming a public extra input. + Passing ``None`` removes the binding. + """ + if not 0 <= index < self.num_extra_inputs: + raise IndexError( + f"Extra input index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_inputs} extra inputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra input channel must be a non-empty string or None") + self._extra_input_channels[index] = channel + return self + + def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Bind an extra output slot to an internal fuser channel. + + A bound slot can feed one or more later operations and is not returned + as a public extra output. Passing ``None`` removes the binding. + """ + if not 0 <= index < self.num_extra_outputs: + raise IndexError( + f"Extra output index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_outputs} extra outputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra output channel must be a non-empty string or None") + self._extra_output_channels[index] = channel + return self + @property def is_fused_op(self) -> bool: return False From 63192abdcaf135d1929f8e2ab73d25d69b9f5870 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 04:28:38 +0000 Subject: [PATCH 02/83] allow for fusions with producer/consumer being part of same fuser with error handling tests Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 468 ++++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 2 + transformer_engine/pytorch/ops/fuser.py | 74 ++- 3 files changed, 534 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index fc262b30cf..67358493b0 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -23,6 +23,7 @@ OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, ) +from transformer_engine.pytorch.ops.fuser import OperationFuser from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.ops.fused import ( @@ -478,6 +479,473 @@ def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None torch.testing.assert_close(x.grad, torch.full_like(x, 2)) torch.testing.assert_close(extra.grad, torch.ones_like(extra)) + def test_moe_style_dispatch_combine_extra_channels( + self, + *, + group_size: int = 4, + hidden_size: int = 32, + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ) -> None: + """Wire Dispatch-style extras through GroupedLinear / ScaledActivation / Combine. + + Stand-in for ``te.Sequential(Dispatch, GroupedLinear, ScaledActivation, + GroupedLinear, Combine)`` once real Dispatch/Combine ops land. ``m_splits`` + and ``probs`` are produced once and fan out to later consumers via named + channels so the public call is ``model(x, m_splits, probs)``. + """ + + class FakeDispatch(te_ops.BasicOperation): + """Stand-in MoE dispatch: passthrough hidden states, emit routing extras. + + Real Dispatch would permute tokens. Extra inputs are the externally + provided ``m_splits`` and ``probs``; matching extra outputs are bound + to internal channels for later consumers. A third extra output is a + stub ``routing_map`` for Combine. + """ + + num_extra_inputs = 2 + num_extra_outputs = 3 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + m_splits, probs = basic_op_extra_inputs[0] + # Stub row-id map: real Dispatch would emit permute indices. + routing_map = torch.arange( + input_.size(0), device=input_.device, dtype=torch.int64 + ) + return input_, [(m_splits, probs, routing_map)] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + # Channel grads from GroupedLinear / ScaledActivation land here. + return ( + grad_output, + [()], + [tuple(basic_op_grad_extra_outputs[0][:2])], + ) + + class FakeCombine(te_ops.BasicOperation): + """Stand-in MoE combine: consumes Dispatch ``routing_map``, identity path. + + Real Combine would unpermute with the routing map. + """ + + num_extra_inputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("FakeCombine uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("FakeCombine uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + routing_map = basic_op_extra_inputs[0][0] + if routing_map is None: + raise RuntimeError("FakeCombine expected routing_map channel") + if int(routing_map.numel()) != int(input_.size(0)): + raise RuntimeError( + "FakeCombine routing_map length does not match tokens" + ) + return input_, [()] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + del basic_op_grad_extra_outputs + return grad_output, [()], [(None,)] + + split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[ + :group_size + ] + num_tokens = int(split_sizes.sum()) + in_shape = (num_tokens, hidden_size) + + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + probs_ref, probs_test = make_reference_and_test_tensors( + (num_tokens,), + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Reference: GroupedLinear + ScaledSReLU + GroupedLinear (no dispatch permute). + # Run the PyTorch reference in the same dtype/device as TE. Cross-device + # float32 GEMMs (CPU vs CUDA) differ enough that probs grads — a + # reduction over hidden — can miss dtype_tols even when channel wiring + # is correct. + fc1_w_refs, fc1_w_tests = [], [] + fc2_w_refs, fc2_w_tests = [], [] + for _ in range(group_size): + w1_ref, w1_test = make_reference_and_test_tensors( + (hidden_size, hidden_size), + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + w2_ref, w2_test = make_reference_and_test_tensors( + (hidden_size, hidden_size), + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + fc1_w_refs.append(w1_test.detach().clone()) + fc1_w_tests.append(w1_test) + fc2_w_refs.append(w2_test.detach().clone()) + fc2_w_tests.append(w2_test) + x_ref = x_test.detach().clone().requires_grad_(True) + probs_ref = probs_test.detach().clone().requires_grad_(True) + dy_ref = dy_test.detach().clone() + xs = torch.split(x_ref, split_sizes.tolist()) + probs = torch.split(probs_ref, split_sizes.tolist()) + ys = [] + for group_idx in range(group_size): + fc1_out = torch.nn.functional.linear(xs[group_idx], fc1_w_refs[group_idx]) + act_out = torch.nn.functional.relu(fc1_out).square() + fc2_in = act_out * probs[group_idx].unsqueeze(-1) + ys.append(torch.nn.functional.linear(fc2_in, fc2_w_refs[group_idx])) + y_ref = torch.cat(ys) + y_ref.backward(dy_ref) + + dispatch = FakeDispatch() + fc1 = te_ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + activation = te_ops.ScaledSReLU() + fc2 = te_ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + combine = FakeCombine() + + # Bind channels: Dispatch fans out, consumers bind by name. + 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") + + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + with torch.no_grad(): + for group_idx in range(group_size): + getattr(fc1, f"weight{group_idx}").copy_(fc1_w_tests[group_idx]) + getattr(fc2, f"weight{group_idx}").copy_(fc2_w_tests[group_idx]) + del fc1_w_tests, fc2_w_tests + + # Only Dispatch's extras remain public: model(x, m_splits, probs). + y_test = model(x_test, split_sizes, probs_test) + y_test.backward(dy_test) + + tols = dtype_tols(dtype) + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + assert_close_grads(probs_test, probs_ref, **tols) + + def test_fused_op_with_internal_producer_consumer(self, size: int = 16) -> None: + """A fused op may contain both a channel producer and its consumer. + + This is the MegaMoE path: Dispatch + MLP + Combine collapse into one + fused kernel that wires channels internally instead of via the fuser. + """ + + class FakeDispatch(te_ops.BasicOperation): + num_extra_inputs = 1 + num_extra_outputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + (route,) = basic_op_extra_inputs[0] + return input_, [(route,)] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + return ( + grad_output, + [()], + [tuple(basic_op_grad_extra_outputs[0])], + ) + + class MegaMoELike(te_ops.FusedOperation): + """Fused stub that owns both producer and consumer of ``route``.""" + + _enabled = True + + def __init__(self, dispatch, consumer) -> None: + super().__init__((dispatch, consumer)) + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + # Consumer slot is intentionally unset: producer is in this fusion. + route = basic_op_extra_inputs[0][0] + assert basic_op_extra_inputs[1][0] is None + out = input_ + route + return out, [(route,), ()] + + def fuse_mega_moe_like(ops, **unused): + if not MegaMoELike._enabled: + return ops + MegaMoELike._enabled = False + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + if isinstance(window[0], FakeDispatch) and isinstance( + window[1], te_ops.AddExtraInput + ): + window = [MegaMoELike(*window)] + else: + out.append(window[0]) + window = window[1:] + window, ops = window + ops[:1], ops[1:] + out.extend(window + ops) + return out + + dispatch = FakeDispatch() + consumer = te_ops.AddExtraInput() + dispatch.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(dispatch, consumer) + + te_ops.register_forward_fusion(fuse_mega_moe_like) + x = torch.rand((size,), requires_grad=True) + route = torch.rand((size,), requires_grad=True) + y = model(x, route) + torch.testing.assert_close(y, x + route) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.ones_like(x)) + torch.testing.assert_close(route.grad, torch.ones_like(route)) + + +class TestExtraTensorChannels: + """Error handling and grad coverage for named extra-tensor channels.""" + + def test_consumer_channel_without_producer(self) -> None: + """Extra input bound to a channel that no earlier op produces.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "missing") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer]) + + def test_consumer_before_producer(self) -> None: + """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" + consumer = te_ops.AddExtraInput() + producer = te_ops.MakeExtraOutput() + consumer.set_extra_input_channel(0, "route") + producer.set_extra_output_channel(0, "route") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer, producer]) + + def test_set_extra_channel_rejects_invalid_index(self) -> None: + """Slot indices must be in range; negatives and OOB are rejected at bind time.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(1, "route") + + def test_set_extra_channel_rejects_invalid_name(self) -> None: + """Channel names must be non-empty strings.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + consumer.set_extra_input_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + + def test_duplicate_extra_output_channel_names(self) -> None: + """Two extra outputs may not publish the same channel name.""" + producer1 = te_ops.MakeExtraOutput() + producer2 = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer1.set_extra_output_channel(0, "route") + producer2.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser([producer1, producer2, consumer]) + + def test_duplicate_extra_output_channels_on_same_op(self) -> None: + """A single op with multiple extras still cannot reuse a channel name.""" + + class DualExtraOutput(te_ops.BasicOperation): + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + return input_, [(input_, input_)] + + def fuser_backward( + self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs + ): + g0, g1 = basic_op_grad_extra_outputs[0] + grad_extra = torch.zeros_like(grad_output) + if g0 is not None: + grad_extra = grad_extra + g0 + if g1 is not None: + grad_extra = grad_extra + g1 + return grad_output + grad_extra, [()], [()] + + producer = DualExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + producer.set_extra_output_channel(1, "route") + consumer.set_extra_input_channel(0, "route") + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser([producer, consumer]) + + def test_unused_extra_output_channel(self) -> None: + """Every produced channel must have at least one consumer.""" + producer = te_ops.MakeExtraOutput() + producer.set_extra_output_channel(0, "orphan") + with pytest.raises(ValueError, match="have no consumers"): + OperationFuser([producer]) + + def test_one_extra_input_has_single_source(self) -> None: + """Each extra-input slot binds to one channel / one producer source. + + Rebinding replaces the previous name; the abandoned producer channel + then fails as unused rather than attaching two sources to one input. + """ + producer_a = te_ops.MakeExtraOutput() + producer_b = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer_a.set_extra_output_channel(0, "a") + producer_b.set_extra_output_channel(0, "b") + consumer.set_extra_input_channel(0, "a") + consumer.set_extra_input_channel(0, "b") + with pytest.raises(ValueError, match="have no consumers"): + OperationFuser([producer_a, producer_b, consumer]) + + # Valid single binding: consumer input 0 is fed only by producer_a. + consumer.set_extra_input_channel(0, "a") + producer_b.set_extra_output_channel(0, None) + fuser = OperationFuser([producer_a, producer_b, consumer]) + assert fuser._basic_op_extra_input_sources[2] == [(0, 0)] + assert fuser.num_extra_inputs == 0 + assert fuser.num_extra_outputs == 1 # producer_b's unbound extra output + + def test_channel_fanout_accumulates_grads(self, size: int = 16) -> None: + """Grads from every consumer of a channel are accumulated into the producer.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + # Forward: x -> x+x -> x+x+x + torch.testing.assert_close(y, 3 * x) + + dy = torch.rand((size,)) + y.backward(dy) + # Main path contributes dy; each AddExtraInput also routes dy back + # through the channel into MakeExtraOutput's extra-output grad, which + # is added again into dx. Total: dy (main) + dy + dy (two consumers). + torch.testing.assert_close(x.grad, 3 * dy) + + def test_mixed_internal_external_grad(self, size: int = 16) -> None: + """Internal channel grads and public extra-input grads both flow correctly.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + torch.testing.assert_close(y, 2 * x + extra) + + dy = torch.rand((size,)) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(extra.grad, dy) + class TestFuser: """Tests for operation fusion infrastructure""" diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..3f4ca2d4b2 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -152,6 +152,8 @@ def __init__( self._scale_bias: bool = scale_bias and bias if self._scale_bias: self.num_extra_inputs = 2 + # BasicOperation.__init__ sized channel lists from the class default (1). + self._extra_input_channels = [None] * self.num_extra_inputs self.wgrad_store = WeightGradStore(delay_wgrad_compute) self.wgrad_accumulation_and_reduce_hooks: list = [] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 33742137db..2655ea4319 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -123,8 +123,9 @@ def forward( basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward # Forward op. Resolve internal channel inputs from outputs of - # earlier basic ops. A fusion may consume an earlier channel, but - # may not contain both its producer and consumer. + # 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] @@ -133,16 +134,25 @@ def forward( continue producer_idx, output_idx = source if producer_idx in basic_op_idxs: - raise RuntimeError( - "An operation fusion contains both producer and consumer " - f"of extra tensor channel " - f"{fuser._basic_op_extra_output_channels[producer_idx][output_idx]!r}" - ) + # fused op will wire the channel itself internally + continue producer_outputs = extra_outputs[producer_idx] if producer_outputs is None: raise RuntimeError( f"Extra tensor channel producer op {producer_idx} has not run" ) + if ( + output_idx >= len(producer_outputs) + or producer_outputs[output_idx] is None + ): + raise RuntimeError( + f"Extra tensor channel producer op {producer_idx} " + f"({type(fuser._basic_ops[producer_idx]).__name__}) " + f"did not emit extra output {output_idx} for " + f"consumer op {idx} " + f"({type(fuser._basic_ops[idx]).__name__}) " + f"input {input_idx}" + ) basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] op_extra_inputs = [ tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs @@ -167,7 +177,12 @@ def forward( basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): - for y in ys: + for output_idx, y in enumerate(ys): + if y is None: + raise RuntimeError( + f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " + f"did not emit extra output {output_idx}" + ) if set_output_requires_grad and ( y.is_floating_point() or y.is_complex() ): @@ -315,6 +330,10 @@ def backward( if source is None or grad is None: continue producer_idx, output_idx = source + # Producer already ran inside this fusion; the fused op + # must apply these grads itself rather than via channel_grads. + if producer_idx in basic_op_idxs: + continue channel = fuser._basic_op_extra_output_channels[producer_idx][output_idx] previous_grad = channel_grads.get(channel) channel_grads[channel] = ( @@ -420,7 +439,8 @@ def __init__( channel_producers: dict[str, tuple[int, int]] = {} consumed_channels: set[str] = set() for op_idx, op in enumerate(basic_ops): - for input_idx, channel in enumerate(op._extra_input_channels): + for input_idx in range(op.num_extra_inputs): + channel = op._extra_input_channels[input_idx] if channel is None: self._external_extra_input_slots.append((op_idx, input_idx)) continue @@ -433,7 +453,9 @@ def __init__( channel ] consumed_channels.add(channel) - for output_idx, channel in enumerate(op._extra_output_channels): + for output_idx, channel in enumerate( + self._basic_op_extra_output_channels[op_idx] + ): if channel is None: self._external_extra_output_slots.append((op_idx, output_idx)) continue @@ -450,6 +472,38 @@ def __init__( channels = ", ".join(repr(channel) for channel in sorted(unused_channels)) raise ValueError(f"Extra tensor channels have no consumers: {channels}") + # Every channel-bound extra input must be wired to a matching producer + # extra output. External slots remain unbound (source is None). + for op_idx, sources in enumerate(self._basic_op_extra_input_sources): + op = basic_ops[op_idx] + for input_idx, source in enumerate(sources): + channel = op._extra_input_channels[input_idx] + if channel is None: + if source is not None: + raise RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is external but has a " + f"producer source {source}" + ) + continue + if source is None: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r} " + f"but has no producer" + ) + producer_idx, output_idx = source + producer_channel = self._basic_op_extra_output_channels[producer_idx][ + output_idx + ] + if producer_channel != channel: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r}, " + f"but producer op {producer_idx} extra output {output_idx} " + f"is bound to {producer_channel!r}" + ) + self.num_extra_inputs = len(self._external_extra_input_slots) self.num_extra_outputs = len(self._external_extra_output_slots) From 3b4b523937b8741784f3f3c1e27433326dfdde0c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 18:18:43 +0000 Subject: [PATCH 03/83] cleanup Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 6 +-- transformer_engine/pytorch/ops/fuser.py | 37 +++++++++++-------- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 3f4ca2d4b2..7faad6536b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -147,13 +147,11 @@ def __init__( delay_wgrad_compute: bool = False, scale_bias: bool = False, ) -> None: - super().__init__() - + # Decide before BasicOperation.__init__ sizes _extra_input_channels. self._scale_bias: bool = scale_bias and bias if self._scale_bias: self.num_extra_inputs = 2 - # BasicOperation.__init__ sized channel lists from the class default (1). - self._extra_input_channels = [None] * self.num_extra_inputs + super().__init__() self.wgrad_store = WeightGradStore(delay_wgrad_compute) self.wgrad_accumulation_and_reduce_hooks: list = [] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 2655ea4319..f0f5f36fe9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,7 +102,7 @@ def forward( for tensor in (input_,) + params_and_extra_inputs: tensor._do_not_clear = True - # Place public extra inputs into their basic-op slots. Slots bound to + # Place user provided extra inputs into their basic-op slots. Slots bound to # internal channels are filled lazily as their producers execute. extra_inputs = params_and_extra_inputs[len(fuser._flat_basic_op_params) :] basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ @@ -183,14 +183,10 @@ def forward( f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " f"did not emit extra output {output_idx}" ) - if set_output_requires_grad and ( - y.is_floating_point() or y.is_complex() - ): - y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs[idx] = ys # Validate extra outputs and flatten only public slots. Outputs bound - # to channels stay internal to the fuser. + # to channels stay internal to the fuser and are not marked. extra_outputs_flat = [] for idx, ys in enumerate(extra_outputs): ys = list(ys) @@ -202,7 +198,11 @@ def forward( f"but got {len(ys)}" ) for output_idx, y in enumerate(ys): + # Output is not bound to a channel consumed by another op, + # so it is a public output. if fuser._basic_op_extra_output_channels[idx][output_idx] is None: + if set_output_requires_grad: + y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs_flat.append(y) # Save context for backward pass @@ -228,14 +228,16 @@ def forward( if fuser.first_op_requiring_backward < fuser._num_basic_ops: is_first_module = FP8GlobalStateManager.is_first_fp8_module() - # Other context + # Other context. Save only the wiring metadata needed by + # backward instead of the whole OperationFuser. func_ctx.backward_ops = fuser._backward_ops func_ctx.basic_ops = fuser._basic_ops func_ctx.basic_op_ctxs = basic_op_ctxs func_ctx.basic_op_num_params = fuser._basic_op_num_params - func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) - func_ctx.fuser = fuser + func_ctx.external_extra_output_slots = fuser._external_extra_output_slots + func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward @@ -273,7 +275,10 @@ def backward( ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None - fuser = func_ctx.fuser + # Channel wiring saved from forward + external_extra_output_slots = func_ctx.external_extra_output_slots + basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels + basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources # Place public extra-output grads into their basic-op slots. Internal # output grads are accumulated from channel consumers during backward. @@ -286,7 +291,7 @@ def backward( [None] * op.num_extra_outputs for op in basic_ops ] for grad, (op_idx, output_idx) in zip( - grad_extra_outputs, fuser._external_extra_output_slots + grad_extra_outputs, external_extra_output_slots ): basic_op_grad_extra_outputs[op_idx][output_idx] = grad @@ -306,7 +311,7 @@ def backward( # each internal channel. for idx in basic_op_idxs: for output_idx, channel in enumerate( - fuser._basic_op_extra_output_channels[idx] + basic_op_extra_output_channels[idx] ): if channel is not None: basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get( @@ -326,7 +331,7 @@ def backward( for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs for input_idx, grad in enumerate(dxs): - source = fuser._basic_op_extra_input_sources[idx][input_idx] + source = basic_op_extra_input_sources[idx][input_idx] if source is None or grad is None: continue producer_idx, output_idx = source @@ -334,7 +339,7 @@ def backward( # must apply these grads itself rather than via channel_grads. if producer_idx in basic_op_idxs: continue - channel = fuser._basic_op_extra_output_channels[producer_idx][output_idx] + channel = basic_op_extra_output_channels[producer_idx][output_idx] previous_grad = channel_grads.get(channel) channel_grads[channel] = ( grad if previous_grad is None else previous_grad + grad @@ -370,7 +375,9 @@ def backward( f"but got {len(dxs)}" ) for input_idx, grad in enumerate(dxs): - if fuser._basic_op_extra_input_sources[idx][input_idx] is None: + # Only append public grad extra inputs to the list + # to be returned to the user. + if basic_op_extra_input_sources[idx][input_idx] is None: grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors From de38ed836a4c54ea84a2c2db2e8519c6e48ba197 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 19:03:26 +0000 Subject: [PATCH 04/83] minor cleanup Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 3 ++- transformer_engine/pytorch/ops/fuser.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 67358493b0..9cecd4531b 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -903,7 +903,8 @@ def test_one_extra_input_has_single_source(self) -> None: fuser = OperationFuser([producer_a, producer_b, consumer]) assert fuser._basic_op_extra_input_sources[2] == [(0, 0)] assert fuser.num_extra_inputs == 0 - assert fuser.num_extra_outputs == 1 # producer_b's unbound extra output + # producer_b's unbound extra output remains public + assert fuser._external_extra_output_slots == [(1, 0)] def test_channel_fanout_accumulates_grads(self, size: int = 16) -> None: """Grads from every consumer of a channel are accumulated into the producer.""" diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index f0f5f36fe9..cc706ab91f 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -510,9 +510,9 @@ def __init__( f"but producer op {producer_idx} extra output {output_idx} " f"is bound to {producer_channel!r}" ) - + # Used by Sequential to determine the number of extra inputs + # needed for each OperationFuser module in the sequence. self.num_extra_inputs = len(self._external_extra_input_slots) - self.num_extra_outputs = len(self._external_extra_output_slots) # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] From 385b0d51b3b074d00108e1ba0373bf6ba6b57678 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 23:10:17 +0000 Subject: [PATCH 05/83] dispatch combine impl Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/moe_ep_reference.py | 792 ++++++++++++++++++ tests/pytorch/distributed/run_ep.py | 160 ++++ .../pytorch/ops/basic/__init__.py | 2 + transformer_engine/pytorch/ops/combine.py | 153 ++++ transformer_engine/pytorch/ops/dispatch.py | 214 +++++ 5 files changed, 1321 insertions(+) create mode 100644 tests/pytorch/distributed/moe_ep_reference.py create mode 100644 transformer_engine/pytorch/ops/combine.py create mode 100644 transformer_engine/pytorch/ops/dispatch.py diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/tests/pytorch/distributed/moe_ep_reference.py new file mode 100644 index 0000000000..3fab88a50b --- /dev/null +++ b/tests/pytorch/distributed/moe_ep_reference.py @@ -0,0 +1,792 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. + +The implementation deliberately favors readable semantics over performance. It +supports a one-rank execution path and a variable-size ``all_to_all_single`` EP +path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. + +The quantized tensor layouts are logical (unswizzled) layouts. A production +kernel may reorder scale factors internally without changing this API contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +class MoeFormat(str, Enum): + """Public and communication formats supported by the reference.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = value + return tuple(result) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Portable data-plus-scale representation for MXFP8 or NVFP4. + + ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` + has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with + two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` + replaces that axis by one scale per block. + """ + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + fmt = _parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + shape = tuple(int(dim) for dim in self.logical_shape) + if not shape or any(dim < 0 for dim in shape): + raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") + axis = _normalize_axis(self.axis, len(shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", shape) + object.__setattr__(self, "axis", axis) + self._validate_storage() + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + def _validate_storage(self) -> None: + if self.data.device != self.scale.device: + raise ValueError("block-scaled data and scale must be on the same device") + + logical_extent = self.logical_shape[self.axis] + scale_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, self.block_size), + ) + if tuple(self.scale.shape) != scale_shape: + raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") + + if self.format is MoeFormat.MXFP8: + expected_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + data_shape = self.logical_shape + if self.data.dtype != expected_dtype: + raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + else: + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: + raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") + data_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, 2), + ) + + if tuple(self.data.shape) != data_shape: + raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") + if self.scale.dtype != expected_scale_dtype: + raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return the logical tensor with block scales applied.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + packed = packed.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" + + levels = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=values.device, + ) + magnitudes = values.abs().unsqueeze(-1) + distances = (magnitudes - levels).abs() + minimum = distances.amin(dim=-1, keepdim=True) + candidates = distances == minimum + codes = torch.arange(8, dtype=torch.int64, device=values.device) + any_code = torch.where(candidates, codes, 8).amin(dim=-1) + even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) + magnitude_code = torch.where(even_code < 8, even_code, any_code) + sign_code = torch.signbit(values).to(torch.int64) << 3 + return magnitude_code | sign_code + + +def quantize_blockwise( + tensor: torch.Tensor, + format: Union[MoeFormat, str], + *, + axis: int = -1, +) -> BlockScaledTensor: + """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. + + MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward + positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and + E4M3 scales rounded to nearest. + """ + + fmt = _parse_format(format) + if fmt is MoeFormat.BF16: + raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") + if not tensor.is_floating_point(): + raise TypeError(f"tensor must be floating point, got {tensor.dtype}") + + axis = _normalize_axis(axis, tensor.ndim) + logical_shape = tuple(tensor.shape) + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + block_count = _ceil_div(logical_extent, block_size) + padded_extent = block_count * block_size + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) + + value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 + scale_float = blocks.abs().amax(dim=-1) / value_limit + if fmt is MoeFormat.MXFP8: + safe_scale = torch.where(scale_float > 0, scale_float, 1.0) + scale_float = torch.where( + scale_float > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(scale_float), + ) + scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + scale_dtype = _require_torch_dtype("float8_e4m3fn") + + scale = scale_float.to(scale_dtype) + scale_for_math = scale.float() + reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) + + if fmt is MoeFormat.MXFP8: + data_dtype = _require_torch_dtype("float8_e4m3fn") + data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] + else: + codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) + low = codes[..., 0::2] + high = codes[..., 1::2] + data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] + + return BlockScaledTensor( + data=data.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format=fmt, + logical_shape=logical_shape, + axis=axis, + ) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +@dataclass(frozen=True) +class _DispatchPlan: + """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" + + send_expert: torch.Tensor # local expert id per sent route + send_weight: torch.Tensor # router weight per sent route + send_token_idx: torch.Tensor # source token per sent route + send_slot_idx: torch.Tensor # source top-k slot per sent route + send_counts: Tuple[int, ...] # routes sent to each rank + recv_counts: Tuple[int, ...] # routes received from each rank + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +def _decode_tensor( + tensor: MoeTensor, + *, + name: str, + expected_shape: Tuple[int, ...], + quantized_axis: int, +) -> torch.Tensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.logical_shape != expected_shape: + raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") + if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): + raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") + return tensor.dequantize() + + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") + if not tensor.is_floating_point(): + raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") + return tensor.float() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=-1).dequantize() + + +class MoeEpReference: + """Reference implementation of routed SwiGLU experts plus EP dispatch. + + Global experts are assigned contiguously: rank ``r`` owns + ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit + initialized process group for multi-rank execution; ``None`` means a + one-rank reference even if the default distributed group is initialized. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + ) -> None: + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + + if ep_group is None: + ep_size, ep_rank = 1, 0 + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("ep_group requires an initialized torch.distributed process group") + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + + for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): + required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + if hidden_size % required_multiple != 0: + raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(" + f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " + f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " + f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " + f"output={self.output_format.value}, combine={self.combine_format.value})" + ) + + def _collective_device(self, device: torch.device) -> torch.device: + """Device the process group can run ``all_to_all_single`` on. + + Gloo only implements all-to-all for CPU tensors, so CUDA tensors are + staged through host memory; NCCL groups communicate in place. + """ + if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return send_counts.clone() + comm_device = self._collective_device(send_counts.device) + staged = send_counts.to(comm_device) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single(recv_counts, staged, group=self.ep_group) + return recv_counts.to(send_counts.device) + + def _all_to_all( + self, + send: torch.Tensor, + send_counts: Sequence[int], + recv_counts: Sequence[int], + ) -> torch.Tensor: + if self.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + output_shape = (sum(recv_counts), *send.shape[1:]) + recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) + dist.all_to_all_single( + recv, + staged, + output_split_sizes=list(recv_counts), + input_split_sizes=list(send_counts), + group=self.ep_group, + ) + return recv.to(send.device) + + def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: + """Route valid ``topk_idx`` entries to destination ranks, stably by rank. + + Backward reuses this so gradient re-dispatch reproduces the exact + forward route order. + """ + + device = topk_idx.device + token_count = topk_idx.shape[0] + flat_expert = topk_idx.reshape(-1).to(torch.int64) + flat_weight = topk_weights.reshape(-1).float() + valid = flat_expert != -1 + invalid_negative = flat_expert < -1 + invalid_high = flat_expert >= self.num_experts + if bool((invalid_negative | invalid_high).any().item()): + bad = flat_expert[invalid_negative | invalid_high][0].item() + raise ValueError(f"topk_idx contains out-of-range expert id {bad}") + + flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) + flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) + expert = flat_expert[valid] + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + order = torch.argsort(destination, stable=True) + + send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), + send_weight=flat_weight[valid].index_select(0, order), + send_token_idx=flat_token[valid].index_select(0, order), + send_slot_idx=flat_slot[valid].index_select(0, order), + send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), + recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), + ) + + def _run_local_experts( + self, + tokens: torch.Tensor, + local_expert_idx: torch.Tensor, + route_weight: torch.Tensor, + fc1_weight: torch.Tensor, + fc2_weight: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output = torch.empty( + (tokens.shape[0], self.hidden_size), + dtype=torch.float32, + device=tokens.device, + ) + fc1_c_rows = [] if self.generate_c else None + for expert in range(self.experts_per_rank): + positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + expert_tokens = tokens.index_select(0, positions) + gate_up = expert_tokens @ fc1_weight[expert] + if fc1_c_rows is not None: + # Raw pre-SwiGLU accumulator: before clamp, no router weight. + fc1_c_rows.append(gate_up.to(torch.bfloat16)) + gate, up = gate_up.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + gate = gate.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up + weights = route_weight.index_select(0, positions).unsqueeze(-1) + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + expert_output = intermediate @ fc2_weight[expert] + if not self.apply_topk_in_fc1: + expert_output = expert_output * weights + expert_output = _format_round_trip(expert_output, self.combine_format) + output.index_copy_(0, positions, expert_output) + fc1_c = None + if fc1_c_rows is not None: + fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) + return output, fc1_c + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: + """Run dispatch, local experts, return routing, top-k reduce, and encode. + + Shapes: + activation: ``(T, H)`` + fc1_weight: ``(E_local, H, 2 * I)`` + fc2_weight: ``(E_local, I, H)`` + topk_idx/topk_weights: ``(T, K)`` + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when constructed with ``generate_c=True``. ``fc1_c`` is the BF16 + pre-SwiGLU FC1 accumulator of every route this rank's experts + processed, ``(local_routes, 2 * I)``, grouped by local expert and + ordered within each expert by (source rank, source token-major route + order); captured before the gate/up clamp, without the router weight. + ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies + the route behind ``fc1_c`` row ``i`` for the backward gradient + re-dispatch. + """ + + if topk_idx.ndim != 2: + raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") + token_count = topk_idx.shape[0] + route_shape = (token_count, self.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: + raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + + device = _tensor_device(activation) + inputs = { + "fc1_weight": _tensor_device(fc1_weight), + "fc2_weight": _tensor_device(fc2_weight), + "topk_idx": topk_idx.device, + "topk_weights": topk_weights.device, + } + for name, input_device in inputs.items(): + if input_device != device: + raise ValueError(f"{name} must be on {device}, got {input_device}") + + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + + plan = self._dispatch_plan(topk_idx, topk_weights) + send_token_idx = plan.send_token_idx + send_slot_idx = plan.send_slot_idx + send_counts, recv_counts = plan.send_counts, plan.recv_counts + send_tokens = activation_float.index_select(0, send_token_idx) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + + route_metadata = None + if self.generate_c: + recv_src_rank = torch.repeat_interleave( + torch.arange(self.ep_size, device=device), + torch.tensor(recv_counts, device=device), + ) + recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) + recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) + # Stable sort by local expert reproduces the fc1_c row order + # (grouped by expert; source order preserved within each group). + fc1_c_order = torch.argsort(recv_expert, stable=True) + route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + # recv rows are ordered by source rank, then that source's token-major + # route order, so the per-expert position grouping below realizes the + # documented fc1_c ordering. + recv_output, fc1_c = self._run_local_experts( + recv_tokens, + recv_expert, + recv_weight, + fc1_float, + fc2_float, + ) + + returned = self._all_to_all(recv_output, recv_counts, send_counts) + combine_plane = torch.zeros( + (token_count * self.top_k, self.hidden_size), + dtype=torch.float32, + device=device, + ) + send_flat_slot = send_token_idx * self.top_k + send_slot_idx + combine_plane.index_copy_(0, send_flat_slot, returned) + reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) + + if self.output_format is MoeFormat.BF16: + output = reduced.to(torch.bfloat16) + else: + output = quantize_blockwise(reduced, self.output_format, axis=-1) + if self.generate_c: + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward pass consuming the ``generate_c=True`` stash. + + ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, + and the FC2 input are all rebuilt from it, so no post-SwiGLU forward + intermediate needs to be saved. ``route_metadata`` alone reconstructs + the mapping between re-dispatched rows and ``fc1_c`` rows and drives + the gradient return scatter. + + Quantization round-trips (input decode, ``combine_format``, + ``output_format``) are treated as straight-through identities; + ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. + + Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, + grad_topk_weights)`` in float32. + """ + + if not self.generate_c: + raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + token_count = topk_idx.shape[0] + if tuple(grad_output.shape) != (token_count, self.hidden_size): + raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + + device = _tensor_device(activation) + two_i = 2 * self.intermediate_size + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + if fc1_c.shape != (int(route_metadata.shape[0]), two_i): + raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") + + # Re-dispatch the FC1 inputs, router weights, and output gradients + # along the identical forward routes. + plan = self._dispatch_plan(topk_idx, topk_weights) + send_counts, recv_counts = plan.send_counts, plan.recv_counts + grad_output_float = grad_output.float() + recv_tokens = self._all_to_all(activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + + # route_metadata rows are in fc1_c order; sorting them by + # (src_rank, src_token, src_slot) reproduces the receive order, giving + # the permutation between re-dispatched rows and fc1_c rows. + metadata = route_metadata.to(device=device, dtype=torch.int64) + local_routes = metadata.shape[0] + if local_routes > 0: + token_span = int(metadata[:, 2].max().item()) + 1 + recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] + perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j + else: + perm = torch.empty((0,), dtype=torch.int64, device=device) + x_rows = torch.empty_like(recv_tokens) + x_rows.index_copy_(0, perm, recv_tokens) + w_rows = torch.empty_like(recv_weight) + w_rows.index_copy_(0, perm, recv_weight) + dy_rows = torch.empty_like(recv_grad) + dy_rows.index_copy_(0, perm, recv_grad) + + c_rows = fc1_c.float() + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + grad_fc1 = torch.zeros_like(fc1_float) + grad_fc2 = torch.zeros_like(fc2_float) + for expert in range(self.experts_per_rank): + positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + c = c_rows.index_select(0, positions) + x = x_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = dy_rows.index_select(0, positions) + + gate, up = c.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + g = gate.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate, up + sig = torch.sigmoid(g) + s = g * sig + h = s * u + + if self.apply_topk_in_fc1: + h_fc2 = h * w + d_y_pre = d_y + else: + h_fc2 = h + d_y_pre = d_y * w + grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) + + d_g = d_h * u * (sig * (1 + g * (1 - sig))) + d_u = d_h * s + if self.gate_up_clamp is not None: + d_gate = d_g * (gate <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1) + grad_fc1[expert] = x.transpose(0, 1) @ d_c + d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) + + # Return the route gradients to their source ranks and scatter-add. + returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) + grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + return ( + grad_activation, + grad_fc1, + grad_fc2, + grad_topk_weights.view(token_count, self.top_k), + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "quantize_blockwise", +] diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 4778498b7d..5b965e10bc 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,8 @@ import torch import torch.distributed as dist +from moe_ep_reference import MoeEpReference +from transformer_engine.pytorch import ops as te_ops from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -121,6 +123,30 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _make_moe_inputs(rank, ep_size, device="cuda"): + """Random activations and top-k routing representative of a router output.""" + generator = torch.Generator(device=device) + generator.manual_seed(2026 + rank) + num_experts = ep_size * NUM_LOCAL_EXPERTS + tokens = torch.randn( + TOKENS_PER_RANK, + HIDDEN_DIM, + generator=generator, + dtype=torch.float32, + device=device, + ).mul_(0.25) + router_logits = torch.randn( + TOKENS_PER_RANK, + num_experts, + generator=generator, + dtype=torch.float32, + device=device, + ) + topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) + topk_weights = torch.softmax(topk_logits, dim=-1) + return topk_idx, tokens.to(torch.bfloat16), topk_weights + + class _Cfg: rank: int world_size: int @@ -554,6 +580,140 @@ def zero_grads(): rtol=5e-2, ) + @_eager_test_include + def test_fusible_dispatch_combine_moe(self): + """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" + if not EAGER: + self.skipTest( + "variable grouped-linear splits require eager EP output sizing" + ) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, self.cfg.ep_size + ) + num_experts = NUM_LOCAL_EXPERTS + intermediate_dim = HIDDEN_DIM + + dispatch_buffer = self._make_buffer() + dispatch = te_ops.Dispatch(dispatch_buffer) + fc1 = te_ops.GroupedLinear( + num_experts, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + num_experts, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(dispatch_buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "m_splits") + dispatch.set_extra_output_channel(1, "token_probs") + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "token_probs") + fc2.set_extra_input_channel(0, "m_splits") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(1234 + self.cfg.rank) + with torch.no_grad(): + for expert_idx in range(num_experts): + getattr(fc1, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + getattr(fc2, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + + ref_fc1_weights = torch.stack( + [ + getattr(fc1, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + ref_fc2_weights = torch.stack( + [ + getattr(fc2, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=intermediate_dim, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + apply_topk_in_fc1=True, + generate_c=True, + ) + ref_output, fc1_c, route_metadata = reference( + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + ) + + test_tokens = tokens.detach().clone().requires_grad_(True) + test_topk_weights = topk_weights.detach().clone().requires_grad_(True) + test_output = model(test_tokens, topk_idx, test_topk_weights) + + grad_output = torch.linspace( + -0.2, + 0.2, + TOKENS_PER_RANK * HIDDEN_DIM, + device=self.cfg.device, + dtype=torch.float32, + ).reshape(TOKENS_PER_RANK, HIDDEN_DIM) + grad_output = grad_output.to(torch.bfloat16) + ref_grads = reference.backward( + grad_output, + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + test_output.backward(grad_output) + torch.cuda.synchronize() + + torch.testing.assert_close( + test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_topk_weights.grad, + ref_grads[3], + atol=5e-2, + rtol=5e-2, + ) + for expert_idx in range(num_experts): + torch.testing.assert_close( + getattr(fc1, f"weight{expert_idx}").grad.float(), + ref_grads[1][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + torch.testing.assert_close( + getattr(fc2, f"weight{expert_idx}").grad.float(), + ref_grads[2][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + @_zero_copy_test_include @_eager_test_include def test_combine_autograd(self): diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..4ea116793f 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -23,6 +23,8 @@ from .basic_linear import BasicLinear from .bias import Bias from .constant_scale import ConstantScale +from .combine import Combine +from .dispatch import Dispatch from .dropout import Dropout from .grouped_linear import GroupedLinear from .identity import Identity diff --git a/transformer_engine/pytorch/ops/combine.py b/transformer_engine/pytorch/ops/combine.py new file mode 100644 index 0000000000..19f78f7356 --- /dev/null +++ b/transformer_engine/pytorch/ops/combine.py @@ -0,0 +1,153 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel combine operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, is_symm_backed +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_grad_buffer( + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError( + f"grad_out shape {tuple(tensor.shape)} does not match {shape}." + ) + if tensor.dtype is not dtype: + raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError("grad_out must be contiguous.") + if tensor.requires_grad: + raise ValueError("grad_out must not require gradients.") + return tensor + + +class Combine(BasicOperation): + """Combine pre-weighted local expert outputs with NCCL EP. + + The operation uses routing state produced by a :class:`Dispatch` with the + same :class:`EpBuffer`. + """ + + def __init__( + self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None + ) -> None: + super().__init__() + self.buffer = buffer + self.num_local_tokens = ( + buffer.max_tokens_per_rank + if num_local_tokens is None + else int(num_local_tokens) + ) + if self.num_local_tokens < 0: + raise ValueError("num_local_tokens must be non-negative.") + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError( + f"NCCL EP requires BF16 combine input, got {input_.dtype}." + ) + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Combine input must have shape (R, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + + expert_out = input_ + if self.buffer.zero_copy: + expert_out = _alloc_io( + tuple(input_.shape), + input_.dtype, + input_.device, + True, + ) + expert_out.copy_(input_) + + result = torch.empty( + self.num_local_tokens, + self.buffer.hidden_dim, + dtype=input_.dtype, + device=input_.device, + ) + torch.ops.transformer_engine_ep.combine( + self.buffer.handle_mem, + expert_out, + result, + ) + + if ctx.requires_grad: + grad_out = kwargs.get("grad_out") + if self.buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes combine gradients per step and cannot use " + "a caller-supplied grad_out" + ) + grad_out = _validate_grad_buffer( + grad_out, + shape=tuple(input_.shape), + dtype=input_.dtype, + device=input_.device, + ) + if ( + self.buffer.zero_copy + and grad_out is not None + and not is_symm_backed(grad_out) + ): + raise ValueError( + "zero-copy Combine grad_out must be symmetric-memory-backed." + ) + ctx.grad_out = grad_out + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.save_for_backward(self.buffer.handle_mem) + + return result + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + grad_output, + grad_input, + ) + return grad_input, () + diff --git a/transformer_engine/pytorch/ops/dispatch.py b/transformer_engine/pytorch/ops/dispatch.py new file mode 100644 index 0000000000..ab89f6e130 --- /dev/null +++ b/transformer_engine/pytorch/ops/dispatch.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel dispatch operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, ep_prepare +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_output_buffer( + name: str, + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous.") + if tensor.requires_grad: + raise ValueError(f"{name} must not require gradients.") + return tensor + + +class Dispatch(BasicOperation): + """Dispatch BF16 tokens to local experts with NCCL EP. + + The extra inputs are routing indices and FP32 routing weights. The extra + outputs are local tokens-per-expert and received routing weights. + """ + + num_extra_inputs: int = 2 + num_extra_outputs: int = 2 + + def __init__(self, buffer: EpBuffer) -> None: + super().__init__() + self.buffer = buffer + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + + if input_.dtype is not torch.bfloat16: + raise NotImplementedError( + f"NCCL EP requires BF16 dispatch input, got {input_.dtype}." + ) + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}.") + expected_route_shape = (input_.shape[0], self.buffer.top_k) + if tuple(topk_idx.shape) != expected_route_shape: + raise ValueError( + f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." + ) + if tuple(topk_weights.shape) != expected_route_shape: + raise ValueError( + f"topk_weights shape must be {expected_route_shape}, " + f"got {tuple(topk_weights.shape)}." + ) + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): + if tensor.device != input_.device: + raise ValueError( + f"{name} must be on {input_.device}, got {tensor.device}." + ) + + recv_tokens = kwargs.get("recv_tokens") + recv_topk_weights = kwargs.get("recv_topk_weights") + if self.buffer.eager and ( + recv_tokens is not None or recv_topk_weights is not None + ): + raise ValueError( + "eager mode sizes dispatch outputs per step and cannot use " + "caller-supplied receive buffers" + ) + + tokens_per_expert = ep_prepare(self.buffer, topk_idx) + rows = ( + self.buffer._host_total_recv_tokens + if self.buffer.eager + else self.buffer.recv_capacity_per_rank + ) + if rows is None: + raise RuntimeError("NCCL EP dispatch receive size is unavailable.") + rows = int(rows) + recv_shape = (rows, self.buffer.hidden_dim) + recv_tokens = _validate_output_buffer( + "recv_tokens", + recv_tokens, + shape=recv_shape, + dtype=self.buffer.payload_dtype, + device=self.buffer.device, + ) + recv_topk_weights = _validate_output_buffer( + "recv_topk_weights", + recv_topk_weights, + shape=(rows,), + dtype=torch.float32, + device=self.buffer.device, + ) + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (rows,), + torch.float32, + self.buffer.device, + self.buffer.zero_copy, + ) + + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.save_for_backward(self.buffer.handle_mem) + + return recv_tokens, [(tokens_per_expert, recv_topk_weights)] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + + grad_recv_weights = basic_op_grad_extra_outputs[0][1] + if grad_recv_weights is None: + grad_recv_weights = torch.zeros( + grad_output.shape[0], + dtype=torch.float32, + device=grad_output.device, + ) + else: + grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous() + + grad_input = torch.empty( + ctx.input_shape, + dtype=ctx.input_dtype, + device=grad_output.device, + ) + grad_topk_weights = torch.empty( + ctx.topk_weights_shape, + dtype=torch.float32, + device=grad_output.device, + ) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + grad_output, + grad_recv_weights, + grad_input, + grad_topk_weights, + ) + return grad_input, [()], [(None, grad_topk_weights)] + From ad3b044f08bd5821850930c4debd39fb3b775c4c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 5 Aug 2026 01:15:06 +0000 Subject: [PATCH 06/83] fusible ops test Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 134 ++++++++++++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 5b965e10bc..2c81782a0e 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -714,6 +714,140 @@ def test_fusible_dispatch_combine_moe(self): rtol=5e-2, ) + @_eager_test_include + def test_fusible_dispatch_combine_moe(self): + """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" + if not EAGER: + self.skipTest( + "variable grouped-linear splits require eager EP output sizing" + ) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, self.cfg.ep_size + ) + num_experts = NUM_LOCAL_EXPERTS + intermediate_dim = HIDDEN_DIM + + dispatch_buffer = self._make_buffer() + dispatch = te_ops.Dispatch(dispatch_buffer) + fc1 = te_ops.GroupedLinear( + num_experts, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + num_experts, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(dispatch_buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "m_splits") + dispatch.set_extra_output_channel(1, "token_probs") + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "token_probs") + fc2.set_extra_input_channel(0, "m_splits") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(1234 + self.cfg.rank) + with torch.no_grad(): + for expert_idx in range(num_experts): + getattr(fc1, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + getattr(fc2, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + + ref_fc1_weights = torch.stack( + [ + getattr(fc1, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + ref_fc2_weights = torch.stack( + [ + getattr(fc2, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=intermediate_dim, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + apply_topk_in_fc1=True, + generate_c=True, + ) + ref_output, fc1_c, route_metadata = reference( + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + ) + + test_tokens = tokens.detach().clone().requires_grad_(True) + test_topk_weights = topk_weights.detach().clone().requires_grad_(True) + test_output = model(test_tokens, topk_idx, test_topk_weights) + + grad_output = torch.linspace( + -0.2, + 0.2, + TOKENS_PER_RANK * HIDDEN_DIM, + device=self.cfg.device, + dtype=torch.float32, + ).reshape(TOKENS_PER_RANK, HIDDEN_DIM) + grad_output = grad_output.to(torch.bfloat16) + ref_grads = reference.backward( + grad_output, + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + test_output.backward(grad_output) + torch.cuda.synchronize() + + torch.testing.assert_close( + test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_topk_weights.grad, + ref_grads[3], + atol=5e-2, + rtol=5e-2, + ) + for expert_idx in range(num_experts): + torch.testing.assert_close( + getattr(fc1, f"weight{expert_idx}").grad.float(), + ref_grads[1][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + torch.testing.assert_close( + getattr(fc2, f"weight{expert_idx}").grad.float(), + ref_grads[2][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + @_zero_copy_test_include @_eager_test_include def test_combine_autograd(self): From 5fb0d3af53e90bf8cf421c705e018307c71676df Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:46:15 +0000 Subject: [PATCH 07/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/moe_ep_reference.py | 90 ++++++++++++++----- tests/pytorch/distributed/run_ep.py | 68 ++++---------- tests/pytorch/test_fusible_ops.py | 16 +--- transformer_engine/pytorch/ops/combine.py | 27 ++---- transformer_engine/pytorch/ops/dispatch.py | 13 +-- transformer_engine/pytorch/ops/fuser.py | 55 +++--------- 6 files changed, 113 insertions(+), 156 deletions(-) diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/tests/pytorch/distributed/moe_ep_reference.py index 3fab88a50b..280c4f0d21 100644 --- a/tests/pytorch/distributed/moe_ep_reference.py +++ b/tests/pytorch/distributed/moe_ep_reference.py @@ -123,7 +123,9 @@ def _validate_storage(self) -> None: expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") data_shape = self.logical_shape if self.data.dtype != expected_dtype: - raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) else: expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -298,7 +300,9 @@ def _decode_tensor( ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: - raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") + raise ValueError( + f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" + ) if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") return tensor.dequantize() @@ -357,11 +361,15 @@ def __init__( ep_size, ep_rank = 1, 0 else: if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("ep_group requires an initialized torch.distributed process group") + raise RuntimeError( + "ep_group requires an initialized torch.distributed process group" + ) ep_size = dist.get_world_size(ep_group) ep_rank = dist.get_rank(ep_group) if num_experts % ep_size != 0: - raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + raise ValueError( + f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" + ) self.num_experts = num_experts self.hidden_size = hidden_size @@ -378,10 +386,18 @@ def __init__( self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) - for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): - required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = ( + 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + ) if hidden_size % required_multiple != 0: - raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" + f" {name}={fmt.value}" + ) def __repr__(self) -> str: return ( @@ -456,7 +472,9 @@ def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") order = torch.argsort(destination, stable=True) - send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) recv_counts_tensor = self._exchange_counts(send_counts_tensor) return _DispatchPlan( send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), @@ -505,7 +523,13 @@ def _run_local_experts( output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: - fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) + fc1_c = ( + torch.cat(fc1_c_rows) + if fc1_c_rows + else torch.empty( + (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device + ) + ) return output, fc1_c def __call__( @@ -543,13 +567,17 @@ def __call__( if tuple(topk_idx.shape) != route_shape: raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") if tuple(topk_weights.shape) != route_shape: - raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) if topk_idx.dtype not in (torch.int32, torch.int64): raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") if not topk_weights.is_floating_point(): raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) device = _tensor_device(activation) inputs = { @@ -602,7 +630,11 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) # recv rows are ordered by source rank, then that source's token-major # route order, so the per-expert position grouping below realizes the # documented fc1_c ordering. @@ -660,10 +692,15 @@ def backward( """ if not self.generate_c: - raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) token_count = topk_idx.shape[0] if tuple(grad_output.shape) != (token_count, self.hidden_size): - raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") + raise ValueError( + f"grad_output shape must be {(token_count, self.hidden_size)}, got" + f" {tuple(grad_output.shape)}" + ) if not grad_output.is_floating_point(): raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") @@ -688,16 +725,23 @@ def backward( quantized_axis=1, ) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): - raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") + raise ValueError( + f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" + f" {tuple(fc1_c.shape)}" + ) # Re-dispatch the FC1 inputs, router weights, and output gradients # along the identical forward routes. plan = self._dispatch_plan(topk_idx, topk_weights) send_counts, recv_counts = plan.send_counts, plan.recv_counts grad_output_float = grad_output.float() - recv_tokens = self._all_to_all(activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_tokens = self._all_to_all( + activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) - recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_grad = self._all_to_all( + grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) # route_metadata rows are in fc1_c order; sorting them by # (src_rank, src_token, src_slot) reproduces the receive order, giving @@ -771,10 +815,16 @@ def backward( # Return the route gradients to their source ranks and scatter-add. returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) - grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation = torch.zeros( + (token_count, self.hidden_size), dtype=torch.float32, device=device + ) grad_activation.index_add_(0, plan.send_token_idx, returned_dx) - grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) - grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + grad_topk_weights = torch.zeros( + (token_count * self.top_k,), dtype=torch.float32, device=device + ) + grad_topk_weights.index_copy_( + 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw + ) return ( grad_activation, grad_fc1, diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 2c81782a0e..a69418088c 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -584,13 +584,9 @@ def zero_grads(): def test_fusible_dispatch_combine_moe(self): """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" if not EAGER: - self.skipTest( - "variable grouped-linear splits require eager EP output sizing" - ) + self.skipTest("variable grouped-linear splits require eager EP output sizing") - topk_idx, tokens, topk_weights = _make_moe_inputs( - self.cfg.rank, self.cfg.ep_size - ) + topk_idx, tokens, topk_weights = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) num_experts = NUM_LOCAL_EXPERTS intermediate_dim = HIDDEN_DIM @@ -626,24 +622,14 @@ def test_fusible_dispatch_combine_moe(self): generator.manual_seed(1234 + self.cfg.rank) with torch.no_grad(): for expert_idx in range(num_experts): - getattr(fc1, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) - getattr(fc2, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) + getattr(fc1, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) + getattr(fc2, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) ref_fc1_weights = torch.stack( - [ - getattr(fc1, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc1, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) ref_fc2_weights = torch.stack( - [ - getattr(fc2, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc2, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) reference = MoeEpReference( num_experts=self.cfg.num_experts, @@ -688,12 +674,8 @@ def test_fusible_dispatch_combine_moe(self): test_output.backward(grad_output) torch.cuda.synchronize() - torch.testing.assert_close( - test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 - ) - torch.testing.assert_close( - test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 - ) + torch.testing.assert_close(test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2) torch.testing.assert_close( test_topk_weights.grad, ref_grads[3], @@ -718,13 +700,9 @@ def test_fusible_dispatch_combine_moe(self): def test_fusible_dispatch_combine_moe(self): """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" if not EAGER: - self.skipTest( - "variable grouped-linear splits require eager EP output sizing" - ) + self.skipTest("variable grouped-linear splits require eager EP output sizing") - topk_idx, tokens, topk_weights = _make_moe_inputs( - self.cfg.rank, self.cfg.ep_size - ) + topk_idx, tokens, topk_weights = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) num_experts = NUM_LOCAL_EXPERTS intermediate_dim = HIDDEN_DIM @@ -760,24 +738,14 @@ def test_fusible_dispatch_combine_moe(self): generator.manual_seed(1234 + self.cfg.rank) with torch.no_grad(): for expert_idx in range(num_experts): - getattr(fc1, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) - getattr(fc2, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) + getattr(fc1, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) + getattr(fc2, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) ref_fc1_weights = torch.stack( - [ - getattr(fc1, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc1, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) ref_fc2_weights = torch.stack( - [ - getattr(fc2, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc2, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) reference = MoeEpReference( num_experts=self.cfg.num_experts, @@ -822,12 +790,8 @@ def test_fusible_dispatch_combine_moe(self): test_output.backward(grad_output) torch.cuda.synchronize() - torch.testing.assert_close( - test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 - ) - torch.testing.assert_close( - test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 - ) + torch.testing.assert_close(test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2) torch.testing.assert_close( test_topk_weights.grad, ref_grads[3], diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 9cecd4531b..dba1cab263 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -523,9 +523,7 @@ def fuser_forward( ): m_splits, probs = basic_op_extra_inputs[0] # Stub row-id map: real Dispatch would emit permute indices. - routing_map = torch.arange( - input_.size(0), device=input_.device, dtype=torch.int64 - ) + routing_map = torch.arange(input_.size(0), device=input_.device, dtype=torch.int64) return input_, [(m_splits, probs, routing_map)] def fuser_backward( @@ -568,9 +566,7 @@ def fuser_forward( if routing_map is None: raise RuntimeError("FakeCombine expected routing_map channel") if int(routing_map.numel()) != int(input_.size(0)): - raise RuntimeError( - "FakeCombine routing_map length does not match tokens" - ) + raise RuntimeError("FakeCombine routing_map length does not match tokens") return input_, [()] def fuser_backward( @@ -583,9 +579,7 @@ def fuser_backward( del basic_op_grad_extra_outputs return grad_output, [()], [(None,)] - split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[ - :group_size - ] + split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[:group_size] num_tokens = int(split_sizes.sum()) in_shape = (num_tokens, hidden_size) @@ -855,9 +849,7 @@ def op_backward(self, *args, **kwargs): def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): return input_, [(input_, input_)] - def fuser_backward( - self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs - ): + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): g0, g1 = basic_op_grad_extra_outputs[0] grad_extra = torch.zeros_like(grad_output) if g0 is not None: diff --git a/transformer_engine/pytorch/ops/combine.py b/transformer_engine/pytorch/ops/combine.py index 19f78f7356..4d4740c517 100644 --- a/transformer_engine/pytorch/ops/combine.py +++ b/transformer_engine/pytorch/ops/combine.py @@ -25,9 +25,7 @@ def _validate_grad_buffer( if tensor is None: return None if tuple(tensor.shape) != shape: - raise ValueError( - f"grad_out shape {tuple(tensor.shape)} does not match {shape}." - ) + raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") if tensor.dtype is not dtype: raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") if tensor.device != device: @@ -46,15 +44,11 @@ class Combine(BasicOperation): same :class:`EpBuffer`. """ - def __init__( - self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None - ) -> None: + def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: super().__init__() self.buffer = buffer self.num_local_tokens = ( - buffer.max_tokens_per_rank - if num_local_tokens is None - else int(num_local_tokens) + buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) ) if self.num_local_tokens < 0: raise ValueError("num_local_tokens must be non-negative.") @@ -70,9 +64,7 @@ def op_forward( ) -> torch.Tensor: del prev_op_grad_output_quantizer, next_op_input_quantizer if input_.dtype is not torch.bfloat16: - raise NotImplementedError( - f"NCCL EP requires BF16 combine input, got {input_.dtype}." - ) + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: raise ValueError( f"Combine input must have shape (R, {self.buffer.hidden_dim}), " @@ -114,14 +106,8 @@ def op_forward( dtype=input_.dtype, device=input_.device, ) - if ( - self.buffer.zero_copy - and grad_out is not None - and not is_symm_backed(grad_out) - ): - raise ValueError( - "zero-copy Combine grad_out must be symmetric-memory-backed." - ) + if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): + raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") ctx.grad_out = grad_out ctx.input_shape = tuple(input_.shape) ctx.input_dtype = input_.dtype @@ -150,4 +136,3 @@ def op_backward( grad_input, ) return grad_input, () - diff --git a/transformer_engine/pytorch/ops/dispatch.py b/transformer_engine/pytorch/ops/dispatch.py index ab89f6e130..9762ed76f7 100644 --- a/transformer_engine/pytorch/ops/dispatch.py +++ b/transformer_engine/pytorch/ops/dispatch.py @@ -73,9 +73,7 @@ def fuser_forward( kwargs = basic_op_kwargs[0] if input_.dtype is not torch.bfloat16: - raise NotImplementedError( - f"NCCL EP requires BF16 dispatch input, got {input_.dtype}." - ) + raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: raise ValueError( f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " @@ -97,15 +95,11 @@ def fuser_forward( raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): if tensor.device != input_.device: - raise ValueError( - f"{name} must be on {input_.device}, got {tensor.device}." - ) + raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") recv_tokens = kwargs.get("recv_tokens") recv_topk_weights = kwargs.get("recv_topk_weights") - if self.buffer.eager and ( - recv_tokens is not None or recv_topk_weights is not None - ): + if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): raise ValueError( "eager mode sizes dispatch outputs per step and cannot use " "caller-supplied receive buffers" @@ -211,4 +205,3 @@ def fuser_backward( grad_topk_weights, ) return grad_input, [()], [(None, grad_topk_weights)] - diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index cc706ab91f..63d3f4a6be 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -108,9 +108,7 @@ def forward( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in fuser._basic_ops ] - for tensor, (op_idx, input_idx) in zip( - extra_inputs, fuser._external_extra_input_slots - ): + for tensor, (op_idx, input_idx) in zip(extra_inputs, fuser._external_extra_input_slots): basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops @@ -127,9 +125,7 @@ def forward( # 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] - ): + for input_idx, source in enumerate(fuser._basic_op_extra_input_sources[idx]): if source is None: continue producer_idx, output_idx = source @@ -141,10 +137,7 @@ def forward( raise RuntimeError( f"Extra tensor channel producer op {producer_idx} has not run" ) - if ( - output_idx >= len(producer_outputs) - or producer_outputs[output_idx] is None - ): + if output_idx >= len(producer_outputs) or producer_outputs[output_idx] is None: raise RuntimeError( f"Extra tensor channel producer op {producer_idx} " f"({type(fuser._basic_ops[producer_idx]).__name__}) " @@ -154,9 +147,7 @@ def forward( f"input {input_idx}" ) basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] - op_extra_inputs = [ - tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs - ] + op_extra_inputs = [tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None prev_op_grad_output_quantizer = None @@ -290,9 +281,7 @@ def backward( basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_outputs for op in basic_ops ] - for grad, (op_idx, output_idx) in zip( - grad_extra_outputs, external_extra_output_slots - ): + for grad, (op_idx, output_idx) in zip(grad_extra_outputs, external_extra_output_slots): basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops @@ -310,13 +299,9 @@ def backward( # Backward op. Supply gradients accumulated from every consumer of # each internal channel. for idx in basic_op_idxs: - for output_idx, channel in enumerate( - basic_op_extra_output_channels[idx] - ): + for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): if channel is not None: - basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get( - channel - ) + basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get(channel) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs ] @@ -341,9 +326,7 @@ def backward( continue channel = basic_op_extra_output_channels[producer_idx][output_idx] previous_grad = channel_grads.get(channel) - channel_grads[channel] = ( - grad if previous_grad is None else previous_grad + grad - ) + channel_grads[channel] = grad if previous_grad is None else previous_grad + grad # Flatten list of parameter gradients grad_params_flat = [] @@ -456,13 +439,9 @@ def __init__( f"Extra tensor channel {channel!r} consumed by op {op_idx} " f"({type(op).__name__}) has no earlier producer" ) - self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[ - channel - ] + self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[channel] consumed_channels.add(channel) - for output_idx, channel in enumerate( - self._basic_op_extra_output_channels[op_idx] - ): + for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): if channel is None: self._external_extra_output_slots.append((op_idx, output_idx)) continue @@ -497,12 +476,10 @@ def __init__( raise ValueError( f"Extra input {input_idx} of op {op_idx} " f"({type(op).__name__}) is bound to channel {channel!r} " - f"but has no producer" + "but has no producer" ) producer_idx, output_idx = source - producer_channel = self._basic_op_extra_output_channels[producer_idx][ - output_idx - ] + producer_channel = self._basic_op_extra_output_channels[producer_idx][output_idx] if producer_channel != channel: raise ValueError( f"Extra input {input_idx} of op {op_idx} " @@ -602,9 +579,7 @@ def maybe_fuse_ops( first_op_requiring_backward = self._num_basic_ops for op_idx in range(self._num_basic_ops): op_inputs = itertools.chain(self._basic_op_params[op_idx], extra_inputs[op_idx]) - if any( - tensor is not None and tensor.requires_grad for tensor in op_inputs - ): + if any(tensor is not None and tensor.requires_grad for tensor in op_inputs): first_op_requiring_backward = op_idx break @@ -694,9 +669,7 @@ def __call__( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in self._basic_ops ] - for tensor, (op_idx, input_idx) in zip( - extra_inputs, self._external_extra_input_slots - ): + for tensor, (op_idx, input_idx) in zip(extra_inputs, self._external_extra_input_slots): basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state From 3af2eccc59de25a97b4baf341853470c8ea6cc2a Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 5 Aug 2026 22:42:49 +0000 Subject: [PATCH 08/83] keep just ops infra changes Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/moe_ep_reference.py | 842 ------------------ tests/pytorch/distributed/run_ep.py | 258 ------ .../pytorch/ops/basic/__init__.py | 2 - transformer_engine/pytorch/ops/combine.py | 138 --- transformer_engine/pytorch/ops/dispatch.py | 207 ----- 5 files changed, 1447 deletions(-) delete mode 100644 tests/pytorch/distributed/moe_ep_reference.py delete mode 100644 transformer_engine/pytorch/ops/combine.py delete mode 100644 transformer_engine/pytorch/ops/dispatch.py diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/tests/pytorch/distributed/moe_ep_reference.py deleted file mode 100644 index 280c4f0d21..0000000000 --- a/tests/pytorch/distributed/moe_ep_reference.py +++ /dev/null @@ -1,842 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. - -The implementation deliberately favors readable semantics over performance. It -supports a one-rank execution path and a variable-size ``all_to_all_single`` EP -path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. - -The quantized tensor layouts are logical (unswizzled) layouts. A production -kernel may reorder scale factors internally without changing this API contract. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from enum import Enum -from typing import Optional, Sequence, Tuple, Union - -import torch -import torch.distributed as dist -import torch.nn.functional as F - - -class MoeFormat(str, Enum): - """Public and communication formats supported by the reference.""" - - BF16 = "bf16" - MXFP8 = "mxfp8" - NVFP4 = "nvfp4" - - -def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: - if isinstance(value, MoeFormat): - return value - try: - return MoeFormat(value.lower()) - except (AttributeError, ValueError) as exc: - choices = ", ".join(item.value for item in MoeFormat) - raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc - - -def _require_torch_dtype(name: str) -> torch.dtype: - dtype = getattr(torch, name, None) - if dtype is None: - raise RuntimeError(f"this PyTorch build does not provide torch.{name}") - return dtype - - -def _normalize_axis(axis: int, ndim: int) -> int: - normalized = axis + ndim if axis < 0 else axis - if normalized < 0 or normalized >= ndim: - raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") - return normalized - - -def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: - result = list(shape) - result[axis] = value - return tuple(result) - - -def _ceil_div(numerator: int, denominator: int) -> int: - return (numerator + denominator - 1) // denominator - - -@dataclass(frozen=True) -class BlockScaledTensor: - """Portable data-plus-scale representation for MXFP8 or NVFP4. - - ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` - has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with - two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` - replaces that axis by one scale per block. - """ - - data: torch.Tensor - scale: torch.Tensor - format: Union[MoeFormat, str] - logical_shape: Tuple[int, ...] - axis: int = -1 - - def __post_init__(self) -> None: - fmt = _parse_format(self.format) - if fmt is MoeFormat.BF16: - raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") - shape = tuple(int(dim) for dim in self.logical_shape) - if not shape or any(dim < 0 for dim in shape): - raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") - axis = _normalize_axis(self.axis, len(shape)) - object.__setattr__(self, "format", fmt) - object.__setattr__(self, "logical_shape", shape) - object.__setattr__(self, "axis", axis) - self._validate_storage() - - @property - def block_size(self) -> int: - return 32 if self.format is MoeFormat.MXFP8 else 16 - - @property - def shape(self) -> Tuple[int, ...]: - return self.logical_shape - - @property - def device(self) -> torch.device: - return self.data.device - - def _validate_storage(self) -> None: - if self.data.device != self.scale.device: - raise ValueError("block-scaled data and scale must be on the same device") - - logical_extent = self.logical_shape[self.axis] - scale_shape = _shape_with_axis( - self.logical_shape, - self.axis, - _ceil_div(logical_extent, self.block_size), - ) - if tuple(self.scale.shape) != scale_shape: - raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") - - if self.format is MoeFormat.MXFP8: - expected_dtype = _require_torch_dtype("float8_e4m3fn") - expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") - data_shape = self.logical_shape - if self.data.dtype != expected_dtype: - raise TypeError( - f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" - ) - else: - expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") - fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) - if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: - raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") - data_shape = _shape_with_axis( - self.logical_shape, - self.axis, - _ceil_div(logical_extent, 2), - ) - - if tuple(self.data.shape) != data_shape: - raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") - if self.scale.dtype != expected_scale_dtype: - raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") - - def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: - """Return the logical tensor with block scales applied.""" - - logical_extent = self.logical_shape[self.axis] - scale = self.scale.movedim(self.axis, -1).float() - expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] - - if self.format is MoeFormat.MXFP8: - values = self.data.movedim(self.axis, -1).float() - else: - packed = self.data - if packed.dtype != torch.uint8: - packed = packed.view(torch.uint8) - packed = packed.movedim(self.axis, -1) - low = packed & 0x0F - high = packed >> 4 - codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] - table = torch.tensor( - [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ], - dtype=torch.float32, - device=packed.device, - ) - values = table[codes.long()] - - return (values * expanded_scale).movedim(-1, self.axis).to(dtype) - - -def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: - """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" - - levels = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], - dtype=torch.float32, - device=values.device, - ) - magnitudes = values.abs().unsqueeze(-1) - distances = (magnitudes - levels).abs() - minimum = distances.amin(dim=-1, keepdim=True) - candidates = distances == minimum - codes = torch.arange(8, dtype=torch.int64, device=values.device) - any_code = torch.where(candidates, codes, 8).amin(dim=-1) - even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) - magnitude_code = torch.where(even_code < 8, even_code, any_code) - sign_code = torch.signbit(values).to(torch.int64) << 3 - return magnitude_code | sign_code - - -def quantize_blockwise( - tensor: torch.Tensor, - format: Union[MoeFormat, str], - *, - axis: int = -1, -) -> BlockScaledTensor: - """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. - - MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward - positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and - E4M3 scales rounded to nearest. - """ - - fmt = _parse_format(format) - if fmt is MoeFormat.BF16: - raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") - if not tensor.is_floating_point(): - raise TypeError(f"tensor must be floating point, got {tensor.dtype}") - - axis = _normalize_axis(axis, tensor.ndim) - logical_shape = tuple(tensor.shape) - moved = tensor.float().movedim(axis, -1) - logical_extent = moved.shape[-1] - block_size = 32 if fmt is MoeFormat.MXFP8 else 16 - block_count = _ceil_div(logical_extent, block_size) - padded_extent = block_count * block_size - if padded_extent != logical_extent: - moved = F.pad(moved, (0, padded_extent - logical_extent)) - blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) - - value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 - scale_float = blocks.abs().amax(dim=-1) / value_limit - if fmt is MoeFormat.MXFP8: - safe_scale = torch.where(scale_float > 0, scale_float, 1.0) - scale_float = torch.where( - scale_float > 0, - torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), - torch.zeros_like(scale_float), - ) - scale_dtype = _require_torch_dtype("float8_e8m0fnu") - else: - scale_dtype = _require_torch_dtype("float8_e4m3fn") - - scale = scale_float.to(scale_dtype) - scale_for_math = scale.float() - reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) - normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) - - if fmt is MoeFormat.MXFP8: - data_dtype = _require_torch_dtype("float8_e4m3fn") - data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] - else: - codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) - low = codes[..., 0::2] - high = codes[..., 1::2] - data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] - - return BlockScaledTensor( - data=data.movedim(-1, axis).contiguous(), - scale=scale.movedim(-1, axis).contiguous(), - format=fmt, - logical_shape=logical_shape, - axis=axis, - ) - - -MoeTensor = Union[torch.Tensor, BlockScaledTensor] - - -@dataclass(frozen=True) -class _DispatchPlan: - """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" - - send_expert: torch.Tensor # local expert id per sent route - send_weight: torch.Tensor # router weight per sent route - send_token_idx: torch.Tensor # source token per sent route - send_slot_idx: torch.Tensor # source top-k slot per sent route - send_counts: Tuple[int, ...] # routes sent to each rank - recv_counts: Tuple[int, ...] # routes received from each rank - - -def _tensor_device(tensor: MoeTensor) -> torch.device: - return tensor.device - - -def _decode_tensor( - tensor: MoeTensor, - *, - name: str, - expected_shape: Tuple[int, ...], - quantized_axis: int, -) -> torch.Tensor: - if isinstance(tensor, BlockScaledTensor): - if tensor.logical_shape != expected_shape: - raise ValueError( - f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" - ) - if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): - raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") - return tensor.dequantize() - - if tuple(tensor.shape) != expected_shape: - raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") - if not tensor.is_floating_point(): - raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") - return tensor.float() - - -def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: - if format is MoeFormat.BF16: - return tensor.to(torch.bfloat16).float() - return quantize_blockwise(tensor, format, axis=-1).dequantize() - - -class MoeEpReference: - """Reference implementation of routed SwiGLU experts plus EP dispatch. - - Global experts are assigned contiguously: rank ``r`` owns - ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit - initialized process group for multi-rank execution; ``None`` means a - one-rank reference even if the default distributed group is initialized. - """ - - def __init__( - self, - *, - num_experts: int, - hidden_size: int, - intermediate_size: int, - top_k: int, - ep_group: Optional[dist.ProcessGroup] = None, - max_tokens_per_rank: Optional[int] = None, - output_format: Union[MoeFormat, str] = MoeFormat.BF16, - combine_format: Union[MoeFormat, str] = MoeFormat.BF16, - apply_topk_in_fc1: bool = True, - gate_up_clamp: Optional[float] = None, - generate_c: bool = False, - ) -> None: - for name, value in ( - ("num_experts", num_experts), - ("hidden_size", hidden_size), - ("intermediate_size", intermediate_size), - ("top_k", top_k), - ): - if not isinstance(value, int) or value <= 0: - raise ValueError(f"{name} must be a positive integer, got {value!r}") - if top_k > num_experts: - raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") - if max_tokens_per_rank is not None and max_tokens_per_rank < 0: - raise ValueError("max_tokens_per_rank must be non-negative") - - if ep_group is None: - ep_size, ep_rank = 1, 0 - else: - if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError( - "ep_group requires an initialized torch.distributed process group" - ) - ep_size = dist.get_world_size(ep_group) - ep_rank = dist.get_rank(ep_group) - if num_experts % ep_size != 0: - raise ValueError( - f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" - ) - - self.num_experts = num_experts - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.top_k = top_k - self.ep_group = ep_group - self.ep_size = ep_size - self.ep_rank = ep_rank - self.experts_per_rank = num_experts // ep_size - self.max_tokens_per_rank = max_tokens_per_rank - self.output_format = _parse_format(output_format) - self.combine_format = _parse_format(combine_format) - self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) - self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) - self.generate_c = bool(generate_c) - - for name, fmt in ( - ("output_format", self.output_format), - ("combine_format", self.combine_format), - ): - required_multiple = ( - 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 - ) - if hidden_size % required_multiple != 0: - raise ValueError( - f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" - f" {name}={fmt.value}" - ) - - def __repr__(self) -> str: - return ( - f"{type(self).__name__}(" - f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " - f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " - f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " - f"output={self.output_format.value}, combine={self.combine_format.value})" - ) - - def _collective_device(self, device: torch.device) -> torch.device: - """Device the process group can run ``all_to_all_single`` on. - - Gloo only implements all-to-all for CPU tensors, so CUDA tensors are - staged through host memory; NCCL groups communicate in place. - """ - if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": - return torch.device("cpu") - return device - - def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: - if self.ep_size == 1: - return send_counts.clone() - comm_device = self._collective_device(send_counts.device) - staged = send_counts.to(comm_device) - recv_counts = torch.empty_like(staged) - dist.all_to_all_single(recv_counts, staged, group=self.ep_group) - return recv_counts.to(send_counts.device) - - def _all_to_all( - self, - send: torch.Tensor, - send_counts: Sequence[int], - recv_counts: Sequence[int], - ) -> torch.Tensor: - if self.ep_size == 1: - return send.clone() - comm_device = self._collective_device(send.device) - staged = send.contiguous().to(comm_device) - output_shape = (sum(recv_counts), *send.shape[1:]) - recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) - dist.all_to_all_single( - recv, - staged, - output_split_sizes=list(recv_counts), - input_split_sizes=list(send_counts), - group=self.ep_group, - ) - return recv.to(send.device) - - def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: - """Route valid ``topk_idx`` entries to destination ranks, stably by rank. - - Backward reuses this so gradient re-dispatch reproduces the exact - forward route order. - """ - - device = topk_idx.device - token_count = topk_idx.shape[0] - flat_expert = topk_idx.reshape(-1).to(torch.int64) - flat_weight = topk_weights.reshape(-1).float() - valid = flat_expert != -1 - invalid_negative = flat_expert < -1 - invalid_high = flat_expert >= self.num_experts - if bool((invalid_negative | invalid_high).any().item()): - bad = flat_expert[invalid_negative | invalid_high][0].item() - raise ValueError(f"topk_idx contains out-of-range expert id {bad}") - - flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) - flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) - expert = flat_expert[valid] - destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") - order = torch.argsort(destination, stable=True) - - send_counts_tensor = torch.bincount( - destination.index_select(0, order), minlength=self.ep_size - ).to(torch.int64) - recv_counts_tensor = self._exchange_counts(send_counts_tensor) - return _DispatchPlan( - send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), - send_weight=flat_weight[valid].index_select(0, order), - send_token_idx=flat_token[valid].index_select(0, order), - send_slot_idx=flat_slot[valid].index_select(0, order), - send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), - recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), - ) - - def _run_local_experts( - self, - tokens: torch.Tensor, - local_expert_idx: torch.Tensor, - route_weight: torch.Tensor, - fc1_weight: torch.Tensor, - fc2_weight: torch.Tensor, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - output = torch.empty( - (tokens.shape[0], self.hidden_size), - dtype=torch.float32, - device=tokens.device, - ) - fc1_c_rows = [] if self.generate_c else None - for expert in range(self.experts_per_rank): - positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() - if positions.numel() == 0: - continue - expert_tokens = tokens.index_select(0, positions) - gate_up = expert_tokens @ fc1_weight[expert] - if fc1_c_rows is not None: - # Raw pre-SwiGLU accumulator: before clamp, no router weight. - fc1_c_rows.append(gate_up.to(torch.bfloat16)) - gate, up = gate_up.split(self.intermediate_size, dim=-1) - if self.gate_up_clamp is not None: - gate = gate.clamp(max=self.gate_up_clamp) - up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) - intermediate = F.silu(gate) * up - weights = route_weight.index_select(0, positions).unsqueeze(-1) - if self.apply_topk_in_fc1: - intermediate = intermediate * weights - expert_output = intermediate @ fc2_weight[expert] - if not self.apply_topk_in_fc1: - expert_output = expert_output * weights - expert_output = _format_round_trip(expert_output, self.combine_format) - output.index_copy_(0, positions, expert_output) - fc1_c = None - if fc1_c_rows is not None: - fc1_c = ( - torch.cat(fc1_c_rows) - if fc1_c_rows - else torch.empty( - (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device - ) - ) - return output, fc1_c - - def __call__( - self, - activation: MoeTensor, - fc1_weight: MoeTensor, - fc2_weight: MoeTensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: - """Run dispatch, local experts, return routing, top-k reduce, and encode. - - Shapes: - activation: ``(T, H)`` - fc1_weight: ``(E_local, H, 2 * I)`` - fc2_weight: ``(E_local, I, H)`` - topk_idx/topk_weights: ``(T, K)`` - - Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` - when constructed with ``generate_c=True``. ``fc1_c`` is the BF16 - pre-SwiGLU FC1 accumulator of every route this rank's experts - processed, ``(local_routes, 2 * I)``, grouped by local expert and - ordered within each expert by (source rank, source token-major route - order); captured before the gate/up clamp, without the router weight. - ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns - ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies - the route behind ``fc1_c`` row ``i`` for the backward gradient - re-dispatch. - """ - - if topk_idx.ndim != 2: - raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") - token_count = topk_idx.shape[0] - route_shape = (token_count, self.top_k) - if tuple(topk_idx.shape) != route_shape: - raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") - if tuple(topk_weights.shape) != route_shape: - raise ValueError( - f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" - ) - if topk_idx.dtype not in (torch.int32, torch.int64): - raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") - if not topk_weights.is_floating_point(): - raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") - if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError( - f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" - ) - - device = _tensor_device(activation) - inputs = { - "fc1_weight": _tensor_device(fc1_weight), - "fc2_weight": _tensor_device(fc2_weight), - "topk_idx": topk_idx.device, - "topk_weights": topk_weights.device, - } - for name, input_device in inputs.items(): - if input_device != device: - raise ValueError(f"{name} must be on {device}, got {input_device}") - - activation_float = _decode_tensor( - activation, - name="activation", - expected_shape=(token_count, self.hidden_size), - quantized_axis=1, - ) - fc1_float = _decode_tensor( - fc1_weight, - name="fc1_weight", - expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), - quantized_axis=1, - ) - fc2_float = _decode_tensor( - fc2_weight, - name="fc2_weight", - expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), - quantized_axis=1, - ) - - plan = self._dispatch_plan(topk_idx, topk_weights) - send_token_idx = plan.send_token_idx - send_slot_idx = plan.send_slot_idx - send_counts, recv_counts = plan.send_counts, plan.recv_counts - send_tokens = activation_float.index_select(0, send_token_idx) - - recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) - recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) - recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) - - route_metadata = None - if self.generate_c: - recv_src_rank = torch.repeat_interleave( - torch.arange(self.ep_size, device=device), - torch.tensor(recv_counts, device=device), - ) - recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) - recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) - # Stable sort by local expert reproduces the fc1_c row order - # (grouped by expert; source order preserved within each group). - fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = ( - torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) - .index_select(0, fc1_c_order) - .to(torch.int32) - ) - # recv rows are ordered by source rank, then that source's token-major - # route order, so the per-expert position grouping below realizes the - # documented fc1_c ordering. - recv_output, fc1_c = self._run_local_experts( - recv_tokens, - recv_expert, - recv_weight, - fc1_float, - fc2_float, - ) - - returned = self._all_to_all(recv_output, recv_counts, send_counts) - combine_plane = torch.zeros( - (token_count * self.top_k, self.hidden_size), - dtype=torch.float32, - device=device, - ) - send_flat_slot = send_token_idx * self.top_k + send_slot_idx - combine_plane.index_copy_(0, send_flat_slot, returned) - reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) - - if self.output_format is MoeFormat.BF16: - output = reduced.to(torch.bfloat16) - else: - output = quantize_blockwise(reduced, self.output_format, axis=-1) - if self.generate_c: - return output, fc1_c, route_metadata - return output - - def backward( - self, - grad_output: torch.Tensor, - activation: MoeTensor, - fc1_weight: MoeTensor, - fc2_weight: MoeTensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - fc1_c: torch.Tensor, - route_metadata: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Backward pass consuming the ``generate_c=True`` stash. - - ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, - and the FC2 input are all rebuilt from it, so no post-SwiGLU forward - intermediate needs to be saved. ``route_metadata`` alone reconstructs - the mapping between re-dispatched rows and ``fc1_c`` rows and drives - the gradient return scatter. - - Quantization round-trips (input decode, ``combine_format``, - ``output_format``) are treated as straight-through identities; - ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. - - Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, - grad_topk_weights)`` in float32. - """ - - if not self.generate_c: - raise RuntimeError( - "backward requires the operator to be constructed with generate_c=True" - ) - token_count = topk_idx.shape[0] - if tuple(grad_output.shape) != (token_count, self.hidden_size): - raise ValueError( - f"grad_output shape must be {(token_count, self.hidden_size)}, got" - f" {tuple(grad_output.shape)}" - ) - if not grad_output.is_floating_point(): - raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") - - device = _tensor_device(activation) - two_i = 2 * self.intermediate_size - activation_float = _decode_tensor( - activation, - name="activation", - expected_shape=(token_count, self.hidden_size), - quantized_axis=1, - ) - fc1_float = _decode_tensor( - fc1_weight, - name="fc1_weight", - expected_shape=(self.experts_per_rank, self.hidden_size, two_i), - quantized_axis=1, - ) - fc2_float = _decode_tensor( - fc2_weight, - name="fc2_weight", - expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), - quantized_axis=1, - ) - if fc1_c.shape != (int(route_metadata.shape[0]), two_i): - raise ValueError( - f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" - f" {tuple(fc1_c.shape)}" - ) - - # Re-dispatch the FC1 inputs, router weights, and output gradients - # along the identical forward routes. - plan = self._dispatch_plan(topk_idx, topk_weights) - send_counts, recv_counts = plan.send_counts, plan.recv_counts - grad_output_float = grad_output.float() - recv_tokens = self._all_to_all( - activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts - ) - recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) - recv_grad = self._all_to_all( - grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts - ) - - # route_metadata rows are in fc1_c order; sorting them by - # (src_rank, src_token, src_slot) reproduces the receive order, giving - # the permutation between re-dispatched rows and fc1_c rows. - metadata = route_metadata.to(device=device, dtype=torch.int64) - local_routes = metadata.shape[0] - if local_routes > 0: - token_span = int(metadata[:, 2].max().item()) + 1 - recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] - perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j - else: - perm = torch.empty((0,), dtype=torch.int64, device=device) - x_rows = torch.empty_like(recv_tokens) - x_rows.index_copy_(0, perm, recv_tokens) - w_rows = torch.empty_like(recv_weight) - w_rows.index_copy_(0, perm, recv_weight) - dy_rows = torch.empty_like(recv_grad) - dy_rows.index_copy_(0, perm, recv_grad) - - c_rows = fc1_c.float() - expert_rows = metadata[:, 0] - d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) - d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) - grad_fc1 = torch.zeros_like(fc1_float) - grad_fc2 = torch.zeros_like(fc2_float) - for expert in range(self.experts_per_rank): - positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() - if positions.numel() == 0: - continue - c = c_rows.index_select(0, positions) - x = x_rows.index_select(0, positions) - w = w_rows.index_select(0, positions).unsqueeze(-1) - d_y = dy_rows.index_select(0, positions) - - gate, up = c.split(self.intermediate_size, dim=-1) - if self.gate_up_clamp is not None: - g = gate.clamp(max=self.gate_up_clamp) - u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) - else: - g, u = gate, up - sig = torch.sigmoid(g) - s = g * sig - h = s * u - - if self.apply_topk_in_fc1: - h_fc2 = h * w - d_y_pre = d_y - else: - h_fc2 = h - d_y_pre = d_y * w - grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre - d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) - if self.apply_topk_in_fc1: - d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) - else: - d_h = d_h_fc2 - d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) - - d_g = d_h * u * (sig * (1 + g * (1 - sig))) - d_u = d_h * s - if self.gate_up_clamp is not None: - d_gate = d_g * (gate <= self.gate_up_clamp) - d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) - else: - d_gate, d_up = d_g, d_u - d_c = torch.cat((d_gate, d_up), dim=-1) - grad_fc1[expert] = x.transpose(0, 1) @ d_c - d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) - - # Return the route gradients to their source ranks and scatter-add. - returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) - returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) - grad_activation = torch.zeros( - (token_count, self.hidden_size), dtype=torch.float32, device=device - ) - grad_activation.index_add_(0, plan.send_token_idx, returned_dx) - grad_topk_weights = torch.zeros( - (token_count * self.top_k,), dtype=torch.float32, device=device - ) - grad_topk_weights.index_copy_( - 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw - ) - return ( - grad_activation, - grad_fc1, - grad_fc2, - grad_topk_weights.view(token_count, self.top_k), - ) - - -__all__ = [ - "BlockScaledTensor", - "MoeEpReference", - "MoeFormat", - "MoeTensor", - "quantize_blockwise", -] diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 49fc6481ba..067f75d725 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,8 +11,6 @@ import torch import torch.distributed as dist -from moe_ep_reference import MoeEpReference -from transformer_engine.pytorch import ops as te_ops from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -123,30 +121,6 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) -def _make_moe_inputs(rank, ep_size, device="cuda"): - """Random activations and top-k routing representative of a router output.""" - generator = torch.Generator(device=device) - generator.manual_seed(2026 + rank) - num_experts = ep_size * NUM_LOCAL_EXPERTS - tokens = torch.randn( - TOKENS_PER_RANK, - HIDDEN_DIM, - generator=generator, - dtype=torch.float32, - device=device, - ).mul_(0.25) - router_logits = torch.randn( - TOKENS_PER_RANK, - num_experts, - generator=generator, - dtype=torch.float32, - device=device, - ) - topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) - topk_weights = torch.softmax(topk_logits, dim=-1) - return topk_idx, tokens.to(torch.bfloat16), topk_weights - - class _Cfg: rank: int world_size: int @@ -617,238 +591,6 @@ def zero_grads(): rtol=5e-2, ) - @_eager_test_include - def test_fusible_dispatch_combine_moe(self): - """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" - if not EAGER: - self.skipTest("variable grouped-linear splits require eager EP output sizing") - - topk_idx, tokens, topk_weights = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) - num_experts = NUM_LOCAL_EXPERTS - intermediate_dim = HIDDEN_DIM - - dispatch_buffer = self._make_buffer() - dispatch = te_ops.Dispatch(dispatch_buffer) - fc1 = te_ops.GroupedLinear( - num_experts, - HIDDEN_DIM, - 2 * intermediate_dim, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) - activation = te_ops.ScaledSwiGLU() - fc2 = te_ops.GroupedLinear( - num_experts, - intermediate_dim, - HIDDEN_DIM, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) - combine = te_ops.Combine(dispatch_buffer, num_local_tokens=TOKENS_PER_RANK) - - dispatch.set_extra_output_channel(0, "m_splits") - dispatch.set_extra_output_channel(1, "token_probs") - fc1.set_extra_input_channel(0, "m_splits") - activation.set_extra_input_channel(0, "token_probs") - fc2.set_extra_input_channel(0, "m_splits") - model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) - - generator = torch.Generator(device=self.cfg.device) - generator.manual_seed(1234 + self.cfg.rank) - with torch.no_grad(): - for expert_idx in range(num_experts): - getattr(fc1, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) - getattr(fc2, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) - - ref_fc1_weights = torch.stack( - [getattr(fc1, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] - ) - ref_fc2_weights = torch.stack( - [getattr(fc2, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] - ) - reference = MoeEpReference( - num_experts=self.cfg.num_experts, - hidden_size=HIDDEN_DIM, - intermediate_size=intermediate_dim, - top_k=TOP_K, - ep_group=self.ep_group, - max_tokens_per_rank=TOKENS_PER_RANK, - apply_topk_in_fc1=True, - generate_c=True, - ) - ref_output, fc1_c, route_metadata = reference( - tokens, - ref_fc1_weights, - ref_fc2_weights, - topk_idx, - topk_weights, - ) - - test_tokens = tokens.detach().clone().requires_grad_(True) - test_topk_weights = topk_weights.detach().clone().requires_grad_(True) - test_output = model(test_tokens, topk_idx, test_topk_weights) - - grad_output = torch.linspace( - -0.2, - 0.2, - TOKENS_PER_RANK * HIDDEN_DIM, - device=self.cfg.device, - dtype=torch.float32, - ).reshape(TOKENS_PER_RANK, HIDDEN_DIM) - grad_output = grad_output.to(torch.bfloat16) - ref_grads = reference.backward( - grad_output, - tokens, - ref_fc1_weights, - ref_fc2_weights, - topk_idx, - topk_weights, - fc1_c, - route_metadata, - ) - test_output.backward(grad_output) - torch.cuda.synchronize() - - torch.testing.assert_close(test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2) - torch.testing.assert_close(test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2) - torch.testing.assert_close( - test_topk_weights.grad, - ref_grads[3], - atol=5e-2, - rtol=5e-2, - ) - for expert_idx in range(num_experts): - torch.testing.assert_close( - getattr(fc1, f"weight{expert_idx}").grad.float(), - ref_grads[1][expert_idx].transpose(0, 1), - atol=5e-2, - rtol=5e-2, - ) - torch.testing.assert_close( - getattr(fc2, f"weight{expert_idx}").grad.float(), - ref_grads[2][expert_idx].transpose(0, 1), - atol=5e-2, - rtol=5e-2, - ) - - @_eager_test_include - def test_fusible_dispatch_combine_moe(self): - """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" - if not EAGER: - self.skipTest("variable grouped-linear splits require eager EP output sizing") - - topk_idx, tokens, topk_weights = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) - num_experts = NUM_LOCAL_EXPERTS - intermediate_dim = HIDDEN_DIM - - dispatch_buffer = self._make_buffer() - dispatch = te_ops.Dispatch(dispatch_buffer) - fc1 = te_ops.GroupedLinear( - num_experts, - HIDDEN_DIM, - 2 * intermediate_dim, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) - activation = te_ops.ScaledSwiGLU() - fc2 = te_ops.GroupedLinear( - num_experts, - intermediate_dim, - HIDDEN_DIM, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) - combine = te_ops.Combine(dispatch_buffer, num_local_tokens=TOKENS_PER_RANK) - - dispatch.set_extra_output_channel(0, "m_splits") - dispatch.set_extra_output_channel(1, "token_probs") - fc1.set_extra_input_channel(0, "m_splits") - activation.set_extra_input_channel(0, "token_probs") - fc2.set_extra_input_channel(0, "m_splits") - model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) - - generator = torch.Generator(device=self.cfg.device) - generator.manual_seed(1234 + self.cfg.rank) - with torch.no_grad(): - for expert_idx in range(num_experts): - getattr(fc1, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) - getattr(fc2, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) - - ref_fc1_weights = torch.stack( - [getattr(fc1, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] - ) - ref_fc2_weights = torch.stack( - [getattr(fc2, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] - ) - reference = MoeEpReference( - num_experts=self.cfg.num_experts, - hidden_size=HIDDEN_DIM, - intermediate_size=intermediate_dim, - top_k=TOP_K, - ep_group=self.ep_group, - max_tokens_per_rank=TOKENS_PER_RANK, - apply_topk_in_fc1=True, - generate_c=True, - ) - ref_output, fc1_c, route_metadata = reference( - tokens, - ref_fc1_weights, - ref_fc2_weights, - topk_idx, - topk_weights, - ) - - test_tokens = tokens.detach().clone().requires_grad_(True) - test_topk_weights = topk_weights.detach().clone().requires_grad_(True) - test_output = model(test_tokens, topk_idx, test_topk_weights) - - grad_output = torch.linspace( - -0.2, - 0.2, - TOKENS_PER_RANK * HIDDEN_DIM, - device=self.cfg.device, - dtype=torch.float32, - ).reshape(TOKENS_PER_RANK, HIDDEN_DIM) - grad_output = grad_output.to(torch.bfloat16) - ref_grads = reference.backward( - grad_output, - tokens, - ref_fc1_weights, - ref_fc2_weights, - topk_idx, - topk_weights, - fc1_c, - route_metadata, - ) - test_output.backward(grad_output) - torch.cuda.synchronize() - - torch.testing.assert_close(test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2) - torch.testing.assert_close(test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2) - torch.testing.assert_close( - test_topk_weights.grad, - ref_grads[3], - atol=5e-2, - rtol=5e-2, - ) - for expert_idx in range(num_experts): - torch.testing.assert_close( - getattr(fc1, f"weight{expert_idx}").grad.float(), - ref_grads[1][expert_idx].transpose(0, 1), - atol=5e-2, - rtol=5e-2, - ) - torch.testing.assert_close( - getattr(fc2, f"weight{expert_idx}").grad.float(), - ref_grads[2][expert_idx].transpose(0, 1), - atol=5e-2, - rtol=5e-2, - ) - @_zero_copy_test_include @_eager_test_include def test_combine_autograd(self): diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 4ea116793f..6def36ffc7 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -23,8 +23,6 @@ from .basic_linear import BasicLinear from .bias import Bias from .constant_scale import ConstantScale -from .combine import Combine -from .dispatch import Dispatch from .dropout import Dropout from .grouped_linear import GroupedLinear from .identity import Identity diff --git a/transformer_engine/pytorch/ops/combine.py b/transformer_engine/pytorch/ops/combine.py deleted file mode 100644 index 4d4740c517..0000000000 --- a/transformer_engine/pytorch/ops/combine.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Fusible NCCL expert-parallel combine operation.""" - -from __future__ import annotations - -from typing import Any, Optional - -import torch - -from ...ep import EpBuffer, _alloc_io, is_symm_backed -from ...tensor import Quantizer -from ..op import BasicOperation, OperationContext - - -def _validate_grad_buffer( - tensor: Optional[torch.Tensor], - *, - shape: tuple[int, ...], - dtype: torch.dtype, - device: torch.device, -) -> Optional[torch.Tensor]: - if tensor is None: - return None - if tuple(tensor.shape) != shape: - raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") - if tensor.dtype is not dtype: - raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") - if tensor.device != device: - raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") - if not tensor.is_contiguous(): - raise ValueError("grad_out must be contiguous.") - if tensor.requires_grad: - raise ValueError("grad_out must not require gradients.") - return tensor - - -class Combine(BasicOperation): - """Combine pre-weighted local expert outputs with NCCL EP. - - The operation uses routing state produced by a :class:`Dispatch` with the - same :class:`EpBuffer`. - """ - - def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: - super().__init__() - self.buffer = buffer - self.num_local_tokens = ( - buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) - ) - if self.num_local_tokens < 0: - raise ValueError("num_local_tokens must be non-negative.") - - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - *, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - **kwargs: Any, - ) -> torch.Tensor: - del prev_op_grad_output_quantizer, next_op_input_quantizer - if input_.dtype is not torch.bfloat16: - raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") - if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: - raise ValueError( - f"Combine input must have shape (R, {self.buffer.hidden_dim}), " - f"got {tuple(input_.shape)}." - ) - - expert_out = input_ - if self.buffer.zero_copy: - expert_out = _alloc_io( - tuple(input_.shape), - input_.dtype, - input_.device, - True, - ) - expert_out.copy_(input_) - - result = torch.empty( - self.num_local_tokens, - self.buffer.hidden_dim, - dtype=input_.dtype, - device=input_.device, - ) - torch.ops.transformer_engine_ep.combine( - self.buffer.handle_mem, - expert_out, - result, - ) - - if ctx.requires_grad: - grad_out = kwargs.get("grad_out") - if self.buffer.eager and grad_out is not None: - raise ValueError( - "eager mode sizes combine gradients per step and cannot use " - "a caller-supplied grad_out" - ) - grad_out = _validate_grad_buffer( - grad_out, - shape=tuple(input_.shape), - dtype=input_.dtype, - device=input_.device, - ) - if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): - raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") - ctx.grad_out = grad_out - ctx.input_shape = tuple(input_.shape) - ctx.input_dtype = input_.dtype - ctx.save_for_backward(self.buffer.handle_mem) - - return result - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - (handle_mem,) = ctx.saved_tensors - grad_output = grad_output.contiguous() - grad_input = ctx.grad_out - if grad_input is None: - grad_input = _alloc_io( - ctx.input_shape, - ctx.input_dtype, - grad_output.device, - self.buffer.zero_copy, - ) - torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, - grad_output, - grad_input, - ) - return grad_input, () diff --git a/transformer_engine/pytorch/ops/dispatch.py b/transformer_engine/pytorch/ops/dispatch.py deleted file mode 100644 index 9762ed76f7..0000000000 --- a/transformer_engine/pytorch/ops/dispatch.py +++ /dev/null @@ -1,207 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Fusible NCCL expert-parallel dispatch operation.""" - -from __future__ import annotations - -from typing import Any, Iterable, Optional - -import torch - -from ...ep import EpBuffer, _alloc_io, ep_prepare -from ...tensor import Quantizer -from ..op import BasicOperation, OperationContext - - -def _validate_output_buffer( - name: str, - tensor: Optional[torch.Tensor], - *, - shape: tuple[int, ...], - dtype: torch.dtype, - device: torch.device, -) -> Optional[torch.Tensor]: - if tensor is None: - return None - if tuple(tensor.shape) != shape: - raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") - if tensor.dtype is not dtype: - raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") - if tensor.device != device: - raise ValueError(f"{name} must be on {device}, got {tensor.device}.") - if not tensor.is_contiguous(): - raise ValueError(f"{name} must be contiguous.") - if tensor.requires_grad: - raise ValueError(f"{name} must not require gradients.") - return tensor - - -class Dispatch(BasicOperation): - """Dispatch BF16 tokens to local experts with NCCL EP. - - The extra inputs are routing indices and FP32 routing weights. The extra - outputs are local tokens-per-expert and received routing weights. - """ - - num_extra_inputs: int = 2 - num_extra_outputs: int = 2 - - def __init__(self, buffer: EpBuffer) -> None: - super().__init__() - self.buffer = buffer - - def op_forward(self, *args: Any, **kwargs: Any) -> None: - raise RuntimeError("Dispatch uses fuser_forward") - - def op_backward(self, *args: Any, **kwargs: Any) -> None: - raise RuntimeError("Dispatch uses fuser_backward") - - def fuser_forward( - self, - basic_op_ctxs: list[OperationContext], - input_: torch.Tensor, - *, - basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - del prev_op_grad_output_quantizer, next_op_input_quantizer - topk_idx, topk_weights = basic_op_extra_inputs[0] - kwargs = basic_op_kwargs[0] - - if input_.dtype is not torch.bfloat16: - raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") - if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: - raise ValueError( - f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " - f"got {tuple(input_.shape)}." - ) - if topk_idx.dtype not in (torch.int32, torch.int64): - raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}.") - expected_route_shape = (input_.shape[0], self.buffer.top_k) - if tuple(topk_idx.shape) != expected_route_shape: - raise ValueError( - f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." - ) - if tuple(topk_weights.shape) != expected_route_shape: - raise ValueError( - f"topk_weights shape must be {expected_route_shape}, " - f"got {tuple(topk_weights.shape)}." - ) - if topk_weights.dtype is not torch.float32: - raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") - for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): - if tensor.device != input_.device: - raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") - - recv_tokens = kwargs.get("recv_tokens") - recv_topk_weights = kwargs.get("recv_topk_weights") - if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): - raise ValueError( - "eager mode sizes dispatch outputs per step and cannot use " - "caller-supplied receive buffers" - ) - - tokens_per_expert = ep_prepare(self.buffer, topk_idx) - rows = ( - self.buffer._host_total_recv_tokens - if self.buffer.eager - else self.buffer.recv_capacity_per_rank - ) - if rows is None: - raise RuntimeError("NCCL EP dispatch receive size is unavailable.") - rows = int(rows) - recv_shape = (rows, self.buffer.hidden_dim) - recv_tokens = _validate_output_buffer( - "recv_tokens", - recv_tokens, - shape=recv_shape, - dtype=self.buffer.payload_dtype, - device=self.buffer.device, - ) - recv_topk_weights = _validate_output_buffer( - "recv_topk_weights", - recv_topk_weights, - shape=(rows,), - dtype=torch.float32, - device=self.buffer.device, - ) - if recv_tokens is None: - recv_tokens = _alloc_io( - recv_shape, - self.buffer.payload_dtype, - self.buffer.device, - self.buffer.zero_copy, - ) - if recv_topk_weights is None: - recv_topk_weights = _alloc_io( - (rows,), - torch.float32, - self.buffer.device, - self.buffer.zero_copy, - ) - - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - input_, - topk_weights, - recv_tokens, - recv_topk_weights, - ) - - ctx = basic_op_ctxs[0] - if ctx.requires_grad: - ctx.input_shape = tuple(input_.shape) - ctx.input_dtype = input_.dtype - ctx.topk_weights_shape = tuple(topk_weights.shape) - ctx.save_for_backward(self.buffer.handle_mem) - - return recv_tokens, [(tokens_per_expert, recv_topk_weights)] - - def fuser_backward( - self, - basic_op_ctxs: list[OperationContext], - grad_output: torch.Tensor, - *, - basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], - ) -> tuple[ - torch.Tensor, - Iterable[Iterable[Optional[torch.Tensor]]], - Iterable[Iterable[Optional[torch.Tensor]]], - ]: - ctx = basic_op_ctxs[0] - (handle_mem,) = ctx.saved_tensors - grad_output = grad_output.contiguous() - - grad_recv_weights = basic_op_grad_extra_outputs[0][1] - if grad_recv_weights is None: - grad_recv_weights = torch.zeros( - grad_output.shape[0], - dtype=torch.float32, - device=grad_output.device, - ) - else: - grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous() - - grad_input = torch.empty( - ctx.input_shape, - dtype=ctx.input_dtype, - device=grad_output.device, - ) - grad_topk_weights = torch.empty( - ctx.topk_weights_shape, - dtype=torch.float32, - device=grad_output.device, - ) - torch.ops.transformer_engine_ep.dispatch_bwd( - handle_mem, - grad_output, - grad_recv_weights, - grad_input, - grad_topk_weights, - ) - return grad_input, [()], [(None, grad_topk_weights)] From d7d6380bf1027bf66e5d0e3eea6e7fd762bf618e Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 5 Aug 2026 23:57:08 +0000 Subject: [PATCH 09/83] cleanup with residual tests Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 112 +++++++-- tests/pytorch/test_fusible_ops.py | 377 ++++++---------------------- 2 files changed, 167 insertions(+), 322 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..7b003091e0 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -113,43 +113,115 @@ quantized compute. Branching operations ^^^^^^^^^^^^^^^^^^^^ -The operation fuser supports very limited branching behavior. While -the operations must be in sequential order, some operations can accept -extra inputs or produce extra outputs. For example, ``AddExtraInput`` -will add an extra input tensor to the intermediate tensor and -``MakeExtraOutput`` will return the intermediate tensor as an extra -output. When calling a ``Sequential`` that contains any of these -branching operations, the extra inputs should be passed in as -arguments and the extra outputs will be returned. +The operation fuser supports limited branching behavior. While the +operations must be in sequential order, basic operations may declare +extra tensor inputs and outputs. By default, an extra tensor slot has +no channel assigned and is part of the public ``Sequential`` interface: +the caller provides extra inputs as arguments, and extra outputs are +returned after the main output. Assigning the same channel name to an +output slot and a later input slot connects them internally instead. .. code-block:: python import torch import transformer_engine.pytorch as te - # Construct MLP with residual connection - fc1 = te.ops.Sequential( + # Keep a residual connection inside one Sequential. + 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), - te.ops.MakeExtraOutput(), # Output residual + make_residual, te.ops.Linear(4096, 28672), te.ops.SwiGLU(), - ) - fc2 = te.ops.Sequential( te.ops.Linear(14336, 4096), - te.ops.AddExtraInput(), # Add residual + add_residual, ) - # Forward pass x = torch.randn(16384, 4096, device="cuda") - y, residual = fc1(x) - y = fc2(y, residual) + y = block(x) .. figure:: ./residual_layernorm_mlp.png :align: center - Operations for an MLP block with a residual connection. Note that - the block has been split into two sections, each with one branching - operation. + Operations for an MLP block with a residual connection. + +Extra tensor channels +""""""""""""""""""""" + +An extra output and one or more later extra inputs can be assigned the +same channel name. This routes the tensor inside the +``OperationFuser`` and removes the bound slots from the public +``Sequential`` interface. In the residual example above, the caller +therefore receives only ``y`` and does not need to pass the residual +back into the block. + +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. + y = moe(x, router_probs) + +The following conditions apply to extra tensor channels: + +- A producer must appear before all of its consumers. Backward edges + and cycles are not supported. +- A channel has exactly one producer, but its output may fan out to + multiple consumers. +- Every named output channel must have at least one consumer, and the + channel names on the producer and consumers must match. +- 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 have no channel assigned and + receives extra outputs that have no channel assigned. Slots assigned + to channels are internal and do not appear in the ``Sequential`` + arguments or return value. + +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 --------------- diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index dba1cab263..f241c24e32 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -437,6 +437,81 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x3, x3_orig + x2 + b) torch.testing.assert_close(x4, x4_orig + x3) +class TestExtraTensorChannels: + """Error handling and grad coverage for named extra-tensor channels.""" + + def test_internal_residual_connection(self, size: int = 16) -> None: + """A channel can keep a residual connection inside a Sequential.""" + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + add_residual = te_ops.AddExtraInput() + residual.set_extra_output_channel(0, "residual") + add_residual.set_extra_input_channel(0, "residual") + + model = te_ops.Sequential(residual, body, add_residual) + x = torch.rand((size,), requires_grad=True) + y = model(x) + + torch.testing.assert_close(y, 2 * x + body.bias) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + + def test_fused_internal_residual_connection(self, size: int = 16) -> None: + """A fused op can implement both ends of an internal channel.""" + + class FusedResidual(te_ops.FusedOperation): + """Fuse MakeExtraOutput, Bias, and AddExtraInput in forward.""" + + _enabled = True + + def __init__(self, residual, body, add_residual) -> None: + super().__init__((residual, body, add_residual)) + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + del basic_op_ctxs + # The consumer slot is internal to this fusion, so the + # OperationFuser deliberately leaves it unset. + assert basic_op_extra_inputs[2][0] is None + return 2 * input_, [(input_,), (), ()] + + def fuse_residual(ops, **unused): + if not FusedResidual._enabled: + return ops + if ( + len(ops) == 3 + and isinstance(ops[0], te_ops.MakeExtraOutput) + and isinstance(ops[1], te_ops.Identity) + and isinstance(ops[2], te_ops.AddExtraInput) + ): + FusedResidual._enabled = False + return [FusedResidual(*ops)] + return ops + + residual = te_ops.MakeExtraOutput() + body = te_ops.Identity() + add_residual = te_ops.AddExtraInput() + residual.set_extra_output_channel(0, "residual") + add_residual.set_extra_input_channel(0, "residual") + model = te_ops.Sequential(residual, body, add_residual) + + te_ops.register_forward_fusion(fuse_residual, prepend=True) + x = torch.rand((size,), requires_grad=True) + y = model(x) + + forward_ops = model._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], FusedResidual) + torch.testing.assert_close(y, 2 * x) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: """An internal extra output can feed multiple later consumers.""" producer = te_ops.MakeExtraOutput() @@ -479,308 +554,6 @@ def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None torch.testing.assert_close(x.grad, torch.full_like(x, 2)) torch.testing.assert_close(extra.grad, torch.ones_like(extra)) - def test_moe_style_dispatch_combine_extra_channels( - self, - *, - group_size: int = 4, - hidden_size: int = 32, - dtype: torch.dtype = torch.float32, - device: torch.device = "cuda", - ) -> None: - """Wire Dispatch-style extras through GroupedLinear / ScaledActivation / Combine. - - Stand-in for ``te.Sequential(Dispatch, GroupedLinear, ScaledActivation, - GroupedLinear, Combine)`` once real Dispatch/Combine ops land. ``m_splits`` - and ``probs`` are produced once and fan out to later consumers via named - channels so the public call is ``model(x, m_splits, probs)``. - """ - - class FakeDispatch(te_ops.BasicOperation): - """Stand-in MoE dispatch: passthrough hidden states, emit routing extras. - - Real Dispatch would permute tokens. Extra inputs are the externally - provided ``m_splits`` and ``probs``; matching extra outputs are bound - to internal channels for later consumers. A third extra output is a - stub ``routing_map`` for Combine. - """ - - num_extra_inputs = 2 - num_extra_outputs = 3 - - def op_forward(self, *args, **kwargs): - raise RuntimeError("FakeDispatch uses fuser_forward") - - def op_backward(self, *args, **kwargs): - raise RuntimeError("FakeDispatch uses fuser_backward") - - def fuser_forward( - self, - basic_op_ctxs, - input_, - *, - basic_op_extra_inputs, - **unused, - ): - m_splits, probs = basic_op_extra_inputs[0] - # Stub row-id map: real Dispatch would emit permute indices. - routing_map = torch.arange(input_.size(0), device=input_.device, dtype=torch.int64) - return input_, [(m_splits, probs, routing_map)] - - def fuser_backward( - self, - basic_op_ctxs, - grad_output, - *, - basic_op_grad_extra_outputs, - ): - # Channel grads from GroupedLinear / ScaledActivation land here. - return ( - grad_output, - [()], - [tuple(basic_op_grad_extra_outputs[0][:2])], - ) - - class FakeCombine(te_ops.BasicOperation): - """Stand-in MoE combine: consumes Dispatch ``routing_map``, identity path. - - Real Combine would unpermute with the routing map. - """ - - num_extra_inputs = 1 - - def op_forward(self, *args, **kwargs): - raise RuntimeError("FakeCombine uses fuser_forward") - - def op_backward(self, *args, **kwargs): - raise RuntimeError("FakeCombine uses fuser_backward") - - def fuser_forward( - self, - basic_op_ctxs, - input_, - *, - basic_op_extra_inputs, - **unused, - ): - routing_map = basic_op_extra_inputs[0][0] - if routing_map is None: - raise RuntimeError("FakeCombine expected routing_map channel") - if int(routing_map.numel()) != int(input_.size(0)): - raise RuntimeError("FakeCombine routing_map length does not match tokens") - return input_, [()] - - def fuser_backward( - self, - basic_op_ctxs, - grad_output, - *, - basic_op_grad_extra_outputs, - ): - del basic_op_grad_extra_outputs - return grad_output, [()], [(None,)] - - split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[:group_size] - num_tokens = int(split_sizes.sum()) - in_shape = (num_tokens, hidden_size) - - x_ref, x_test = make_reference_and_test_tensors( - in_shape, - min=-0.25, - max=0.25, - test_dtype=dtype, - test_device=device, - ) - probs_ref, probs_test = make_reference_and_test_tensors( - (num_tokens,), - min=-0.25, - max=0.25, - test_dtype=dtype, - test_device=device, - ) - dy_ref, dy_test = make_reference_and_test_tensors( - in_shape, - min=-0.25, - max=0.25, - test_dtype=dtype, - test_device=device, - requires_grad=False, - ) - - # Reference: GroupedLinear + ScaledSReLU + GroupedLinear (no dispatch permute). - # Run the PyTorch reference in the same dtype/device as TE. Cross-device - # float32 GEMMs (CPU vs CUDA) differ enough that probs grads — a - # reduction over hidden — can miss dtype_tols even when channel wiring - # is correct. - fc1_w_refs, fc1_w_tests = [], [] - fc2_w_refs, fc2_w_tests = [], [] - for _ in range(group_size): - w1_ref, w1_test = make_reference_and_test_tensors( - (hidden_size, hidden_size), - min=-0.25, - max=0.25, - test_dtype=dtype, - test_device=device, - ) - w2_ref, w2_test = make_reference_and_test_tensors( - (hidden_size, hidden_size), - min=-0.25, - max=0.25, - test_dtype=dtype, - test_device=device, - ) - fc1_w_refs.append(w1_test.detach().clone()) - fc1_w_tests.append(w1_test) - fc2_w_refs.append(w2_test.detach().clone()) - fc2_w_tests.append(w2_test) - x_ref = x_test.detach().clone().requires_grad_(True) - probs_ref = probs_test.detach().clone().requires_grad_(True) - dy_ref = dy_test.detach().clone() - xs = torch.split(x_ref, split_sizes.tolist()) - probs = torch.split(probs_ref, split_sizes.tolist()) - ys = [] - for group_idx in range(group_size): - fc1_out = torch.nn.functional.linear(xs[group_idx], fc1_w_refs[group_idx]) - act_out = torch.nn.functional.relu(fc1_out).square() - fc2_in = act_out * probs[group_idx].unsqueeze(-1) - ys.append(torch.nn.functional.linear(fc2_in, fc2_w_refs[group_idx])) - y_ref = torch.cat(ys) - y_ref.backward(dy_ref) - - dispatch = FakeDispatch() - fc1 = te_ops.GroupedLinear( - group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype - ) - activation = te_ops.ScaledSReLU() - fc2 = te_ops.GroupedLinear( - group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype - ) - combine = FakeCombine() - - # Bind channels: Dispatch fans out, consumers bind by name. - 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") - - model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) - with torch.no_grad(): - for group_idx in range(group_size): - getattr(fc1, f"weight{group_idx}").copy_(fc1_w_tests[group_idx]) - getattr(fc2, f"weight{group_idx}").copy_(fc2_w_tests[group_idx]) - del fc1_w_tests, fc2_w_tests - - # Only Dispatch's extras remain public: model(x, m_splits, probs). - y_test = model(x_test, split_sizes, probs_test) - y_test.backward(dy_test) - - tols = dtype_tols(dtype) - assert_close(y_test, y_ref, **tols) - assert_close_grads(x_test, x_ref, **tols) - assert_close_grads(probs_test, probs_ref, **tols) - - def test_fused_op_with_internal_producer_consumer(self, size: int = 16) -> None: - """A fused op may contain both a channel producer and its consumer. - - This is the MegaMoE path: Dispatch + MLP + Combine collapse into one - fused kernel that wires channels internally instead of via the fuser. - """ - - class FakeDispatch(te_ops.BasicOperation): - num_extra_inputs = 1 - num_extra_outputs = 1 - - def op_forward(self, *args, **kwargs): - raise RuntimeError("FakeDispatch uses fuser_forward") - - def op_backward(self, *args, **kwargs): - raise RuntimeError("FakeDispatch uses fuser_backward") - - def fuser_forward( - self, - basic_op_ctxs, - input_, - *, - basic_op_extra_inputs, - **unused, - ): - (route,) = basic_op_extra_inputs[0] - return input_, [(route,)] - - def fuser_backward( - self, - basic_op_ctxs, - grad_output, - *, - basic_op_grad_extra_outputs, - ): - return ( - grad_output, - [()], - [tuple(basic_op_grad_extra_outputs[0])], - ) - - class MegaMoELike(te_ops.FusedOperation): - """Fused stub that owns both producer and consumer of ``route``.""" - - _enabled = True - - def __init__(self, dispatch, consumer) -> None: - super().__init__((dispatch, consumer)) - - def fuser_forward( - self, - basic_op_ctxs, - input_, - *, - basic_op_extra_inputs, - **unused, - ): - # Consumer slot is intentionally unset: producer is in this fusion. - route = basic_op_extra_inputs[0][0] - assert basic_op_extra_inputs[1][0] is None - out = input_ + route - return out, [(route,), ()] - - def fuse_mega_moe_like(ops, **unused): - if not MegaMoELike._enabled: - return ops - MegaMoELike._enabled = False - out = [] - window, ops = ops[:2], ops[2:] - while len(window) == 2: - if isinstance(window[0], FakeDispatch) and isinstance( - window[1], te_ops.AddExtraInput - ): - window = [MegaMoELike(*window)] - else: - out.append(window[0]) - window = window[1:] - window, ops = window + ops[:1], ops[1:] - out.extend(window + ops) - return out - - dispatch = FakeDispatch() - consumer = te_ops.AddExtraInput() - dispatch.set_extra_output_channel(0, "route") - consumer.set_extra_input_channel(0, "route") - model = te_ops.Sequential(dispatch, consumer) - - te_ops.register_forward_fusion(fuse_mega_moe_like) - x = torch.rand((size,), requires_grad=True) - route = torch.rand((size,), requires_grad=True) - y = model(x, route) - torch.testing.assert_close(y, x + route) - y.sum().backward() - torch.testing.assert_close(x.grad, torch.ones_like(x)) - torch.testing.assert_close(route.grad, torch.ones_like(route)) - - -class TestExtraTensorChannels: - """Error handling and grad coverage for named extra-tensor channels.""" - def test_consumer_channel_without_producer(self) -> None: """Extra input bound to a channel that no earlier op produces.""" consumer = te_ops.AddExtraInput() From 74f563a2118ba7c3a67da38785816c5b20a2934b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:58:14 +0000 Subject: [PATCH 10/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index f241c24e32..1174ff0d11 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -437,6 +437,7 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x3, x3_orig + x2 + b) torch.testing.assert_close(x4, x4_orig + x3) + class TestExtraTensorChannels: """Error handling and grad coverage for named extra-tensor channels.""" From 87e2b36774679a3b0b29b0ce5e138fec0ff734d4 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Thu, 6 Aug 2026 01:13:12 +0000 Subject: [PATCH 11/83] address review comment Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 13 +++++++++++++ transformer_engine/pytorch/ops/fuser.py | 5 +++++ transformer_engine/pytorch/ops/op.py | 26 +++++++++++++++++++++++-- 3 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 1174ff0d11..1616d830ff 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -597,6 +597,19 @@ def test_set_extra_channel_rejects_invalid_name(self) -> None: with pytest.raises(ValueError, match="non-empty string"): producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + def test_set_extra_channel_rejects_mutation_after_fuser_construction(self) -> None: + """Channel routing is immutable after it has been captured by a fuser.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + fuser = OperationFuser([producer, consumer]) + assert fuser.num_extra_inputs == 0 + with pytest.raises(RuntimeError, match="cannot be changed"): + producer.set_extra_output_channel(0, None) + with pytest.raises(RuntimeError, match="cannot be changed"): + consumer.set_extra_input_channel(0, None) + def test_duplicate_extra_output_channel_names(self) -> None: """Two extra outputs may not publish the same channel name.""" producer1 = te_ops.MakeExtraOutput() diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 63d3f4a6be..c4d7e1ec10 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -506,6 +506,11 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) + # Channel routing is structural fuser state. Prevent the basic ops from + # changing their bindings after this fuser captures them. + for op in self._basic_ops: + op._lock_extra_channels() + @staticmethod def _apply_fusions( ops: Iterable[FusibleOperation], diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 3159ae0f0d..d9f9624ab0 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -191,6 +191,7 @@ def __init__(self) -> None: # Unbound slots remain public inputs/outputs, preserving the original API. self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs + self._extra_channels_locked = False # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None @@ -201,7 +202,8 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp A bound slot receives the matching extra output from an earlier operation in the same fuser instead of consuming a public extra input. - Passing ``None`` removes the binding. + Passing ``None`` removes the binding. Bindings cannot be changed after + the operation has been attached to an ``OperationFuser``. """ if not 0 <= index < self.num_extra_inputs: raise IndexError( @@ -210,6 +212,9 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp ) if channel is not None and (not isinstance(channel, str) or not channel): raise ValueError("Extra input channel must be a non-empty string or None") + if self._extra_input_channels[index] == channel: + return self + self._assert_extra_channels_mutable() self._extra_input_channels[index] = channel return self @@ -217,7 +222,9 @@ def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicO """Bind an extra output slot to an internal fuser channel. A bound slot can feed one or more later operations and is not returned - as a public extra output. Passing ``None`` removes the binding. + as a public extra output. Passing ``None`` removes the binding. Bindings + cannot be changed after the operation has been attached to an + ``OperationFuser``. """ if not 0 <= index < self.num_extra_outputs: raise IndexError( @@ -226,9 +233,24 @@ def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicO ) if channel is not None and (not isinstance(channel, str) or not channel): raise ValueError("Extra output channel must be a non-empty string or None") + if self._extra_output_channels[index] == channel: + return self + self._assert_extra_channels_mutable() self._extra_output_channels[index] = channel return self + def _assert_extra_channels_mutable(self) -> None: + """Check that channel routing has not been captured by a fuser.""" + if self._extra_channels_locked: + raise RuntimeError( + "Extra tensor channels cannot be changed after an operation has been " + "attached to an OperationFuser" + ) + + def _lock_extra_channels(self) -> None: + """Prevent changes after a fuser has captured the channel routing.""" + self._extra_channels_locked = True + @property def is_fused_op(self) -> bool: return False From 80601dc96f918863d21c14d46f35c25c36857739 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 7 Aug 2026 22:53:15 +0000 Subject: [PATCH 12/83] update to cleaner documentation Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 94 ++++++++++++++++++++++------- 1 file changed, 72 insertions(+), 22 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index 7b003091e0..ffec6baa4e 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -114,19 +114,60 @@ Branching operations ^^^^^^^^^^^^^^^^^^^^ The operation fuser supports limited branching behavior. While the -operations must be in sequential order, basic operations may declare -extra tensor inputs and outputs. By default, an extra tensor slot has -no channel assigned and is part of the public ``Sequential`` interface: -the caller provides extra inputs as arguments, and extra outputs are -returned after the main output. Assigning the same channel name to an -output slot and a later input slot connects them internally instead. +operations must be in sequential order, some operations can accept +extra inputs or produce extra outputs. For example, ``AddExtraInput`` +adds an extra input tensor to the intermediate tensor, and +``MakeExtraOutput`` returns the intermediate tensor as an extra output. +When calling a ``Sequential`` that contains any of these branching +operations, the extra inputs should be passed as arguments and the +extra outputs will be returned after the main output. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + # Construct an MLP with a residual connection. + fc1 = te.ops.Sequential( + te.ops.LayerNorm(4096), + te.ops.MakeExtraOutput(), # Output the residual. + te.ops.Linear(4096, 28672), + te.ops.SwiGLU(), + ) + fc2 = te.ops.Sequential( + te.ops.Linear(14336, 4096), + te.ops.AddExtraInput(), # Add the residual. + ) + + # Pass the extra output from fc1 as the extra input to fc2. + x = torch.randn(16384, 4096, device="cuda") + y, residual = fc1(x) + y = fc2(y, residual) + +.. figure:: ./residual_layernorm_mlp.png + :align: center + + Operations for an MLP block with a residual connection. The block + is split into two sections so that the caller can pass the extra + output from the first section to the second. + +Extra tensor channels +""""""""""""""""""""" + +Extra inputs and 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``. Slots bound to channels are removed from the public +``Sequential`` interface. + +With a channel, the residual block above can be expressed using one +``Sequential``: .. code-block:: python import torch import transformer_engine.pytorch as te - # Keep a residual connection inside one Sequential. make_residual = te.ops.MakeExtraOutput() add_residual = te.ops.AddExtraInput() make_residual.set_extra_output_channel(0, "residual") @@ -141,24 +182,10 @@ output slot and a later input slot connects them internally instead. add_residual, ) + # The residual is routed internally, so the caller receives only y. x = torch.randn(16384, 4096, device="cuda") y = block(x) -.. figure:: ./residual_layernorm_mlp.png - :align: center - - Operations for an MLP block with a residual connection. - -Extra tensor channels -""""""""""""""""""""" - -An extra output and one or more later extra inputs can be assigned the -same channel name. This routes the tensor inside the -``OperationFuser`` and removes the bound slots from the public -``Sequential`` interface. In the residual example above, the caller -therefore receives only ``y`` and does not need to pass the residual -back into the block. - 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 @@ -201,6 +228,29 @@ routing map. ``Combine`` consumes the routing map. # Channels supply all later extra inputs internally. y = 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 From 5070e347ff58b5a583d5a7bdb11460d03ab62f81 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 7 Aug 2026 23:27:15 +0000 Subject: [PATCH 13/83] address review comments Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 312 ++++++++++++++++++------ transformer_engine/pytorch/ops/fuser.py | 177 +++++++++----- transformer_engine/pytorch/ops/op.py | 25 +- 3 files changed, 365 insertions(+), 149 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 1616d830ff..61b569dbdc 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -457,11 +457,16 @@ def test_internal_residual_connection(self, size: int = 16) -> None: y.sum().backward() torch.testing.assert_close(x.grad, torch.full_like(x, 2)) - def test_fused_internal_residual_connection(self, size: int = 16) -> None: - """A fused op can implement both ends of an internal channel.""" + @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) + def test_fused_internal_residual_connection( + self, + fusion_kind: str, + size: int = 16, + ) -> None: + """Forward, backward, and joint fusions can own an internal channel.""" class FusedResidual(te_ops.FusedOperation): - """Fuse MakeExtraOutput, Bias, and AddExtraInput in forward.""" + """Fuse MakeExtraOutput, Bias, and AddExtraInput.""" _enabled = True @@ -480,7 +485,24 @@ def fuser_forward( # The consumer slot is internal to this fusion, so the # OperationFuser deliberately leaves it unset. assert basic_op_extra_inputs[2][0] is None - return 2 * input_, [(input_,), (), ()] + return 2 * input_ + self.basic_ops[1].bias, [(input_,), (), ()] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + del basic_op_ctxs + # The fusion owns the internal residual edge, including its + # contribution to the input gradient. + assert basic_op_grad_extra_outputs[0][0] is None + return ( + 2 * grad_output, + [(), (grad_output,), ()], + [(), (), (grad_output,)], + ) def fuse_residual(ops, **unused): if not FusedResidual._enabled: @@ -488,7 +510,7 @@ def fuse_residual(ops, **unused): if ( len(ops) == 3 and isinstance(ops[0], te_ops.MakeExtraOutput) - and isinstance(ops[1], te_ops.Identity) + and isinstance(ops[1], te_ops.Bias) and isinstance(ops[2], te_ops.AddExtraInput) ): FusedResidual._enabled = False @@ -496,22 +518,40 @@ def fuse_residual(ops, **unused): return ops residual = te_ops.MakeExtraOutput() - body = te_ops.Identity() + body = te_ops.Bias(size=size, device="cpu") add_residual = te_ops.AddExtraInput() residual.set_extra_output_channel(0, "residual") add_residual.set_extra_input_channel(0, "residual") model = te_ops.Sequential(residual, body, add_residual) - te_ops.register_forward_fusion(fuse_residual, prepend=True) + if fusion_kind == "forward": + te_ops.register_forward_fusion(fuse_residual, prepend=True) + elif fusion_kind == "backward": + te_ops.register_backward_fusion(fuse_residual, prepend=True) + else: + te_ops.register_forward_backward_fusion(fuse_residual, prepend=True) x = torch.rand((size,), requires_grad=True) y = model(x) forward_ops = model._module_groups[0]._forward_ops - assert len(forward_ops) == 1 - assert isinstance(forward_ops[0][0], FusedResidual) - torch.testing.assert_close(y, 2 * x) - y.sum().backward() - torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + backward_ops = model._module_groups[0]._backward_ops + if fusion_kind in ("forward", "forward_backward"): + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], FusedResidual) + else: + assert len(forward_ops) == 3 + if fusion_kind in ("backward", "forward_backward"): + assert len(backward_ops) == 1 + assert isinstance(backward_ops[0][0], FusedResidual) + else: + assert len(backward_ops) == 3 + if fusion_kind == "forward_backward": + assert backward_ops[0][0] is forward_ops[0][0] + torch.testing.assert_close(y, 2 * x + body.bias) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(body.bias.grad, dy) def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: """An internal extra output can feed multiple later consumers.""" @@ -528,9 +568,10 @@ def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: # Main path: x -> x + route -> x + route + route. torch.testing.assert_close(y, 3 * x) - y.sum().backward() + dy = torch.rand_like(y) + y.backward(dy) # The channel fan-out contributes two independent gradient paths. - torch.testing.assert_close(x.grad, torch.full_like(x, 3)) + torch.testing.assert_close(x.grad, 3 * dy) # Internal slots are unavailable before forward, so grad discovery # must tolerate them when no public input requires gradients. @@ -551,16 +592,28 @@ def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None y = model(x, extra) torch.testing.assert_close(y, 2 * x + extra) - y.sum().backward() - torch.testing.assert_close(x.grad, torch.full_like(x, 2)) - torch.testing.assert_close(extra.grad, torch.ones_like(extra)) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(extra.grad, dy) - def test_consumer_channel_without_producer(self) -> None: - """Extra input bound to a channel that no earlier op produces.""" - consumer = te_ops.AddExtraInput() - consumer.set_extra_input_channel(0, "missing") - with pytest.raises(ValueError, match="has no earlier producer"): - OperationFuser([consumer]) + def test_external_named_extra_input_fanout(self, size: int = 16) -> None: + """One public tensor supplies every unmatched input with the same channel.""" + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + consumer1.set_extra_input_channel(0, "external") + consumer2.set_extra_input_channel(0, "external") + model = te_ops.Sequential(consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + torch.testing.assert_close(y, x + 2 * extra) + + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, dy) + torch.testing.assert_close(extra.grad, 2 * dy) def test_consumer_before_producer(self) -> None: """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" @@ -653,19 +706,17 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp with pytest.raises(ValueError, match="multiple producers"): OperationFuser([producer, consumer]) - def test_unused_extra_output_channel(self) -> None: - """Every produced channel must have at least one consumer.""" + def test_named_extra_output_without_consumer_is_public(self, size: int = 16) -> None: + """A named output remains public when its fuser has no consumer.""" producer = te_ops.MakeExtraOutput() producer.set_extra_output_channel(0, "orphan") - with pytest.raises(ValueError, match="have no consumers"): - OperationFuser([producer]) + x = torch.rand((size,), requires_grad=True) + y, extra = producer(x) + torch.testing.assert_close(y, x) + torch.testing.assert_close(extra, x) def test_one_extra_input_has_single_source(self) -> None: - """Each extra-input slot binds to one channel / one producer source. - - Rebinding replaces the previous name; the abandoned producer channel - then fails as unused rather than attaching two sources to one input. - """ + """Rebinding selects one source and leaves the other output public.""" producer_a = te_ops.MakeExtraOutput() producer_b = te_ops.MakeExtraOutput() consumer = te_ops.AddExtraInput() @@ -673,58 +724,171 @@ def test_one_extra_input_has_single_source(self) -> None: producer_b.set_extra_output_channel(0, "b") consumer.set_extra_input_channel(0, "a") consumer.set_extra_input_channel(0, "b") - with pytest.raises(ValueError, match="have no consumers"): - OperationFuser([producer_a, producer_b, consumer]) - - # Valid single binding: consumer input 0 is fed only by producer_a. - consumer.set_extra_input_channel(0, "a") - producer_b.set_extra_output_channel(0, None) fuser = OperationFuser([producer_a, producer_b, consumer]) - assert fuser._basic_op_extra_input_sources[2] == [(0, 0)] + assert fuser._basic_op_extra_input_sources[2] == [(1, 0)] assert fuser.num_extra_inputs == 0 - # producer_b's unbound extra output remains public - assert fuser._external_extra_output_slots == [(1, 0)] + assert fuser._external_extra_output_slots == [(0, 0)] - def test_channel_fanout_accumulates_grads(self, size: int = 16) -> None: - """Grads from every consumer of a channel are accumulated into the producer.""" - producer = te_ops.MakeExtraOutput() - consumer1 = te_ops.AddExtraInput() - consumer2 = te_ops.AddExtraInput() - producer.set_extra_output_channel(0, "route") - consumer1.set_extra_input_channel(0, "route") - consumer2.set_extra_input_channel(0, "route") - model = te_ops.Sequential(producer, consumer1, consumer2) + def test_mixed_channel_outputs_accept_generators(self, size: int = 16) -> None: + """Generator outputs support mixed internal and public channel slots.""" + + class DualExtraOutput(te_ops.BasicOperation): + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_ctxs, basic_op_extra_inputs + outputs = (2 * input_, 3 * input_) + return input_, (iter(outputs) for _ in range(1)) + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + grad_internal, grad_public = basic_op_grad_extra_outputs[0] + return ( + grad_output + 2 * grad_internal + 3 * grad_public, + (iter(()) for _ in range(1)), + (iter(()) for _ in range(1)), + ) + + producer = DualExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "internal") + producer.set_extra_output_channel(1, "public") + consumer.set_extra_input_channel(0, "internal") + model = te_ops.Sequential(producer, consumer) x = torch.rand((size,), requires_grad=True) - y = model(x) - # Forward: x -> x+x -> x+x+x + y, public = model(x) torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(public, 3 * x) - dy = torch.rand((size,)) - y.backward(dy) - # Main path contributes dy; each AddExtraInput also routes dy back - # through the channel into MakeExtraOutput's extra-output grad, which - # is added again into dx. Total: dy (main) + dy + dy (two consumers). - torch.testing.assert_close(x.grad, 3 * dy) + dy = torch.rand_like(y) + dpublic = torch.rand_like(public) + torch.autograd.backward((y, public), (dy, dpublic)) + torch.testing.assert_close(x.grad, 3 * dy + 3 * dpublic) - def test_mixed_internal_external_grad(self, size: int = 16) -> None: - """Internal channel grads and public extra-input grads both flow correctly.""" - producer = te_ops.MakeExtraOutput() - internal_consumer = te_ops.AddExtraInput() - external_consumer = te_ops.AddExtraInput() - producer.set_extra_output_channel(0, "route") - internal_consumer.set_extra_input_channel(0, "route") - model = te_ops.Sequential(producer, internal_consumer, external_consumer) + def test_fresh_internal_output_preserves_grad_requirement(self) -> None: + """A fresh internal tensor requests its gradient from a scaled activation.""" - x = torch.rand((size,), requires_grad=True) - extra = torch.rand((size,), requires_grad=True) - y = model(x, extra) - torch.testing.assert_close(y, 2 * x + extra) + class MakeScale(te_ops.BasicOperation): + num_extra_outputs = 1 - dy = torch.rand((size,)) - y.backward(dy) - torch.testing.assert_close(x.grad, 2 * dy) - torch.testing.assert_close(extra.grad, dy) + def op_forward(self, *args, **kwargs): + raise RuntimeError("MakeScale uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("MakeScale uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_extra_inputs + basic_op_ctxs[0].save_for_backward(input_) + return input_, [(input_.square().mean(dim=-1),)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + (input_,) = basic_op_ctxs[0].saved_tensors + grad_scale = basic_op_grad_extra_outputs[0][0] + assert grad_scale is not None + grad_input = grad_output + grad_scale.unsqueeze(-1) * 2 * input_ / input_.size(-1) + return grad_input, [()], [()] + + producer = MakeScale() + activation = te_ops.ScaledSReLU() + producer.set_extra_output_channel(0, "scale") + activation.set_extra_input_channel(0, "scale") + model = te_ops.Sequential(producer, activation) + + x_ref = torch.randn((5, 8), device="cuda", requires_grad=True) + x_test = x_ref.detach().clone().requires_grad_(True) + scale_ref = x_ref.square().mean(dim=-1) + y_ref = torch.nn.functional.relu(x_ref).square() * scale_ref.unsqueeze(-1) + y_test = model(x_test) + torch.testing.assert_close(y_test, y_ref) + + dy = torch.rand_like(y_ref) + y_ref.backward(dy) + y_test.backward(dy) + torch.testing.assert_close(x_test.grad, x_ref.grad) + + + def test_grouped_linear_scale_bias_channels(self) -> None: + """Both GroupedLinear extra inputs can be supplied by channels.""" + + class RouteExtras(te_ops.BasicOperation): + num_extra_inputs = 2 + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("RouteExtras uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("RouteExtras uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_ctxs + return input_, [basic_op_extra_inputs[0]] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + return grad_output, [()], [basic_op_grad_extra_outputs[0]] + + group_size, in_features, out_features = 2, 8, 6 + split_sizes = torch.tensor((3, 2), dtype=torch.int32, device="cuda") + num_tokens = int(split_sizes.sum()) + x = torch.randn((num_tokens, in_features), device="cuda", requires_grad=True) + scales = torch.randn((num_tokens,), device="cuda", requires_grad=True) + + producer = RouteExtras() + linear = te_ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=True, + scale_bias=True, + device="cuda", + dtype=torch.float32, + ) + producer.set_extra_output_channel(0, "split_sizes") + producer.set_extra_output_channel(1, "bias_scales") + linear.set_extra_input_channel(0, "split_sizes") + linear.set_extra_input_channel(1, "bias_scales") + model = te_ops.Sequential(producer, linear) + + x_ref = x.detach().clone().requires_grad_(True) + scales_ref = scales.detach().clone().requires_grad_(True) + ys_ref = [] + for group_idx, (x_group, scale_group) in enumerate( + zip( + torch.split(x_ref, split_sizes.tolist()), + torch.split(scales_ref, split_sizes.tolist()), + ) + ): + weight = getattr(linear, f"weight{group_idx}") + bias = getattr(linear, f"bias{group_idx}") + ys_ref.append( + torch.nn.functional.linear(x_group, weight) + + scale_group.unsqueeze(-1) * bias + ) + y_ref = torch.cat(ys_ref) + y_test = model(x, split_sizes, scales) + dy = torch.rand_like(y_test) + grads_ref = torch.autograd.grad( + y_ref, + (x_ref, scales_ref, *linear.parameters()), + dy, + ) + y_test.backward(dy) + + tols = dtype_tols(torch.float16) # Grouped GEMM uses TF32 for FP32 inputs. + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(x.grad, grads_ref[0], **tols) + torch.testing.assert_close(scales.grad, grads_ref[1], **tols) + for param, grad_ref in zip(linear.parameters(), grads_ref[2:]): + torch.testing.assert_close(param.grad, grad_ref, **tols) class TestFuser: diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index c4d7e1ec10..d4e9015b9e 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -108,8 +108,9 @@ def forward( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in fuser._basic_ops ] - for tensor, (op_idx, input_idx) in zip(extra_inputs, fuser._external_extra_input_slots): - basic_op_extra_inputs[op_idx][input_idx] = tensor + for tensor, slots in zip(extra_inputs, fuser._external_extra_input_slots): + for op_idx, input_idx in slots: + basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops x = input_ @@ -167,34 +168,40 @@ def forward( next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) + fused_op_extra_outputs = tuple(tuple(ys) for ys in fused_op_extra_outputs) + if len(fused_op_extra_outputs) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate extra outputs for " + f"{len(basic_op_idxs)} basic operations, " + f"but got {len(fused_op_extra_outputs)}" + ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): + num_extra_outputs = fuser._basic_ops[idx].num_extra_outputs + if len(ys) != num_extra_outputs: + raise RuntimeError( + f"Expected op {idx} to generate {num_extra_outputs} extra outputs, " + f"but got {len(ys)}" + ) for output_idx, y in enumerate(ys): if y is None: raise RuntimeError( f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " f"did not emit extra output {output_idx}" ) + if ( + set_output_requires_grad + and idx >= fuser.first_op_requiring_backward + and (y.is_floating_point() or y.is_complex()) + ): + y.requires_grad_(True) extra_outputs[idx] = ys - # Validate extra outputs and flatten only public slots. Outputs bound - # to channels stay internal to the fuser and are not marked. - extra_outputs_flat = [] - for idx, ys in enumerate(extra_outputs): - ys = list(ys) - num_extra_outputs = fuser._basic_ops[idx].num_extra_outputs - if len(ys) != num_extra_outputs: - raise RuntimeError( - f"Expected op {idx} to generate " - "{num_extra_outputs} extra inputs, " - f"but got {len(ys)}" - ) - for output_idx, y in enumerate(ys): - # Output is not bound to a channel consumed by another op, - # so it is a public output. - if fuser._basic_op_extra_output_channels[idx][output_idx] is None: - if set_output_requires_grad: - y.requires_grad_(idx >= fuser.first_op_requiring_backward) - extra_outputs_flat.append(y) + # Flatten public extra outputs. Matched channels stay internal, while + # unnamed slots and named channels without consumers remain public. + extra_outputs_flat = [ + extra_outputs[op_idx][output_idx] + for op_idx, output_idx in fuser._external_extra_output_slots + ] # Save context for backward pass if func_ctx is not None: @@ -226,8 +233,12 @@ def forward( func_ctx.basic_op_ctxs = basic_op_ctxs func_ctx.basic_op_num_params = fuser._basic_op_num_params func_ctx.num_extra_outputs = len(extra_outputs_flat) + func_ctx.external_extra_input_slots = fuser._external_extra_input_slots func_ctx.external_extra_output_slots = fuser._external_extra_output_slots func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_output_is_internal = ( + fuser._basic_op_extra_output_is_internal + ) func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module @@ -269,6 +280,7 @@ def backward( # Channel wiring saved from forward external_extra_output_slots = func_ctx.external_extra_output_slots basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels + basic_op_extra_output_is_internal = func_ctx.basic_op_extra_output_is_internal basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources # Place public extra-output grads into their basic-op slots. Internal @@ -300,7 +312,7 @@ def backward( # each internal channel. for idx in basic_op_idxs: for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): - if channel is not None: + if basic_op_extra_output_is_internal[idx][output_idx]: basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get(channel) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs @@ -310,6 +322,22 @@ def backward( dx, basic_op_grad_extra_outputs=op_grad_extra_outputs, ) + fused_op_grad_params = tuple(tuple(grads) for grads in fused_op_grad_params) + fused_op_grad_extra_inputs = tuple( + tuple(grads) for grads in fused_op_grad_extra_inputs + ) + if len(fused_op_grad_params) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate parameter grads for " + f"{len(basic_op_idxs)} basic operations, but got " + f"{len(fused_op_grad_params)}" + ) + if len(fused_op_grad_extra_inputs) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate extra-input grads for " + f"{len(basic_op_idxs)} basic operations, but got " + f"{len(fused_op_grad_extra_inputs)}" + ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None @@ -344,24 +372,27 @@ def backward( grad_params_flat.extend(dparams) # Flatten list of parameter gradients - grad_extra_inputs_flat = [] for idx, dxs in enumerate(grad_extra_inputs): num_extra_inputs = basic_ops[idx].num_extra_inputs if dxs is None: - dxs = [None for _ in range(num_extra_inputs)] - else: - dxs = list(dxs) - if len(dxs) != num_extra_inputs: + grad_extra_inputs[idx] = (None,) * num_extra_inputs + elif len(dxs) != num_extra_inputs: raise RuntimeError( f"Expected op {idx} to generate grads " f"for {num_extra_inputs} extra inputs, " f"but got {len(dxs)}" ) - for input_idx, grad in enumerate(dxs): - # Only append public grad extra inputs to the list - # to be returned to the user. - if basic_op_extra_input_sources[idx][input_idx] is None: - grad_extra_inputs_flat.append(grad) + + # One public tensor may fan out to several unmatched slots sharing a + # channel name, so sum their gradients before returning to autograd. + grad_extra_inputs_flat = [] + for slots in func_ctx.external_extra_input_slots: + grad = None + for op_idx, input_idx in slots: + slot_grad = grad_extra_inputs[op_idx][input_idx] + if slot_grad is not None: + grad = slot_grad if grad is None else grad + slot_grad + grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -401,6 +432,8 @@ class OperationFuser: def __init__( self, ops: list[FusibleOperation], + *, + lock_extra_channels: bool = True, ) -> None: # Get list of basic operations @@ -421,29 +454,17 @@ def __init__( self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ list(op._extra_output_channels) for op in basic_ops ] - self._external_extra_input_slots: list[tuple[int, int]] = [] + self._basic_op_extra_output_is_internal: list[list[bool]] = [ + [False] * op.num_extra_outputs for op in basic_ops + ] + self._external_extra_input_slots: list[list[tuple[int, int]]] = [] self._external_extra_output_slots: list[tuple[int, int]] = [] - # Resolve named channels in pipeline order. Channels deliberately only - # connect an output to later inputs, which keeps execution acyclic. + # Find channel producers and reject ambiguous names. channel_producers: dict[str, tuple[int, int]] = {} - consumed_channels: set[str] = set() for op_idx, op in enumerate(basic_ops): - for input_idx in range(op.num_extra_inputs): - channel = op._extra_input_channels[input_idx] - if channel is None: - self._external_extra_input_slots.append((op_idx, input_idx)) - continue - if channel not in channel_producers: - raise ValueError( - f"Extra tensor channel {channel!r} consumed by op {op_idx} " - f"({type(op).__name__}) has no earlier producer" - ) - self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[channel] - consumed_channels.add(channel) for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): if channel is None: - self._external_extra_output_slots.append((op_idx, output_idx)) continue if channel in channel_producers: producer_idx, _ = channel_producers[channel] @@ -453,10 +474,41 @@ def __init__( ) channel_producers[channel] = (op_idx, output_idx) - unused_channels = channel_producers.keys() - consumed_channels - if unused_channels: - channels = ", ".join(repr(channel) for channel in sorted(unused_channels)) - raise ValueError(f"Extra tensor channels have no consumers: {channels}") + # Resolve inputs. A channel with an earlier producer is internal. A + # channel without a producer is public, with one positional tensor + # fanning out to every public slot that shares the channel name. + external_input_channels: dict[str, int] = {} + consumed_channels: set[str] = set() + for op_idx, op in enumerate(basic_ops): + for input_idx, channel in enumerate(op._extra_input_channels): + if channel is None: + self._external_extra_input_slots.append([(op_idx, input_idx)]) + continue + producer = channel_producers.get(channel) + if producer is None: + group_idx = external_input_channels.get(channel) + if group_idx is None: + group_idx = len(self._external_extra_input_slots) + external_input_channels[channel] = group_idx + self._external_extra_input_slots.append([]) + self._external_extra_input_slots[group_idx].append((op_idx, input_idx)) + continue + producer_idx, _ = producer + if producer_idx >= op_idx: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no earlier producer" + ) + self._basic_op_extra_input_sources[op_idx][input_idx] = producer + consumed_channels.add(channel) + + # Unnamed outputs and named outputs without local consumers are public. + for op_idx, op in enumerate(basic_ops): + for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): + if channel is not None and channel in consumed_channels: + self._basic_op_extra_output_is_internal[op_idx][output_idx] = True + else: + self._external_extra_output_slots.append((op_idx, output_idx)) # Every channel-bound extra input must be wired to a matching producer # extra output. External slots remain unbound (source is None). @@ -473,11 +525,7 @@ def __init__( ) continue if source is None: - raise ValueError( - f"Extra input {input_idx} of op {op_idx} " - f"({type(op).__name__}) is bound to channel {channel!r} " - "but has no producer" - ) + continue producer_idx, output_idx = source producer_channel = self._basic_op_extra_output_channels[producer_idx][output_idx] if producer_channel != channel: @@ -506,10 +554,10 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) - # Channel routing is structural fuser state. Prevent the basic ops from - # changing their bindings after this fuser captures them. - for op in self._basic_ops: - op._lock_extra_channels() + # Persistent fusers capture channel routing as structural state. + if lock_extra_channels: + for op in self._basic_ops: + op._lock_extra_channels() @staticmethod def _apply_fusions( @@ -674,8 +722,9 @@ def __call__( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in self._basic_ops ] - for tensor, (op_idx, input_idx) in zip(extra_inputs, self._external_extra_input_slots): - basic_op_extra_inputs[op_idx][input_idx] = tensor + for tensor, slots in zip(extra_inputs, self._external_extra_input_slots): + for op_idx, input_idx in slots: + basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state recipe = None diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index d9f9624ab0..d668a7c904 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -198,12 +198,14 @@ def __init__(self) -> None: self._quantizers: Optional[dict[str, list[Quantizer]]] = None def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: - """Bind an extra input slot to an internal fuser channel. + """Assign a channel name to an extra input slot. - A bound slot receives the matching extra output from an earlier - operation in the same fuser instead of consuming a public extra input. - Passing ``None`` removes the binding. Bindings cannot be changed after - the operation has been attached to an ``OperationFuser``. + The slot receives the matching extra output from an earlier operation + in the same fuser. If there is no producer in the fuser, the slot + remains public; one public tensor fans out to all input slots with the + same channel name. Passing ``None`` removes the name. Channels cannot + be changed after the operation has been attached to a persistent + ``OperationFuser``. """ if not 0 <= index < self.num_extra_inputs: raise IndexError( @@ -219,11 +221,12 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp return self def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: - """Bind an extra output slot to an internal fuser channel. + """Assign a channel name to an extra output slot. - A bound slot can feed one or more later operations and is not returned - as a public extra output. Passing ``None`` removes the binding. Bindings - cannot be changed after the operation has been attached to an + The slot feeds matching extra inputs on later operations in the same + fuser. If there are no consumers in the fuser, the slot remains a + public extra output. Passing ``None`` removes the name. Channels cannot + be changed after the operation has been attached to a persistent ``OperationFuser``. """ if not 0 <= index < self.num_extra_outputs: @@ -595,7 +598,7 @@ def forward( """Apply operation""" from .fuser import OperationFuser - return OperationFuser([self])( + return OperationFuser([self], lock_extra_channels=False)( input, *extra_inputs, basic_op_kwargs=[kwargs], @@ -830,7 +833,7 @@ def forward( basic_op_kwargs = [{} for _ in range(len(self.basic_ops))] from .fuser import OperationFuser - return OperationFuser([self])( + return OperationFuser([self], lock_extra_channels=False)( input, *extra_inputs, basic_op_kwargs=basic_op_kwargs, From ae41ad3261809282c6f6464e494ad6c02ed64854 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:28:23 +0000 Subject: [PATCH 14/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 4 +--- transformer_engine/pytorch/ops/fuser.py | 8 ++------ 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 61b569dbdc..5362fac716 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -814,7 +814,6 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp y_test.backward(dy) torch.testing.assert_close(x_test.grad, x_ref.grad) - def test_grouped_linear_scale_bias_channels(self) -> None: """Both GroupedLinear extra inputs can be supplied by channels.""" @@ -870,8 +869,7 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp weight = getattr(linear, f"weight{group_idx}") bias = getattr(linear, f"bias{group_idx}") ys_ref.append( - torch.nn.functional.linear(x_group, weight) - + scale_group.unsqueeze(-1) * bias + torch.nn.functional.linear(x_group, weight) + scale_group.unsqueeze(-1) * bias ) y_ref = torch.cat(ys_ref) y_test = model(x, split_sizes, scales) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index d4e9015b9e..7c641ce7fa 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -236,9 +236,7 @@ def forward( func_ctx.external_extra_input_slots = fuser._external_extra_input_slots func_ctx.external_extra_output_slots = fuser._external_extra_output_slots func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels - func_ctx.basic_op_extra_output_is_internal = ( - fuser._basic_op_extra_output_is_internal - ) + func_ctx.basic_op_extra_output_is_internal = fuser._basic_op_extra_output_is_internal func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module @@ -323,9 +321,7 @@ def backward( basic_op_grad_extra_outputs=op_grad_extra_outputs, ) fused_op_grad_params = tuple(tuple(grads) for grads in fused_op_grad_params) - fused_op_grad_extra_inputs = tuple( - tuple(grads) for grads in fused_op_grad_extra_inputs - ) + fused_op_grad_extra_inputs = tuple(tuple(grads) for grads in fused_op_grad_extra_inputs) if len(fused_op_grad_params) != len(basic_op_idxs): raise RuntimeError( f"Expected {type(op).__name__} to generate parameter grads for " From f82cbedd5b2eb85241e1b8980c0b880d88c042fd Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 00:58:32 +0000 Subject: [PATCH 15/83] some cleanup Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 89 +++++++++------ .../pytorch/ops/basic/activation.py | 4 +- .../pytorch/ops/basic/add_extra_input.py | 4 +- .../pytorch/ops/basic/grouped_linear.py | 2 +- .../pytorch/ops/basic/make_extra_output.py | 4 +- .../pytorch/ops/basic/swiglu.py | 4 +- .../fused/forward_linear_bias_activation.py | 4 +- .../ops/fused/forward_linear_bias_add.py | 4 +- .../ops/fused/forward_linear_scale_add.py | 4 +- .../pytorch/ops/fused/grouped_mlp.py | 4 +- .../ops/fused/userbuffers_forward_linear.py | 4 +- transformer_engine/pytorch/ops/fuser.py | 107 +++++++----------- transformer_engine/pytorch/ops/op.py | 28 ++--- 13 files changed, 126 insertions(+), 136 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 5362fac716..dcac932db6 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -451,11 +451,15 @@ def test_internal_residual_connection(self, size: int = 16) -> None: model = te_ops.Sequential(residual, body, add_residual) x = torch.rand((size,), requires_grad=True) - y = model(x) + y, residual_out = model(x) torch.testing.assert_close(y, 2 * x + body.bias) - y.sum().backward() - torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + torch.testing.assert_close(residual_out, x) + dy = torch.rand_like(y) + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + torch.testing.assert_close(x.grad, 2 * dy + dresidual) + torch.testing.assert_close(body.bias.grad, dy) @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) def test_fused_internal_residual_connection( @@ -495,11 +499,12 @@ def fuser_backward( basic_op_grad_extra_outputs, ): del basic_op_ctxs - # The fusion owns the internal residual edge, including its - # contribution to the input gradient. - assert basic_op_grad_extra_outputs[0][0] is None + # The fusion owns the internal residual edge. The fuser also + # supplies the gradient from the public residual output. + grad_residual = basic_op_grad_extra_outputs[0][0] return ( - 2 * grad_output, + 2 * grad_output + + (torch.zeros_like(grad_output) if grad_residual is None else grad_residual), [(), (grad_output,), ()], [(), (), (grad_output,)], ) @@ -531,7 +536,7 @@ def fuse_residual(ops, **unused): else: te_ops.register_forward_backward_fusion(fuse_residual, prepend=True) x = torch.rand((size,), requires_grad=True) - y = model(x) + y, residual_out = model(x) forward_ops = model._module_groups[0]._forward_ops backward_ops = model._module_groups[0]._backward_ops @@ -549,8 +554,9 @@ def fuse_residual(ops, **unused): assert backward_ops[0][0] is forward_ops[0][0] torch.testing.assert_close(y, 2 * x + body.bias) dy = torch.rand_like(y) - y.backward(dy) - torch.testing.assert_close(x.grad, 2 * dy) + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + torch.testing.assert_close(x.grad, 2 * dy + dresidual) torch.testing.assert_close(body.bias.grad, dy) def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: @@ -564,19 +570,23 @@ def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: model = te_ops.Sequential(producer, consumer1, consumer2) x = torch.rand((size,), requires_grad=True) - y = model(x) + y, route = model(x) # Main path: x -> x + route -> x + route + route. torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(route, x) dy = torch.rand_like(y) - y.backward(dy) + droute = torch.rand_like(route) + torch.autograd.backward((y, route), (dy, droute)) # The channel fan-out contributes two independent gradient paths. - torch.testing.assert_close(x.grad, 3 * dy) + torch.testing.assert_close(x.grad, 3 * dy + droute) # Internal slots are unavailable before forward, so grad discovery # must tolerate them when no public input requires gradients. x_no_grad = x.detach() - torch.testing.assert_close(model(x_no_grad), 3 * x_no_grad) + y_no_grad, route_no_grad = model(x_no_grad) + torch.testing.assert_close(y_no_grad, 3 * x_no_grad) + torch.testing.assert_close(route_no_grad, x_no_grad) def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: """Unbound slots remain public when other slots use internal channels.""" @@ -589,16 +599,17 @@ def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None x = torch.rand((size,), requires_grad=True) extra = torch.rand((size,), requires_grad=True) - y = model(x, extra) + y, route = model(x, extra) torch.testing.assert_close(y, 2 * x + extra) + torch.testing.assert_close(route, x) dy = torch.rand_like(y) y.backward(dy) torch.testing.assert_close(x.grad, 2 * dy) torch.testing.assert_close(extra.grad, dy) - def test_external_named_extra_input_fanout(self, size: int = 16) -> None: - """One public tensor supplies every unmatched input with the same channel.""" + def test_external_named_extra_inputs_remain_separate(self, size: int = 16) -> None: + """Unmatched inputs with the same channel require separate public tensors.""" consumer1 = te_ops.AddExtraInput() consumer2 = te_ops.AddExtraInput() consumer1.set_extra_input_channel(0, "external") @@ -606,14 +617,18 @@ def test_external_named_extra_input_fanout(self, size: int = 16) -> None: model = te_ops.Sequential(consumer1, consumer2) x = torch.rand((size,), requires_grad=True) - extra = torch.rand((size,), requires_grad=True) - y = model(x, extra) - torch.testing.assert_close(y, x + 2 * extra) + extra1 = torch.rand((size,), requires_grad=True) + extra2 = torch.rand((size,), requires_grad=True) + with pytest.raises(ValueError, match="Expected 2 extra inputs but got 1"): + model(x, extra1) + y = model(x, extra1, extra2) + torch.testing.assert_close(y, x + extra1 + extra2) dy = torch.rand_like(y) y.backward(dy) torch.testing.assert_close(x.grad, dy) - torch.testing.assert_close(extra.grad, 2 * dy) + torch.testing.assert_close(extra1.grad, dy) + torch.testing.assert_close(extra2.grad, dy) def test_consumer_before_producer(self) -> None: """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" @@ -715,7 +730,7 @@ def test_named_extra_output_without_consumer_is_public(self, size: int = 16) -> torch.testing.assert_close(y, x) torch.testing.assert_close(extra, x) - def test_one_extra_input_has_single_source(self) -> None: + def test_one_extra_input_has_single_source(self, size: int = 16) -> None: """Rebinding selects one source and leaves the other output public.""" producer_a = te_ops.MakeExtraOutput() producer_b = te_ops.MakeExtraOutput() @@ -727,10 +742,15 @@ def test_one_extra_input_has_single_source(self) -> None: fuser = OperationFuser([producer_a, producer_b, consumer]) assert fuser._basic_op_extra_input_sources[2] == [(1, 0)] assert fuser.num_extra_inputs == 0 - assert fuser._external_extra_output_slots == [(0, 0)] - def test_mixed_channel_outputs_accept_generators(self, size: int = 16) -> None: - """Generator outputs support mixed internal and public channel slots.""" + x = torch.rand((size,)) + y, output_a, output_b = fuser(x) + torch.testing.assert_close(y, 2 * x) + torch.testing.assert_close(output_a, x) + torch.testing.assert_close(output_b, x) + + def test_mixed_channel_outputs_are_public(self, size: int = 16) -> None: + """Both internally consumed and unconsumed channel outputs are public.""" class DualExtraOutput(te_ops.BasicOperation): num_extra_outputs = 2 @@ -743,16 +763,15 @@ def op_backward(self, *args, **kwargs): def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): del basic_op_ctxs, basic_op_extra_inputs - outputs = (2 * input_, 3 * input_) - return input_, (iter(outputs) for _ in range(1)) + return input_, [(2 * input_, 3 * input_)] def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): del basic_op_ctxs grad_internal, grad_public = basic_op_grad_extra_outputs[0] return ( grad_output + 2 * grad_internal + 3 * grad_public, - (iter(()) for _ in range(1)), - (iter(()) for _ in range(1)), + [()], + [()], ) producer = DualExtraOutput() @@ -763,14 +782,16 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp model = te_ops.Sequential(producer, consumer) x = torch.rand((size,), requires_grad=True) - y, public = model(x) + y, internal, public = model(x) torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(internal, 2 * x) torch.testing.assert_close(public, 3 * x) dy = torch.rand_like(y) + dinternal = torch.rand_like(internal) dpublic = torch.rand_like(public) - torch.autograd.backward((y, public), (dy, dpublic)) - torch.testing.assert_close(x.grad, 3 * dy + 3 * dpublic) + torch.autograd.backward((y, internal, public), (dy, dinternal, dpublic)) + torch.testing.assert_close(x.grad, 3 * dy + 2 * dinternal + 3 * dpublic) def test_fresh_internal_output_preserves_grad_requirement(self) -> None: """A fresh internal tensor requests its gradient from a scaled activation.""" @@ -806,7 +827,7 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp x_test = x_ref.detach().clone().requires_grad_(True) scale_ref = x_ref.square().mean(dim=-1) y_ref = torch.nn.functional.relu(x_ref).square() * scale_ref.unsqueeze(-1) - y_test = model(x_test) + y_test, _scale_test = model(x_test) torch.testing.assert_close(y_test, y_ref) dy = torch.rand_like(y_ref) @@ -872,7 +893,7 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp torch.nn.functional.linear(x_group, weight) + scale_group.unsqueeze(-1) * bias ) y_ref = torch.cat(ys_ref) - y_test = model(x, split_sizes, scales) + y_test, _split_sizes_test, _scales_test = model(x, split_sizes, scales) dy = torch.rand_like(y_test) grads_ref = torch.autograd.grad( y_ref, diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..5159ea50e9 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -392,7 +392,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], # pylint: disable=unused-argument next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: if self.activation_recompute_in_mlp: raise RuntimeError( f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " diff --git a/transformer_engine/pytorch/ops/basic/add_extra_input.py b/transformer_engine/pytorch/ops/basic/add_extra_input.py index fc3ca9cade..9af399f2dc 100644 --- a/transformer_engine/pytorch/ops/basic/add_extra_input.py +++ b/transformer_engine/pytorch/ops/basic/add_extra_input.py @@ -5,7 +5,7 @@ """Fusible operation for adding extra input tensor.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -67,7 +67,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: extra_input = basic_op_extra_inputs[0][0] if self._in_place: extra_input = extra_input.detach() diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 7faad6536b..60e0b72ab9 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1003,7 +1003,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: num_groups = self.num_groups weight_param = self.weight if self.single_grouped_weight else self.weight0 device = weight_param.device diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 0d9c870262..5ad00e2fd1 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -5,7 +5,7 @@ """Make extra tensor output in operation fuser.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -72,7 +72,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: return input_, [(input_,)] def fuser_backward( diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 02f330ede3..a1c65e0998 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -5,7 +5,7 @@ """Fusible operation for SwiGLU and variants.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -422,7 +422,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: if self.activation_recompute_in_mlp: raise RuntimeError( f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 8df929f799..b6f56effef 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + activation.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -59,7 +59,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 5376a7d264..f32eb0f7ac 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + add.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -57,7 +57,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index abeb39adfa..2fb6c01dca 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + scale + add.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -47,7 +47,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations linear_op = self.basic_ops[0] diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 0113833647..7c7502f4af 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -6,7 +6,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence import functools import os from importlib.metadata import PackageNotFoundError, version as get_pkg_version @@ -975,7 +975,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations fc1_op, activation_op, fc2_op = self.basic_ops fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 3a8ff5438d..93c4024880 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -5,7 +5,7 @@ """Linear layer forward with Userbuffers communication.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -286,7 +286,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 7c641ce7fa..3232f2d86e 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -108,9 +108,11 @@ def forward( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in fuser._basic_ops ] - for tensor, slots in zip(extra_inputs, fuser._external_extra_input_slots): - for op_idx, input_idx in slots: - basic_op_extra_inputs[op_idx][input_idx] = tensor + for tensor, (op_idx, input_idx) in zip( + extra_inputs, + fuser._external_extra_input_slots, + ): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops x = input_ @@ -134,21 +136,8 @@ def forward( # fused op will wire the channel itself internally continue producer_outputs = extra_outputs[producer_idx] - if producer_outputs is None: - raise RuntimeError( - f"Extra tensor channel producer op {producer_idx} has not run" - ) - if output_idx >= len(producer_outputs) or producer_outputs[output_idx] is None: - raise RuntimeError( - f"Extra tensor channel producer op {producer_idx} " - f"({type(fuser._basic_ops[producer_idx]).__name__}) " - f"did not emit extra output {output_idx} for " - f"consumer op {idx} " - f"({type(fuser._basic_ops[idx]).__name__}) " - f"input {input_idx}" - ) basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] - op_extra_inputs = [tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs] + extra_inputs = [basic_op_extra_inputs[idx] for idx in basic_op_idxs] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None prev_op_grad_output_quantizer = None @@ -163,12 +152,11 @@ def forward( x, fused_op_extra_outputs = op.fuser_forward( [basic_op_ctxs[idx] for idx in basic_op_idxs], x, - basic_op_extra_inputs=op_extra_inputs, + basic_op_extra_inputs=extra_inputs, prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) - fused_op_extra_outputs = tuple(tuple(ys) for ys in fused_op_extra_outputs) if len(fused_op_extra_outputs) != len(basic_op_idxs): raise RuntimeError( f"Expected {type(op).__name__} to generate extra outputs for " @@ -196,12 +184,9 @@ def forward( y.requires_grad_(True) extra_outputs[idx] = ys - # Flatten public extra outputs. Matched channels stay internal, while - # unnamed slots and named channels without consumers remain public. - extra_outputs_flat = [ - extra_outputs[op_idx][output_idx] - for op_idx, output_idx in fuser._external_extra_output_slots - ] + # Flatten all extra outputs in basic-operation order. Channel-bound + # outputs remain visible even when they are also consumed internally. + extra_outputs_flat = [y for ys in extra_outputs for y in ys] # Save context for backward pass if func_ctx is not None: @@ -234,15 +219,13 @@ def forward( func_ctx.basic_op_num_params = fuser._basic_op_num_params func_ctx.num_extra_outputs = len(extra_outputs_flat) func_ctx.external_extra_input_slots = fuser._external_extra_input_slots - func_ctx.external_extra_output_slots = fuser._external_extra_output_slots func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels func_ctx.basic_op_extra_output_is_internal = fuser._basic_op_extra_output_is_internal func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward - all_extra_outputs = [y for ys in extra_outputs for y in ys] - for tensor in [x] + all_extra_outputs: + for tensor in [x] + extra_outputs_flat: tensor._do_not_clear = True if set_output_requires_grad: @@ -276,23 +259,25 @@ def backward( ctx._saved_tensors_range = None # Channel wiring saved from forward - external_extra_output_slots = func_ctx.external_extra_output_slots basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels basic_op_extra_output_is_internal = func_ctx.basic_op_extra_output_is_internal basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources - # Place public extra-output grads into their basic-op slots. Internal - # output grads are accumulated from channel consumers during backward. + # Place caller-provided extra-output grads into their basic-op slots. + # Gradients from internal channel consumers are added during backward. if len(grad_extra_outputs) != func_ctx.num_extra_outputs: raise ValueError( f"Expected grads for {func_ctx.num_extra_outputs} extra tensor outputs, " f"but got {len(grad_extra_outputs)}" ) - basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ - [None] * op.num_extra_outputs for op in basic_ops - ] - for grad, (op_idx, output_idx) in zip(grad_extra_outputs, external_extra_output_slots): - basic_op_grad_extra_outputs[op_idx][output_idx] = grad + remaining_grad_extra_outputs = grad_extra_outputs + basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [] + for op in basic_ops: + grads, remaining_grad_extra_outputs = _split_tuple( + remaining_grad_extra_outputs, + op.num_extra_outputs, + ) + basic_op_grad_extra_outputs.append(list(grads)) # Apply backward ops dx = grad_output @@ -311,7 +296,12 @@ def backward( for idx in basic_op_idxs: for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): if basic_op_extra_output_is_internal[idx][output_idx]: - basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get(channel) + channel_grad = channel_grads.get(channel) + if channel_grad is not None: + output_grad = basic_op_grad_extra_outputs[idx][output_idx] + basic_op_grad_extra_outputs[idx][output_idx] = ( + channel_grad if output_grad is None else output_grad + channel_grad + ) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs ] @@ -379,16 +369,11 @@ def backward( f"but got {len(dxs)}" ) - # One public tensor may fan out to several unmatched slots sharing a - # channel name, so sum their gradients before returning to autograd. - grad_extra_inputs_flat = [] - for slots in func_ctx.external_extra_input_slots: - grad = None - for op_idx, input_idx in slots: - slot_grad = grad_extra_inputs[op_idx][input_idx] - if slot_grad is not None: - grad = slot_grad if grad is None else grad + slot_grad - grad_extra_inputs_flat.append(grad) + # Collect the gradient for each public extra input. + grad_extra_inputs_flat = [ + grad_extra_inputs[op_idx][input_idx] + for op_idx, input_idx in func_ctx.external_extra_input_slots + ] # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -453,8 +438,7 @@ def __init__( self._basic_op_extra_output_is_internal: list[list[bool]] = [ [False] * op.num_extra_outputs for op in basic_ops ] - self._external_extra_input_slots: list[list[tuple[int, int]]] = [] - self._external_extra_output_slots: list[tuple[int, int]] = [] + self._external_extra_input_slots: list[tuple[int, int]] = [] # Find channel producers and reject ambiguous names. channel_producers: dict[str, tuple[int, int]] = {} @@ -470,24 +454,18 @@ def __init__( ) channel_producers[channel] = (op_idx, output_idx) - # Resolve inputs. A channel with an earlier producer is internal. A - # channel without a producer is public, with one positional tensor - # fanning out to every public slot that shares the channel name. - external_input_channels: dict[str, int] = {} + # Resolve inputs. A channel with an earlier producer is internal. + # Every input without an earlier producer is a separate public input. consumed_channels: set[str] = set() for op_idx, op in enumerate(basic_ops): for input_idx, channel in enumerate(op._extra_input_channels): if channel is None: - self._external_extra_input_slots.append([(op_idx, input_idx)]) + self._external_extra_input_slots.append((op_idx, input_idx)) continue producer = channel_producers.get(channel) + # If no producer for named channel, this is a public input. if producer is None: - group_idx = external_input_channels.get(channel) - if group_idx is None: - group_idx = len(self._external_extra_input_slots) - external_input_channels[channel] = group_idx - self._external_extra_input_slots.append([]) - self._external_extra_input_slots[group_idx].append((op_idx, input_idx)) + self._external_extra_input_slots.append((op_idx, input_idx)) continue producer_idx, _ = producer if producer_idx >= op_idx: @@ -498,13 +476,11 @@ def __init__( self._basic_op_extra_input_sources[op_idx][input_idx] = producer consumed_channels.add(channel) - # Unnamed outputs and named outputs without local consumers are public. + # All extra outputs remain public, including outputs consumed internally. for op_idx, op in enumerate(basic_ops): for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): if channel is not None and channel in consumed_channels: self._basic_op_extra_output_is_internal[op_idx][output_idx] = True - else: - self._external_extra_output_slots.append((op_idx, output_idx)) # Every channel-bound extra input must be wired to a matching producer # extra output. External slots remain unbound (source is None). @@ -718,9 +694,8 @@ def __call__( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in self._basic_ops ] - for tensor, slots in zip(extra_inputs, self._external_extra_input_slots): - for op_idx, input_idx in slots: - basic_op_extra_inputs[op_idx][input_idx] = tensor + for tensor, (op_idx, input_idx) in zip(extra_inputs, self._external_extra_input_slots): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state recipe = None diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index d668a7c904..25daf3a4be 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import dataclasses import pickle from typing import Any, Optional @@ -89,7 +89,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: """Forward pass This op is either a basic op or the fusion of basic ops, so @@ -118,7 +118,7 @@ def fuser_forward( ------- torch.Tensor: Output tensor. - Iterable of torch.Tensor: + Sequence of torch.Tensor: Extra tensor outputs from basic operations. """ @@ -198,14 +198,11 @@ def __init__(self) -> None: self._quantizers: Optional[dict[str, list[Quantizer]]] = None def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: - """Assign a channel name to an extra input slot. - - The slot receives the matching extra output from an earlier operation - in the same fuser. If there is no producer in the fuser, the slot - remains public; one public tensor fans out to all input slots with the - same channel name. Passing ``None`` removes the name. Channels cannot - be changed after the operation has been attached to a persistent - ``OperationFuser``. + """Bind an extra input slot to an internal fuser channel. + + A bound slot receives the matching extra output from an earlier + operation in the same fuser instead of consuming a public extra input. + Passing ``None`` removes the binding. """ if not 0 <= index < self.num_extra_inputs: raise IndexError( @@ -221,13 +218,10 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp return self def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: - """Assign a channel name to an extra output slot. + """Bind an extra output slot to an internal fuser channel. - The slot feeds matching extra inputs on later operations in the same - fuser. If there are no consumers in the fuser, the slot remains a - public extra output. Passing ``None`` removes the name. Channels cannot - be changed after the operation has been attached to a persistent - ``OperationFuser``. + A bound slot can feed one or more later operations and is not returned + as a public extra output. Passing ``None`` removes the binding. """ if not 0 <= index < self.num_extra_outputs: raise IndexError( From 5a4e1ec2b61cd345b12cc3c195c0bd65c3d41b68 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 20:06:45 +0000 Subject: [PATCH 16/83] update docs Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 35 +++++++++++++++++------------ 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index ffec6baa4e..0a285528f3 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -157,8 +157,11 @@ Extra tensor channels Extra inputs and 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``. Slots bound to channels are removed from the public -``Sequential`` interface. +``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``: @@ -182,9 +185,9 @@ With a channel, the residual block above can be expressed using one add_residual, ) - # The residual is routed internally, so the caller receives only y. + # The residual is routed internally and is also returned to the caller. x = torch.randn(16384, 4096, device="cuda") - y = block(x) + y, residual = block(x) Channels are also useful for mixture-of-experts blocks. The following example assumes custom ``Dispatch`` and ``Combine`` basic operations. @@ -225,8 +228,9 @@ routing map. ``Combine`` consumes the 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. - y = moe(x, 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 @@ -255,17 +259,20 @@ The following conditions apply to extra tensor channels: - A producer must appear before all of its consumers. Backward edges and cycles are not supported. -- A channel has exactly one producer, but its output may fan out to - multiple consumers. -- Every named output channel must have at least one consumer, and the - channel names on the producer and consumers must match. +- 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 have no channel assigned and - receives extra outputs that have no channel assigned. Slots assigned - to channels are internal and do not appear in the ``Sequential`` - arguments or return value. +- 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-connected basic operations may still be replaced by registered ``FusedOperation`` implementations. If a fused operation contains both From 0a479c737dbd7be2b9d2c18ed3c33bef9299bf35 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 21:04:48 +0000 Subject: [PATCH 17/83] pin channels through channel version Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 67 ++++++++++++++++++-- transformer_engine/pytorch/ops/fuser.py | 24 +++++-- transformer_engine/pytorch/ops/op.py | 22 ++----- transformer_engine/pytorch/ops/sequential.py | 8 +++ 4 files changed, 92 insertions(+), 29 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index dcac932db6..e415bbc886 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -665,18 +665,73 @@ def test_set_extra_channel_rejects_invalid_name(self) -> None: with pytest.raises(ValueError, match="non-empty string"): producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] - def test_set_extra_channel_rejects_mutation_after_fuser_construction(self) -> None: - """Channel routing is immutable after it has been captured by a fuser.""" + def test_extra_channel_change_invalidates_existing_fuser(self, size: int = 16) -> None: + """Channel changes invalidate an existing fuser but allow constructing another.""" producer = te_ops.MakeExtraOutput() consumer = te_ops.AddExtraInput() producer.set_extra_output_channel(0, "route") consumer.set_extra_input_channel(0, "route") fuser = OperationFuser([producer, consumer]) assert fuser.num_extra_inputs == 0 - with pytest.raises(RuntimeError, match="cannot be changed"): - producer.set_extra_output_channel(0, None) - with pytest.raises(RuntimeError, match="cannot be changed"): - consumer.set_extra_input_channel(0, None) + + consumer.set_extra_input_channel(0, None) + x = torch.rand((size,)) + extra = torch.rand_like(x) + with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): + fuser(x) + + new_fuser = OperationFuser([producer, consumer]) + y, route = new_fuser(x, extra) + torch.testing.assert_close(y, x + extra) + torch.testing.assert_close(route, x) + + def test_extra_channel_change_rebuilds_sequential(self, size: int = 16) -> None: + """Sequential rebuilds its routing after a channel configuration change.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer) + + x = torch.rand((size,)) + y, route = model(x) + torch.testing.assert_close(y, 2 * x) + torch.testing.assert_close(route, x) + + consumer.set_extra_input_channel(0, None) + extra = torch.rand_like(x) + y, route = model(x, extra) + torch.testing.assert_close(y, x + extra) + torch.testing.assert_close(route, x) + + @pytest.mark.parametrize("mutation", ("insert", "replace", "delete")) + def test_sequential_structure_change_does_not_leave_channel_locks( + self, + mutation: str, + size: int = 16, + ) -> None: + """Discarding cached fusers does not prevent channel reconfiguration.""" + producer = te_ops.MakeExtraOutput() + middle = te_ops.Identity() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, middle, consumer) + + x = torch.rand((size,)) + model(x) + if mutation == "insert": + model.insert(1, te_ops.Identity()) + elif mutation == "replace": + model[1] = te_ops.Identity() + else: + del model[1] + + consumer.set_extra_input_channel(0, None) + extra = torch.rand_like(x) + y, route = model(x, extra) + torch.testing.assert_close(y, x + extra) + torch.testing.assert_close(route, x) def test_duplicate_extra_output_channel_names(self) -> None: """Two extra outputs may not publish the same channel name.""" diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 3232f2d86e..f8b7085518 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -413,8 +413,6 @@ class OperationFuser: def __init__( self, ops: list[FusibleOperation], - *, - lock_extra_channels: bool = True, ) -> None: # Get list of basic operations @@ -426,6 +424,9 @@ def __init__( basic_ops.append(op) self._num_basic_ops: int = len(basic_ops) self._basic_ops: list[BasicOperation] = basic_ops + self._basic_op_extra_channels_versions = [ + op._extra_channels_version for op in self._basic_ops + ] # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) @@ -526,10 +527,15 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) - # Persistent fusers capture channel routing as structural state. - if lock_extra_channels: - for op in self._basic_ops: - op._lock_extra_channels() + def has_stale_op_channels(self) -> bool: + """Whether an operation's extra tensor channels have changed.""" + return any( + op._extra_channels_version != version + for op, version in zip( + self._basic_ops, + self._basic_op_extra_channels_versions, + ) + ) @staticmethod def _apply_fusions( @@ -679,6 +685,12 @@ def __call__( *extra_inputs: torch.Tensor, basic_op_kwargs: Optional[list[dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: + if self.has_stale_op_channels(): + raise RuntimeError( + "Extra tensor channels changed after this OperationFuser captured " + "its routing. Construct a new OperationFuser." + ) + # Verify extra input count if len(extra_inputs) != self.num_extra_inputs: raise ValueError( diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 25daf3a4be..85e5641078 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -191,7 +191,7 @@ def __init__(self) -> None: # Unbound slots remain public inputs/outputs, preserving the original API. self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs - self._extra_channels_locked = False + self._extra_channels_version = 0 # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None @@ -213,8 +213,8 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp raise ValueError("Extra input channel must be a non-empty string or None") if self._extra_input_channels[index] == channel: return self - self._assert_extra_channels_mutable() self._extra_input_channels[index] = channel + self._extra_channels_version += 1 return self def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: @@ -232,22 +232,10 @@ def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicO raise ValueError("Extra output channel must be a non-empty string or None") if self._extra_output_channels[index] == channel: return self - self._assert_extra_channels_mutable() self._extra_output_channels[index] = channel + self._extra_channels_version += 1 return self - def _assert_extra_channels_mutable(self) -> None: - """Check that channel routing has not been captured by a fuser.""" - if self._extra_channels_locked: - raise RuntimeError( - "Extra tensor channels cannot be changed after an operation has been " - "attached to an OperationFuser" - ) - - def _lock_extra_channels(self) -> None: - """Prevent changes after a fuser has captured the channel routing.""" - self._extra_channels_locked = True - @property def is_fused_op(self) -> bool: return False @@ -592,7 +580,7 @@ def forward( """Apply operation""" from .fuser import OperationFuser - return OperationFuser([self], lock_extra_channels=False)( + return OperationFuser([self])( input, *extra_inputs, basic_op_kwargs=[kwargs], @@ -827,7 +815,7 @@ def forward( basic_op_kwargs = [{} for _ in range(len(self.basic_ops))] from .fuser import OperationFuser - return OperationFuser([self], lock_extra_channels=False)( + return OperationFuser([self])( input, *extra_inputs, basic_op_kwargs=basic_op_kwargs, diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index cb5dfecb9f..9c3708ca96 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -179,6 +179,14 @@ def forward( or grouped MLP. """ + # Channel routing is captured when an OperationFuser is constructed. + # Rebuild groups if a contained operation's channel configuration changed. + if self._module_groups is not None and any( + isinstance(group, OperationFuser) and group.has_stale_op_channels() + for group in self._module_groups + ): + self._module_groups = None + # Create module groups if needed if self._module_groups is None: self._module_groups = self._make_module_groups(self._modules.values()) From d6799982ee0bcf9278dbf13ec87cc02eb4368802 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 21:24:33 +0000 Subject: [PATCH 18/83] unecessary handling removal Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fuser.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index f8b7085518..6ecc1963cc 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -171,11 +171,6 @@ def forward( f"but got {len(ys)}" ) for output_idx, y in enumerate(ys): - if y is None: - raise RuntimeError( - f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " - f"did not emit extra output {output_idx}" - ) if ( set_output_requires_grad and idx >= fuser.first_op_requiring_backward From 8f7ba95841caf8711389acccc0ada2100cde977e Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 21:37:26 +0000 Subject: [PATCH 19/83] simplify Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fuser.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 6ecc1963cc..6e0c063af6 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -293,9 +293,8 @@ def backward( if basic_op_extra_output_is_internal[idx][output_idx]: channel_grad = channel_grads.get(channel) if channel_grad is not None: - output_grad = basic_op_grad_extra_outputs[idx][output_idx] basic_op_grad_extra_outputs[idx][output_idx] = ( - channel_grad if output_grad is None else output_grad + channel_grad + basic_op_grad_extra_outputs[idx][output_idx] + channel_grad ) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs From c62bb157fac41433b9db401ecb2a2200d3d7ffbb Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 22:15:20 +0000 Subject: [PATCH 20/83] doc update + extra_grad = None case Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 28 +++++++++++----------- tests/pytorch/test_fusible_ops.py | 31 +++++++++++++++++++------ transformer_engine/pytorch/ops/fuser.py | 5 +++- 3 files changed, 42 insertions(+), 22 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index 0a285528f3..bb68c426fd 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -113,33 +113,33 @@ quantized compute. Branching operations ^^^^^^^^^^^^^^^^^^^^ -The operation fuser supports limited branching behavior. While the -operations must be in sequential order, some operations can accept +The operation fuser supports very limited branching behavior. While +the operations must be in sequential order, some operations can accept extra inputs or produce extra outputs. For example, ``AddExtraInput`` -adds an extra input tensor to the intermediate tensor, and -``MakeExtraOutput`` returns the intermediate tensor as an extra output. -When calling a ``Sequential`` that contains any of these branching -operations, the extra inputs should be passed as arguments and the -extra outputs will be returned after the main output. +will add an extra input tensor to the intermediate tensor and +``MakeExtraOutput`` will return the intermediate tensor as an extra +output. When calling a ``Sequential`` that contains any of these +branching operations, the extra inputs should be passed in as +arguments and the extra outputs will be returned. .. code-block:: python import torch import transformer_engine.pytorch as te - # Construct an MLP with a residual connection. + # Construct MLP with residual connection fc1 = te.ops.Sequential( te.ops.LayerNorm(4096), - te.ops.MakeExtraOutput(), # Output the residual. + te.ops.MakeExtraOutput(), # Output residual te.ops.Linear(4096, 28672), te.ops.SwiGLU(), ) fc2 = te.ops.Sequential( te.ops.Linear(14336, 4096), - te.ops.AddExtraInput(), # Add the residual. + te.ops.AddExtraInput(), # Add residual ) - # Pass the extra output from fc1 as the extra input to fc2. + # Forward pass x = torch.randn(16384, 4096, device="cuda") y, residual = fc1(x) y = fc2(y, residual) @@ -147,9 +147,9 @@ extra outputs will be returned after the main output. .. figure:: ./residual_layernorm_mlp.png :align: center - Operations for an MLP block with a residual connection. The block - is split into two sections so that the caller can pass the extra - output from the first section to the second. + Operations for an MLP block with a residual connection. Note that + the block has been split into two sections, each with one branching + operation. Extra tensor channels """"""""""""""""""""" diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index e415bbc886..ef584d3737 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -441,7 +441,12 @@ def test_extra_tensors(self, size: int = 16) -> None: class TestExtraTensorChannels: """Error handling and grad coverage for named extra-tensor channels.""" - def test_internal_residual_connection(self, size: int = 16) -> None: + @pytest.mark.parametrize("with_extra_grad", (True, False)) + def test_internal_residual_connection( + self, + with_extra_grad: bool, + size: int = 16, + ) -> None: """A channel can keep a residual connection inside a Sequential.""" residual = te_ops.MakeExtraOutput() body = te_ops.Bias(size=size, device="cpu") @@ -456,15 +461,22 @@ def test_internal_residual_connection(self, size: int = 16) -> None: torch.testing.assert_close(y, 2 * x + body.bias) torch.testing.assert_close(residual_out, x) dy = torch.rand_like(y) - dresidual = torch.rand_like(residual_out) - torch.autograd.backward((y, residual_out), (dy, dresidual)) - torch.testing.assert_close(x.grad, 2 * dy + dresidual) + if with_extra_grad: + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + expected_dx = 2 * dy + dresidual + else: + y.backward(dy) + expected_dx = 2 * dy + torch.testing.assert_close(x.grad, expected_dx) torch.testing.assert_close(body.bias.grad, dy) @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_fused_internal_residual_connection( self, fusion_kind: str, + with_extra_grad: bool, size: int = 16, ) -> None: """Forward, backward, and joint fusions can own an internal channel.""" @@ -554,9 +566,14 @@ def fuse_residual(ops, **unused): assert backward_ops[0][0] is forward_ops[0][0] torch.testing.assert_close(y, 2 * x + body.bias) dy = torch.rand_like(y) - dresidual = torch.rand_like(residual_out) - torch.autograd.backward((y, residual_out), (dy, dresidual)) - torch.testing.assert_close(x.grad, 2 * dy + dresidual) + if with_extra_grad: + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + expected_dx = 2 * dy + dresidual + else: + y.backward(dy) + expected_dx = 2 * dy + torch.testing.assert_close(x.grad, expected_dx) torch.testing.assert_close(body.bias.grad, dy) def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 6e0c063af6..16deeeccf0 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -293,8 +293,11 @@ def backward( if basic_op_extra_output_is_internal[idx][output_idx]: channel_grad = channel_grads.get(channel) if channel_grad is not None: + output_grad = basic_op_grad_extra_outputs[idx][output_idx] basic_op_grad_extra_outputs[idx][output_idx] = ( - basic_op_grad_extra_outputs[idx][output_idx] + channel_grad + channel_grad + if output_grad is None + else output_grad + channel_grad ) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs From a93b82029f0f13acb5c9f6efdc6dead7109932b3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:18:01 +0000 Subject: [PATCH 21/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/fuser.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 16deeeccf0..6ecc1963cc 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -295,9 +295,7 @@ def backward( if channel_grad is not None: output_grad = basic_op_grad_extra_outputs[idx][output_idx] basic_op_grad_extra_outputs[idx][output_idx] = ( - channel_grad - if output_grad is None - else output_grad + channel_grad + channel_grad if output_grad is None else output_grad + channel_grad ) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs From 35b73b1601bdeb3d7cefc6960ae0ba2dc235a5d9 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 22:29:53 +0000 Subject: [PATCH 22/83] test cleanup Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 112 +++++++++++------------- transformer_engine/pytorch/ops/fuser.py | 3 +- 2 files changed, 50 insertions(+), 65 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index ef584d3737..d6d2dbae82 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -438,6 +438,38 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x4, x4_orig + x3) +class _DualExtraOutput(te_ops.BasicOperation): + """Test helper: one op with two scaled extra outputs.""" + + num_extra_outputs = 2 + + def __init__(self, scales: tuple[float, float] = (1.0, 1.0)) -> None: + super().__init__() + self._scales = scales + + def op_forward(self, *args, **kwargs): + raise RuntimeError("_DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("_DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_ctxs, basic_op_extra_inputs + s0, s1 = self._scales + return input_, [(s0 * input_, s1 * input_)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + s0, s1 = self._scales + g0, g1 = basic_op_grad_extra_outputs[0] + grad_extra = torch.zeros_like(grad_output) + if g0 is not None: + grad_extra = grad_extra + s0 * g0 + if g1 is not None: + grad_extra = grad_extra + s1 * g1 + return grad_output + grad_extra, [()], [()] + + class TestExtraTensorChannels: """Error handling and grad coverage for named extra-tensor channels.""" @@ -530,6 +562,7 @@ def fuse_residual(ops, **unused): and isinstance(ops[1], te_ops.Bias) and isinstance(ops[2], te_ops.AddExtraInput) ): + # We want to enable this fusion just for this test. FusedResidual._enabled = False return [FusedResidual(*ops)] return ops @@ -750,48 +783,24 @@ def test_sequential_structure_change_does_not_leave_channel_locks( torch.testing.assert_close(y, x + extra) torch.testing.assert_close(route, x) - def test_duplicate_extra_output_channel_names(self) -> None: - """Two extra outputs may not publish the same channel name.""" - producer1 = te_ops.MakeExtraOutput() - producer2 = te_ops.MakeExtraOutput() + @pytest.mark.parametrize("layout", ("two_ops", "same_op")) + def test_duplicate_extra_output_channel_names(self, layout: str) -> None: + """A channel name may have at most one producer, across ops or slots.""" consumer = te_ops.AddExtraInput() - producer1.set_extra_output_channel(0, "route") - producer2.set_extra_output_channel(0, "route") - consumer.set_extra_input_channel(0, "route") - with pytest.raises(ValueError, match="multiple producers"): - OperationFuser([producer1, producer2, consumer]) - - def test_duplicate_extra_output_channels_on_same_op(self) -> None: - """A single op with multiple extras still cannot reuse a channel name.""" - - class DualExtraOutput(te_ops.BasicOperation): - num_extra_outputs = 2 - - def op_forward(self, *args, **kwargs): - raise RuntimeError("DualExtraOutput uses fuser_forward") - - def op_backward(self, *args, **kwargs): - raise RuntimeError("DualExtraOutput uses fuser_backward") - - def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): - return input_, [(input_, input_)] - - def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): - g0, g1 = basic_op_grad_extra_outputs[0] - grad_extra = torch.zeros_like(grad_output) - if g0 is not None: - grad_extra = grad_extra + g0 - if g1 is not None: - grad_extra = grad_extra + g1 - return grad_output + grad_extra, [()], [()] - - producer = DualExtraOutput() - consumer = te_ops.AddExtraInput() - producer.set_extra_output_channel(0, "route") - producer.set_extra_output_channel(1, "route") consumer.set_extra_input_channel(0, "route") + if layout == "two_ops": + producer1 = te_ops.MakeExtraOutput() + producer2 = te_ops.MakeExtraOutput() + producer1.set_extra_output_channel(0, "route") + producer2.set_extra_output_channel(0, "route") + ops = [producer1, producer2, consumer] + else: + producer = _DualExtraOutput() + producer.set_extra_output_channel(0, "route") + producer.set_extra_output_channel(1, "route") + ops = [producer, consumer] with pytest.raises(ValueError, match="multiple producers"): - OperationFuser([producer, consumer]) + OperationFuser(ops) def test_named_extra_output_without_consumer_is_public(self, size: int = 16) -> None: """A named output remains public when its fuser has no consumer.""" @@ -823,30 +832,7 @@ def test_one_extra_input_has_single_source(self, size: int = 16) -> None: def test_mixed_channel_outputs_are_public(self, size: int = 16) -> None: """Both internally consumed and unconsumed channel outputs are public.""" - - class DualExtraOutput(te_ops.BasicOperation): - num_extra_outputs = 2 - - def op_forward(self, *args, **kwargs): - raise RuntimeError("DualExtraOutput uses fuser_forward") - - def op_backward(self, *args, **kwargs): - raise RuntimeError("DualExtraOutput uses fuser_backward") - - def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): - del basic_op_ctxs, basic_op_extra_inputs - return input_, [(2 * input_, 3 * input_)] - - def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): - del basic_op_ctxs - grad_internal, grad_public = basic_op_grad_extra_outputs[0] - return ( - grad_output + 2 * grad_internal + 3 * grad_public, - [()], - [()], - ) - - producer = DualExtraOutput() + producer = _DualExtraOutput(scales=(2.0, 3.0)) consumer = te_ops.AddExtraInput() producer.set_extra_output_channel(0, "internal") producer.set_extra_output_channel(1, "public") diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 6ecc1963cc..e3b342521c 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -179,8 +179,7 @@ def forward( y.requires_grad_(True) extra_outputs[idx] = ys - # Flatten all extra outputs in basic-operation order. Channel-bound - # outputs remain visible even when they are also consumed internally. + # Flatten list of extra outputs extra_outputs_flat = [y for ys in extra_outputs for y in ys] # Save context for backward pass From 6801a6d2e1d6fed61410119dd27c9640665288e6 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 22:46:34 +0000 Subject: [PATCH 23/83] no need to check staleness in every forward call Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 4 ++++ tests/pytorch/test_fusible_ops.py | 8 ++++++-- transformer_engine/pytorch/ops/sequential.py | 8 -------- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index bb68c426fd..c8f1c1281e 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -273,6 +273,10 @@ The following conditions apply to extra tensor channels: 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 diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index d6d2dbae82..944171b402 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -735,8 +735,8 @@ def test_extra_channel_change_invalidates_existing_fuser(self, size: int = 16) - torch.testing.assert_close(y, x + extra) torch.testing.assert_close(route, x) - def test_extra_channel_change_rebuilds_sequential(self, size: int = 16) -> None: - """Sequential rebuilds its routing after a channel configuration change.""" + def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> None: + """Sequential does not auto-rebuild after a channel configuration change.""" producer = te_ops.MakeExtraOutput() consumer = te_ops.AddExtraInput() producer.set_extra_output_channel(0, "route") @@ -750,6 +750,10 @@ def test_extra_channel_change_rebuilds_sequential(self, size: int = 16) -> None: consumer.set_extra_input_channel(0, None) extra = torch.rand_like(x) + with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): + model(x, extra) + + model = te_ops.Sequential(producer, consumer) y, route = model(x, extra) torch.testing.assert_close(y, x + extra) torch.testing.assert_close(route, x) diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index 9c3708ca96..cb5dfecb9f 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -179,14 +179,6 @@ def forward( or grouped MLP. """ - # Channel routing is captured when an OperationFuser is constructed. - # Rebuild groups if a contained operation's channel configuration changed. - if self._module_groups is not None and any( - isinstance(group, OperationFuser) and group.has_stale_op_channels() - for group in self._module_groups - ): - self._module_groups = None - # Create module groups if needed if self._module_groups is None: self._module_groups = self._make_module_groups(self._modules.values()) From 6688e8ab1d467b17e4a625c61cc35d51d34b3fb7 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 23:08:21 +0000 Subject: [PATCH 24/83] remove redundant tests Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 138 ++---------------------------- 1 file changed, 7 insertions(+), 131 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 944171b402..4bad009c2e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -667,18 +667,16 @@ def test_external_named_extra_inputs_remain_separate(self, size: int = 16) -> No model = te_ops.Sequential(consumer1, consumer2) x = torch.rand((size,), requires_grad=True) - extra1 = torch.rand((size,), requires_grad=True) - extra2 = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) with pytest.raises(ValueError, match="Expected 2 extra inputs but got 1"): - model(x, extra1) - y = model(x, extra1, extra2) - torch.testing.assert_close(y, x + extra1 + extra2) + model(x, extra) + y = model(x, extra, extra) + torch.testing.assert_close(y, x + 2 * extra) dy = torch.rand_like(y) y.backward(dy) torch.testing.assert_close(x.grad, dy) - torch.testing.assert_close(extra1.grad, dy) - torch.testing.assert_close(extra2.grad, dy) + torch.testing.assert_close(extra.grad, 2 * dy) def test_consumer_before_producer(self) -> None: """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" @@ -715,26 +713,6 @@ def test_set_extra_channel_rejects_invalid_name(self) -> None: with pytest.raises(ValueError, match="non-empty string"): producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] - def test_extra_channel_change_invalidates_existing_fuser(self, size: int = 16) -> None: - """Channel changes invalidate an existing fuser but allow constructing another.""" - producer = te_ops.MakeExtraOutput() - consumer = te_ops.AddExtraInput() - producer.set_extra_output_channel(0, "route") - consumer.set_extra_input_channel(0, "route") - fuser = OperationFuser([producer, consumer]) - assert fuser.num_extra_inputs == 0 - - consumer.set_extra_input_channel(0, None) - x = torch.rand((size,)) - extra = torch.rand_like(x) - with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): - fuser(x) - - new_fuser = OperationFuser([producer, consumer]) - y, route = new_fuser(x, extra) - torch.testing.assert_close(y, x + extra) - torch.testing.assert_close(route, x) - def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> None: """Sequential does not auto-rebuild after a channel configuration change.""" producer = te_ops.MakeExtraOutput() @@ -758,35 +736,6 @@ def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> N torch.testing.assert_close(y, x + extra) torch.testing.assert_close(route, x) - @pytest.mark.parametrize("mutation", ("insert", "replace", "delete")) - def test_sequential_structure_change_does_not_leave_channel_locks( - self, - mutation: str, - size: int = 16, - ) -> None: - """Discarding cached fusers does not prevent channel reconfiguration.""" - producer = te_ops.MakeExtraOutput() - middle = te_ops.Identity() - consumer = te_ops.AddExtraInput() - producer.set_extra_output_channel(0, "route") - consumer.set_extra_input_channel(0, "route") - model = te_ops.Sequential(producer, middle, consumer) - - x = torch.rand((size,)) - model(x) - if mutation == "insert": - model.insert(1, te_ops.Identity()) - elif mutation == "replace": - model[1] = te_ops.Identity() - else: - del model[1] - - consumer.set_extra_input_channel(0, None) - extra = torch.rand_like(x) - y, route = model(x, extra) - torch.testing.assert_close(y, x + extra) - torch.testing.assert_close(route, x) - @pytest.mark.parametrize("layout", ("two_ops", "same_op")) def test_duplicate_extra_output_channel_names(self, layout: str) -> None: """A channel name may have at most one producer, across ops or slots.""" @@ -857,7 +806,8 @@ def test_mixed_channel_outputs_are_public(self, size: int = 16) -> None: def test_fresh_internal_output_preserves_grad_requirement(self) -> None: """A fresh internal tensor requests its gradient from a scaled activation.""" - + # A BasicOperation with one extra output that is freshly computed instead of + # retrieved from a previous op's tensor. class MakeScale(te_ops.BasicOperation): num_extra_outputs = 1 @@ -897,80 +847,6 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp y_test.backward(dy) torch.testing.assert_close(x_test.grad, x_ref.grad) - def test_grouped_linear_scale_bias_channels(self) -> None: - """Both GroupedLinear extra inputs can be supplied by channels.""" - - class RouteExtras(te_ops.BasicOperation): - num_extra_inputs = 2 - num_extra_outputs = 2 - - def op_forward(self, *args, **kwargs): - raise RuntimeError("RouteExtras uses fuser_forward") - - def op_backward(self, *args, **kwargs): - raise RuntimeError("RouteExtras uses fuser_backward") - - def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): - del basic_op_ctxs - return input_, [basic_op_extra_inputs[0]] - - def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): - del basic_op_ctxs - return grad_output, [()], [basic_op_grad_extra_outputs[0]] - - group_size, in_features, out_features = 2, 8, 6 - split_sizes = torch.tensor((3, 2), dtype=torch.int32, device="cuda") - num_tokens = int(split_sizes.sum()) - x = torch.randn((num_tokens, in_features), device="cuda", requires_grad=True) - scales = torch.randn((num_tokens,), device="cuda", requires_grad=True) - - producer = RouteExtras() - linear = te_ops.GroupedLinear( - group_size, - in_features, - out_features, - bias=True, - scale_bias=True, - device="cuda", - dtype=torch.float32, - ) - producer.set_extra_output_channel(0, "split_sizes") - producer.set_extra_output_channel(1, "bias_scales") - linear.set_extra_input_channel(0, "split_sizes") - linear.set_extra_input_channel(1, "bias_scales") - model = te_ops.Sequential(producer, linear) - - x_ref = x.detach().clone().requires_grad_(True) - scales_ref = scales.detach().clone().requires_grad_(True) - ys_ref = [] - for group_idx, (x_group, scale_group) in enumerate( - zip( - torch.split(x_ref, split_sizes.tolist()), - torch.split(scales_ref, split_sizes.tolist()), - ) - ): - weight = getattr(linear, f"weight{group_idx}") - bias = getattr(linear, f"bias{group_idx}") - ys_ref.append( - torch.nn.functional.linear(x_group, weight) + scale_group.unsqueeze(-1) * bias - ) - y_ref = torch.cat(ys_ref) - y_test, _split_sizes_test, _scales_test = model(x, split_sizes, scales) - dy = torch.rand_like(y_test) - grads_ref = torch.autograd.grad( - y_ref, - (x_ref, scales_ref, *linear.parameters()), - dy, - ) - y_test.backward(dy) - - tols = dtype_tols(torch.float16) # Grouped GEMM uses TF32 for FP32 inputs. - torch.testing.assert_close(y_test, y_ref, **tols) - torch.testing.assert_close(x.grad, grads_ref[0], **tols) - torch.testing.assert_close(scales.grad, grads_ref[1], **tols) - for param, grad_ref in zip(linear.parameters(), grads_ref[2:]): - torch.testing.assert_close(param.grad, grad_ref, **tols) - class TestFuser: """Tests for operation fusion infrastructure""" From 12430c286cbc18f4b1351c9d413d8db28414ca74 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:09:43 +0000 Subject: [PATCH 25/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 4bad009c2e..b50c9c3fc1 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -806,7 +806,8 @@ def test_mixed_channel_outputs_are_public(self, size: int = 16) -> None: def test_fresh_internal_output_preserves_grad_requirement(self) -> None: """A fresh internal tensor requests its gradient from a scaled activation.""" - # A BasicOperation with one extra output that is freshly computed instead of + + # A BasicOperation with one extra output that is freshly computed instead of # retrieved from a previous op's tensor. class MakeScale(te_ops.BasicOperation): num_extra_outputs = 1 From b1895500a0bbc84ba906cf18431e3b264d2b3a9f Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 23:26:20 +0000 Subject: [PATCH 26/83] revert from bad names Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fuser.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index e3b342521c..1a6ee8e670 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -205,8 +205,7 @@ def forward( if fuser.first_op_requiring_backward < fuser._num_basic_ops: is_first_module = FP8GlobalStateManager.is_first_fp8_module() - # Other context. Save only the wiring metadata needed by - # backward instead of the whole OperationFuser. + # Other context func_ctx.backward_ops = fuser._backward_ops func_ctx.basic_ops = fuser._basic_ops func_ctx.basic_op_ctxs = basic_op_ctxs @@ -264,11 +263,10 @@ def backward( f"Expected grads for {func_ctx.num_extra_outputs} extra tensor outputs, " f"but got {len(grad_extra_outputs)}" ) - remaining_grad_extra_outputs = grad_extra_outputs basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [] for op in basic_ops: - grads, remaining_grad_extra_outputs = _split_tuple( - remaining_grad_extra_outputs, + grads, grad_extra_outputs = _split_tuple( + grad_extra_outputs, op.num_extra_outputs, ) basic_op_grad_extra_outputs.append(list(grads)) From a4cc11246a9b0f804d8f480a565b8dc0fc62bb4b Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 23:31:36 +0000 Subject: [PATCH 27/83] keep simple Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fuser.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 1a6ee8e670..f0d8e7d52a 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -294,13 +294,11 @@ def backward( basic_op_grad_extra_outputs[idx][output_idx] = ( channel_grad if output_grad is None else output_grad + channel_grad ) - op_grad_extra_outputs = [ - tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs - ] + grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( [basic_op_ctxs[idx] for idx in basic_op_idxs], dx, - basic_op_grad_extra_outputs=op_grad_extra_outputs, + basic_op_grad_extra_outputs=grad_extra_outputs, ) fused_op_grad_params = tuple(tuple(grads) for grads in fused_op_grad_params) fused_op_grad_extra_inputs = tuple(tuple(grads) for grads in fused_op_grad_extra_inputs) From 5ba605505b885f2b9a241ddb19bd0a21e1654632 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 23:37:21 +0000 Subject: [PATCH 28/83] unecessary checks Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fuser.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index f0d8e7d52a..cf1df530e1 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -300,20 +300,6 @@ def backward( dx, basic_op_grad_extra_outputs=grad_extra_outputs, ) - fused_op_grad_params = tuple(tuple(grads) for grads in fused_op_grad_params) - fused_op_grad_extra_inputs = tuple(tuple(grads) for grads in fused_op_grad_extra_inputs) - if len(fused_op_grad_params) != len(basic_op_idxs): - raise RuntimeError( - f"Expected {type(op).__name__} to generate parameter grads for " - f"{len(basic_op_idxs)} basic operations, but got " - f"{len(fused_op_grad_params)}" - ) - if len(fused_op_grad_extra_inputs) != len(basic_op_idxs): - raise RuntimeError( - f"Expected {type(op).__name__} to generate extra-input grads for " - f"{len(basic_op_idxs)} basic operations, but got " - f"{len(fused_op_grad_extra_inputs)}" - ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None From 76826dc2bcf7c9867a6c9ecbb1a6277086668d1c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 9 Aug 2026 23:43:21 +0000 Subject: [PATCH 29/83] minor doc Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index c8f1c1281e..123fd1cd76 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -154,7 +154,7 @@ arguments and the extra outputs will be returned. Extra tensor channels """"""""""""""""""""" -Extra inputs and outputs may optionally specify a channel. Assigning +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 From 60ca03051c0207f23b5f3bb4737e1818509b68d5 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 10 Aug 2026 06:27:45 +0000 Subject: [PATCH 30/83] basic op and refrence implementation Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/moe_ep_reference.py | 842 ++++++++++++++++++ .../pytorch/ops/basic/combine.py | 138 +++ .../pytorch/ops/basic/dispatch.py | 207 +++++ 3 files changed, 1187 insertions(+) create mode 100644 tests/pytorch/distributed/moe_ep_reference.py create mode 100644 transformer_engine/pytorch/ops/basic/combine.py create mode 100644 transformer_engine/pytorch/ops/basic/dispatch.py diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/tests/pytorch/distributed/moe_ep_reference.py new file mode 100644 index 0000000000..280c4f0d21 --- /dev/null +++ b/tests/pytorch/distributed/moe_ep_reference.py @@ -0,0 +1,842 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. + +The implementation deliberately favors readable semantics over performance. It +supports a one-rank execution path and a variable-size ``all_to_all_single`` EP +path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. + +The quantized tensor layouts are logical (unswizzled) layouts. A production +kernel may reorder scale factors internally without changing this API contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +class MoeFormat(str, Enum): + """Public and communication formats supported by the reference.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = value + return tuple(result) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Portable data-plus-scale representation for MXFP8 or NVFP4. + + ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` + has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with + two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` + replaces that axis by one scale per block. + """ + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + fmt = _parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + shape = tuple(int(dim) for dim in self.logical_shape) + if not shape or any(dim < 0 for dim in shape): + raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") + axis = _normalize_axis(self.axis, len(shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", shape) + object.__setattr__(self, "axis", axis) + self._validate_storage() + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + def _validate_storage(self) -> None: + if self.data.device != self.scale.device: + raise ValueError("block-scaled data and scale must be on the same device") + + logical_extent = self.logical_shape[self.axis] + scale_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, self.block_size), + ) + if tuple(self.scale.shape) != scale_shape: + raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") + + if self.format is MoeFormat.MXFP8: + expected_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + data_shape = self.logical_shape + if self.data.dtype != expected_dtype: + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) + else: + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: + raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") + data_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, 2), + ) + + if tuple(self.data.shape) != data_shape: + raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") + if self.scale.dtype != expected_scale_dtype: + raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return the logical tensor with block scales applied.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + packed = packed.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" + + levels = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=values.device, + ) + magnitudes = values.abs().unsqueeze(-1) + distances = (magnitudes - levels).abs() + minimum = distances.amin(dim=-1, keepdim=True) + candidates = distances == minimum + codes = torch.arange(8, dtype=torch.int64, device=values.device) + any_code = torch.where(candidates, codes, 8).amin(dim=-1) + even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) + magnitude_code = torch.where(even_code < 8, even_code, any_code) + sign_code = torch.signbit(values).to(torch.int64) << 3 + return magnitude_code | sign_code + + +def quantize_blockwise( + tensor: torch.Tensor, + format: Union[MoeFormat, str], + *, + axis: int = -1, +) -> BlockScaledTensor: + """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. + + MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward + positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and + E4M3 scales rounded to nearest. + """ + + fmt = _parse_format(format) + if fmt is MoeFormat.BF16: + raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") + if not tensor.is_floating_point(): + raise TypeError(f"tensor must be floating point, got {tensor.dtype}") + + axis = _normalize_axis(axis, tensor.ndim) + logical_shape = tuple(tensor.shape) + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + block_count = _ceil_div(logical_extent, block_size) + padded_extent = block_count * block_size + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) + + value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 + scale_float = blocks.abs().amax(dim=-1) / value_limit + if fmt is MoeFormat.MXFP8: + safe_scale = torch.where(scale_float > 0, scale_float, 1.0) + scale_float = torch.where( + scale_float > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(scale_float), + ) + scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + scale_dtype = _require_torch_dtype("float8_e4m3fn") + + scale = scale_float.to(scale_dtype) + scale_for_math = scale.float() + reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) + + if fmt is MoeFormat.MXFP8: + data_dtype = _require_torch_dtype("float8_e4m3fn") + data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] + else: + codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) + low = codes[..., 0::2] + high = codes[..., 1::2] + data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] + + return BlockScaledTensor( + data=data.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format=fmt, + logical_shape=logical_shape, + axis=axis, + ) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +@dataclass(frozen=True) +class _DispatchPlan: + """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" + + send_expert: torch.Tensor # local expert id per sent route + send_weight: torch.Tensor # router weight per sent route + send_token_idx: torch.Tensor # source token per sent route + send_slot_idx: torch.Tensor # source top-k slot per sent route + send_counts: Tuple[int, ...] # routes sent to each rank + recv_counts: Tuple[int, ...] # routes received from each rank + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +def _decode_tensor( + tensor: MoeTensor, + *, + name: str, + expected_shape: Tuple[int, ...], + quantized_axis: int, +) -> torch.Tensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.logical_shape != expected_shape: + raise ValueError( + f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" + ) + if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): + raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") + return tensor.dequantize() + + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") + if not tensor.is_floating_point(): + raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") + return tensor.float() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=-1).dequantize() + + +class MoeEpReference: + """Reference implementation of routed SwiGLU experts plus EP dispatch. + + Global experts are assigned contiguously: rank ``r`` owns + ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit + initialized process group for multi-rank execution; ``None`` means a + one-rank reference even if the default distributed group is initialized. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + ) -> None: + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + + if ep_group is None: + ep_size, ep_rank = 1, 0 + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "ep_group requires an initialized torch.distributed process group" + ) + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size != 0: + raise ValueError( + f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" + ) + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = ( + 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + ) + if hidden_size % required_multiple != 0: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" + f" {name}={fmt.value}" + ) + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(" + f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " + f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " + f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " + f"output={self.output_format.value}, combine={self.combine_format.value})" + ) + + def _collective_device(self, device: torch.device) -> torch.device: + """Device the process group can run ``all_to_all_single`` on. + + Gloo only implements all-to-all for CPU tensors, so CUDA tensors are + staged through host memory; NCCL groups communicate in place. + """ + if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return send_counts.clone() + comm_device = self._collective_device(send_counts.device) + staged = send_counts.to(comm_device) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single(recv_counts, staged, group=self.ep_group) + return recv_counts.to(send_counts.device) + + def _all_to_all( + self, + send: torch.Tensor, + send_counts: Sequence[int], + recv_counts: Sequence[int], + ) -> torch.Tensor: + if self.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + output_shape = (sum(recv_counts), *send.shape[1:]) + recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) + dist.all_to_all_single( + recv, + staged, + output_split_sizes=list(recv_counts), + input_split_sizes=list(send_counts), + group=self.ep_group, + ) + return recv.to(send.device) + + def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: + """Route valid ``topk_idx`` entries to destination ranks, stably by rank. + + Backward reuses this so gradient re-dispatch reproduces the exact + forward route order. + """ + + device = topk_idx.device + token_count = topk_idx.shape[0] + flat_expert = topk_idx.reshape(-1).to(torch.int64) + flat_weight = topk_weights.reshape(-1).float() + valid = flat_expert != -1 + invalid_negative = flat_expert < -1 + invalid_high = flat_expert >= self.num_experts + if bool((invalid_negative | invalid_high).any().item()): + bad = flat_expert[invalid_negative | invalid_high][0].item() + raise ValueError(f"topk_idx contains out-of-range expert id {bad}") + + flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) + flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) + expert = flat_expert[valid] + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + order = torch.argsort(destination, stable=True) + + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), + send_weight=flat_weight[valid].index_select(0, order), + send_token_idx=flat_token[valid].index_select(0, order), + send_slot_idx=flat_slot[valid].index_select(0, order), + send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), + recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), + ) + + def _run_local_experts( + self, + tokens: torch.Tensor, + local_expert_idx: torch.Tensor, + route_weight: torch.Tensor, + fc1_weight: torch.Tensor, + fc2_weight: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output = torch.empty( + (tokens.shape[0], self.hidden_size), + dtype=torch.float32, + device=tokens.device, + ) + fc1_c_rows = [] if self.generate_c else None + for expert in range(self.experts_per_rank): + positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + expert_tokens = tokens.index_select(0, positions) + gate_up = expert_tokens @ fc1_weight[expert] + if fc1_c_rows is not None: + # Raw pre-SwiGLU accumulator: before clamp, no router weight. + fc1_c_rows.append(gate_up.to(torch.bfloat16)) + gate, up = gate_up.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + gate = gate.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up + weights = route_weight.index_select(0, positions).unsqueeze(-1) + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + expert_output = intermediate @ fc2_weight[expert] + if not self.apply_topk_in_fc1: + expert_output = expert_output * weights + expert_output = _format_round_trip(expert_output, self.combine_format) + output.index_copy_(0, positions, expert_output) + fc1_c = None + if fc1_c_rows is not None: + fc1_c = ( + torch.cat(fc1_c_rows) + if fc1_c_rows + else torch.empty( + (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device + ) + ) + return output, fc1_c + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: + """Run dispatch, local experts, return routing, top-k reduce, and encode. + + Shapes: + activation: ``(T, H)`` + fc1_weight: ``(E_local, H, 2 * I)`` + fc2_weight: ``(E_local, I, H)`` + topk_idx/topk_weights: ``(T, K)`` + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when constructed with ``generate_c=True``. ``fc1_c`` is the BF16 + pre-SwiGLU FC1 accumulator of every route this rank's experts + processed, ``(local_routes, 2 * I)``, grouped by local expert and + ordered within each expert by (source rank, source token-major route + order); captured before the gate/up clamp, without the router weight. + ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies + the route behind ``fc1_c`` row ``i`` for the backward gradient + re-dispatch. + """ + + if topk_idx.ndim != 2: + raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") + token_count = topk_idx.shape[0] + route_shape = (token_count, self.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) + + device = _tensor_device(activation) + inputs = { + "fc1_weight": _tensor_device(fc1_weight), + "fc2_weight": _tensor_device(fc2_weight), + "topk_idx": topk_idx.device, + "topk_weights": topk_weights.device, + } + for name, input_device in inputs.items(): + if input_device != device: + raise ValueError(f"{name} must be on {device}, got {input_device}") + + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + + plan = self._dispatch_plan(topk_idx, topk_weights) + send_token_idx = plan.send_token_idx + send_slot_idx = plan.send_slot_idx + send_counts, recv_counts = plan.send_counts, plan.recv_counts + send_tokens = activation_float.index_select(0, send_token_idx) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + + route_metadata = None + if self.generate_c: + recv_src_rank = torch.repeat_interleave( + torch.arange(self.ep_size, device=device), + torch.tensor(recv_counts, device=device), + ) + recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) + recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) + # Stable sort by local expert reproduces the fc1_c row order + # (grouped by expert; source order preserved within each group). + fc1_c_order = torch.argsort(recv_expert, stable=True) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) + # recv rows are ordered by source rank, then that source's token-major + # route order, so the per-expert position grouping below realizes the + # documented fc1_c ordering. + recv_output, fc1_c = self._run_local_experts( + recv_tokens, + recv_expert, + recv_weight, + fc1_float, + fc2_float, + ) + + returned = self._all_to_all(recv_output, recv_counts, send_counts) + combine_plane = torch.zeros( + (token_count * self.top_k, self.hidden_size), + dtype=torch.float32, + device=device, + ) + send_flat_slot = send_token_idx * self.top_k + send_slot_idx + combine_plane.index_copy_(0, send_flat_slot, returned) + reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) + + if self.output_format is MoeFormat.BF16: + output = reduced.to(torch.bfloat16) + else: + output = quantize_blockwise(reduced, self.output_format, axis=-1) + if self.generate_c: + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward pass consuming the ``generate_c=True`` stash. + + ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, + and the FC2 input are all rebuilt from it, so no post-SwiGLU forward + intermediate needs to be saved. ``route_metadata`` alone reconstructs + the mapping between re-dispatched rows and ``fc1_c`` rows and drives + the gradient return scatter. + + Quantization round-trips (input decode, ``combine_format``, + ``output_format``) are treated as straight-through identities; + ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. + + Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, + grad_topk_weights)`` in float32. + """ + + if not self.generate_c: + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) + token_count = topk_idx.shape[0] + if tuple(grad_output.shape) != (token_count, self.hidden_size): + raise ValueError( + f"grad_output shape must be {(token_count, self.hidden_size)}, got" + f" {tuple(grad_output.shape)}" + ) + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + + device = _tensor_device(activation) + two_i = 2 * self.intermediate_size + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + if fc1_c.shape != (int(route_metadata.shape[0]), two_i): + raise ValueError( + f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" + f" {tuple(fc1_c.shape)}" + ) + + # Re-dispatch the FC1 inputs, router weights, and output gradients + # along the identical forward routes. + plan = self._dispatch_plan(topk_idx, topk_weights) + send_counts, recv_counts = plan.send_counts, plan.recv_counts + grad_output_float = grad_output.float() + recv_tokens = self._all_to_all( + activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_grad = self._all_to_all( + grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + + # route_metadata rows are in fc1_c order; sorting them by + # (src_rank, src_token, src_slot) reproduces the receive order, giving + # the permutation between re-dispatched rows and fc1_c rows. + metadata = route_metadata.to(device=device, dtype=torch.int64) + local_routes = metadata.shape[0] + if local_routes > 0: + token_span = int(metadata[:, 2].max().item()) + 1 + recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] + perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j + else: + perm = torch.empty((0,), dtype=torch.int64, device=device) + x_rows = torch.empty_like(recv_tokens) + x_rows.index_copy_(0, perm, recv_tokens) + w_rows = torch.empty_like(recv_weight) + w_rows.index_copy_(0, perm, recv_weight) + dy_rows = torch.empty_like(recv_grad) + dy_rows.index_copy_(0, perm, recv_grad) + + c_rows = fc1_c.float() + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + grad_fc1 = torch.zeros_like(fc1_float) + grad_fc2 = torch.zeros_like(fc2_float) + for expert in range(self.experts_per_rank): + positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + c = c_rows.index_select(0, positions) + x = x_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = dy_rows.index_select(0, positions) + + gate, up = c.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + g = gate.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate, up + sig = torch.sigmoid(g) + s = g * sig + h = s * u + + if self.apply_topk_in_fc1: + h_fc2 = h * w + d_y_pre = d_y + else: + h_fc2 = h + d_y_pre = d_y * w + grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) + + d_g = d_h * u * (sig * (1 + g * (1 - sig))) + d_u = d_h * s + if self.gate_up_clamp is not None: + d_gate = d_g * (gate <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1) + grad_fc1[expert] = x.transpose(0, 1) @ d_c + d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) + + # Return the route gradients to their source ranks and scatter-add. + returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + grad_activation = torch.zeros( + (token_count, self.hidden_size), dtype=torch.float32, device=device + ) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_topk_weights = torch.zeros( + (token_count * self.top_k,), dtype=torch.float32, device=device + ) + grad_topk_weights.index_copy_( + 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw + ) + return ( + grad_activation, + grad_fc1, + grad_fc2, + grad_topk_weights.view(token_count, self.top_k), + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "quantize_blockwise", +] diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py new file mode 100644 index 0000000000..4d4740c517 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel combine operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, is_symm_backed +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_grad_buffer( + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError("grad_out must be contiguous.") + if tensor.requires_grad: + raise ValueError("grad_out must not require gradients.") + return tensor + + +class Combine(BasicOperation): + """Combine pre-weighted local expert outputs with NCCL EP. + + The operation uses routing state produced by a :class:`Dispatch` with the + same :class:`EpBuffer`. + """ + + def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: + super().__init__() + self.buffer = buffer + self.num_local_tokens = ( + buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) + ) + if self.num_local_tokens < 0: + raise ValueError("num_local_tokens must be non-negative.") + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Combine input must have shape (R, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + + expert_out = input_ + if self.buffer.zero_copy: + expert_out = _alloc_io( + tuple(input_.shape), + input_.dtype, + input_.device, + True, + ) + expert_out.copy_(input_) + + result = torch.empty( + self.num_local_tokens, + self.buffer.hidden_dim, + dtype=input_.dtype, + device=input_.device, + ) + torch.ops.transformer_engine_ep.combine( + self.buffer.handle_mem, + expert_out, + result, + ) + + if ctx.requires_grad: + grad_out = kwargs.get("grad_out") + if self.buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes combine gradients per step and cannot use " + "a caller-supplied grad_out" + ) + grad_out = _validate_grad_buffer( + grad_out, + shape=tuple(input_.shape), + dtype=input_.dtype, + device=input_.device, + ) + if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): + raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") + ctx.grad_out = grad_out + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.save_for_backward(self.buffer.handle_mem) + + return result + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + grad_output, + grad_input, + ) + return grad_input, () diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py new file mode 100644 index 0000000000..9762ed76f7 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel dispatch operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, ep_prepare +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_output_buffer( + name: str, + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous.") + if tensor.requires_grad: + raise ValueError(f"{name} must not require gradients.") + return tensor + + +class Dispatch(BasicOperation): + """Dispatch BF16 tokens to local experts with NCCL EP. + + The extra inputs are routing indices and FP32 routing weights. The extra + outputs are local tokens-per-expert and received routing weights. + """ + + num_extra_inputs: int = 2 + num_extra_outputs: int = 2 + + def __init__(self, buffer: EpBuffer) -> None: + super().__init__() + self.buffer = buffer + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}.") + expected_route_shape = (input_.shape[0], self.buffer.top_k) + if tuple(topk_idx.shape) != expected_route_shape: + raise ValueError( + f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." + ) + if tuple(topk_weights.shape) != expected_route_shape: + raise ValueError( + f"topk_weights shape must be {expected_route_shape}, " + f"got {tuple(topk_weights.shape)}." + ) + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): + if tensor.device != input_.device: + raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") + + recv_tokens = kwargs.get("recv_tokens") + recv_topk_weights = kwargs.get("recv_topk_weights") + if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): + raise ValueError( + "eager mode sizes dispatch outputs per step and cannot use " + "caller-supplied receive buffers" + ) + + tokens_per_expert = ep_prepare(self.buffer, topk_idx) + rows = ( + self.buffer._host_total_recv_tokens + if self.buffer.eager + else self.buffer.recv_capacity_per_rank + ) + if rows is None: + raise RuntimeError("NCCL EP dispatch receive size is unavailable.") + rows = int(rows) + recv_shape = (rows, self.buffer.hidden_dim) + recv_tokens = _validate_output_buffer( + "recv_tokens", + recv_tokens, + shape=recv_shape, + dtype=self.buffer.payload_dtype, + device=self.buffer.device, + ) + recv_topk_weights = _validate_output_buffer( + "recv_topk_weights", + recv_topk_weights, + shape=(rows,), + dtype=torch.float32, + device=self.buffer.device, + ) + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (rows,), + torch.float32, + self.buffer.device, + self.buffer.zero_copy, + ) + + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.save_for_backward(self.buffer.handle_mem) + + return recv_tokens, [(tokens_per_expert, recv_topk_weights)] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + + grad_recv_weights = basic_op_grad_extra_outputs[0][1] + if grad_recv_weights is None: + grad_recv_weights = torch.zeros( + grad_output.shape[0], + dtype=torch.float32, + device=grad_output.device, + ) + else: + grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous() + + grad_input = torch.empty( + ctx.input_shape, + dtype=ctx.input_dtype, + device=grad_output.device, + ) + grad_topk_weights = torch.empty( + ctx.topk_weights_shape, + dtype=torch.float32, + device=grad_output.device, + ) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + grad_output, + grad_recv_weights, + grad_input, + grad_topk_weights, + ) + return grad_input, [()], [(None, grad_topk_weights)] From 8ec3a39dc05d297e4c0f00b4fae301340c951fb2 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 10 Aug 2026 07:22:58 +0000 Subject: [PATCH 31/83] compare fused and unfused ops and they are matching Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 150 +++++++++++ tests/pytorch/test_ep_reference.py | 109 ++++++++ transformer_engine/pytorch/ep.py | 6 + .../pytorch/ep_reference.py | 111 ++++++-- .../pytorch/ops/basic/__init__.py | 2 + .../pytorch/ops/fused/__init__.py | 1 + .../pytorch/ops/fused/moe_ep.py | 247 ++++++++++++++++++ 7 files changed, 606 insertions(+), 20 deletions(-) create mode 100644 tests/pytorch/test_ep_reference.py rename tests/pytorch/distributed/moe_ep_reference.py => transformer_engine/pytorch/ep_reference.py (89%) create mode 100644 transformer_engine/pytorch/ops/fused/moe_ep.py diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 067f75d725..750a31e32e 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist +from transformer_engine.pytorch import ops as te_ops from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -23,6 +24,7 @@ _ep_combine_raw, _ep_dispatch_raw, ) +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" @@ -37,6 +39,7 @@ HIDDEN_DIM = 32 TOP_K = 2 TOKENS_PER_RANK = 4 +INTERMEDIATE_DIM = 16 def _zero_copy_test_include(fn): @@ -121,6 +124,31 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _make_moe_inputs(rank, ep_size, device="cuda"): + """Deterministic BF16 activations and FP32 top-k router weights.""" + generator = torch.Generator(device=device) + generator.manual_seed(2026 + rank) + tokens = ( + torch.randn( + TOKENS_PER_RANK, + HIDDEN_DIM, + generator=generator, + dtype=torch.float32, + device=device, + ) + * 0.25 + ).to(torch.bfloat16) + router_logits = torch.randn( + TOKENS_PER_RANK, + ep_size * NUM_LOCAL_EXPERTS, + generator=generator, + dtype=torch.float32, + device=device, + ) + topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) + return topk_idx, tokens, torch.softmax(topk_logits, dim=-1) + + class _Cfg: rank: int world_size: int @@ -232,6 +260,39 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) + def _make_moe_model(self, *, fusion_barrier=False): + buffer = self._make_buffer() + dispatch = te_ops.Dispatch(buffer) + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * INTERMEDIATE_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + INTERMEDIATE_DIM, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "tokens_per_expert") + dispatch.set_extra_output_channel(1, "routing_weights") + fc1.set_extra_input_channel(0, "tokens_per_expert") + activation.set_extra_input_channel(0, "routing_weights") + fc2.set_extra_input_channel(0, "tokens_per_expert") + ops = [dispatch, fc1, activation, fc2] + if fusion_barrier: + ops.append(te_ops.Identity()) + ops.append(combine) + return te_ops.Sequential(*ops), fc1, fc2 + # Prepare @_eager_test_include @@ -419,6 +480,95 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) + @_eager_test_include + def test_bf16_moe_sequential_fusion(self): + """Reference-backed fusion matches the unfused BF16 EP MoE sequence.""" + if not EAGER: + self.skipTest("variable-size reference comparison requires eager EP mode") + + fused, fused_fc1, fused_fc2 = self._make_moe_model() + unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fusion_barrier=True) + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(3100 + self.cfg.rank) + with torch.no_grad(): + for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for expert in range(NUM_LOCAL_EXPERTS): + weight = ( + torch.randn( + getattr(fused_op, f"weight{expert}").shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + getattr(fused_op, f"weight{expert}").copy_(weight) + getattr(unfused_op, f"weight{expert}").copy_(weight) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + fused_tokens = tokens.detach().clone().requires_grad_(True) + unfused_tokens = tokens.detach().clone().requires_grad_(True) + fused_topk_weights = topk_weights.detach().clone().requires_grad_(True) + unfused_topk_weights = topk_weights.detach().clone().requires_grad_(True) + + fused_out, fused_counts, fused_recv_weights = fused( + fused_tokens, + topk_idx, + fused_topk_weights, + ) + unfused_out, unfused_counts, unfused_recv_weights = unfused( + unfused_tokens, + topk_idx, + unfused_topk_weights, + ) + + fused_forward_ops = fused._module_groups[0]._forward_ops + unfused_forward_ops = unfused._module_groups[0]._forward_ops + self.assertEqual(len(fused_forward_ops), 1) + self.assertIsInstance(fused_forward_ops[0][0], FusedMoeEp) + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops)) + self.assertEqual(fused_out.dtype, torch.bfloat16) + self.assertEqual(unfused_out.dtype, torch.bfloat16) + torch.testing.assert_close(fused_counts, unfused_counts, rtol=0, atol=0) + self.assertEqual(fused_recv_weights.dtype, torch.float32) + self.assertEqual(unfused_recv_weights.dtype, torch.float32) + + dy = ( + torch.randn( + fused_out.shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + fused_out.backward(dy) + unfused_out.backward(dy) + torch.cuda.synchronize() + + # The two BF16 paths use different grouped-GEMM and reduction orders. + # Their largest observed forward absolute error is ~3.1e-5, at values + # close enough to zero that the relative tolerance does not apply. + tolerances = {"rtol": 1.6e-2, "atol": 5e-5} + torch.testing.assert_close(fused_out, unfused_out, **tolerances) + torch.testing.assert_close(fused_tokens.grad, unfused_tokens.grad, **tolerances) + torch.testing.assert_close( + fused_topk_weights.grad, + unfused_topk_weights.grad, + **tolerances, + ) + for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for expert in range(NUM_LOCAL_EXPERTS): + fused_grad = getattr(fused_op, f"weight{expert}").grad + unfused_grad = getattr(unfused_op, f"weight{expert}").grad + self.assertEqual(fused_grad.dtype, torch.bfloat16) + self.assertEqual(unfused_grad.dtype, torch.bfloat16) + torch.testing.assert_close(fused_grad, unfused_grad, **tolerances) + @_zero_copy_test_include def test_zero_copy_pool_auto_alloc(self): """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO diff --git a/tests/pytorch/test_ep_reference.py b/tests/pytorch/test_ep_reference.py new file mode 100644 index 0000000000..47a0731ff6 --- /dev/null +++ b/tests/pytorch/test_ep_reference.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the pure PyTorch MoE EP reference.""" + +from types import SimpleNamespace + +import pytest +import torch + +from transformer_engine.pytorch import ops as te_ops +from transformer_engine.pytorch.ep_reference import MoeEpReference +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp + + +@pytest.mark.parametrize("compute_dtype", (torch.float32, torch.bfloat16)) +def test_moe_ep_reference_compute_dtype(compute_dtype): + """The configured dtype controls MLP, combine, and non-router gradients.""" + generator = torch.Generator().manual_seed(1234) + activation = torch.randn(4, 8, generator=generator, dtype=torch.bfloat16) + fc1_weight = torch.randn(2, 8, 8, generator=generator, dtype=torch.bfloat16) + fc2_weight = torch.randn(2, 4, 8, generator=generator, dtype=torch.bfloat16) + topk_idx = torch.tensor([[0], [1], [0], [1]], dtype=torch.int64) + topk_weights = torch.ones(4, 1, dtype=torch.float32) + + reference = MoeEpReference( + num_experts=2, + hidden_size=8, + intermediate_size=4, + top_k=1, + generate_c=True, + compute_dtype=compute_dtype, + ) + output, fc1_c, route_metadata = reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + grads = reference.backward( + torch.ones_like(output), + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + assert output.dtype is torch.bfloat16 + assert grads[0].dtype is compute_dtype + assert grads[1].dtype is compute_dtype + assert grads[2].dtype is compute_dtype + assert grads[3].dtype is torch.float32 + + +def test_moe_ep_reference_default_compute_dtype_is_fp32(): + """Preserve the reference's pre-existing FP32 compute default.""" + reference = MoeEpReference( + num_experts=1, + hidden_size=8, + intermediate_size=4, + top_k=1, + ) + assert reference.compute_dtype is torch.float32 + + +def test_single_rank_bf16_moe_fusion_forward_backward(): + """Exercise the fuser contract without requiring the NCCL EP backend.""" + buffer = SimpleNamespace( + num_local_experts=2, + hidden_dim=8, + top_k=1, + max_tokens_per_rank=4, + payload_dtype=torch.bfloat16, + eager=True, + ) + dispatch = te_ops.Dispatch(buffer) + fc1 = te_ops.GroupedLinear(2, 8, 8, bias=False, device="cpu", dtype=torch.bfloat16) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear(2, 4, 8, bias=False, device="cpu", dtype=torch.bfloat16) + combine = te_ops.Combine(buffer, num_local_tokens=4) + dispatch.set_extra_output_channel(0, "tokens_per_expert") + dispatch.set_extra_output_channel(1, "routing_weights") + fc1.set_extra_input_channel(0, "tokens_per_expert") + activation.set_extra_input_channel(0, "routing_weights") + fc2.set_extra_input_channel(0, "tokens_per_expert") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + + x = torch.randn(4, 8, dtype=torch.bfloat16, requires_grad=True) + topk_idx = torch.tensor([[0], [1], [0], [1]], dtype=torch.int64) + topk_weights = torch.ones(4, 1, dtype=torch.float32, requires_grad=True) + output, counts, recv_weights = model(x, topk_idx, topk_weights) + torch.autograd.backward( + (output, recv_weights), + (torch.ones_like(output), torch.ones_like(recv_weights)), + ) + + assert isinstance(model._module_groups[0]._forward_ops[0][0], FusedMoeEp) + assert counts.dtype is torch.int64 + assert output.dtype is torch.bfloat16 + assert x.grad.dtype is torch.bfloat16 + assert topk_weights.grad.dtype is torch.float32 + for op in (fc1, fc2): + for expert in range(op.num_groups): + assert getattr(op, f"weight{expert}").grad.dtype is torch.bfloat16 diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 42f605367a..28c4a49a13 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -21,6 +21,7 @@ __all__ = [ "EpBuffer", "ep_bootstrap", + "get_ep_group", "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", @@ -175,6 +176,11 @@ def is_ep_bootstrapped() -> bool: return _BOOTSTRAPPED +def get_ep_group() -> Optional[dist.ProcessGroup]: + """Return the process group registered by :func:`ep_bootstrap`.""" + return _EP_GROUP + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/transformer_engine/pytorch/ep_reference.py similarity index 89% rename from tests/pytorch/distributed/moe_ep_reference.py rename to transformer_engine/pytorch/ep_reference.py index 280c4f0d21..efb8d5c471 100644 --- a/tests/pytorch/distributed/moe_ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -297,6 +297,7 @@ def _decode_tensor( name: str, expected_shape: Tuple[int, ...], quantized_axis: int, + dtype: torch.dtype, ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: @@ -305,19 +306,24 @@ def _decode_tensor( ) if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") - return tensor.dequantize() + return tensor.dequantize(dtype=dtype) if tuple(tensor.shape) != expected_shape: raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") if not tensor.is_floating_point(): raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") - return tensor.float() + return tensor.to(dtype=dtype) -def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: +def _format_round_trip( + tensor: torch.Tensor, + format: MoeFormat, + *, + dtype: torch.dtype, +) -> torch.Tensor: if format is MoeFormat.BF16: - return tensor.to(torch.bfloat16).float() - return quantize_blockwise(tensor, format, axis=-1).dequantize() + return tensor.to(torch.bfloat16).to(dtype) + return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) class MoeEpReference: @@ -343,6 +349,7 @@ def __init__( apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, generate_c: bool = False, + compute_dtype: torch.dtype = torch.float32, ) -> None: for name, value in ( ("num_experts", num_experts), @@ -356,6 +363,11 @@ def __init__( raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") if max_tokens_per_rank is not None and max_tokens_per_rank < 0: raise ValueError("max_tokens_per_rank must be non-negative") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError( + "compute_dtype must be torch.float32 or torch.bfloat16, " + f"got {compute_dtype}" + ) if ep_group is None: ep_size, ep_rank = 1, 0 @@ -385,6 +397,7 @@ def __init__( self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) + self.compute_dtype = compute_dtype for name, fmt in ( ("output_format", self.output_format), @@ -405,7 +418,8 @@ def __repr__(self) -> str: f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " - f"output={self.output_format.value}, combine={self.combine_format.value})" + f"output={self.output_format.value}, combine={self.combine_format.value}, " + f"compute_dtype={self.compute_dtype})" ) def _collective_device(self, device: torch.device) -> torch.device: @@ -495,7 +509,7 @@ def _run_local_experts( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: output = torch.empty( (tokens.shape[0], self.hidden_size), - dtype=torch.float32, + dtype=self.compute_dtype, device=tokens.device, ) fc1_c_rows = [] if self.generate_c else None @@ -513,13 +527,21 @@ def _run_local_experts( gate = gate.clamp(max=self.gate_up_clamp) up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) intermediate = F.silu(gate) * up - weights = route_weight.index_select(0, positions).unsqueeze(-1) + weights = ( + route_weight.index_select(0, positions) + .to(dtype=self.compute_dtype) + .unsqueeze(-1) + ) if self.apply_topk_in_fc1: intermediate = intermediate * weights expert_output = intermediate @ fc2_weight[expert] if not self.apply_topk_in_fc1: expert_output = expert_output * weights - expert_output = _format_round_trip(expert_output, self.combine_format) + expert_output = _format_round_trip( + expert_output, + self.combine_format, + dtype=self.compute_dtype, + ) output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: @@ -539,7 +561,13 @@ def __call__( fc2_weight: MoeTensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, - ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: + *, + return_dispatch_metadata: bool = False, + ) -> Union[ + MoeTensor, + Tuple[MoeTensor, torch.Tensor, torch.Tensor], + Tuple[MoeTensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + ]: """Run dispatch, local experts, return routing, top-k reduce, and encode. Shapes: @@ -595,18 +623,21 @@ def __call__( name="activation", expected_shape=(token_count, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) fc1_float = _decode_tensor( fc1_weight, name="fc1_weight", expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), quantized_axis=1, + dtype=self.compute_dtype, ) fc2_float = _decode_tensor( fc2_weight, name="fc2_weight", expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) plan = self._dispatch_plan(topk_idx, topk_weights) @@ -620,6 +651,7 @@ def __call__( recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) route_metadata = None + grouped_recv_weight = None if self.generate_c: recv_src_rank = torch.repeat_interleave( torch.arange(self.ep_size, device=device), @@ -630,6 +662,7 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) + grouped_recv_weight = recv_weight.index_select(0, fc1_c_order) route_metadata = ( torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) .index_select(0, fc1_c_order) @@ -649,7 +682,7 @@ def __call__( returned = self._all_to_all(recv_output, recv_counts, send_counts) combine_plane = torch.zeros( (token_count * self.top_k, self.hidden_size), - dtype=torch.float32, + dtype=self.compute_dtype, device=device, ) send_flat_slot = send_token_idx * self.top_k + send_slot_idx @@ -661,7 +694,15 @@ def __call__( else: output = quantize_blockwise(reduced, self.output_format, axis=-1) if self.generate_c: + if return_dispatch_metadata: + tokens_per_expert = torch.bincount( + recv_expert.to(torch.int64), + minlength=self.experts_per_rank, + ) + return output, fc1_c, route_metadata, tokens_per_expert, grouped_recv_weight return output, fc1_c, route_metadata + if return_dispatch_metadata: + raise ValueError("return_dispatch_metadata=True requires generate_c=True") return output def backward( @@ -674,6 +715,7 @@ def backward( topk_weights: torch.Tensor, fc1_c: torch.Tensor, route_metadata: torch.Tensor, + grad_recv_topk_weights: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Backward pass consuming the ``generate_c=True`` stash. @@ -688,7 +730,8 @@ def backward( ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, - grad_topk_weights)`` in float32. + grad_topk_weights)``. Activation and expert weight gradients use + ``compute_dtype``; router-weight gradients remain float32. """ if not self.generate_c: @@ -703,6 +746,13 @@ def backward( ) if not grad_output.is_floating_point(): raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + if grad_recv_topk_weights is not None: + expected_shape = (int(route_metadata.shape[0]),) + if tuple(grad_recv_topk_weights.shape) != expected_shape: + raise ValueError( + "grad_recv_topk_weights shape must be " + f"{expected_shape}, got {tuple(grad_recv_topk_weights.shape)}" + ) device = _tensor_device(activation) two_i = 2 * self.intermediate_size @@ -711,18 +761,21 @@ def backward( name="activation", expected_shape=(token_count, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) fc1_float = _decode_tensor( fc1_weight, name="fc1_weight", expected_shape=(self.experts_per_rank, self.hidden_size, two_i), quantized_axis=1, + dtype=self.compute_dtype, ) fc2_float = _decode_tensor( fc2_weight, name="fc2_weight", expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): raise ValueError( @@ -734,7 +787,7 @@ def backward( # along the identical forward routes. plan = self._dispatch_plan(topk_idx, topk_weights) send_counts, recv_counts = plan.send_counts, plan.recv_counts - grad_output_float = grad_output.float() + grad_output_float = grad_output.to(dtype=self.compute_dtype) recv_tokens = self._all_to_all( activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts ) @@ -761,9 +814,13 @@ def backward( dy_rows = torch.empty_like(recv_grad) dy_rows.index_copy_(0, perm, recv_grad) - c_rows = fc1_c.float() + c_rows = fc1_c.to(dtype=self.compute_dtype) expert_rows = metadata[:, 0] - d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_x_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=self.compute_dtype, + device=device, + ) d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) grad_fc1 = torch.zeros_like(fc1_float) grad_fc2 = torch.zeros_like(fc2_float) @@ -773,7 +830,11 @@ def backward( continue c = c_rows.index_select(0, positions) x = x_rows.index_select(0, positions) - w = w_rows.index_select(0, positions).unsqueeze(-1) + w = ( + w_rows.index_select(0, positions) + .to(dtype=self.compute_dtype) + .unsqueeze(-1) + ) d_y = dy_rows.index_select(0, positions) gate, up = c.split(self.intermediate_size, dim=-1) @@ -796,10 +857,12 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).float() else: d_h = d_h_fc2 - d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) + d_w_rows[positions] = ( + d_y * (h @ fc2_float[expert]) + ).sum(dim=-1).float() d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s @@ -814,9 +877,17 @@ def backward( # Return the route gradients to their source ranks and scatter-add. returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) - returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + if grad_recv_topk_weights is not None: + d_w_rows = d_w_rows + grad_recv_topk_weights.to( + device=device, + dtype=torch.float32, + ) + recv_dw = d_w_rows.index_select(0, perm) + returned_dw = self._all_to_all(recv_dw, recv_counts, send_counts) grad_activation = torch.zeros( - (token_count, self.hidden_size), dtype=torch.float32, device=device + (token_count, self.hidden_size), + dtype=self.compute_dtype, + device=device, ) grad_activation.index_add_(0, plan.send_token_idx, returned_dx) grad_topk_weights = torch.zeros( diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..531dde233f 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -22,7 +22,9 @@ from .all_reduce import AllReduce from .basic_linear import BasicLinear from .bias import Bias +from .combine import Combine from .constant_scale import ConstantScale +from .dispatch import Dispatch from .dropout import Dropout from .grouped_linear import GroupedLinear from .identity import Identity diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index dc9dcd6dc3..7aa45a89f5 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -35,3 +35,4 @@ GroupedMLP_CuTeGEMMGLU, GroupedMLP_CuTeGEMMUnary, ) +from .moe_ep import FusedMoeEp # pylint: disable=wrong-import-position diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py new file mode 100644 index 0000000000..5d5b861f8f --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -0,0 +1,247 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Reference-backed BF16 expert-parallel MoE fusion.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any, Optional + +import torch + +from ...ep import get_ep_group +from ...ep_reference import MoeEpReference +from ...quantization import Recipe +from ...tensor import Quantizer +from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU +from ..fuser import register_forward_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext + + +def _weight_list(op: GroupedLinear) -> list[torch.Tensor]: + """Return per-expert weights in their registered order.""" + return [getattr(op, f"weight{idx}") for idx in range(op.num_groups)] + + +def _reference_weights(op: GroupedLinear) -> torch.Tensor: + """Pack ``(out, in)`` expert weights into reference ``(E, in, out)`` layout.""" + return torch.stack([weight.transpose(0, 1) for weight in _weight_list(op)]) + + +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weights = _weight_list(op) if not op.single_grouped_weight else [] + return ( + not op.use_bias + and not op._scale_bias + and not op.single_grouped_weight + and not op.single_grouped_bias + and not op._accumulate_into_main_grad + and not op._is_distributed_weight() + and not op.wgrad_store.delay_wgrad_compute() + and bool(weights) + and all(weight.dtype is torch.bfloat16 for weight in weights) + ) + + +def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: + if recipe is not None or len(window) != 5: + return False + dispatch, fc1, activation, fc2, combine = window + if not ( + isinstance(dispatch, Dispatch) + and isinstance(fc1, GroupedLinear) + and isinstance(activation, ScaledSwiGLU) + and isinstance(fc2, GroupedLinear) + and isinstance(combine, Combine) + ): + return False + if dispatch.buffer is not combine.buffer: + return False + buffer = dispatch.buffer + if not buffer.eager or buffer.payload_dtype is not torch.bfloat16: + return False + if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): + return False + if activation.activation_recompute_in_mlp or activation.glu_interleave_size is not None: + return False + return ( + fc1.num_groups == buffer.num_local_experts + and fc2.num_groups == buffer.num_local_experts + and fc1.in_features == buffer.hidden_dim + and fc2.out_features == buffer.hidden_dim + and fc1.out_features == 2 * fc2.in_features + ) + + +class FusedMoeEp(FusedOperation): + """Joint BF16 fusion implemented with :class:`MoeEpReference`.""" + + def __init__( + self, + *, + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, + combine: Combine, + ) -> None: + super().__init__([dispatch, fc1, activation, fc2, combine]) + ep_group = get_ep_group() + ep_size = 1 if ep_group is None else ep_group.size() + self._reference = MoeEpReference( + num_experts=dispatch.buffer.num_local_experts * ep_size, + hidden_size=dispatch.buffer.hidden_dim, + intermediate_size=fc2.in_features, + top_k=dispatch.buffer.top_k, + ep_group=ep_group, + max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, + compute_dtype=torch.bfloat16, + apply_topk_in_fc1=True, + generate_c=True, + ) + + @property + def dispatch(self) -> Dispatch: + return self.basic_ops[0] + + @property + def fc1(self) -> GroupedLinear: + return self.basic_ops[1] + + @property + def fc2(self) -> GroupedLinear: + return self.basic_ops[3] + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") + if any(kwargs for kwargs in basic_op_kwargs): + raise NotImplementedError("FusedMoeEp does not support per-operation output buffers.") + + topk_idx, topk_weights = basic_op_extra_inputs[0] + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + fc1_weight = _reference_weights(self.fc1) + fc2_weight = _reference_weights(self.fc2) + output, fc1_c, route_metadata, tokens_per_expert, recv_topk_weights = self._reference( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + return_dispatch_metadata=True, + ) + + if any(ctx.requires_grad for ctx in basic_op_ctxs): + basic_op_ctxs[0].save_for_backward( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + return output, [ + (tokens_per_expert, recv_topk_weights), + (), + (), + (), + (), + ] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) = basic_op_ctxs[0].saved_tensors + grad_recv_topk_weights = basic_op_grad_extra_outputs[0][1] + grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._reference.backward( + grad_output, + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + grad_recv_topk_weights=grad_recv_topk_weights, + ) + + fc1_param_grads = [ + grad_fc1[idx].transpose(0, 1) if weight.requires_grad else None + for idx, weight in enumerate(_weight_list(self.fc1)) + ] + fc2_param_grads = [ + grad_fc2[idx].transpose(0, 1) if weight.requires_grad else None + for idx, weight in enumerate(_weight_list(self.fc2)) + ] + return ( + grad_input, + [(), fc1_param_grads, (), fc2_param_grads, ()], + [(None, grad_topk_weights), (None,), (None,), (None,), ()], + ) + + +def fuse_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused: Any, +) -> list[FusibleOperation]: + """Fuse supported five-op BF16 EP MoE sequences.""" + del unused + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + window = ops[idx : idx + 5] + if _matches(window, recipe): + dispatch, fc1, activation, fc2, combine = window + out.append( + FusedMoeEp( + dispatch=dispatch, + fc1=fc1, + activation=activation, + fc2=fc2, + combine=combine, + ) + ) + idx += 5 + else: + out.append(ops[idx]) + idx += 1 + return out + + +register_forward_backward_fusion(fuse_ops, prepend=True) + + +__all__ = ["FusedMoeEp"] From 63a4ea31df6378afb34aa20ebd5f86ba8c0ff287 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 10 Aug 2026 21:00:27 +0000 Subject: [PATCH 32/83] fix lint Signed-off-by: Varun Thumbe --- .../pytorch/ops/fused/forward_linear_bias_activation.py | 2 +- transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py | 2 +- .../pytorch/ops/fused/forward_linear_scale_add.py | 2 +- .../pytorch/ops/fused/userbuffers_forward_linear.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index b6f56effef..6fa63675b4 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + activation.""" from __future__ import annotations -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Any, Optional import torch diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index f32eb0f7ac..28586360f5 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + add.""" from __future__ import annotations -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Any, Optional import torch diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index 2fb6c01dca..277263e0ec 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + scale + add.""" from __future__ import annotations -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Any, Optional import torch diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 93c4024880..cb5f8a16e6 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -5,7 +5,7 @@ """Linear layer forward with Userbuffers communication.""" from __future__ import annotations -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Any, Optional import torch From 87c1cf6ad9044d9f01a8e419cec46d4135d56b1a Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 10 Aug 2026 22:07:17 -0700 Subject: [PATCH 33/83] Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 --- transformer_engine/pytorch/ops/fuser.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index cf1df530e1..07e37156d2 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -123,7 +123,7 @@ def forward( for idx in basic_op_idxs: basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward - # Forward op. Resolve internal channel inputs from outputs of + # 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 @@ -137,6 +137,8 @@ def forward( continue producer_outputs = extra_outputs[producer_idx] basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] + + # Prepare args for op forward extra_inputs = [basic_op_extra_inputs[idx] for idx in basic_op_idxs] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None From 307ab15a0049d4e744a88d6a84a88df8be4fb966 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 10 Aug 2026 22:10:53 -0700 Subject: [PATCH 34/83] Update docs/examples/op_fuser/op_fuser.rst Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 --- docs/examples/op_fuser/op_fuser.rst | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index 123fd1cd76..1466e38121 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -154,14 +154,7 @@ arguments and the extra outputs will be returned. 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. +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``. With a channel, the residual block above can be expressed using one ``Sequential``: From 97a91cf96f2436dd6d2e9cfebb18bde05ea66c30 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 10 Aug 2026 22:11:24 -0700 Subject: [PATCH 35/83] Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 --- transformer_engine/pytorch/ops/fuser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 07e37156d2..3c95845c80 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -176,7 +176,7 @@ def forward( if ( set_output_requires_grad and idx >= fuser.first_op_requiring_backward - and (y.is_floating_point() or y.is_complex()) + and y.is_floating_point() ): y.requires_grad_(True) extra_outputs[idx] = ys From 6468a14991c5b8aac317fc670e1556a1e96be5cb Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 11 Aug 2026 06:38:10 +0000 Subject: [PATCH 36/83] address review comments + extra output being configurable to be outputted Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 64 ++++++--- tests/pytorch/test_fusible_ops.py | 175 ++++++++++++++++++------ transformer_engine/pytorch/ops/fuser.py | 95 +++++++++---- transformer_engine/pytorch/ops/op.py | 41 ++++-- 4 files changed, 272 insertions(+), 103 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index 1466e38121..f7e9acc0a0 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -154,7 +154,12 @@ arguments and the extra outputs will be returned. Extra tensor channels """"""""""""""""""""" -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``. +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 and may optionally +be returned from the ``Sequential``. Extra input tensors with a specified +channel are accessed internally instead of being provided as arguments +to ``Sequential``. With a channel, the residual block above can be expressed using one ``Sequential``: @@ -166,7 +171,9 @@ With a channel, the residual block above can be expressed using one make_residual = te.ops.MakeExtraOutput() add_residual = te.ops.AddExtraInput() - make_residual.set_extra_output_channel(0, "residual") + make_residual.set_extra_output_channel( + 0, "residual", output_to_caller=False + ) add_residual.set_extra_input_channel(0, "residual") block = te.ops.Sequential( @@ -178,9 +185,9 @@ With a channel, the residual block above can be expressed using one add_residual, ) - # The residual is routed internally and is also returned to the caller. + # The residual is routed internally and omitted from the public outputs. x = torch.randn(16384, 4096, device="cuda") - y, residual = block(x) + y = block(x) Channels are also useful for mixture-of-experts blocks. The following example assumes custom ``Dispatch`` and ``Combine`` basic operations. @@ -209,9 +216,15 @@ routing map. ``Combine`` consumes the routing map. # 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") + dispatch.set_extra_output_channel( + 0, "m_splits", output_to_caller=False + ) + dispatch.set_extra_output_channel( + 1, "probs", output_to_caller=False + ) + dispatch.set_extra_output_channel( + 2, "routing_map", output_to_caller=False + ) fc1.set_extra_input_channel(0, "m_splits") activation.set_extra_input_channel(0, "probs") @@ -221,9 +234,9 @@ routing map. ``Combine`` consumes the 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 supply all later extra inputs internally. The channel outputs + # are not returned because output_to_caller=False. + y = moe(x, router_probs) Channels cannot connect operations in different ``OperationFuser`` instances. In particular, an ordinary PyTorch module inside a @@ -250,22 +263,25 @@ 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. +- Every named extra input must have a matching producer earlier in the + same fuser. Missing producers, backward edges, and cycles are not + supported. Leave an extra input unnamed when the caller should provide it. - 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 named output does not require a consumer. It is returned as a public + extra output by default. - 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. +- The caller passes unnamed extra inputs. Named, channel-connected extra + input slots do not appear in the ``Sequential`` arguments. +- ``set_extra_output_channel`` accepts ``output_to_caller`` (``True`` by + default). Public extra outputs are returned in their original + basic-operation and slot order. Gradients supplied for a returned output + are combined with gradients from its internal channel consumers. +- Set ``output_to_caller=False`` for a channel tensor that should remain + internal. Removing a channel binding with ``channel=None`` restores that + output as public. - 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 @@ -275,7 +291,11 @@ 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. +tensor and its gradient between those basic operations. For a non-public +channel fully owned by one forward fusion, ``fuser_forward`` may return +``None`` in the corresponding basic-operation output slot. A tensor is +still required when the output is public or when a consumer is outside +that forward fusion. Developer guide --------------- diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index b50c9c3fc1..57a3b29b94 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -504,11 +504,11 @@ def test_internal_residual_connection( torch.testing.assert_close(body.bias.grad, dy) @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) - @pytest.mark.parametrize("with_extra_grad", (True, False)) + @pytest.mark.parametrize("output_to_caller", (True, False)) def test_fused_internal_residual_connection( self, fusion_kind: str, - with_extra_grad: bool, + output_to_caller: bool, size: int = 16, ) -> None: """Forward, backward, and joint fusions can own an internal channel.""" @@ -533,7 +533,8 @@ def fuser_forward( # The consumer slot is internal to this fusion, so the # OperationFuser deliberately leaves it unset. assert basic_op_extra_inputs[2][0] is None - return 2 * input_ + self.basic_ops[1].bias, [(input_,), (), ()] + residual_out = input_ if output_to_caller else None + return 2 * input_ + self.basic_ops[1].bias, [(residual_out,), (), ()] def fuser_backward( self, @@ -544,7 +545,7 @@ def fuser_backward( ): del basic_op_ctxs # The fusion owns the internal residual edge. The fuser also - # supplies the gradient from the public residual output. + # supplies a gradient when the residual is a public output. grad_residual = basic_op_grad_extra_outputs[0][0] return ( 2 * grad_output @@ -570,7 +571,11 @@ def fuse_residual(ops, **unused): residual = te_ops.MakeExtraOutput() body = te_ops.Bias(size=size, device="cpu") add_residual = te_ops.AddExtraInput() - residual.set_extra_output_channel(0, "residual") + residual.set_extra_output_channel( + 0, + "residual", + output_to_caller=output_to_caller, + ) add_residual.set_extra_input_channel(0, "residual") model = te_ops.Sequential(residual, body, add_residual) @@ -581,7 +586,12 @@ def fuse_residual(ops, **unused): else: te_ops.register_forward_backward_fusion(fuse_residual, prepend=True) x = torch.rand((size,), requires_grad=True) - y, residual_out = model(x) + outputs = model(x) + if output_to_caller: + y, residual_out = outputs + else: + assert isinstance(outputs, torch.Tensor) + y = outputs forward_ops = model._module_groups[0]._forward_ops backward_ops = model._module_groups[0]._backward_ops @@ -599,7 +609,7 @@ def fuse_residual(ops, **unused): assert backward_ops[0][0] is forward_ops[0][0] torch.testing.assert_close(y, 2 * x + body.bias) dy = torch.rand_like(y) - if with_extra_grad: + if output_to_caller: dresidual = torch.rand_like(residual_out) torch.autograd.backward((y, residual_out), (dy, dresidual)) expected_dx = 2 * dy + dresidual @@ -638,6 +648,25 @@ def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: torch.testing.assert_close(y_no_grad, 3 * x_no_grad) torch.testing.assert_close(route_no_grad, x_no_grad) + def test_internal_extra_tensor_channel_can_be_hidden(self, size: int = 16) -> None: + """A non-public channel still propagates forward and backward.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route", output_to_caller=False) + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + assert isinstance(y, torch.Tensor) + torch.testing.assert_close(y, 3 * x) + + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 3 * dy) + def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: """Unbound slots remain public when other slots use internal channels.""" producer = te_ops.MakeExtraOutput() @@ -658,25 +687,12 @@ def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None torch.testing.assert_close(x.grad, 2 * dy) torch.testing.assert_close(extra.grad, dy) - def test_external_named_extra_inputs_remain_separate(self, size: int = 16) -> None: - """Unmatched inputs with the same channel require separate public tensors.""" - consumer1 = te_ops.AddExtraInput() - consumer2 = te_ops.AddExtraInput() - consumer1.set_extra_input_channel(0, "external") - consumer2.set_extra_input_channel(0, "external") - model = te_ops.Sequential(consumer1, consumer2) - - x = torch.rand((size,), requires_grad=True) - extra = torch.rand((size,), requires_grad=True) - with pytest.raises(ValueError, match="Expected 2 extra inputs but got 1"): - model(x, extra) - y = model(x, extra, extra) - torch.testing.assert_close(y, x + 2 * extra) - - dy = torch.rand_like(y) - y.backward(dy) - torch.testing.assert_close(x.grad, dy) - torch.testing.assert_close(extra.grad, 2 * dy) + def test_named_extra_input_requires_producer(self) -> None: + """A named input cannot fall back to a caller-provided tensor.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "missing") + with pytest.raises(ValueError, match="has no producer"): + OperationFuser([consumer]) def test_consumer_before_producer(self) -> None: """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" @@ -712,6 +728,12 @@ def test_set_extra_channel_rejects_invalid_name(self) -> None: consumer.set_extra_input_channel(0, "") with pytest.raises(ValueError, match="non-empty string"): producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + with pytest.raises(TypeError, match="output_to_caller must be a bool"): + producer.set_extra_output_channel( + 0, + "route", + output_to_caller=1, # type: ignore[arg-type] + ) def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> None: """Sequential does not auto-rebuild after a channel configuration change.""" @@ -736,6 +758,10 @@ def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> N torch.testing.assert_close(y, x + extra) torch.testing.assert_close(route, x) + producer.set_extra_output_channel(0, "route", output_to_caller=False) + with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): + model(x, extra) + @pytest.mark.parametrize("layout", ("two_ops", "same_op")) def test_duplicate_extra_output_channel_names(self, layout: str) -> None: """A channel name may have at most one producer, across ops or slots.""" @@ -783,29 +809,72 @@ def test_one_extra_input_has_single_source(self, size: int = 16) -> None: torch.testing.assert_close(output_a, x) torch.testing.assert_close(output_b, x) - def test_mixed_channel_outputs_are_public(self, size: int = 16) -> None: - """Both internally consumed and unconsumed channel outputs are public.""" + def test_mixed_public_and_hidden_channel_outputs(self, size: int = 16) -> None: + """Only configured public outputs are returned, in slot order.""" producer = _DualExtraOutput(scales=(2.0, 3.0)) consumer = te_ops.AddExtraInput() - producer.set_extra_output_channel(0, "internal") + producer.set_extra_output_channel(0, "internal", output_to_caller=False) producer.set_extra_output_channel(1, "public") consumer.set_extra_input_channel(0, "internal") model = te_ops.Sequential(producer, consumer) x = torch.rand((size,), requires_grad=True) - y, internal, public = model(x) + y, public = model(x) torch.testing.assert_close(y, 3 * x) - torch.testing.assert_close(internal, 2 * x) torch.testing.assert_close(public, 3 * x) dy = torch.rand_like(y) - dinternal = torch.rand_like(internal) dpublic = torch.rand_like(public) - torch.autograd.backward((y, internal, public), (dy, dinternal, dpublic)) - torch.testing.assert_close(x.grad, 3 * dy + 2 * dinternal + 3 * dpublic) + torch.autograd.backward((y, public), (dy, dpublic)) + torch.testing.assert_close(x.grad, 3 * dy + 3 * dpublic) + + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_fused_op_cannot_omit_required_channel_output( + self, + output_to_caller: bool, + size: int = 16, + ) -> None: + """A fusion must materialize public outputs and cross-fusion channels.""" + + class FusedProducer(te_ops.FusedOperation): + _enabled = True + + def __init__(self, producer) -> None: + super().__init__((producer,)) + + def fuser_forward(self, basic_op_ctxs, input_, **unused): + del basic_op_ctxs + return input_, [(None,)] + + def fuse_producer(ops, **unused): + if ( + FusedProducer._enabled + and len(ops) == 2 + and isinstance(ops[0], te_ops.MakeExtraOutput) + and isinstance(ops[1], te_ops.AddExtraInput) + ): + FusedProducer._enabled = False + return [FusedProducer(ops[0]), ops[1]] + return ops + + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel( + 0, + "route", + output_to_caller=output_to_caller, + ) + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer) + te_ops.register_forward_fusion(fuse_producer, prepend=True) + + x = torch.rand((size,), requires_grad=True) + error = "is public" if output_to_caller else "outside its forward fusion" + with pytest.raises(RuntimeError, match=error): + model(x) def test_fresh_internal_output_preserves_grad_requirement(self) -> None: - """A fresh internal tensor requests its gradient from a scaled activation.""" + """A freshly computed internal channel still receives a consumer gradient.""" # A BasicOperation with one extra output that is freshly computed instead of # retrieved from a previous op's tensor. @@ -830,17 +899,39 @@ def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outp grad_input = grad_output + grad_scale.unsqueeze(-1) * 2 * input_ / input_.size(-1) return grad_input, [()], [()] + class ScaleByExtra(te_ops.BasicOperation): + num_extra_inputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("ScaleByExtra uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("ScaleByExtra uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + scale = basic_op_extra_inputs[0][0] + basic_op_ctxs[0].save_for_backward(input_, scale) + return input_ * scale.unsqueeze(-1), [()] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_grad_extra_outputs + input_, scale = basic_op_ctxs[0].saved_tensors + grad_input = grad_output * scale.unsqueeze(-1) + grad_scale = (grad_output * input_).sum(dim=-1) + return grad_input, [()], [(grad_scale,)] + producer = MakeScale() - activation = te_ops.ScaledSReLU() - producer.set_extra_output_channel(0, "scale") - activation.set_extra_input_channel(0, "scale") - model = te_ops.Sequential(producer, activation) + consumer = ScaleByExtra() + producer.set_extra_output_channel(0, "scale", output_to_caller=False) + consumer.set_extra_input_channel(0, "scale") + model = te_ops.Sequential(producer, consumer) - x_ref = torch.randn((5, 8), device="cuda", requires_grad=True) + x_ref = torch.randn((5, 8), requires_grad=True) x_test = x_ref.detach().clone().requires_grad_(True) scale_ref = x_ref.square().mean(dim=-1) - y_ref = torch.nn.functional.relu(x_ref).square() * scale_ref.unsqueeze(-1) - y_test, _scale_test = model(x_test) + y_ref = x_ref * scale_ref.unsqueeze(-1) + y_test = model(x_test) + assert isinstance(y_test, torch.Tensor) torch.testing.assert_close(y_test, y_ref) dy = torch.rand_like(y_ref) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 3c95845c80..50c1c11ffa 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -116,7 +116,9 @@ def forward( # Apply forward ops x = input_ - extra_outputs = [None] * fuser._num_basic_ops + extra_outputs: list[Optional[Sequence[Optional[torch.Tensor]]]] = [ + None + ] * fuser._num_basic_ops for op, basic_op_idxs in fuser._forward_ops: # Set if backward op is required @@ -173,6 +175,26 @@ def forward( f"but got {len(ys)}" ) for output_idx, y in enumerate(ys): + if y is None: + # Extra output can be None if it is not required by any operations outside the fusion + # and is not required to be outputted to the caller. + output_to_caller = fuser._basic_op_extra_output_to_caller[idx][output_idx] + consumers = fuser._basic_op_extra_output_consumers[idx][output_idx] + needed_outside_fusion = any( + consumer_idx not in basic_op_idxs for consumer_idx in consumers + ) + if output_to_caller: + raise RuntimeError( + f"Op {idx} extra output {output_idx} is public, " + f"but {type(op).__name__} returned None" + ) + if needed_outside_fusion: + raise RuntimeError( + f"Op {idx} extra output {output_idx} is required by an " + "operation outside its forward fusion, " + f"but {type(op).__name__} returned None" + ) + continue if ( set_output_requires_grad and idx >= fuser.first_op_requiring_backward @@ -181,8 +203,11 @@ def forward( y.requires_grad_(True) extra_outputs[idx] = ys - # Flatten list of extra outputs - extra_outputs_flat = [y for ys in extra_outputs for y in ys] + # Collect caller-visible extra outputs in basic-op and slot order. + extra_outputs_flat = [ + extra_outputs[op_idx][output_idx] + for op_idx, output_idx in fuser._public_extra_output_slots + ] # Save context for backward pass if func_ctx is not None: @@ -214,13 +239,17 @@ def forward( func_ctx.basic_op_num_params = fuser._basic_op_num_params func_ctx.num_extra_outputs = len(extra_outputs_flat) func_ctx.external_extra_input_slots = fuser._external_extra_input_slots + func_ctx.public_extra_output_slots = fuser._public_extra_output_slots func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels - func_ctx.basic_op_extra_output_is_internal = fuser._basic_op_extra_output_is_internal + func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward - for tensor in [x] + extra_outputs_flat: + for tensor in itertools.chain( + (x,), + (y for ys in extra_outputs for y in ys if y is not None), + ): tensor._do_not_clear = True if set_output_requires_grad: @@ -255,7 +284,7 @@ def backward( # Channel wiring saved from forward basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels - basic_op_extra_output_is_internal = func_ctx.basic_op_extra_output_is_internal + basic_op_extra_output_consumers = func_ctx.basic_op_extra_output_consumers basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources # Place caller-provided extra-output grads into their basic-op slots. @@ -265,13 +294,14 @@ def backward( f"Expected grads for {func_ctx.num_extra_outputs} extra tensor outputs, " f"but got {len(grad_extra_outputs)}" ) - basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [] - for op in basic_ops: - grads, grad_extra_outputs = _split_tuple( - grad_extra_outputs, - op.num_extra_outputs, - ) - basic_op_grad_extra_outputs.append(list(grads)) + basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_outputs for op in basic_ops + ] + for grad, (op_idx, output_idx) in zip( + grad_extra_outputs, + func_ctx.public_extra_output_slots, + ): + basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops dx = grad_output @@ -289,7 +319,7 @@ def backward( # each internal channel. for idx in basic_op_idxs: for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): - if basic_op_extra_output_is_internal[idx][output_idx]: + if basic_op_extra_output_consumers[idx][output_idx]: channel_grad = channel_grads.get(channel) if channel_grad is not None: output_grad = basic_op_grad_extra_outputs[idx][output_idx] @@ -414,10 +444,14 @@ def __init__( self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ list(op._extra_output_channels) for op in basic_ops ] - self._basic_op_extra_output_is_internal: list[list[bool]] = [ - [False] * op.num_extra_outputs for op in basic_ops + self._basic_op_extra_output_to_caller: list[list[bool]] = [ + list(op._extra_output_to_caller) for op in basic_ops + ] + self._basic_op_extra_output_consumers: list[list[list[int]]] = [ + [[] for _ in range(op.num_extra_outputs)] for op in basic_ops ] self._external_extra_input_slots: list[tuple[int, int]] = [] + self._public_extra_output_slots: list[tuple[int, int]] = [] # Find channel producers and reject ambiguous names. channel_producers: dict[str, tuple[int, int]] = {} @@ -433,19 +467,19 @@ def __init__( ) channel_producers[channel] = (op_idx, output_idx) - # Resolve inputs. A channel with an earlier producer is internal. - # Every input without an earlier producer is a separate public input. - consumed_channels: set[str] = set() + # Resolve inputs. Named inputs must have an earlier producer; + # unnamed inputs remain public. for op_idx, op in enumerate(basic_ops): for input_idx, channel in enumerate(op._extra_input_channels): if channel is None: self._external_extra_input_slots.append((op_idx, input_idx)) continue producer = channel_producers.get(channel) - # If no producer for named channel, this is a public input. if producer is None: - self._external_extra_input_slots.append((op_idx, input_idx)) - continue + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no producer" + ) producer_idx, _ = producer if producer_idx >= op_idx: raise ValueError( @@ -453,13 +487,14 @@ def __init__( f"({type(op).__name__}) has no earlier producer" ) self._basic_op_extra_input_sources[op_idx][input_idx] = producer - consumed_channels.add(channel) + producer_idx, output_idx = producer + self._basic_op_extra_output_consumers[producer_idx][output_idx].append(op_idx) - # All extra outputs remain public, including outputs consumed internally. + # Record caller-visible outputs in stable basic-op and slot order. for op_idx, op in enumerate(basic_ops): - for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): - if channel is not None and channel in consumed_channels: - self._basic_op_extra_output_is_internal[op_idx][output_idx] = True + for output_idx in range(op.num_extra_outputs): + if self._basic_op_extra_output_to_caller[op_idx][output_idx]: + self._public_extra_output_slots.append((op_idx, output_idx)) # Every channel-bound extra input must be wired to a matching producer # extra output. External slots remain unbound (source is None). @@ -476,7 +511,11 @@ def __init__( ) continue if source is None: - continue + raise RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r} " + "without a producer source" + ) producer_idx, output_idx = source producer_channel = self._basic_op_extra_output_channels[producer_idx][output_idx] if producer_channel != channel: diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 85e5641078..2325002645 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -85,11 +85,11 @@ def fuser_forward( basic_op_ctxs: list[OperationContext], input_: torch.Tensor, *, - basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + basic_op_extra_inputs: Sequence[Sequence[Optional[torch.Tensor]]], prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: """Forward pass This op is either a basic op or the fusion of basic ops, so @@ -104,8 +104,9 @@ def fuser_forward( Contexts for basic operations input_: torch.Tensor Input tensor - basic_op_extra_inputs: list of torch.Tensor - Extra tensor inputs to basic operations + basic_op_extra_inputs: sequence of sequences of torch.Tensor + Extra tensor inputs to basic operations. An internal input + owned by this fused operation may be ``None``. prev_op_grad_output_quantizer: Quantizer, optional The grad_output_quantizer of the preceeding operation next_op_input_quantizer: Quantizer, optional @@ -118,8 +119,9 @@ def fuser_forward( ------- torch.Tensor: Output tensor. - Sequence of torch.Tensor: - Extra tensor outputs from basic operations. + Sequence of sequences of torch.Tensor: + Extra tensor outputs from basic operations. A non-public + channel owned by this fused operation may be ``None``. """ raise NotImplementedError( @@ -131,7 +133,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]], ) -> tuple[ torch.Tensor, Iterable[Iterable[Optional[torch.Tensor]]], @@ -191,6 +193,7 @@ def __init__(self) -> None: # Unbound slots remain public inputs/outputs, preserving the original API. self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs + self._extra_output_to_caller: list[bool] = [True] * self.num_extra_outputs self._extra_channels_version = 0 # Objects for quantization @@ -217,11 +220,19 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp self._extra_channels_version += 1 return self - def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + def set_extra_output_channel( + self, + index: int, + channel: Optional[str], + *, + output_to_caller: bool = True, + ) -> BasicOperation: """Bind an extra output slot to an internal fuser channel. - A bound slot can feed one or more later operations and is not returned - as a public extra output. Passing ``None`` removes the binding. + A bound slot can feed one or more later operations. By default, the + output is also returned to the caller. Set ``output_to_caller=False`` + to keep it internal to the fuser. Passing ``channel=None`` removes the + binding and restores the output as public. """ if not 0 <= index < self.num_extra_outputs: raise IndexError( @@ -230,9 +241,17 @@ def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicO ) if channel is not None and (not isinstance(channel, str) or not channel): raise ValueError("Extra output channel must be a non-empty string or None") - if self._extra_output_channels[index] == channel: + if not isinstance(output_to_caller, bool): + raise TypeError("output_to_caller must be a bool") + if channel is None: + output_to_caller = True + if ( + self._extra_output_channels[index] == channel + and self._extra_output_to_caller[index] == output_to_caller + ): return self self._extra_output_channels[index] = channel + self._extra_output_to_caller[index] = output_to_caller self._extra_channels_version += 1 return self From 2eb21edfc07edac1ac2a3c913b18f2952f926e19 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 11 Aug 2026 08:43:21 +0000 Subject: [PATCH 37/83] cleanup Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 7 +++---- tests/pytorch/test_fusible_ops.py | 10 +--------- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index f7e9acc0a0..5b9bd838ca 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -156,8 +156,8 @@ Extra tensor channels 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 and may optionally -be returned from the ``Sequential``. Extra input tensors with a specified +specified channel can be consumed by later operations in the same ``Sequential`` +and may optionally be returned to the caller. Extra input tensors with a specified channel are accessed internally instead of being provided as arguments to ``Sequential``. @@ -264,8 +264,7 @@ be placed in the same ``OperationFuser``. The following conditions apply to extra tensor channels: - Every named extra input must have a matching producer earlier in the - same fuser. Missing producers, backward edges, and cycles are not - supported. Leave an extra input unnamed when the caller should provide it. + same fuser. Leave an extra input unnamed when the caller should provide it. - 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 returned as a public diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 57a3b29b94..a06c6da06f 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -564,6 +564,7 @@ def fuse_residual(ops, **unused): and isinstance(ops[2], te_ops.AddExtraInput) ): # We want to enable this fusion just for this test. + # Hence disable it after fusing it once in the test. FusedResidual._enabled = False return [FusedResidual(*ops)] return ops @@ -781,15 +782,6 @@ def test_duplicate_extra_output_channel_names(self, layout: str) -> None: with pytest.raises(ValueError, match="multiple producers"): OperationFuser(ops) - def test_named_extra_output_without_consumer_is_public(self, size: int = 16) -> None: - """A named output remains public when its fuser has no consumer.""" - producer = te_ops.MakeExtraOutput() - producer.set_extra_output_channel(0, "orphan") - x = torch.rand((size,), requires_grad=True) - y, extra = producer(x) - torch.testing.assert_close(y, x) - torch.testing.assert_close(extra, x) - def test_one_extra_input_has_single_source(self, size: int = 16) -> None: """Rebinding selects one source and leaves the other output public.""" producer_a = te_ops.MakeExtraOutput() From 212a4602c22f2e8009f11cbb66212caf0c32e0dd Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 11 Aug 2026 19:51:38 +0000 Subject: [PATCH 38/83] Restore ep_reference to its originally created content. --- transformer_engine/pytorch/ep_reference.py | 111 ++++----------------- 1 file changed, 20 insertions(+), 91 deletions(-) diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index efb8d5c471..280c4f0d21 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -297,7 +297,6 @@ def _decode_tensor( name: str, expected_shape: Tuple[int, ...], quantized_axis: int, - dtype: torch.dtype, ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: @@ -306,24 +305,19 @@ def _decode_tensor( ) if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") - return tensor.dequantize(dtype=dtype) + return tensor.dequantize() if tuple(tensor.shape) != expected_shape: raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") if not tensor.is_floating_point(): raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") - return tensor.to(dtype=dtype) + return tensor.float() -def _format_round_trip( - tensor: torch.Tensor, - format: MoeFormat, - *, - dtype: torch.dtype, -) -> torch.Tensor: +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: if format is MoeFormat.BF16: - return tensor.to(torch.bfloat16).to(dtype) - return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=-1).dequantize() class MoeEpReference: @@ -349,7 +343,6 @@ def __init__( apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, generate_c: bool = False, - compute_dtype: torch.dtype = torch.float32, ) -> None: for name, value in ( ("num_experts", num_experts), @@ -363,11 +356,6 @@ def __init__( raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") if max_tokens_per_rank is not None and max_tokens_per_rank < 0: raise ValueError("max_tokens_per_rank must be non-negative") - if compute_dtype not in (torch.float32, torch.bfloat16): - raise ValueError( - "compute_dtype must be torch.float32 or torch.bfloat16, " - f"got {compute_dtype}" - ) if ep_group is None: ep_size, ep_rank = 1, 0 @@ -397,7 +385,6 @@ def __init__( self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) - self.compute_dtype = compute_dtype for name, fmt in ( ("output_format", self.output_format), @@ -418,8 +405,7 @@ def __repr__(self) -> str: f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " - f"output={self.output_format.value}, combine={self.combine_format.value}, " - f"compute_dtype={self.compute_dtype})" + f"output={self.output_format.value}, combine={self.combine_format.value})" ) def _collective_device(self, device: torch.device) -> torch.device: @@ -509,7 +495,7 @@ def _run_local_experts( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: output = torch.empty( (tokens.shape[0], self.hidden_size), - dtype=self.compute_dtype, + dtype=torch.float32, device=tokens.device, ) fc1_c_rows = [] if self.generate_c else None @@ -527,21 +513,13 @@ def _run_local_experts( gate = gate.clamp(max=self.gate_up_clamp) up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) intermediate = F.silu(gate) * up - weights = ( - route_weight.index_select(0, positions) - .to(dtype=self.compute_dtype) - .unsqueeze(-1) - ) + weights = route_weight.index_select(0, positions).unsqueeze(-1) if self.apply_topk_in_fc1: intermediate = intermediate * weights expert_output = intermediate @ fc2_weight[expert] if not self.apply_topk_in_fc1: expert_output = expert_output * weights - expert_output = _format_round_trip( - expert_output, - self.combine_format, - dtype=self.compute_dtype, - ) + expert_output = _format_round_trip(expert_output, self.combine_format) output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: @@ -561,13 +539,7 @@ def __call__( fc2_weight: MoeTensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, - *, - return_dispatch_metadata: bool = False, - ) -> Union[ - MoeTensor, - Tuple[MoeTensor, torch.Tensor, torch.Tensor], - Tuple[MoeTensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], - ]: + ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: """Run dispatch, local experts, return routing, top-k reduce, and encode. Shapes: @@ -623,21 +595,18 @@ def __call__( name="activation", expected_shape=(token_count, self.hidden_size), quantized_axis=1, - dtype=self.compute_dtype, ) fc1_float = _decode_tensor( fc1_weight, name="fc1_weight", expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), quantized_axis=1, - dtype=self.compute_dtype, ) fc2_float = _decode_tensor( fc2_weight, name="fc2_weight", expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), quantized_axis=1, - dtype=self.compute_dtype, ) plan = self._dispatch_plan(topk_idx, topk_weights) @@ -651,7 +620,6 @@ def __call__( recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) route_metadata = None - grouped_recv_weight = None if self.generate_c: recv_src_rank = torch.repeat_interleave( torch.arange(self.ep_size, device=device), @@ -662,7 +630,6 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - grouped_recv_weight = recv_weight.index_select(0, fc1_c_order) route_metadata = ( torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) .index_select(0, fc1_c_order) @@ -682,7 +649,7 @@ def __call__( returned = self._all_to_all(recv_output, recv_counts, send_counts) combine_plane = torch.zeros( (token_count * self.top_k, self.hidden_size), - dtype=self.compute_dtype, + dtype=torch.float32, device=device, ) send_flat_slot = send_token_idx * self.top_k + send_slot_idx @@ -694,15 +661,7 @@ def __call__( else: output = quantize_blockwise(reduced, self.output_format, axis=-1) if self.generate_c: - if return_dispatch_metadata: - tokens_per_expert = torch.bincount( - recv_expert.to(torch.int64), - minlength=self.experts_per_rank, - ) - return output, fc1_c, route_metadata, tokens_per_expert, grouped_recv_weight return output, fc1_c, route_metadata - if return_dispatch_metadata: - raise ValueError("return_dispatch_metadata=True requires generate_c=True") return output def backward( @@ -715,7 +674,6 @@ def backward( topk_weights: torch.Tensor, fc1_c: torch.Tensor, route_metadata: torch.Tensor, - grad_recv_topk_weights: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Backward pass consuming the ``generate_c=True`` stash. @@ -730,8 +688,7 @@ def backward( ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, - grad_topk_weights)``. Activation and expert weight gradients use - ``compute_dtype``; router-weight gradients remain float32. + grad_topk_weights)`` in float32. """ if not self.generate_c: @@ -746,13 +703,6 @@ def backward( ) if not grad_output.is_floating_point(): raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") - if grad_recv_topk_weights is not None: - expected_shape = (int(route_metadata.shape[0]),) - if tuple(grad_recv_topk_weights.shape) != expected_shape: - raise ValueError( - "grad_recv_topk_weights shape must be " - f"{expected_shape}, got {tuple(grad_recv_topk_weights.shape)}" - ) device = _tensor_device(activation) two_i = 2 * self.intermediate_size @@ -761,21 +711,18 @@ def backward( name="activation", expected_shape=(token_count, self.hidden_size), quantized_axis=1, - dtype=self.compute_dtype, ) fc1_float = _decode_tensor( fc1_weight, name="fc1_weight", expected_shape=(self.experts_per_rank, self.hidden_size, two_i), quantized_axis=1, - dtype=self.compute_dtype, ) fc2_float = _decode_tensor( fc2_weight, name="fc2_weight", expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), quantized_axis=1, - dtype=self.compute_dtype, ) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): raise ValueError( @@ -787,7 +734,7 @@ def backward( # along the identical forward routes. plan = self._dispatch_plan(topk_idx, topk_weights) send_counts, recv_counts = plan.send_counts, plan.recv_counts - grad_output_float = grad_output.to(dtype=self.compute_dtype) + grad_output_float = grad_output.float() recv_tokens = self._all_to_all( activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts ) @@ -814,13 +761,9 @@ def backward( dy_rows = torch.empty_like(recv_grad) dy_rows.index_copy_(0, perm, recv_grad) - c_rows = fc1_c.to(dtype=self.compute_dtype) + c_rows = fc1_c.float() expert_rows = metadata[:, 0] - d_x_rows = torch.zeros( - (local_routes, self.hidden_size), - dtype=self.compute_dtype, - device=device, - ) + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) grad_fc1 = torch.zeros_like(fc1_float) grad_fc2 = torch.zeros_like(fc2_float) @@ -830,11 +773,7 @@ def backward( continue c = c_rows.index_select(0, positions) x = x_rows.index_select(0, positions) - w = ( - w_rows.index_select(0, positions) - .to(dtype=self.compute_dtype) - .unsqueeze(-1) - ) + w = w_rows.index_select(0, positions).unsqueeze(-1) d_y = dy_rows.index_select(0, positions) gate, up = c.split(self.intermediate_size, dim=-1) @@ -857,12 +796,10 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).float() + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) else: d_h = d_h_fc2 - d_w_rows[positions] = ( - d_y * (h @ fc2_float[expert]) - ).sum(dim=-1).float() + d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s @@ -877,17 +814,9 @@ def backward( # Return the route gradients to their source ranks and scatter-add. returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) - if grad_recv_topk_weights is not None: - d_w_rows = d_w_rows + grad_recv_topk_weights.to( - device=device, - dtype=torch.float32, - ) - recv_dw = d_w_rows.index_select(0, perm) - returned_dw = self._all_to_all(recv_dw, recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) grad_activation = torch.zeros( - (token_count, self.hidden_size), - dtype=self.compute_dtype, - device=device, + (token_count, self.hidden_size), dtype=torch.float32, device=device ) grad_activation.index_add_(0, plan.send_token_idx, returned_dx) grad_topk_weights = torch.zeros( From c7dacdfe17c781e9af45429fa9fffc8946f2fa67 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 11 Aug 2026 19:55:56 +0000 Subject: [PATCH 39/83] Support BF16 grouped MLP in ep_reference and keep MOE extras internal. Use compute_dtype for apples-to-apples TE comparisons, and set output_to_caller=False so fused MOE does not materialize internal dispatch extras for the Sequential caller. Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 11 +-- tests/pytorch/test_ep_reference.py | 12 +-- transformer_engine/pytorch/ep_reference.py | 77 ++++++++++++++----- .../pytorch/ops/fused/moe_ep.py | 12 +-- 4 files changed, 73 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 750a31e32e..05f5888f09 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -282,8 +282,8 @@ def _make_moe_model(self, *, fusion_barrier=False): ) combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) - dispatch.set_extra_output_channel(0, "tokens_per_expert") - dispatch.set_extra_output_channel(1, "routing_weights") + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") @@ -515,12 +515,12 @@ def test_bf16_moe_sequential_fusion(self): fused_topk_weights = topk_weights.detach().clone().requires_grad_(True) unfused_topk_weights = topk_weights.detach().clone().requires_grad_(True) - fused_out, fused_counts, fused_recv_weights = fused( + fused_out = fused( fused_tokens, topk_idx, fused_topk_weights, ) - unfused_out, unfused_counts, unfused_recv_weights = unfused( + unfused_out = unfused( unfused_tokens, topk_idx, unfused_topk_weights, @@ -533,9 +533,6 @@ def test_bf16_moe_sequential_fusion(self): self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops)) self.assertEqual(fused_out.dtype, torch.bfloat16) self.assertEqual(unfused_out.dtype, torch.bfloat16) - torch.testing.assert_close(fused_counts, unfused_counts, rtol=0, atol=0) - self.assertEqual(fused_recv_weights.dtype, torch.float32) - self.assertEqual(unfused_recv_weights.dtype, torch.float32) dy = ( torch.randn( diff --git a/tests/pytorch/test_ep_reference.py b/tests/pytorch/test_ep_reference.py index 47a0731ff6..0484b13a53 100644 --- a/tests/pytorch/test_ep_reference.py +++ b/tests/pytorch/test_ep_reference.py @@ -83,8 +83,8 @@ def test_single_rank_bf16_moe_fusion_forward_backward(): activation = te_ops.ScaledSwiGLU() fc2 = te_ops.GroupedLinear(2, 4, 8, bias=False, device="cpu", dtype=torch.bfloat16) combine = te_ops.Combine(buffer, num_local_tokens=4) - dispatch.set_extra_output_channel(0, "tokens_per_expert") - dispatch.set_extra_output_channel(1, "routing_weights") + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") @@ -93,14 +93,10 @@ def test_single_rank_bf16_moe_fusion_forward_backward(): x = torch.randn(4, 8, dtype=torch.bfloat16, requires_grad=True) topk_idx = torch.tensor([[0], [1], [0], [1]], dtype=torch.int64) topk_weights = torch.ones(4, 1, dtype=torch.float32, requires_grad=True) - output, counts, recv_weights = model(x, topk_idx, topk_weights) - torch.autograd.backward( - (output, recv_weights), - (torch.ones_like(output), torch.ones_like(recv_weights)), - ) + output = model(x, topk_idx, topk_weights) + output.backward(torch.ones_like(output)) assert isinstance(model._module_groups[0]._forward_ops[0][0], FusedMoeEp) - assert counts.dtype is torch.int64 assert output.dtype is torch.bfloat16 assert x.grad.dtype is torch.bfloat16 assert topk_weights.grad.dtype is torch.float32 diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 280c4f0d21..9261b05f06 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -297,6 +297,7 @@ def _decode_tensor( name: str, expected_shape: Tuple[int, ...], quantized_axis: int, + dtype: torch.dtype, ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: @@ -305,19 +306,24 @@ def _decode_tensor( ) if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") - return tensor.dequantize() + return tensor.dequantize(dtype=dtype) if tuple(tensor.shape) != expected_shape: raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") if not tensor.is_floating_point(): raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") - return tensor.float() + return tensor.to(dtype=dtype) -def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: +def _format_round_trip( + tensor: torch.Tensor, + format: MoeFormat, + *, + dtype: torch.dtype, +) -> torch.Tensor: if format is MoeFormat.BF16: - return tensor.to(torch.bfloat16).float() - return quantize_blockwise(tensor, format, axis=-1).dequantize() + return tensor.to(torch.bfloat16).to(dtype) + return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) class MoeEpReference: @@ -343,6 +349,7 @@ def __init__( apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, generate_c: bool = False, + compute_dtype: torch.dtype = torch.float32, ) -> None: for name, value in ( ("num_experts", num_experts), @@ -356,6 +363,11 @@ def __init__( raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") if max_tokens_per_rank is not None and max_tokens_per_rank < 0: raise ValueError("max_tokens_per_rank must be non-negative") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError( + "compute_dtype must be torch.float32 or torch.bfloat16, " + f"got {compute_dtype}" + ) if ep_group is None: ep_size, ep_rank = 1, 0 @@ -385,6 +397,7 @@ def __init__( self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) + self.compute_dtype = compute_dtype for name, fmt in ( ("output_format", self.output_format), @@ -405,7 +418,8 @@ def __repr__(self) -> str: f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " - f"output={self.output_format.value}, combine={self.combine_format.value})" + f"output={self.output_format.value}, combine={self.combine_format.value}, " + f"compute_dtype={self.compute_dtype})" ) def _collective_device(self, device: torch.device) -> torch.device: @@ -495,7 +509,7 @@ def _run_local_experts( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: output = torch.empty( (tokens.shape[0], self.hidden_size), - dtype=torch.float32, + dtype=self.compute_dtype, device=tokens.device, ) fc1_c_rows = [] if self.generate_c else None @@ -513,13 +527,21 @@ def _run_local_experts( gate = gate.clamp(max=self.gate_up_clamp) up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) intermediate = F.silu(gate) * up - weights = route_weight.index_select(0, positions).unsqueeze(-1) + weights = ( + route_weight.index_select(0, positions) + .to(dtype=self.compute_dtype) + .unsqueeze(-1) + ) if self.apply_topk_in_fc1: intermediate = intermediate * weights expert_output = intermediate @ fc2_weight[expert] if not self.apply_topk_in_fc1: expert_output = expert_output * weights - expert_output = _format_round_trip(expert_output, self.combine_format) + expert_output = _format_round_trip( + expert_output, + self.combine_format, + dtype=self.compute_dtype, + ) output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: @@ -595,18 +617,21 @@ def __call__( name="activation", expected_shape=(token_count, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) fc1_float = _decode_tensor( fc1_weight, name="fc1_weight", expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), quantized_axis=1, + dtype=self.compute_dtype, ) fc2_float = _decode_tensor( fc2_weight, name="fc2_weight", expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) plan = self._dispatch_plan(topk_idx, topk_weights) @@ -649,7 +674,7 @@ def __call__( returned = self._all_to_all(recv_output, recv_counts, send_counts) combine_plane = torch.zeros( (token_count * self.top_k, self.hidden_size), - dtype=torch.float32, + dtype=self.compute_dtype, device=device, ) send_flat_slot = send_token_idx * self.top_k + send_slot_idx @@ -688,7 +713,8 @@ def backward( ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, - grad_topk_weights)`` in float32. + grad_topk_weights)``. Activation and expert weight gradients use + ``compute_dtype``; router-weight gradients remain float32. """ if not self.generate_c: @@ -711,18 +737,21 @@ def backward( name="activation", expected_shape=(token_count, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) fc1_float = _decode_tensor( fc1_weight, name="fc1_weight", expected_shape=(self.experts_per_rank, self.hidden_size, two_i), quantized_axis=1, + dtype=self.compute_dtype, ) fc2_float = _decode_tensor( fc2_weight, name="fc2_weight", expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), quantized_axis=1, + dtype=self.compute_dtype, ) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): raise ValueError( @@ -734,7 +763,7 @@ def backward( # along the identical forward routes. plan = self._dispatch_plan(topk_idx, topk_weights) send_counts, recv_counts = plan.send_counts, plan.recv_counts - grad_output_float = grad_output.float() + grad_output_float = grad_output.to(dtype=self.compute_dtype) recv_tokens = self._all_to_all( activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts ) @@ -761,9 +790,13 @@ def backward( dy_rows = torch.empty_like(recv_grad) dy_rows.index_copy_(0, perm, recv_grad) - c_rows = fc1_c.float() + c_rows = fc1_c.to(dtype=self.compute_dtype) expert_rows = metadata[:, 0] - d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_x_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=self.compute_dtype, + device=device, + ) d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) grad_fc1 = torch.zeros_like(fc1_float) grad_fc2 = torch.zeros_like(fc2_float) @@ -773,7 +806,11 @@ def backward( continue c = c_rows.index_select(0, positions) x = x_rows.index_select(0, positions) - w = w_rows.index_select(0, positions).unsqueeze(-1) + w = ( + w_rows.index_select(0, positions) + .to(dtype=self.compute_dtype) + .unsqueeze(-1) + ) d_y = dy_rows.index_select(0, positions) gate, up = c.split(self.intermediate_size, dim=-1) @@ -796,10 +833,12 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).float() else: d_h = d_h_fc2 - d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) + d_w_rows[positions] = ( + d_y * (h @ fc2_float[expert]) + ).sum(dim=-1).float() d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s @@ -816,7 +855,9 @@ def backward( returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) grad_activation = torch.zeros( - (token_count, self.hidden_size), dtype=torch.float32, device=device + (token_count, self.hidden_size), + dtype=self.compute_dtype, + device=device, ) grad_activation.index_add_(0, plan.send_token_idx, returned_dx) grad_topk_weights = torch.zeros( diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 5d5b861f8f..ea1027bf06 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -123,7 +123,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: del prev_op_grad_output_quantizer, next_op_input_quantizer if input_.dtype is not torch.bfloat16: raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") @@ -135,13 +135,12 @@ def fuser_forward( raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") fc1_weight = _reference_weights(self.fc1) fc2_weight = _reference_weights(self.fc2) - output, fc1_c, route_metadata, tokens_per_expert, recv_topk_weights = self._reference( + output, fc1_c, route_metadata = self._reference( input_, fc1_weight, fc2_weight, topk_idx, topk_weights, - return_dispatch_metadata=True, ) if any(ctx.requires_grad for ctx in basic_op_ctxs): @@ -155,8 +154,10 @@ def fuser_forward( route_metadata, ) + # Dispatch extras are channel-bound with output_to_caller=False and are + # only consumed by ops inside this fusion, so they need not be materialized. return output, [ - (tokens_per_expert, recv_topk_weights), + (None, None), (), (), (), @@ -174,6 +175,7 @@ def fuser_backward( Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: + del basic_op_grad_extra_outputs ( input_, fc1_weight, @@ -183,7 +185,6 @@ def fuser_backward( fc1_c, route_metadata, ) = basic_op_ctxs[0].saved_tensors - grad_recv_topk_weights = basic_op_grad_extra_outputs[0][1] grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._reference.backward( grad_output, input_, @@ -193,7 +194,6 @@ def fuser_backward( topk_weights, fc1_c, route_metadata, - grad_recv_topk_weights=grad_recv_topk_weights, ) fc1_param_grads = [ From 00b48cfc60b188f2f337de50adf8cf3a923a9a70 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 11 Aug 2026 21:31:51 +0000 Subject: [PATCH 40/83] update test to have internal extra_out Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 43 ++++--- tests/pytorch/test_ep_reference.py | 105 ------------------ transformer_engine/pytorch/ep_reference.py | 36 +++--- .../pytorch/ops/fused/moe_ep.py | 27 +++++ 4 files changed, 76 insertions(+), 135 deletions(-) delete mode 100644 tests/pytorch/test_ep_reference.py diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 05f5888f09..2af094ea98 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -260,7 +260,14 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) - def _make_moe_model(self, *, fusion_barrier=False): + def _make_moe_model(self, *, fuse_ops=True): + """Build a BF16 EP MoE Sequential. + + With ``fuse_ops=True``, dispatch routing extras stay internal + (``output_to_caller=False``) so :class:`FusedMoeEp` can claim the + sequence. With ``fuse_ops=False``, those extras are returned to the + caller, which blocks fusion. + """ buffer = self._make_buffer() dispatch = te_ops.Dispatch(buffer) fc1 = te_ops.GroupedLinear( @@ -282,16 +289,16 @@ def _make_moe_model(self, *, fusion_barrier=False): ) combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) - dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) - dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) + dispatch.set_extra_output_channel( + 0, "tokens_per_expert", output_to_caller=not fuse_ops + ) + dispatch.set_extra_output_channel( + 1, "routing_weights", output_to_caller=not fuse_ops + ) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") - ops = [dispatch, fc1, activation, fc2] - if fusion_barrier: - ops.append(te_ops.Identity()) - ops.append(combine) - return te_ops.Sequential(*ops), fc1, fc2 + return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 # Prepare @@ -482,12 +489,16 @@ def test_caller_provides_grad_expert_out(self): @_eager_test_include def test_bf16_moe_sequential_fusion(self): - """Reference-backed fusion matches the unfused BF16 EP MoE sequence.""" + """Reference-backed fusion matches the unfused BF16 EP MoE sequence. + + ``fuse_ops=True`` keeps dispatch routing extras internal so fusion + fires; ``fuse_ops=False`` returns them to the caller and blocks it. + """ if not EAGER: self.skipTest("variable-size reference comparison requires eager EP mode") - fused, fused_fc1, fused_fc2 = self._make_moe_model() - unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fusion_barrier=True) + fused, fused_fc1, fused_fc2 = self._make_moe_model(fuse_ops=True) + unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fuse_ops=False) generator = torch.Generator(device=self.cfg.device) generator.manual_seed(3100 + self.cfg.rank) with torch.no_grad(): @@ -520,7 +531,7 @@ def test_bf16_moe_sequential_fusion(self): topk_idx, fused_topk_weights, ) - unfused_out = unfused( + unfused_out, tokens_per_expert, recv_topk_weights = unfused( unfused_tokens, topk_idx, unfused_topk_weights, @@ -531,8 +542,12 @@ def test_bf16_moe_sequential_fusion(self): self.assertEqual(len(fused_forward_ops), 1) self.assertIsInstance(fused_forward_ops[0][0], FusedMoeEp) self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops)) + self.assertIsInstance(fused_out, torch.Tensor) self.assertEqual(fused_out.dtype, torch.bfloat16) self.assertEqual(unfused_out.dtype, torch.bfloat16) + self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,)) + self.assertEqual(tokens_per_expert.dtype, torch.int64) + self.assertEqual(recv_topk_weights.dtype, torch.float32) dy = ( torch.randn( @@ -548,8 +563,8 @@ def test_bf16_moe_sequential_fusion(self): torch.cuda.synchronize() # The two BF16 paths use different grouped-GEMM and reduction orders. - # Their largest observed forward absolute error is ~3.1e-5, at values - # close enough to zero that the relative tolerance does not apply. + # Observed forward abs error reaches ~3e-5 on near-zero values where + # rtol does not apply, so atol must stay above that (1e-5 is too tight). tolerances = {"rtol": 1.6e-2, "atol": 5e-5} torch.testing.assert_close(fused_out, unfused_out, **tolerances) torch.testing.assert_close(fused_tokens.grad, unfused_tokens.grad, **tolerances) diff --git a/tests/pytorch/test_ep_reference.py b/tests/pytorch/test_ep_reference.py deleted file mode 100644 index 0484b13a53..0000000000 --- a/tests/pytorch/test_ep_reference.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Tests for the pure PyTorch MoE EP reference.""" - -from types import SimpleNamespace - -import pytest -import torch - -from transformer_engine.pytorch import ops as te_ops -from transformer_engine.pytorch.ep_reference import MoeEpReference -from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp - - -@pytest.mark.parametrize("compute_dtype", (torch.float32, torch.bfloat16)) -def test_moe_ep_reference_compute_dtype(compute_dtype): - """The configured dtype controls MLP, combine, and non-router gradients.""" - generator = torch.Generator().manual_seed(1234) - activation = torch.randn(4, 8, generator=generator, dtype=torch.bfloat16) - fc1_weight = torch.randn(2, 8, 8, generator=generator, dtype=torch.bfloat16) - fc2_weight = torch.randn(2, 4, 8, generator=generator, dtype=torch.bfloat16) - topk_idx = torch.tensor([[0], [1], [0], [1]], dtype=torch.int64) - topk_weights = torch.ones(4, 1, dtype=torch.float32) - - reference = MoeEpReference( - num_experts=2, - hidden_size=8, - intermediate_size=4, - top_k=1, - generate_c=True, - compute_dtype=compute_dtype, - ) - output, fc1_c, route_metadata = reference( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - ) - grads = reference.backward( - torch.ones_like(output), - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - fc1_c, - route_metadata, - ) - - assert output.dtype is torch.bfloat16 - assert grads[0].dtype is compute_dtype - assert grads[1].dtype is compute_dtype - assert grads[2].dtype is compute_dtype - assert grads[3].dtype is torch.float32 - - -def test_moe_ep_reference_default_compute_dtype_is_fp32(): - """Preserve the reference's pre-existing FP32 compute default.""" - reference = MoeEpReference( - num_experts=1, - hidden_size=8, - intermediate_size=4, - top_k=1, - ) - assert reference.compute_dtype is torch.float32 - - -def test_single_rank_bf16_moe_fusion_forward_backward(): - """Exercise the fuser contract without requiring the NCCL EP backend.""" - buffer = SimpleNamespace( - num_local_experts=2, - hidden_dim=8, - top_k=1, - max_tokens_per_rank=4, - payload_dtype=torch.bfloat16, - eager=True, - ) - dispatch = te_ops.Dispatch(buffer) - fc1 = te_ops.GroupedLinear(2, 8, 8, bias=False, device="cpu", dtype=torch.bfloat16) - activation = te_ops.ScaledSwiGLU() - fc2 = te_ops.GroupedLinear(2, 4, 8, bias=False, device="cpu", dtype=torch.bfloat16) - combine = te_ops.Combine(buffer, num_local_tokens=4) - dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) - dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) - fc1.set_extra_input_channel(0, "tokens_per_expert") - activation.set_extra_input_channel(0, "routing_weights") - fc2.set_extra_input_channel(0, "tokens_per_expert") - model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) - - x = torch.randn(4, 8, dtype=torch.bfloat16, requires_grad=True) - topk_idx = torch.tensor([[0], [1], [0], [1]], dtype=torch.int64) - topk_weights = torch.ones(4, 1, dtype=torch.float32, requires_grad=True) - output = model(x, topk_idx, topk_weights) - output.backward(torch.ones_like(output)) - - assert isinstance(model._module_groups[0]._forward_ops[0][0], FusedMoeEp) - assert output.dtype is torch.bfloat16 - assert x.grad.dtype is torch.bfloat16 - assert topk_weights.grad.dtype is torch.float32 - for op in (fc1, fc2): - for expert in range(op.num_groups): - assert getattr(op, f"weight{expert}").grad.dtype is torch.bfloat16 diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 9261b05f06..987df16732 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -527,11 +527,7 @@ def _run_local_experts( gate = gate.clamp(max=self.gate_up_clamp) up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) intermediate = F.silu(gate) * up - weights = ( - route_weight.index_select(0, positions) - .to(dtype=self.compute_dtype) - .unsqueeze(-1) - ) + weights = route_weight.index_select(0, positions).unsqueeze(-1) if self.apply_topk_in_fc1: intermediate = intermediate * weights expert_output = intermediate @ fc2_weight[expert] @@ -642,7 +638,9 @@ def __call__( recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) - recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( + dtype=self.compute_dtype + ) route_metadata = None if self.generate_c: @@ -714,7 +712,9 @@ def backward( Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, grad_topk_weights)``. Activation and expert weight gradients use - ``compute_dtype``; router-weight gradients remain float32. + ``compute_dtype``. Router weights are cast to ``compute_dtype`` for the + activation scaling path; the returned + ``grad_topk_weights`` remain float32. """ if not self.generate_c: @@ -767,7 +767,10 @@ def backward( recv_tokens = self._all_to_all( activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts ) - recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + # Same FP32-on-wire / compute_dtype-for-activation cast as forward. + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( + dtype=self.compute_dtype + ) recv_grad = self._all_to_all( grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts ) @@ -806,11 +809,7 @@ def backward( continue c = c_rows.index_select(0, positions) x = x_rows.index_select(0, positions) - w = ( - w_rows.index_select(0, positions) - .to(dtype=self.compute_dtype) - .unsqueeze(-1) - ) + w = w_rows.index_select(0, positions).unsqueeze(-1) d_y = dy_rows.index_select(0, positions) gate, up = c.split(self.intermediate_size, dim=-1) @@ -833,12 +832,17 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).float() + d_w_rows[positions] = ( + (d_h_fc2 * h).sum(dim=-1).to(dtype=self.compute_dtype).float() + ) else: d_h = d_h_fc2 d_w_rows[positions] = ( - d_y * (h @ fc2_float[expert]) - ).sum(dim=-1).float() + (d_y * (h @ fc2_float[expert])) + .sum(dim=-1) + .to(dtype=self.compute_dtype) + .float() + ) d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index ea1027bf06..ba7a44a731 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -45,6 +45,31 @@ def _grouped_linear_supported(op: GroupedLinear) -> bool: ) +def _routing_extras_internal( + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`MoeEpReference`, so it can only replace the sequence + when those two outputs feed exactly these ops and are not returned to the + caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: if recipe is not None or len(window) != 5: return False @@ -66,6 +91,8 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo return False if activation.activation_recompute_in_mlp or activation.glu_interleave_size is not None: return False + if not _routing_extras_internal(dispatch, fc1, activation, fc2): + return False return ( fc1.num_groups == buffer.num_local_experts and fc2.num_groups == buffer.num_local_experts From db244954149a3bd8b34a3fd1890e4f429ab96e39 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:36:17 +0000 Subject: [PATCH 41/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_ep.py | 8 ++------ transformer_engine/pytorch/ep_reference.py | 12 +++--------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 2af094ea98..d8aab26e11 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -289,12 +289,8 @@ def _make_moe_model(self, *, fuse_ops=True): ) combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) - dispatch.set_extra_output_channel( - 0, "tokens_per_expert", output_to_caller=not fuse_ops - ) - dispatch.set_extra_output_channel( - 1, "routing_weights", output_to_caller=not fuse_ops - ) + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 987df16732..8a8d58ad0f 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -365,8 +365,7 @@ def __init__( raise ValueError("max_tokens_per_rank must be non-negative") if compute_dtype not in (torch.float32, torch.bfloat16): raise ValueError( - "compute_dtype must be torch.float32 or torch.bfloat16, " - f"got {compute_dtype}" + f"compute_dtype must be torch.float32 or torch.bfloat16, got {compute_dtype}" ) if ep_group is None: @@ -832,16 +831,11 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = ( - (d_h_fc2 * h).sum(dim=-1).to(dtype=self.compute_dtype).float() - ) + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).to(dtype=self.compute_dtype).float() else: d_h = d_h_fc2 d_w_rows[positions] = ( - (d_y * (h @ fc2_float[expert])) - .sum(dim=-1) - .to(dtype=self.compute_dtype) - .float() + (d_y * (h @ fc2_float[expert])).sum(dim=-1).to(dtype=self.compute_dtype).float() ) d_g = d_h * u * (sig * (1 + g * (1 - sig))) From 3c37643d926b21b50e9103fe7e6007cc896ab9f3 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 12 Aug 2026 03:11:55 +0000 Subject: [PATCH 42/83] extra output grad can be None, equivalent to zero Signed-off-by: Varun Thumbe address review comment Signed-off-by: Varun Thumbe address review comments Signed-off-by: Varun Thumbe --- docs/examples/op_fuser/op_fuser.rst | 8 ++- tests/pytorch/test_fusible_ops.py | 67 ++++++++++++++++--- .../pytorch/ops/basic/make_extra_output.py | 7 +- .../pytorch/ops/fused/backward_add_rmsnorm.py | 8 ++- .../pytorch/ops/fused/backward_linear_add.py | 10 +-- .../pytorch/ops/fused/moe_ep.py | 25 +++++++ transformer_engine/pytorch/ops/fuser.py | 20 ++---- transformer_engine/pytorch/ops/op.py | 28 ++++++-- 8 files changed, 133 insertions(+), 40 deletions(-) diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index 5b9bd838ca..b7019b6237 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -277,14 +277,16 @@ The following conditions apply to extra tensor channels: - ``set_extra_output_channel`` accepts ``output_to_caller`` (``True`` by default). Public extra outputs are returned in their original basic-operation and slot order. Gradients supplied for a returned output - are combined with gradients from its internal channel consumers. + are combined with gradients from its internal channel consumers. An + unused extra-output gradient may be ``None`` and is treated as zero. - Set ``output_to_caller=False`` for a channel tensor that should remain internal. Removing a channel binding with ``channel=None`` restores that output as public. - Channel bindings are captured when an ``OperationFuser`` (or the - fusers inside a ``Sequential``) is first constructed. Changing + fusers inside a ``Sequential``) is constructed. Rebinding a channel with ``set_extra_input_channel`` / ``set_extra_output_channel`` afterward - requires constructing a new ``OperationFuser`` or ``Sequential``. + invalidates any fuser that captured it, so a new ``OperationFuser`` or + ``Sequential`` must be constructed before the next forward pass. Channel-connected basic operations may still be replaced by registered ``FusedOperation`` implementations. If a fused operation contains both diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index a06c6da06f..36c02a3154 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -503,6 +503,34 @@ def test_internal_residual_connection( torch.testing.assert_close(x.grad, expected_dx) torch.testing.assert_close(body.bias.grad, dy) + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_unconsumed_extra_output( + self, + output_to_caller: bool, + size: int = 16, + ) -> None: + """Unused public or hidden-unconsumed MakeExtraOutput treats None grad as zero.""" + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + if not output_to_caller: + residual.set_extra_output_channel(0, "unused", output_to_caller=False) + + model = te_ops.Sequential(residual, body) + x = torch.rand((size,), requires_grad=True) + outputs = model(x) + if output_to_caller: + y, residual_out = outputs + torch.testing.assert_close(residual_out, x) + else: + assert isinstance(outputs, torch.Tensor) + y = outputs + + torch.testing.assert_close(y, x + body.bias) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, dy) + torch.testing.assert_close(body.bias.grad, dy) + @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) @pytest.mark.parametrize("output_to_caller", (True, False)) def test_fused_internal_residual_connection( @@ -737,7 +765,7 @@ def test_set_extra_channel_rejects_invalid_name(self) -> None: ) def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> None: - """Sequential does not auto-rebuild after a channel configuration change.""" + """Rebinding a channel invalidates the fuser that captured it.""" producer = te_ops.MakeExtraOutput() consumer = te_ops.AddExtraInput() producer.set_extra_output_channel(0, "route") @@ -902,14 +930,21 @@ def op_backward(self, *args, **kwargs): def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): scale = basic_op_extra_inputs[0][0] - basic_op_ctxs[0].save_for_backward(input_, scale) + ctx = basic_op_ctxs[0] + # Match scaled activations: only compute scale grads when the + # fuser marked this fresh internal channel as requiring grad. + ctx.extra_input_requires_grad = scale.requires_grad + ctx.save_for_backward(input_, scale) return input_ * scale.unsqueeze(-1), [()] def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): del basic_op_grad_extra_outputs - input_, scale = basic_op_ctxs[0].saved_tensors + ctx = basic_op_ctxs[0] + input_, scale = ctx.saved_tensors grad_input = grad_output * scale.unsqueeze(-1) - grad_scale = (grad_output * input_).sum(dim=-1) + grad_scale = ( + (grad_output * input_).sum(dim=-1) if ctx.extra_input_requires_grad else None + ) return grad_input, [()], [(grad_scale,)] producer = MakeScale() @@ -3456,6 +3491,7 @@ def test_backward_activation_bias( @pytest.mark.parametrize("in_shape", ((-1,), (6, 16, -1))) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("zero_centered_gamma", (False, True)) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_backward_add_rmsnorm( self, *, @@ -3465,6 +3501,7 @@ def test_backward_add_rmsnorm( device: torch.device = "cuda", eps: float = 0.3, zero_centered_gamma: bool, + with_extra_grad: bool, ) -> None: """Fused backward RMNorm + add""" @@ -3503,7 +3540,10 @@ def test_backward_add_rmsnorm( else: y1_ref = x_ref / torch.sqrt(eps + var_ref) * w_ref y2_ref = x_ref - (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + if with_extra_grad: + (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + else: + (y1_ref * dy1_ref).sum().backward() # Implementation with fusible operations model = te_ops.Sequential( @@ -3520,7 +3560,10 @@ def test_backward_add_rmsnorm( model[1].weight.copy_(w_test) del w_test y1_test, y2_test = model(x_test) - (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + if with_extra_grad: + (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + else: + (y1_test * dy1_test).sum().backward() # Check that backward operations have been fused backward_ops = model._module_groups[0]._backward_ops @@ -3542,6 +3585,7 @@ def test_backward_add_rmsnorm( @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_backward_linear_add( self, *, @@ -3551,6 +3595,7 @@ def test_backward_linear_add( device: torch.device = "cuda", quantization: Optional[str], quantized_weight: bool = False, + with_extra_grad: bool, ) -> None: """Backward dgrad GEMM + add""" @@ -3598,7 +3643,10 @@ def test_backward_linear_add( # Plain PyTorch implementation y1_ref = torch.nn.functional.linear(x_ref, w_ref) y2_ref = x_ref - (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + if with_extra_grad: + (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + else: + (y1_ref * dy1_ref).sum().backward() # Implementation with fusible operations recipe = make_recipe(quantization) @@ -3618,7 +3666,10 @@ def test_backward_linear_add( del w_test with te.autocast(enabled=quantized_compute, recipe=recipe): y1_test, y2_test = model(x_test) - (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + if with_extra_grad: + (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + else: + (y1_test * dy1_test).sum().backward() # Check that backward operations have been fused backward_ops = model._module_groups[0]._backward_ops diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 5ad00e2fd1..47b7c6495d 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -80,14 +80,17 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: grad_extra_output = basic_op_grad_extra_outputs[0][0] - if self._in_place: + if grad_extra_output is None: + # Extra output is not consumed, so its gradient is zero + grad_input = grad_output + elif self._in_place: grad_extra_output += grad_output grad_input = grad_extra_output else: diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index a3c81e60c8..483749da47 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -33,7 +33,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, list[tuple[Optional[torch.Tensor], ...]], @@ -56,7 +56,11 @@ def fuser_backward( extra_grad = basic_op_grad_extra_outputs[0][0] dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) w = maybe_dequantize(rmsnorm_op.weight, dtype).view((inner_dim,)) - add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) + if extra_grad is None: + # Extra output is not consumed, so its gradient is zero + add = torch.zeros_like(dy) + else: + add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) # Compute RMSNorm backward pass dx, dw = tex.rmsnorm_bwd_add( diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index 382fecfd07..60b1320b2e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -40,7 +40,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, list[tuple[Optional[torch.Tensor], ...]], @@ -67,19 +67,21 @@ def fuser_backward( else: accumulate_into_main_grad = False - # Linear backward pass + # Linear backward pass. Skip in-place add when there is no + # extra-output gradient. grad_input = basic_op_grad_extra_outputs[0][0] + accumulate_into_grad_input = grad_input is not None grad_input, grad_weight = BasicLinear._functional_backward( grad_output=grad_output, input=x_local, weight=w, input_requires_grad=linear_op_ctx.input_requires_grad, weight_requires_grad=linear_op_ctx.weight_requires_grad, - dtype=grad_input.dtype, + dtype=None if grad_input is None else grad_input.dtype, grad_weight=grad_weight, accumulate_into_grad_weight=accumulate_into_main_grad, grad_input=grad_input, - accumulate_into_grad_input=True, + accumulate_into_grad_input=accumulate_into_grad_input, tensor_parallel_mode=linear_op.tensor_parallel_mode, tensor_parallel_group=linear_op.tensor_parallel_group, sequence_parallel=linear_op.sequence_parallel, diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index ba7a44a731..dd0f46c7ab 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -70,6 +70,31 @@ def _routing_extras_internal( ) +def _routing_extras_internal( + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`MoeEpReference`, so it can only replace the sequence + when those two outputs feed exactly these ops and are not returned to the + caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: if recipe is not None or len(window) != 5: return False diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 50c1c11ffa..9f346bd4c8 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -432,9 +432,11 @@ def __init__( basic_ops.append(op) self._num_basic_ops: int = len(basic_ops) self._basic_ops: list[BasicOperation] = basic_ops - self._basic_op_extra_channels_versions = [ - op._extra_channels_version for op in self._basic_ops - ] + # Capture channel routing from each basic op. If any op later rebinds a + # channel it flips this flag, essentially invalidating this fuser's forward pass. + self._channels_stale: bool = False + for op in self._basic_ops: + op._capturing_op_fusers.add(self) # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) @@ -544,16 +546,6 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) - def has_stale_op_channels(self) -> bool: - """Whether an operation's extra tensor channels have changed.""" - return any( - op._extra_channels_version != version - for op, version in zip( - self._basic_ops, - self._basic_op_extra_channels_versions, - ) - ) - @staticmethod def _apply_fusions( ops: Iterable[FusibleOperation], @@ -702,7 +694,7 @@ def __call__( *extra_inputs: torch.Tensor, basic_op_kwargs: Optional[list[dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: - if self.has_stale_op_channels(): + if self._channels_stale: raise RuntimeError( "Extra tensor channels changed after this OperationFuser captured " "its routing. Construct a new OperationFuser." diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 2325002645..bb02a5f0ee 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -10,6 +10,7 @@ import dataclasses import pickle from typing import Any, Optional +import weakref import torch @@ -136,8 +137,8 @@ def fuser_backward( basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]], ) -> tuple[ torch.Tensor, - Iterable[Iterable[Optional[torch.Tensor]]], - Iterable[Iterable[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], ]: """Backward pass @@ -161,9 +162,9 @@ def fuser_backward( ------- torch.Tensor: Loss gradient w.r.t. operation input - Iterable of iterable of torch.Tensor: + Sequence of sequences of torch.Tensor: Loss gradients w.r.t. parameters for basic operations - Iterable of iterable of torch.Tensor: + Sequence of sequences of torch.Tensor: Loss gradients w.r.t. extra tensor inputs to basic operations @@ -194,12 +195,25 @@ def __init__(self) -> None: self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs self._extra_output_to_caller: list[bool] = [True] * self.num_extra_outputs - self._extra_channels_version = 0 + # OperationFusers that have captured this op's channel routing. Weak + # refs avoid cycles and drop automatically when a fuser is discarded. + self._capturing_op_fusers: weakref.WeakSet = weakref.WeakSet() # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + def _invalidate_capturing_op_fusers(self) -> None: + """Mark fusers that captured this op's channels as stale. + + Channel routing is cheap to change, but any OperationFuser that + already captured it must be rebuilt. Flipping a flag here keeps the + per-call check O(1) instead of rescanning versions every forward. + """ + for fuser in self._capturing_op_fusers: + fuser._channels_stale = True + self._capturing_op_fusers.clear() + def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: """Bind an extra input slot to an internal fuser channel. @@ -217,7 +231,7 @@ def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOp if self._extra_input_channels[index] == channel: return self self._extra_input_channels[index] = channel - self._extra_channels_version += 1 + self._invalidate_capturing_op_fusers() return self def set_extra_output_channel( @@ -252,7 +266,7 @@ def set_extra_output_channel( return self self._extra_output_channels[index] = channel self._extra_output_to_caller[index] = output_to_caller - self._extra_channels_version += 1 + self._invalidate_capturing_op_fusers() return self @property From 5b555ce4ea66691dcd530379d06193f116d18938 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 17 Aug 2026 16:43:08 +0000 Subject: [PATCH 43/83] revert commit for nccl Signed-off-by: Varun Thumbe --- 3rdparty/nccl-extensions | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions index 705ca8eb38..9f47d6eb3b 160000 --- a/3rdparty/nccl-extensions +++ b/3rdparty/nccl-extensions @@ -1 +1 @@ -Subproject commit 705ca8eb38297f3a8c4af6adf4185176a1116cb9 +Subproject commit 9f47d6eb3b60962d8157a579b4caaaa4ae6b19f4 From 14d9ef9139786c9b7190cf0158af9c7edec617c5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:43:47 +0000 Subject: [PATCH 44/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_ep.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 03d3a13eb0..7a73a81223 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -162,6 +162,8 @@ def _make_moe_inputs(rank, ep_size, device="cuda"): ) topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) return topk_idx, tokens, torch.softmax(topk_logits, dim=-1) + + def _degroup_mxfp8(recv_grouped, valid_counts=None): """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert From 42c622e9996bbb380f0c67285e1caf4593db0e43 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 17 Aug 2026 16:45:07 +0000 Subject: [PATCH 45/83] more fix Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 16 ---------------- transformer_engine/pytorch/ops/op.py | 1 - 2 files changed, 17 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index f4be756a6a..3a0ad7dd3f 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -791,26 +791,10 @@ def test_extra_channels_lock_after_capture(self, size: int = 16) -> None: torch.testing.assert_close(y, 2 * x) torch.testing.assert_close(route, x) -<<<<<<< HEAD - consumer.set_extra_input_channel(0, None) - extra = torch.rand_like(x) - with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): - model(x, extra) - - model = te_ops.Sequential(producer, consumer) - y, route = model(x, extra) - torch.testing.assert_close(y, x + extra) - torch.testing.assert_close(route, x) - - producer.set_extra_output_channel(0, "route", output_to_caller=False) - with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): - model(x, extra) -======= with pytest.raises(RuntimeError, match="already captured its channel routing"): consumer.set_extra_input_channel(0, None) with pytest.raises(RuntimeError, match="already captured its channel routing"): producer.set_extra_output_channel(0, "route", output_to_caller=False) ->>>>>>> nvidia_origin/main @pytest.mark.parametrize("layout", ("two_ops", "same_op")) def test_duplicate_extra_output_channel_names(self, layout: str) -> None: diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index e7e2fdff2c..d057d46816 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -10,7 +10,6 @@ import dataclasses import pickle from typing import Any, Optional -import weakref import torch From e40eb74720d0f0d3a1f000032c5cc95b768c4b2f Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 19 Aug 2026 09:49:06 -0700 Subject: [PATCH 46/83] cudnn kernel Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 215 ++++++++++------ transformer_engine/pytorch/ep_reference.py | 243 +++++++++++++----- .../pytorch/ops/fused/moe_ep.py | 228 ++++++++++++---- 3 files changed, 497 insertions(+), 189 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 7a73a81223..fe354029a4 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Multi-process PyTorch EP tests, launched via torchrun (one process per GPU).""" +from contextlib import nullcontext import os import sys import unittest @@ -11,6 +12,7 @@ import torch import torch.distributed as dist +import transformer_engine.pytorch as te from transformer_engine.pytorch import ops as te_ops from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( @@ -26,7 +28,9 @@ _ep_combine_raw, _ep_dispatch_raw, ) -from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp +from transformer_engine.pytorch.ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp, _pack_grouped_linear_weights +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" @@ -305,33 +309,47 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) - def _make_moe_model(self, *, fuse_ops=True): - """Build a BF16 EP MoE Sequential. + def _make_moe_model( + self, + *, + fuse_ops=True, + intermediate_dim=INTERMEDIATE_DIM, + recipe=None, + ): + """Build an EP MoE Sequential. - With ``fuse_ops=True``, dispatch routing extras stay internal - (``output_to_caller=False``) so :class:`FusedMoeEp` can claim the - sequence. With ``fuse_ops=False``, those extras are returned to the - caller, which blocks fusion. + With ``fuse_ops=True``, dispatch routing extras stay internal so + MegaMoE can claim the sequence when its gates pass. With + ``fuse_ops=False``, those extras are returned to the caller and the + Sequential stays unfused. Pass an MXFP8 ``recipe`` to construct + natively quantized GroupedLinear weights via ``quantized_model_init``. """ buffer = self._make_buffer() dispatch = te_ops.Dispatch(buffer) - fc1 = te_ops.GroupedLinear( - NUM_LOCAL_EXPERTS, - HIDDEN_DIM, - 2 * INTERMEDIATE_DIM, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) - activation = te_ops.ScaledSwiGLU() - fc2 = te_ops.GroupedLinear( - NUM_LOCAL_EXPERTS, - INTERMEDIATE_DIM, - HIDDEN_DIM, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, + init_ctx = ( + te.quantized_model_init(enabled=True, recipe=recipe) + if recipe is not None + else nullcontext() ) + with init_ctx: + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) @@ -593,98 +611,142 @@ def test_caller_provides_grad_expert_out(self): self.assertGreater(gbuf.abs().sum().item(), 0.0) @_eager_test_include - def test_bf16_moe_sequential_fusion(self): - """Reference-backed fusion matches the unfused BF16 EP MoE sequence. + def test_bf16_moe_sequential_vs_reference(self): + """Fused MegaMoE Sequential matches the MXFP8-QDQ PyTorch reference.""" + self._run_moe_sequential_vs_reference(quantization=None) + + @_eager_test_include + def test_mxfp8_moe_sequential_vs_reference(self): + """Fused MegaMoE Sequential under MXFP8 autocast matches the QDQ reference.""" + self._run_moe_sequential_vs_reference(quantization="mxfp8") - ``fuse_ops=True`` keeps dispatch routing extras internal so fusion - fires; ``fuse_ops=False`` returns them to the caller and blocks it. + def _run_moe_sequential_vs_reference(self, *, quantization): + """Compare Sequential (MegaMoE when supported) to ``MoeEpReference``. + + MegaMoE quantizes GEMM operands to MXFP8 and accumulates in FP32. + Combine and public output are BF16 on the current device path. The + reference QDQ those operands then uses FP32 matmuls because PyTorch + cannot run MXFP8 GEMMs. """ if not EAGER: self.skipTest("variable-size reference comparison requires eager EP mode") - - fused, fused_fc1, fused_fc2 = self._make_moe_model(fuse_ops=True) - unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fuse_ops=False) + if torch.cuda.get_device_capability() != (10, 7): + self.skipTest("MegaMoE fusion requires Rubin SM107") + try: + from cudnn.moe_ep import MoeEp # noqa: F401 + except ImportError: + self.skipTest("cudnn.moe_ep.MoeEp is not installed") + + intermediate_dim = 128 + recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None + model, fc1, fc2 = self._make_moe_model( + fuse_ops=True, intermediate_dim=intermediate_dim, recipe=recipe + ) generator = torch.Generator(device=self.cfg.device) generator.manual_seed(3100 + self.cfg.rank) with torch.no_grad(): - for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for op in (fc1, fc2): for expert in range(NUM_LOCAL_EXPERTS): weight = ( torch.randn( - getattr(fused_op, f"weight{expert}").shape, + getattr(op, f"weight{expert}").shape, generator=generator, dtype=torch.float32, device=self.cfg.device, ) * 0.1 ).to(torch.bfloat16) - getattr(fused_op, f"weight{expert}").copy_(weight) - getattr(unfused_op, f"weight{expert}").copy_(weight) + getattr(op, f"weight{expert}").copy_(weight) + if quantization == "mxfp8": + self.assertIsInstance( + getattr(op, f"weight{expert}"), + MXFP8Tensor, + ) topk_idx, tokens, topk_weights = _make_moe_inputs( self.cfg.rank, self.cfg.ep_size, self.cfg.device, ) - fused_tokens = tokens.detach().clone().requires_grad_(True) - unfused_tokens = tokens.detach().clone().requires_grad_(True) - fused_topk_weights = topk_weights.detach().clone().requires_grad_(True) - unfused_topk_weights = topk_weights.detach().clone().requires_grad_(True) - - fused_out = fused( - fused_tokens, - topk_idx, - fused_topk_weights, + seq_tokens = tokens.detach().clone().requires_grad_(True) + seq_topk_weights = topk_weights.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() ) - unfused_out, tokens_per_expert, recv_topk_weights = unfused( - unfused_tokens, + with autocast_ctx: + seq_out = model(seq_tokens, topk_idx, seq_topk_weights) + + forward_ops = model._module_groups[0]._forward_ops + self.assertEqual(len(forward_ops), 1) + self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + self.assertIsInstance(seq_out, torch.Tensor) + self.assertEqual(seq_out.dtype, torch.bfloat16) + + fc1_weight = _reference_weights(fc1).detach() + fc2_weight = _reference_weights(fc2).detach() + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=intermediate_dim, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + output_format=MoeFormat.BF16, + combine_format=MoeFormat.BF16, + apply_topk_in_fc1=True, + generate_c=True, + compute_dtype=torch.float32, + gemm_format=MoeFormat.MXFP8, + ) + ref_out, fc1_c, route_metadata = reference( + tokens.detach(), + fc1_weight, + fc2_weight, topk_idx, - unfused_topk_weights, + topk_weights.detach(), ) - - fused_forward_ops = fused._module_groups[0]._forward_ops - unfused_forward_ops = unfused._module_groups[0]._forward_ops - self.assertEqual(len(fused_forward_ops), 1) - self.assertIsInstance(fused_forward_ops[0][0], FusedMoeEp) - self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops)) - self.assertIsInstance(fused_out, torch.Tensor) - self.assertEqual(fused_out.dtype, torch.bfloat16) - self.assertEqual(unfused_out.dtype, torch.bfloat16) - self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,)) - self.assertEqual(tokens_per_expert.dtype, torch.int64) - self.assertEqual(recv_topk_weights.dtype, torch.float32) + self.assertEqual(ref_out.dtype, torch.bfloat16) dy = ( torch.randn( - fused_out.shape, + seq_out.shape, generator=generator, dtype=torch.float32, device=self.cfg.device, ) * 0.1 ).to(torch.bfloat16) - fused_out.backward(dy) - unfused_out.backward(dy) + seq_out.backward(dy) + grad_tokens, grad_fc1, grad_fc2, grad_topk_weights = reference.backward( + dy, + tokens.detach(), + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + fc1_c, + route_metadata, + ) torch.cuda.synchronize() - # The two BF16 paths use different grouped-GEMM and reduction orders. - # Observed forward abs error reaches ~3e-5 on near-zero values where - # rtol does not apply, so atol must stay above that (1e-5 is too tight). - tolerances = {"rtol": 1.6e-2, "atol": 5e-5} - torch.testing.assert_close(fused_out, unfused_out, **tolerances) - torch.testing.assert_close(fused_tokens.grad, unfused_tokens.grad, **tolerances) + # Kernel MXFP8 GEMM vs QDQ+FP32 matmul: grouped-MLP-style MXFP8 tols. + tolerances = {"rtol": 0.125, "atol": 0.25} + torch.testing.assert_close(seq_out, ref_out, **tolerances) + torch.testing.assert_close( + seq_tokens.grad, grad_tokens.to(dtype=seq_tokens.dtype), **tolerances + ) torch.testing.assert_close( - fused_topk_weights.grad, - unfused_topk_weights.grad, - **tolerances, + seq_topk_weights.grad, grad_topk_weights.float(), **tolerances ) - for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for op, ref_grad in ((fc1, grad_fc1), (fc2, grad_fc2)): for expert in range(NUM_LOCAL_EXPERTS): - fused_grad = getattr(fused_op, f"weight{expert}").grad - unfused_grad = getattr(unfused_op, f"weight{expert}").grad - self.assertEqual(fused_grad.dtype, torch.bfloat16) - self.assertEqual(unfused_grad.dtype, torch.bfloat16) - torch.testing.assert_close(fused_grad, unfused_grad, **tolerances) + seq_grad = getattr(op, f"weight{expert}").grad + self.assertEqual(seq_grad.dtype, torch.bfloat16) + torch.testing.assert_close( + seq_grad, + ref_grad[expert].transpose(0, 1).to(dtype=seq_grad.dtype), + **tolerances, + ) @_zero_copy_test_include @_mxfp8_align_test @@ -992,3 +1054,4 @@ def _init_distributed(): release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) + diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 8a8d58ad0f..cdb089059b 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2028 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT """Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. @@ -326,6 +326,43 @@ def _format_round_trip( return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) +def _qdq( + tensor: torch.Tensor, + format: Optional[MoeFormat], + *, + axis: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Simulate an MXFP8/NVFP4 GEMM operand in PyTorch. + + The MegaMoE kernel quantizes once and feeds the quantized values into the + GEMM (FP32 accumulate). PyTorch has no MXFP8 matmul, so the reference + round-trips through block-scale storage and dequantizes back to ``dtype``. + """ + if format is None or format is MoeFormat.BF16: + return tensor.to(dtype=dtype) + return quantize_blockwise(tensor, format, axis=axis).dequantize(dtype=dtype) + + +def _swiglu_scale_fp32( + gate: torch.Tensor, + up: torch.Tensor, + weights: torch.Tensor, + *, + apply_scale: bool, + dtype: torch.dtype, +) -> torch.Tensor: + """SiLU(gate)*up[*weights] in fp32, then one cast to ``dtype``. + + Matches ``tex.scaled_swiglu``: promote, multiply, and store once. Stepwise + BF16 ``F.silu(g) * up * w`` extra-rounds between those muls. + """ + out = F.silu(gate.float()) * up.float() + if apply_scale: + out = out * weights.float() + return out.to(dtype=dtype) + + class MoeEpReference: """Reference implementation of routed SwiGLU experts plus EP dispatch. @@ -350,6 +387,7 @@ def __init__( gate_up_clamp: Optional[float] = None, generate_c: bool = False, compute_dtype: torch.dtype = torch.float32, + gemm_format: Optional[Union[MoeFormat, str]] = None, ) -> None: for name, value in ( ("num_experts", num_experts), @@ -397,6 +435,9 @@ def __init__( self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) self.compute_dtype = compute_dtype + self.gemm_format = None if gemm_format is None else _parse_format(gemm_format) + if self.gemm_format is MoeFormat.BF16: + self.gemm_format = None for name, fmt in ( ("output_format", self.output_format), @@ -418,9 +459,27 @@ def __repr__(self) -> str: f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " f"output={self.output_format.value}, combine={self.combine_format.value}, " + f"gemm_format={None if self.gemm_format is None else self.gemm_format.value}, " f"compute_dtype={self.compute_dtype})" ) + def _gemm_dtype(self) -> torch.dtype: + """FP32 accumulate when simulating an MXFP8 GEMM; otherwise ``compute_dtype``.""" + return torch.float32 if self.gemm_format is MoeFormat.MXFP8 else self.compute_dtype + + def _stage_gemm_operand( + self, + tensor: torch.Tensor, + *, + already_quantized: bool, + axis: int, + ) -> torch.Tensor: + """Apply the kernel's pre-GEMM quantize, then dequant for a PyTorch matmul.""" + dtype = self._gemm_dtype() + if already_quantized or self.gemm_format is None: + return tensor.to(dtype=dtype) + return _qdq(tensor, self.gemm_format, axis=axis, dtype=dtype) + def _collective_device(self, device: torch.device) -> torch.device: """Device the process group can run ``all_to_all_single`` on. @@ -523,15 +582,30 @@ def _run_local_experts( fc1_c_rows.append(gate_up.to(torch.bfloat16)) gate, up = gate_up.split(self.intermediate_size, dim=-1) if self.gate_up_clamp is not None: - gate = gate.clamp(max=self.gate_up_clamp) - up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) - intermediate = F.silu(gate) * up + gate = gate.float().clamp(max=self.gate_up_clamp) + up = up.float().clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) weights = route_weight.index_select(0, positions).unsqueeze(-1) - if self.apply_topk_in_fc1: - intermediate = intermediate * weights + # FP32 SiLU*up[*scale], then one cast. MegaMoE then quantizes this + # intermediate for the FC2 MXFP8 GEMM; the reference QDQ mimics that. + intermediate = _swiglu_scale_fp32( + gate, + up, + weights, + apply_scale=self.apply_topk_in_fc1, + dtype=self._gemm_dtype(), + ) + if self.gemm_format is not None: + intermediate = _qdq( + intermediate, + self.gemm_format, + axis=-1, + dtype=self._gemm_dtype(), + ) expert_output = intermediate @ fc2_weight[expert] if not self.apply_topk_in_fc1: - expert_output = expert_output * weights + expert_output = (expert_output.float() * weights.float()).to( + dtype=self.compute_dtype + ) expert_output = _format_round_trip( expert_output, self.combine_format, @@ -607,26 +681,42 @@ def __call__( if input_device != device: raise ValueError(f"{name} must be on {device}, got {input_device}") - activation_float = _decode_tensor( - activation, - name="activation", - expected_shape=(token_count, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, + activation_float = self._stage_gemm_operand( + _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(activation, BlockScaledTensor), + axis=-1, ) - fc1_float = _decode_tensor( - fc1_weight, - name="fc1_weight", - expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), - quantized_axis=1, - dtype=self.compute_dtype, + fc1_float = self._stage_gemm_operand( + _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=( + self.experts_per_rank, + self.hidden_size, + 2 * self.intermediate_size, + ), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc1_weight, BlockScaledTensor), + axis=1, ) - fc2_float = _decode_tensor( - fc2_weight, - name="fc2_weight", - expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, + fc2_float = self._stage_gemm_operand( + _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc2_weight, BlockScaledTensor), + axis=1, ) plan = self._dispatch_plan(topk_idx, topk_weights) @@ -669,13 +759,15 @@ def __call__( ) returned = self._all_to_all(recv_output, recv_counts, send_counts) + # Combine payload is combine_format (BF16 round-trip above); reduce in + # fp32 so top-k summation does not extra-round in compute_dtype. combine_plane = torch.zeros( (token_count * self.top_k, self.hidden_size), - dtype=self.compute_dtype, + dtype=torch.float32, device=device, ) send_flat_slot = send_token_idx * self.top_k + send_slot_idx - combine_plane.index_copy_(0, send_flat_slot, returned) + combine_plane.index_copy_(0, send_flat_slot, returned.float()) reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) if self.output_format is MoeFormat.BF16: @@ -711,9 +803,10 @@ def backward( Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, grad_topk_weights)``. Activation and expert weight gradients use - ``compute_dtype``. Router weights are cast to ``compute_dtype`` for the - activation scaling path; the returned - ``grad_topk_weights`` remain float32. + ``compute_dtype``. SwiGLU and router-weight scaling recompute in + fp32 and round once into ``compute_dtype`` before the FC2 GEMMs, + matching ``tex.scaled_swiglu``. The returned ``grad_topk_weights`` + remain float32. """ if not self.generate_c: @@ -731,26 +824,38 @@ def backward( device = _tensor_device(activation) two_i = 2 * self.intermediate_size - activation_float = _decode_tensor( - activation, - name="activation", - expected_shape=(token_count, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, + activation_float = self._stage_gemm_operand( + _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(activation, BlockScaledTensor), + axis=-1, ) - fc1_float = _decode_tensor( - fc1_weight, - name="fc1_weight", - expected_shape=(self.experts_per_rank, self.hidden_size, two_i), - quantized_axis=1, - dtype=self.compute_dtype, + fc1_float = self._stage_gemm_operand( + _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc1_weight, BlockScaledTensor), + axis=1, ) - fc2_float = _decode_tensor( - fc2_weight, - name="fc2_weight", - expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, + fc2_float = self._stage_gemm_operand( + _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc2_weight, BlockScaledTensor), + axis=1, ) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): raise ValueError( @@ -808,44 +913,48 @@ def backward( continue c = c_rows.index_select(0, positions) x = x_rows.index_select(0, positions) - w = w_rows.index_select(0, positions).unsqueeze(-1) + w = w_rows.index_select(0, positions).unsqueeze(-1).float() d_y = dy_rows.index_select(0, positions) gate, up = c.split(self.intermediate_size, dim=-1) if self.gate_up_clamp is not None: - g = gate.clamp(max=self.gate_up_clamp) - u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + g = gate.float().clamp(max=self.gate_up_clamp) + u = up.float().clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) else: - g, u = gate, up + g, u = gate.float(), up.float() sig = torch.sigmoid(g) s = g * sig h = s * u if self.apply_topk_in_fc1: - h_fc2 = h * w - d_y_pre = d_y + h_fc2 = (h * w).to(dtype=self._gemm_dtype()) + d_y_pre = d_y.to(dtype=self._gemm_dtype()) else: - h_fc2 = h - d_y_pre = d_y * w + h_fc2 = h.to(dtype=self._gemm_dtype()) + d_y_pre = (d_y.float() * w).to(dtype=self._gemm_dtype()) + if self.gemm_format is not None: + h_fc2 = _qdq(h_fc2, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) + d_y_pre = _qdq(d_y_pre, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre - d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + d_h_fc2 = (d_y_pre @ fc2_float[expert].transpose(0, 1)).float() if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).to(dtype=self.compute_dtype).float() + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) else: d_h = d_h_fc2 - d_w_rows[positions] = ( - (d_y * (h @ fc2_float[expert])).sum(dim=-1).to(dtype=self.compute_dtype).float() - ) + d_w_rows[positions] = (d_y.float() * (h @ fc2_float[expert].float())).sum(dim=-1) d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s if self.gate_up_clamp is not None: - d_gate = d_g * (gate <= self.gate_up_clamp) - d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + d_gate = d_g * (gate.float() <= self.gate_up_clamp) + up_f = up.float() + d_up = d_u * ((up_f >= -self.gate_up_clamp) & (up_f <= self.gate_up_clamp)) else: d_gate, d_up = d_g, d_u - d_c = torch.cat((d_gate, d_up), dim=-1) + d_c = torch.cat((d_gate, d_up), dim=-1).to(dtype=self._gemm_dtype()) + if self.gemm_format is not None: + d_c = _qdq(d_c, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) grad_fc1[expert] = x.transpose(0, 1) @ d_c d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) @@ -854,10 +963,11 @@ def backward( returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) grad_activation = torch.zeros( (token_count, self.hidden_size), - dtype=self.compute_dtype, + dtype=torch.float32, device=device, ) - grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx.float()) + grad_activation = grad_activation.to(dtype=self.compute_dtype) grad_topk_weights = torch.zeros( (token_count * self.top_k,), dtype=torch.float32, device=device ) @@ -879,3 +989,4 @@ def backward( "MoeTensor", "quantize_blockwise", ] + diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index dd0f46c7ab..122eb75e95 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -2,7 +2,8 @@ # # See LICENSE for license information. -"""Reference-backed BF16 expert-parallel MoE fusion.""" +"""MegaMoE-backed expert-parallel MoE fusion (cudnn.moe_ep.MoeEp). +""" from __future__ import annotations @@ -11,10 +12,11 @@ import torch +from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import get_ep_group -from ...ep_reference import MoeEpReference from ...quantization import Recipe from ...tensor import Quantizer +from ...tensor.mxfp8_tensor import MXFP8Tensor from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU from ..fuser import register_forward_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -25,13 +27,102 @@ def _weight_list(op: GroupedLinear) -> list[torch.Tensor]: return [getattr(op, f"weight{idx}") for idx in range(op.num_groups)] -def _reference_weights(op: GroupedLinear) -> torch.Tensor: - """Pack ``(out, in)`` expert weights into reference ``(E, in, out)`` layout.""" - return torch.stack([weight.transpose(0, 1) for weight in _weight_list(op)]) +def _mxfp8_weight_k_major(weight: MXFP8Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Transpose a TE ``(out, in)`` MXFP8 weight to MegaMoE ``(in, out)``. + + Rowwise block scales travel with the ``in`` axis (K after the transpose). + GEMM-swizzled CuBLAS layouts are not a MegaMoE input; GroupedLinear stores + compact unswizzled scales under quantized model init. + """ + if weight._with_gemm_swizzled_scales: + raise NotImplementedError( + "FusedMoeEp requires unswizzled MXFP8 weight scales, " + "got a GEMM-swizzled tensor" + ) + if weight._rowwise_data is None or weight._rowwise_scale_inv is None: + raise ValueError("MXFP8 weight is missing rowwise data or scales") + out_features, in_features = weight.size() + if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: + raise ValueError( + f"MXFP8 weight K={in_features} is not divisible by " + f"{MXFP8_BLOCK_SCALING_SIZE}" + ) + data = ( + weight._rowwise_data.view(torch.float8_e4m3fn) + .transpose(0, 1) + .contiguous() + ) + scale = ( + weight._rowwise_scale_inv[:out_features, : in_features // MXFP8_BLOCK_SCALING_SIZE] + .view(torch.float8_e8m0fnu) + .transpose(0, 1) + .contiguous() + ) + return data, scale + + +def _pack_grouped_linear_weights(op: GroupedLinear, *, block_scaled_cls: Optional[type] = None): + """Pack TE ``(out, in)`` expert weights into MegaMoE ``(E, in, out)``. + + Dense BF16 stays a dense tensor. MXFP8 stays MXFP8: payloads and compact + scales are restacked without dequantizing. ``block_scaled_cls`` selects the + ``BlockScaledTensor`` type (cuDNN MegaMoE vs the PyTorch reference). + """ + weights = _weight_list(op) + if not weights: + raise ValueError("GroupedLinear has no per-expert weights to pack") + if all(isinstance(weight, MXFP8Tensor) for weight in weights): + if block_scaled_cls is None: + from cudnn.moe_ep import BlockScaledTensor as block_scaled_cls + data = [] + scale = [] + for weight in weights: + expert_data, expert_scale = _mxfp8_weight_k_major(weight) + data.append(expert_data) + scale.append(expert_scale) + packed_data = torch.stack(data) + return block_scaled_cls( + data=packed_data, + scale=torch.stack(scale), + format="mxfp8", + logical_shape=tuple(packed_data.shape), + axis=1, + ) + if any(isinstance(weight, MXFP8Tensor) for weight in weights): + raise TypeError("cannot mix MXFP8 and dense expert weights") + return torch.stack([weight.transpose(0, 1).contiguous() for weight in weights]) + + +def _flatten_moe_weight(weight) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Split a MegaMoE weight into tensors that ``save_for_backward`` can hold.""" + if isinstance(weight, torch.Tensor): + return weight, None + return weight.data, weight.scale + + +def _restore_moe_weight(data: torch.Tensor, scale: Optional[torch.Tensor], block_scaled_cls: type): + """Rebuild the MegaMoE weight saved by :func:`_flatten_moe_weight`.""" + if scale is None: + return data + return block_scaled_cls( + data=data, + scale=scale, + format="mxfp8", + logical_shape=tuple(data.shape), + axis=1, + ) def _grouped_linear_supported(op: GroupedLinear) -> bool: weights = _weight_list(op) if not op.single_grouped_weight else [] + + def _weight_ok(weight: torch.Tensor) -> bool: + if weight.dtype is not torch.bfloat16: + return False + if isinstance(weight, MXFP8Tensor): + return True + return not hasattr(weight, "dequantize") + return ( not op.use_bias and not op._scale_bias @@ -41,10 +132,19 @@ def _grouped_linear_supported(op: GroupedLinear) -> bool: and not op._is_distributed_weight() and not op.wgrad_store.delay_wgrad_compute() and bool(weights) - and all(weight.dtype is torch.bfloat16 for weight in weights) + and all(_weight_ok(weight) for weight in weights) ) +def _import_cudnn_moe_ep(): + """Return ``cudnn.moe_ep.MoeEp`` or ``None`` if the package is missing.""" + try: + from cudnn.moe_ep import MoeEp + except ImportError: + return None + return MoeEp + + def _routing_extras_internal( dispatch: Dispatch, fc1: GroupedLinear, @@ -54,9 +154,9 @@ def _routing_extras_internal( """Whether the dispatch routing extras stay inside the fusion. The fused op keeps tokens-per-expert and the received routing weights - internal to :class:`MoeEpReference`, so it can only replace the sequence - when those two outputs feed exactly these ops and are not returned to the - caller. + internal to :class:`cudnn.moe_ep.MoeEp`, so it can only replace the + sequence when those two outputs feed exactly these ops and are not + returned to the caller. """ tokens_per_expert, routing_weights = dispatch._extra_output_channels if tokens_per_expert is None or routing_weights is None: @@ -70,33 +170,29 @@ def _routing_extras_internal( ) -def _routing_extras_internal( - dispatch: Dispatch, - fc1: GroupedLinear, - activation: ScaledSwiGLU, - fc2: GroupedLinear, -) -> bool: - """Whether the dispatch routing extras stay inside the fusion. - - The fused op keeps tokens-per-expert and the received routing weights - internal to :class:`MoeEpReference`, so it can only replace the sequence - when those two outputs feed exactly these ops and are not returned to the - caller. - """ - tokens_per_expert, routing_weights = dispatch._extra_output_channels - if tokens_per_expert is None or routing_weights is None: +def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: + """Static MegaMoE capability gates that can be checked before first launch.""" + if _import_cudnn_moe_ep() is None: return False - if any(dispatch._extra_output_to_caller): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): return False - return ( - fc1._extra_input_channels[0] == tokens_per_expert - and fc2._extra_input_channels[0] == tokens_per_expert - and activation._extra_input_channels[0] == routing_weights - ) + if buffer.max_tokens_per_rank is None or buffer.max_tokens_per_rank <= 0: + return False + if buffer.hidden_dim % 128 != 0 or fc2.in_features % 128 != 0: + return False + if buffer.top_k > 32: + return False + ep_group = get_ep_group() + ep_size = 1 if ep_group is None else ep_group.size() + if ep_size > 16: + return False + return True def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: - if recipe is not None or len(window) != 5: + if len(window) != 5: + return False + if recipe is not None and not recipe.mxfp8(): return False dispatch, fc1, activation, fc2, combine = window if not ( @@ -118,6 +214,8 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo return False if not _routing_extras_internal(dispatch, fc1, activation, fc2): return False + if not _megamoe_supported(buffer, fc1, fc2): + return False return ( fc1.num_groups == buffer.num_local_experts and fc2.num_groups == buffer.num_local_experts @@ -128,7 +226,7 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo class FusedMoeEp(FusedOperation): - """Joint BF16 fusion implemented with :class:`MoeEpReference`.""" + """Joint EP MoE fusion implemented with :class:`cudnn.moe_ep.MoeEp`.""" def __init__( self, @@ -140,16 +238,25 @@ def __init__( combine: Combine, ) -> None: super().__init__([dispatch, fc1, activation, fc2, combine]) + moe_ep_cls = _import_cudnn_moe_ep() + if moe_ep_cls is None: + raise ImportError( + "FusedMoeEp requires cudnn.moe_ep.MoeEp. Install the in-tree " + "cuDNN frontend with: pip install --force-reinstall " + "'./cudnn_frontend[moe_ep]'" + ) + from cudnn.moe_ep import BlockScaledTensor + ep_group = get_ep_group() ep_size = 1 if ep_group is None else ep_group.size() - self._reference = MoeEpReference( + self._block_scaled_cls = BlockScaledTensor + self._moe = moe_ep_cls( num_experts=dispatch.buffer.num_local_experts * ep_size, hidden_size=dispatch.buffer.hidden_dim, intermediate_size=fc2.in_features, top_k=dispatch.buffer.top_k, ep_group=ep_group, max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, - compute_dtype=torch.bfloat16, apply_topk_in_fc1=True, generate_c=True, ) @@ -185,9 +292,14 @@ def fuser_forward( topk_idx, topk_weights = basic_op_extra_inputs[0] if topk_weights.dtype is not torch.float32: raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") - fc1_weight = _reference_weights(self.fc1) - fc2_weight = _reference_weights(self.fc2) - output, fc1_c, route_metadata = self._reference( + with torch.no_grad(): + fc1_weight = _pack_grouped_linear_weights( + self.fc1, block_scaled_cls=self._block_scaled_cls + ) + fc2_weight = _pack_grouped_linear_weights( + self.fc2, block_scaled_cls=self._block_scaled_cls + ) + output, fc1_c, route_metadata = self._moe( input_, fc1_weight, fc2_weight, @@ -196,10 +308,14 @@ def fuser_forward( ) if any(ctx.requires_grad for ctx in basic_op_ctxs): + fc1_data, fc1_scale = _flatten_moe_weight(fc1_weight) + fc2_data, fc2_scale = _flatten_moe_weight(fc2_weight) basic_op_ctxs[0].save_for_backward( input_, - fc1_weight, - fc2_weight, + fc1_data, + fc1_scale, + fc2_data, + fc2_scale, topk_idx, topk_weights, fc1_c, @@ -230,14 +346,18 @@ def fuser_backward( del basic_op_grad_extra_outputs ( input_, - fc1_weight, - fc2_weight, + fc1_data, + fc1_scale, + fc2_data, + fc2_scale, topk_idx, topk_weights, fc1_c, route_metadata, ) = basic_op_ctxs[0].saved_tensors - grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._reference.backward( + fc1_weight = _restore_moe_weight(fc1_data, fc1_scale, self._block_scaled_cls) + fc2_weight = _restore_moe_weight(fc2_data, fc2_scale, self._block_scaled_cls) + grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._moe.backward( grad_output, input_, fc1_weight, @@ -248,18 +368,27 @@ def fuser_backward( route_metadata, ) + # MegaMoE returns float32 grads in (E, in, out); TE parameters are (out, in). fc1_param_grads = [ - grad_fc1[idx].transpose(0, 1) if weight.requires_grad else None + ( + grad_fc1[idx].transpose(0, 1).to(dtype=weight.dtype) + if weight.requires_grad + else None + ) for idx, weight in enumerate(_weight_list(self.fc1)) ] fc2_param_grads = [ - grad_fc2[idx].transpose(0, 1) if weight.requires_grad else None + ( + grad_fc2[idx].transpose(0, 1).to(dtype=weight.dtype) + if weight.requires_grad + else None + ) for idx, weight in enumerate(_weight_list(self.fc2)) ] return ( - grad_input, + grad_input.to(dtype=input_.dtype), [(), fc1_param_grads, (), fc2_param_grads, ()], - [(None, grad_topk_weights), (None,), (None,), (None,), ()], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), ()], ) @@ -269,7 +398,11 @@ def fuse_ops( recipe: Optional[Recipe] = None, **unused: Any, ) -> list[FusibleOperation]: - """Fuse supported five-op BF16 EP MoE sequences.""" + """Fuse supported five-op EP MoE sequences into MegaMoE. + + Unfused Sequential (NCCL dispatch/combine + GroupedLinear + ScaledSwiGLU) + is the default. MegaMoE is claimed only when ``_matches`` succeeds. + """ del unused out: list[FusibleOperation] = [] idx = 0 @@ -297,3 +430,4 @@ def fuse_ops( __all__ = ["FusedMoeEp"] + From c7eceaaea987aceb0cf1d255338f532bb93e8829 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:52:26 +0000 Subject: [PATCH 47/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_ep.py | 5 +--- transformer_engine/pytorch/ep_reference.py | 1 - .../pytorch/ops/fused/moe_ep.py | 28 ++++--------------- 3 files changed, 7 insertions(+), 27 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index fe354029a4..c6ec49ab6f 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -735,9 +735,7 @@ def _run_moe_sequential_vs_reference(self, *, quantization): torch.testing.assert_close( seq_tokens.grad, grad_tokens.to(dtype=seq_tokens.dtype), **tolerances ) - torch.testing.assert_close( - seq_topk_weights.grad, grad_topk_weights.float(), **tolerances - ) + torch.testing.assert_close(seq_topk_weights.grad, grad_topk_weights.float(), **tolerances) for op, ref_grad in ((fc1, grad_fc1), (fc2, grad_fc2)): for expert in range(NUM_LOCAL_EXPERTS): seq_grad = getattr(op, f"weight{expert}").grad @@ -1054,4 +1052,3 @@ def _init_distributed(): release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) - diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index cdb089059b..d38c6a80bf 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -989,4 +989,3 @@ def backward( "MoeTensor", "quantize_blockwise", ] - diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 122eb75e95..b06166f980 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -2,8 +2,7 @@ # # See LICENSE for license information. -"""MegaMoE-backed expert-parallel MoE fusion (cudnn.moe_ep.MoeEp). -""" +"""MegaMoE-backed expert-parallel MoE fusion (cudnn.moe_ep.MoeEp).""" from __future__ import annotations @@ -36,22 +35,16 @@ def _mxfp8_weight_k_major(weight: MXFP8Tensor) -> tuple[torch.Tensor, torch.Tens """ if weight._with_gemm_swizzled_scales: raise NotImplementedError( - "FusedMoeEp requires unswizzled MXFP8 weight scales, " - "got a GEMM-swizzled tensor" + "FusedMoeEp requires unswizzled MXFP8 weight scales, got a GEMM-swizzled tensor" ) if weight._rowwise_data is None or weight._rowwise_scale_inv is None: raise ValueError("MXFP8 weight is missing rowwise data or scales") out_features, in_features = weight.size() if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: raise ValueError( - f"MXFP8 weight K={in_features} is not divisible by " - f"{MXFP8_BLOCK_SCALING_SIZE}" + f"MXFP8 weight K={in_features} is not divisible by {MXFP8_BLOCK_SCALING_SIZE}" ) - data = ( - weight._rowwise_data.view(torch.float8_e4m3fn) - .transpose(0, 1) - .contiguous() - ) + data = weight._rowwise_data.view(torch.float8_e4m3fn).transpose(0, 1).contiguous() scale = ( weight._rowwise_scale_inv[:out_features, : in_features // MXFP8_BLOCK_SCALING_SIZE] .view(torch.float8_e8m0fnu) @@ -370,19 +363,11 @@ def fuser_backward( # MegaMoE returns float32 grads in (E, in, out); TE parameters are (out, in). fc1_param_grads = [ - ( - grad_fc1[idx].transpose(0, 1).to(dtype=weight.dtype) - if weight.requires_grad - else None - ) + (grad_fc1[idx].transpose(0, 1).to(dtype=weight.dtype) if weight.requires_grad else None) for idx, weight in enumerate(_weight_list(self.fc1)) ] fc2_param_grads = [ - ( - grad_fc2[idx].transpose(0, 1).to(dtype=weight.dtype) - if weight.requires_grad - else None - ) + (grad_fc2[idx].transpose(0, 1).to(dtype=weight.dtype) if weight.requires_grad else None) for idx, weight in enumerate(_weight_list(self.fc2)) ] return ( @@ -430,4 +415,3 @@ def fuse_ops( __all__ = ["FusedMoeEp"] - From fb8fc831aa62cfb15e81b99ce951c9202e3d9ba8 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 19 Aug 2026 19:36:24 +0000 Subject: [PATCH 48/83] remove the previous artifact Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fuser.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 7009a2818c..fd66529ba8 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -432,11 +432,6 @@ def __init__( basic_ops.append(op) self._num_basic_ops: int = len(basic_ops) self._basic_ops: list[BasicOperation] = basic_ops - # Capture channel routing from each basic op. If any op later rebinds a - # channel it flips this flag, essentially invalidating this fuser's forward pass. - self._channels_stale: bool = False - for op in self._basic_ops: - op._capturing_op_fusers.add(self) # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) From 76d7fdd2d690fe42be63aac5444b9dec4b3bf91a Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 19 Aug 2026 21:21:44 +0000 Subject: [PATCH 49/83] minor change Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 40 ++++++++++++------- .../pytorch/ops/fused/moe_ep.py | 2 +- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index c6ec49ab6f..4066e7d4bf 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -178,6 +178,14 @@ def _degroup_mxfp8(recv_grouped, valid_counts=None): return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) +def _reference_weights(op): + """Pack GroupedLinear weights into MoeEpReference ``(E, in, out)`` layout.""" + packed = _pack_grouped_linear_weights(op, block_scaled_cls=BlockScaledTensor) + if isinstance(packed, torch.Tensor): + return packed.detach() + return packed + + class _Cfg: rank: int world_size: int @@ -612,7 +620,7 @@ def test_caller_provides_grad_expert_out(self): @_eager_test_include def test_bf16_moe_sequential_vs_reference(self): - """Fused MegaMoE Sequential matches the MXFP8-QDQ PyTorch reference.""" + """Unfused BF16 Sequential matches a BF16 ``MoeEpReference``.""" self._run_moe_sequential_vs_reference(quantization=None) @_eager_test_include @@ -621,12 +629,11 @@ def test_mxfp8_moe_sequential_vs_reference(self): self._run_moe_sequential_vs_reference(quantization="mxfp8") def _run_moe_sequential_vs_reference(self, *, quantization): - """Compare Sequential (MegaMoE when supported) to ``MoeEpReference``. + """Compare Sequential to ``MoeEpReference``. - MegaMoE quantizes GEMM operands to MXFP8 and accumulates in FP32. - Combine and public output are BF16 on the current device path. The - reference QDQ those operands then uses FP32 matmuls because PyTorch - cannot run MXFP8 GEMMs. + MXFP8 Sequential uses MegaMoE when supported; the reference QDQ GEMM + operands to MXFP8 and matmuls in FP32. BF16 Sequential stays unfused + and the reference runs FP32 GEMMs with no QDQ. """ if not EAGER: self.skipTest("variable-size reference comparison requires eager EP mode") @@ -634,9 +641,10 @@ def _run_moe_sequential_vs_reference(self, *, quantization): self.skipTest("MegaMoE fusion requires Rubin SM107") try: from cudnn.moe_ep import MoeEp # noqa: F401 - except ImportError: - self.skipTest("cudnn.moe_ep.MoeEp is not installed") + except ImportError as exc: + self.skipTest(f"cudnn.moe_ep.MoeEp is not installed ({type(exc).__name__}: {exc})") + # MegaMoE SM107 requires intermediate_size % 256 == 0. intermediate_dim = 128 recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None model, fc1, fc2 = self._make_moe_model( @@ -677,8 +685,9 @@ def _run_moe_sequential_vs_reference(self, *, quantization): seq_out = model(seq_tokens, topk_idx, seq_topk_weights) forward_ops = model._module_groups[0]._forward_ops - self.assertEqual(len(forward_ops), 1) - self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + if quantization == "mxfp8": + self.assertEqual(len(forward_ops), 1) + self.assertIsInstance(forward_ops[0][0], FusedMoeEp) self.assertIsInstance(seq_out, torch.Tensor) self.assertEqual(seq_out.dtype, torch.bfloat16) @@ -695,8 +704,8 @@ def _run_moe_sequential_vs_reference(self, *, quantization): combine_format=MoeFormat.BF16, apply_topk_in_fc1=True, generate_c=True, - compute_dtype=torch.float32, - gemm_format=MoeFormat.MXFP8, + compute_dtype=torch.float32 if quantization == "mxfp8" else torch.bfloat16, + gemm_format=MoeFormat.MXFP8 if quantization == "mxfp8" else MoeFormat.BF16, ) ref_out, fc1_c, route_metadata = reference( tokens.detach(), @@ -728,9 +737,10 @@ def _run_moe_sequential_vs_reference(self, *, quantization): route_metadata, ) torch.cuda.synchronize() - - # Kernel MXFP8 GEMM vs QDQ+FP32 matmul: grouped-MLP-style MXFP8 tols. - tolerances = {"rtol": 0.125, "atol": 0.25} + if quantization == "bf16": + tolerances = {"rtol": 1.6e-2, "atol": 5e-5} + else: + tolerances = {"rtol": 0.125, "atol": 0.25} torch.testing.assert_close(seq_out, ref_out, **tolerances) torch.testing.assert_close( seq_tokens.grad, grad_tokens.to(dtype=seq_tokens.dtype), **tolerances diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index b06166f980..4f271caeb8 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -171,7 +171,7 @@ def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: return False if buffer.max_tokens_per_rank is None or buffer.max_tokens_per_rank <= 0: return False - if buffer.hidden_dim % 128 != 0 or fc2.in_features % 128 != 0: + if buffer.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: return False if buffer.top_k > 32: return False From c94616291b0864f5c64a32e3590e13371ea65536 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 19 Aug 2026 23:35:20 +0000 Subject: [PATCH 50/83] minor bug fix Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 4066e7d4bf..2e5d068116 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -183,7 +183,13 @@ def _reference_weights(op): packed = _pack_grouped_linear_weights(op, block_scaled_cls=BlockScaledTensor) if isinstance(packed, torch.Tensor): return packed.detach() - return packed + return BlockScaledTensor( + data=packed.data.detach(), + scale=packed.scale.detach(), + format=packed.format, + logical_shape=packed.logical_shape, + axis=packed.axis, + ) class _Cfg: @@ -645,7 +651,7 @@ def _run_moe_sequential_vs_reference(self, *, quantization): self.skipTest(f"cudnn.moe_ep.MoeEp is not installed ({type(exc).__name__}: {exc})") # MegaMoE SM107 requires intermediate_size % 256 == 0. - intermediate_dim = 128 + intermediate_dim = 256 recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None model, fc1, fc2 = self._make_moe_model( fuse_ops=True, intermediate_dim=intermediate_dim, recipe=recipe @@ -691,8 +697,8 @@ def _run_moe_sequential_vs_reference(self, *, quantization): self.assertIsInstance(seq_out, torch.Tensor) self.assertEqual(seq_out.dtype, torch.bfloat16) - fc1_weight = _reference_weights(fc1).detach() - fc2_weight = _reference_weights(fc2).detach() + fc1_weight = _reference_weights(fc1) + fc2_weight = _reference_weights(fc2) reference = MoeEpReference( num_experts=self.cfg.num_experts, hidden_size=HIDDEN_DIM, From 80f60632ccdf1bca913496a77e85c731f794178c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Thu, 20 Aug 2026 00:03:09 +0000 Subject: [PATCH 51/83] single_groupe_weight + remove copies Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 98 ++++++----- .../pytorch/ops/fused/moe_ep.py | 163 ++++++++++-------- 2 files changed, 152 insertions(+), 109 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 2e5d068116..567fb7d48b 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -345,24 +345,34 @@ def _make_moe_model( if recipe is not None else nullcontext() ) - with init_ctx: - fc1 = te_ops.GroupedLinear( - NUM_LOCAL_EXPERTS, - HIDDEN_DIM, - 2 * intermediate_dim, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) - activation = te_ops.ScaledSwiGLU() - fc2 = te_ops.GroupedLinear( - NUM_LOCAL_EXPERTS, - intermediate_dim, - HIDDEN_DIM, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - ) + previous_single_param = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" + try: + with init_ctx: + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + ) + finally: + if previous_single_param is None: + del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] + else: + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) @@ -626,7 +636,7 @@ def test_caller_provides_grad_expert_out(self): @_eager_test_include def test_bf16_moe_sequential_vs_reference(self): - """Unfused BF16 Sequential matches a BF16 ``MoeEpReference``.""" + """Fused BF16 MegaMoE Sequential matches a BF16 ``MoeEpReference``.""" self._run_moe_sequential_vs_reference(quantization=None) @_eager_test_include @@ -637,9 +647,9 @@ def test_mxfp8_moe_sequential_vs_reference(self): def _run_moe_sequential_vs_reference(self, *, quantization): """Compare Sequential to ``MoeEpReference``. - MXFP8 Sequential uses MegaMoE when supported; the reference QDQ GEMM - operands to MXFP8 and matmuls in FP32. BF16 Sequential stays unfused - and the reference runs FP32 GEMMs with no QDQ. + Sequential uses MegaMoE when supported. The MXFP8 reference QDQ GEMM + operands to MXFP8 and matmuls in FP32; the BF16 reference runs FP32 + GEMMs with no QDQ. """ if not EAGER: self.skipTest("variable-size reference comparison requires eager EP mode") @@ -660,22 +670,22 @@ def _run_moe_sequential_vs_reference(self, *, quantization): generator.manual_seed(3100 + self.cfg.rank) with torch.no_grad(): for op in (fc1, fc2): + weights = op.weight.quantized_tensors + if weights is None: + weights = op.weight.split_into_quantized_tensors() for expert in range(NUM_LOCAL_EXPERTS): weight = ( torch.randn( - getattr(op, f"weight{expert}").shape, + weights[expert].shape, generator=generator, dtype=torch.float32, device=self.cfg.device, ) * 0.1 ).to(torch.bfloat16) - getattr(op, f"weight{expert}").copy_(weight) + weights[expert].copy_(weight) if quantization == "mxfp8": - self.assertIsInstance( - getattr(op, f"weight{expert}"), - MXFP8Tensor, - ) + self.assertIsInstance(weights[expert], MXFP8Tensor) topk_idx, tokens, topk_weights = _make_moe_inputs( self.cfg.rank, @@ -691,14 +701,20 @@ def _run_moe_sequential_vs_reference(self, *, quantization): seq_out = model(seq_tokens, topk_idx, seq_topk_weights) forward_ops = model._module_groups[0]._forward_ops - if quantization == "mxfp8": - self.assertEqual(len(forward_ops), 1) - self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + self.assertEqual(len(forward_ops), 1) + self.assertIsInstance(forward_ops[0][0], FusedMoeEp) self.assertIsInstance(seq_out, torch.Tensor) self.assertEqual(seq_out.dtype, torch.bfloat16) fc1_weight = _reference_weights(fc1) fc2_weight = _reference_weights(fc2) + for op, packed in ((fc1, fc1_weight), (fc2, fc2_weight)): + packed_data = packed.data if isinstance(packed, BlockScaledTensor) else packed + self.assertEqual(packed_data.data_ptr(), op.weight.rowwise_data.data_ptr()) + self.assertTrue(packed_data.permute(0, 2, 1).is_contiguous()) + if isinstance(packed, BlockScaledTensor): + self.assertEqual(packed.scale.data_ptr(), op.weight.scale_inv.data_ptr()) + self.assertTrue(packed.scale.permute(0, 2, 1).is_contiguous()) reference = MoeEpReference( num_experts=self.cfg.num_experts, hidden_size=HIDDEN_DIM, @@ -753,14 +769,18 @@ def _run_moe_sequential_vs_reference(self, *, quantization): ) torch.testing.assert_close(seq_topk_weights.grad, grad_topk_weights.float(), **tolerances) for op, ref_grad in ((fc1, grad_fc1), (fc2, grad_fc2)): - for expert in range(NUM_LOCAL_EXPERTS): - seq_grad = getattr(op, f"weight{expert}").grad - self.assertEqual(seq_grad.dtype, torch.bfloat16) - torch.testing.assert_close( - seq_grad, - ref_grad[expert].transpose(0, 1).to(dtype=seq_grad.dtype), - **tolerances, - ) + seq_grad = op.weight.grad + self.assertEqual(seq_grad.dtype, torch.bfloat16) + self.assertEqual( + tuple(seq_grad.shape), + (NUM_LOCAL_EXPERTS, op.out_features, op.in_features), + ) + self.assertTrue(seq_grad.is_contiguous()) + torch.testing.assert_close( + seq_grad, + ref_grad.transpose(1, 2).to(dtype=seq_grad.dtype), + **tolerances, + ) @_zero_copy_test_include @_mxfp8_align_test diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 4f271caeb8..2e007edba8 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -14,76 +14,76 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import get_ep_group from ...quantization import Recipe -from ...tensor import Quantizer -from ...tensor.mxfp8_tensor import MXFP8Tensor +from ...tensor import GroupedTensor, Quantizer from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU from ..fuser import register_forward_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext -def _weight_list(op: GroupedLinear) -> list[torch.Tensor]: - """Return per-expert weights in their registered order.""" - return [getattr(op, f"weight{idx}") for idx in range(op.num_groups)] +def _grouped_weight(op: GroupedLinear) -> GroupedTensor: + """Return the single packed ``(E, out, in)`` parameter required by MegaMoE.""" + if not op.single_grouped_weight or not isinstance(op.weight, GroupedTensor): + raise ValueError("FusedMoeEp requires GroupedLinear(single_grouped_weight=True)") + return op.weight -def _mxfp8_weight_k_major(weight: MXFP8Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Transpose a TE ``(out, in)`` MXFP8 weight to MegaMoE ``(in, out)``. +def _pack_grouped_linear_weights(op: GroupedLinear, *, block_scaled_cls: Optional[type] = None): + """View a packed TE ``(E, out, in)`` weight as MegaMoE ``(E, in, out)``. - Rowwise block scales travel with the ``in`` axis (K after the transpose). - GEMM-swizzled CuBLAS layouts are not a MegaMoE input; GroupedLinear stores - compact unswizzled scales under quantized model init. + The permutation is intentionally not made contiguous. MegaMoE internally + permutes back to ``(E, out, in)`` before requesting contiguous storage, so + retaining this view lets that request reuse GroupedLinear's original packed + buffer. ``block_scaled_cls`` selects the ``BlockScaledTensor`` type (cuDNN + MegaMoE vs the PyTorch reference). """ - if weight._with_gemm_swizzled_scales: - raise NotImplementedError( - "FusedMoeEp requires unswizzled MXFP8 weight scales, got a GEMM-swizzled tensor" - ) - if weight._rowwise_data is None or weight._rowwise_scale_inv is None: - raise ValueError("MXFP8 weight is missing rowwise data or scales") - out_features, in_features = weight.size() - if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: + weight = _grouped_weight(op) + num_groups = op.num_groups + out_features = op.out_features + in_features = op.in_features + expected_data_numel = num_groups * out_features * in_features + if weight.rowwise_data is None or weight.rowwise_data.numel() != expected_data_numel: raise ValueError( - f"MXFP8 weight K={in_features} is not divisible by {MXFP8_BLOCK_SCALING_SIZE}" + "GroupedLinear weight must have compact rowwise storage with shape " + f"({num_groups}, {out_features}, {in_features})" ) - data = weight._rowwise_data.view(torch.float8_e4m3fn).transpose(0, 1).contiguous() - scale = ( - weight._rowwise_scale_inv[:out_features, : in_features // MXFP8_BLOCK_SCALING_SIZE] - .view(torch.float8_e8m0fnu) - .transpose(0, 1) - .contiguous() - ) - return data, scale - - -def _pack_grouped_linear_weights(op: GroupedLinear, *, block_scaled_cls: Optional[type] = None): - """Pack TE ``(out, in)`` expert weights into MegaMoE ``(E, in, out)``. - Dense BF16 stays a dense tensor. MXFP8 stays MXFP8: payloads and compact - scales are restacked without dequantizing. ``block_scaled_cls`` selects the - ``BlockScaledTensor`` type (cuDNN MegaMoE vs the PyTorch reference). - """ - weights = _weight_list(op) - if not weights: - raise ValueError("GroupedLinear has no per-expert weights to pack") - if all(isinstance(weight, MXFP8Tensor) for weight in weights): + data_nk = weight.rowwise_data.view(num_groups, out_features, in_features) + if weight.quantizer is not None: + recipe = weight.quantizer._get_compatible_recipe() + if recipe is None or not recipe.mxfp8(): + raise TypeError("FusedMoeEp only supports dense BF16 or MXFP8 grouped weights") + if weight._with_gemm_swizzled_scales: + raise NotImplementedError( + "FusedMoeEp requires unswizzled MXFP8 weight scales, got a GEMM-swizzled tensor" + ) + if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: + raise ValueError( + f"MXFP8 weight K={in_features} is not divisible by " + f"{MXFP8_BLOCK_SCALING_SIZE}" + ) + scale_cols = in_features // MXFP8_BLOCK_SCALING_SIZE + expected_scale_numel = num_groups * out_features * scale_cols + if weight.scale_inv is None or weight.scale_inv.numel() != expected_scale_numel: + raise ValueError( + "GroupedLinear MXFP8 scales must have compact rowwise storage with shape " + f"({num_groups}, {out_features}, {scale_cols})" + ) if block_scaled_cls is None: from cudnn.moe_ep import BlockScaledTensor as block_scaled_cls - data = [] - scale = [] - for weight in weights: - expert_data, expert_scale = _mxfp8_weight_k_major(weight) - data.append(expert_data) - scale.append(expert_scale) - packed_data = torch.stack(data) + packed_data = data_nk.view(torch.float8_e4m3fn).permute(0, 2, 1) + packed_scale = ( + weight.scale_inv.view(num_groups, out_features, scale_cols) + .view(torch.float8_e8m0fnu) + .permute(0, 2, 1) + ) return block_scaled_cls( data=packed_data, - scale=torch.stack(scale), + scale=packed_scale, format="mxfp8", logical_shape=tuple(packed_data.shape), axis=1, ) - if any(isinstance(weight, MXFP8Tensor) for weight in weights): - raise TypeError("cannot mix MXFP8 and dense expert weights") - return torch.stack([weight.transpose(0, 1).contiguous() for weight in weights]) + return data_nk.permute(0, 2, 1) def _flatten_moe_weight(weight) -> tuple[torch.Tensor, Optional[torch.Tensor]]: @@ -106,26 +106,54 @@ def _restore_moe_weight(data: torch.Tensor, scale: Optional[torch.Tensor], block ) -def _grouped_linear_supported(op: GroupedLinear) -> bool: - weights = _weight_list(op) if not op.single_grouped_weight else [] +def _grouped_weight_grad(op: GroupedLinear, grad: torch.Tensor) -> list[Optional[torch.Tensor]]: + """Convert a MegaMoE ``(E, in, out)`` wgrad to one TE ``(E, out, in)`` grad.""" + weight = _grouped_weight(op) + expected_shape = (op.num_groups, op.in_features, op.out_features) + if tuple(grad.shape) != expected_shape: + raise RuntimeError( + f"MegaMoE weight gradient must have shape {expected_shape}, got {tuple(grad.shape)}" + ) + if not weight.requires_grad: + return [None] + + # copy_ performs the float32-to-parameter-dtype conversion while writing + # directly into the contiguous layout expected by the grouped parameter. + param_grad = torch.empty( + (op.num_groups, op.out_features, op.in_features), + dtype=weight.dtype, + device=grad.device, + ) + param_grad.copy_(grad.transpose(1, 2)) + return [param_grad] + - def _weight_ok(weight: torch.Tensor) -> bool: - if weight.dtype is not torch.bfloat16: - return False - if isinstance(weight, MXFP8Tensor): - return True - return not hasattr(weight, "dequantize") +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weight = op.weight if op.single_grouped_weight else None + weight_ok = ( + isinstance(weight, GroupedTensor) + and weight.dtype is torch.bfloat16 + and weight.rowwise_data is not None + ) + if weight_ok and weight.quantizer is not None: + recipe = weight.quantizer._get_compatible_recipe() + weight_ok = ( + recipe is not None + and recipe.mxfp8() + and not weight._with_gemm_swizzled_scales + and weight.rowwise_data is not None + and weight.scale_inv is not None + ) return ( not op.use_bias and not op._scale_bias - and not op.single_grouped_weight + and op.single_grouped_weight and not op.single_grouped_bias and not op._accumulate_into_main_grad and not op._is_distributed_weight() and not op.wgrad_store.delay_wgrad_compute() - and bool(weights) - and all(_weight_ok(weight) for weight in weights) + and weight_ok ) @@ -361,15 +389,10 @@ def fuser_backward( route_metadata, ) - # MegaMoE returns float32 grads in (E, in, out); TE parameters are (out, in). - fc1_param_grads = [ - (grad_fc1[idx].transpose(0, 1).to(dtype=weight.dtype) if weight.requires_grad else None) - for idx, weight in enumerate(_weight_list(self.fc1)) - ] - fc2_param_grads = [ - (grad_fc2[idx].transpose(0, 1).to(dtype=weight.dtype) if weight.requires_grad else None) - for idx, weight in enumerate(_weight_list(self.fc2)) - ] + # MegaMoE returns float32 grads in (E, in, out); each GroupedLinear has + # one packed (E, out, in) parameter. + fc1_param_grads = _grouped_weight_grad(self.fc1, grad_fc1) + fc2_param_grads = _grouped_weight_grad(self.fc2, grad_fc2) return ( grad_input.to(dtype=input_.dtype), [(), fc1_param_grads, (), fc2_param_grads, ()], From a779c6c50da7d7c45f0140943d3be0bd33582dec Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:04:19 +0000 Subject: [PATCH 52/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/fused/moe_ep.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 2e007edba8..14747a726a 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -58,8 +58,7 @@ def _pack_grouped_linear_weights(op: GroupedLinear, *, block_scaled_cls: Optiona ) if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: raise ValueError( - f"MXFP8 weight K={in_features} is not divisible by " - f"{MXFP8_BLOCK_SCALING_SIZE}" + f"MXFP8 weight K={in_features} is not divisible by {MXFP8_BLOCK_SCALING_SIZE}" ) scale_cols = in_features // MXFP8_BLOCK_SCALING_SIZE expected_scale_numel = num_groups * out_features * scale_cols From f197433d61ed54c8ff83f23afeebb5aaabe916a8 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 21 Aug 2026 09:27:19 +0000 Subject: [PATCH 53/83] cpu init for megatron needs Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_tensor.py | 30 ++++ .../pytorch/tensor/grouped_tensor.py | 133 ++++++++++++++++++ 2 files changed, 163 insertions(+) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index ac8c493290..ed1f5c225d 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -187,6 +187,36 @@ def test_basic_construction_all_same_shape(self) -> None: assert grouped_tensor.get_common_last_dim() == 512 assert grouped_tensor.has_data() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_module_to_moves_grouped_parameter_storage(self) -> None: + """Module.to preserves GroupedTensor parameters initialized on CPU.""" + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=2, + shapes=[(4, 8), (4, 8)], + quantizer=None, + device="cpu", + dtype=torch.float32, + ) + values = torch.arange(grouped_tensor.numel(), dtype=torch.float32) + grouped_tensor.rowwise_data.copy_(values) + + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(grouped_tensor)) + original_parameter = module.weight + expected = values.to(device="cuda", dtype=torch.bfloat16) + + module.to(device="cuda", dtype=torch.bfloat16) + + assert module.weight is original_parameter + assert isinstance(module.weight, GroupedTensor) + assert module.weight.device.type == "cuda" + assert module.weight.dtype == torch.bfloat16 + assert module.weight.rowwise_data.device.type == "cuda" + assert module.weight.rowwise_data.dtype == torch.bfloat16 + torch.testing.assert_close(module.weight.rowwise_data, expected, rtol=0, atol=0) + members = module.weight.split_into_quantized_tensors() + assert all(member.device.type == "cuda" for member in members) + def test_basic_construction_varying_first_dim(self) -> None: """Test GroupedTensor construction with varying first dimension""" num_tensors = 3 diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 0cc03602a1..051a40d33c 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -173,6 +173,69 @@ def __new__( ) return instance + def __tensor_flatten__(self): + """Expose grouped backing buffers to PyTorch's wrapper-subclass protocol.""" + tensor_attrs = ( + "rowwise_data", + "columnwise_data", + "scale_inv", + "columnwise_scale_inv", + "amax", + "columnwise_amax", + "scale", + "first_dims", + "last_dims", + "tensor_offsets", + ) + present = [name for name in tensor_attrs if getattr(self, name) is not None] + context = { + "cls": type(self), + "logical_shape": self.logical_shape, + "fake_dtype": self.fake_dtype, + "num_tensors": self.num_tensors, + "tensor_shapes": self.tensor_shapes, + "quantizer": self.quantizer, + "offsets": self.offsets, + "scale_inv_offsets": self.scale_inv_offsets, + "columnwise_scale_inv_offsets": self.columnwise_scale_inv_offsets, + "requires_grad": self.requires_grad, + "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + } + return present, context + + @staticmethod + def __tensor_unflatten__(inner_tensors, context, outer_size, outer_stride): + """Rebuild a GroupedTensor from PyTorch-managed backing buffers.""" + return context["cls"]( + shape=context["logical_shape"], + dtype=context["fake_dtype"], + num_tensors=context["num_tensors"], + shapes=context["tensor_shapes"], + quantizer=context["quantizer"], + data=inner_tensors.get("rowwise_data"), + columnwise_data=inner_tensors.get("columnwise_data"), + scale_inv=inner_tensors.get("scale_inv"), + columnwise_scale_inv=inner_tensors.get("columnwise_scale_inv"), + amax=inner_tensors.get("amax"), + columnwise_amax=inner_tensors.get("columnwise_amax"), + scale=inner_tensors.get("scale"), + first_dims=inner_tensors.get("first_dims"), + last_dims=inner_tensors.get("last_dims"), + tensor_offsets=inner_tensors.get("tensor_offsets"), + offsets=context["offsets"], + scale_inv_offsets=context["scale_inv_offsets"], + columnwise_scale_inv_offsets=context["columnwise_scale_inv_offsets"], + requires_grad=context["requires_grad"], + stride=list(outer_stride) if outer_stride is not None else None, + with_gemm_swizzled_scales=context["with_gemm_swizzled_scales"], + row_scaled_nvfp4=context["row_scaled_nvfp4"], + nvfp4_use_4over6=context["nvfp4_use_4over6"], + nvfp4_e4m3_max=context["nvfp4_e4m3_max"], + ) + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): """Dispatch by dequantizing grouped members, then requantizing writes.""" @@ -229,6 +292,76 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: return make_wrapper_like(src, requires_grad=False) return make_wrapper_like(src, requires_grad=src.requires_grad) + # Module.to()/cuda() use aten._to_copy on parameters. Rebuild the + # wrapper around migrated grouped storage instead of taking the default + # dequantize-and-stack path, which loses the GroupedTensor subclass and + # leaves its backing buffers on the original device. + if func == torch.ops.aten._to_copy.default: + src = args[0] + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") + + target_device = kwargs.get("device", src.device) + target_dtype = kwargs.get("dtype", src.dtype) + target_layout = kwargs.get("layout", src.layout) + pin_memory = kwargs.get("pin_memory", False) + non_blocking = kwargs.get("non_blocking", False) + + if target_layout != torch.strided: + raise NotImplementedError( + f"{cls.__name__} only supports strided layout, got {target_layout}" + ) + if pin_memory: + raise NotImplementedError(f"{cls.__name__} does not support pin_memory=True") + if src.quantizer is not None and target_dtype != src.dtype: + raise NotImplementedError( + f"{cls.__name__} cannot change the logical dtype of quantized storage " + f"from {src.dtype} to {target_dtype}" + ) + + def move_storage( + tensor: Optional[torch.Tensor], *, convert_dtype: bool = False + ) -> Optional[torch.Tensor]: + if tensor is None: + return None + dtype = target_dtype if convert_dtype else tensor.dtype + return tensor.to( + device=target_device, + dtype=dtype, + non_blocking=non_blocking, + copy=True, + ) + + convert_data_dtype = src.quantizer is None + return type(src)( + shape=src.logical_shape, + dtype=target_dtype, + num_tensors=src.num_tensors, + shapes=src.tensor_shapes, + quantizer=src.quantizer, + data=move_storage(src.rowwise_data, convert_dtype=convert_data_dtype), + columnwise_data=move_storage( + src.columnwise_data, convert_dtype=convert_data_dtype + ), + scale_inv=move_storage(src.scale_inv), + columnwise_scale_inv=move_storage(src.columnwise_scale_inv), + amax=move_storage(src.amax), + columnwise_amax=move_storage(src.columnwise_amax), + scale=move_storage(src.scale), + first_dims=move_storage(src.first_dims), + last_dims=move_storage(src.last_dims), + tensor_offsets=move_storage(src.tensor_offsets), + offsets=src.offsets, + scale_inv_offsets=src.scale_inv_offsets, + columnwise_scale_inv_offsets=src.columnwise_scale_inv_offsets, + requires_grad=src.requires_grad, + stride=list(src.stride()), + with_gemm_swizzled_scales=src._with_gemm_swizzled_scales, + row_scaled_nvfp4=src.row_scaled_nvfp4, + nvfp4_use_4over6=src.nvfp4_use_4over6, + nvfp4_e4m3_max=src.nvfp4_e4m3_max, + ) + # Parameter construction may invoke aten.expand on tensor subclasses. # Handle this explicitly so grouped parameters can be created safely. if func == torch.ops.aten.expand.default: From 075fc445eb4b4cd0b4d3d9f373d7afd06d7bc732 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:28:55 +0000 Subject: [PATCH 54/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/tensor/grouped_tensor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 051a40d33c..65615d2662 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -340,9 +340,7 @@ def move_storage( shapes=src.tensor_shapes, quantizer=src.quantizer, data=move_storage(src.rowwise_data, convert_dtype=convert_data_dtype), - columnwise_data=move_storage( - src.columnwise_data, convert_dtype=convert_data_dtype - ), + columnwise_data=move_storage(src.columnwise_data, convert_dtype=convert_data_dtype), scale_inv=move_storage(src.scale_inv), columnwise_scale_inv=move_storage(src.columnwise_scale_inv), amax=move_storage(src.amax), From cc01b975de54118ccfb938bac219176ce5452b11 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sat, 22 Aug 2026 22:59:21 +0000 Subject: [PATCH 55/83] untested commit for quantizer roles for dispatch/combine along with passing epbuffer from dipsatch to combine op via extra input Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 99 +++++++- transformer_engine/pytorch/ops/_common.py | 31 +++ .../pytorch/ops/basic/combine.py | 230 ++++++++++++++---- .../pytorch/ops/basic/dispatch.py | 164 ++++++++++--- .../pytorch/ops/fused/moe_ep.py | 128 ++++++++-- 5 files changed, 557 insertions(+), 95 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 567fb7d48b..a0fe93791a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -30,6 +30,8 @@ ) from transformer_engine.pytorch.ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp, _pack_grouped_linear_weights +from transformer_engine.pytorch.ops.op import OperationContext +from transformer_engine.pytorch.tensor import GroupedTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" @@ -373,16 +375,29 @@ def _make_moe_model( del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] else: os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param - combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) - combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) + combine = te_ops.Combine(num_local_tokens=TOKENS_PER_RANK) dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) + dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") + combine.set_extra_input_channel(0, "ep_handle") + combine.set_extra_input_channel(1, "tokens_per_expert") return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 + def _make_dispatch_combine_model(self, buffer): + """Build the minimal channel-routed Dispatch -> Combine pipeline.""" + dispatch = te_ops.Dispatch(buffer) + combine = te_ops.Combine(num_local_tokens=TOKENS_PER_RANK) + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) + dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) + combine.set_extra_input_channel(0, "ep_handle") + combine.set_extra_input_channel(1, "tokens_per_expert") + return te_ops.Sequential(dispatch, combine) + # Prepare @_eager_test_include @@ -494,6 +509,30 @@ def test_primitive_dispatch_combine_identity(self): # Autograd + @_eager_test_include + @_zero_copy_test_include + def test_basic_ops_channel_routed_identity(self): + """Combine consumes Dispatch's handle/count channels without receiving EpBuffer.""" + buffer = self._make_buffer() + model = self._make_dispatch_combine_model(buffer) + topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + tokens = tokens.detach().clone().requires_grad_(True) + output = model(tokens, topk_idx, weights) + (0.5 * output.float().square().sum()).backward() + torch.cuda.synchronize() + torch.testing.assert_close( + output.float(), + tokens.detach().float() * TOP_K, + atol=5e-2, + rtol=5e-2, + ) + torch.testing.assert_close( + tokens.grad.float(), + tokens.detach().float() * (TOP_K**2), + atol=1e-1, + rtol=1e-1, + ) + @_zero_copy_test_include def test_dispatch_autograd(self): """0.5*||recv_tokens||^2 ; grad_tokens equals TOP_K * tokens. Covers the @@ -564,6 +603,62 @@ def test_dispatch_mxfp8(self): self.assertTrue(is_symm_backed(recv_mx.scale_inv)) self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_basic_dispatch_mxfp8_and_recipe_change(self): + """The basic op uses its recipe-owned input quantizer and refreshes it.""" + self._require_mxfp8_shapes() + buffer = self._make_buffer(alignment=128) + dispatch = te_ops.Dispatch(buffer) + topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + + recv_bf16, _counts, _recv_weights, _handle = dispatch(tokens, topk_idx, weights) + self.assertIsInstance(recv_bf16, torch.Tensor) + self.assertNotIsInstance(recv_bf16, MXFP8Tensor) + + with te.autocast(enabled=True, recipe=MXFP8BlockScaling()): + recv_mx, counts, _recv_weights, handle = dispatch(tokens, topk_idx, weights) + self.assertIsInstance(recv_mx, GroupedTensor) + self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) + self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, weights, counts) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_basic_combine_backward_mxfp8(self): + """Combine's recipe-owned grad-output quantizer drives scaled combine_bwd.""" + self._require_mxfp8_shapes() + buffer = self._make_buffer(alignment=128) + topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_tokens, _recv_weights, counts = ep_dispatch( + buffer, + tokens, + topk_idx, + weights, + ) + combine = te_ops.Combine(num_local_tokens=TOKENS_PER_RANK) + combine.reset_recipe_state(recipe=MXFP8BlockScaling()) + combine.pre_fuser_forward(requires_grad=True) + ctx = OperationContext() + output, _ = combine.fuser_forward( + [ctx], + recv_tokens, + basic_op_extra_inputs=[(buffer.handle_mem, counts)], + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + basic_op_kwargs=[{}], + ) + ctx.saved_tensors = ctx.to_save + grad_input, _, _ = combine.fuser_backward( + [ctx], + torch.ones_like(output), + basic_op_grad_extra_outputs=[()], + ) + self.assertIsInstance(grad_input, GroupedTensor) + self.assertIsNotNone(grad_input.rowwise_data) + self.assertIsNotNone(grad_input.scale_inv) + @_zero_copy_test_include @_mxfp8_align_test def test_caller_provides_dispatch_recv_mxfp8(self): diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index f39115c4c6..33ea9bb144 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,6 +11,7 @@ import torch from transformer_engine_torch import FP8TensorMeta +from ..constants import MXFP8_BLOCK_SCALING_SIZE from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..quantized_tensor import QuantizedTensorStorage, Quantizer @@ -22,6 +23,8 @@ NVFP4Quantizer, ) from ..tensor.float8_tensor import Float8Tensor +from ..tensor.mxfp8_tensor import MXFP8Tensor +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..utils import canonicalize_dtype @@ -91,6 +94,34 @@ def maybe_dequantize( return tensor +def quantize_mxfp8_for_ep( + input_: torch.Tensor, + quantizer: MXFP8Quantizer, +) -> tuple[MXFP8Tensor | MXFP8TensorStorage, torch.Tensor]: + """Apply an MXFP8 quantizer and expose compact rowwise scales for EP.""" + if isinstance(input_, (MXFP8Tensor, MXFP8TensorStorage)): + quantized = input_ + else: + quantized = quantizer(input_) + if quantized._with_gemm_swizzled_scales: + raise ValueError("NCCL EP requires unswizzled MXFP8 scales.") + data = quantized.rowwise_data + scale_inv = quantized.scale_inv + if data is None or scale_inv is None: + raise ValueError("NCCL EP requires rowwise MXFP8 data and scales.") + rows, hidden = input_.shape + scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE + if scale_cols * scale_inv.element_size() % 16: + raise ValueError( + f"MXFP8 EP transport requires hidden size divisible by " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {hidden}." + ) + scale_inv = scale_inv[:rows, :scale_cols] + if not scale_inv.is_contiguous(): + raise ValueError("NCCL EP requires compact contiguous MXFP8 scales.") + return quantized, scale_inv + + def maybe_autocast_dtype( *, device_type: str = "cuda", diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 4d4740c517..bfb63b3034 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -6,12 +6,21 @@ from __future__ import annotations -from typing import Any, Optional +from typing import Any, Iterable, Optional import torch +import transformer_engine_torch as tex -from ...ep import EpBuffer, _alloc_io, is_symm_backed -from ...tensor import Quantizer +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ...ep import ( + _alloc_io, + _make_grouped_mxfp8, + _scale_alloc_io, + is_symm_backed, +) +from ...quantization import QuantizerRole, Recipe +from ...tensor import MXFP8Quantizer, Quantizer +from .._common import is_quantized_tensor, maybe_dequantize, quantize_mxfp8_for_ep from ..op import BasicOperation, OperationContext @@ -40,39 +49,108 @@ def _validate_grad_buffer( class Combine(BasicOperation): """Combine pre-weighted local expert outputs with NCCL EP. - The operation uses routing state produced by a :class:`Dispatch` with the - same :class:`EpBuffer`. + The operation consumes the routing handle and tokens-per-expert produced by + a preceding :class:`Dispatch` through extra-tensor channels. """ - def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: + num_extra_inputs: int = 2 + + def __init__(self, *, num_local_tokens: int) -> None: super().__init__() - self.buffer = buffer - self.num_local_tokens = ( - buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) - ) + self.num_local_tokens = int(num_local_tokens) if self.num_local_tokens < 0: raise ValueError("num_local_tokens must be non-negative.") - def op_forward( + def num_quantizers(self, mode: str) -> int: + return 1 if mode == "backward" else 0 + + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + if mode == "backward": + name = getattr(self, "name", "") or "" + return [ + QuantizerRole( + module_type="combine", + tensor_type="grad_output", + name=name, + ) + ] + return None + + def pre_fuser_forward(self, *, requires_grad: bool) -> None: + super().pre_fuser_forward(requires_grad=requires_grad) + quantizer = self.get_quantizer("backward", 0) + if quantizer is not None: + quantizer.set_usage(rowwise=True, columnwise=False) + quantizer.optimize_for_gemm = False + + def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: + super().reset_recipe_state(recipe=recipe) + quantizer = self.get_quantizer("backward", 0) + if quantizer is not None: + quantizer.internal = True + + def _resolve_grad_output_quantizer( self, - ctx: OperationContext, + prev_op_grad_output_quantizer: Optional[Quantizer], + ) -> Optional[MXFP8Quantizer]: + quantizer = self.get_quantizer("backward", 0) + if ( + quantizer is not None + and prev_op_grad_output_quantizer is not None + and quantizer is not prev_op_grad_output_quantizer + ): + raise ValueError( + "Combine grad_output_quantizer and previous operation grad-output " + "quantizer must be the same object when both are set." + ) + if quantizer is None: + quantizer = prev_op_grad_output_quantizer + if quantizer is None: + return None + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + "NCCL EP Combine backward supports MXFP8Quantizer only, got " + f"{type(quantizer).__name__}." + ) + if quantizer.dtype != DType.kFloat8E4M3: + raise NotImplementedError("NCCL EP Combine backward supports E4M3 MXFP8 only.") + return quantizer + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Combine uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Combine uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], input_: torch.Tensor, *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], - **kwargs: Any, - ) -> torch.Tensor: - del prev_op_grad_output_quantizer, next_op_input_quantizer + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, list[tuple[()]]]: + grad_output_quantizer = self._resolve_grad_output_quantizer( + prev_op_grad_output_quantizer + ) + handle_mem, tokens_per_expert = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] if input_.dtype is not torch.bfloat16: raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") - if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + if input_.ndim != 2: + raise ValueError(f"Combine input must be 2D, got shape {tuple(input_.shape)}.") + if handle_mem.dtype is not torch.uint8 or handle_mem.device != input_.device: + raise ValueError("Combine routing handle must be a uint8 tensor on the input device.") + if tokens_per_expert.dtype is not torch.int64 or tokens_per_expert.device != input_.device: raise ValueError( - f"Combine input must have shape (R, {self.buffer.hidden_dim}), " - f"got {tuple(input_.shape)}." + "Combine tokens_per_expert must be an int64 tensor on the input device." ) expert_out = input_ - if self.buffer.zero_copy: + zero_copy = bool(tex.ep_get_zero_copy()) + if zero_copy: expert_out = _alloc_io( tuple(input_.shape), input_.dtype, @@ -83,56 +161,106 @@ def op_forward( result = torch.empty( self.num_local_tokens, - self.buffer.hidden_dim, + input_.shape[-1], dtype=input_.dtype, device=input_.device, ) torch.ops.transformer_engine_ep.combine( - self.buffer.handle_mem, + handle_mem, expert_out, result, ) + ctx = basic_op_ctxs[0] if ctx.requires_grad: grad_out = kwargs.get("grad_out") - if self.buffer.eager and grad_out is not None: - raise ValueError( - "eager mode sizes combine gradients per step and cannot use " - "a caller-supplied grad_out" + if grad_output_quantizer is None: + grad_out = _validate_grad_buffer( + grad_out, + shape=tuple(input_.shape), + dtype=input_.dtype, + device=input_.device, ) - grad_out = _validate_grad_buffer( - grad_out, - shape=tuple(input_.shape), - dtype=input_.dtype, - device=input_.device, - ) - if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): + elif grad_out is not None: + if grad_out.device != input_.device: + raise ValueError( + f"grad_out must be on {input_.device}, got {grad_out.device}." + ) + if not grad_out.is_contiguous(): + raise ValueError("MXFP8 grad_out storage must be contiguous.") + if grad_out.requires_grad: + raise ValueError("grad_out must not require gradients.") + if zero_copy and grad_out is not None and not is_symm_backed(grad_out): raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") ctx.grad_out = grad_out ctx.input_shape = tuple(input_.shape) ctx.input_dtype = input_.dtype - ctx.save_for_backward(self.buffer.handle_mem) + ctx.zero_copy = zero_copy + ctx.grad_output_quantizer = grad_output_quantizer + ctx.save_for_backward(handle_mem, tokens_per_expert) - return result + if next_op_input_quantizer is not None and not is_quantized_tensor(result): + result = next_op_input_quantizer(result) + return result, [()] - def op_backward( + def fuser_backward( self, - ctx: OperationContext, + basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - (handle_mem,) = ctx.saved_tensors - grad_output = grad_output.contiguous() - grad_input = ctx.grad_out - if grad_input is None: - grad_input = _alloc_io( - ctx.input_shape, - ctx.input_dtype, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + ctx = basic_op_ctxs[0] + handle_mem, tokens_per_expert = ctx.saved_tensors + quantizer = ctx.grad_output_quantizer + if quantizer is None: + grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + ctx.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + grad_output, + grad_input, + ) + else: + quantized_grad, grad_scale_inv = quantize_mxfp8_for_ep( + grad_output, quantizer + ) + rows, hidden = ctx.input_shape + scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE + grad_data, grad_input_scale_inv = _scale_alloc_io( + ctx.grad_out, + rows, + hidden, + scale_cols, + quantized_grad.rowwise_data.dtype, + grad_scale_inv.dtype, grad_output.device, - self.buffer.zero_copy, + ctx.zero_copy, ) - torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, - grad_output, - grad_input, - ) - return grad_input, () + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + quantized_grad.rowwise_data.view(torch.float8_e4m3fn), + grad_data.view(torch.float8_e4m3fn), + grad_scale_inv, + grad_input_scale_inv, + ) + grad_input = _make_grouped_mxfp8( + grad_data, + grad_input_scale_inv, + tokens_per_expert, + quantized_grad._fp8_dtype, + ctx.input_dtype, + ) + return grad_input, [()], [(None, None)] diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 9762ed76f7..7884d9cb12 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -10,8 +10,21 @@ import torch -from ...ep import EpBuffer, _alloc_io, ep_prepare -from ...tensor import Quantizer +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ...ep import ( + EpBuffer, + _alloc_io, + _make_grouped_mxfp8, + _scale_alloc_io, + ep_prepare, +) +from ...quantization import QuantizerRole, Recipe +from ...tensor import MXFP8Quantizer, Quantizer +from .._common import ( + is_quantized_tensor, + maybe_dequantize, + quantize_mxfp8_for_ep, +) from ..op import BasicOperation, OperationContext @@ -46,12 +59,69 @@ class Dispatch(BasicOperation): """ num_extra_inputs: int = 2 - num_extra_outputs: int = 2 + # tokens-per-expert, received routing weights, and the opaque NCCL EP + # routing handle consumed by Combine. + num_extra_outputs: int = 3 def __init__(self, buffer: EpBuffer) -> None: super().__init__() self.buffer = buffer + def num_quantizers(self, mode: str) -> int: + return 1 if mode == "forward" else 0 + + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + if mode == "forward": + name = getattr(self, "name", "") or "" + return [ + QuantizerRole( + module_type="dispatch", + tensor_type="input", + name=name, + ) + ] + return None + + def pre_fuser_forward(self, *, requires_grad: bool) -> None: + super().pre_fuser_forward(requires_grad=requires_grad) + quantizer = self.get_quantizer("forward", 0) + if quantizer is not None: + quantizer.set_usage(rowwise=True, columnwise=False) + quantizer.optimize_for_gemm = False + + def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: + super().reset_recipe_state(recipe=recipe) + quantizer = self.get_quantizer("forward", 0) + if quantizer is not None: + quantizer.internal = True + + def _resolve_input_quantizer( + self, + next_op_input_quantizer: Optional[Quantizer], + ) -> Optional[MXFP8Quantizer]: + quantizer = self.get_quantizer("forward", 0) + if ( + quantizer is not None + and next_op_input_quantizer is not None + and quantizer is not next_op_input_quantizer + ): + raise ValueError( + "Dispatch input_quantizer and next operation input quantizer " + "must be the same object when both are set." + ) + if quantizer is None: + quantizer = next_op_input_quantizer + if quantizer is None: + return None + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + "NCCL EP Dispatch supports MXFP8Quantizer only, got " + f"{type(quantizer).__name__}." + ) + if quantizer.dtype != DType.kFloat8E4M3: + raise NotImplementedError("NCCL EP Dispatch supports E4M3 MXFP8 only.") + return quantizer + def op_forward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Dispatch uses fuser_forward") @@ -68,7 +138,7 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - del prev_op_grad_output_quantizer, next_op_input_quantizer + input_quantizer = self._resolve_input_quantizer(next_op_input_quantizer) topk_idx, topk_weights = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] @@ -115,13 +185,14 @@ def fuser_forward( raise RuntimeError("NCCL EP dispatch receive size is unavailable.") rows = int(rows) recv_shape = (rows, self.buffer.hidden_dim) - recv_tokens = _validate_output_buffer( - "recv_tokens", - recv_tokens, - shape=recv_shape, - dtype=self.buffer.payload_dtype, - device=self.buffer.device, - ) + if input_quantizer is None: + recv_tokens = _validate_output_buffer( + "recv_tokens", + recv_tokens, + shape=recv_shape, + dtype=self.buffer.payload_dtype, + device=self.buffer.device, + ) recv_topk_weights = _validate_output_buffer( "recv_topk_weights", recv_topk_weights, @@ -129,13 +200,6 @@ def fuser_forward( dtype=torch.float32, device=self.buffer.device, ) - if recv_tokens is None: - recv_tokens = _alloc_io( - recv_shape, - self.buffer.payload_dtype, - self.buffer.device, - self.buffer.zero_copy, - ) if recv_topk_weights is None: recv_topk_weights = _alloc_io( (rows,), @@ -144,23 +208,64 @@ def fuser_forward( self.buffer.zero_copy, ) - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - input_, - topk_weights, - recv_tokens, - recv_topk_weights, - ) + if input_quantizer is None: + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + else: + quantized_input, input_scale_inv = quantize_mxfp8_for_ep( + input_, input_quantizer + ) + scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE + recv_data, recv_scale_inv = _scale_alloc_io( + recv_tokens, + rows, + self.buffer.hidden_dim, + scale_cols, + quantized_input.rowwise_data.dtype, + input_scale_inv.dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + quantized_input.rowwise_data.view(torch.float8_e4m3fn), + topk_weights, + recv_data.view(torch.float8_e4m3fn), + recv_topk_weights, + input_scale_inv, + recv_scale_inv, + ) + recv_tokens = _make_grouped_mxfp8( + recv_data, + recv_scale_inv, + tokens_per_expert, + quantized_input._fp8_dtype, + input_.dtype, + ) ctx = basic_op_ctxs[0] if ctx.requires_grad: ctx.input_shape = tuple(input_.shape) ctx.input_dtype = input_.dtype ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer ctx.save_for_backward(self.buffer.handle_mem) - return recv_tokens, [(tokens_per_expert, recv_topk_weights)] + return recv_tokens, [(tokens_per_expert, recv_topk_weights, self.buffer.handle_mem)] def fuser_backward( self, @@ -175,7 +280,7 @@ def fuser_backward( ]: ctx = basic_op_ctxs[0] (handle_mem,) = ctx.saved_tensors - grad_output = grad_output.contiguous() + grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() grad_recv_weights = basic_op_grad_extra_outputs[0][1] if grad_recv_weights is None: @@ -204,4 +309,7 @@ def fuser_backward( grad_input, grad_topk_weights, ) + quantizer = ctx.prev_op_grad_output_quantizer + if quantizer is not None and not is_quantized_tensor(grad_input): + grad_input = quantizer(grad_input) return grad_input, [()], [(None, grad_topk_weights)] diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 14747a726a..202a6903cf 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -11,10 +11,15 @@ import torch -from ...constants import MXFP8_BLOCK_SCALING_SIZE +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE from ...ep import get_ep_group from ...quantization import Recipe -from ...tensor import GroupedTensor, Quantizer +from ...tensor import GroupedTensor, MXFP8Quantizer, Quantizer +from .._common import ( + is_quantized_tensor, + maybe_dequantize, + quantize_mxfp8_for_ep, +) from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU from ..fuser import register_forward_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -105,6 +110,64 @@ def _restore_moe_weight(data: torch.Tensor, scale: Optional[torch.Tensor], block ) +def _pack_activation( + input_: torch.Tensor, + quantizer: Optional[MXFP8Quantizer], + block_scaled_cls: type, +): + """Represent a TE activation in the public cuDNN MoeEp format.""" + if quantizer is None: + return input_ + quantized, scale_inv = quantize_mxfp8_for_ep(input_, quantizer) + return block_scaled_cls( + data=quantized.rowwise_data.view(torch.float8_e4m3fn), + scale=scale_inv.view(torch.float8_e8m0fnu), + format="mxfp8", + logical_shape=tuple(input_.shape), + axis=1, + ) + + +def _validate_internal_quantizers( + dispatch: Dispatch, + fc1: GroupedLinear, + fc2: GroupedLinear, + combine: Combine, +) -> Optional[MXFP8Quantizer]: + """Validate quantizers hidden inside the five-op MegaMoE fusion.""" + dispatch_quantizer = dispatch.get_input_quantizer() + quantizers = { + "Dispatch input": dispatch_quantizer, + "Combine grad_output": combine.get_grad_output_quantizer(), + } + for op_name, op in (("FC1", fc1), ("FC2", fc2)): + for group_idx in range(op.num_groups): + quantizers[f"{op_name} input {group_idx}"] = op.get_quantizer( + "forward", + 2 * group_idx, + ) + quantizers[f"{op_name} grad_output {group_idx}"] = op.get_quantizer( + "backward", + group_idx, + ) + active = {name: quantizer for name, quantizer in quantizers.items() if quantizer is not None} + for name, quantizer in active.items(): + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + f"FusedMoeEp supports MXFP8 internal quantizers only; {name} uses " + f"{type(quantizer).__name__}." + ) + if quantizer.dtype != DType.kFloat8E4M3: + raise NotImplementedError(f"FusedMoeEp requires E4M3 MXFP8 for {name}.") + if active and len(active) != len(quantizers): + missing = ", ".join(name for name, quantizer in quantizers.items() if quantizer is None) + raise ValueError( + "FusedMoeEp requires either an all-BF16 boundary or a complete MXFP8 " + f"quantizer set; missing {missing}." + ) + return dispatch_quantizer or fc1.get_quantizer("forward", 0) + + def _grouped_weight_grad(op: GroupedLinear, grad: torch.Tensor) -> list[Optional[torch.Tensor]]: """Convert a MegaMoE ``(E, in, out)`` wgrad to one TE ``(E, out, in)`` grad.""" weight = _grouped_weight(op) @@ -170,6 +233,7 @@ def _routing_extras_internal( fc1: GroupedLinear, activation: ScaledSwiGLU, fc2: GroupedLinear, + combine: Combine, ) -> bool: """Whether the dispatch routing extras stay inside the fusion. @@ -178,8 +242,8 @@ def _routing_extras_internal( sequence when those two outputs feed exactly these ops and are not returned to the caller. """ - tokens_per_expert, routing_weights = dispatch._extra_output_channels - if tokens_per_expert is None or routing_weights is None: + tokens_per_expert, routing_weights, ep_handle = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None or ep_handle is None: return False if any(dispatch._extra_output_to_caller): return False @@ -187,6 +251,8 @@ def _routing_extras_internal( fc1._extra_input_channels[0] == tokens_per_expert and fc2._extra_input_channels[0] == tokens_per_expert and activation._extra_input_channels[0] == routing_weights + and combine._extra_input_channels[0] == ep_handle + and combine._extra_input_channels[1] == tokens_per_expert ) @@ -223,16 +289,16 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo and isinstance(combine, Combine) ): return False - if dispatch.buffer is not combine.buffer: - return False buffer = dispatch.buffer + if combine.num_local_tokens != buffer.max_tokens_per_rank: + return False if not buffer.eager or buffer.payload_dtype is not torch.bfloat16: return False if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): return False if activation.activation_recompute_in_mlp or activation.glu_interleave_size is not None: return False - if not _routing_extras_internal(dispatch, fc1, activation, fc2): + if not _routing_extras_internal(dispatch, fc1, activation, fc2, combine): return False if not _megamoe_supported(buffer, fc1, fc2): return False @@ -279,6 +345,8 @@ def __init__( max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, apply_topk_in_fc1=True, generate_c=True, + combine_format="bf16", + output_format="bf16", ) @property @@ -303,7 +371,6 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: - del prev_op_grad_output_quantizer, next_op_input_quantizer if input_.dtype is not torch.bfloat16: raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") if any(kwargs for kwargs in basic_op_kwargs): @@ -313,6 +380,17 @@ def fuser_forward( if topk_weights.dtype is not torch.float32: raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") with torch.no_grad(): + input_quantizer = _validate_internal_quantizers( + self.dispatch, + self.fc1, + self.fc2, + self.basic_ops[4], + ) + activation = _pack_activation( + input_, + input_quantizer, + self._block_scaled_cls, + ) fc1_weight = _pack_grouped_linear_weights( self.fc1, block_scaled_cls=self._block_scaled_cls ) @@ -320,7 +398,7 @@ def fuser_forward( self.fc2, block_scaled_cls=self._block_scaled_cls ) output, fc1_c, route_metadata = self._moe( - input_, + activation, fc1_weight, fc2_weight, topk_idx, @@ -328,10 +406,12 @@ def fuser_forward( ) if any(ctx.requires_grad for ctx in basic_op_ctxs): + input_data, input_scale = _flatten_moe_weight(activation) fc1_data, fc1_scale = _flatten_moe_weight(fc1_weight) fc2_data, fc2_scale = _flatten_moe_weight(fc2_weight) basic_op_ctxs[0].save_for_backward( - input_, + input_data, + input_scale, fc1_data, fc1_scale, fc2_data, @@ -341,11 +421,17 @@ def fuser_forward( fc1_c, route_metadata, ) + basic_op_ctxs[0].input_dtype = input_.dtype + basic_op_ctxs[0].prev_op_grad_output_quantizer = ( + prev_op_grad_output_quantizer + ) # Dispatch extras are channel-bound with output_to_caller=False and are # only consumed by ops inside this fusion, so they need not be materialized. + if next_op_input_quantizer is not None and not is_quantized_tensor(output): + output = next_op_input_quantizer(output) return output, [ - (None, None), + (None, None, None), (), (), (), @@ -365,7 +451,8 @@ def fuser_backward( ]: del basic_op_grad_extra_outputs ( - input_, + input_data, + input_scale, fc1_data, fc1_scale, fc2_data, @@ -375,8 +462,17 @@ def fuser_backward( fc1_c, route_metadata, ) = basic_op_ctxs[0].saved_tensors + input_ = _restore_moe_weight( + input_data, + input_scale, + self._block_scaled_cls, + ) fc1_weight = _restore_moe_weight(fc1_data, fc1_scale, self._block_scaled_cls) fc2_weight = _restore_moe_weight(fc2_data, fc2_scale, self._block_scaled_cls) + grad_output = maybe_dequantize( + grad_output, + basic_op_ctxs[0].input_dtype, + ) grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._moe.backward( grad_output, input_, @@ -392,10 +488,14 @@ def fuser_backward( # one packed (E, out, in) parameter. fc1_param_grads = _grouped_weight_grad(self.fc1, grad_fc1) fc2_param_grads = _grouped_weight_grad(self.fc2, grad_fc2) + grad_input = grad_input.to(dtype=basic_op_ctxs[0].input_dtype) + grad_input_quantizer = basic_op_ctxs[0].prev_op_grad_output_quantizer + if grad_input_quantizer is not None: + grad_input = grad_input_quantizer(grad_input) return ( - grad_input.to(dtype=input_.dtype), + grad_input, [(), fc1_param_grads, (), fc2_param_grads, ()], - [(None, grad_topk_weights.float()), (None,), (None,), (None,), ()], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), (None, None)], ) From 609c9b403658bcf89845c7bdabc18a58526d2aaa Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 22 Aug 2026 23:00:32 +0000 Subject: [PATCH 56/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/_common.py | 2 +- transformer_engine/pytorch/ops/basic/combine.py | 12 +++--------- transformer_engine/pytorch/ops/basic/dispatch.py | 7 ++----- transformer_engine/pytorch/ops/fused/moe_ep.py | 4 +--- 4 files changed, 7 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 33ea9bb144..5c4167938c 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -113,7 +113,7 @@ def quantize_mxfp8_for_ep( scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE if scale_cols * scale_inv.element_size() % 16: raise ValueError( - f"MXFP8 EP transport requires hidden size divisible by " + "MXFP8 EP transport requires hidden size divisible by " f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {hidden}." ) scale_inv = scale_inv[:rows, :scale_cols] diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index bfb63b3034..96580f4f70 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -132,9 +132,7 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, list[tuple[()]]]: - grad_output_quantizer = self._resolve_grad_output_quantizer( - prev_op_grad_output_quantizer - ) + grad_output_quantizer = self._resolve_grad_output_quantizer(prev_op_grad_output_quantizer) handle_mem, tokens_per_expert = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] if input_.dtype is not torch.bfloat16: @@ -183,9 +181,7 @@ def fuser_forward( ) elif grad_out is not None: if grad_out.device != input_.device: - raise ValueError( - f"grad_out must be on {input_.device}, got {grad_out.device}." - ) + raise ValueError(f"grad_out must be on {input_.device}, got {grad_out.device}.") if not grad_out.is_contiguous(): raise ValueError("MXFP8 grad_out storage must be contiguous.") if grad_out.requires_grad: @@ -234,9 +230,7 @@ def fuser_backward( grad_input, ) else: - quantized_grad, grad_scale_inv = quantize_mxfp8_for_ep( - grad_output, quantizer - ) + quantized_grad, grad_scale_inv = quantize_mxfp8_for_ep(grad_output, quantizer) rows, hidden = ctx.input_shape scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE grad_data, grad_input_scale_inv = _scale_alloc_io( diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 7884d9cb12..9263b41d51 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -115,8 +115,7 @@ def _resolve_input_quantizer( return None if not isinstance(quantizer, MXFP8Quantizer): raise TypeError( - "NCCL EP Dispatch supports MXFP8Quantizer only, got " - f"{type(quantizer).__name__}." + f"NCCL EP Dispatch supports MXFP8Quantizer only, got {type(quantizer).__name__}." ) if quantizer.dtype != DType.kFloat8E4M3: raise NotImplementedError("NCCL EP Dispatch supports E4M3 MXFP8 only.") @@ -225,9 +224,7 @@ def fuser_forward( recv_topk_weights, ) else: - quantized_input, input_scale_inv = quantize_mxfp8_for_ep( - input_, input_quantizer - ) + quantized_input, input_scale_inv = quantize_mxfp8_for_ep(input_, input_quantizer) scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE recv_data, recv_scale_inv = _scale_alloc_io( recv_tokens, diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 202a6903cf..ace2643177 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -422,9 +422,7 @@ def fuser_forward( route_metadata, ) basic_op_ctxs[0].input_dtype = input_.dtype - basic_op_ctxs[0].prev_op_grad_output_quantizer = ( - prev_op_grad_output_quantizer - ) + basic_op_ctxs[0].prev_op_grad_output_quantizer = prev_op_grad_output_quantizer # Dispatch extras are channel-bound with output_to_caller=False and are # only consumed by ops inside this fusion, so they need not be materialized. From 31407af032209865059064c69cc62c6e9971927b Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 23 Aug 2026 09:37:28 +0000 Subject: [PATCH 57/83] cleanups Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 76 +++-- transformer_engine/pytorch/ops/_common.py | 12 +- .../pytorch/ops/basic/combine.py | 271 +++++++++------ .../pytorch/ops/basic/dispatch.py | 322 ++++++++++++------ .../pytorch/ops/fused/moe_ep.py | 248 ++++++-------- 5 files changed, 547 insertions(+), 382 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index a0fe93791a..6e18975fd2 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -29,7 +29,7 @@ _ep_dispatch_raw, ) from transformer_engine.pytorch.ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat -from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp, _pack_grouped_linear_weights +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp, _pack_cudnn_weights from transformer_engine.pytorch.ops.op import OperationContext from transformer_engine.pytorch.tensor import GroupedTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor @@ -182,7 +182,7 @@ def _degroup_mxfp8(recv_grouped, valid_counts=None): def _reference_weights(op): """Pack GroupedLinear weights into MoeEpReference ``(E, in, out)`` layout.""" - packed = _pack_grouped_linear_weights(op, block_scaled_cls=BlockScaledTensor) + packed = _pack_cudnn_weights(op, block_scaled_cls=BlockScaledTensor) if isinstance(packed, torch.Tensor): return packed.detach() return BlockScaledTensor( @@ -375,27 +375,31 @@ def _make_moe_model( del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] else: os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param - combine = te_ops.Combine(num_local_tokens=TOKENS_PER_RANK) + combine = te_ops.Combine() dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) + dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") combine.set_extra_input_channel(0, "ep_handle") combine.set_extra_input_channel(1, "tokens_per_expert") + combine.set_extra_input_channel(2, "routing_indices") return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 def _make_dispatch_combine_model(self, buffer): """Build the minimal channel-routed Dispatch -> Combine pipeline.""" dispatch = te_ops.Dispatch(buffer) - combine = te_ops.Combine(num_local_tokens=TOKENS_PER_RANK) + combine = te_ops.Combine() dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) + dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) combine.set_extra_input_channel(0, "ep_handle") combine.set_extra_input_channel(1, "tokens_per_expert") + combine.set_extra_input_channel(2, "routing_indices") return te_ops.Sequential(dispatch, combine) # Prepare @@ -606,22 +610,48 @@ def test_dispatch_mxfp8(self): @_eager_test_include @_zero_copy_test_include @_mxfp8_align_test - def test_basic_dispatch_mxfp8_and_recipe_change(self): - """The basic op uses its recipe-owned input quantizer and refreshes it.""" + def test_basic_dispatch_prequantized_mxfp8(self): + """A fresh basic Dispatch accepts an already-quantized MXFP8 input.""" self._require_mxfp8_shapes() buffer = self._make_buffer(alignment=128) dispatch = te_ops.Dispatch(buffer) topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + quantized_tokens = self._mxfp8_quantizer().quantize(tokens) - recv_bf16, _counts, _recv_weights, _handle = dispatch(tokens, topk_idx, weights) - self.assertIsInstance(recv_bf16, torch.Tensor) - self.assertNotIsInstance(recv_bf16, MXFP8Tensor) + recv_tokens, counts, _recv_weights, handle, _routing_indices = dispatch( + quantized_tokens, topk_idx, weights + ) + + self.assertIsInstance(recv_tokens, GroupedTensor) + self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) + torch.cuda.synchronize() + self.assertEqual(counts.shape, (NUM_LOCAL_EXPERTS,)) + recv_data = _degroup_mxfp8(recv_tokens).float() + self.assertTrue(torch.isfinite(recv_data).all()) + self.assertGreater(recv_data.abs().sum().item(), 0.0) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_basic_dispatch_recipe_mxfp8(self): + """A fresh basic Dispatch quantizes BF16 input under MXFP8 autocast.""" + self._require_mxfp8_shapes() + buffer = self._make_buffer(alignment=128) + dispatch = te_ops.Dispatch(buffer) + topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) with te.autocast(enabled=True, recipe=MXFP8BlockScaling()): - recv_mx, counts, _recv_weights, handle = dispatch(tokens, topk_idx, weights) - self.assertIsInstance(recv_mx, GroupedTensor) + recv_tokens, counts, _recv_weights, handle, _routing_indices = dispatch( + tokens, topk_idx, weights + ) + + self.assertIsInstance(recv_tokens, GroupedTensor) self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) - self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, weights, counts) + torch.cuda.synchronize() + self.assertEqual(counts.shape, (NUM_LOCAL_EXPERTS,)) + recv_data = _degroup_mxfp8(recv_tokens).float() + self.assertTrue(torch.isfinite(recv_data).all()) + self.assertGreater(recv_data.abs().sum().item(), 0.0) @_eager_test_include @_zero_copy_test_include @@ -637,14 +667,14 @@ def test_basic_combine_backward_mxfp8(self): topk_idx, weights, ) - combine = te_ops.Combine(num_local_tokens=TOKENS_PER_RANK) + combine = te_ops.Combine() combine.reset_recipe_state(recipe=MXFP8BlockScaling()) combine.pre_fuser_forward(requires_grad=True) ctx = OperationContext() output, _ = combine.fuser_forward( [ctx], recv_tokens, - basic_op_extra_inputs=[(buffer.handle_mem, counts)], + basic_op_extra_inputs=[(buffer.handle_mem, counts, topk_idx)], prev_op_grad_output_quantizer=None, next_op_input_quantizer=None, basic_op_kwargs=[{}], @@ -731,8 +761,8 @@ def test_caller_provides_grad_expert_out(self): @_eager_test_include def test_bf16_moe_sequential_vs_reference(self): - """Fused BF16 MegaMoE Sequential matches a BF16 ``MoeEpReference``.""" - self._run_moe_sequential_vs_reference(quantization=None) + """Unfused BF16 Sequential matches a BF16 ``MoeEpReference``.""" + self._run_moe_sequential_vs_reference(quantization="bf16") @_eager_test_include def test_mxfp8_moe_sequential_vs_reference(self): @@ -742,9 +772,9 @@ def test_mxfp8_moe_sequential_vs_reference(self): def _run_moe_sequential_vs_reference(self, *, quantization): """Compare Sequential to ``MoeEpReference``. - Sequential uses MegaMoE when supported. The MXFP8 reference QDQ GEMM - operands to MXFP8 and matmuls in FP32; the BF16 reference runs FP32 - GEMMs with no QDQ. + Sequential uses MegaMoE for MXFP8 when supported. The MXFP8 reference + QDQ GEMM operands to MXFP8 and matmuls in FP32; the BF16 reference runs + FP32 GEMMs with no QDQ. """ if not EAGER: self.skipTest("variable-size reference comparison requires eager EP mode") @@ -796,8 +826,12 @@ def _run_moe_sequential_vs_reference(self, *, quantization): seq_out = model(seq_tokens, topk_idx, seq_topk_weights) forward_ops = model._module_groups[0]._forward_ops - self.assertEqual(len(forward_ops), 1) - self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + if quantization == "mxfp8": + self.assertEqual(len(forward_ops), 1) + self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + else: + self.assertEqual(len(forward_ops), 5) + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in forward_ops)) self.assertIsInstance(seq_out, torch.Tensor) self.assertEqual(seq_out.dtype, torch.bfloat16) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 5c4167938c..ddb7bceefd 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -95,18 +95,20 @@ def maybe_dequantize( def quantize_mxfp8_for_ep( - input_: torch.Tensor, - quantizer: MXFP8Quantizer, + input_: torch.Tensor | MXFP8TensorStorage, + quantizer: Optional[MXFP8Quantizer], ) -> tuple[MXFP8Tensor | MXFP8TensorStorage, torch.Tensor]: - """Apply an MXFP8 quantizer and expose compact rowwise scales for EP.""" + """Return an MXFP8 input and its compact rowwise scales for EP.""" if isinstance(input_, (MXFP8Tensor, MXFP8TensorStorage)): quantized = input_ else: + if quantizer is None: + raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") quantized = quantizer(input_) if quantized._with_gemm_swizzled_scales: raise ValueError("NCCL EP requires unswizzled MXFP8 scales.") - data = quantized.rowwise_data - scale_inv = quantized.scale_inv + data = quantized._rowwise_data + scale_inv = quantized._rowwise_scale_inv if data is None or scale_inv is None: raise ValueError("NCCL EP requires rowwise MXFP8 data and scales.") rows, hidden = input_.shape diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 96580f4f70..1f65f2c741 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -46,20 +46,41 @@ def _validate_grad_buffer( return tensor +def _validate_combine_inputs( + input_: torch.Tensor, + handle_mem: torch.Tensor, + tokens_per_expert: torch.Tensor, + topk_idx: torch.Tensor, +) -> tuple[tuple[int, int], int]: + """Validate the expert output and routing metadata consumed by Combine.""" + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") + if input_.ndim != 2: + raise ValueError(f"Combine input must be 2D, got shape {tuple(input_.shape)}.") + if handle_mem.dtype is not torch.uint8 or handle_mem.device != input_.device: + raise ValueError("Combine routing handle must be a uint8 tensor on the input device.") + if tokens_per_expert.dtype is not torch.int64 or tokens_per_expert.device != input_.device: + raise ValueError("Combine tokens_per_expert must be an int64 tensor on the input device.") + if topk_idx.ndim != 2: + raise ValueError( + f"Combine routing indices must be 2D, got shape {tuple(topk_idx.shape)}." + ) + if topk_idx.dtype not in (torch.int32, torch.int64) or topk_idx.device != input_.device: + raise ValueError( + "Combine routing indices must be an int32 or int64 tensor on the input device." + ) + return tuple(input_.shape), topk_idx.shape[0] + + class Combine(BasicOperation): """Combine pre-weighted local expert outputs with NCCL EP. - The operation consumes the routing handle and tokens-per-expert produced by - a preceding :class:`Dispatch` through extra-tensor channels. + The operation consumes the routing handle, tokens-per-expert, and routing + indices produced by a preceding :class:`Dispatch` through extra-tensor + channels. """ - num_extra_inputs: int = 2 - - def __init__(self, *, num_local_tokens: int) -> None: - super().__init__() - self.num_local_tokens = int(num_local_tokens) - if self.num_local_tokens < 0: - raise ValueError("num_local_tokens must be non-negative.") + num_extra_inputs: int = 3 def num_quantizers(self, mode: str) -> int: return 1 if mode == "backward" else 0 @@ -122,6 +143,115 @@ def op_forward(self, *args: Any, **kwargs: Any) -> None: def op_backward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Combine uses fuser_backward") + @staticmethod + def _stage_expert_output(input_: torch.Tensor, *, zero_copy: bool) -> torch.Tensor: + """Copy expert output into symmetric memory when zero-copy IO is enabled.""" + if not zero_copy: + return input_ + expert_out = _alloc_io(tuple(input_.shape), input_.dtype, input_.device, True) + expert_out.copy_(input_) + return expert_out + + @staticmethod + def _combine_forward_impl( + handle_mem: torch.Tensor, + expert_out: torch.Tensor, + *, + num_local_tokens: int, + ) -> torch.Tensor: + """Allocate and populate the local-token result.""" + result = torch.empty( + num_local_tokens, + expert_out.shape[-1], + dtype=expert_out.dtype, + device=expert_out.device, + ) + torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) + return result + + @staticmethod + def _prepare_grad_buffer( + grad_out: Optional[torch.Tensor], + grad_output_quantizer: Optional[MXFP8Quantizer], + *, + input_shape: tuple[int, int], + input_dtype: torch.dtype, + device: torch.device, + zero_copy: bool, + ) -> Optional[torch.Tensor]: + """Validate caller storage for the expert-output gradient.""" + if grad_output_quantizer is None: + grad_out = _validate_grad_buffer( + grad_out, + shape=input_shape, + dtype=input_dtype, + device=device, + ) + elif grad_out is not None: + if grad_out.device != device: + raise ValueError(f"grad_out must be on {device}, got {grad_out.device}.") + if not grad_out.is_contiguous(): + raise ValueError("MXFP8 grad_out storage must be contiguous.") + if grad_out.requires_grad: + raise ValueError("grad_out must not require gradients.") + if zero_copy and grad_out is not None and not is_symm_backed(grad_out): + raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") + return grad_out + + @staticmethod + def _combine_backward_impl( + ctx: OperationContext, + handle_mem: torch.Tensor, + tokens_per_expert: torch.Tensor, + grad_output: torch.Tensor, + ) -> torch.Tensor: + """Route a local-token gradient in the selected transport format. + + Keep transport-specific setup in this function so future quantized + formats can be added without branching in ``fuser_backward``. + """ + quantizer = ctx.grad_output_quantizer + if quantizer is None: + grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + ctx.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd(handle_mem, grad_output, grad_input) + return grad_input + + quantized_grad, grad_scale_inv = quantize_mxfp8_for_ep(grad_output, quantizer) + num_recv_tokens, hidden = ctx.input_shape + scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE + grad_data, grad_input_scale_inv = _scale_alloc_io( + ctx.grad_out, + num_recv_tokens, + hidden, + scale_cols, + quantized_grad._rowwise_data.dtype, + grad_scale_inv.dtype, + grad_output.device, + ctx.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + quantized_grad._rowwise_data.view(torch.float8_e4m3fn), + grad_data.view(torch.float8_e4m3fn), + grad_scale_inv, + grad_input_scale_inv, + ) + return _make_grouped_mxfp8( + grad_data, + grad_input_scale_inv, + tokens_per_expert, + quantized_grad._fp8_dtype, + ctx.input_dtype, + ) + def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -132,69 +262,44 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, list[tuple[()]]]: + # Resolve backward quantization and unpack Dispatch routing metadata. grad_output_quantizer = self._resolve_grad_output_quantizer(prev_op_grad_output_quantizer) - handle_mem, tokens_per_expert = basic_op_extra_inputs[0] + handle_mem, tokens_per_expert, topk_idx = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] - if input_.dtype is not torch.bfloat16: - raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") - if input_.ndim != 2: - raise ValueError(f"Combine input must be 2D, got shape {tuple(input_.shape)}.") - if handle_mem.dtype is not torch.uint8 or handle_mem.device != input_.device: - raise ValueError("Combine routing handle must be a uint8 tensor on the input device.") - if tokens_per_expert.dtype is not torch.int64 or tokens_per_expert.device != input_.device: - raise ValueError( - "Combine tokens_per_expert must be an int64 tensor on the input device." - ) + input_shape, num_local_tokens = _validate_combine_inputs( + input_, + handle_mem, + tokens_per_expert, + topk_idx, + ) - expert_out = input_ + # Stage zero-copy input if needed, then restore local-token order. zero_copy = bool(tex.ep_get_zero_copy()) - if zero_copy: - expert_out = _alloc_io( - tuple(input_.shape), - input_.dtype, - input_.device, - True, - ) - expert_out.copy_(input_) - - result = torch.empty( - self.num_local_tokens, - input_.shape[-1], - dtype=input_.dtype, - device=input_.device, - ) - torch.ops.transformer_engine_ep.combine( + expert_out = self._stage_expert_output(input_, zero_copy=zero_copy) + result = self._combine_forward_impl( handle_mem, expert_out, - result, + num_local_tokens=num_local_tokens, ) + # Preserve routing state and optional caller storage for backward. ctx = basic_op_ctxs[0] if ctx.requires_grad: - grad_out = kwargs.get("grad_out") - if grad_output_quantizer is None: - grad_out = _validate_grad_buffer( - grad_out, - shape=tuple(input_.shape), - dtype=input_.dtype, - device=input_.device, - ) - elif grad_out is not None: - if grad_out.device != input_.device: - raise ValueError(f"grad_out must be on {input_.device}, got {grad_out.device}.") - if not grad_out.is_contiguous(): - raise ValueError("MXFP8 grad_out storage must be contiguous.") - if grad_out.requires_grad: - raise ValueError("grad_out must not require gradients.") - if zero_copy and grad_out is not None and not is_symm_backed(grad_out): - raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") - ctx.grad_out = grad_out - ctx.input_shape = tuple(input_.shape) + ctx.grad_out = self._prepare_grad_buffer( + kwargs.get("grad_out"), + grad_output_quantizer, + input_shape=input_shape, + input_dtype=input_.dtype, + device=input_.device, + zero_copy=zero_copy, + ) + ctx.input_shape = input_shape ctx.input_dtype = input_.dtype ctx.zero_copy = zero_copy ctx.grad_output_quantizer = grad_output_quantizer ctx.save_for_backward(handle_mem, tokens_per_expert) + # Hand off to the next op in its requested representation. if next_op_input_quantizer is not None and not is_quantized_tensor(result): result = next_op_input_quantizer(result) return result, [()] @@ -213,48 +318,10 @@ def fuser_backward( del basic_op_grad_extra_outputs ctx = basic_op_ctxs[0] handle_mem, tokens_per_expert = ctx.saved_tensors - quantizer = ctx.grad_output_quantizer - if quantizer is None: - grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() - grad_input = ctx.grad_out - if grad_input is None: - grad_input = _alloc_io( - ctx.input_shape, - ctx.input_dtype, - grad_output.device, - ctx.zero_copy, - ) - torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, - grad_output, - grad_input, - ) - else: - quantized_grad, grad_scale_inv = quantize_mxfp8_for_ep(grad_output, quantizer) - rows, hidden = ctx.input_shape - scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE - grad_data, grad_input_scale_inv = _scale_alloc_io( - ctx.grad_out, - rows, - hidden, - scale_cols, - quantized_grad.rowwise_data.dtype, - grad_scale_inv.dtype, - grad_output.device, - ctx.zero_copy, - ) - torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, - quantized_grad.rowwise_data.view(torch.float8_e4m3fn), - grad_data.view(torch.float8_e4m3fn), - grad_scale_inv, - grad_input_scale_inv, - ) - grad_input = _make_grouped_mxfp8( - grad_data, - grad_input_scale_inv, - tokens_per_expert, - quantized_grad._fp8_dtype, - ctx.input_dtype, - ) - return grad_input, [()], [(None, None)] + grad_input = self._combine_backward_impl( + ctx, + handle_mem, + tokens_per_expert, + grad_output, + ) + return grad_input, [()], [(None, None, None)] diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 9263b41d51..2b34c9a940 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -20,6 +20,7 @@ ) from ...quantization import QuantizerRole, Recipe from ...tensor import MXFP8Quantizer, Quantizer +from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .._common import ( is_quantized_tensor, maybe_dequantize, @@ -44,24 +45,91 @@ def _validate_output_buffer( raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") if tensor.device != device: raise ValueError(f"{name} must be on {device}, got {tensor.device}.") - if not tensor.is_contiguous(): - raise ValueError(f"{name} must be contiguous.") if tensor.requires_grad: raise ValueError(f"{name} must not require gradients.") return tensor +def _validate_dispatch_input( + input_: torch.Tensor | MXFP8TensorStorage, + buffer: EpBuffer, +) -> tuple[tuple[int, int], torch.dtype, bool]: + """Validate the local token matrix and identify its transport format.""" + is_mxfp8 = isinstance(input_, MXFP8TensorStorage) + if is_quantized_tensor(input_) and not is_mxfp8: + raise TypeError( + "NCCL EP Dispatch supports BF16 and MXFP8 inputs, " + f"got {type(input_).__name__}." + ) + + input_shape = tuple(input_.shape) + input_dtype = input_.dtype if isinstance(input_, torch.Tensor) else input_._dtype + expected_hidden = buffer.hidden_dim + if len(input_shape) != 2 or input_shape[-1] != expected_hidden: + raise ValueError( + f"Dispatch input must have shape (T, {expected_hidden}), got {input_shape}." + ) + if input_.device != buffer.device: + raise ValueError(f"Dispatch input must be on {buffer.device}, got {input_.device}.") + if input_dtype is not torch.bfloat16: + raise TypeError( + "Dispatch input must be BF16 or an MXFP8 tensor representing BF16 values, " + f"got logical dtype {input_dtype}." + ) + + if is_mxfp8: + if input_._fp8_dtype != DType.kFloat8E4M3: + raise NotImplementedError( + f"NCCL EP Dispatch supports E4M3 MXFP8 only, got {input_._fp8_dtype}." + ) + + return input_shape, input_dtype, is_mxfp8 + + +def _validate_routing_inputs( + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + device: torch.device, +) -> None: + """Validate routing properties not checked by the native binding.""" + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + + +def _validate_mxfp8_output_buffer( + tensor: Optional[torch.Tensor], + *, + device: torch.device, +) -> None: + """Validate properties common to packed MXFP8 caller buffers. + + ``_scale_alloc_io`` validates contiguity and byte capacity after the data + and scale sizes are known. + """ + if tensor is None: + return + if tensor.device != device: + raise ValueError(f"recv_tokens must be on {device}, got {tensor.device}.") + if tensor.requires_grad: + raise ValueError("recv_tokens must not require gradients.") + + class Dispatch(BasicOperation): - """Dispatch BF16 tokens to local experts with NCCL EP. + """Dispatch BF16 or MXFP8 tokens to local experts with NCCL EP. The extra inputs are routing indices and FP32 routing weights. The extra - outputs are local tokens-per-expert and received routing weights. + outputs are local tokens-per-expert, received routing weights, the routing + handle, and routing indices for recovering the local token shape. """ num_extra_inputs: int = 2 # tokens-per-expert, received routing weights, and the opaque NCCL EP - # routing handle consumed by Combine. - num_extra_outputs: int = 3 + # routing handle and routing indices consumed by Combine. + num_extra_outputs: int = 4 def __init__(self, buffer: EpBuffer) -> None: super().__init__() @@ -127,64 +195,31 @@ def op_forward(self, *args: Any, **kwargs: Any) -> None: def op_backward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Dispatch uses fuser_backward") - def fuser_forward( - self, - basic_op_ctxs: list[OperationContext], - input_: torch.Tensor, - *, - basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - input_quantizer = self._resolve_input_quantizer(next_op_input_quantizer) - topk_idx, topk_weights = basic_op_extra_inputs[0] - kwargs = basic_op_kwargs[0] - - if input_.dtype is not torch.bfloat16: - raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") - if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: - raise ValueError( - f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " - f"got {tuple(input_.shape)}." - ) - if topk_idx.dtype not in (torch.int32, torch.int64): - raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}.") - expected_route_shape = (input_.shape[0], self.buffer.top_k) - if tuple(topk_idx.shape) != expected_route_shape: - raise ValueError( - f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." - ) - if tuple(topk_weights.shape) != expected_route_shape: - raise ValueError( - f"topk_weights shape must be {expected_route_shape}, " - f"got {tuple(topk_weights.shape)}." - ) - if topk_weights.dtype is not torch.float32: - raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") - for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): - if tensor.device != input_.device: - raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") - - recv_tokens = kwargs.get("recv_tokens") - recv_topk_weights = kwargs.get("recv_topk_weights") - if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): - raise ValueError( - "eager mode sizes dispatch outputs per step and cannot use " - "caller-supplied receive buffers" - ) - + def _prepare_routing(self, topk_idx: torch.Tensor) -> tuple[torch.Tensor, int]: + """Prepare NCCL routing state and return the receive row count.""" tokens_per_expert = ep_prepare(self.buffer, topk_idx) - rows = ( + num_recv_tokens = ( self.buffer._host_total_recv_tokens if self.buffer.eager else self.buffer.recv_capacity_per_rank ) - if rows is None: + if num_recv_tokens is None: raise RuntimeError("NCCL EP dispatch receive size is unavailable.") - rows = int(rows) - recv_shape = (rows, self.buffer.hidden_dim) - if input_quantizer is None: + return tokens_per_expert, int(num_recv_tokens) + + def _prepare_output_buffers( + self, + recv_tokens: Optional[torch.Tensor], + recv_topk_weights: Optional[torch.Tensor], + *, + num_recv_tokens: int, + use_mxfp8: bool, + ) -> tuple[Optional[torch.Tensor], torch.Tensor]: + """Validate caller buffers and allocate any missing dispatch outputs.""" + recv_shape = (num_recv_tokens, self.buffer.hidden_dim) + if use_mxfp8: + _validate_mxfp8_output_buffer(recv_tokens, device=self.buffer.device) + else: recv_tokens = _validate_output_buffer( "recv_tokens", recv_tokens, @@ -192,29 +227,52 @@ def fuser_forward( dtype=self.buffer.payload_dtype, device=self.buffer.device, ) + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + recv_topk_weights = _validate_output_buffer( "recv_topk_weights", recv_topk_weights, - shape=(rows,), + shape=(num_recv_tokens,), dtype=torch.float32, device=self.buffer.device, ) if recv_topk_weights is None: recv_topk_weights = _alloc_io( - (rows,), + (num_recv_tokens,), torch.float32, self.buffer.device, self.buffer.zero_copy, ) + return recv_tokens, recv_topk_weights - if input_quantizer is None: + def _dispatch_impl( + self, + input_: torch.Tensor | MXFP8TensorStorage, + input_dtype: torch.dtype, + input_quantizer: Optional[MXFP8Quantizer], + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + tokens_per_expert: torch.Tensor, + recv_tokens: Optional[torch.Tensor], + recv_topk_weights: torch.Tensor, + *, + num_recv_tokens: int, + use_mxfp8: bool, + ) -> torch.Tensor | MXFP8TensorStorage: + """Route tokens in the selected transport format. + + Keep transport-specific setup in this function so future quantized + formats can be added without branching in ``fuser_forward``. + """ + if not use_mxfp8: if recv_tokens is None: - recv_tokens = _alloc_io( - recv_shape, - self.buffer.payload_dtype, - self.buffer.device, - self.buffer.zero_copy, - ) + raise RuntimeError("BF16 dispatch receive storage was not allocated.") torch.ops.transformer_engine_ep.dispatch( self.buffer.handle_mem, topk_idx, @@ -223,46 +281,106 @@ def fuser_forward( recv_tokens, recv_topk_weights, ) - else: - quantized_input, input_scale_inv = quantize_mxfp8_for_ep(input_, input_quantizer) - scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE - recv_data, recv_scale_inv = _scale_alloc_io( - recv_tokens, - rows, - self.buffer.hidden_dim, - scale_cols, - quantized_input.rowwise_data.dtype, - input_scale_inv.dtype, - self.buffer.device, - self.buffer.zero_copy, - ) - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - quantized_input.rowwise_data.view(torch.float8_e4m3fn), - topk_weights, - recv_data.view(torch.float8_e4m3fn), - recv_topk_weights, - input_scale_inv, - recv_scale_inv, - ) - recv_tokens = _make_grouped_mxfp8( - recv_data, - recv_scale_inv, - tokens_per_expert, - quantized_input._fp8_dtype, - input_.dtype, + return recv_tokens + + quantized_input, input_scale_inv = quantize_mxfp8_for_ep(input_, input_quantizer) + scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE + recv_data, recv_scale_inv = _scale_alloc_io( + recv_tokens, + num_recv_tokens, + self.buffer.hidden_dim, + scale_cols, + quantized_input._rowwise_data.dtype, + input_scale_inv.dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + quantized_input._rowwise_data.view(torch.float8_e4m3fn), + topk_weights, + recv_data.view(torch.float8_e4m3fn), + recv_topk_weights, + input_scale_inv, + recv_scale_inv, + ) + return _make_grouped_mxfp8( + recv_data, + recv_scale_inv, + tokens_per_expert, + quantized_input._fp8_dtype, + input_dtype, + ) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor | MXFP8TensorStorage, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + # Resolve the transport format. A BF16 input uses MXFP8 transport when + # an input quantizer is active; an existing MXFP8 input is passed through. + input_quantizer = self._resolve_input_quantizer(next_op_input_quantizer) + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + input_shape, input_dtype, input_is_mxfp8 = _validate_dispatch_input(input_, self.buffer) + use_mxfp8 = input_is_mxfp8 or input_quantizer is not None + _validate_routing_inputs( + topk_idx, + topk_weights, + device=self.buffer.device, + ) + + # Eager mode discovers the receive size at runtime, so persistent + # caller-owned output buffers cannot be used. + recv_tokens = kwargs.get("recv_tokens") + recv_topk_weights = kwargs.get("recv_topk_weights") + if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): + raise ValueError( + "eager mode sizes dispatch outputs per step and cannot use " + "caller-supplied receive buffers" ) + # Prepare the routing handle, then size and allocate the receive outputs. + tokens_per_expert, num_recv_tokens = self._prepare_routing(topk_idx) + recv_tokens, recv_topk_weights = self._prepare_output_buffers( + recv_tokens, + recv_topk_weights, + num_recv_tokens=num_recv_tokens, + use_mxfp8=use_mxfp8, + ) + + # Launch NCCL EP using the selected transport representation. + output = self._dispatch_impl( + input_, + input_dtype, + input_quantizer, + topk_idx, + topk_weights, + tokens_per_expert, + recv_tokens, + recv_topk_weights, + num_recv_tokens=num_recv_tokens, + use_mxfp8=use_mxfp8, + ) + + # Save only shape/dtype metadata and the opaque routing handle for backward. ctx = basic_op_ctxs[0] if ctx.requires_grad: - ctx.input_shape = tuple(input_.shape) - ctx.input_dtype = input_.dtype + ctx.input_shape = input_shape + ctx.input_dtype = input_dtype ctx.topk_weights_shape = tuple(topk_weights.shape) ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer ctx.save_for_backward(self.buffer.handle_mem) - return recv_tokens, [(tokens_per_expert, recv_topk_weights, self.buffer.handle_mem)] + return output, [ + (tokens_per_expert, recv_topk_weights, self.buffer.handle_mem, topk_idx) + ] def fuser_backward( self, @@ -277,7 +395,7 @@ def fuser_backward( ]: ctx = basic_op_ctxs[0] (handle_mem,) = ctx.saved_tensors - grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() + grad_output = maybe_dequantize(grad_output, ctx.input_dtype) grad_recv_weights = basic_op_grad_extra_outputs[0][1] if grad_recv_weights is None: @@ -287,7 +405,7 @@ def fuser_backward( device=grad_output.device, ) else: - grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous() + grad_recv_weights = grad_recv_weights.to(dtype=torch.float32) grad_input = torch.empty( ctx.input_shape, @@ -307,6 +425,6 @@ def fuser_backward( grad_topk_weights, ) quantizer = ctx.prev_op_grad_output_quantizer - if quantizer is not None and not is_quantized_tensor(grad_input): + if quantizer is not None: grad_input = quantizer(grad_input) return grad_input, [()], [(None, grad_topk_weights)] diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index ace2643177..0d3e2b458e 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -10,8 +10,9 @@ from typing import Any, Optional import torch +import transformer_engine_torch as tex -from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import get_ep_group from ...quantization import Recipe from ...tensor import GroupedTensor, MXFP8Quantizer, Quantizer @@ -25,80 +26,12 @@ from ..op import FusedOperation, FusibleOperation, OperationContext -def _grouped_weight(op: GroupedLinear) -> GroupedTensor: - """Return the single packed ``(E, out, in)`` parameter required by MegaMoE.""" - if not op.single_grouped_weight or not isinstance(op.weight, GroupedTensor): - raise ValueError("FusedMoeEp requires GroupedLinear(single_grouped_weight=True)") - return op.weight - - -def _pack_grouped_linear_weights(op: GroupedLinear, *, block_scaled_cls: Optional[type] = None): - """View a packed TE ``(E, out, in)`` weight as MegaMoE ``(E, in, out)``. - - The permutation is intentionally not made contiguous. MegaMoE internally - permutes back to ``(E, out, in)`` before requesting contiguous storage, so - retaining this view lets that request reuse GroupedLinear's original packed - buffer. ``block_scaled_cls`` selects the ``BlockScaledTensor`` type (cuDNN - MegaMoE vs the PyTorch reference). - """ - weight = _grouped_weight(op) - num_groups = op.num_groups - out_features = op.out_features - in_features = op.in_features - expected_data_numel = num_groups * out_features * in_features - if weight.rowwise_data is None or weight.rowwise_data.numel() != expected_data_numel: - raise ValueError( - "GroupedLinear weight must have compact rowwise storage with shape " - f"({num_groups}, {out_features}, {in_features})" - ) - - data_nk = weight.rowwise_data.view(num_groups, out_features, in_features) - if weight.quantizer is not None: - recipe = weight.quantizer._get_compatible_recipe() - if recipe is None or not recipe.mxfp8(): - raise TypeError("FusedMoeEp only supports dense BF16 or MXFP8 grouped weights") - if weight._with_gemm_swizzled_scales: - raise NotImplementedError( - "FusedMoeEp requires unswizzled MXFP8 weight scales, got a GEMM-swizzled tensor" - ) - if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: - raise ValueError( - f"MXFP8 weight K={in_features} is not divisible by {MXFP8_BLOCK_SCALING_SIZE}" - ) - scale_cols = in_features // MXFP8_BLOCK_SCALING_SIZE - expected_scale_numel = num_groups * out_features * scale_cols - if weight.scale_inv is None or weight.scale_inv.numel() != expected_scale_numel: - raise ValueError( - "GroupedLinear MXFP8 scales must have compact rowwise storage with shape " - f"({num_groups}, {out_features}, {scale_cols})" - ) - if block_scaled_cls is None: - from cudnn.moe_ep import BlockScaledTensor as block_scaled_cls - packed_data = data_nk.view(torch.float8_e4m3fn).permute(0, 2, 1) - packed_scale = ( - weight.scale_inv.view(num_groups, out_features, scale_cols) - .view(torch.float8_e8m0fnu) - .permute(0, 2, 1) - ) - return block_scaled_cls( - data=packed_data, - scale=packed_scale, - format="mxfp8", - logical_shape=tuple(packed_data.shape), - axis=1, - ) - return data_nk.permute(0, 2, 1) - - -def _flatten_moe_weight(weight) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Split a MegaMoE weight into tensors that ``save_for_backward`` can hold.""" - if isinstance(weight, torch.Tensor): - return weight, None - return weight.data, weight.scale - - -def _restore_moe_weight(data: torch.Tensor, scale: Optional[torch.Tensor], block_scaled_cls: type): - """Rebuild the MegaMoE weight saved by :func:`_flatten_moe_weight`.""" +def _pack_as_cudnn_moe_tensor( + data: torch.Tensor, + scale: Optional[torch.Tensor], + block_scaled_cls: type, +): + """Represent payload and scales using the public cuDNN MoE tensor type.""" if scale is None: return data return block_scaled_cls( @@ -110,67 +43,78 @@ def _restore_moe_weight(data: torch.Tensor, scale: Optional[torch.Tensor], block ) -def _pack_activation( +def _pack_cudnn_activation( input_: torch.Tensor, quantizer: Optional[MXFP8Quantizer], block_scaled_cls: type, ): - """Represent a TE activation in the public cuDNN MoeEp format.""" + """Pack the dispatch input in the public cuDNN MoE activation layout.""" if quantizer is None: return input_ - quantized, scale_inv = quantize_mxfp8_for_ep(input_, quantizer) - return block_scaled_cls( - data=quantized.rowwise_data.view(torch.float8_e4m3fn), - scale=scale_inv.view(torch.float8_e8m0fnu), - format="mxfp8", - logical_shape=tuple(input_.shape), - axis=1, + quantized, scale = quantize_mxfp8_for_ep(input_, quantizer) + return _pack_as_cudnn_moe_tensor( + quantized._rowwise_data.view(torch.float8_e4m3fn), + scale.view(torch.float8_e8m0fnu), + block_scaled_cls, ) -def _validate_internal_quantizers( - dispatch: Dispatch, - fc1: GroupedLinear, - fc2: GroupedLinear, - combine: Combine, -) -> Optional[MXFP8Quantizer]: - """Validate quantizers hidden inside the five-op MegaMoE fusion.""" - dispatch_quantizer = dispatch.get_input_quantizer() - quantizers = { - "Dispatch input": dispatch_quantizer, - "Combine grad_output": combine.get_grad_output_quantizer(), - } - for op_name, op in (("FC1", fc1), ("FC2", fc2)): - for group_idx in range(op.num_groups): - quantizers[f"{op_name} input {group_idx}"] = op.get_quantizer( - "forward", - 2 * group_idx, - ) - quantizers[f"{op_name} grad_output {group_idx}"] = op.get_quantizer( - "backward", - group_idx, - ) - active = {name: quantizer for name, quantizer in quantizers.items() if quantizer is not None} - for name, quantizer in active.items(): - if not isinstance(quantizer, MXFP8Quantizer): - raise TypeError( - f"FusedMoeEp supports MXFP8 internal quantizers only; {name} uses " - f"{type(quantizer).__name__}." - ) - if quantizer.dtype != DType.kFloat8E4M3: - raise NotImplementedError(f"FusedMoeEp requires E4M3 MXFP8 for {name}.") - if active and len(active) != len(quantizers): - missing = ", ".join(name for name, quantizer in quantizers.items() if quantizer is None) - raise ValueError( - "FusedMoeEp requires either an all-BF16 boundary or a complete MXFP8 " - f"quantizer set; missing {missing}." +def _pack_cudnn_weights( + op: GroupedLinear, + *, + block_scaled_cls: Optional[type] = None, +): + """Pack a GroupedLinear ``(E, out, in)`` weight as cuDNN ``(E, in, out)``. + + The permutation is intentionally not made contiguous. MegaMoE internally + permutes back to ``(E, out, in)`` before requesting contiguous storage, so + retaining this view lets quantized-model-init weights reuse their original + packed buffer. Dense parameters are quantized with GroupedLinear's weight + quantizer, matching its normal MXFP8 forward path. + """ + weight = op.weight + num_groups = op.num_groups + out_features = op.out_features + in_features = op.in_features + weight_quantizer = op.get_quantizer("forward", 1) + if weight_quantizer is None and weight.quantizer is None: + return weight.rowwise_data.view( + num_groups, + out_features, + in_features, + ).permute(0, 2, 1) + + if weight_quantizer is not None and weight.quantizer is None: + weight_quantizer.set_usage(rowwise=True) + weight = tex.group_quantize( + weight.rowwise_data.view(weight.logical_shape), + weight_quantizer, + op.num_groups, + None, ) - return dispatch_quantizer or fc1.get_quantizer("forward", 0) + scale_cols = in_features // MXFP8_BLOCK_SCALING_SIZE + data = ( + weight.rowwise_data.view(num_groups, out_features, in_features) + .view(torch.float8_e4m3fn) + .permute(0, 2, 1) + ) + scale = ( + weight.scale_inv.view(num_groups, out_features, scale_cols) + .view(torch.float8_e8m0fnu) + .permute(0, 2, 1) + ) + if block_scaled_cls is None: + from cudnn.moe_ep import BlockScaledTensor as block_scaled_cls + return _pack_as_cudnn_moe_tensor( + data, + scale, + block_scaled_cls, + ) def _grouped_weight_grad(op: GroupedLinear, grad: torch.Tensor) -> list[Optional[torch.Tensor]]: """Convert a MegaMoE ``(E, in, out)`` wgrad to one TE ``(E, out, in)`` grad.""" - weight = _grouped_weight(op) + weight = op.weight expected_shape = (op.num_groups, op.in_features, op.out_features) if tuple(grad.shape) != expected_shape: raise RuntimeError( @@ -242,8 +186,15 @@ def _routing_extras_internal( sequence when those two outputs feed exactly these ops and are not returned to the caller. """ - tokens_per_expert, routing_weights, ep_handle = dispatch._extra_output_channels - if tokens_per_expert is None or routing_weights is None or ep_handle is None: + tokens_per_expert, routing_weights, ep_handle, routing_indices = ( + dispatch._extra_output_channels + ) + if ( + tokens_per_expert is None + or routing_weights is None + or ep_handle is None + or routing_indices is None + ): return False if any(dispatch._extra_output_to_caller): return False @@ -253,6 +204,7 @@ def _routing_extras_internal( and activation._extra_input_channels[0] == routing_weights and combine._extra_input_channels[0] == ep_handle and combine._extra_input_channels[1] == tokens_per_expert + and combine._extra_input_channels[2] == routing_indices ) @@ -278,7 +230,7 @@ def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: if len(window) != 5: return False - if recipe is not None and not recipe.mxfp8(): + if recipe is None or not recipe.mxfp8(): return False dispatch, fc1, activation, fc2, combine = window if not ( @@ -290,8 +242,6 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo ): return False buffer = dispatch.buffer - if combine.num_local_tokens != buffer.max_tokens_per_rank: - return False if not buffer.eager or buffer.payload_dtype is not torch.bfloat16: return False if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): @@ -379,24 +329,18 @@ def fuser_forward( topk_idx, topk_weights = basic_op_extra_inputs[0] if topk_weights.dtype is not torch.float32: raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") - with torch.no_grad(): - input_quantizer = _validate_internal_quantizers( - self.dispatch, - self.fc1, - self.fc2, - self.basic_ops[4], - ) - activation = _pack_activation( - input_, - input_quantizer, - self._block_scaled_cls, - ) - fc1_weight = _pack_grouped_linear_weights( - self.fc1, block_scaled_cls=self._block_scaled_cls - ) - fc2_weight = _pack_grouped_linear_weights( - self.fc2, block_scaled_cls=self._block_scaled_cls - ) + input_quantizer = self.dispatch.get_quantizer("forward", 0) + activation = _pack_cudnn_activation( + input_, + input_quantizer, + self._block_scaled_cls, + ) + fc1_weight = _pack_cudnn_weights( + self.fc1, block_scaled_cls=self._block_scaled_cls + ) + fc2_weight = _pack_cudnn_weights( + self.fc2, block_scaled_cls=self._block_scaled_cls + ) output, fc1_c, route_metadata = self._moe( activation, fc1_weight, @@ -406,9 +350,9 @@ def fuser_forward( ) if any(ctx.requires_grad for ctx in basic_op_ctxs): - input_data, input_scale = _flatten_moe_weight(activation) - fc1_data, fc1_scale = _flatten_moe_weight(fc1_weight) - fc2_data, fc2_scale = _flatten_moe_weight(fc2_weight) + input_data, input_scale = activation.data, activation.scale + fc1_data, fc1_scale = fc1_weight.data, fc1_weight.scale + fc2_data, fc2_scale = fc2_weight.data, fc2_weight.scale basic_op_ctxs[0].save_for_backward( input_data, input_scale, @@ -429,7 +373,7 @@ def fuser_forward( if next_op_input_quantizer is not None and not is_quantized_tensor(output): output = next_op_input_quantizer(output) return output, [ - (None, None, None), + (None, None, None, None), (), (), (), @@ -460,13 +404,13 @@ def fuser_backward( fc1_c, route_metadata, ) = basic_op_ctxs[0].saved_tensors - input_ = _restore_moe_weight( + input_ = _pack_as_cudnn_moe_tensor( input_data, input_scale, self._block_scaled_cls, ) - fc1_weight = _restore_moe_weight(fc1_data, fc1_scale, self._block_scaled_cls) - fc2_weight = _restore_moe_weight(fc2_data, fc2_scale, self._block_scaled_cls) + fc1_weight = _pack_as_cudnn_moe_tensor(fc1_data, fc1_scale, self._block_scaled_cls) + fc2_weight = _pack_as_cudnn_moe_tensor(fc2_data, fc2_scale, self._block_scaled_cls) grad_output = maybe_dequantize( grad_output, basic_op_ctxs[0].input_dtype, From 1d5d2ede425f082021707397f6e97430b19c4197 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 09:38:37 +0000 Subject: [PATCH 58/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/basic/combine.py | 4 +--- transformer_engine/pytorch/ops/basic/dispatch.py | 7 ++----- transformer_engine/pytorch/ops/fused/moe_ep.py | 12 +++--------- 3 files changed, 6 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 1f65f2c741..f723a87e1b 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -62,9 +62,7 @@ def _validate_combine_inputs( if tokens_per_expert.dtype is not torch.int64 or tokens_per_expert.device != input_.device: raise ValueError("Combine tokens_per_expert must be an int64 tensor on the input device.") if topk_idx.ndim != 2: - raise ValueError( - f"Combine routing indices must be 2D, got shape {tuple(topk_idx.shape)}." - ) + raise ValueError(f"Combine routing indices must be 2D, got shape {tuple(topk_idx.shape)}.") if topk_idx.dtype not in (torch.int32, torch.int64) or topk_idx.device != input_.device: raise ValueError( "Combine routing indices must be an int32 or int64 tensor on the input device." diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 2b34c9a940..22273c4ec6 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -58,8 +58,7 @@ def _validate_dispatch_input( is_mxfp8 = isinstance(input_, MXFP8TensorStorage) if is_quantized_tensor(input_) and not is_mxfp8: raise TypeError( - "NCCL EP Dispatch supports BF16 and MXFP8 inputs, " - f"got {type(input_).__name__}." + f"NCCL EP Dispatch supports BF16 and MXFP8 inputs, got {type(input_).__name__}." ) input_shape = tuple(input_.shape) @@ -378,9 +377,7 @@ def fuser_forward( ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer ctx.save_for_backward(self.buffer.handle_mem) - return output, [ - (tokens_per_expert, recv_topk_weights, self.buffer.handle_mem, topk_idx) - ] + return output, [(tokens_per_expert, recv_topk_weights, self.buffer.handle_mem, topk_idx)] def fuser_backward( self, diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 0d3e2b458e..ee2a8b6033 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -186,9 +186,7 @@ def _routing_extras_internal( sequence when those two outputs feed exactly these ops and are not returned to the caller. """ - tokens_per_expert, routing_weights, ep_handle, routing_indices = ( - dispatch._extra_output_channels - ) + tokens_per_expert, routing_weights, ep_handle, routing_indices = dispatch._extra_output_channels if ( tokens_per_expert is None or routing_weights is None @@ -335,12 +333,8 @@ def fuser_forward( input_quantizer, self._block_scaled_cls, ) - fc1_weight = _pack_cudnn_weights( - self.fc1, block_scaled_cls=self._block_scaled_cls - ) - fc2_weight = _pack_cudnn_weights( - self.fc2, block_scaled_cls=self._block_scaled_cls - ) + fc1_weight = _pack_cudnn_weights(self.fc1, block_scaled_cls=self._block_scaled_cls) + fc2_weight = _pack_cudnn_weights(self.fc2, block_scaled_cls=self._block_scaled_cls) output, fc1_c, route_metadata = self._moe( activation, fc1_weight, From d13c17fc7f7d8ddcb282b553fa1e9048b57a8009 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 23 Aug 2026 10:13:11 +0000 Subject: [PATCH 59/83] more cleanup Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 48 +++++++++---------- .../pytorch/ops/basic/dispatch.py | 10 ++-- .../pytorch/ops/fused/moe_ep.py | 24 ++++++++-- 3 files changed, 48 insertions(+), 34 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6e18975fd2..ae818c2b17 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -29,7 +29,11 @@ _ep_dispatch_raw, ) from transformer_engine.pytorch.ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat -from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp, _pack_cudnn_weights +from transformer_engine.pytorch.ops.fused.moe_ep import ( + FusedMoeEp, + _cudnn_megamoe_supported, + _pack_cudnn_weights, +) from transformer_engine.pytorch.ops.op import OperationContext from transformer_engine.pytorch.tensor import GroupedTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor @@ -328,19 +332,16 @@ def _moe_step(self, buffer, topk_idx, tokens, w): def _make_moe_model( self, *, - fuse_ops=True, intermediate_dim=INTERMEDIATE_DIM, recipe=None, ): """Build an EP MoE Sequential. - With ``fuse_ops=True``, dispatch routing extras stay internal so - MegaMoE can claim the sequence when its gates pass. With - ``fuse_ops=False``, those extras are returned to the caller and the - Sequential stays unfused. Pass an MXFP8 ``recipe`` to construct - natively quantized GroupedLinear weights via ``quantized_model_init``. + Routing extras stay internal so MegaMoE can claim the sequence when + its gates pass. Pass an MXFP8 ``recipe`` to construct natively + quantized GroupedLinear weights via ``quantized_model_init``. """ - buffer = self._make_buffer() + buffer = self._make_buffer(alignment=128 if recipe is not None else 0) dispatch = te_ops.Dispatch(buffer) init_ctx = ( te.quantized_model_init(enabled=True, recipe=recipe) @@ -377,8 +378,8 @@ def _make_moe_model( os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param combine = te_ops.Combine() - dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) - dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) fc1.set_extra_input_channel(0, "tokens_per_expert") @@ -761,20 +762,20 @@ def test_caller_provides_grad_expert_out(self): @_eager_test_include def test_bf16_moe_sequential_vs_reference(self): - """Unfused BF16 Sequential matches a BF16 ``MoeEpReference``.""" + """BF16 Sequential matches its reference, with MegaMoE when available.""" self._run_moe_sequential_vs_reference(quantization="bf16") @_eager_test_include + @_mxfp8_align_test def test_mxfp8_moe_sequential_vs_reference(self): - """Fused MegaMoE Sequential under MXFP8 autocast matches the QDQ reference.""" + """MXFP8 Sequential matches the QDQ reference, with MegaMoE when available.""" self._run_moe_sequential_vs_reference(quantization="mxfp8") def _run_moe_sequential_vs_reference(self, *, quantization): - """Compare Sequential to ``MoeEpReference``. + """Compare a Sequential MoE to ``MoeEpReference``. - Sequential uses MegaMoE for MXFP8 when supported. The MXFP8 reference - QDQ GEMM operands to MXFP8 and matmuls in FP32; the BF16 reference runs - FP32 GEMMs with no QDQ. + MegaMoE is used when its cuDNN and hardware requirements are met; + otherwise the basic EP and grouped-MLP operations run unfused. """ if not EAGER: self.skipTest("variable-size reference comparison requires eager EP mode") @@ -789,7 +790,8 @@ def _run_moe_sequential_vs_reference(self, *, quantization): intermediate_dim = 256 recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None model, fc1, fc2 = self._make_moe_model( - fuse_ops=True, intermediate_dim=intermediate_dim, recipe=recipe + intermediate_dim=intermediate_dim, + recipe=recipe, ) generator = torch.Generator(device=self.cfg.device) generator.manual_seed(3100 + self.cfg.rank) @@ -826,11 +828,10 @@ def _run_moe_sequential_vs_reference(self, *, quantization): seq_out = model(seq_tokens, topk_idx, seq_topk_weights) forward_ops = model._module_groups[0]._forward_ops - if quantization == "mxfp8": + if _cudnn_megamoe_supported(): self.assertEqual(len(forward_ops), 1) self.assertIsInstance(forward_ops[0][0], FusedMoeEp) else: - self.assertEqual(len(forward_ops), 5) self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in forward_ops)) self.assertIsInstance(seq_out, torch.Tensor) self.assertEqual(seq_out.dtype, torch.bfloat16) @@ -855,9 +856,7 @@ def _run_moe_sequential_vs_reference(self, *, quantization): combine_format=MoeFormat.BF16, apply_topk_in_fc1=True, generate_c=True, - compute_dtype=torch.float32 if quantization == "mxfp8" else torch.bfloat16, - gemm_format=MoeFormat.MXFP8 if quantization == "mxfp8" else MoeFormat.BF16, - ) + compute_dtype=torch.float32) ref_out, fc1_c, route_metadata = reference( tokens.detach(), fc1_weight, @@ -888,10 +887,7 @@ def _run_moe_sequential_vs_reference(self, *, quantization): route_metadata, ) torch.cuda.synchronize() - if quantization == "bf16": - tolerances = {"rtol": 1.6e-2, "atol": 5e-5} - else: - tolerances = {"rtol": 0.125, "atol": 0.25} + tolerances = {"rtol": 0.125, "atol": 0.25} torch.testing.assert_close(seq_out, ref_out, **tolerances) torch.testing.assert_close( seq_tokens.grad, grad_tokens.to(dtype=seq_tokens.dtype), **tolerances diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 2b34c9a940..c9b0d7866e 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -171,11 +171,13 @@ def _resolve_input_quantizer( if ( quantizer is not None and next_op_input_quantizer is not None - and quantizer is not next_op_input_quantizer + and type(quantizer) is not type(next_op_input_quantizer) ): - raise ValueError( - "Dispatch input_quantizer and next operation input quantizer " - "must be the same object when both are set." + raise TypeError( + "Dispatch input quantizer and next operation input quantizer " + "must have the same type, got " + f"{type(quantizer).__name__} and " + f"{type(next_op_input_quantizer).__name__}." ) if quantizer is None: quantizer = next_op_input_quantizer diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 0d3e2b458e..0a979569d1 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -7,8 +7,10 @@ from __future__ import annotations from collections.abc import Iterable, Sequence +from importlib.metadata import PackageNotFoundError, version as get_pkg_version from typing import Any, Optional +from packaging.version import Version as PkgVersion import torch import transformer_engine_torch as tex @@ -26,6 +28,15 @@ from ..op import FusedOperation, FusibleOperation, OperationContext +def _cudnn_megamoe_supported() -> bool: + """Whether the installed cuDNN frontend includes the public MegaMoE API.""" + return True + # try: + # return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.28.0") + # except PackageNotFoundError: + # return False + + def _pack_as_cudnn_moe_tensor( data: torch.Tensor, scale: Optional[torch.Tensor], @@ -210,6 +221,8 @@ def _routing_extras_internal( def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: """Static MegaMoE capability gates that can be checked before first launch.""" + if not _cudnn_megamoe_supported(): + return False if _import_cudnn_moe_ep() is None: return False if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): @@ -230,7 +243,7 @@ def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: if len(window) != 5: return False - if recipe is None or not recipe.mxfp8(): + if recipe is not None and not recipe.mxfp8(): return False dispatch, fc1, activation, fc2, combine = window if not ( @@ -350,9 +363,12 @@ def fuser_forward( ) if any(ctx.requires_grad for ctx in basic_op_ctxs): - input_data, input_scale = activation.data, activation.scale - fc1_data, fc1_scale = fc1_weight.data, fc1_weight.scale - fc2_data, fc2_scale = fc2_weight.data, fc2_weight.scale + input_data = activation if isinstance(activation, torch.Tensor) else activation.data + input_scale = None if isinstance(activation, torch.Tensor) else activation.scale + fc1_data = fc1_weight if isinstance(fc1_weight, torch.Tensor) else fc1_weight.data + fc1_scale = None if isinstance(fc1_weight, torch.Tensor) else fc1_weight.scale + fc2_data = fc2_weight if isinstance(fc2_weight, torch.Tensor) else fc2_weight.data + fc2_scale = None if isinstance(fc2_weight, torch.Tensor) else fc2_weight.scale basic_op_ctxs[0].save_for_backward( input_data, input_scale, From 8e18c46df2574a4da515b7e5bdd85f8792a46eae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 10:14:40 +0000 Subject: [PATCH 60/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_ep.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index ae818c2b17..dcf79c38ca 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -856,7 +856,8 @@ def _run_moe_sequential_vs_reference(self, *, quantization): combine_format=MoeFormat.BF16, apply_topk_in_fc1=True, generate_c=True, - compute_dtype=torch.float32) + compute_dtype=torch.float32, + ) ref_out, fc1_c, route_metadata = reference( tokens.detach(), fc1_weight, From 895d1a13c7fcda455b8d818f535ef88ad8528fd7 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 23 Aug 2026 22:31:45 +0000 Subject: [PATCH 61/83] refactor tests Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 740 +++++++++-------- transformer_engine/pytorch/ep_reference.py | 777 +++++++++++------- transformer_engine/pytorch/ops/_common.py | 11 +- .../pytorch/ops/basic/combine.py | 10 +- .../pytorch/ops/basic/dispatch.py | 5 + .../pytorch/ops/fused/moe_ep.py | 193 +++-- 6 files changed, 1053 insertions(+), 683 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index dcf79c38ca..ce3b069bea 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -28,14 +28,18 @@ _ep_combine_raw, _ep_dispatch_raw, ) -from transformer_engine.pytorch.ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat +from transformer_engine.pytorch.ep_reference import ( + BlockScaledTensor, + MoeEpReference, + MoeFormat, + quantize_blockwise, +) from transformer_engine.pytorch.ops.fused.moe_ep import ( FusedMoeEp, _cudnn_megamoe_supported, + _pack_cudnn_activation, _pack_cudnn_weights, ) -from transformer_engine.pytorch.ops.op import OperationContext -from transformer_engine.pytorch.tensor import GroupedTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" @@ -51,7 +55,6 @@ # NVTE_EP_TOKENS_PER_RANK. HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512")) TOP_K = 2 -INTERMEDIATE_DIM = 16 TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32")) @@ -225,16 +228,24 @@ def _make_cfg() -> _Cfg: return cfg -class TestEP(unittest.TestCase): +class _EpTestCase(unittest.TestCase): + """Shared NCCL EP process state and pass selection.""" + cfg: _Cfg ep_group: dist.ProcessGroup @classmethod def setUpClass(cls): + if hasattr(_EpTestCase, "cfg"): + cls.cfg = _EpTestCase.cfg + cls.ep_group = _EpTestCase.ep_group + return if _device_sm() < 90: raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") - cls.cfg = _make_cfg() - cls.ep_group = _build_ep_group() + _EpTestCase.cfg = _make_cfg() + _EpTestCase.ep_group = _build_ep_group() + cls.cfg = _EpTestCase.cfg + cls.ep_group = _EpTestCase.ep_group ep_bootstrap( cls.ep_group, num_experts=cls.cfg.num_experts, @@ -290,6 +301,10 @@ def _make_buffer( combine_bwd_quant_recipe=combine_bwd_quant_recipe, ) + +class TestEP(_EpTestCase): + """NCCL EP tests inherited from nvidia_origin/main.""" + def _expert_out(self, expert_out): """Stage the combine input into symm-mem under zero-copy (combine requires it).""" if not ZERO_COPY: @@ -329,80 +344,6 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) - def _make_moe_model( - self, - *, - intermediate_dim=INTERMEDIATE_DIM, - recipe=None, - ): - """Build an EP MoE Sequential. - - Routing extras stay internal so MegaMoE can claim the sequence when - its gates pass. Pass an MXFP8 ``recipe`` to construct natively - quantized GroupedLinear weights via ``quantized_model_init``. - """ - buffer = self._make_buffer(alignment=128 if recipe is not None else 0) - dispatch = te_ops.Dispatch(buffer) - init_ctx = ( - te.quantized_model_init(enabled=True, recipe=recipe) - if recipe is not None - else nullcontext() - ) - previous_single_param = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") - os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" - try: - with init_ctx: - fc1 = te_ops.GroupedLinear( - NUM_LOCAL_EXPERTS, - HIDDEN_DIM, - 2 * intermediate_dim, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - single_grouped_weight=True, - ) - activation = te_ops.ScaledSwiGLU() - fc2 = te_ops.GroupedLinear( - NUM_LOCAL_EXPERTS, - intermediate_dim, - HIDDEN_DIM, - bias=False, - device=self.cfg.device, - dtype=torch.bfloat16, - single_grouped_weight=True, - ) - finally: - if previous_single_param is None: - del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] - else: - os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param - combine = te_ops.Combine() - - dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) - dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) - dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) - dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) - fc1.set_extra_input_channel(0, "tokens_per_expert") - activation.set_extra_input_channel(0, "routing_weights") - fc2.set_extra_input_channel(0, "tokens_per_expert") - combine.set_extra_input_channel(0, "ep_handle") - combine.set_extra_input_channel(1, "tokens_per_expert") - combine.set_extra_input_channel(2, "routing_indices") - return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 - - def _make_dispatch_combine_model(self, buffer): - """Build the minimal channel-routed Dispatch -> Combine pipeline.""" - dispatch = te_ops.Dispatch(buffer) - combine = te_ops.Combine() - dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) - dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) - dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) - dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) - combine.set_extra_input_channel(0, "ep_handle") - combine.set_extra_input_channel(1, "tokens_per_expert") - combine.set_extra_input_channel(2, "routing_indices") - return te_ops.Sequential(dispatch, combine) - # Prepare @_eager_test_include @@ -514,30 +455,6 @@ def test_primitive_dispatch_combine_identity(self): # Autograd - @_eager_test_include - @_zero_copy_test_include - def test_basic_ops_channel_routed_identity(self): - """Combine consumes Dispatch's handle/count channels without receiving EpBuffer.""" - buffer = self._make_buffer() - model = self._make_dispatch_combine_model(buffer) - topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) - tokens = tokens.detach().clone().requires_grad_(True) - output = model(tokens, topk_idx, weights) - (0.5 * output.float().square().sum()).backward() - torch.cuda.synchronize() - torch.testing.assert_close( - output.float(), - tokens.detach().float() * TOP_K, - atol=5e-2, - rtol=5e-2, - ) - torch.testing.assert_close( - tokens.grad.float(), - tokens.detach().float() * (TOP_K**2), - atol=1e-1, - rtol=1e-1, - ) - @_zero_copy_test_include def test_dispatch_autograd(self): """0.5*||recv_tokens||^2 ; grad_tokens equals TOP_K * tokens. Covers the @@ -608,88 +525,6 @@ def test_dispatch_mxfp8(self): self.assertTrue(is_symm_backed(recv_mx.scale_inv)) self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc) - @_eager_test_include - @_zero_copy_test_include - @_mxfp8_align_test - def test_basic_dispatch_prequantized_mxfp8(self): - """A fresh basic Dispatch accepts an already-quantized MXFP8 input.""" - self._require_mxfp8_shapes() - buffer = self._make_buffer(alignment=128) - dispatch = te_ops.Dispatch(buffer) - topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) - quantized_tokens = self._mxfp8_quantizer().quantize(tokens) - - recv_tokens, counts, _recv_weights, handle, _routing_indices = dispatch( - quantized_tokens, topk_idx, weights - ) - - self.assertIsInstance(recv_tokens, GroupedTensor) - self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) - torch.cuda.synchronize() - self.assertEqual(counts.shape, (NUM_LOCAL_EXPERTS,)) - recv_data = _degroup_mxfp8(recv_tokens).float() - self.assertTrue(torch.isfinite(recv_data).all()) - self.assertGreater(recv_data.abs().sum().item(), 0.0) - - @_eager_test_include - @_zero_copy_test_include - @_mxfp8_align_test - def test_basic_dispatch_recipe_mxfp8(self): - """A fresh basic Dispatch quantizes BF16 input under MXFP8 autocast.""" - self._require_mxfp8_shapes() - buffer = self._make_buffer(alignment=128) - dispatch = te_ops.Dispatch(buffer) - topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) - - with te.autocast(enabled=True, recipe=MXFP8BlockScaling()): - recv_tokens, counts, _recv_weights, handle, _routing_indices = dispatch( - tokens, topk_idx, weights - ) - - self.assertIsInstance(recv_tokens, GroupedTensor) - self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) - torch.cuda.synchronize() - self.assertEqual(counts.shape, (NUM_LOCAL_EXPERTS,)) - recv_data = _degroup_mxfp8(recv_tokens).float() - self.assertTrue(torch.isfinite(recv_data).all()) - self.assertGreater(recv_data.abs().sum().item(), 0.0) - - @_eager_test_include - @_zero_copy_test_include - @_mxfp8_align_test - def test_basic_combine_backward_mxfp8(self): - """Combine's recipe-owned grad-output quantizer drives scaled combine_bwd.""" - self._require_mxfp8_shapes() - buffer = self._make_buffer(alignment=128) - topk_idx, tokens, weights = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) - recv_tokens, _recv_weights, counts = ep_dispatch( - buffer, - tokens, - topk_idx, - weights, - ) - combine = te_ops.Combine() - combine.reset_recipe_state(recipe=MXFP8BlockScaling()) - combine.pre_fuser_forward(requires_grad=True) - ctx = OperationContext() - output, _ = combine.fuser_forward( - [ctx], - recv_tokens, - basic_op_extra_inputs=[(buffer.handle_mem, counts, topk_idx)], - prev_op_grad_output_quantizer=None, - next_op_input_quantizer=None, - basic_op_kwargs=[{}], - ) - ctx.saved_tensors = ctx.to_save - grad_input, _, _ = combine.fuser_backward( - [ctx], - torch.ones_like(output), - basic_op_grad_extra_outputs=[()], - ) - self.assertIsInstance(grad_input, GroupedTensor) - self.assertIsNotNone(grad_input.rowwise_data) - self.assertIsNotNone(grad_input.scale_inv) - @_zero_copy_test_include @_mxfp8_align_test def test_caller_provides_dispatch_recv_mxfp8(self): @@ -760,154 +595,6 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) - @_eager_test_include - def test_bf16_moe_sequential_vs_reference(self): - """BF16 Sequential matches its reference, with MegaMoE when available.""" - self._run_moe_sequential_vs_reference(quantization="bf16") - - @_eager_test_include - @_mxfp8_align_test - def test_mxfp8_moe_sequential_vs_reference(self): - """MXFP8 Sequential matches the QDQ reference, with MegaMoE when available.""" - self._run_moe_sequential_vs_reference(quantization="mxfp8") - - def _run_moe_sequential_vs_reference(self, *, quantization): - """Compare a Sequential MoE to ``MoeEpReference``. - - MegaMoE is used when its cuDNN and hardware requirements are met; - otherwise the basic EP and grouped-MLP operations run unfused. - """ - if not EAGER: - self.skipTest("variable-size reference comparison requires eager EP mode") - if torch.cuda.get_device_capability() != (10, 7): - self.skipTest("MegaMoE fusion requires Rubin SM107") - try: - from cudnn.moe_ep import MoeEp # noqa: F401 - except ImportError as exc: - self.skipTest(f"cudnn.moe_ep.MoeEp is not installed ({type(exc).__name__}: {exc})") - - # MegaMoE SM107 requires intermediate_size % 256 == 0. - intermediate_dim = 256 - recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None - model, fc1, fc2 = self._make_moe_model( - intermediate_dim=intermediate_dim, - recipe=recipe, - ) - generator = torch.Generator(device=self.cfg.device) - generator.manual_seed(3100 + self.cfg.rank) - with torch.no_grad(): - for op in (fc1, fc2): - weights = op.weight.quantized_tensors - if weights is None: - weights = op.weight.split_into_quantized_tensors() - for expert in range(NUM_LOCAL_EXPERTS): - weight = ( - torch.randn( - weights[expert].shape, - generator=generator, - dtype=torch.float32, - device=self.cfg.device, - ) - * 0.1 - ).to(torch.bfloat16) - weights[expert].copy_(weight) - if quantization == "mxfp8": - self.assertIsInstance(weights[expert], MXFP8Tensor) - - topk_idx, tokens, topk_weights = _make_moe_inputs( - self.cfg.rank, - self.cfg.ep_size, - self.cfg.device, - ) - seq_tokens = tokens.detach().clone().requires_grad_(True) - seq_topk_weights = topk_weights.detach().clone().requires_grad_(True) - autocast_ctx = ( - te.autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() - ) - with autocast_ctx: - seq_out = model(seq_tokens, topk_idx, seq_topk_weights) - - forward_ops = model._module_groups[0]._forward_ops - if _cudnn_megamoe_supported(): - self.assertEqual(len(forward_ops), 1) - self.assertIsInstance(forward_ops[0][0], FusedMoeEp) - else: - self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in forward_ops)) - self.assertIsInstance(seq_out, torch.Tensor) - self.assertEqual(seq_out.dtype, torch.bfloat16) - - fc1_weight = _reference_weights(fc1) - fc2_weight = _reference_weights(fc2) - for op, packed in ((fc1, fc1_weight), (fc2, fc2_weight)): - packed_data = packed.data if isinstance(packed, BlockScaledTensor) else packed - self.assertEqual(packed_data.data_ptr(), op.weight.rowwise_data.data_ptr()) - self.assertTrue(packed_data.permute(0, 2, 1).is_contiguous()) - if isinstance(packed, BlockScaledTensor): - self.assertEqual(packed.scale.data_ptr(), op.weight.scale_inv.data_ptr()) - self.assertTrue(packed.scale.permute(0, 2, 1).is_contiguous()) - reference = MoeEpReference( - num_experts=self.cfg.num_experts, - hidden_size=HIDDEN_DIM, - intermediate_size=intermediate_dim, - top_k=TOP_K, - ep_group=self.ep_group, - max_tokens_per_rank=TOKENS_PER_RANK, - output_format=MoeFormat.BF16, - combine_format=MoeFormat.BF16, - apply_topk_in_fc1=True, - generate_c=True, - compute_dtype=torch.float32, - ) - ref_out, fc1_c, route_metadata = reference( - tokens.detach(), - fc1_weight, - fc2_weight, - topk_idx, - topk_weights.detach(), - ) - self.assertEqual(ref_out.dtype, torch.bfloat16) - - dy = ( - torch.randn( - seq_out.shape, - generator=generator, - dtype=torch.float32, - device=self.cfg.device, - ) - * 0.1 - ).to(torch.bfloat16) - seq_out.backward(dy) - grad_tokens, grad_fc1, grad_fc2, grad_topk_weights = reference.backward( - dy, - tokens.detach(), - fc1_weight, - fc2_weight, - topk_idx, - topk_weights.detach(), - fc1_c, - route_metadata, - ) - torch.cuda.synchronize() - tolerances = {"rtol": 0.125, "atol": 0.25} - torch.testing.assert_close(seq_out, ref_out, **tolerances) - torch.testing.assert_close( - seq_tokens.grad, grad_tokens.to(dtype=seq_tokens.dtype), **tolerances - ) - torch.testing.assert_close(seq_topk_weights.grad, grad_topk_weights.float(), **tolerances) - for op, ref_grad in ((fc1, grad_fc1), (fc2, grad_fc2)): - seq_grad = op.weight.grad - self.assertEqual(seq_grad.dtype, torch.bfloat16) - self.assertEqual( - tuple(seq_grad.shape), - (NUM_LOCAL_EXPERTS, op.out_features, op.in_features), - ) - self.assertTrue(seq_grad.is_contiguous()) - torch.testing.assert_close( - seq_grad, - ref_grad.transpose(1, 2).to(dtype=seq_grad.dtype), - **tolerances, - ) - @_zero_copy_test_include @_mxfp8_align_test def test_combine_bwd_mxfp8_caller_grad_out(self): @@ -1188,6 +875,382 @@ def test_combine_autograd(self): torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) +class TestMoeEpSequential(_EpTestCase): + """Integration tests for Dispatch -> expert MLP -> Combine sequences.""" + + def _mxfp8_quantizer(self): + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + + def _require_mxfp8_shapes(self): + if HIDDEN_DIM % 512 != 0 or TOKENS_PER_RANK % 32 != 0: + self.skipTest( + "MXFP8 needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0 " + "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)" + ) + + def _make_dispatch_combine_ops(self, *, mxfp8): + buffer = self._make_buffer(alignment=128 if mxfp8 else 0) + return buffer, te_ops.Dispatch(buffer), te_ops.Combine() + + def _run_dispatch_combine_identity(self, *, mxfp8): + """Route, apply top-k weights, and combine back to local token order.""" + if mxfp8: + self._require_mxfp8_shapes() + buffer, dispatch, combine = self._make_dispatch_combine_ops(mxfp8=mxfp8) + topk_idx, tokens, topk_weights = _make_identity_inputs( + self.cfg.rank, + self.cfg.ep_size, + ) + recipe = MXFP8BlockScaling() if mxfp8 else None + with te.autocast(enabled=mxfp8, recipe=recipe): + recv_tokens, tokens_per_expert, recv_weights, handle, routing_indices = dispatch( + tokens, + topk_idx, + topk_weights, + ) + if mxfp8: + recv_tokens = _degroup_mxfp8(recv_tokens) + recv_weights = recv_weights[: recv_tokens.shape[0]] + weighted_expert_output = ( + recv_tokens.float() * recv_weights.float().unsqueeze(-1) + ).to(torch.bfloat16) + output = combine( + weighted_expert_output, + handle, + tokens_per_expert, + routing_indices, + ) + torch.cuda.synchronize() + torch.testing.assert_close(output, tokens, atol=5e-2, rtol=5e-2) + self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) + + @_eager_test_include + @_zero_copy_test_include + def test_dispatch_combine_identity_bf16(self): + """Dispatch and Combine basic ops form an identity in BF16.""" + self._run_dispatch_combine_identity(mxfp8=False) + + @_eager_test_include + @_zero_copy_test_include + @_mxfp8_align_test + def test_dispatch_combine_identity_mxfp8(self): + """Dispatch and Combine basic ops form an identity with MXFP8 transport.""" + self._run_dispatch_combine_identity(mxfp8=True) + + def _make_megamoe_model( + self, + *, + recipe, + accumulate_into_main_grad=False, + delay_wgrad_compute=False, + ): + """Build the exact five-op sequence recognized by MegaMoE fusion.""" + buffer = self._make_buffer(alignment=128 if recipe is not None else 0) + dispatch = te_ops.Dispatch(buffer) + init_ctx = ( + te.quantized_model_init(enabled=True, recipe=recipe) + if recipe is not None + else nullcontext() + ) + previous_single_param = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" + try: + with init_ctx: + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * 256, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + 256, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + finally: + if previous_single_param is None: + del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] + else: + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param + combine = te_ops.Combine() + + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) + dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) + dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) + fc1.set_extra_input_channel(0, "tokens_per_expert") + activation.set_extra_input_channel(0, "routing_weights") + fc2.set_extra_input_channel(0, "tokens_per_expert") + combine.set_extra_input_channel(0, "ep_handle") + combine.set_extra_input_channel(1, "tokens_per_expert") + combine.set_extra_input_channel(2, "routing_indices") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + return model, fc1, fc2 + + @_eager_test_include + def test_megamoe_bf16_numerics(self): + self._run_megamoe_vs_reference(quantization="bf16") + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_mxfp8_numerics(self): + self._run_megamoe_vs_reference(quantization="mxfp8") + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_main_grad_accumulation(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + accumulate_into_main_grad=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_main_grad_overwrite(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + accumulate_into_main_grad=True, + overwrite_main_grad=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_delayed_wgrad(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + delay_wgrad_compute=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_delayed_main_grad(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + accumulate_into_main_grad=True, + delay_wgrad_compute=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_prequantized_input(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + prequantized_input=True, + ) + + def _run_megamoe_vs_reference( + self, + *, + quantization, + accumulate_into_main_grad=False, + overwrite_main_grad=False, + delay_wgrad_compute=False, + prequantized_input=False, + ): + """Compare the five-op MoE sequence with the PyTorch EP reference. + + The fuser selects MegaMoE when its runtime gates pass. Otherwise this + exercises the same sequence as separate NCCL EP and grouped-MLP ops. + """ + if not EAGER: + self.skipTest("MoE reference comparison requires eager EP mode") + + recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None + model, fc1, fc2 = self._make_megamoe_model( + recipe=recipe, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(3100 + self.cfg.rank) + with torch.no_grad(): + for op in (fc1, fc2): + weights = op.weight.quantized_tensors + if weights is None: + weights = op.weight.split_into_quantized_tensors() + for expert in range(NUM_LOCAL_EXPERTS): + weight = ( + torch.randn( + weights[expert].shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + weights[expert].copy_(weight) + if quantization == "mxfp8": + self.assertIsInstance(weights[expert], MXFP8Tensor) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + seq_tokens = tokens.detach().clone().requires_grad_(True) + seq_topk_weights = topk_weights.detach().clone().requires_grad_(True) + main_grad_sentinel = 0.5 + if accumulate_into_main_grad: + for op in (fc1, fc2): + op.weight.main_grad = torch.full( + op.weight.size(), + main_grad_sentinel, + dtype=torch.float32, + device=op.weight.device, + ) + op.weight.overwrite_main_grad = overwrite_main_grad + op.weight.zero_out_wgrad = False + op.weight.grad_added_to_main_grad = False + autocast_ctx = ( + te.autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() + ) + with autocast_ctx: + model_input = ( + self._mxfp8_quantizer()(seq_tokens) if prequantized_input else seq_tokens + ) + seq_out = model(model_input, topk_idx, seq_topk_weights) + + forward_ops = model._module_groups[0]._forward_ops + fused = len(forward_ops) == 1 and isinstance(forward_ops[0][0], FusedMoeEp) + if fused: + self.assertTrue(_cudnn_megamoe_supported()) + else: + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in forward_ops)) + self.assertEqual(seq_out.dtype, torch.bfloat16) + + fc1_weight = _reference_weights(fc1) + fc2_weight = _reference_weights(fc2) + emulate_mxfp8 = fused or quantization == "mxfp8" + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=256, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + output_format=MoeFormat.BF16, + combine_format=MoeFormat.BF16, + apply_topk_in_fc1=True, + generate_c=True, + intermediate_format=MoeFormat.MXFP8 if emulate_mxfp8 else None, + backward_operand_format=MoeFormat.MXFP8 if emulate_mxfp8 and not fused else None, + backward_wgrad_mode="operands" if fused else "none", + token_padding_size=256, + ) + if emulate_mxfp8: + reference_activation = _pack_cudnn_activation( + tokens.detach(), + self._mxfp8_quantizer(), + BlockScaledTensor, + ) + else: + reference_activation = tokens.detach() + if emulate_mxfp8 and not isinstance(fc1_weight, BlockScaledTensor): + fc1_weight = quantize_blockwise(fc1_weight, MoeFormat.MXFP8, axis=1) + if emulate_mxfp8 and not isinstance(fc2_weight, BlockScaledTensor): + fc2_weight = quantize_blockwise(fc2_weight, MoeFormat.MXFP8, axis=1) + reference_outputs = reference( + reference_activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + ) + if fused: + ref_out, fc1_c, route_metadata, wgrad_forward_stash = reference_outputs + else: + ref_out, fc1_c, route_metadata = reference_outputs + wgrad_forward_stash = None + + dy = ( + torch.randn( + seq_out.shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + seq_out.backward(dy) + if delay_wgrad_compute: + fc1.backward_dw() + fc2.backward_dw() + reference_grads = reference.backward( + dy, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + fc1_c, + route_metadata, + wgrad_forward_stash=wgrad_forward_stash, + ) + if fused: + grad_tokens, grad_topk_weights, wgrad_operands = reference_grads + grad_fc1, grad_fc2 = wgrad_operands.dense_wgrads() + reference_wgrads = (grad_fc1, grad_fc2) + else: + grad_tokens, grad_topk_weights = reference_grads + reference_wgrads = (None, None) + + torch.cuda.synchronize() + tolerances = {"rtol": 0.125, "atol": 0.25} + torch.testing.assert_close(seq_out, ref_out, **tolerances) + torch.testing.assert_close( + seq_tokens.grad, + grad_tokens.to(dtype=seq_tokens.dtype), + **tolerances, + ) + torch.testing.assert_close(seq_topk_weights.grad, grad_topk_weights.float(), **tolerances) + for op, ref_grad in zip((fc1, fc2), reference_wgrads): + expected_grad = None if ref_grad is None else ref_grad.transpose(1, 2) + if accumulate_into_main_grad: + if expected_grad is not None and not overwrite_main_grad: + expected_grad = expected_grad + main_grad_sentinel + if expected_grad is not None: + torch.testing.assert_close( + op.weight.main_grad, + expected_grad.to(dtype=op.weight.main_grad.dtype), + **tolerances, + ) + else: + self.assertTrue(torch.isfinite(op.weight.main_grad).all()) + self.assertFalse( + torch.all(op.weight.main_grad == main_grad_sentinel).item() + ) + self.assertTrue(op.weight.grad_added_to_main_grad) + self.assertIsNotNone(op.weight.grad) + continue + seq_grad = op.weight.grad + self.assertEqual(seq_grad.dtype, torch.bfloat16) + self.assertEqual( + tuple(seq_grad.shape), + (NUM_LOCAL_EXPERTS, op.out_features, op.in_features), + ) + self.assertTrue(seq_grad.is_contiguous()) + self.assertTrue(torch.isfinite(seq_grad).all()) + if expected_grad is not None: + torch.testing.assert_close( + seq_grad, + expected_grad.to(dtype=seq_grad.dtype), + **tolerances, + ) + + def _init_distributed(): dist.init_process_group(backend="nccl") torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) @@ -1205,7 +1268,12 @@ def _init_distributed(): name_filter = os.environ.get("NVTE_EP_TEST_FILTER") if name_filter: loader.testMethodPrefix = name_filter - suite = loader.loadTestsFromTestCase(TestEP) + suite = unittest.TestSuite( + ( + loader.loadTestsFromTestCase(TestEP), + loader.loadTestsFromTestCase(TestMoeEpSequential), + ) + ) runner = unittest.TextTestRunner(stream=sys.stdout, verbosity=2) result = runner.run(suite) dist.barrier() diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index d38c6a80bf..56c1c6396e 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -1,4 +1,4 @@ -# Copyright (c) 2028 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT """Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. @@ -123,9 +123,7 @@ def _validate_storage(self) -> None: expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") data_shape = self.logical_shape if self.data.dtype != expected_dtype: - raise TypeError( - f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" - ) + raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") else: expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -275,6 +273,68 @@ def quantize_blockwise( MoeTensor = Union[torch.Tensor, BlockScaledTensor] +@dataclass(frozen=True) +class WgradForwardStashReference: + """Logical reference for the caller-owned forward wgrad stash. + + Unlike the production object, ``fc1_a`` bundles its logical E8M0 scales + with the E4M3 payload. It represents the padded, expert-concatenated + ``x.T`` operand after input MXFP8 staging and token-axis requantization. + """ + + fc1_a: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + +@dataclass(frozen=True) +class WgradOperandsReference: + """Logical MXFP8 operands and dense expert-weight-gradient oracle. + + The K dimension is a concatenation of local experts. Each expert's valid + routes come first, followed by zero rows up to its 256-route boundary. + Production scale tensors use a blocked physical layout; these reference + tensors keep ordinary logical scales so their represented values are easy + to inspect. + """ + + fc1_a: BlockScaledTensor + fc1_b: BlockScaledTensor + fc2_a: BlockScaledTensor + fc2_b: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + def dense_wgrads(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Return dense ``dW1=x.T@dC`` and ``dW2=h.T@(p*dY)`` per expert.""" + + a1 = self.fc1_a.dequantize() + b1 = self.fc1_b.dequantize() + a2 = self.fc2_a.dequantize() + b2 = self.fc2_b.dequantize() + expert_count = int(self.expert_offsets.numel()) + dw1 = torch.zeros( + (expert_count, a1.shape[0], b1.shape[1]), + dtype=torch.float32, + device=a1.device, + ) + dw2 = torch.zeros( + (expert_count, a2.shape[0], b2.shape[1]), + dtype=torch.float32, + device=a2.device, + ) + begin = 0 + for expert, end_tensor in enumerate(self.expert_offsets): + end = int(end_tensor.item()) + if end > begin: + dw1[expert] = a1[:, begin:end] @ b1[begin:end] + dw2[expert] = a2[:, begin:end] @ b2[begin:end] + begin = end + return dw1, dw2 + + @dataclass(frozen=True) class _DispatchPlan: """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" @@ -297,70 +357,84 @@ def _decode_tensor( name: str, expected_shape: Tuple[int, ...], quantized_axis: int, - dtype: torch.dtype, ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: - raise ValueError( - f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" - ) + raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") - return tensor.dequantize(dtype=dtype) + return tensor.dequantize() if tuple(tensor.shape) != expected_shape: raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") if not tensor.is_floating_point(): raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") - return tensor.to(dtype=dtype) + return tensor.float() -def _format_round_trip( +def _format_round_trip_axis( tensor: torch.Tensor, format: MoeFormat, *, - dtype: torch.dtype, + axis: int, ) -> torch.Tensor: if format is MoeFormat.BF16: - return tensor.to(torch.bfloat16).to(dtype) - return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=axis).dequantize() -def _qdq( +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + return _format_round_trip_axis(tensor, format, axis=-1) + + +def forward_combine_round_trip( tensor: torch.Tensor, - format: Optional[MoeFormat], - *, - axis: int, - dtype: torch.dtype, + format: MoeFormat, ) -> torch.Tensor: - """Simulate an MXFP8/NVFP4 GEMM operand in PyTorch. + """Model GLU combine conversion directly from its FP32 accumulator.""" - The MegaMoE kernel quantizes once and feeds the quantized values into the - GEMM (FP32 accumulate). PyTorch has no MXFP8 matmul, so the reference - round-trips through block-scale storage and dequantizes back to ``dtype``. - """ - if format is None or format is MoeFormat.BF16: - return tensor.to(dtype=dtype) - return quantize_blockwise(tensor, format, axis=axis).dequantize(dtype=dtype) + return _format_round_trip(tensor, format) -def _swiglu_scale_fp32( - gate: torch.Tensor, - up: torch.Tensor, - weights: torch.Tensor, - *, - apply_scale: bool, - dtype: torch.dtype, +def backward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, ) -> torch.Tensor: - """SiLU(gate)*up[*weights] in fp32, then one cast to ``dtype``. + """Model dGLU combine conversion directly from its FP32 accumulator.""" - Matches ``tex.scaled_swiglu``: promote, multiply, and store once. Stepwise - BF16 ``F.silu(g) * up * w`` extra-rounds between those muls. - """ - out = F.silu(gate.float()) * up.float() - if apply_scale: - out = out * weights.float() - return out.to(dtype=dtype) + return _format_round_trip(tensor, format) + + +def _padded_expert_rows( + rows: torch.Tensor, + expert_rows: torch.Tensor, + valid_counts: Sequence[int], + padded_ends: Sequence[int], +) -> torch.Tensor: + """Place compact expert-grouped rows at the start of padded ranges.""" + + padded_extent = int(padded_ends[-1]) if padded_ends else 0 + padded = torch.zeros( + (padded_extent, *rows.shape[1:]), + dtype=rows.dtype, + device=rows.device, + ) + begin = 0 + for expert, (count, end) in enumerate(zip(valid_counts, padded_ends)): + positions = torch.nonzero( + expert_rows == expert, + as_tuple=False, + ).flatten() + if int(positions.numel()) != int(count): + raise ValueError( + f"expert {expert} has {positions.numel()} rows, expected {count}" + ) + if count: + padded[begin : begin + count].copy_( + rows.index_select(0, positions) + ) + begin = int(end) + return padded class MoeEpReference: @@ -370,6 +444,12 @@ class MoeEpReference: ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit initialized process group for multi-rank execution; ``None`` means a one-rank reference even if the default distributed group is initialized. + + ``intermediate_format`` optionally applies a post-SwiGLU, pre-FC2 format + round trip to model fused kernels that materialize their FC2 input in low + precision. ``None`` preserves the raw mathematical reference semantics. + ``backward_operand_format`` additionally models dGLU staging of grad-output + and transposed weights along their backward reduction dimensions. """ def __init__( @@ -383,11 +463,13 @@ def __init__( max_tokens_per_rank: Optional[int] = None, output_format: Union[MoeFormat, str] = MoeFormat.BF16, combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + intermediate_format: Optional[Union[MoeFormat, str]] = None, + backward_operand_format: Optional[Union[MoeFormat, str]] = None, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, generate_c: bool = False, - compute_dtype: torch.dtype = torch.float32, - gemm_format: Optional[Union[MoeFormat, str]] = None, + backward_wgrad_mode: str = "none", + token_padding_size: int = 128, ) -> None: for name, value in ( ("num_experts", num_experts), @@ -401,24 +483,34 @@ def __init__( raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") if max_tokens_per_rank is not None and max_tokens_per_rank < 0: raise ValueError("max_tokens_per_rank must be non-negative") - if compute_dtype not in (torch.float32, torch.bfloat16): + if backward_wgrad_mode not in ("none", "operands"): + raise ValueError( + "backward_wgrad_mode must be 'none' or 'operands'" + ) + if backward_wgrad_mode == "operands" and not generate_c: + raise ValueError( + "backward_wgrad_mode='operands' requires generate_c=True" + ) + if not isinstance(token_padding_size, int) or token_padding_size <= 0: + raise ValueError("token_padding_size must be a positive integer") + if ( + backward_wgrad_mode == "operands" + and token_padding_size != 256 + ): raise ValueError( - f"compute_dtype must be torch.float32 or torch.bfloat16, got {compute_dtype}" + "backward_wgrad_mode='operands' requires " + "token_padding_size=256" ) if ep_group is None: ep_size, ep_rank = 1, 0 else: if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError( - "ep_group requires an initialized torch.distributed process group" - ) + raise RuntimeError("ep_group requires an initialized torch.distributed process group") ep_size = dist.get_world_size(ep_group) ep_rank = dist.get_rank(ep_group) if num_experts % ep_size != 0: - raise ValueError( - f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" - ) + raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") self.num_experts = num_experts self.hidden_size = hidden_size @@ -431,26 +523,24 @@ def __init__( self.max_tokens_per_rank = max_tokens_per_rank self.output_format = _parse_format(output_format) self.combine_format = _parse_format(combine_format) + self.intermediate_format = ( + None if intermediate_format is None else _parse_format(intermediate_format) + ) + self.backward_operand_format = ( + None + if backward_operand_format is None + else _parse_format(backward_operand_format) + ) self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) - self.compute_dtype = compute_dtype - self.gemm_format = None if gemm_format is None else _parse_format(gemm_format) - if self.gemm_format is MoeFormat.BF16: - self.gemm_format = None - - for name, fmt in ( - ("output_format", self.output_format), - ("combine_format", self.combine_format), - ): - required_multiple = ( - 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 - ) + self.backward_wgrad_mode = backward_wgrad_mode + self.token_padding_size = token_padding_size + + for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): + required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 if hidden_size % required_multiple != 0: - raise ValueError( - f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" - f" {name}={fmt.value}" - ) + raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") def __repr__(self) -> str: return ( @@ -458,28 +548,9 @@ def __repr__(self) -> str: f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " - f"output={self.output_format.value}, combine={self.combine_format.value}, " - f"gemm_format={None if self.gemm_format is None else self.gemm_format.value}, " - f"compute_dtype={self.compute_dtype})" + f"output={self.output_format.value}, combine={self.combine_format.value})" ) - def _gemm_dtype(self) -> torch.dtype: - """FP32 accumulate when simulating an MXFP8 GEMM; otherwise ``compute_dtype``.""" - return torch.float32 if self.gemm_format is MoeFormat.MXFP8 else self.compute_dtype - - def _stage_gemm_operand( - self, - tensor: torch.Tensor, - *, - already_quantized: bool, - axis: int, - ) -> torch.Tensor: - """Apply the kernel's pre-GEMM quantize, then dequant for a PyTorch matmul.""" - dtype = self._gemm_dtype() - if already_quantized or self.gemm_format is None: - return tensor.to(dtype=dtype) - return _qdq(tensor, self.gemm_format, axis=axis, dtype=dtype) - def _collective_device(self, device: torch.device) -> torch.device: """Device the process group can run ``all_to_all_single`` on. @@ -544,9 +615,7 @@ def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") order = torch.argsort(destination, stable=True) - send_counts_tensor = torch.bincount( - destination.index_select(0, order), minlength=self.ep_size - ).to(torch.int64) + send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) recv_counts_tensor = self._exchange_counts(send_counts_tensor) return _DispatchPlan( send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), @@ -567,7 +636,7 @@ def _run_local_experts( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: output = torch.empty( (tokens.shape[0], self.hidden_size), - dtype=self.compute_dtype, + dtype=torch.float32, device=tokens.device, ) fc1_c_rows = [] if self.generate_c else None @@ -582,45 +651,31 @@ def _run_local_experts( fc1_c_rows.append(gate_up.to(torch.bfloat16)) gate, up = gate_up.split(self.intermediate_size, dim=-1) if self.gate_up_clamp is not None: - gate = gate.float().clamp(max=self.gate_up_clamp) - up = up.float().clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + gate = gate.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up weights = route_weight.index_select(0, positions).unsqueeze(-1) - # FP32 SiLU*up[*scale], then one cast. MegaMoE then quantizes this - # intermediate for the FC2 MXFP8 GEMM; the reference QDQ mimics that. - intermediate = _swiglu_scale_fp32( - gate, - up, - weights, - apply_scale=self.apply_topk_in_fc1, - dtype=self._gemm_dtype(), - ) - if self.gemm_format is not None: - intermediate = _qdq( + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + if self.intermediate_format is not None: + intermediate = _format_round_trip( intermediate, - self.gemm_format, - axis=-1, - dtype=self._gemm_dtype(), + self.intermediate_format, ) expert_output = intermediate @ fc2_weight[expert] - if not self.apply_topk_in_fc1: - expert_output = (expert_output.float() * weights.float()).to( - dtype=self.compute_dtype - ) - expert_output = _format_round_trip( + expert_output = forward_combine_round_trip( expert_output, self.combine_format, - dtype=self.compute_dtype, ) + if not self.apply_topk_in_fc1: + # The upstream training kernel leaves scores out of dispatch + # and applies them in standalone TopkReduce after the combine + # wire-format round trip. + expert_output = expert_output * weights output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: - fc1_c = ( - torch.cat(fc1_c_rows) - if fc1_c_rows - else torch.empty( - (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device - ) - ) + fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) return output, fc1_c def __call__( @@ -630,7 +685,16 @@ def __call__( fc2_weight: MoeTensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, - ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: + ) -> Union[ + MoeTensor, + Tuple[MoeTensor, torch.Tensor, torch.Tensor], + Tuple[ + MoeTensor, + torch.Tensor, + torch.Tensor, + WgradForwardStashReference, + ], + ]: """Run dispatch, local experts, return routing, top-k reduce, and encode. Shapes: @@ -640,7 +704,9 @@ def __call__( topk_idx/topk_weights: ``(T, K)`` Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` - when constructed with ``generate_c=True``. ``fc1_c`` is the BF16 + when constructed with ``generate_c=True``. In wgrad operand mode, a + fourth :class:`WgradForwardStashReference` item is returned. ``fc1_c`` + is the BF16 pre-SwiGLU FC1 accumulator of every route this rank's experts processed, ``(local_routes, 2 * I)``, grouped by local expert and ordered within each expert by (source rank, source token-major route @@ -658,17 +724,13 @@ def __call__( if tuple(topk_idx.shape) != route_shape: raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") if tuple(topk_weights.shape) != route_shape: - raise ValueError( - f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" - ) + raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") if topk_idx.dtype not in (torch.int32, torch.int64): raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") if not topk_weights.is_floating_point(): raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError( - f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" - ) + raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") device = _tensor_device(activation) inputs = { @@ -681,57 +743,73 @@ def __call__( if input_device != device: raise ValueError(f"{name} must be on {device}, got {input_device}") - activation_float = self._stage_gemm_operand( - _decode_tensor( - activation, - name="activation", - expected_shape=(token_count, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, - ), - already_quantized=isinstance(activation, BlockScaledTensor), - axis=-1, + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, ) - fc1_float = self._stage_gemm_operand( - _decode_tensor( - fc1_weight, - name="fc1_weight", - expected_shape=( - self.experts_per_rank, - self.hidden_size, - 2 * self.intermediate_size, - ), - quantized_axis=1, - dtype=self.compute_dtype, - ), - already_quantized=isinstance(fc1_weight, BlockScaledTensor), - axis=1, + # The Rubin path first stages plain activation along H, then its + # forward column requantization forms x.T scales along routed K. + wgrad_activation_float = None + if self.backward_wgrad_mode == "operands": + wgrad_activation_float = _format_round_trip( + activation_float, + MoeFormat.MXFP8, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, ) - fc2_float = self._stage_gemm_operand( - _decode_tensor( - fc2_weight, - name="fc2_weight", - expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, - ), - already_quantized=isinstance(fc2_weight, BlockScaledTensor), - axis=1, + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, ) + if self.backward_wgrad_mode == "operands": + # Plain forward operands are staged to the same public MXFP8 + # reduction-axis representation before the Rubin GEMMs. + fc1_float = _format_round_trip_axis( + fc1_float, + MoeFormat.MXFP8, + axis=1, + ) + fc2_float = _format_round_trip_axis( + fc2_float, + MoeFormat.MXFP8, + axis=1, + ) plan = self._dispatch_plan(topk_idx, topk_weights) send_token_idx = plan.send_token_idx send_slot_idx = plan.send_slot_idx send_counts, recv_counts = plan.send_counts, plan.recv_counts - send_tokens = activation_float.index_select(0, send_token_idx) + forward_activation_float = ( + wgrad_activation_float + if wgrad_activation_float is not None + else activation_float + ) + send_tokens = forward_activation_float.index_select( + 0, + send_token_idx, + ) recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_wgrad_tokens = None + if wgrad_activation_float is not None: + recv_wgrad_tokens = self._all_to_all( + wgrad_activation_float.index_select(0, send_token_idx), + send_counts, + recv_counts, + ) recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) - recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( - dtype=self.compute_dtype - ) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) route_metadata = None + fc1_c_order = None if self.generate_c: recv_src_rank = torch.repeat_interleave( torch.arange(self.ep_size, device=device), @@ -742,11 +820,7 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = ( - torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) - .index_select(0, fc1_c_order) - .to(torch.int32) - ) + route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) # recv rows are ordered by source rank, then that source's token-major # route order, so the per-expert position grouping below realizes the # documented fc1_c ordering. @@ -759,15 +833,13 @@ def __call__( ) returned = self._all_to_all(recv_output, recv_counts, send_counts) - # Combine payload is combine_format (BF16 round-trip above); reduce in - # fp32 so top-k summation does not extra-round in compute_dtype. combine_plane = torch.zeros( (token_count * self.top_k, self.hidden_size), dtype=torch.float32, device=device, ) send_flat_slot = send_token_idx * self.top_k + send_slot_idx - combine_plane.index_copy_(0, send_flat_slot, returned.float()) + combine_plane.index_copy_(0, send_flat_slot, returned) reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) if self.output_format is MoeFormat.BF16: @@ -775,20 +847,78 @@ def __call__( else: output = quantize_blockwise(reduced, self.output_format, axis=-1) if self.generate_c: + if self.backward_wgrad_mode == "operands": + if recv_wgrad_tokens is None or fc1_c_order is None: + raise RuntimeError("wgrad forward staging was not built") + valid_counts = tuple( + int(value) + for value in torch.bincount( + recv_expert, + minlength=self.experts_per_rank, + ).cpu().tolist() + ) + padded_ends = [] + total = 0 + for count in valid_counts: + total += _ceil_div( + count, + self.token_padding_size, + ) * self.token_padding_size + padded_ends.append(total) + ordered_tokens = recv_wgrad_tokens.index_select( + 0, + fc1_c_order, + ) + metadata_experts = route_metadata[:, 0].to(torch.int64) + padded_x = _padded_expert_rows( + ordered_tokens, + metadata_experts, + valid_counts, + padded_ends, + ) + wgrad_stash = WgradForwardStashReference( + fc1_a=quantize_blockwise( + padded_x.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + expert_offsets=torch.tensor( + padded_ends, + dtype=torch.int32, + device=device, + ), + valid_route_counts=torch.tensor( + valid_counts, + dtype=torch.int32, + device=device, + ), + route_metadata=route_metadata, + ) + return output, fc1_c, route_metadata, wgrad_stash return output, fc1_c, route_metadata return output def backward( self, grad_output: torch.Tensor, - activation: MoeTensor, fc1_weight: MoeTensor, fc2_weight: MoeTensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, fc1_c: torch.Tensor, route_metadata: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + *, + wgrad_forward_stash: Optional[ + WgradForwardStashReference + ] = None, + ) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[ + torch.Tensor, + torch.Tensor, + WgradOperandsReference, + ], + ]: """Backward pass consuming the ``generate_c=True`` stash. ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, @@ -801,82 +931,94 @@ def backward( ``output_format``) are treated as straight-through identities; ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. - Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, - grad_topk_weights)``. Activation and expert weight gradients use - ``compute_dtype``. SwiGLU and router-weight scaling recompute in - fp32 and round once into ``compute_dtype`` before the FC2 GEMMs, - matching ``tex.scaled_swiglu``. The returned ``grad_topk_weights`` - remain float32. + Returns ``(grad_activation, grad_topk_weights)`` in float32. In wgrad + operand mode, a third :class:`WgradOperandsReference` item models the + caller-owned grouped-GEMM operands. """ if not self.generate_c: - raise RuntimeError( - "backward requires the operator to be constructed with generate_c=True" + raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + if self.backward_wgrad_mode == "operands": + if not isinstance( + wgrad_forward_stash, + WgradForwardStashReference, + ): + raise TypeError( + "wgrad_forward_stash must be a " + "WgradForwardStashReference" + ) + if not torch.equal( + wgrad_forward_stash.route_metadata, + route_metadata, + ): + raise ValueError( + "wgrad_forward_stash route identity does not match " + "route_metadata" + ) + elif wgrad_forward_stash is not None: + raise ValueError( + "wgrad_forward_stash is only accepted in operands mode" ) token_count = topk_idx.shape[0] if tuple(grad_output.shape) != (token_count, self.hidden_size): - raise ValueError( - f"grad_output shape must be {(token_count, self.hidden_size)}, got" - f" {tuple(grad_output.shape)}" - ) + raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") if not grad_output.is_floating_point(): raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") - device = _tensor_device(activation) + device = _tensor_device(fc1_weight) two_i = 2 * self.intermediate_size - activation_float = self._stage_gemm_operand( - _decode_tensor( - activation, - name="activation", - expected_shape=(token_count, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, - ), - already_quantized=isinstance(activation, BlockScaledTensor), - axis=-1, + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, ) - fc1_float = self._stage_gemm_operand( - _decode_tensor( - fc1_weight, - name="fc1_weight", - expected_shape=(self.experts_per_rank, self.hidden_size, two_i), - quantized_axis=1, - dtype=self.compute_dtype, - ), - already_quantized=isinstance(fc1_weight, BlockScaledTensor), - axis=1, - ) - fc2_float = self._stage_gemm_operand( - _decode_tensor( - fc2_weight, - name="fc2_weight", - expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), - quantized_axis=1, - dtype=self.compute_dtype, - ), - already_quantized=isinstance(fc2_weight, BlockScaledTensor), - axis=1, + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, ) + semantic_fc2_float = fc2_float + effective_backward_format = self.backward_operand_format + if ( + effective_backward_format is None + and self.backward_wgrad_mode == "operands" + ): + effective_backward_format = MoeFormat.MXFP8 + if effective_backward_format is not None: + # The dGLU adapter requantizes both transposed weights along the + # backward GEMM reduction dimension. + fc1_float = _format_round_trip_axis( + fc1_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + fc2_float = _format_round_trip_axis( + fc2_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): - raise ValueError( - f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" - f" {tuple(fc1_c.shape)}" - ) + raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") - # Re-dispatch the FC1 inputs, router weights, and output gradients - # along the identical forward routes. + # Re-dispatch router weights and output gradients along the identical + # forward routes. plan = self._dispatch_plan(topk_idx, topk_weights) send_counts, recv_counts = plan.send_counts, plan.recv_counts - grad_output_float = grad_output.to(dtype=self.compute_dtype) - recv_tokens = self._all_to_all( - activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts - ) - # Same FP32-on-wire / compute_dtype-for-activation cast as forward. - recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( - dtype=self.compute_dtype - ) - recv_grad = self._all_to_all( - grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + semantic_grad_output = grad_output.float() + grad_output_float = semantic_grad_output + if effective_backward_format is not None: + grad_output_float = _format_round_trip( + grad_output_float, + effective_backward_format, + ) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_semantic_grad = self._all_to_all( + semantic_grad_output.index_select(0, plan.send_token_idx), + send_counts, + recv_counts, ) # route_metadata rows are in fc1_c order; sorting them by @@ -890,96 +1032,151 @@ def backward( perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j else: perm = torch.empty((0,), dtype=torch.int64, device=device) - x_rows = torch.empty_like(recv_tokens) - x_rows.index_copy_(0, perm, recv_tokens) w_rows = torch.empty_like(recv_weight) w_rows.index_copy_(0, perm, recv_weight) dy_rows = torch.empty_like(recv_grad) dy_rows.index_copy_(0, perm, recv_grad) + semantic_dy_rows = torch.empty_like(recv_semantic_grad) + semantic_dy_rows.index_copy_(0, perm, recv_semantic_grad) - c_rows = fc1_c.to(dtype=self.compute_dtype) + c_rows = fc1_c.float() expert_rows = metadata[:, 0] - d_x_rows = torch.zeros( + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + h_rows = torch.zeros( + (local_routes, self.intermediate_size), + dtype=torch.float32, + device=device, + ) + weighted_dy_rows = torch.zeros( (local_routes, self.hidden_size), - dtype=self.compute_dtype, + dtype=torch.float32, + device=device, + ) + dc_rows = torch.zeros( + (local_routes, two_i), + dtype=torch.float32, device=device, ) - d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) - grad_fc1 = torch.zeros_like(fc1_float) - grad_fc2 = torch.zeros_like(fc2_float) for expert in range(self.experts_per_rank): positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() if positions.numel() == 0: continue c = c_rows.index_select(0, positions) - x = x_rows.index_select(0, positions) - w = w_rows.index_select(0, positions).unsqueeze(-1).float() + w = w_rows.index_select(0, positions).unsqueeze(-1) d_y = dy_rows.index_select(0, positions) + semantic_d_y = semantic_dy_rows.index_select(0, positions) gate, up = c.split(self.intermediate_size, dim=-1) if self.gate_up_clamp is not None: - g = gate.float().clamp(max=self.gate_up_clamp) - u = up.float().clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + g = gate.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) else: - g, u = gate.float(), up.float() + g, u = gate, up sig = torch.sigmoid(g) s = g * sig h = s * u + h_rows.index_copy_(0, positions, h) + weighted_dy_rows.index_copy_(0, positions, d_y * w) if self.apply_topk_in_fc1: - h_fc2 = (h * w).to(dtype=self._gemm_dtype()) - d_y_pre = d_y.to(dtype=self._gemm_dtype()) + d_y_pre = d_y else: - h_fc2 = h.to(dtype=self._gemm_dtype()) - d_y_pre = (d_y.float() * w).to(dtype=self._gemm_dtype()) - if self.gemm_format is not None: - h_fc2 = _qdq(h_fc2, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) - d_y_pre = _qdq(d_y_pre, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) - grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre - d_h_fc2 = (d_y_pre @ fc2_float[expert].transpose(0, 1)).float() + d_y_pre = d_y * w + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + semantic_d_h = ( + semantic_d_y + @ semantic_fc2_float[expert].transpose(0, 1) + ) + d_w_rows[positions] = (semantic_d_h * h).sum(dim=-1) else: d_h = d_h_fc2 - d_w_rows[positions] = (d_y.float() * (h @ fc2_float[expert].float())).sum(dim=-1) + d_w_rows[positions] = ( + semantic_d_y * (h @ semantic_fc2_float[expert]) + ).sum(dim=-1) d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s if self.gate_up_clamp is not None: - d_gate = d_g * (gate.float() <= self.gate_up_clamp) - up_f = up.float() - d_up = d_u * ((up_f >= -self.gate_up_clamp) & (up_f <= self.gate_up_clamp)) + d_gate = d_g * (gate <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) else: d_gate, d_up = d_g, d_u - d_c = torch.cat((d_gate, d_up), dim=-1).to(dtype=self._gemm_dtype()) - if self.gemm_format is not None: - d_c = _qdq(d_c, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) - grad_fc1[expert] = x.transpose(0, 1) @ d_c - d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) + d_c = torch.cat((d_gate, d_up), dim=-1) + dc_rows.index_copy_(0, positions, d_c) + if self.intermediate_format is not None: + d_c = _format_round_trip(d_c, self.intermediate_format) + d_x = d_c @ fc1_float[expert].transpose(0, 1) + d_x_rows.index_copy_( + 0, + positions, + backward_combine_round_trip(d_x, self.combine_format), + ) # Return the route gradients to their source ranks and scatter-add. returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) - grad_activation = torch.zeros( - (token_count, self.hidden_size), - dtype=torch.float32, - device=device, - ) - grad_activation.index_add_(0, plan.send_token_idx, returned_dx.float()) - grad_activation = grad_activation.to(dtype=self.compute_dtype) - grad_topk_weights = torch.zeros( - (token_count * self.top_k,), dtype=torch.float32, device=device - ) - grad_topk_weights.index_copy_( - 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw - ) - return ( - grad_activation, - grad_fc1, - grad_fc2, - grad_topk_weights.view(token_count, self.top_k), + grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) + grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + grad_topk_weights = grad_topk_weights.view( + token_count, + self.top_k, ) + if self.backward_wgrad_mode == "operands": + stash = wgrad_forward_stash + padded_ends = tuple( + int(value) + for value in stash.expert_offsets.cpu().tolist() + ) + valid_counts = tuple( + int(value) + for value in stash.valid_route_counts.cpu().tolist() + ) + padded_dc = _padded_expert_rows( + dc_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_h = _padded_expert_rows( + h_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_weighted_dy = _padded_expert_rows( + weighted_dy_rows, + expert_rows, + valid_counts, + padded_ends, + ) + operands = WgradOperandsReference( + fc1_a=stash.fc1_a, + fc1_b=quantize_blockwise( + padded_dc, + MoeFormat.MXFP8, + axis=0, + ), + fc2_a=quantize_blockwise( + padded_h.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + fc2_b=quantize_blockwise( + padded_weighted_dy, + MoeFormat.MXFP8, + axis=0, + ), + expert_offsets=stash.expert_offsets, + valid_route_counts=stash.valid_route_counts, + route_metadata=stash.route_metadata, + ) + return grad_activation, grad_topk_weights, operands + return grad_activation, grad_topk_weights __all__ = [ @@ -987,5 +1184,9 @@ def backward( "MoeEpReference", "MoeFormat", "MoeTensor", + "WgradForwardStashReference", + "WgradOperandsReference", + "backward_combine_round_trip", + "forward_combine_round_trip", "quantize_blockwise", -] +] \ No newline at end of file diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index ddb7bceefd..670da0e1da 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -106,21 +106,16 @@ def quantize_mxfp8_for_ep( raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") quantized = quantizer(input_) if quantized._with_gemm_swizzled_scales: - raise ValueError("NCCL EP requires unswizzled MXFP8 scales.") + raise ValueError("EP requires unswizzled MXFP8 scales.") data = quantized._rowwise_data scale_inv = quantized._rowwise_scale_inv if data is None or scale_inv is None: - raise ValueError("NCCL EP requires rowwise MXFP8 data and scales.") + raise ValueError("EP requires rowwise MXFP8 data and scales.") rows, hidden = input_.shape scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE - if scale_cols * scale_inv.element_size() % 16: - raise ValueError( - "MXFP8 EP transport requires hidden size divisible by " - f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {hidden}." - ) scale_inv = scale_inv[:rows, :scale_cols] if not scale_inv.is_contiguous(): - raise ValueError("NCCL EP requires compact contiguous MXFP8 scales.") + raise ValueError("EP requires compact contiguous MXFP8 scales.") return quantized, scale_inv diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index f723a87e1b..3555634904 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -116,11 +116,13 @@ def _resolve_grad_output_quantizer( if ( quantizer is not None and prev_op_grad_output_quantizer is not None - and quantizer is not prev_op_grad_output_quantizer + and type(quantizer) is not type(prev_op_grad_output_quantizer) ): - raise ValueError( - "Combine grad_output_quantizer and previous operation grad-output " - "quantizer must be the same object when both are set." + raise TypeError( + "Combine grad-output quantizer and previous operation grad-output " + "quantizer must have the same type, got " + f"{type(quantizer).__name__} and " + f"{type(prev_op_grad_output_quantizer).__name__}." ) if quantizer is None: quantizer = prev_op_grad_output_quantizer diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index abd4065290..144e0cb49f 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -286,6 +286,11 @@ def _dispatch_impl( quantized_input, input_scale_inv = quantize_mxfp8_for_ep(input_, input_quantizer) scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE + if scale_cols * input_scale_inv.element_size() % 16: + raise ValueError( + "MXFP8 NCCL EP transport requires hidden size divisible by " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {self.buffer.hidden_dim}." + ) recv_data, recv_scale_inv = _scale_alloc_io( recv_tokens, num_recv_tokens, diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index c2183a1a2e..3c1993de21 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -7,6 +7,7 @@ from __future__ import annotations from collections.abc import Iterable, Sequence +import functools from importlib.metadata import PackageNotFoundError, version as get_pkg_version from typing import Any, Optional @@ -14,14 +15,19 @@ import torch import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE from ...ep import get_ep_group from ...quantization import Recipe from ...tensor import GroupedTensor, MXFP8Quantizer, Quantizer +from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, quantize_mxfp8_for_ep, + view_main_grad_as_grouped_buffer, ) from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU from ..fuser import register_forward_backward_fusion @@ -30,11 +36,10 @@ def _cudnn_megamoe_supported() -> bool: """Whether the installed cuDNN frontend includes the public MegaMoE API.""" - return True - # try: - # return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.28.0") - # except PackageNotFoundError: - # return False + try: + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.28.0") + except PackageNotFoundError: + return False def _pack_as_cudnn_moe_tensor( @@ -55,12 +60,12 @@ def _pack_as_cudnn_moe_tensor( def _pack_cudnn_activation( - input_: torch.Tensor, + input_: torch.Tensor | MXFP8TensorStorage, quantizer: Optional[MXFP8Quantizer], block_scaled_cls: type, ): """Pack the dispatch input in the public cuDNN MoE activation layout.""" - if quantizer is None: + if quantizer is None and not isinstance(input_, MXFP8TensorStorage): return input_ quantized, scale = quantize_mxfp8_for_ep(input_, quantizer) return _pack_as_cudnn_moe_tensor( @@ -123,26 +128,99 @@ def _pack_cudnn_weights( ) -def _grouped_weight_grad(op: GroupedLinear, grad: torch.Tensor) -> list[Optional[torch.Tensor]]: - """Convert a MegaMoE ``(E, in, out)`` wgrad to one TE ``(E, out, in)`` grad.""" +def _launch_grouped_wgrad_from_operands( + layer_operands: list[torch.Tensor], + unused: None, + output: torch.Tensor | GroupedTensor, + *, + offsets: torch.Tensor, + accumulate: bool, +) -> None: + """Compute one TE-layout grouped wgrad directly from MegaMoE's operands.""" + del unused + from cudnn import grouped_gemm_wgrad_wrapper_sm100 + + a_tensor, sfa_tensor, b_tensor, sfb_tensor = layer_operands + output_data = output.rowwise_data if isinstance(output, GroupedTensor) else output + # MegaMoE exports dW in (in, out) layout. Swapping operands computes + # B.T @ A.T directly into TE's contiguous (out, in) parameter layout. + grouped_gemm_wgrad_wrapper_sm100( + a_tensor=b_tensor.transpose(0, 1), + b_tensor=a_tensor.transpose(0, 1), + sfa_tensor=sfb_tensor, + sfb_tensor=sfa_tensor, + offsets_tensor=offsets, + output_mode="dense", + wgrad_tensor=output_data, + wgrad_dtype=output_data.dtype, + acc_dtype=torch.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, + accumulate_on_output=accumulate, + ) + + +def _compute_grouped_weight_grad( + op: GroupedLinear, + operands, + prefix: str, +) -> list[Optional[torch.Tensor]]: + """Launch or defer one operand-mode wgrad with GroupedLinear semantics.""" weight = op.weight - expected_shape = (op.num_groups, op.in_features, op.out_features) - if tuple(grad.shape) != expected_shape: - raise RuntimeError( - f"MegaMoE weight gradient must have shape {expected_shape}, got {tuple(grad.shape)}" - ) if not weight.requires_grad: return [None] - # copy_ performs the float32-to-parameter-dtype conversion while writing - # directly into the contiguous layout expected by the grouped parameter. - param_grad = torch.empty( - (op.num_groups, op.out_features, op.in_features), - dtype=weight.dtype, - device=grad.device, + weight_shape = (op.out_features, op.in_features) + output_shape = (op.num_groups, *weight_shape) + accumulate = False + if op._accumulate_into_main_grad: + output_data = get_main_grad_from_param( + weight, + op_label=f"FusedMoeEp {prefix.upper()}", + ) + output_data = view_main_grad_as_grouped_buffer( + output_data, + op.num_groups, + weight_shape, + label=f"FusedMoeEp {prefix.upper()} weight", + ) + accumulate = get_accumulate_flag_in_param(weight) + else: + output_data = torch.empty( + output_shape, + dtype=weight.dtype, + device=weight.device, + ) + + layer_operands = [ + getattr(operands, f"{prefix}_a"), + getattr(operands, f"{prefix}_sfa"), + getattr(operands, f"{prefix}_b"), + getattr(operands, f"{prefix}_sfb"), + ] + launch = functools.partial( + _launch_grouped_wgrad_from_operands, + offsets=operands.expert_offsets, + accumulate=accumulate, ) - param_grad.copy_(grad.transpose(1, 2)) - return [param_grad] + delay_wgrad = op.wgrad_store is not None and op.wgrad_store.delay_wgrad_compute() + if delay_wgrad: + grouped_output = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=op.num_groups, + tensor_shape=weight_shape, + rowwise_data=output_data, + dtype=output_data.dtype, + ) + op.wgrad_store.put([layer_operands, None, grouped_output], launch) + else: + launch(layer_operands, None, output_data) + + if op._accumulate_into_main_grad: + return get_dummy_wgrads_for_params([weight]) + if delay_wgrad: + return [None] + return [output_data] def _grouped_linear_supported(op: GroupedLinear) -> bool: @@ -167,9 +245,7 @@ def _grouped_linear_supported(op: GroupedLinear) -> bool: and not op._scale_bias and op.single_grouped_weight and not op.single_grouped_bias - and not op._accumulate_into_main_grad and not op._is_distributed_weight() - and not op.wgrad_store.delay_wgrad_compute() and weight_ok ) @@ -306,6 +382,9 @@ def __init__( max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, apply_topk_in_fc1=True, generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=256, + sf_padding_size=128, combine_format="bf16", output_format="bf16", ) @@ -325,15 +404,29 @@ def fc2(self) -> GroupedLinear: def fuser_forward( self, basic_op_ctxs: list[OperationContext], - input_: torch.Tensor, + input_: torch.Tensor | MXFP8TensorStorage, *, basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: - if input_.dtype is not torch.bfloat16: - raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") + is_mxfp8 = isinstance(input_, MXFP8TensorStorage) + if is_quantized_tensor(input_) and not is_mxfp8: + raise TypeError( + "FusedMoeEp supports BF16 and MXFP8 inputs, " + f"got {type(input_).__name__}." + ) + input_dtype = input_.dtype if isinstance(input_, torch.Tensor) else input_._dtype + if input_dtype is not torch.bfloat16: + raise TypeError( + "FusedMoeEp input must be BF16 or an MXFP8 tensor representing BF16 values, " + f"got logical dtype {input_dtype}." + ) + if is_mxfp8 and input_._fp8_dtype != DType.kFloat8E4M3: + raise NotImplementedError( + f"FusedMoeEp supports E4M3 MXFP8 only, got {input_._fp8_dtype}." + ) if any(kwargs for kwargs in basic_op_kwargs): raise NotImplementedError("FusedMoeEp does not support per-operation output buffers.") @@ -348,7 +441,7 @@ def fuser_forward( ) fc1_weight = _pack_cudnn_weights(self.fc1, block_scaled_cls=self._block_scaled_cls) fc2_weight = _pack_cudnn_weights(self.fc2, block_scaled_cls=self._block_scaled_cls) - output, fc1_c, route_metadata = self._moe( + output, fc1_c, route_metadata, wgrad_forward_stash = self._moe( activation, fc1_weight, fc2_weight, @@ -357,15 +450,11 @@ def fuser_forward( ) if any(ctx.requires_grad for ctx in basic_op_ctxs): - input_data = activation if isinstance(activation, torch.Tensor) else activation.data - input_scale = None if isinstance(activation, torch.Tensor) else activation.scale fc1_data = fc1_weight if isinstance(fc1_weight, torch.Tensor) else fc1_weight.data fc1_scale = None if isinstance(fc1_weight, torch.Tensor) else fc1_weight.scale fc2_data = fc2_weight if isinstance(fc2_weight, torch.Tensor) else fc2_weight.data fc2_scale = None if isinstance(fc2_weight, torch.Tensor) else fc2_weight.scale basic_op_ctxs[0].save_for_backward( - input_data, - input_scale, fc1_data, fc1_scale, fc2_data, @@ -374,8 +463,13 @@ def fuser_forward( topk_weights, fc1_c, route_metadata, + wgrad_forward_stash.fc1_a, + wgrad_forward_stash.fc1_sfa, + wgrad_forward_stash.expert_offsets, + wgrad_forward_stash.valid_route_counts, + wgrad_forward_stash.route_metadata, ) - basic_op_ctxs[0].input_dtype = input_.dtype + basic_op_ctxs[0].input_dtype = input_dtype basic_op_ctxs[0].prev_op_grad_output_quantizer = prev_op_grad_output_quantizer # Dispatch extras are channel-bound with output_to_caller=False and are @@ -403,8 +497,6 @@ def fuser_backward( ]: del basic_op_grad_extra_outputs ( - input_data, - input_scale, fc1_data, fc1_scale, fc2_data, @@ -413,33 +505,40 @@ def fuser_backward( topk_weights, fc1_c, route_metadata, + wgrad_fc1_a, + wgrad_fc1_sfa, + wgrad_expert_offsets, + wgrad_valid_route_counts, + wgrad_route_metadata, ) = basic_op_ctxs[0].saved_tensors - input_ = _pack_as_cudnn_moe_tensor( - input_data, - input_scale, - self._block_scaled_cls, - ) fc1_weight = _pack_as_cudnn_moe_tensor(fc1_data, fc1_scale, self._block_scaled_cls) fc2_weight = _pack_as_cudnn_moe_tensor(fc2_data, fc2_scale, self._block_scaled_cls) grad_output = maybe_dequantize( grad_output, basic_op_ctxs[0].input_dtype, ) - grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._moe.backward( + from cudnn.moe_ep import MoeEpWgradForwardStash + + wgrad_forward_stash = MoeEpWgradForwardStash( + fc1_a=wgrad_fc1_a, + fc1_sfa=wgrad_fc1_sfa, + expert_offsets=wgrad_expert_offsets, + valid_route_counts=wgrad_valid_route_counts, + route_metadata=wgrad_route_metadata, + ) + grad_input, grad_topk_weights, wgrad_operands = self._moe.backward( grad_output, - input_, fc1_weight, fc2_weight, topk_idx, topk_weights, fc1_c, route_metadata, + wgrad_forward_stash=wgrad_forward_stash, ) - # MegaMoE returns float32 grads in (E, in, out); each GroupedLinear has - # one packed (E, out, in) parameter. - fc1_param_grads = _grouped_weight_grad(self.fc1, grad_fc1) - fc2_param_grads = _grouped_weight_grad(self.fc2, grad_fc2) + fc1_param_grads = _compute_grouped_weight_grad(self.fc1, wgrad_operands, "fc1") + fc2_param_grads = _compute_grouped_weight_grad(self.fc2, wgrad_operands, "fc2") grad_input = grad_input.to(dtype=basic_op_ctxs[0].input_dtype) grad_input_quantizer = basic_op_ctxs[0].prev_op_grad_output_quantizer if grad_input_quantizer is not None: From a4279c2e0456d13b2a62f029cc097fd5b589653d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 23 Aug 2026 22:32:54 +0000 Subject: [PATCH 62/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_ep.py | 14 +- transformer_engine/pytorch/ep_reference.py | 179 ++++++++++-------- .../pytorch/ops/fused/moe_ep.py | 3 +- 3 files changed, 101 insertions(+), 95 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index ce3b069bea..b78e40a42a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -913,9 +913,9 @@ def _run_dispatch_combine_identity(self, *, mxfp8): if mxfp8: recv_tokens = _degroup_mxfp8(recv_tokens) recv_weights = recv_weights[: recv_tokens.shape[0]] - weighted_expert_output = ( - recv_tokens.float() * recv_weights.float().unsqueeze(-1) - ).to(torch.bfloat16) + weighted_expert_output = (recv_tokens.float() * recv_weights.float().unsqueeze(-1)).to( + torch.bfloat16 + ) output = combine( weighted_expert_output, handle, @@ -1119,9 +1119,7 @@ def _run_megamoe_vs_reference( te.autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() ) with autocast_ctx: - model_input = ( - self._mxfp8_quantizer()(seq_tokens) if prequantized_input else seq_tokens - ) + model_input = self._mxfp8_quantizer()(seq_tokens) if prequantized_input else seq_tokens seq_out = model(model_input, topk_idx, seq_topk_weights) forward_ops = model._module_groups[0]._forward_ops @@ -1229,9 +1227,7 @@ def _run_megamoe_vs_reference( ) else: self.assertTrue(torch.isfinite(op.weight.main_grad).all()) - self.assertFalse( - torch.all(op.weight.main_grad == main_grad_sentinel).item() - ) + self.assertFalse(torch.all(op.weight.main_grad == main_grad_sentinel).item()) self.assertTrue(op.weight.grad_added_to_main_grad) self.assertIsNotNone(op.weight.grad) continue diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 56c1c6396e..981b04e1cf 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -123,7 +123,9 @@ def _validate_storage(self) -> None: expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") data_shape = self.logical_shape if self.data.dtype != expected_dtype: - raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) else: expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -360,7 +362,9 @@ def _decode_tensor( ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: - raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") + raise ValueError( + f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" + ) if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") return tensor.dequantize() @@ -426,13 +430,9 @@ def _padded_expert_rows( as_tuple=False, ).flatten() if int(positions.numel()) != int(count): - raise ValueError( - f"expert {expert} has {positions.numel()} rows, expected {count}" - ) + raise ValueError(f"expert {expert} has {positions.numel()} rows, expected {count}") if count: - padded[begin : begin + count].copy_( - rows.index_select(0, positions) - ) + padded[begin : begin + count].copy_(rows.index_select(0, positions)) begin = int(end) return padded @@ -484,33 +484,27 @@ def __init__( if max_tokens_per_rank is not None and max_tokens_per_rank < 0: raise ValueError("max_tokens_per_rank must be non-negative") if backward_wgrad_mode not in ("none", "operands"): - raise ValueError( - "backward_wgrad_mode must be 'none' or 'operands'" - ) + raise ValueError("backward_wgrad_mode must be 'none' or 'operands'") if backward_wgrad_mode == "operands" and not generate_c: - raise ValueError( - "backward_wgrad_mode='operands' requires generate_c=True" - ) + raise ValueError("backward_wgrad_mode='operands' requires generate_c=True") if not isinstance(token_padding_size, int) or token_padding_size <= 0: raise ValueError("token_padding_size must be a positive integer") - if ( - backward_wgrad_mode == "operands" - and token_padding_size != 256 - ): - raise ValueError( - "backward_wgrad_mode='operands' requires " - "token_padding_size=256" - ) + if backward_wgrad_mode == "operands" and token_padding_size != 256: + raise ValueError("backward_wgrad_mode='operands' requires token_padding_size=256") if ep_group is None: ep_size, ep_rank = 1, 0 else: if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("ep_group requires an initialized torch.distributed process group") + raise RuntimeError( + "ep_group requires an initialized torch.distributed process group" + ) ep_size = dist.get_world_size(ep_group) ep_rank = dist.get_rank(ep_group) if num_experts % ep_size != 0: - raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + raise ValueError( + f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" + ) self.num_experts = num_experts self.hidden_size = hidden_size @@ -527,9 +521,7 @@ def __init__( None if intermediate_format is None else _parse_format(intermediate_format) ) self.backward_operand_format = ( - None - if backward_operand_format is None - else _parse_format(backward_operand_format) + None if backward_operand_format is None else _parse_format(backward_operand_format) ) self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) @@ -537,10 +529,18 @@ def __init__( self.backward_wgrad_mode = backward_wgrad_mode self.token_padding_size = token_padding_size - for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): - required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = ( + 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + ) if hidden_size % required_multiple != 0: - raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" + f" {name}={fmt.value}" + ) def __repr__(self) -> str: return ( @@ -615,7 +615,9 @@ def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") order = torch.argsort(destination, stable=True) - send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) recv_counts_tensor = self._exchange_counts(send_counts_tensor) return _DispatchPlan( send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), @@ -675,7 +677,13 @@ def _run_local_experts( output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: - fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) + fc1_c = ( + torch.cat(fc1_c_rows) + if fc1_c_rows + else torch.empty( + (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device + ) + ) return output, fc1_c def __call__( @@ -724,13 +732,17 @@ def __call__( if tuple(topk_idx.shape) != route_shape: raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") if tuple(topk_weights.shape) != route_shape: - raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) if topk_idx.dtype not in (torch.int32, torch.int64): raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") if not topk_weights.is_floating_point(): raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) device = _tensor_device(activation) inputs = { @@ -788,9 +800,7 @@ def __call__( send_slot_idx = plan.send_slot_idx send_counts, recv_counts = plan.send_counts, plan.recv_counts forward_activation_float = ( - wgrad_activation_float - if wgrad_activation_float is not None - else activation_float + wgrad_activation_float if wgrad_activation_float is not None else activation_float ) send_tokens = forward_activation_float.index_select( 0, @@ -820,7 +830,11 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) # recv rows are ordered by source rank, then that source's token-major # route order, so the per-expert position grouping below realizes the # documented fc1_c ordering. @@ -855,15 +869,20 @@ def __call__( for value in torch.bincount( recv_expert, minlength=self.experts_per_rank, - ).cpu().tolist() + ) + .cpu() + .tolist() ) padded_ends = [] total = 0 for count in valid_counts: - total += _ceil_div( - count, - self.token_padding_size, - ) * self.token_padding_size + total += ( + _ceil_div( + count, + self.token_padding_size, + ) + * self.token_padding_size + ) padded_ends.append(total) ordered_tokens = recv_wgrad_tokens.index_select( 0, @@ -908,9 +927,7 @@ def backward( fc1_c: torch.Tensor, route_metadata: torch.Tensor, *, - wgrad_forward_stash: Optional[ - WgradForwardStashReference - ] = None, + wgrad_forward_stash: Optional[WgradForwardStashReference] = None, ) -> Union[ Tuple[torch.Tensor, torch.Tensor], Tuple[ @@ -937,31 +954,28 @@ def backward( """ if not self.generate_c: - raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) if self.backward_wgrad_mode == "operands": if not isinstance( wgrad_forward_stash, WgradForwardStashReference, ): - raise TypeError( - "wgrad_forward_stash must be a " - "WgradForwardStashReference" - ) + raise TypeError("wgrad_forward_stash must be a WgradForwardStashReference") if not torch.equal( wgrad_forward_stash.route_metadata, route_metadata, ): - raise ValueError( - "wgrad_forward_stash route identity does not match " - "route_metadata" - ) + raise ValueError("wgrad_forward_stash route identity does not match route_metadata") elif wgrad_forward_stash is not None: - raise ValueError( - "wgrad_forward_stash is only accepted in operands mode" - ) + raise ValueError("wgrad_forward_stash is only accepted in operands mode") token_count = topk_idx.shape[0] if tuple(grad_output.shape) != (token_count, self.hidden_size): - raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") + raise ValueError( + f"grad_output shape must be {(token_count, self.hidden_size)}, got" + f" {tuple(grad_output.shape)}" + ) if not grad_output.is_floating_point(): raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") @@ -981,10 +995,7 @@ def backward( ) semantic_fc2_float = fc2_float effective_backward_format = self.backward_operand_format - if ( - effective_backward_format is None - and self.backward_wgrad_mode == "operands" - ): + if effective_backward_format is None and self.backward_wgrad_mode == "operands": effective_backward_format = MoeFormat.MXFP8 if effective_backward_format is not None: # The dGLU adapter requantizes both transposed weights along the @@ -1000,7 +1011,10 @@ def backward( axis=1, ).transpose(1, 2) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): - raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") + raise ValueError( + f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" + f" {tuple(fc1_c.shape)}" + ) # Re-dispatch router weights and output gradients along the identical # forward routes. @@ -1014,7 +1028,9 @@ def backward( effective_backward_format, ) recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) - recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_grad = self._all_to_all( + grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) recv_semantic_grad = self._all_to_all( semantic_grad_output.index_select(0, plan.send_token_idx), send_counts, @@ -1086,16 +1102,11 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - semantic_d_h = ( - semantic_d_y - @ semantic_fc2_float[expert].transpose(0, 1) - ) + semantic_d_h = semantic_d_y @ semantic_fc2_float[expert].transpose(0, 1) d_w_rows[positions] = (semantic_d_h * h).sum(dim=-1) else: d_h = d_h_fc2 - d_w_rows[positions] = ( - semantic_d_y * (h @ semantic_fc2_float[expert]) - ).sum(dim=-1) + d_w_rows[positions] = (semantic_d_y * (h @ semantic_fc2_float[expert])).sum(dim=-1) d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s @@ -1118,24 +1129,24 @@ def backward( # Return the route gradients to their source ranks and scatter-add. returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) - grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation = torch.zeros( + (token_count, self.hidden_size), dtype=torch.float32, device=device + ) grad_activation.index_add_(0, plan.send_token_idx, returned_dx) - grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) - grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + grad_topk_weights = torch.zeros( + (token_count * self.top_k,), dtype=torch.float32, device=device + ) + grad_topk_weights.index_copy_( + 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw + ) grad_topk_weights = grad_topk_weights.view( token_count, self.top_k, ) if self.backward_wgrad_mode == "operands": stash = wgrad_forward_stash - padded_ends = tuple( - int(value) - for value in stash.expert_offsets.cpu().tolist() - ) - valid_counts = tuple( - int(value) - for value in stash.valid_route_counts.cpu().tolist() - ) + padded_ends = tuple(int(value) for value in stash.expert_offsets.cpu().tolist()) + valid_counts = tuple(int(value) for value in stash.valid_route_counts.cpu().tolist()) padded_dc = _padded_expert_rows( dc_rows, expert_rows, @@ -1189,4 +1200,4 @@ def backward( "backward_combine_round_trip", "forward_combine_round_trip", "quantize_blockwise", -] \ No newline at end of file +] diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 3c1993de21..0ae9884984 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -414,8 +414,7 @@ def fuser_forward( is_mxfp8 = isinstance(input_, MXFP8TensorStorage) if is_quantized_tensor(input_) and not is_mxfp8: raise TypeError( - "FusedMoeEp supports BF16 and MXFP8 inputs, " - f"got {type(input_).__name__}." + f"FusedMoeEp supports BF16 and MXFP8 inputs, got {type(input_).__name__}." ) input_dtype = input_.dtype if isinstance(input_, torch.Tensor) else input_._dtype if input_dtype is not torch.bfloat16: From ffb4d621dd7697f8eed8bace8621eff19b33132a Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 24 Aug 2026 05:26:08 +0000 Subject: [PATCH 63/83] some fix Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 3 - transformer_engine/pytorch/ops/_common.py | 48 ++- .../pytorch/ops/basic/activation.py | 2 +- .../pytorch/ops/basic/combine.py | 260 ++++++--------- .../pytorch/ops/basic/dispatch.py | 302 ++++++------------ .../pytorch/ops/basic/swiglu.py | 4 +- .../pytorch/ops/fused/moe_ep.py | 4 +- 7 files changed, 249 insertions(+), 374 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index b78e40a42a..f0695c112a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -1066,9 +1066,6 @@ def _run_megamoe_vs_reference( The fuser selects MegaMoE when its runtime gates pass. Otherwise this exercises the same sequence as separate NCCL EP and grouped-MLP ops. """ - if not EAGER: - self.skipTest("MoE reference comparison requires eager EP mode") - recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None model, fc1, fc2 = self._make_megamoe_model( recipe=recipe, diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 670da0e1da..b7341ac462 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,7 +11,7 @@ import torch from transformer_engine_torch import FP8TensorMeta -from ..constants import MXFP8_BLOCK_SCALING_SIZE +from ..constants import DType, MXFP8_BLOCK_SCALING_SIZE from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..quantized_tensor import QuantizedTensorStorage, Quantizer @@ -46,6 +46,30 @@ def get_fused_normalization_quantizer( return None +def validate_buffer( + name: str, + buffer: Optional[torch.Tensor], + *, + shape: Optional[tuple[int, ...] | list[int]] = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + contiguous: Optional[bool] = None, +) -> Optional[torch.Tensor]: + """Validate requested buffer properties and return detached storage.""" + if buffer is None: + return None + if shape is not None and tuple(buffer.shape) != tuple(shape): + raise ValueError(f"{name} shape {tuple(buffer.shape)} does not match {tuple(shape)}.") + if dtype is not None and buffer.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {buffer.dtype}.") + if device is not None and buffer.device != device: + raise ValueError(f"{name} must be on {device}, got {buffer.device}.") + if contiguous is not None and buffer.is_contiguous() != contiguous: + requirement = "contiguous" if contiguous else "non-contiguous" + raise ValueError(f"{name} must be {requirement}.") + return buffer.detach() + + def validate_or_alloc_output( buffer: Optional[torch.Tensor], shape: tuple[int, ...] | list[int], @@ -94,17 +118,31 @@ def maybe_dequantize( return tensor -def quantize_mxfp8_for_ep( - input_: torch.Tensor | MXFP8TensorStorage, - quantizer: Optional[MXFP8Quantizer], +def quantize_for_ep( + input_: torch.Tensor | QuantizedTensorStorage, + quantizer: Optional[Quantizer], ) -> tuple[MXFP8Tensor | MXFP8TensorStorage, torch.Tensor]: """Return an MXFP8 input and its compact rowwise scales for EP.""" - if isinstance(input_, (MXFP8Tensor, MXFP8TensorStorage)): + if quantizer is not None: + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + f"EP MXFP8 transport requires MXFP8Quantizer, got {type(quantizer).__name__}." + ) + if quantizer.dtype != DType.kFloat8E4M3: + raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") + + if isinstance(input_, MXFP8TensorStorage): quantized = input_ + elif isinstance(input_, QuantizedTensorStorage): + raise TypeError( + f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}." + ) else: if quantizer is None: raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") quantized = quantizer(input_) + if quantized._fp8_dtype != DType.kFloat8E4M3: + raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") if quantized._with_gemm_swizzled_scales: raise ValueError("EP requires unswizzled MXFP8 scales.") data = quantized._rowwise_data diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a974d41ef9..ea8d29e0f7 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -418,7 +418,7 @@ def fuser_forward( dtype = extra_input.dtype x = maybe_dequantize(input_.contiguous(), dtype) - scales = maybe_dequantize(extra_input, dtype) + scales = extra_input y = self._scaled_unary_forward(x, scales) ctx = basic_op_ctxs[0] diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 3555634904..9c0c6d047d 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -11,7 +11,7 @@ import torch import transformer_engine_torch as tex -from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import ( _alloc_io, _make_grouped_mxfp8, @@ -20,32 +20,16 @@ ) from ...quantization import QuantizerRole, Recipe from ...tensor import MXFP8Quantizer, Quantizer -from .._common import is_quantized_tensor, maybe_dequantize, quantize_mxfp8_for_ep +from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from .._common import ( + is_quantized_tensor, + maybe_dequantize, + quantize_for_ep, + validate_buffer, +) from ..op import BasicOperation, OperationContext -def _validate_grad_buffer( - tensor: Optional[torch.Tensor], - *, - shape: tuple[int, ...], - dtype: torch.dtype, - device: torch.device, -) -> Optional[torch.Tensor]: - if tensor is None: - return None - if tuple(tensor.shape) != shape: - raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") - if tensor.dtype is not dtype: - raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") - if tensor.device != device: - raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") - if not tensor.is_contiguous(): - raise ValueError("grad_out must be contiguous.") - if tensor.requires_grad: - raise ValueError("grad_out must not require gradients.") - return tensor - - def _validate_combine_inputs( input_: torch.Tensor, handle_mem: torch.Tensor, @@ -61,12 +45,6 @@ def _validate_combine_inputs( raise ValueError("Combine routing handle must be a uint8 tensor on the input device.") if tokens_per_expert.dtype is not torch.int64 or tokens_per_expert.device != input_.device: raise ValueError("Combine tokens_per_expert must be an int64 tensor on the input device.") - if topk_idx.ndim != 2: - raise ValueError(f"Combine routing indices must be 2D, got shape {tuple(topk_idx.shape)}.") - if topk_idx.dtype not in (torch.int32, torch.int64) or topk_idx.device != input_.device: - raise ValueError( - "Combine routing indices must be an int32 or int64 tensor on the input device." - ) return tuple(input_.shape), topk_idx.shape[0] @@ -101,78 +79,18 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: if quantizer is not None: quantizer.set_usage(rowwise=True, columnwise=False) quantizer.optimize_for_gemm = False - - def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: - super().reset_recipe_state(recipe=recipe) - quantizer = self.get_quantizer("backward", 0) - if quantizer is not None: quantizer.internal = True - def _resolve_grad_output_quantizer( - self, - prev_op_grad_output_quantizer: Optional[Quantizer], - ) -> Optional[MXFP8Quantizer]: - quantizer = self.get_quantizer("backward", 0) - if ( - quantizer is not None - and prev_op_grad_output_quantizer is not None - and type(quantizer) is not type(prev_op_grad_output_quantizer) - ): - raise TypeError( - "Combine grad-output quantizer and previous operation grad-output " - "quantizer must have the same type, got " - f"{type(quantizer).__name__} and " - f"{type(prev_op_grad_output_quantizer).__name__}." - ) - if quantizer is None: - quantizer = prev_op_grad_output_quantizer - if quantizer is None: - return None - if not isinstance(quantizer, MXFP8Quantizer): - raise TypeError( - "NCCL EP Combine backward supports MXFP8Quantizer only, got " - f"{type(quantizer).__name__}." - ) - if quantizer.dtype != DType.kFloat8E4M3: - raise NotImplementedError("NCCL EP Combine backward supports E4M3 MXFP8 only.") - return quantizer - def op_forward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Combine uses fuser_forward") def op_backward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Combine uses fuser_backward") - @staticmethod - def _stage_expert_output(input_: torch.Tensor, *, zero_copy: bool) -> torch.Tensor: - """Copy expert output into symmetric memory when zero-copy IO is enabled.""" - if not zero_copy: - return input_ - expert_out = _alloc_io(tuple(input_.shape), input_.dtype, input_.device, True) - expert_out.copy_(input_) - return expert_out - - @staticmethod - def _combine_forward_impl( - handle_mem: torch.Tensor, - expert_out: torch.Tensor, - *, - num_local_tokens: int, - ) -> torch.Tensor: - """Allocate and populate the local-token result.""" - result = torch.empty( - num_local_tokens, - expert_out.shape[-1], - dtype=expert_out.dtype, - device=expert_out.device, - ) - torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) - return result - @staticmethod def _prepare_grad_buffer( grad_out: Optional[torch.Tensor], - grad_output_quantizer: Optional[MXFP8Quantizer], + grad_output_quantizer: Optional[Quantizer], *, input_shape: tuple[int, int], input_dtype: torch.dtype, @@ -181,75 +99,59 @@ def _prepare_grad_buffer( ) -> Optional[torch.Tensor]: """Validate caller storage for the expert-output gradient.""" if grad_output_quantizer is None: - grad_out = _validate_grad_buffer( + grad_out = validate_buffer( + "grad_out", grad_out, shape=input_shape, dtype=input_dtype, device=device, ) - elif grad_out is not None: - if grad_out.device != device: - raise ValueError(f"grad_out must be on {device}, got {grad_out.device}.") - if not grad_out.is_contiguous(): - raise ValueError("MXFP8 grad_out storage must be contiguous.") - if grad_out.requires_grad: - raise ValueError("grad_out must not require gradients.") + else: + grad_out = validate_buffer( + "MXFP8 grad_out storage", + grad_out, + device=device, + contiguous=True, + ) if zero_copy and grad_out is not None and not is_symm_backed(grad_out): raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") return grad_out @staticmethod - def _combine_backward_impl( - ctx: OperationContext, - handle_mem: torch.Tensor, - tokens_per_expert: torch.Tensor, - grad_output: torch.Tensor, - ) -> torch.Tensor: - """Route a local-token gradient in the selected transport format. - - Keep transport-specific setup in this function so future quantized - formats can be added without branching in ``fuser_backward``. - """ - quantizer = ctx.grad_output_quantizer - if quantizer is None: - grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() - grad_input = ctx.grad_out - if grad_input is None: - grad_input = _alloc_io( - ctx.input_shape, - ctx.input_dtype, - grad_output.device, - ctx.zero_copy, - ) - torch.ops.transformer_engine_ep.combine_bwd(handle_mem, grad_output, grad_input) - return grad_input + def _prepare_grad_buffers( + grad_out: Optional[torch.Tensor], + quantized_grad: Optional[MXFP8TensorStorage], + grad_scale_inv: Optional[torch.Tensor], + *, + input_shape: tuple[int, int], + input_dtype: torch.dtype, + device: torch.device, + zero_copy: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Allocate the expert-output gradient data and optional scale storage.""" + if quantized_grad is None: + if grad_out is None: + grad_out = _alloc_io(input_shape, input_dtype, device, zero_copy) + return grad_out, None - quantized_grad, grad_scale_inv = quantize_mxfp8_for_ep(grad_output, quantizer) - num_recv_tokens, hidden = ctx.input_shape + if grad_scale_inv is None: + raise RuntimeError("MXFP8 Combine gradient scales are unavailable.") + num_recv_tokens, hidden = input_shape scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE - grad_data, grad_input_scale_inv = _scale_alloc_io( - ctx.grad_out, + if scale_cols * grad_scale_inv.element_size() % 16: + raise ValueError( + "MXFP8 NCCL EP transport requires hidden size divisible by " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {hidden}." + ) + return _scale_alloc_io( + grad_out, num_recv_tokens, hidden, scale_cols, quantized_grad._rowwise_data.dtype, grad_scale_inv.dtype, - grad_output.device, - ctx.zero_copy, - ) - torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, - quantized_grad._rowwise_data.view(torch.float8_e4m3fn), - grad_data.view(torch.float8_e4m3fn), - grad_scale_inv, - grad_input_scale_inv, - ) - return _make_grouped_mxfp8( - grad_data, - grad_input_scale_inv, - tokens_per_expert, - quantized_grad._fp8_dtype, - ctx.input_dtype, + device, + zero_copy, ) def fuser_forward( @@ -262,8 +164,11 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, list[tuple[()]]]: - # Resolve backward quantization and unpack Dispatch routing metadata. - grad_output_quantizer = self._resolve_grad_output_quantizer(prev_op_grad_output_quantizer) + # Combine's transport format is selected by Combine's own grad-output quantizer. + # If the preceding op expects a different gradient quantized format, + # it is requantized in that op's backward implementation (e.g., GroupedLinear backward). + del prev_op_grad_output_quantizer + grad_output_quantizer = self.get_quantizer("backward", 0) handle_mem, tokens_per_expert, topk_idx = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] input_shape, num_local_tokens = _validate_combine_inputs( @@ -272,16 +177,19 @@ def fuser_forward( tokens_per_expert, topk_idx, ) - # Stage zero-copy input if needed, then restore local-token order. zero_copy = bool(tex.ep_get_zero_copy()) - expert_out = self._stage_expert_output(input_, zero_copy=zero_copy) - result = self._combine_forward_impl( - handle_mem, - expert_out, - num_local_tokens=num_local_tokens, + expert_out = input_ + if zero_copy: + expert_out = _alloc_io(tuple(input_.shape), input_.dtype, input_.device, True) + expert_out.copy_(input_) + result = torch.empty( + num_local_tokens, + expert_out.shape[-1], + dtype=expert_out.dtype, + device=expert_out.device, ) - + torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) # Preserve routing state and optional caller storage for backward. ctx = basic_op_ctxs[0] if ctx.requires_grad: @@ -296,7 +204,6 @@ def fuser_forward( ctx.input_shape = input_shape ctx.input_dtype = input_.dtype ctx.zero_copy = zero_copy - ctx.grad_output_quantizer = grad_output_quantizer ctx.save_for_backward(handle_mem, tokens_per_expert) # Hand off to the next op in its requested representation. @@ -318,10 +225,47 @@ def fuser_backward( del basic_op_grad_extra_outputs ctx = basic_op_ctxs[0] handle_mem, tokens_per_expert = ctx.saved_tensors - grad_input = self._combine_backward_impl( - ctx, - handle_mem, - tokens_per_expert, - grad_output, + grad_output_quantizer = self.get_quantizer("backward", 0) + grad_scale_inv = None + # Prepare grad_output (Quantize if necessary) + if grad_output_quantizer is None: + grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() + quantized_grad = None + elif isinstance(grad_output_quantizer, MXFP8Quantizer): + quantized_grad, grad_scale_inv = quantize_for_ep( + grad_output, + grad_output_quantizer, + ) + grad_output = quantized_grad + else: + raise TypeError( + "NCCL EP Combine backward supports MXFP8Quantizer only, got " + f"{type(grad_output_quantizer).__name__}." + ) + grad_input, grad_input_scale_inv = self._prepare_grad_buffers( + ctx.grad_out, + quantized_grad, + grad_scale_inv, + input_shape=ctx.input_shape, + input_dtype=ctx.input_dtype, + device=grad_output.device, + zero_copy=ctx.zero_copy, ) + if quantized_grad is None: + torch.ops.transformer_engine_ep.combine_bwd(handle_mem, grad_output, grad_input) + else: + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + quantized_grad._rowwise_data.view(torch.float8_e4m3fn), + grad_input.view(torch.float8_e4m3fn), + grad_scale_inv, + grad_input_scale_inv, + ) + grad_input = _make_grouped_mxfp8( + grad_input, + grad_input_scale_inv, + tokens_per_expert, + quantized_grad._fp8_dtype, + ctx.input_dtype, + ) return grad_input, [()], [(None, None, None)] diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 144e0cb49f..c8fadb10ed 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -10,7 +10,7 @@ import torch -from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import ( EpBuffer, _alloc_io, @@ -22,67 +22,26 @@ from ...tensor import MXFP8Quantizer, Quantizer from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .._common import ( - is_quantized_tensor, maybe_dequantize, - quantize_mxfp8_for_ep, + quantize_for_ep, + validate_buffer, ) from ..op import BasicOperation, OperationContext -def _validate_output_buffer( - name: str, - tensor: Optional[torch.Tensor], - *, - shape: tuple[int, ...], - dtype: torch.dtype, - device: torch.device, -) -> Optional[torch.Tensor]: - if tensor is None: - return None - if tuple(tensor.shape) != shape: - raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") - if tensor.dtype is not dtype: - raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") - if tensor.device != device: - raise ValueError(f"{name} must be on {device}, got {tensor.device}.") - if tensor.requires_grad: - raise ValueError(f"{name} must not require gradients.") - return tensor - - def _validate_dispatch_input( input_: torch.Tensor | MXFP8TensorStorage, buffer: EpBuffer, -) -> tuple[tuple[int, int], torch.dtype, bool]: - """Validate the local token matrix and identify its transport format.""" - is_mxfp8 = isinstance(input_, MXFP8TensorStorage) - if is_quantized_tensor(input_) and not is_mxfp8: - raise TypeError( - f"NCCL EP Dispatch supports BF16 and MXFP8 inputs, got {type(input_).__name__}." - ) - +) -> tuple[int, int]: + """Validate the local token matrix.""" input_shape = tuple(input_.shape) - input_dtype = input_.dtype if isinstance(input_, torch.Tensor) else input_._dtype - expected_hidden = buffer.hidden_dim - if len(input_shape) != 2 or input_shape[-1] != expected_hidden: + if len(input_shape) != 2 or input_shape[-1] != buffer.hidden_dim: raise ValueError( - f"Dispatch input must have shape (T, {expected_hidden}), got {input_shape}." + f"Dispatch input must have shape (T, {buffer.hidden_dim}), got {input_shape}." ) if input_.device != buffer.device: raise ValueError(f"Dispatch input must be on {buffer.device}, got {input_.device}.") - if input_dtype is not torch.bfloat16: - raise TypeError( - "Dispatch input must be BF16 or an MXFP8 tensor representing BF16 values, " - f"got logical dtype {input_dtype}." - ) - - if is_mxfp8: - if input_._fp8_dtype != DType.kFloat8E4M3: - raise NotImplementedError( - f"NCCL EP Dispatch supports E4M3 MXFP8 only, got {input_._fp8_dtype}." - ) - - return input_shape, input_dtype, is_mxfp8 + return input_shape def _validate_routing_inputs( @@ -99,26 +58,8 @@ def _validate_routing_inputs( raise ValueError(f"{name} must be on {device}, got {tensor.device}.") -def _validate_mxfp8_output_buffer( - tensor: Optional[torch.Tensor], - *, - device: torch.device, -) -> None: - """Validate properties common to packed MXFP8 caller buffers. - - ``_scale_alloc_io`` validates contiguity and byte capacity after the data - and scale sizes are known. - """ - if tensor is None: - return - if tensor.device != device: - raise ValueError(f"recv_tokens must be on {device}, got {tensor.device}.") - if tensor.requires_grad: - raise ValueError("recv_tokens must not require gradients.") - - class Dispatch(BasicOperation): - """Dispatch BF16 or MXFP8 tokens to local experts with NCCL EP. + """Dispatch floating-point or MXFP8 tokens to local experts with NCCL EP. The extra inputs are routing indices and FP32 routing weights. The extra outputs are local tokens-per-expert, received routing weights, the routing @@ -135,6 +76,7 @@ def __init__(self, buffer: EpBuffer) -> None: self.buffer = buffer def num_quantizers(self, mode: str) -> int: + # quantized dispatch_bwd/combine is not supported. return 1 if mode == "forward" else 0 def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: @@ -153,43 +95,12 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: super().pre_fuser_forward(requires_grad=requires_grad) quantizer = self.get_quantizer("forward", 0) if quantizer is not None: + # We just need data, scales for dispatch, and grouped tensor + # will be recreated after dispatch op. quantizer.set_usage(rowwise=True, columnwise=False) quantizer.optimize_for_gemm = False - - def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: - super().reset_recipe_state(recipe=recipe) - quantizer = self.get_quantizer("forward", 0) - if quantizer is not None: quantizer.internal = True - def _resolve_input_quantizer( - self, - next_op_input_quantizer: Optional[Quantizer], - ) -> Optional[MXFP8Quantizer]: - quantizer = self.get_quantizer("forward", 0) - if ( - quantizer is not None - and next_op_input_quantizer is not None - and type(quantizer) is not type(next_op_input_quantizer) - ): - raise TypeError( - "Dispatch input quantizer and next operation input quantizer " - "must have the same type, got " - f"{type(quantizer).__name__} and " - f"{type(next_op_input_quantizer).__name__}." - ) - if quantizer is None: - quantizer = next_op_input_quantizer - if quantizer is None: - return None - if not isinstance(quantizer, MXFP8Quantizer): - raise TypeError( - f"NCCL EP Dispatch supports MXFP8Quantizer only, got {type(quantizer).__name__}." - ) - if quantizer.dtype != DType.kFloat8E4M3: - raise NotImplementedError("NCCL EP Dispatch supports E4M3 MXFP8 only.") - return quantizer - def op_forward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Dispatch uses fuser_forward") @@ -204,24 +115,44 @@ def _prepare_routing(self, topk_idx: torch.Tensor) -> tuple[torch.Tensor, int]: if self.buffer.eager else self.buffer.recv_capacity_per_rank ) - if num_recv_tokens is None: - raise RuntimeError("NCCL EP dispatch receive size is unavailable.") return tokens_per_expert, int(num_recv_tokens) def _prepare_output_buffers( self, recv_tokens: Optional[torch.Tensor], recv_topk_weights: Optional[torch.Tensor], + quantized_input: Optional[MXFP8TensorStorage], + input_scale_inv: Optional[torch.Tensor], *, num_recv_tokens: int, - use_mxfp8: bool, - ) -> tuple[Optional[torch.Tensor], torch.Tensor]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: """Validate caller buffers and allocate any missing dispatch outputs.""" recv_shape = (num_recv_tokens, self.buffer.hidden_dim) - if use_mxfp8: - _validate_mxfp8_output_buffer(recv_tokens, device=self.buffer.device) + recv_scale_inv = None + if quantized_input is not None: + recv_tokens = validate_buffer( + "recv_tokens", + recv_tokens, + device=self.buffer.device, + ) + scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE + if scale_cols * input_scale_inv.element_size() % 16: + raise ValueError( + "MXFP8 NCCL EP transport requires hidden size divisible by " + f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {self.buffer.hidden_dim}." + ) + recv_tokens, recv_scale_inv = _scale_alloc_io( + recv_tokens, + num_recv_tokens, + self.buffer.hidden_dim, + scale_cols, + quantized_input._rowwise_data.dtype, + input_scale_inv.dtype, + self.buffer.device, + self.buffer.zero_copy, + ) else: - recv_tokens = _validate_output_buffer( + recv_tokens = validate_buffer( "recv_tokens", recv_tokens, shape=recv_shape, @@ -236,7 +167,7 @@ def _prepare_output_buffers( self.buffer.zero_copy, ) - recv_topk_weights = _validate_output_buffer( + recv_topk_weights = validate_buffer( "recv_topk_weights", recv_topk_weights, shape=(num_recv_tokens,), @@ -250,74 +181,7 @@ def _prepare_output_buffers( self.buffer.device, self.buffer.zero_copy, ) - return recv_tokens, recv_topk_weights - - def _dispatch_impl( - self, - input_: torch.Tensor | MXFP8TensorStorage, - input_dtype: torch.dtype, - input_quantizer: Optional[MXFP8Quantizer], - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - tokens_per_expert: torch.Tensor, - recv_tokens: Optional[torch.Tensor], - recv_topk_weights: torch.Tensor, - *, - num_recv_tokens: int, - use_mxfp8: bool, - ) -> torch.Tensor | MXFP8TensorStorage: - """Route tokens in the selected transport format. - - Keep transport-specific setup in this function so future quantized - formats can be added without branching in ``fuser_forward``. - """ - if not use_mxfp8: - if recv_tokens is None: - raise RuntimeError("BF16 dispatch receive storage was not allocated.") - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - input_, - topk_weights, - recv_tokens, - recv_topk_weights, - ) - return recv_tokens - - quantized_input, input_scale_inv = quantize_mxfp8_for_ep(input_, input_quantizer) - scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE - if scale_cols * input_scale_inv.element_size() % 16: - raise ValueError( - "MXFP8 NCCL EP transport requires hidden size divisible by " - f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {self.buffer.hidden_dim}." - ) - recv_data, recv_scale_inv = _scale_alloc_io( - recv_tokens, - num_recv_tokens, - self.buffer.hidden_dim, - scale_cols, - quantized_input._rowwise_data.dtype, - input_scale_inv.dtype, - self.buffer.device, - self.buffer.zero_copy, - ) - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - quantized_input._rowwise_data.view(torch.float8_e4m3fn), - topk_weights, - recv_data.view(torch.float8_e4m3fn), - recv_topk_weights, - input_scale_inv, - recv_scale_inv, - ) - return _make_grouped_mxfp8( - recv_data, - recv_scale_inv, - tokens_per_expert, - quantized_input._fp8_dtype, - input_dtype, - ) + return recv_tokens, recv_scale_inv, recv_topk_weights def fuser_forward( self, @@ -329,19 +193,35 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - # Resolve the transport format. A BF16 input uses MXFP8 transport when - # an input quantizer is active; an existing MXFP8 input is passed through. - input_quantizer = self._resolve_input_quantizer(next_op_input_quantizer) + + # Dispatch uses unquantized transport without an input quantizer and + # MXFP8 transport with an MXFP8 input quantizer. + input_quantizer = self.get_quantizer("forward", 0) topk_idx, topk_weights = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] - input_shape, input_dtype, input_is_mxfp8 = _validate_dispatch_input(input_, self.buffer) - use_mxfp8 = input_is_mxfp8 or input_quantizer is not None + input_shape = _validate_dispatch_input(input_, self.buffer) _validate_routing_inputs( topk_idx, topk_weights, device=self.buffer.device, ) - + # Prepare the routing handle, then size and allocate the receive outputs. + tokens_per_expert, num_recv_tokens = self._prepare_routing(topk_idx) + # Prepare the input + input_scale_inv = None + if input_quantizer is None: + # Only BF16 dispatch is supported for now. + input_ = maybe_dequantize(input_, torch.bfloat16) + quantized_input = None + elif isinstance(input_quantizer, MXFP8Quantizer): + quantized_input, input_scale_inv = quantize_for_ep(input_, input_quantizer) + input_ = quantized_input + else: + raise TypeError( + "NCCL EP Dispatch supports MXFP8Quantizer only, got " + f"{type(input_quantizer).__name__}." + ) + # Prepare the output buffers # Eager mode discovers the receive size at runtime, so persistent # caller-owned output buffers cannot be used. recv_tokens = kwargs.get("recv_tokens") @@ -351,35 +231,50 @@ def fuser_forward( "eager mode sizes dispatch outputs per step and cannot use " "caller-supplied receive buffers" ) - - # Prepare the routing handle, then size and allocate the receive outputs. - tokens_per_expert, num_recv_tokens = self._prepare_routing(topk_idx) - recv_tokens, recv_topk_weights = self._prepare_output_buffers( + recv_tokens, recv_scale_inv, recv_topk_weights = self._prepare_output_buffers( recv_tokens, recv_topk_weights, + quantized_input, + input_scale_inv, num_recv_tokens=num_recv_tokens, - use_mxfp8=use_mxfp8, ) - # Launch NCCL EP using the selected transport representation. - output = self._dispatch_impl( - input_, - input_dtype, - input_quantizer, - topk_idx, - topk_weights, - tokens_per_expert, - recv_tokens, - recv_topk_weights, - num_recv_tokens=num_recv_tokens, - use_mxfp8=use_mxfp8, - ) - + if quantized_input is None: + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + output = recv_tokens + else: + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + quantized_input._rowwise_data.view(torch.float8_e4m3fn), + topk_weights, + recv_tokens.view(torch.float8_e4m3fn), + recv_topk_weights, + input_scale_inv, + recv_scale_inv, + ) + output = _make_grouped_mxfp8( + recv_tokens, + recv_scale_inv, + tokens_per_expert, + quantized_input._fp8_dtype, + torch.bfloat16, + ) + # If next_op_input_quantizer is different from input_quantizer, + # we need to requantize the data, which is handled in grouped_linear anyway. + # We won't get any fusion benefit, so don't do it here. # Save only shape/dtype metadata and the opaque routing handle for backward. ctx = basic_op_ctxs[0] if ctx.requires_grad: ctx.input_shape = input_shape - ctx.input_dtype = input_dtype + ctx.input_dtype = torch.bfloat16 ctx.topk_weights_shape = tuple(topk_weights.shape) ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer ctx.save_for_backward(self.buffer.handle_mem) @@ -399,6 +294,7 @@ def fuser_backward( ]: ctx = basic_op_ctxs[0] (handle_mem,) = ctx.saved_tensors + # Only BF16 Dispatch_bwd is supported for now. grad_output = maybe_dequantize(grad_output, ctx.input_dtype) grad_recv_weights = basic_op_grad_extra_outputs[0][1] diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 598126e2cb..bea9b31f99 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -446,9 +446,9 @@ def fuser_forward( else: dtype = extra_input.dtype - # Make sure inputs are in correct dtype + # Prepare the activation input in the compute dtype. input_ = maybe_dequantize(input_, dtype) - scales = maybe_dequantize(extra_input, dtype) + scales = extra_input out = self._scaled_glu_forward(input_, scales) # Save state for backward pass diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 0ae9884984..abb5128252 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -26,7 +26,7 @@ get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, - quantize_mxfp8_for_ep, + quantize_for_ep, view_main_grad_as_grouped_buffer, ) from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU @@ -67,7 +67,7 @@ def _pack_cudnn_activation( """Pack the dispatch input in the public cuDNN MoE activation layout.""" if quantizer is None and not isinstance(input_, MXFP8TensorStorage): return input_ - quantized, scale = quantize_mxfp8_for_ep(input_, quantizer) + quantized, scale = quantize_for_ep(input_, quantizer) return _pack_as_cudnn_moe_tensor( quantized._rowwise_data.view(torch.float8_e4m3fn), scale.view(torch.float8_e8m0fnu), From d449da9c31855af23e107bc6286d4ee797508f06 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Aug 2026 05:29:35 +0000 Subject: [PATCH 64/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/_common.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index b7341ac462..78abfb9373 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -134,9 +134,7 @@ def quantize_for_ep( if isinstance(input_, MXFP8TensorStorage): quantized = input_ elif isinstance(input_, QuantizedTensorStorage): - raise TypeError( - f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}." - ) + raise TypeError(f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}.") else: if quantizer is None: raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") From f888050bd2a385139465ade15a7c3a5f54da7b03 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 24 Aug 2026 19:59:46 +0000 Subject: [PATCH 65/83] update fused kernel import issue and view issue in delay wgrad Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fused/moe_ep.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index abb5128252..ce94cfcb75 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -138,10 +138,12 @@ def _launch_grouped_wgrad_from_operands( ) -> None: """Compute one TE-layout grouped wgrad directly from MegaMoE's operands.""" del unused - from cudnn import grouped_gemm_wgrad_wrapper_sm100 + from cudnn.gemm.cutedsl.grouped.wgrad import grouped_gemm_wgrad_wrapper_sm100 a_tensor, sfa_tensor, b_tensor, sfb_tensor = layer_operands - output_data = output.rowwise_data if isinstance(output, GroupedTensor) else output + output_data = ( + output.rowwise_data.view(output.shape) if isinstance(output, GroupedTensor) else output + ) # MegaMoE exports dW in (in, out) layout. Swapping operands computes # B.T @ A.T directly into TE's contiguous (out, in) parameter layout. grouped_gemm_wgrad_wrapper_sm100( @@ -545,7 +547,7 @@ def fuser_backward( return ( grad_input, [(), fc1_param_grads, (), fc2_param_grads, ()], - [(None, grad_topk_weights.float()), (None,), (None,), (None,), (None, None)], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), (None, None, None)], ) From b882b350ea67ea2e4af57435b5aeda8f80b222bb Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 25 Aug 2026 23:40:59 +0000 Subject: [PATCH 66/83] refactor based on ep.py + reduce cpu overheads Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 30 ++++ transformer_engine/pytorch/ep.py | 137 +++++++++++----- transformer_engine/pytorch/ep_reference.py | 22 +-- transformer_engine/pytorch/ops/_common.py | 8 +- .../pytorch/ops/basic/combine.py | 99 +++-------- .../pytorch/ops/basic/dispatch.py | 155 ++---------------- 6 files changed, 172 insertions(+), 279 deletions(-) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index a3219113d3..357b3bcff5 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -1031,6 +1031,13 @@ def test_megamoe_bf16_numerics(self): def test_megamoe_mxfp8_numerics(self): self._run_megamoe_vs_reference(quantization="mxfp8") + @_eager_test_include + def test_megamoe_main_grad_accumulation_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + accumulate_into_main_grad=True, + ) + @_eager_test_include @_mxfp8_align_test def test_megamoe_main_grad_accumulation(self): @@ -1039,6 +1046,14 @@ def test_megamoe_main_grad_accumulation(self): accumulate_into_main_grad=True, ) + @_eager_test_include + def test_megamoe_main_grad_overwrite_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + accumulate_into_main_grad=True, + overwrite_main_grad=True, + ) + @_eager_test_include @_mxfp8_align_test def test_megamoe_main_grad_overwrite(self): @@ -1048,6 +1063,13 @@ def test_megamoe_main_grad_overwrite(self): overwrite_main_grad=True, ) + @_eager_test_include + def test_megamoe_delayed_wgrad_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + delay_wgrad_compute=True, + ) + @_eager_test_include @_mxfp8_align_test def test_megamoe_delayed_wgrad(self): @@ -1056,6 +1078,14 @@ def test_megamoe_delayed_wgrad(self): delay_wgrad_compute=True, ) + @_eager_test_include + def test_megamoe_delayed_main_grad_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + accumulate_into_main_grad=True, + delay_wgrad_compute=True, + ) + @_eager_test_include @_mxfp8_align_test def test_megamoe_delayed_main_grad(self): diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 8a7ed7fc7f..3d157dfcff 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -17,7 +17,8 @@ from .cpu_offload import mark_not_offload from .distributed import symm_mem_alloc, release_symm_mem_pool -from .quantized_tensor import QuantizedTensor +from .quantized_tensor import QuantizedTensor, QuantizedTensorStorage +from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage # Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of # common.recipe (the concrete recipe classes are imported lazily where used). @@ -195,6 +196,11 @@ def get_ep_group() -> Optional[dist.ProcessGroup]: return _EP_GROUP +def _ep_is_eager() -> bool: + """Return whether the bootstrapped EP group uses variable-size eager buffers.""" + return _EAGER + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. @@ -551,7 +557,7 @@ class _DispatchState: def _ep_prepare_and_dispatch_fwd( - tokens: torch.Tensor, + tokens: torch.Tensor | MXFP8TensorStorage, topk_weights: torch.Tensor, topk_idx: torch.Tensor, buffer: "EpBuffer", @@ -572,7 +578,7 @@ def _ep_prepare_and_dispatch_fwd( num_recv_tokens = buffer.recv_capacity_per_rank payload_dtype = buffer.payload_dtype is_scaled = tokens_scale_inv is not None - tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens + tokens_data = tokens._rowwise_data if isinstance(tokens, MXFP8TensorStorage) else tokens assert tokens_data.dim() == 2, "EP dispatch tokens must be 2D [num_tokens, hidden]" hidden = tokens_data.shape[-1] if is_scaled and tokens._fp8_dtype != tex.DType.kFloat8E4M3: @@ -650,7 +656,7 @@ def _ep_prepare_and_dispatch_fwd( recv_scale_inv, tokens_per_expert, tokens._fp8_dtype, - tokens.dtype, + tokens._dtype, ) return recv_out, recv_topk_weights, state return recv_tokens, recv_topk_weights, state @@ -700,10 +706,10 @@ class _EpPrepareAndDispatch(torch.autograd.Function): recv-count, so no Python runs between the count read and the dispatch launch; caller-supplied buffers and zero-copy are then forbidden. Otherwise the recv outputs are allocated here to the static recv capacity (caller-supplied or symm-mem-backed under zero-copy) and passed in. When - ``tokens_scale_inv`` is set (MXFP8 for now), ``tokens`` is the quantized tensor kept as the - autograd operand so grad reaches the pre-quant input, and recv is returned as a per-expert - GroupedTensor. The compute lives in ``_ep_prepare_and_dispatch_fwd`` / ``_ep_dispatch_bwd``; - this wrapper only bridges autograd context handling.""" + When MXFP8 is configured, forward quantizes ``tokens`` to lightweight storage locally while + keeping the high-precision tensor as the autograd operand. The compute lives in + ``_ep_prepare_and_dispatch_fwd`` / ``_ep_dispatch_bwd``; this wrapper only bridges autograd + context handling.""" @staticmethod def forward( # type: ignore[override] @@ -714,12 +720,20 @@ def forward( # type: ignore[override] buffer: "EpBuffer", recv_tokens: Optional[torch.Tensor] = None, recv_topk_weights: Optional[torch.Tensor] = None, - tokens_scale_inv: Optional[torch.Tensor] = None, ): """Only tokens and topk_weights are differentiable, so the non-diff buffer tensors ride on the buffer object to keep the autograd operand list short.""" + tokens_scale_inv = None + if buffer.dispatch_fwd_quant_recipe is not None: + tokens, tokens_scale_inv = _quantize_mxfp8(tokens) recv_out, recv_topk_weights, state = _ep_prepare_and_dispatch_fwd( - tokens, topk_weights, topk_idx, buffer, recv_tokens, recv_topk_weights, tokens_scale_inv + tokens, + topk_weights, + topk_idx, + buffer, + recv_tokens, + recv_topk_weights, + tokens_scale_inv, ) ctx.state = state # Detach so the long-lived buffers aren't tracked as differentiable outputs; autograd @@ -745,7 +759,6 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] None, # buffer None, # recv_tokens None, # recv_topk_weights - None, # tokens_scale_inv ) @@ -769,17 +782,23 @@ class _CombineState: def _ep_combine_fwd( expert_out: torch.Tensor, grad_out: Optional[torch.Tensor], - buffer: "EpBuffer", + *, + handle_mem: torch.Tensor, + token_counts: torch.Tensor, num_local_tokens: int, + hidden_dim: int, bwd_quant_recipe, + eager: bool, + zero_copy: bool, ): - """Run combine and return ``(result, _CombineState)``. Eager mode is not graph-capturable, so - it calls the backend op directly and skips the torch.library dispatch. No autograd; the caller - owns context handling.""" - handle_mem = buffer.handle_mem - eager = buffer.eager + """Run combine from explicit routing state and return ``(result, _CombineState)``. + + This is the shared implementation for the public ``ep_combine`` wrapper and + fusible operations. Eager mode calls the backend directly to avoid the + ``torch.library`` dispatch overhead. + """ device = expert_out.device - result = torch.empty(num_local_tokens, buffer.hidden_dim, dtype=expert_out.dtype, device=device) + result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) if eager: tex.ep_combine(handle_mem, expert_out, result) else: @@ -788,25 +807,31 @@ def _ep_combine_fwd( handle_mem=handle_mem, grad_out=grad_out, bwd_quant_recipe=bwd_quant_recipe, - token_counts=buffer.tokens_per_expert, + token_counts=token_counts, expert_out_shape=expert_out.shape, expert_out_dtype=expert_out.dtype, device=device, eager=eager, - zero_copy=buffer.zero_copy, + zero_copy=zero_copy, ) return result, state -def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): +def _ep_combine_bwd( + state: "_CombineState", + g_result: torch.Tensor, + quantized_grad: Optional[QuantizedTensorStorage] = None, + grad_scale_inv: Optional[torch.Tensor] = None, +): """Scatter the result-grad to expert positions and return the expert_out grad. High-precision sends the grad as-is; a quantized recipe (MXFP8 today) quantizes it and returns a per-expert - GroupedTensor. No autograd; the caller owns context handling.""" - if not g_result.is_contiguous(): + GroupedTensor. A fusible-op caller may instead provide its fuser-quantized grad and compact + scales. No autograd; the caller owns context handling.""" + if quantized_grad is None and not g_result.is_contiguous(): g_result = g_result.contiguous() handle_mem = state.handle_mem - if state.bwd_quant_recipe is None: + if state.bwd_quant_recipe is None and quantized_grad is None: grad_expert_out = state.grad_out if grad_expert_out is None: grad_expert_out = _alloc_io( @@ -817,16 +842,21 @@ def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): else: torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) else: - mx, g_scale_inv = _quantize_mxfp8(g_result) + if quantized_grad is None: + mx, grad_scale_inv = _quantize_mxfp8(g_result) + else: + mx = quantized_grad + if grad_scale_inv is None: + raise ValueError("MXFP8 combine backward requires compact rowwise scales.") g_data = mx._rowwise_data recv_pr, hidden = state.expert_out_shape[0], state.expert_out_shape[-1] ge_data, ge_scale_inv = _scale_alloc_io( state.grad_out, recv_pr, hidden, - g_scale_inv.shape[-1], + grad_scale_inv.shape[-1], g_data.dtype, - g_scale_inv.dtype, + grad_scale_inv.dtype, state.device, state.zero_copy, ) @@ -834,10 +864,10 @@ def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): g_data_fp8 = g_data.view(torch.float8_e4m3fn) ge_data_fp8 = ge_data.view(torch.float8_e4m3fn) if state.eager: - tex.ep_combine_bwd(handle_mem, g_data_fp8, ge_data_fp8, g_scale_inv, ge_scale_inv) + tex.ep_combine_bwd(handle_mem, g_data_fp8, ge_data_fp8, grad_scale_inv, ge_scale_inv) else: torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, g_data_fp8, ge_data_fp8, g_scale_inv, ge_scale_inv + handle_mem, g_data_fp8, ge_data_fp8, grad_scale_inv, ge_scale_inv ) grad_expert_out = _make_grouped_mxfp8( ge_data, ge_scale_inv, state.token_counts, mx._fp8_dtype, state.expert_out_dtype @@ -867,7 +897,15 @@ def forward( # type: ignore[override] """Combine fwd; stashes the backward state on ctx. When ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8.""" result, ctx.state = _ep_combine_fwd( - expert_out, grad_out, buffer, num_local_tokens, bwd_quant_recipe + expert_out, + grad_out, + handle_mem=buffer.handle_mem, + token_counts=buffer.tokens_per_expert, + num_local_tokens=num_local_tokens, + hidden_dim=buffer.hidden_dim, + bwd_quant_recipe=bwd_quant_recipe, + eager=buffer.eager, + zero_copy=buffer.zero_copy, ) return result @@ -916,17 +954,37 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso return torch.empty(*shape, dtype=dtype, device=device) +def _as_mxfp8_storage(tensor: QuantizedTensorStorage): + """Return a lightweight MXFP8 storage view without a ``torch.Tensor`` wrapper.""" + if type(tensor) is MXFP8TensorStorage: + return tensor + if not isinstance(tensor, MXFP8TensorStorage): + raise TypeError(f"Expected MXFP8 tensor storage, got {type(tensor).__name__}.") + return MXFP8TensorStorage( + rowwise_data=tensor._rowwise_data, + rowwise_scale_inv=tensor._rowwise_scale_inv, + columnwise_data=tensor._columnwise_data, + columnwise_scale_inv=tensor._columnwise_scale_inv, + fp8_dtype=tensor._fp8_dtype, + quantizer=tensor._quantizer, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + fake_dtype=tensor._dtype, + ) + + def _quantize_mxfp8(x: torch.Tensor): - """Quantize a high-precision tensor to MXFP8 and return ``(quantized_tensor, scale_inv)`` where - ``scale_inv`` is the compact ``[T, H/block]`` scale the EP backend routes. The quantized tensor - is returned so callers can keep it as the autograd operand; its ``_rowwise_data`` is the fp8 - payload and ``scale_inv.shape[-1]`` the scale-column count. EP routes and returns E4M3 data in - both directions, so quantize to E4M3 regardless of pass. Strips the GEMM scale row padding to - the compact ``[T, H/block]`` layout; requires a 16-byte-aligned scale row.""" + """Quantize to lightweight MXFP8 storage and return it with compact transport scales. + + ``scale_inv`` has shape ``[T, H/block]``. EP routes and returns E4M3 data in both directions, + so quantize to E4M3 regardless of pass. The GEMM scale-row padding is stripped and each compact + scale row must remain 16-byte aligned. + """ from .constants import MXFP8_BLOCK_SCALING_SIZE from .tensor.mxfp8_tensor import MXFP8Quantizer - mx = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False).quantize(x) + quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + quantizer.internal = True + mx = quantizer.quantize(x) if mx._with_gemm_swizzled_scales: raise RuntimeError( "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact." @@ -1055,9 +1113,6 @@ def ep_dispatch( "and cannot use caller-supplied recv_tokens / recv_topk_weights" ) - # Quantize up front (before prepare) so the quant kernels overlap the eager count sync and the - # quantized tensor stays the autograd operand; grad reaches the pre-quant input. - tokens_scale_inv = None if buffer.dispatch_fwd_quant_recipe is not None: from ..common.recipe import MXFP8BlockScaling @@ -1066,7 +1121,6 @@ def ep_dispatch( "EP block-scaled dispatch supports MXFP8BlockScaling only; got " f"{type(buffer.dispatch_fwd_quant_recipe).__name__}." ) - tokens, tokens_scale_inv = _quantize_mxfp8(tokens) # Fused prepare + dispatch in one C++ op. Eager sizes the recv outputs from the per-step host # recv-count (allocated in C++, no caller buffers); non-eager sizes them to the static recv @@ -1078,7 +1132,6 @@ def ep_dispatch( buffer, recv_tokens, recv_topk_weights, - tokens_scale_inv, ) return recv_out, recv_topk_weights, buffer.tokens_per_expert diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 981b04e1cf..0a02d33b9f 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -123,9 +123,7 @@ def _validate_storage(self) -> None: expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") data_shape = self.logical_shape if self.data.dtype != expected_dtype: - raise TypeError( - f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" - ) + raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") else: expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -615,9 +613,7 @@ def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") order = torch.argsort(destination, stable=True) - send_counts_tensor = torch.bincount( - destination.index_select(0, order), minlength=self.ep_size - ).to(torch.int64) + send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) recv_counts_tensor = self._exchange_counts(send_counts_tensor) return _DispatchPlan( send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), @@ -732,17 +728,13 @@ def __call__( if tuple(topk_idx.shape) != route_shape: raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") if tuple(topk_weights.shape) != route_shape: - raise ValueError( - f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" - ) + raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") if topk_idx.dtype not in (torch.int32, torch.int64): raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") if not topk_weights.is_floating_point(): raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError( - f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" - ) + raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") device = _tensor_device(activation) inputs = { @@ -830,11 +822,7 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = ( - torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) - .index_select(0, fc1_c_order) - .to(torch.int32) - ) + route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) # recv rows are ordered by source rank, then that source's token-major # route order, so the per-expert position grouping below realizes the # documented fc1_c ordering. diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 78abfb9373..8eae6d7cf7 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -12,6 +12,7 @@ from transformer_engine_torch import FP8TensorMeta from ..constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ..ep import _as_mxfp8_storage from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..quantized_tensor import QuantizedTensorStorage, Quantizer @@ -23,7 +24,6 @@ NVFP4Quantizer, ) from ..tensor.float8_tensor import Float8Tensor -from ..tensor.mxfp8_tensor import MXFP8Tensor from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..utils import canonicalize_dtype @@ -121,7 +121,7 @@ def maybe_dequantize( def quantize_for_ep( input_: torch.Tensor | QuantizedTensorStorage, quantizer: Optional[Quantizer], -) -> tuple[MXFP8Tensor | MXFP8TensorStorage, torch.Tensor]: +) -> tuple[MXFP8TensorStorage, torch.Tensor]: """Return an MXFP8 input and its compact rowwise scales for EP.""" if quantizer is not None: if not isinstance(quantizer, MXFP8Quantizer): @@ -138,7 +138,11 @@ def quantize_for_ep( else: if quantizer is None: raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") + if not quantizer.internal: + quantizer = quantizer.copy() + quantizer.internal = True quantized = quantizer(input_) + quantized = _as_mxfp8_storage(quantized) if quantized._fp8_dtype != DType.kFloat8E4M3: raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") if quantized._with_gemm_swizzled_scales: diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 9c0c6d047d..b08d302dff 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -11,16 +11,15 @@ import torch import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import ( _alloc_io, - _make_grouped_mxfp8, - _scale_alloc_io, + _ep_combine_bwd, + _ep_combine_fwd, + _ep_is_eager, is_symm_backed, ) -from ...quantization import QuantizerRole, Recipe +from ...quantization import QuantizerRole from ...tensor import MXFP8Quantizer, Quantizer -from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .._common import ( is_quantized_tensor, maybe_dequantize, @@ -117,43 +116,6 @@ def _prepare_grad_buffer( raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") return grad_out - @staticmethod - def _prepare_grad_buffers( - grad_out: Optional[torch.Tensor], - quantized_grad: Optional[MXFP8TensorStorage], - grad_scale_inv: Optional[torch.Tensor], - *, - input_shape: tuple[int, int], - input_dtype: torch.dtype, - device: torch.device, - zero_copy: bool, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Allocate the expert-output gradient data and optional scale storage.""" - if quantized_grad is None: - if grad_out is None: - grad_out = _alloc_io(input_shape, input_dtype, device, zero_copy) - return grad_out, None - - if grad_scale_inv is None: - raise RuntimeError("MXFP8 Combine gradient scales are unavailable.") - num_recv_tokens, hidden = input_shape - scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE - if scale_cols * grad_scale_inv.element_size() % 16: - raise ValueError( - "MXFP8 NCCL EP transport requires hidden size divisible by " - f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {hidden}." - ) - return _scale_alloc_io( - grad_out, - num_recv_tokens, - hidden, - scale_cols, - quantized_grad._rowwise_data.dtype, - grad_scale_inv.dtype, - device, - zero_copy, - ) - def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -183,17 +145,11 @@ def fuser_forward( if zero_copy: expert_out = _alloc_io(tuple(input_.shape), input_.dtype, input_.device, True) expert_out.copy_(input_) - result = torch.empty( - num_local_tokens, - expert_out.shape[-1], - dtype=expert_out.dtype, - device=expert_out.device, - ) - torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result) # Preserve routing state and optional caller storage for backward. ctx = basic_op_ctxs[0] + grad_out = None if ctx.requires_grad: - ctx.grad_out = self._prepare_grad_buffer( + grad_out = self._prepare_grad_buffer( kwargs.get("grad_out"), grad_output_quantizer, input_shape=input_shape, @@ -201,10 +157,20 @@ def fuser_forward( device=input_.device, zero_copy=zero_copy, ) - ctx.input_shape = input_shape ctx.input_dtype = input_.dtype - ctx.zero_copy = zero_copy - ctx.save_for_backward(handle_mem, tokens_per_expert) + result, combine_state = _ep_combine_fwd( + expert_out, + grad_out, + handle_mem=handle_mem, + token_counts=tokens_per_expert, + num_local_tokens=num_local_tokens, + hidden_dim=expert_out.shape[-1], + bwd_quant_recipe=grad_output_quantizer, + eager=_ep_is_eager(), + zero_copy=zero_copy, + ) + if ctx.requires_grad: + ctx.combine_state = combine_state # Hand off to the next op in its requested representation. if next_op_input_quantizer is not None and not is_quantized_tensor(result): @@ -224,7 +190,6 @@ def fuser_backward( ]: del basic_op_grad_extra_outputs ctx = basic_op_ctxs[0] - handle_mem, tokens_per_expert = ctx.saved_tensors grad_output_quantizer = self.get_quantizer("backward", 0) grad_scale_inv = None # Prepare grad_output (Quantize if necessary) @@ -242,30 +207,10 @@ def fuser_backward( "NCCL EP Combine backward supports MXFP8Quantizer only, got " f"{type(grad_output_quantizer).__name__}." ) - grad_input, grad_input_scale_inv = self._prepare_grad_buffers( - ctx.grad_out, + grad_input = _ep_combine_bwd( + ctx.combine_state, + grad_output, quantized_grad, grad_scale_inv, - input_shape=ctx.input_shape, - input_dtype=ctx.input_dtype, - device=grad_output.device, - zero_copy=ctx.zero_copy, ) - if quantized_grad is None: - torch.ops.transformer_engine_ep.combine_bwd(handle_mem, grad_output, grad_input) - else: - torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, - quantized_grad._rowwise_data.view(torch.float8_e4m3fn), - grad_input.view(torch.float8_e4m3fn), - grad_scale_inv, - grad_input_scale_inv, - ) - grad_input = _make_grouped_mxfp8( - grad_input, - grad_input_scale_inv, - tokens_per_expert, - quantized_grad._fp8_dtype, - ctx.input_dtype, - ) return grad_input, [()], [(None, None, None)] diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index c8fadb10ed..43a96a63d1 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -10,21 +10,17 @@ import torch -from ...constants import MXFP8_BLOCK_SCALING_SIZE from ...ep import ( EpBuffer, - _alloc_io, - _make_grouped_mxfp8, - _scale_alloc_io, - ep_prepare, + _ep_dispatch_bwd, + _ep_prepare_and_dispatch_fwd, ) -from ...quantization import QuantizerRole, Recipe +from ...quantization import QuantizerRole from ...tensor import MXFP8Quantizer, Quantizer from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .._common import ( maybe_dequantize, quantize_for_ep, - validate_buffer, ) from ..op import BasicOperation, OperationContext @@ -107,82 +103,6 @@ def op_forward(self, *args: Any, **kwargs: Any) -> None: def op_backward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("Dispatch uses fuser_backward") - def _prepare_routing(self, topk_idx: torch.Tensor) -> tuple[torch.Tensor, int]: - """Prepare NCCL routing state and return the receive row count.""" - tokens_per_expert = ep_prepare(self.buffer, topk_idx) - num_recv_tokens = ( - self.buffer._host_total_recv_tokens - if self.buffer.eager - else self.buffer.recv_capacity_per_rank - ) - return tokens_per_expert, int(num_recv_tokens) - - def _prepare_output_buffers( - self, - recv_tokens: Optional[torch.Tensor], - recv_topk_weights: Optional[torch.Tensor], - quantized_input: Optional[MXFP8TensorStorage], - input_scale_inv: Optional[torch.Tensor], - *, - num_recv_tokens: int, - ) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: - """Validate caller buffers and allocate any missing dispatch outputs.""" - recv_shape = (num_recv_tokens, self.buffer.hidden_dim) - recv_scale_inv = None - if quantized_input is not None: - recv_tokens = validate_buffer( - "recv_tokens", - recv_tokens, - device=self.buffer.device, - ) - scale_cols = self.buffer.hidden_dim // MXFP8_BLOCK_SCALING_SIZE - if scale_cols * input_scale_inv.element_size() % 16: - raise ValueError( - "MXFP8 NCCL EP transport requires hidden size divisible by " - f"{16 * MXFP8_BLOCK_SCALING_SIZE}, got {self.buffer.hidden_dim}." - ) - recv_tokens, recv_scale_inv = _scale_alloc_io( - recv_tokens, - num_recv_tokens, - self.buffer.hidden_dim, - scale_cols, - quantized_input._rowwise_data.dtype, - input_scale_inv.dtype, - self.buffer.device, - self.buffer.zero_copy, - ) - else: - recv_tokens = validate_buffer( - "recv_tokens", - recv_tokens, - shape=recv_shape, - dtype=self.buffer.payload_dtype, - device=self.buffer.device, - ) - if recv_tokens is None: - recv_tokens = _alloc_io( - recv_shape, - self.buffer.payload_dtype, - self.buffer.device, - self.buffer.zero_copy, - ) - - recv_topk_weights = validate_buffer( - "recv_topk_weights", - recv_topk_weights, - shape=(num_recv_tokens,), - dtype=torch.float32, - device=self.buffer.device, - ) - if recv_topk_weights is None: - recv_topk_weights = _alloc_io( - (num_recv_tokens,), - torch.float32, - self.buffer.device, - self.buffer.zero_copy, - ) - return recv_tokens, recv_scale_inv, recv_topk_weights - def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -199,29 +119,24 @@ def fuser_forward( input_quantizer = self.get_quantizer("forward", 0) topk_idx, topk_weights = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] - input_shape = _validate_dispatch_input(input_, self.buffer) + _validate_dispatch_input(input_, self.buffer) _validate_routing_inputs( topk_idx, topk_weights, device=self.buffer.device, ) - # Prepare the routing handle, then size and allocate the receive outputs. - tokens_per_expert, num_recv_tokens = self._prepare_routing(topk_idx) # Prepare the input input_scale_inv = None if input_quantizer is None: # Only BF16 dispatch is supported for now. input_ = maybe_dequantize(input_, torch.bfloat16) - quantized_input = None elif isinstance(input_quantizer, MXFP8Quantizer): - quantized_input, input_scale_inv = quantize_for_ep(input_, input_quantizer) - input_ = quantized_input + input_, input_scale_inv = quantize_for_ep(input_, input_quantizer) else: raise TypeError( "NCCL EP Dispatch supports MXFP8Quantizer only, got " f"{type(input_quantizer).__name__}." ) - # Prepare the output buffers # Eager mode discovers the receive size at runtime, so persistent # caller-owned output buffers cannot be used. recv_tokens = kwargs.get("recv_tokens") @@ -231,53 +146,24 @@ def fuser_forward( "eager mode sizes dispatch outputs per step and cannot use " "caller-supplied receive buffers" ) - recv_tokens, recv_scale_inv, recv_topk_weights = self._prepare_output_buffers( + output, recv_topk_weights, dispatch_state = _ep_prepare_and_dispatch_fwd( + input_, + topk_weights, + topk_idx, + self.buffer, recv_tokens, recv_topk_weights, - quantized_input, input_scale_inv, - num_recv_tokens=num_recv_tokens, ) - # Launch NCCL EP using the selected transport representation. - if quantized_input is None: - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - input_, - topk_weights, - recv_tokens, - recv_topk_weights, - ) - output = recv_tokens - else: - torch.ops.transformer_engine_ep.dispatch( - self.buffer.handle_mem, - topk_idx, - quantized_input._rowwise_data.view(torch.float8_e4m3fn), - topk_weights, - recv_tokens.view(torch.float8_e4m3fn), - recv_topk_weights, - input_scale_inv, - recv_scale_inv, - ) - output = _make_grouped_mxfp8( - recv_tokens, - recv_scale_inv, - tokens_per_expert, - quantized_input._fp8_dtype, - torch.bfloat16, - ) + tokens_per_expert = self.buffer.tokens_per_expert # If next_op_input_quantizer is different from input_quantizer, # we need to requantize the data, which is handled in grouped_linear anyway. # We won't get any fusion benefit, so don't do it here. - # Save only shape/dtype metadata and the opaque routing handle for backward. ctx = basic_op_ctxs[0] if ctx.requires_grad: - ctx.input_shape = input_shape ctx.input_dtype = torch.bfloat16 - ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.dispatch_state = dispatch_state ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer - ctx.save_for_backward(self.buffer.handle_mem) return output, [(tokens_per_expert, recv_topk_weights, self.buffer.handle_mem, topk_idx)] @@ -293,7 +179,6 @@ def fuser_backward( Iterable[Iterable[Optional[torch.Tensor]]], ]: ctx = basic_op_ctxs[0] - (handle_mem,) = ctx.saved_tensors # Only BF16 Dispatch_bwd is supported for now. grad_output = maybe_dequantize(grad_output, ctx.input_dtype) @@ -307,22 +192,10 @@ def fuser_backward( else: grad_recv_weights = grad_recv_weights.to(dtype=torch.float32) - grad_input = torch.empty( - ctx.input_shape, - dtype=ctx.input_dtype, - device=grad_output.device, - ) - grad_topk_weights = torch.empty( - ctx.topk_weights_shape, - dtype=torch.float32, - device=grad_output.device, - ) - torch.ops.transformer_engine_ep.dispatch_bwd( - handle_mem, + grad_input, grad_topk_weights = _ep_dispatch_bwd( + ctx.dispatch_state, grad_output, grad_recv_weights, - grad_input, - grad_topk_weights, ) quantizer = ctx.prev_op_grad_output_quantizer if quantizer is not None: From be9a0dc27e332cdb877841cb3f49c044c68e6363 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 23:42:47 +0000 Subject: [PATCH 67/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ep_reference.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 0a02d33b9f..981b04e1cf 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -123,7 +123,9 @@ def _validate_storage(self) -> None: expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") data_shape = self.logical_shape if self.data.dtype != expected_dtype: - raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) else: expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -613,7 +615,9 @@ def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") order = torch.argsort(destination, stable=True) - send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) recv_counts_tensor = self._exchange_counts(send_counts_tensor) return _DispatchPlan( send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), @@ -728,13 +732,17 @@ def __call__( if tuple(topk_idx.shape) != route_shape: raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") if tuple(topk_weights.shape) != route_shape: - raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) if topk_idx.dtype not in (torch.int32, torch.int64): raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") if not topk_weights.is_floating_point(): raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) device = _tensor_device(activation) inputs = { @@ -822,7 +830,11 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) # recv rows are ordered by source rank, then that source's token-major # route order, so the per-expert position grouping below realizes the # documented fc1_c ordering. From f6c6306e06833df813482bbfa23512bfb015b85b Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 26 Aug 2026 20:17:50 +0000 Subject: [PATCH 68/83] limit the recv tokens per rank Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fused/moe_ep.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index ce94cfcb75..f30fe9f71d 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -305,14 +305,15 @@ def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: return False if buffer.max_tokens_per_rank is None or buffer.max_tokens_per_rank <= 0: return False + if ( + buffer.recv_capacity_per_rank is not None + and buffer.recv_capacity_per_rank <= 0 + ): + return False if buffer.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: return False if buffer.top_k > 32: return False - ep_group = get_ep_group() - ep_size = 1 if ep_group is None else ep_group.size() - if ep_size > 16: - return False return True @@ -331,7 +332,7 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo ): return False buffer = dispatch.buffer - if not buffer.eager or buffer.payload_dtype is not torch.bfloat16: + if buffer.payload_dtype is not torch.bfloat16: return False if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): return False @@ -382,6 +383,8 @@ def __init__( top_k=dispatch.buffer.top_k, ep_group=ep_group, max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, + max_recv_size_per_rank=dispatch.buffer.recv_capacity_per_rank, + drop_on_overflow=False, apply_topk_in_fc1=True, generate_c=True, backward_wgrad_mode="operands", From 38fd3650dcaebed84d95dfe14f90ffc6e934366f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 20:19:09 +0000 Subject: [PATCH 69/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/fused/moe_ep.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index f30fe9f71d..3921733422 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -305,10 +305,7 @@ def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: return False if buffer.max_tokens_per_rank is None or buffer.max_tokens_per_rank <= 0: return False - if ( - buffer.recv_capacity_per_rank is not None - and buffer.recv_capacity_per_rank <= 0 - ): + if buffer.recv_capacity_per_rank is not None and buffer.recv_capacity_per_rank <= 0: return False if buffer.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: return False From fedcda0705128d68c2b19efa81b0c0528c0f4840 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 18:27:22 +0000 Subject: [PATCH 70/83] address review comments + rework design for ops(epconfig + kwarg epbuffer) and allow mxfp8 combine through env variable Signed-off-by: Varun Thumbe --- examples/pytorch/ep/bench/ep_bench.py | 4 +- examples/pytorch/ep/ep_moe.py | 4 +- tests/pytorch/distributed/run_ep.py | 131 +++++++++-- transformer_engine/pytorch/ep.py | 217 ++++++++++++------ transformer_engine/pytorch/ops/_common.py | 82 +++---- .../pytorch/ops/basic/__init__.py | 4 +- .../pytorch/ops/basic/combine.py | 108 ++++----- .../pytorch/ops/basic/dispatch.py | 66 +++--- .../pytorch/ops/fused/moe_ep.py | 85 ++++--- 9 files changed, 433 insertions(+), 268 deletions(-) diff --git a/examples/pytorch/ep/bench/ep_bench.py b/examples/pytorch/ep/bench/ep_bench.py index f80a57015c..7091422894 100644 --- a/examples/pytorch/ep/bench/ep_bench.py +++ b/examples/pytorch/ep/bench/ep_bench.py @@ -30,6 +30,7 @@ from transformer_engine.pytorch.ep import ( EpBuffer, + EpConfig, ep_bootstrap, ep_combine, ep_dispatch, @@ -202,13 +203,14 @@ def main(): else None ) - buffer = EpBuffer( + config = EpConfig( top_k=K, max_tokens_per_rank=T, recv_capacity_per_rank=recv_pr, hidden_dim=H, num_local_experts=num_local_experts, ) + buffer = EpBuffer(config) tokens = tokens_hbm topk_w = topk_w_hbm diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py index b47c334225..45993aa4cc 100644 --- a/examples/pytorch/ep/ep_moe.py +++ b/examples/pytorch/ep/ep_moe.py @@ -16,6 +16,7 @@ from transformer_engine.pytorch.ep import ( EpBuffer, + EpConfig, ep_bootstrap, ep_combine, ep_dispatch, @@ -177,13 +178,14 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, else None ) - buffer = EpBuffer( + config = EpConfig( top_k=args.top_k, max_tokens_per_rank=T, recv_capacity_per_rank=recv_pr, hidden_dim=args.hidden, num_local_experts=num_local_experts, ) + buffer = EpBuffer(config) recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w, recv_tokens=recv_tokens) expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 357b3bcff5..6455bbc5a1 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -4,6 +4,7 @@ """Multi-process PyTorch EP tests, launched via torchrun (one process per GPU).""" from contextlib import nullcontext +from dataclasses import replace import os import sys import unittest @@ -17,6 +18,7 @@ from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, + EpConfig, ep_bootstrap, ep_finalize, ep_prepare, @@ -37,6 +39,7 @@ from transformer_engine.pytorch.ops.fused.moe_ep import ( FusedMoeEp, _cudnn_megamoe_supported, + _get_megamoe_combine_format, _pack_cudnn_activation, _pack_cudnn_weights, ) @@ -290,13 +293,18 @@ def _make_buffer( dispatch_fwd_quant_recipe=None, combine_bwd_quant_recipe=None, ): - return EpBuffer( + config = EpConfig( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, alignment=alignment, + zero_copy=ZERO_COPY, + drop_on_overflow=OVERFLOW, + ) + return EpBuffer( + config, dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, combine_bwd_quant_recipe=combine_bwd_quant_recipe, ) @@ -911,9 +919,78 @@ def _require_mxfp8_shapes(self): "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)" ) + def test_runtime_buffer_config_mismatch(self): + buffer = self._make_buffer() + wrong_config = replace(buffer.config, hidden_dim=buffer.config.hidden_dim + 1) + topk_idx, tokens, topk_weights = _make_identity_inputs( + self.cfg.rank, + self.cfg.ep_size, + ) + with self.assertRaisesRegex(ValueError, "runtime buffer config"): + te_ops.MoeDispatch(wrong_config)( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) + expert_out = torch.empty( + self.cfg.recv_capacity_per_rank, + HIDDEN_DIM, + dtype=torch.bfloat16, + device=self.cfg.device, + ) + with self.assertRaisesRegex(ValueError, "runtime buffer config"): + te_ops.MoeCombine(wrong_config)(expert_out, buffer=buffer) + + def test_megamoe_combine_format_env(self): + old_value = os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE") + try: + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = "0" + self.assertEqual(_get_megamoe_combine_format(), "bf16") + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = "1" + self.assertEqual(_get_megamoe_combine_format(), "mxfp8") + finally: + if old_value is None: + os.environ.pop("NVTE_MEGAMOE_MXFP8_COMBINE", None) + else: + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = old_value + + @_mxfp8_align_test + def test_role_quantizer_requires_matching_buffer_recipe(self): + self._require_mxfp8_shapes() + buffer = self._make_buffer(alignment=128) + dispatch = te_ops.MoeDispatch(buffer.config) + combine = te_ops.MoeCombine(buffer.config) + topk_idx, tokens, topk_weights = _make_identity_inputs( + self.cfg.rank, + self.cfg.ep_size, + ) + recipe = MXFP8BlockScaling() + with te.autocast(enabled=True, recipe=recipe): + with self.assertRaisesRegex(ValueError, "does not have an MXFP8BlockScaling recipe"): + dispatch(tokens, topk_idx, topk_weights, buffer=buffer) + expert_out = torch.empty( + self.cfg.recv_capacity_per_rank, + HIDDEN_DIM, + dtype=torch.bfloat16, + device=self.cfg.device, + requires_grad=True, + ) + with self.assertRaisesRegex(ValueError, "does not have an MXFP8BlockScaling recipe"): + combine(expert_out, buffer=buffer) + def _make_dispatch_combine_ops(self, *, mxfp8): - buffer = self._make_buffer(alignment=128 if mxfp8 else 0) - return buffer, te_ops.Dispatch(buffer), te_ops.Combine() + recipe = MXFP8BlockScaling() if mxfp8 else None + buffer = self._make_buffer( + alignment=128 if mxfp8 else 0, + dispatch_fwd_quant_recipe=recipe, + combine_bwd_quant_recipe=recipe, + ) + return ( + buffer, + te_ops.MoeDispatch(buffer.config), + te_ops.MoeCombine(buffer.config), + ) def _run_dispatch_combine_identity(self, *, mxfp8): """Route, apply top-k weights, and combine back to local token order.""" @@ -926,10 +1003,11 @@ def _run_dispatch_combine_identity(self, *, mxfp8): ) recipe = MXFP8BlockScaling() if mxfp8 else None with te.autocast(enabled=mxfp8, recipe=recipe): - recv_tokens, tokens_per_expert, recv_weights, handle, routing_indices = dispatch( + recv_tokens, tokens_per_expert, recv_weights = dispatch( tokens, topk_idx, topk_weights, + buffer=buffer, ) if mxfp8: recv_tokens = _degroup_mxfp8(recv_tokens) @@ -939,25 +1017,23 @@ def _run_dispatch_combine_identity(self, *, mxfp8): ) output = combine( weighted_expert_output, - handle, - tokens_per_expert, - routing_indices, + buffer=buffer, ) torch.cuda.synchronize() torch.testing.assert_close(output, tokens, atol=5e-2, rtol=5e-2) - self.assertEqual(handle.data_ptr(), buffer.handle_mem.data_ptr()) + self.assertEqual(tokens_per_expert.data_ptr(), buffer.tokens_per_expert.data_ptr()) @_eager_test_include @_zero_copy_test_include def test_dispatch_combine_identity_bf16(self): - """Dispatch and Combine basic ops form an identity in BF16.""" + """MoeDispatch and MoeCombine basic ops form an identity in BF16.""" self._run_dispatch_combine_identity(mxfp8=False) @_eager_test_include @_zero_copy_test_include @_mxfp8_align_test def test_dispatch_combine_identity_mxfp8(self): - """Dispatch and Combine basic ops form an identity with MXFP8 transport.""" + """MoeDispatch and MoeCombine basic ops form an identity with MXFP8 transport.""" self._run_dispatch_combine_identity(mxfp8=True) def _make_megamoe_model( @@ -968,8 +1044,12 @@ def _make_megamoe_model( delay_wgrad_compute=False, ): """Build the exact five-op sequence recognized by MegaMoE fusion.""" - buffer = self._make_buffer(alignment=128 if recipe is not None else 0) - dispatch = te_ops.Dispatch(buffer) + buffer = self._make_buffer( + alignment=128 if recipe is not None else 0, + dispatch_fwd_quant_recipe=recipe, + combine_bwd_quant_recipe=recipe, + ) + dispatch = te_ops.MoeDispatch(buffer.config) init_ctx = ( te.quantized_model_init(enabled=True, recipe=recipe) if recipe is not None @@ -1007,20 +1087,15 @@ def _make_megamoe_model( del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] else: os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param - combine = te_ops.Combine() + combine = te_ops.MoeCombine(buffer.config) dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) - dispatch.set_extra_output_channel(2, "ep_handle", output_to_caller=False) - dispatch.set_extra_output_channel(3, "routing_indices", output_to_caller=False) fc1.set_extra_input_channel(0, "tokens_per_expert") activation.set_extra_input_channel(0, "routing_weights") fc2.set_extra_input_channel(0, "tokens_per_expert") - combine.set_extra_input_channel(0, "ep_handle") - combine.set_extra_input_channel(1, "tokens_per_expert") - combine.set_extra_input_channel(2, "routing_indices") model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) - return model, fc1, fc2 + return model, fc1, fc2, buffer @_eager_test_include def test_megamoe_bf16_numerics(self): @@ -1118,7 +1193,7 @@ def _run_megamoe_vs_reference( exercises the same sequence as separate NCCL EP and grouped-MLP ops. """ recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None - model, fc1, fc2 = self._make_megamoe_model( + model, fc1, fc2, buffer = self._make_megamoe_model( recipe=recipe, accumulate_into_main_grad=accumulate_into_main_grad, delay_wgrad_compute=delay_wgrad_compute, @@ -1168,7 +1243,15 @@ def _run_megamoe_vs_reference( ) with autocast_ctx: model_input = self._mxfp8_quantizer()(seq_tokens) if prequantized_input else seq_tokens - seq_out = model(model_input, topk_idx, seq_topk_weights) + seq_out = model( + model_input, + topk_idx, + seq_topk_weights, + op_kwargs={ + 0: {"buffer": buffer}, + 4: {"buffer": buffer}, + }, + ) forward_ops = model._module_groups[0]._forward_ops fused = len(forward_ops) == 1 and isinstance(forward_ops[0][0], FusedMoeEp) @@ -1189,7 +1272,11 @@ def _run_megamoe_vs_reference( ep_group=self.ep_group, max_tokens_per_rank=TOKENS_PER_RANK, output_format=MoeFormat.BF16, - combine_format=MoeFormat.BF16, + combine_format=( + MoeFormat.MXFP8 + if fused and os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE", "0") == "1" + else MoeFormat.BF16 + ), apply_topk_in_fc1=True, generate_c=True, intermediate_format=MoeFormat.MXFP8 if emulate_mxfp8 else None, diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 3d157dfcff..2df71f334f 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -24,8 +24,10 @@ # common.recipe (the concrete recipe classes are imported lazily where used). if TYPE_CHECKING: from ..common.recipe import Recipe + from .quantized_tensor import Quantizer __all__ = [ + "EpConfig", "EpBuffer", "ep_bootstrap", "get_ep_group", @@ -83,11 +85,27 @@ def _check_nccl_runtime_version() -> None: # omitted); ep_dispatch reads it to size the recv outputs from the per-step # recv-token total instead of a fixed recv_capacity_per_rank. _EAGER = False +_BOOTSTRAP_SETTINGS: Optional[dict[str, object]] = None + + +@dataclass(frozen=True) +class EpConfig: + """Immutable configuration shared by an EP buffer and its MoE operations.""" + + top_k: int + hidden_dim: int + num_local_experts: int + max_tokens_per_rank: int + recv_capacity_per_rank: Optional[int] + alignment: int = 0 + payload_dtype: torch.dtype = torch.bfloat16 + zero_copy: bool = False + drop_on_overflow: bool = False def _atexit_finalize() -> None: """Best-effort teardown at interpreter shutdown; swallows errors.""" - global _BOOTSTRAPPED, _EP_GROUP, _EAGER + global _BOOTSTRAPPED, _EP_GROUP, _EAGER, _BOOTSTRAP_SETTINGS if _BOOTSTRAPPED: try: tex.ep_finalize() @@ -99,6 +117,7 @@ def _atexit_finalize() -> None: _BOOTSTRAPPED = False _EP_GROUP = None _EAGER = False + _BOOTSTRAP_SETTINGS = None def ep_bootstrap( @@ -134,7 +153,7 @@ def ep_bootstrap( ``drop_on_overflow`` drops tokens exceeding ``recv_capacity_per_rank`` instead of trapping. Requires ``recv_capacity_per_rank``. """ - global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER + global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER, _BOOTSTRAP_SETTINGS eager = recv_capacity_per_rank is None if _BOOTSTRAPPED: raise RuntimeError("ep_bootstrap was already called in this process") @@ -181,6 +200,18 @@ def ep_bootstrap( _BOOTSTRAPPED = True _EP_GROUP = ep_group _EAGER = bool(eager) + _BOOTSTRAP_SETTINGS = { + "top_k": int(num_topk), + "hidden_dim": int(hidden_dim), + "num_local_experts": int(num_experts) // ep_group.size(), + "max_tokens_per_rank": int(max_tokens_per_rank), + "recv_capacity_per_rank": ( + None if recv_capacity_per_rank is None else int(recv_capacity_per_rank) + ), + "payload_dtype": max_token_dtype, + "zero_copy": bool(zero_copy), + "drop_on_overflow": bool(drop_on_overflow), + } if not _ATEXIT_REGISTERED: atexit.register(_atexit_finalize) _ATEXIT_REGISTERED = True @@ -211,7 +242,7 @@ def ep_finalize() -> None: a caller that used ``symm_mem_alloc(use_pool=True)`` does not need a separate ``release_symm_mem_pool()`` before destroying the PG. """ - global _BOOTSTRAPPED, _EP_GROUP, _EAGER + global _BOOTSTRAPPED, _EP_GROUP, _EAGER, _BOOTSTRAP_SETTINGS if not _BOOTSTRAPPED: return try: @@ -222,6 +253,7 @@ def ep_finalize() -> None: _BOOTSTRAPPED = False _EP_GROUP = None _EAGER = False + _BOOTSTRAP_SETTINGS = None def is_symm_backed(t: torch.Tensor) -> bool: @@ -248,22 +280,17 @@ def is_symm_backed(t: torch.Tensor) -> bool: class EpBuffer: - """Per-microbatch EP layer state: handle_mem, tokens_per_expert, and shape/dtype config. + """Per-microbatch EP routing state allocated from an immutable :class:`EpConfig`. + Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). """ __slots__ = ( + "config", "handle_mem", - "top_k", - "alignment", - "max_tokens_per_rank", - "recv_capacity_per_rank", - "hidden_dim", - "num_local_experts", - "payload_dtype", "device", "tokens_per_expert", - "zero_copy", + "num_local_tokens", "eager", "total_recv_tokens", "dispatch_fwd_quant_recipe", @@ -272,49 +299,43 @@ class EpBuffer: def __init__( self, - top_k: int, - max_tokens_per_rank: int, - hidden_dim: int, - num_local_experts: int, - recv_capacity_per_rank: Optional[int] = None, - alignment: int = 0, - payload_dtype: torch.dtype = torch.bfloat16, + config: EpConfig, + *, device: Optional[torch.device] = None, dispatch_fwd_quant_recipe: Optional["Recipe"] = None, combine_bwd_quant_recipe: Optional["Recipe"] = None, ) -> None: if not _BOOTSTRAPPED: raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") + if not isinstance(config, EpConfig): + raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + if _BOOTSTRAP_SETTINGS is None: + raise RuntimeError("EP bootstrap configuration is unavailable.") + mismatches = { + name: (getattr(config, name), expected) + for name, expected in _BOOTSTRAP_SETTINGS.items() + if getattr(config, name) != expected + } + if mismatches: + details = ", ".join( + f"{name}={actual!r} (bootstrapped {expected!r})" + for name, (actual, expected) in mismatches.items() + ) + raise ValueError(f"EpConfig does not match ep_bootstrap: {details}.") if device is None: device = torch.device("cuda", torch.cuda.current_device()) - alignment = int(alignment) - if alignment > 1 and (alignment & (alignment - 1)) != 0: - raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") - self.eager = _EAGER - if not self.eager and recv_capacity_per_rank is None: - raise ValueError( - "EpBuffer requires recv_capacity_per_rank unless the EP group was " - "bootstrapped in eager mode (recv_capacity_per_rank omitted)." - ) - self.top_k = int(top_k) - self.alignment = alignment - self.max_tokens_per_rank = int(max_tokens_per_rank) - self.recv_capacity_per_rank = ( - None if recv_capacity_per_rank is None else int(recv_capacity_per_rank) - ) - self.hidden_dim = int(hidden_dim) - self.num_local_experts = int(num_local_experts) - self.payload_dtype = payload_dtype + self.config = config + self.eager = config.recv_capacity_per_rank is None self.device = device - self.zero_copy = bool(tex.ep_get_zero_copy()) self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe self.combine_bwd_quant_recipe = combine_bwd_quant_recipe - size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) + size_bytes = tex.ep_handle_mem_size(config.top_k, config.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) self.tokens_per_expert = torch.empty( - self.num_local_experts, dtype=torch.int64, device=device + config.num_local_experts, dtype=torch.int64, device=device ) + self.num_local_tokens = config.max_tokens_per_rank # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) # Per-step recv-token total (int64 [1]), written by ep_prepare. Eager reads it @@ -329,6 +350,38 @@ def __init__( self.total_recv_tokens = torch.empty(1, dtype=torch.int64, device=device) mark_not_offload(self.total_recv_tokens) + @property + def top_k(self) -> int: + return self.config.top_k + + @property + def alignment(self) -> int: + return self.config.alignment + + @property + def max_tokens_per_rank(self) -> int: + return self.config.max_tokens_per_rank + + @property + def recv_capacity_per_rank(self) -> Optional[int]: + return self.config.recv_capacity_per_rank + + @property + def hidden_dim(self) -> int: + return self.config.hidden_dim + + @property + def num_local_experts(self) -> int: + return self.config.num_local_experts + + @property + def payload_dtype(self) -> torch.dtype: + return self.config.payload_dtype + + @property + def zero_copy(self) -> bool: + return self.config.zero_copy + # torch.library custom ops (so they don't graph-break under torch.compile) @@ -725,7 +778,15 @@ def forward( # type: ignore[override] the buffer object to keep the autograd operand list short.""" tokens_scale_inv = None if buffer.dispatch_fwd_quant_recipe is not None: - tokens, tokens_scale_inv = _quantize_mxfp8(tokens) + from .tensor.mxfp8_tensor import MXFP8Quantizer + # Only MXFP8 Quantizer is supported for EP dispatch + quantizer = MXFP8Quantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + quantizer.internal = True + tokens, tokens_scale_inv = quantize_for_ep(tokens, quantizer) recv_out, recv_topk_weights, state = _ep_prepare_and_dispatch_fwd( tokens, topk_weights, @@ -843,7 +904,15 @@ def _ep_combine_bwd( torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) else: if quantized_grad is None: - mx, grad_scale_inv = _quantize_mxfp8(g_result) + from .tensor.mxfp8_tensor import MXFP8Quantizer + # Only MXFP8 Quantizer is supported for EP combine bwd + quantizer = MXFP8Quantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + quantizer.internal = True + mx, grad_scale_inv = quantize_for_ep(g_result, quantizer) else: mx = quantized_grad if grad_scale_inv is None: @@ -972,33 +1041,46 @@ def _as_mxfp8_storage(tensor: QuantizedTensorStorage): ) -def _quantize_mxfp8(x: torch.Tensor): - """Quantize to lightweight MXFP8 storage and return it with compact transport scales. - - ``scale_inv`` has shape ``[T, H/block]``. EP routes and returns E4M3 data in both directions, - so quantize to E4M3 regardless of pass. The GEMM scale-row padding is stripped and each compact - scale row must remain 16-byte aligned. - """ - from .constants import MXFP8_BLOCK_SCALING_SIZE +def quantize_for_ep( + input_: torch.Tensor | QuantizedTensorStorage, + quantizer: Optional["Quantizer"], +) -> tuple[MXFP8TensorStorage, torch.Tensor]: + """Return an MXFP8 input and its compact rowwise scales for EP.""" + from .constants import DType, MXFP8_BLOCK_SCALING_SIZE from .tensor.mxfp8_tensor import MXFP8Quantizer - quantizer = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) - quantizer.internal = True - mx = quantizer.quantize(x) - if mx._with_gemm_swizzled_scales: - raise RuntimeError( - "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact." + if quantizer is not None: + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + f"EP MXFP8 transport requires MXFP8Quantizer, got {type(quantizer).__name__}." + ) + if quantizer.dtype != DType.kFloat8E4M3: + raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") + + if isinstance(input_, MXFP8TensorStorage): + quantized = input_ + elif isinstance(input_, QuantizedTensorStorage): + raise TypeError( + f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}." ) - data = mx._rowwise_data - scale_inv = mx._rowwise_scale_inv + else: + if quantizer is None: + raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") + quantized = quantizer(input_) + quantized = _as_mxfp8_storage(quantized) + if quantized._fp8_dtype != DType.kFloat8E4M3: + raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") + if quantized._with_gemm_swizzled_scales: + raise ValueError("EP requires unswizzled MXFP8 scales.") + data = quantized._rowwise_data + scale_inv = quantized._rowwise_scale_inv if data is None or scale_inv is None: - raise ValueError("MXFP8 tokens must carry rowwise data and scale_inv for EP dispatch.") - t_flat = x.shape[0] - hidden = x.shape[-1] - cols = hidden // MXFP8_BLOCK_SCALING_SIZE + raise ValueError("EP requires rowwise MXFP8 data and scales.") + rows, hidden = input_.shape + scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE # The backend forwards each token's scale row with a 16-byte-aligned store, so the row # (cols * dtype bytes) must be a multiple of 16. - scale_row_bytes = cols * scale_inv.element_size() + scale_row_bytes = scale_cols * scale_inv.element_size() if scale_row_bytes % 16 != 0: raise ValueError( f"MXFP8 dispatch requires a 16-byte-aligned scale row; hidden={hidden} gives " @@ -1008,13 +1090,10 @@ def _quantize_mxfp8(x: torch.Tensor): # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block] # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the # slice stays contiguous; assert rather than force a copy. - scale_inv = scale_inv[:t_flat, :cols] + scale_inv = scale_inv[:rows, :scale_cols] if not scale_inv.is_contiguous(): - raise ValueError( - "MXFP8 dispatch requires compact contiguous scales [T, H/block]; got a " - f"non-contiguous [{t_flat}, {cols}] slice." - ) - return mx, scale_inv + raise ValueError("EP requires compact contiguous MXFP8 scales.") + return quantized, scale_inv def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, device, zero_copy): diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 8eae6d7cf7..546f65e1f8 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,8 +11,7 @@ import torch from transformer_engine_torch import FP8TensorMeta -from ..constants import DType, MXFP8_BLOCK_SCALING_SIZE -from ..ep import _as_mxfp8_storage +from ..ep import EpBuffer, EpConfig from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..quantized_tensor import QuantizedTensorStorage, Quantizer @@ -24,7 +23,6 @@ NVFP4Quantizer, ) from ..tensor.float8_tensor import Float8Tensor -from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..utils import canonicalize_dtype @@ -70,6 +68,43 @@ def validate_buffer( return buffer.detach() +def validate_ep_buffer( + op_name: str, + expected_config: EpConfig, + buffer: object, +) -> EpBuffer: + """Validate a runtime EP buffer against an operation's immutable config.""" + if not isinstance(buffer, EpBuffer): + raise TypeError(f"{op_name} requires buffer=EpBuffer(...), got {type(buffer).__name__}.") + if buffer.config != expected_config: + raise ValueError( + f"{op_name} runtime buffer config does not match its initialized config: " + f"{buffer.config!r} != {expected_config!r}." + ) + return buffer + + +def validate_ep_comms_recipe( + op_name: str, + quantizer: Optional[Quantizer], + buffer_recipe: object, +) -> None: + """Require the buffer recipe to match the Op's quantizer role.""" + if isinstance(quantizer, MXFP8Quantizer): + from transformer_engine.common.recipe import MXFP8BlockScaling + + if not isinstance(buffer_recipe, MXFP8BlockScaling): + raise ValueError( + f"{op_name} selected MXFP8 Comms from its quantizer role, but the " + "runtime EpBuffer does not have an MXFP8BlockScaling recipe." + ) + elif buffer_recipe is not None: + raise ValueError( + f"{op_name} selected BF16 Comms from its quantizer role, but the " + f"runtime EpBuffer has recipe {type(buffer_recipe).__name__}." + ) + + def validate_or_alloc_output( buffer: Optional[torch.Tensor], shape: tuple[int, ...] | list[int], @@ -118,47 +153,6 @@ def maybe_dequantize( return tensor -def quantize_for_ep( - input_: torch.Tensor | QuantizedTensorStorage, - quantizer: Optional[Quantizer], -) -> tuple[MXFP8TensorStorage, torch.Tensor]: - """Return an MXFP8 input and its compact rowwise scales for EP.""" - if quantizer is not None: - if not isinstance(quantizer, MXFP8Quantizer): - raise TypeError( - f"EP MXFP8 transport requires MXFP8Quantizer, got {type(quantizer).__name__}." - ) - if quantizer.dtype != DType.kFloat8E4M3: - raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") - - if isinstance(input_, MXFP8TensorStorage): - quantized = input_ - elif isinstance(input_, QuantizedTensorStorage): - raise TypeError(f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}.") - else: - if quantizer is None: - raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") - if not quantizer.internal: - quantizer = quantizer.copy() - quantizer.internal = True - quantized = quantizer(input_) - quantized = _as_mxfp8_storage(quantized) - if quantized._fp8_dtype != DType.kFloat8E4M3: - raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") - if quantized._with_gemm_swizzled_scales: - raise ValueError("EP requires unswizzled MXFP8 scales.") - data = quantized._rowwise_data - scale_inv = quantized._rowwise_scale_inv - if data is None or scale_inv is None: - raise ValueError("EP requires rowwise MXFP8 data and scales.") - rows, hidden = input_.shape - scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE - scale_inv = scale_inv[:rows, :scale_cols] - if not scale_inv.is_contiguous(): - raise ValueError("EP requires compact contiguous MXFP8 scales.") - return quantized, scale_inv - - def maybe_autocast_dtype( *, device_type: str = "cuda", diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index f82e5264fc..aa327aabe5 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -22,9 +22,9 @@ from .all_reduce import AllReduce from .basic_linear import BasicLinear from .bias import Bias -from .combine import Combine +from .combine import MoeCombine from .constant_scale import ConstantScale -from .dispatch import Dispatch +from .dispatch import MoeDispatch from .dropout import Dropout from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported from .identity import Identity diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index b08d302dff..e52040f701 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -9,64 +9,71 @@ from typing import Any, Iterable, Optional import torch -import transformer_engine_torch as tex from ...ep import ( + EpBuffer, + EpConfig, _alloc_io, _ep_combine_bwd, _ep_combine_fwd, - _ep_is_eager, is_symm_backed, + quantize_for_ep, ) from ...quantization import QuantizerRole from ...tensor import MXFP8Quantizer, Quantizer from .._common import ( - is_quantized_tensor, maybe_dequantize, - quantize_for_ep, validate_buffer, + validate_ep_buffer, + validate_ep_comms_recipe, ) from ..op import BasicOperation, OperationContext def _validate_combine_inputs( input_: torch.Tensor, - handle_mem: torch.Tensor, - tokens_per_expert: torch.Tensor, - topk_idx: torch.Tensor, -) -> tuple[tuple[int, int], int]: - """Validate the expert output and routing metadata consumed by Combine.""" + buffer: EpBuffer, +) -> tuple[int, int]: + """Validate the expert output and routing metadata consumed by MoeCombine.""" if input_.dtype is not torch.bfloat16: raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") if input_.ndim != 2: - raise ValueError(f"Combine input must be 2D, got shape {tuple(input_.shape)}.") - if handle_mem.dtype is not torch.uint8 or handle_mem.device != input_.device: - raise ValueError("Combine routing handle must be a uint8 tensor on the input device.") - if tokens_per_expert.dtype is not torch.int64 or tokens_per_expert.device != input_.device: - raise ValueError("Combine tokens_per_expert must be an int64 tensor on the input device.") - return tuple(input_.shape), topk_idx.shape[0] - - -class Combine(BasicOperation): + raise ValueError(f"MoeCombine input must be 2D, got shape {tuple(input_.shape)}.") + if buffer.handle_mem.dtype is not torch.uint8 or buffer.handle_mem.device != input_.device: + raise ValueError("MoeCombine routing handle must be a uint8 tensor on the input device.") + if ( + buffer.tokens_per_expert.dtype is not torch.int64 + or buffer.tokens_per_expert.device != input_.device + ): + raise ValueError("MoeCombine tokens_per_expert must be an int64 tensor on the input device.") + return tuple(input_.shape) + + +class MoeCombine(BasicOperation): """Combine pre-weighted local expert outputs with NCCL EP. - The operation consumes the routing handle, tokens-per-expert, and routing - indices produced by a preceding :class:`Dispatch` through extra-tensor - channels. + The operation consumes routing state from a runtime :class:`EpBuffer`. """ - num_extra_inputs: int = 3 + num_extra_inputs: int = 0 + + def __init__(self, config: EpConfig) -> None: + super().__init__() + if not isinstance(config, EpConfig): + raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + self.config = config def num_quantizers(self, mode: str) -> int: return 1 if mode == "backward" else 0 def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: if mode == "backward": + # combine backward dispatches grad_output name = getattr(self, "name", "") or "" return [ QuantizerRole( module_type="combine", - tensor_type="grad_output", + tensor_type="dispatch_grad_output", name=name, ) ] @@ -81,10 +88,10 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: quantizer.internal = True def op_forward(self, *args: Any, **kwargs: Any) -> None: - raise RuntimeError("Combine uses fuser_forward") + raise RuntimeError("MoeCombine uses fuser_forward") def op_backward(self, *args: Any, **kwargs: Any) -> None: - raise RuntimeError("Combine uses fuser_backward") + raise RuntimeError("MoeCombine uses fuser_backward") @staticmethod def _prepare_grad_buffer( @@ -97,7 +104,7 @@ def _prepare_grad_buffer( zero_copy: bool, ) -> Optional[torch.Tensor]: """Validate caller storage for the expert-output gradient.""" - if grad_output_quantizer is None: + if not isinstance(grad_output_quantizer, MXFP8Quantizer): grad_out = validate_buffer( "grad_out", grad_out, @@ -113,7 +120,7 @@ def _prepare_grad_buffer( contiguous=True, ) if zero_copy and grad_out is not None and not is_symm_backed(grad_out): - raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") + raise ValueError("zero-copy MoeCombine grad_out must be symmetric-memory-backed.") return grad_out def fuser_forward( @@ -129,18 +136,21 @@ def fuser_forward( # Combine's transport format is selected by Combine's own grad-output quantizer. # If the preceding op expects a different gradient quantized format, # it is requantized in that op's backward implementation (e.g., GroupedLinear backward). - del prev_op_grad_output_quantizer + del basic_op_extra_inputs, prev_op_grad_output_quantizer, next_op_input_quantizer grad_output_quantizer = self.get_quantizer("backward", 0) - handle_mem, tokens_per_expert, topk_idx = basic_op_extra_inputs[0] + transport_quantizer = ( + grad_output_quantizer if isinstance(grad_output_quantizer, MXFP8Quantizer) else None + ) kwargs = basic_op_kwargs[0] - input_shape, num_local_tokens = _validate_combine_inputs( - input_, - handle_mem, - tokens_per_expert, - topk_idx, + buffer = validate_ep_buffer("MoeCombine", self.config, kwargs.get("buffer")) + validate_ep_comms_recipe( + "MoeCombine", + grad_output_quantizer, + buffer.combine_bwd_quant_recipe, ) + input_shape = _validate_combine_inputs(input_, buffer) # Stage zero-copy input if needed, then restore local-token order. - zero_copy = bool(tex.ep_get_zero_copy()) + zero_copy = buffer.zero_copy expert_out = input_ if zero_copy: expert_out = _alloc_io(tuple(input_.shape), input_.dtype, input_.device, True) @@ -151,7 +161,7 @@ def fuser_forward( if ctx.requires_grad: grad_out = self._prepare_grad_buffer( kwargs.get("grad_out"), - grad_output_quantizer, + transport_quantizer, input_shape=input_shape, input_dtype=input_.dtype, device=input_.device, @@ -161,20 +171,17 @@ def fuser_forward( result, combine_state = _ep_combine_fwd( expert_out, grad_out, - handle_mem=handle_mem, - token_counts=tokens_per_expert, - num_local_tokens=num_local_tokens, + handle_mem=buffer.handle_mem, + token_counts=buffer.tokens_per_expert, + num_local_tokens=buffer.num_local_tokens, hidden_dim=expert_out.shape[-1], - bwd_quant_recipe=grad_output_quantizer, - eager=_ep_is_eager(), + bwd_quant_recipe=transport_quantizer, + eager=buffer.eager, zero_copy=zero_copy, ) if ctx.requires_grad: ctx.combine_state = combine_state - # Hand off to the next op in its requested representation. - if next_op_input_quantizer is not None and not is_quantized_tensor(result): - result = next_op_input_quantizer(result) return result, [()] def fuser_backward( @@ -193,24 +200,19 @@ def fuser_backward( grad_output_quantizer = self.get_quantizer("backward", 0) grad_scale_inv = None # Prepare grad_output (Quantize if necessary) - if grad_output_quantizer is None: - grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() - quantized_grad = None - elif isinstance(grad_output_quantizer, MXFP8Quantizer): + if isinstance(grad_output_quantizer, MXFP8Quantizer): quantized_grad, grad_scale_inv = quantize_for_ep( grad_output, grad_output_quantizer, ) grad_output = quantized_grad else: - raise TypeError( - "NCCL EP Combine backward supports MXFP8Quantizer only, got " - f"{type(grad_output_quantizer).__name__}." - ) + grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() + quantized_grad = None grad_input = _ep_combine_bwd( ctx.combine_state, grad_output, quantized_grad, grad_scale_inv, ) - return grad_input, [()], [(None, None, None)] + return grad_input, [()], [()] diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 43a96a63d1..5fe06640fe 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -12,15 +12,18 @@ from ...ep import ( EpBuffer, + EpConfig, _ep_dispatch_bwd, _ep_prepare_and_dispatch_fwd, + quantize_for_ep, ) from ...quantization import QuantizerRole from ...tensor import MXFP8Quantizer, Quantizer from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .._common import ( maybe_dequantize, - quantize_for_ep, + validate_ep_buffer, + validate_ep_comms_recipe, ) from ..op import BasicOperation, OperationContext @@ -33,10 +36,10 @@ def _validate_dispatch_input( input_shape = tuple(input_.shape) if len(input_shape) != 2 or input_shape[-1] != buffer.hidden_dim: raise ValueError( - f"Dispatch input must have shape (T, {buffer.hidden_dim}), got {input_shape}." + f"MoeDispatch input must have shape (T, {buffer.hidden_dim}), got {input_shape}." ) if input_.device != buffer.device: - raise ValueError(f"Dispatch input must be on {buffer.device}, got {input_.device}.") + raise ValueError(f"MoeDispatch input must be on {buffer.device}, got {input_.device}.") return input_shape @@ -54,22 +57,22 @@ def _validate_routing_inputs( raise ValueError(f"{name} must be on {device}, got {tensor.device}.") -class Dispatch(BasicOperation): +class MoeDispatch(BasicOperation): """Dispatch floating-point or MXFP8 tokens to local experts with NCCL EP. The extra inputs are routing indices and FP32 routing weights. The extra - outputs are local tokens-per-expert, received routing weights, the routing - handle, and routing indices for recovering the local token shape. + outputs are local tokens-per-expert and received routing weights. """ num_extra_inputs: int = 2 - # tokens-per-expert, received routing weights, and the opaque NCCL EP - # routing handle and routing indices consumed by Combine. - num_extra_outputs: int = 4 + # tokens-per-expert and received routing weights consumed by the expert MLP. + num_extra_outputs: int = 2 - def __init__(self, buffer: EpBuffer) -> None: + def __init__(self, config: EpConfig) -> None: super().__init__() - self.buffer = buffer + if not isinstance(config, EpConfig): + raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + self.config = config def num_quantizers(self, mode: str) -> int: # quantized dispatch_bwd/combine is not supported. @@ -81,7 +84,7 @@ def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: return [ QuantizerRole( module_type="dispatch", - tensor_type="input", + tensor_type="dispatch_input", name=name, ) ] @@ -98,10 +101,10 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: quantizer.internal = True def op_forward(self, *args: Any, **kwargs: Any) -> None: - raise RuntimeError("Dispatch uses fuser_forward") + raise RuntimeError("MoeDispatch uses fuser_forward") def op_backward(self, *args: Any, **kwargs: Any) -> None: - raise RuntimeError("Dispatch uses fuser_backward") + raise RuntimeError("MoeDispatch uses fuser_backward") def fuser_forward( self, @@ -113,35 +116,37 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - + del next_op_input_quantizer # Dispatch uses unquantized transport without an input quantizer and # MXFP8 transport with an MXFP8 input quantizer. input_quantizer = self.get_quantizer("forward", 0) topk_idx, topk_weights = basic_op_extra_inputs[0] kwargs = basic_op_kwargs[0] - _validate_dispatch_input(input_, self.buffer) + buffer = validate_ep_buffer("MoeDispatch", self.config, kwargs.get("buffer")) + validate_ep_comms_recipe( + "MoeDispatch", + input_quantizer, + buffer.dispatch_fwd_quant_recipe, + ) + input_shape = _validate_dispatch_input(input_, buffer) + buffer.num_local_tokens = input_shape[0] _validate_routing_inputs( topk_idx, topk_weights, - device=self.buffer.device, + device=buffer.device, ) # Prepare the input input_scale_inv = None - if input_quantizer is None: - # Only BF16 dispatch is supported for now. - input_ = maybe_dequantize(input_, torch.bfloat16) - elif isinstance(input_quantizer, MXFP8Quantizer): + if isinstance(input_quantizer, MXFP8Quantizer): input_, input_scale_inv = quantize_for_ep(input_, input_quantizer) else: - raise TypeError( - "NCCL EP Dispatch supports MXFP8Quantizer only, got " - f"{type(input_quantizer).__name__}." - ) + # Only BF16 dispatch is supported for now. + input_ = maybe_dequantize(input_, torch.bfloat16) # Eager mode discovers the receive size at runtime, so persistent # caller-owned output buffers cannot be used. recv_tokens = kwargs.get("recv_tokens") recv_topk_weights = kwargs.get("recv_topk_weights") - if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): + if buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): raise ValueError( "eager mode sizes dispatch outputs per step and cannot use " "caller-supplied receive buffers" @@ -150,12 +155,12 @@ def fuser_forward( input_, topk_weights, topk_idx, - self.buffer, + buffer, recv_tokens, recv_topk_weights, input_scale_inv, ) - tokens_per_expert = self.buffer.tokens_per_expert + tokens_per_expert = buffer.tokens_per_expert # If next_op_input_quantizer is different from input_quantizer, # we need to requantize the data, which is handled in grouped_linear anyway. # We won't get any fusion benefit, so don't do it here. @@ -165,7 +170,7 @@ def fuser_forward( ctx.dispatch_state = dispatch_state ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer - return output, [(tokens_per_expert, recv_topk_weights, self.buffer.handle_mem, topk_idx)] + return output, [(tokens_per_expert, recv_topk_weights)] def fuser_backward( self, @@ -197,7 +202,4 @@ def fuser_backward( grad_output, grad_recv_weights, ) - quantizer = ctx.prev_op_grad_output_quantizer - if quantizer is not None: - grad_input = quantizer(grad_input) return grad_input, [()], [(None, grad_topk_weights)] diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 3921733422..8b3290b102 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -9,6 +9,7 @@ from collections.abc import Iterable, Sequence import functools from importlib.metadata import PackageNotFoundError, version as get_pkg_version +import os from typing import Any, Optional from packaging.version import Version as PkgVersion @@ -16,7 +17,7 @@ import transformer_engine_torch as tex from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE -from ...ep import get_ep_group +from ...ep import get_ep_group, quantize_for_ep from ...quantization import Recipe from ...tensor import GroupedTensor, MXFP8Quantizer, Quantizer from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -26,10 +27,9 @@ get_main_grad_from_param, is_quantized_tensor, maybe_dequantize, - quantize_for_ep, view_main_grad_as_grouped_buffer, ) -from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU +from ..basic import GroupedLinear, MoeCombine, MoeDispatch, ScaledSwiGLU from ..fuser import register_forward_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -42,6 +42,12 @@ def _cudnn_megamoe_supported() -> bool: return False +def _get_megamoe_combine_format() -> str: + """Return the MegaMoE combine wire format selected by the environment.""" + enabled = int(os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE", "0")) + return "mxfp8" if enabled > 0 else "bf16" + + def _pack_as_cudnn_moe_tensor( data: torch.Tensor, scale: Optional[torch.Tensor], @@ -262,11 +268,10 @@ def _import_cudnn_moe_ep(): def _routing_extras_internal( - dispatch: Dispatch, + dispatch: MoeDispatch, fc1: GroupedLinear, activation: ScaledSwiGLU, fc2: GroupedLinear, - combine: Combine, ) -> bool: """Whether the dispatch routing extras stay inside the fusion. @@ -275,13 +280,8 @@ def _routing_extras_internal( sequence when those two outputs feed exactly these ops and are not returned to the caller. """ - tokens_per_expert, routing_weights, ep_handle, routing_indices = dispatch._extra_output_channels - if ( - tokens_per_expert is None - or routing_weights is None - or ep_handle is None - or routing_indices is None - ): + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: return False if any(dispatch._extra_output_to_caller): return False @@ -289,13 +289,10 @@ def _routing_extras_internal( fc1._extra_input_channels[0] == tokens_per_expert and fc2._extra_input_channels[0] == tokens_per_expert and activation._extra_input_channels[0] == routing_weights - and combine._extra_input_channels[0] == ep_handle - and combine._extra_input_channels[1] == tokens_per_expert - and combine._extra_input_channels[2] == routing_indices ) -def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: +def _megamoe_supported(config, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: """Static MegaMoE capability gates that can be checked before first launch.""" if not _cudnn_megamoe_supported(): return False @@ -303,13 +300,13 @@ def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: return False if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): return False - if buffer.max_tokens_per_rank is None or buffer.max_tokens_per_rank <= 0: + if config.max_tokens_per_rank <= 0: return False - if buffer.recv_capacity_per_rank is not None and buffer.recv_capacity_per_rank <= 0: + if config.recv_capacity_per_rank is None or config.recv_capacity_per_rank <= 0: return False - if buffer.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: + if config.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: return False - if buffer.top_k > 32: + if config.top_k > 32: return False return True @@ -321,29 +318,29 @@ def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bo return False dispatch, fc1, activation, fc2, combine = window if not ( - isinstance(dispatch, Dispatch) + isinstance(dispatch, MoeDispatch) and isinstance(fc1, GroupedLinear) and isinstance(activation, ScaledSwiGLU) and isinstance(fc2, GroupedLinear) - and isinstance(combine, Combine) + and isinstance(combine, MoeCombine) ): return False - buffer = dispatch.buffer - if buffer.payload_dtype is not torch.bfloat16: + config = dispatch.config + if combine.config != config or config.payload_dtype is not torch.bfloat16: return False if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): return False if activation.activation_recompute_in_mlp or activation.glu_interleave_size is not None: return False - if not _routing_extras_internal(dispatch, fc1, activation, fc2, combine): + if not _routing_extras_internal(dispatch, fc1, activation, fc2): return False - if not _megamoe_supported(buffer, fc1, fc2): + if not _megamoe_supported(config, fc1, fc2): return False return ( - fc1.num_groups == buffer.num_local_experts - and fc2.num_groups == buffer.num_local_experts - and fc1.in_features == buffer.hidden_dim - and fc2.out_features == buffer.hidden_dim + fc1.num_groups == config.num_local_experts + and fc2.num_groups == config.num_local_experts + and fc1.in_features == config.hidden_dim + and fc2.out_features == config.hidden_dim and fc1.out_features == 2 * fc2.in_features ) @@ -354,11 +351,11 @@ class FusedMoeEp(FusedOperation): def __init__( self, *, - dispatch: Dispatch, + dispatch: MoeDispatch, fc1: GroupedLinear, activation: ScaledSwiGLU, fc2: GroupedLinear, - combine: Combine, + combine: MoeCombine, ) -> None: super().__init__([dispatch, fc1, activation, fc2, combine]) moe_ep_cls = _import_cudnn_moe_ep() @@ -372,27 +369,29 @@ def __init__( ep_group = get_ep_group() ep_size = 1 if ep_group is None else ep_group.size() + config = dispatch.config + combine_format = _get_megamoe_combine_format() self._block_scaled_cls = BlockScaledTensor self._moe = moe_ep_cls( - num_experts=dispatch.buffer.num_local_experts * ep_size, - hidden_size=dispatch.buffer.hidden_dim, + num_experts=config.num_local_experts * ep_size, + hidden_size=config.hidden_dim, intermediate_size=fc2.in_features, - top_k=dispatch.buffer.top_k, + top_k=config.top_k, ep_group=ep_group, - max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, - max_recv_size_per_rank=dispatch.buffer.recv_capacity_per_rank, - drop_on_overflow=False, + max_tokens_per_rank=config.max_tokens_per_rank, + max_recv_size_per_rank=config.recv_capacity_per_rank, + drop_on_overflow=config.drop_on_overflow, apply_topk_in_fc1=True, generate_c=True, backward_wgrad_mode="operands", token_padding_size=256, sf_padding_size=128, - combine_format="bf16", + combine_format=combine_format, output_format="bf16", ) @property - def dispatch(self) -> Dispatch: + def dispatch(self) -> MoeDispatch: return self.basic_ops[0] @property @@ -428,8 +427,6 @@ def fuser_forward( raise NotImplementedError( f"FusedMoeEp supports E4M3 MXFP8 only, got {input_._fp8_dtype}." ) - if any(kwargs for kwargs in basic_op_kwargs): - raise NotImplementedError("FusedMoeEp does not support per-operation output buffers.") topk_idx, topk_weights = basic_op_extra_inputs[0] if topk_weights.dtype is not torch.float32: @@ -478,7 +475,7 @@ def fuser_forward( if next_op_input_quantizer is not None and not is_quantized_tensor(output): output = next_op_input_quantizer(output) return output, [ - (None, None, None, None), + (None, None), (), (), (), @@ -547,7 +544,7 @@ def fuser_backward( return ( grad_input, [(), fc1_param_grads, (), fc2_param_grads, ()], - [(None, grad_topk_weights.float()), (None,), (None,), (None,), (None, None, None)], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), ()], ) From a584be88e46943c3cb1259b77f8c086893d2b240 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 18:28:35 +0000 Subject: [PATCH 71/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ep.py | 6 +++--- transformer_engine/pytorch/ops/basic/combine.py | 4 +++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 2df71f334f..3969b3026c 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -779,6 +779,7 @@ def forward( # type: ignore[override] tokens_scale_inv = None if buffer.dispatch_fwd_quant_recipe is not None: from .tensor.mxfp8_tensor import MXFP8Quantizer + # Only MXFP8 Quantizer is supported for EP dispatch quantizer = MXFP8Quantizer( tex.DType.kFloat8E4M3, @@ -905,6 +906,7 @@ def _ep_combine_bwd( else: if quantized_grad is None: from .tensor.mxfp8_tensor import MXFP8Quantizer + # Only MXFP8 Quantizer is supported for EP combine bwd quantizer = MXFP8Quantizer( tex.DType.kFloat8E4M3, @@ -1060,9 +1062,7 @@ def quantize_for_ep( if isinstance(input_, MXFP8TensorStorage): quantized = input_ elif isinstance(input_, QuantizedTensorStorage): - raise TypeError( - f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}." - ) + raise TypeError(f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}.") else: if quantizer is None: raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index e52040f701..d34491a45d 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -45,7 +45,9 @@ def _validate_combine_inputs( buffer.tokens_per_expert.dtype is not torch.int64 or buffer.tokens_per_expert.device != input_.device ): - raise ValueError("MoeCombine tokens_per_expert must be an int64 tensor on the input device.") + raise ValueError( + "MoeCombine tokens_per_expert must be an int64 tensor on the input device." + ) return tuple(input_.shape) From bf9dd3f0013e6eded4d4d677756a6aaede46327c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 18:43:40 +0000 Subject: [PATCH 72/83] address review comments Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/basic/combine.py | 2 ++ transformer_engine/pytorch/ops/basic/dispatch.py | 3 +-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index e52040f701..38bf4338ff 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -148,6 +148,8 @@ def fuser_forward( grad_output_quantizer, buffer.combine_bwd_quant_recipe, ) + # Only BF16 combine forward is supported for now. + input_ = maybe_dequantize(input_, torch.bfloat16) input_shape = _validate_combine_inputs(input_, buffer) # Stage zero-copy input if needed, then restore local-token order. zero_copy = buffer.zero_copy diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 5fe06640fe..661d77ff76 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -166,7 +166,6 @@ def fuser_forward( # We won't get any fusion benefit, so don't do it here. ctx = basic_op_ctxs[0] if ctx.requires_grad: - ctx.input_dtype = torch.bfloat16 ctx.dispatch_state = dispatch_state ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer @@ -185,7 +184,7 @@ def fuser_backward( ]: ctx = basic_op_ctxs[0] # Only BF16 Dispatch_bwd is supported for now. - grad_output = maybe_dequantize(grad_output, ctx.input_dtype) + grad_output = maybe_dequantize(grad_output, torch.bfloat16) grad_recv_weights = basic_op_grad_extra_outputs[0][1] if grad_recv_weights is None: From 1d1e11acb16b533545379acd965c55732151f355 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 18:46:17 +0000 Subject: [PATCH 73/83] t_flat was a better name Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ep.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 3969b3026c..a468aa8aaf 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -1076,7 +1076,7 @@ def quantize_for_ep( scale_inv = quantized._rowwise_scale_inv if data is None or scale_inv is None: raise ValueError("EP requires rowwise MXFP8 data and scales.") - rows, hidden = input_.shape + t_flat, hidden = input_.shape scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE # The backend forwards each token's scale row with a 16-byte-aligned store, so the row # (cols * dtype bytes) must be a multiple of 16. @@ -1090,7 +1090,7 @@ def quantize_for_ep( # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block] # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the # slice stays contiguous; assert rather than force a copy. - scale_inv = scale_inv[:rows, :scale_cols] + scale_inv = scale_inv[:t_flat, :scale_cols] if not scale_inv.is_contiguous(): raise ValueError("EP requires compact contiguous MXFP8 scales.") return quantized, scale_inv From 5674825692b2eeab20bf668c0d510a24b9840f23 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 18:50:06 +0000 Subject: [PATCH 74/83] simplify Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ep.py | 23 ++++------------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index a468aa8aaf..57445863d9 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -1025,24 +1025,6 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso return torch.empty(*shape, dtype=dtype, device=device) -def _as_mxfp8_storage(tensor: QuantizedTensorStorage): - """Return a lightweight MXFP8 storage view without a ``torch.Tensor`` wrapper.""" - if type(tensor) is MXFP8TensorStorage: - return tensor - if not isinstance(tensor, MXFP8TensorStorage): - raise TypeError(f"Expected MXFP8 tensor storage, got {type(tensor).__name__}.") - return MXFP8TensorStorage( - rowwise_data=tensor._rowwise_data, - rowwise_scale_inv=tensor._rowwise_scale_inv, - columnwise_data=tensor._columnwise_data, - columnwise_scale_inv=tensor._columnwise_scale_inv, - fp8_dtype=tensor._fp8_dtype, - quantizer=tensor._quantizer, - with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, - fake_dtype=tensor._dtype, - ) - - def quantize_for_ep( input_: torch.Tensor | QuantizedTensorStorage, quantizer: Optional["Quantizer"], @@ -1066,8 +1048,11 @@ def quantize_for_ep( else: if quantizer is None: raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") + if not quantized.internal: + quantizer = quantizer.copy() + quantizer.internal = True quantized = quantizer(input_) - quantized = _as_mxfp8_storage(quantized) + if quantized._fp8_dtype != DType.kFloat8E4M3: raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") if quantized._with_gemm_swizzled_scales: From 6646ca3fe3c89beb58736fc60d25a0bf15a48a25 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 18:52:00 +0000 Subject: [PATCH 75/83] clean cursor mess Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ep.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 57445863d9..c0d2594c96 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -853,12 +853,10 @@ def _ep_combine_fwd( eager: bool, zero_copy: bool, ): - """Run combine from explicit routing state and return ``(result, _CombineState)``. + """Run combine and return ``(result, _CombineState)``. Eager mode is not graph-capturable, so + it calls the backend op directly and skips the torch.library dispatch. No autograd; the caller + owns context handling.""" - This is the shared implementation for the public ``ep_combine`` wrapper and - fusible operations. Eager mode calls the backend directly to avoid the - ``torch.library`` dispatch overhead. - """ device = expert_out.device result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) if eager: From 62d9b0662eb85af9e6990366fa302c5546dbe4c2 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 18:57:55 +0000 Subject: [PATCH 76/83] restore comments Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ep.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index c0d2594c96..e918e219c1 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -885,8 +885,7 @@ def _ep_combine_bwd( ): """Scatter the result-grad to expert positions and return the expert_out grad. High-precision sends the grad as-is; a quantized recipe (MXFP8 today) quantizes it and returns a per-expert - GroupedTensor. A fusible-op caller may instead provide its fuser-quantized grad and compact - scales. No autograd; the caller owns context handling.""" + GroupedTensor. Optionally pre-mxfp8-quantized grad and scales can be provided.""" if quantized_grad is None and not g_result.is_contiguous(): g_result = g_result.contiguous() handle_mem = state.handle_mem @@ -1027,7 +1026,14 @@ def quantize_for_ep( input_: torch.Tensor | QuantizedTensorStorage, quantizer: Optional["Quantizer"], ) -> tuple[MXFP8TensorStorage, torch.Tensor]: - """Return an MXFP8 input and its compact rowwise scales for EP.""" + """Return E4M3 MXFP8 storage and compact rowwise scales for EP transport. + + High-precision input is quantized with ``quantizer``; existing MXFP8 storage is accepted + as-is. The returned storage owns the FP8 payload in ``_rowwise_data``, while ``scale_inv`` has + the compact ``[T, H/block]`` layout routed by the EP backend. EP transports E4M3 in both + directions, so other FP8 formats are rejected. GEMM scale-row padding is stripped, and each + compact scale row must remain contiguous and 16-byte aligned. + """ from .constants import DType, MXFP8_BLOCK_SCALING_SIZE from .tensor.mxfp8_tensor import MXFP8Quantizer @@ -1046,7 +1052,7 @@ def quantize_for_ep( else: if quantizer is None: raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") - if not quantized.internal: + if not quantizer.internal: quantizer = quantizer.copy() quantizer.internal = True quantized = quantizer(input_) From e8d1bb729d3ad188b0cef7e20959653b75819e48 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 19:03:36 +0000 Subject: [PATCH 77/83] bf16 hardcoding for combine as well Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/basic/combine.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 7c562cfd60..fbe86300dd 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -157,7 +157,7 @@ def fuser_forward( zero_copy = buffer.zero_copy expert_out = input_ if zero_copy: - expert_out = _alloc_io(tuple(input_.shape), input_.dtype, input_.device, True) + expert_out = _alloc_io(tuple(input_.shape), torch.bfloat16, input_.device, True) expert_out.copy_(input_) # Preserve routing state and optional caller storage for backward. ctx = basic_op_ctxs[0] @@ -167,11 +167,10 @@ def fuser_forward( kwargs.get("grad_out"), transport_quantizer, input_shape=input_shape, - input_dtype=input_.dtype, + input_dtype=torch.bfloat16, device=input_.device, zero_copy=zero_copy, ) - ctx.input_dtype = input_.dtype result, combine_state = _ep_combine_fwd( expert_out, grad_out, @@ -211,7 +210,7 @@ def fuser_backward( ) grad_output = quantized_grad else: - grad_output = maybe_dequantize(grad_output, ctx.input_dtype).contiguous() + grad_output = maybe_dequantize(grad_output, torch.bfloat16).contiguous() quantized_grad = None grad_input = _ep_combine_bwd( ctx.combine_state, From cb4c57ea1f56fe3f4abb5cade3a8228d33e63242 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 19:07:16 +0000 Subject: [PATCH 78/83] scale type should be autocast type if autocast enabled Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/basic/swiglu.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 9824aab07c..b2898a8d74 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -527,10 +527,11 @@ def fuser_forward( ) extra_input = basic_op_extra_inputs[0][0] - + scale_dtype = None # Determine compute dtype if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") + scale_dtype = dtype elif isinstance(input_, torch.Tensor): dtype = input_.dtype else: @@ -538,7 +539,7 @@ def fuser_forward( # Prepare the activation input in the compute dtype. input_ = maybe_dequantize(input_, dtype) - scales = extra_input + scales = extra_input if scale_dtype is None else maybe_dequantize(extra_input, scale_dtype) out = self._scaled_glu_forward(input_, scales) # Save state for backward pass From ebed56a932c949cec53b737dceb78387c1dbd629 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 19:11:41 +0000 Subject: [PATCH 79/83] update copyrights Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/basic/combine.py | 2 +- transformer_engine/pytorch/ops/basic/dispatch.py | 2 +- transformer_engine/pytorch/ops/fused/moe_ep.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index fbe86300dd..3d1ca4c132 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 661d77ff76..06c8800e8c 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index 8b3290b102..f5b64118e8 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. From a44091ead68f54a058d2a2076e722acfc3af8c77 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 28 Aug 2026 19:18:11 +0000 Subject: [PATCH 80/83] match the right syntax Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ep_reference.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py index 981b04e1cf..8614b762d2 100644 --- a/transformer_engine/pytorch/ep_reference.py +++ b/transformer_engine/pytorch/ep_reference.py @@ -1,5 +1,6 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. """Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. From ae6dc49848a4dad26d524c92f981f351bbef0a52 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sat, 29 Aug 2026 00:43:17 +0000 Subject: [PATCH 81/83] better changes Signed-off-by: Varun Thumbe --- examples/pytorch/ep/bench/ep_bench.py | 4 +- examples/pytorch/ep/ep_moe.py | 5 +- tests/pytorch/distributed/run_ep.py | 136 ++++++++++++++---- transformer_engine/pytorch/ep.py | 108 +++++++------- transformer_engine/pytorch/ops/_common.py | 40 +++++- .../pytorch/ops/basic/combine.py | 63 ++------ .../pytorch/ops/basic/dispatch.py | 16 +-- .../pytorch/ops/fused/moe_ep.py | 9 +- 8 files changed, 212 insertions(+), 169 deletions(-) diff --git a/examples/pytorch/ep/bench/ep_bench.py b/examples/pytorch/ep/bench/ep_bench.py index 7091422894..f80a57015c 100644 --- a/examples/pytorch/ep/bench/ep_bench.py +++ b/examples/pytorch/ep/bench/ep_bench.py @@ -30,7 +30,6 @@ from transformer_engine.pytorch.ep import ( EpBuffer, - EpConfig, ep_bootstrap, ep_combine, ep_dispatch, @@ -203,14 +202,13 @@ def main(): else None ) - config = EpConfig( + buffer = EpBuffer( top_k=K, max_tokens_per_rank=T, recv_capacity_per_rank=recv_pr, hidden_dim=H, num_local_experts=num_local_experts, ) - buffer = EpBuffer(config) tokens = tokens_hbm topk_w = topk_w_hbm diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py index 45993aa4cc..084496da48 100644 --- a/examples/pytorch/ep/ep_moe.py +++ b/examples/pytorch/ep/ep_moe.py @@ -16,7 +16,6 @@ from transformer_engine.pytorch.ep import ( EpBuffer, - EpConfig, ep_bootstrap, ep_combine, ep_dispatch, @@ -178,14 +177,13 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, else None ) - config = EpConfig( + buffer = EpBuffer( top_k=args.top_k, max_tokens_per_rank=T, recv_capacity_per_rank=recv_pr, hidden_dim=args.hidden, num_local_experts=num_local_experts, ) - buffer = EpBuffer(config) recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w, recv_tokens=recv_tokens) expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts) @@ -254,3 +252,4 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, if __name__ == "__main__": main() sys.exit(0) + diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6455bbc5a1..bca61b2db2 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -21,6 +21,8 @@ EpConfig, ep_bootstrap, ep_finalize, + get_ep_drop_on_overflow, + get_ep_group, ep_prepare, ep_dispatch, ep_combine, @@ -286,24 +288,51 @@ def setUp(self): ): self.skipTest("not exercised in overflow mode") - def _make_buffer( + def _make_config( self, alignment=0, top_k=TOP_K, - dispatch_fwd_quant_recipe=None, - combine_bwd_quant_recipe=None, ): - config = EpConfig( + return EpConfig( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, + ep_group=self.ep_group, alignment=alignment, zero_copy=ZERO_COPY, drop_on_overflow=OVERFLOW, ) + + def _make_buffer_from_config( + self, + config, + *, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, + ): return EpBuffer( + top_k=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + hidden_dim=config.hidden_dim, + num_local_experts=config.num_local_experts, + recv_capacity_per_rank=config.recv_capacity_per_rank, + alignment=config.alignment, + payload_dtype=config.payload_dtype, + dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, + combine_bwd_quant_recipe=combine_bwd_quant_recipe, + ) + + def _make_buffer( + self, + alignment=0, + top_k=TOP_K, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, + ): + config = self._make_config(alignment=alignment, top_k=top_k) + return self._make_buffer_from_config( config, dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, combine_bwd_quant_recipe=combine_bwd_quant_recipe, @@ -313,6 +342,10 @@ def _make_buffer( class TestEP(_EpTestCase): """NCCL EP tests inherited from nvidia_origin/main.""" + def test_bootstrap_accessors(self): + self.assertIs(get_ep_group(), self.ep_group) + self.assertEqual(get_ep_drop_on_overflow(), OVERFLOW) + def _expert_out(self, expert_out): """Stage the combine input into symm-mem under zero-copy (combine requires it).""" if not ZERO_COPY: @@ -920,19 +953,53 @@ def _require_mxfp8_shapes(self): ) def test_runtime_buffer_config_mismatch(self): - buffer = self._make_buffer() - wrong_config = replace(buffer.config, hidden_dim=buffer.config.hidden_dim + 1) + config = self._make_config() + buffer = self._make_buffer_from_config(config) topk_idx, tokens, topk_weights = _make_identity_inputs( self.cfg.rank, self.cfg.ep_size, ) - with self.assertRaisesRegex(ValueError, "runtime buffer config"): + dispatch = te_ops.MoeDispatch(config) + replacements = { + "top_k": config.top_k + 1, + "hidden_dim": config.hidden_dim + 1, + "num_local_experts": config.num_local_experts + 1, + "max_tokens_per_rank": config.max_tokens_per_rank + 1, + "recv_capacity_per_rank": config.recv_capacity_per_rank + 1, + "alignment": 1, + "payload_dtype": torch.float16, + "zero_copy": not config.zero_copy, + } + for field_name, wrong_value in replacements.items(): + with self.subTest(field_name=field_name): + original = getattr(buffer, field_name) + setattr(buffer, field_name, wrong_value) + try: + with self.assertRaisesRegex(ValueError, field_name): + dispatch( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) + finally: + setattr(buffer, field_name, original) + + wrong_config = replace(config, drop_on_overflow=not config.drop_on_overflow) + with self.assertRaisesRegex(ValueError, "drop_on_overflow"): te_ops.MoeDispatch(wrong_config)( tokens, topk_idx, topk_weights, buffer=buffer, ) + with self.assertRaisesRegex(ValueError, "ep_group"): + te_ops.MoeDispatch(replace(config, ep_group=None))( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) expert_out = torch.empty( self.cfg.recv_capacity_per_rank, HIDDEN_DIM, @@ -940,7 +1007,12 @@ def test_runtime_buffer_config_mismatch(self): device=self.cfg.device, ) with self.assertRaisesRegex(ValueError, "runtime buffer config"): - te_ops.MoeCombine(wrong_config)(expert_out, buffer=buffer) + original = buffer.hidden_dim + buffer.hidden_dim += 1 + try: + te_ops.MoeCombine(config)(expert_out, buffer=buffer) + finally: + buffer.hidden_dim = original def test_megamoe_combine_format_env(self): old_value = os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE") @@ -958,9 +1030,10 @@ def test_megamoe_combine_format_env(self): @_mxfp8_align_test def test_role_quantizer_requires_matching_buffer_recipe(self): self._require_mxfp8_shapes() - buffer = self._make_buffer(alignment=128) - dispatch = te_ops.MoeDispatch(buffer.config) - combine = te_ops.MoeCombine(buffer.config) + config = self._make_config(alignment=128) + buffer = self._make_buffer_from_config(config) + dispatch = te_ops.MoeDispatch(config) + combine = te_ops.MoeCombine(config) topk_idx, tokens, topk_weights = _make_identity_inputs( self.cfg.rank, self.cfg.ep_size, @@ -981,15 +1054,16 @@ def test_role_quantizer_requires_matching_buffer_recipe(self): def _make_dispatch_combine_ops(self, *, mxfp8): recipe = MXFP8BlockScaling() if mxfp8 else None - buffer = self._make_buffer( - alignment=128 if mxfp8 else 0, + config = self._make_config(alignment=128 if mxfp8 else 0) + buffer = self._make_buffer_from_config( + config, dispatch_fwd_quant_recipe=recipe, combine_bwd_quant_recipe=recipe, ) return ( buffer, - te_ops.MoeDispatch(buffer.config), - te_ops.MoeCombine(buffer.config), + te_ops.MoeDispatch(config), + te_ops.MoeCombine(config), ) def _run_dispatch_combine_identity(self, *, mxfp8): @@ -1009,28 +1083,26 @@ def _run_dispatch_combine_identity(self, *, mxfp8): topk_weights, buffer=buffer, ) - if mxfp8: - recv_tokens = _degroup_mxfp8(recv_tokens) - recv_weights = recv_weights[: recv_tokens.shape[0]] - weighted_expert_output = (recv_tokens.float() * recv_weights.float().unsqueeze(-1)).to( - torch.bfloat16 - ) - output = combine( - weighted_expert_output, - buffer=buffer, - ) + if mxfp8: + recv_tokens = _degroup_mxfp8(recv_tokens) + recv_weights = recv_weights[: recv_tokens.shape[0]] + weighted_expert_output = (recv_tokens.float() * recv_weights.float().unsqueeze(-1)).to( + torch.bfloat16 + ) + output = combine( + weighted_expert_output, + buffer=buffer, + ) torch.cuda.synchronize() torch.testing.assert_close(output, tokens, atol=5e-2, rtol=5e-2) self.assertEqual(tokens_per_expert.data_ptr(), buffer.tokens_per_expert.data_ptr()) @_eager_test_include - @_zero_copy_test_include def test_dispatch_combine_identity_bf16(self): """MoeDispatch and MoeCombine basic ops form an identity in BF16.""" self._run_dispatch_combine_identity(mxfp8=False) @_eager_test_include - @_zero_copy_test_include @_mxfp8_align_test def test_dispatch_combine_identity_mxfp8(self): """MoeDispatch and MoeCombine basic ops form an identity with MXFP8 transport.""" @@ -1044,12 +1116,13 @@ def _make_megamoe_model( delay_wgrad_compute=False, ): """Build the exact five-op sequence recognized by MegaMoE fusion.""" - buffer = self._make_buffer( - alignment=128 if recipe is not None else 0, + config = self._make_config(alignment=128 if recipe is not None else 0) + buffer = self._make_buffer_from_config( + config, dispatch_fwd_quant_recipe=recipe, combine_bwd_quant_recipe=recipe, ) - dispatch = te_ops.MoeDispatch(buffer.config) + dispatch = te_ops.MoeDispatch(config) init_ctx = ( te.quantized_model_init(enabled=True, recipe=recipe) if recipe is not None @@ -1087,7 +1160,7 @@ def _make_megamoe_model( del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] else: os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param - combine = te_ops.MoeCombine(buffer.config) + combine = te_ops.MoeCombine(config) dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) @@ -1413,3 +1486,4 @@ def _init_distributed(): release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) + diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index e918e219c1..cb5f79a4be 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -30,6 +30,7 @@ "EpConfig", "EpBuffer", "ep_bootstrap", + "get_ep_drop_on_overflow", "get_ep_group", "is_ep_bootstrapped", "ep_finalize", @@ -88,15 +89,16 @@ def _check_nccl_runtime_version() -> None: _BOOTSTRAP_SETTINGS: Optional[dict[str, object]] = None -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class EpConfig: - """Immutable configuration shared by an EP buffer and its MoE operations.""" + """Immutable configuration shared by EP MoE operations.""" top_k: int hidden_dim: int num_local_experts: int max_tokens_per_rank: int recv_capacity_per_rank: Optional[int] + ep_group: dist.ProcessGroup alignment: int = 0 payload_dtype: torch.dtype = torch.bfloat16 zero_copy: bool = False @@ -227,6 +229,13 @@ def get_ep_group() -> Optional[dist.ProcessGroup]: return _EP_GROUP +def get_ep_drop_on_overflow() -> Optional[bool]: + """Return the overflow policy registered by :func:`ep_bootstrap`.""" + if _BOOTSTRAP_SETTINGS is None: + return None + return bool(_BOOTSTRAP_SETTINGS["drop_on_overflow"]) + + def _ep_is_eager() -> bool: """Return whether the bootstrapped EP group uses variable-size eager buffers.""" return _EAGER @@ -280,17 +289,24 @@ def is_symm_backed(t: torch.Tensor) -> bool: class EpBuffer: - """Per-microbatch EP routing state allocated from an immutable :class:`EpConfig`. + """Per-microbatch EP layer state: handle_mem, tokens_per_expert, and shape/dtype config. Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). """ __slots__ = ( - "config", "handle_mem", + "top_k", + "alignment", + "max_tokens_per_rank", + "recv_capacity_per_rank", + "hidden_dim", + "num_local_experts", + "payload_dtype", "device", "tokens_per_expert", "num_local_tokens", + "zero_copy", "eager", "total_recv_tokens", "dispatch_fwd_quant_recipe", @@ -299,43 +315,50 @@ class EpBuffer: def __init__( self, - config: EpConfig, - *, + top_k: int, + max_tokens_per_rank: int, + hidden_dim: int, + num_local_experts: int, + recv_capacity_per_rank: Optional[int] = None, + alignment: int = 0, + payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, dispatch_fwd_quant_recipe: Optional["Recipe"] = None, combine_bwd_quant_recipe: Optional["Recipe"] = None, ) -> None: if not _BOOTSTRAPPED: raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") - if not isinstance(config, EpConfig): - raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") - if _BOOTSTRAP_SETTINGS is None: - raise RuntimeError("EP bootstrap configuration is unavailable.") - mismatches = { - name: (getattr(config, name), expected) - for name, expected in _BOOTSTRAP_SETTINGS.items() - if getattr(config, name) != expected - } - if mismatches: - details = ", ".join( - f"{name}={actual!r} (bootstrapped {expected!r})" - for name, (actual, expected) in mismatches.items() - ) - raise ValueError(f"EpConfig does not match ep_bootstrap: {details}.") if device is None: device = torch.device("cuda", torch.cuda.current_device()) - self.config = config - self.eager = config.recv_capacity_per_rank is None + alignment = int(alignment) + if alignment > 1 and (alignment & (alignment - 1)) != 0: + raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") + self.eager = _EAGER + if not self.eager and recv_capacity_per_rank is None: + raise ValueError( + "EpBuffer requires recv_capacity_per_rank unless the EP group was " + "bootstrapped in eager mode (recv_capacity_per_rank omitted)." + ) + self.top_k = int(top_k) + self.alignment = alignment + self.max_tokens_per_rank = int(max_tokens_per_rank) + self.recv_capacity_per_rank = ( + None if recv_capacity_per_rank is None else int(recv_capacity_per_rank) + ) + self.hidden_dim = int(hidden_dim) + self.num_local_experts = int(num_local_experts) + self.payload_dtype = payload_dtype self.device = device + self.zero_copy = bool(tex.ep_get_zero_copy()) self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe self.combine_bwd_quant_recipe = combine_bwd_quant_recipe - size_bytes = tex.ep_handle_mem_size(config.top_k, config.alignment) + size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) self.tokens_per_expert = torch.empty( - config.num_local_experts, dtype=torch.int64, device=device + self.num_local_experts, dtype=torch.int64, device=device ) - self.num_local_tokens = config.max_tokens_per_rank + self.num_local_tokens = self.max_tokens_per_rank # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) # Per-step recv-token total (int64 [1]), written by ep_prepare. Eager reads it @@ -350,38 +373,6 @@ def __init__( self.total_recv_tokens = torch.empty(1, dtype=torch.int64, device=device) mark_not_offload(self.total_recv_tokens) - @property - def top_k(self) -> int: - return self.config.top_k - - @property - def alignment(self) -> int: - return self.config.alignment - - @property - def max_tokens_per_rank(self) -> int: - return self.config.max_tokens_per_rank - - @property - def recv_capacity_per_rank(self) -> Optional[int]: - return self.config.recv_capacity_per_rank - - @property - def hidden_dim(self) -> int: - return self.config.hidden_dim - - @property - def num_local_experts(self) -> int: - return self.config.num_local_experts - - @property - def payload_dtype(self) -> torch.dtype: - return self.config.payload_dtype - - @property - def zero_copy(self) -> bool: - return self.config.zero_copy - # torch.library custom ops (so they don't graph-break under torch.compile) @@ -1249,3 +1240,4 @@ def ep_combine( num_local_tokens, bwd_quant_recipe, ) + diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 546f65e1f8..dea22d170c 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,7 +11,12 @@ import torch from transformer_engine_torch import FP8TensorMeta -from ..ep import EpBuffer, EpConfig +from ..ep import ( + EpBuffer, + EpConfig, + get_ep_drop_on_overflow, + get_ep_group, +) from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..quantized_tensor import QuantizedTensorStorage, Quantizer @@ -76,10 +81,37 @@ def validate_ep_buffer( """Validate a runtime EP buffer against an operation's immutable config.""" if not isinstance(buffer, EpBuffer): raise TypeError(f"{op_name} requires buffer=EpBuffer(...), got {type(buffer).__name__}.") - if buffer.config != expected_config: + + mismatches = { + name: (getattr(buffer, name), getattr(expected_config, name)) + for name in ( + "top_k", + "hidden_dim", + "num_local_experts", + "max_tokens_per_rank", + "recv_capacity_per_rank", + "alignment", + "payload_dtype", + "zero_copy", + ) + if getattr(buffer, name) != getattr(expected_config, name) + } + ep_group = get_ep_group() + if expected_config.ep_group is not ep_group: + mismatches["ep_group"] = (ep_group, expected_config.ep_group) + drop_on_overflow = get_ep_drop_on_overflow() + if expected_config.drop_on_overflow != drop_on_overflow: + mismatches["drop_on_overflow"] = ( + drop_on_overflow, + expected_config.drop_on_overflow, + ) + if mismatches: + details = ", ".join( + f"{name}={actual!r} (expected {expected!r})" + for name, (actual, expected) in mismatches.items() + ) raise ValueError( - f"{op_name} runtime buffer config does not match its initialized config: " - f"{buffer.config!r} != {expected_config!r}." + f"{op_name} runtime buffer config does not match its initialized config: {details}." ) return buffer diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 3d1ca4c132..8a8dd6d244 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -13,17 +13,14 @@ from ...ep import ( EpBuffer, EpConfig, - _alloc_io, _ep_combine_bwd, _ep_combine_fwd, - is_symm_backed, quantize_for_ep, ) from ...quantization import QuantizerRole from ...tensor import MXFP8Quantizer, Quantizer from .._common import ( maybe_dequantize, - validate_buffer, validate_ep_buffer, validate_ep_comms_recipe, ) @@ -63,6 +60,8 @@ def __init__(self, config: EpConfig) -> None: super().__init__() if not isinstance(config, EpConfig): raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + if config.zero_copy: + raise NotImplementedError("MoeCombine does not support zero-copy EP.") self.config = config def num_quantizers(self, mode: str) -> int: @@ -95,36 +94,6 @@ def op_forward(self, *args: Any, **kwargs: Any) -> None: def op_backward(self, *args: Any, **kwargs: Any) -> None: raise RuntimeError("MoeCombine uses fuser_backward") - @staticmethod - def _prepare_grad_buffer( - grad_out: Optional[torch.Tensor], - grad_output_quantizer: Optional[Quantizer], - *, - input_shape: tuple[int, int], - input_dtype: torch.dtype, - device: torch.device, - zero_copy: bool, - ) -> Optional[torch.Tensor]: - """Validate caller storage for the expert-output gradient.""" - if not isinstance(grad_output_quantizer, MXFP8Quantizer): - grad_out = validate_buffer( - "grad_out", - grad_out, - shape=input_shape, - dtype=input_dtype, - device=device, - ) - else: - grad_out = validate_buffer( - "MXFP8 grad_out storage", - grad_out, - device=device, - contiguous=True, - ) - if zero_copy and grad_out is not None and not is_symm_backed(grad_out): - raise ValueError("zero-copy MoeCombine grad_out must be symmetric-memory-backed.") - return grad_out - def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -152,35 +121,18 @@ def fuser_forward( ) # Only BF16 combine forward is supported for now. input_ = maybe_dequantize(input_, torch.bfloat16) - input_shape = _validate_combine_inputs(input_, buffer) - # Stage zero-copy input if needed, then restore local-token order. - zero_copy = buffer.zero_copy - expert_out = input_ - if zero_copy: - expert_out = _alloc_io(tuple(input_.shape), torch.bfloat16, input_.device, True) - expert_out.copy_(input_) - # Preserve routing state and optional caller storage for backward. + _validate_combine_inputs(input_, buffer) ctx = basic_op_ctxs[0] - grad_out = None - if ctx.requires_grad: - grad_out = self._prepare_grad_buffer( - kwargs.get("grad_out"), - transport_quantizer, - input_shape=input_shape, - input_dtype=torch.bfloat16, - device=input_.device, - zero_copy=zero_copy, - ) result, combine_state = _ep_combine_fwd( - expert_out, - grad_out, + input_, + None, handle_mem=buffer.handle_mem, token_counts=buffer.tokens_per_expert, num_local_tokens=buffer.num_local_tokens, - hidden_dim=expert_out.shape[-1], + hidden_dim=input_.shape[-1], bwd_quant_recipe=transport_quantizer, eager=buffer.eager, - zero_copy=zero_copy, + zero_copy=False, ) if ctx.requires_grad: ctx.combine_state = combine_state @@ -219,3 +171,4 @@ def fuser_backward( grad_scale_inv, ) return grad_input, [()], [()] + diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 06c8800e8c..6332960ccf 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -72,6 +72,8 @@ def __init__(self, config: EpConfig) -> None: super().__init__() if not isinstance(config, EpConfig): raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + if config.zero_copy: + raise NotImplementedError("MoeDispatch does not support zero-copy EP.") self.config = config def num_quantizers(self, mode: str) -> int: @@ -142,22 +144,13 @@ def fuser_forward( else: # Only BF16 dispatch is supported for now. input_ = maybe_dequantize(input_, torch.bfloat16) - # Eager mode discovers the receive size at runtime, so persistent - # caller-owned output buffers cannot be used. - recv_tokens = kwargs.get("recv_tokens") - recv_topk_weights = kwargs.get("recv_topk_weights") - if buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): - raise ValueError( - "eager mode sizes dispatch outputs per step and cannot use " - "caller-supplied receive buffers" - ) output, recv_topk_weights, dispatch_state = _ep_prepare_and_dispatch_fwd( input_, topk_weights, topk_idx, buffer, - recv_tokens, - recv_topk_weights, + None, + None, input_scale_inv, ) tokens_per_expert = buffer.tokens_per_expert @@ -202,3 +195,4 @@ def fuser_backward( grad_recv_weights, ) return grad_input, [()], [(None, grad_topk_weights)] + diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index f5b64118e8..c452519d42 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -17,7 +17,7 @@ import transformer_engine_torch as tex from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE -from ...ep import get_ep_group, quantize_for_ep +from ...ep import quantize_for_ep from ...quantization import Recipe from ...tensor import GroupedTensor, MXFP8Quantizer, Quantizer from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -37,7 +37,7 @@ def _cudnn_megamoe_supported() -> bool: """Whether the installed cuDNN frontend includes the public MegaMoE API.""" try: - return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.28.0") + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.29.0") except PackageNotFoundError: return False @@ -367,9 +367,9 @@ def __init__( ) from cudnn.moe_ep import BlockScaledTensor - ep_group = get_ep_group() - ep_size = 1 if ep_group is None else ep_group.size() config = dispatch.config + ep_group = config.ep_group + ep_size = ep_group.size() combine_format = _get_megamoe_combine_format() self._block_scaled_cls = BlockScaledTensor self._moe = moe_ep_cls( @@ -586,3 +586,4 @@ def fuse_ops( __all__ = ["FusedMoeEp"] + From dc59ae3466de7fa8df6c6d79296a985d7baf577e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:44:26 +0000 Subject: [PATCH 82/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/pytorch/ep/ep_moe.py | 1 - tests/pytorch/distributed/run_ep.py | 1 - transformer_engine/pytorch/ep.py | 1 - transformer_engine/pytorch/ops/basic/combine.py | 1 - transformer_engine/pytorch/ops/basic/dispatch.py | 1 - transformer_engine/pytorch/ops/fused/moe_ep.py | 1 - 6 files changed, 6 deletions(-) diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py index 084496da48..b47c334225 100644 --- a/examples/pytorch/ep/ep_moe.py +++ b/examples/pytorch/ep/ep_moe.py @@ -252,4 +252,3 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, if __name__ == "__main__": main() sys.exit(0) - diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index bca61b2db2..44e96da124 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -1486,4 +1486,3 @@ def _init_distributed(): release_symm_mem_pool() dist.destroy_process_group() sys.exit(0 if result.wasSuccessful() else 1) - diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index cb5f79a4be..34a3b0df68 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -1240,4 +1240,3 @@ def ep_combine( num_local_tokens, bwd_quant_recipe, ) - diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py index 8a8dd6d244..1ad8602803 100644 --- a/transformer_engine/pytorch/ops/basic/combine.py +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -171,4 +171,3 @@ def fuser_backward( grad_scale_inv, ) return grad_input, [()], [()] - diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py index 6332960ccf..73c52889c7 100644 --- a/transformer_engine/pytorch/ops/basic/dispatch.py +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -195,4 +195,3 @@ def fuser_backward( grad_recv_weights, ) return grad_input, [()], [(None, grad_topk_weights)] - diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index c452519d42..b64441ef4d 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -586,4 +586,3 @@ def fuse_ops( __all__ = ["FusedMoeEp"] - From f83257e1992aa9e1785bf55b75d6f1138f35999d Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sat, 29 Aug 2026 06:18:48 +0000 Subject: [PATCH 83/83] minor review comments pending Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/basic/swiglu.py | 7 +++---- transformer_engine/pytorch/ops/fused/moe_ep.py | 8 +------- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index b2898a8d74..71cb586079 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -532,14 +532,13 @@ def fuser_forward( if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") scale_dtype = dtype - elif isinstance(input_, torch.Tensor): - dtype = input_.dtype else: - dtype = extra_input.dtype + dtype = input_.dtype + scale_dtype = extra_input.dtype # Prepare the activation input in the compute dtype. input_ = maybe_dequantize(input_, dtype) - scales = extra_input if scale_dtype is None else maybe_dequantize(extra_input, scale_dtype) + scales = maybe_dequantize(extra_input, scale_dtype) out = self._scaled_glu_forward(input_, scales) # Save state for backward pass diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py index b64441ef4d..cc6f516084 100644 --- a/transformer_engine/pytorch/ops/fused/moe_ep.py +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -470,10 +470,6 @@ def fuser_forward( basic_op_ctxs[0].input_dtype = input_dtype basic_op_ctxs[0].prev_op_grad_output_quantizer = prev_op_grad_output_quantizer - # Dispatch extras are channel-bound with output_to_caller=False and are - # only consumed by ops inside this fusion, so they need not be materialized. - if next_op_input_quantizer is not None and not is_quantized_tensor(output): - output = next_op_input_quantizer(output) return output, [ (None, None), (), @@ -538,9 +534,7 @@ def fuser_backward( fc1_param_grads = _compute_grouped_weight_grad(self.fc1, wgrad_operands, "fc1") fc2_param_grads = _compute_grouped_weight_grad(self.fc2, wgrad_operands, "fc2") grad_input = grad_input.to(dtype=basic_op_ctxs[0].input_dtype) - grad_input_quantizer = basic_op_ctxs[0].prev_op_grad_output_quantizer - if grad_input_quantizer is not None: - grad_input = grad_input_quantizer(grad_input) + return ( grad_input, [(), fc1_param_grads, (), fc2_param_grads, ()],