Reduce Grouped MLP Fuser CPU Overhead

tests/pytorch/test_fusible_ops.py

@@ -146,6 +146,80 @@ def maybe_skip_quantization(
             pytest.skip("NVFP4 quantization is only supported with BF16 data")
 
 
+def test_operation_fuser_caches_plans_by_grad_requirement(monkeypatch) -> None:
+    """Cache and restore fusion plans for checkpoint forward and recompute."""
+
+    # Count fusion-plan construction without depending on any particular real
+    # fusion implementation. Each distinct fusion configuration invokes this
+    # hook once, while a cache hit must bypass it entirely.
+    fusion_calls = 0
+
+    def track_fusion(ops, *, recipe):  # pylint: disable=unused-argument
+        nonlocal fusion_calls
+        fusion_calls += 1
+        # Preserve the operation list so this hook observes plan construction
+        # without changing the topology under test.
+        return ops
+
+    # The fusion registries are class attributes shared by every OperationFuser.
+    # pytest's monkeypatch fixture restores all three after the test, preventing
+    # this synthetic fusion function from leaking into other tests. Keep only a
+    # joint forward-backward fusion hook so each plan build has one countable
+    # callback and no registered TE fusion can affect the result.
+    monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion])
+    monkeypatch.setattr(OperationFuser, "forward_fusion_functions", [])
+    monkeypatch.setattr(OperationFuser, "backward_fusion_functions", [])
+
+    # One Identity op is enough to exercise the cache. With one basic op,
+    # first_op_requiring_backward has an intentionally simple interpretation:
+    #   0: backward starts at the Identity op;
+    #   1: the boundary is past the only op, so no backward work is required.
+    fuser = OperationFuser([te_ops.Identity()])
+    x = torch.ones(1, requires_grad=True)
+    # maybe_fuse_ops expects one extra-input collection per basic op. Identity
+    # has no extra inputs, so its collection is an empty tuple.
+    extra_inputs = [()]
+
+    # Phase 1: the original checkpointed forward runs with grad disabled. This
+    # is the first invocation, so the fuser must construct and cache the no-grad
+    # configuration. The runtime backward boundary is past the only op.
+    fuser.maybe_fuse_ops(False, None, x, extra_inputs)
+    assert fusion_calls == 1
+    assert fuser.first_op_requiring_backward == 1
+    no_grad_forward_ops = fuser._forward_ops
+    no_grad_backward_ops = fuser._backward_ops
+
+    # Phase 2: backward replays the checkpointed region with grad enabled. The
+    # backward boundary is part of the fusion key, allowing future fusion rules
+    # to choose a training-specific topology. The first grad-enabled invocation
+    # therefore constructs and caches a second configuration.
+    fuser.maybe_fuse_ops(True, None, x, extra_inputs)
+    assert fusion_calls == 2
+    assert fuser.first_op_requiring_backward == 0
+    grad_forward_ops = fuser._forward_ops
+    grad_backward_ops = fuser._backward_ops
+    assert grad_forward_ops is not no_grad_forward_ops
+    assert grad_backward_ops is not no_grad_backward_ops
+
+    # Phase 3: the next checkpointed forward must select the exact no-grad lists
+    # cached in phase 1. Before the cache was added, every boundary transition
+    # rebuilt the fused operations and called track_fusion again.
+    fuser.maybe_fuse_ops(False, None, x, extra_inputs)
+    assert fusion_calls == 2
+    assert fuser.first_op_requiring_backward == 1
+    assert fuser._forward_ops is no_grad_forward_ops
+    assert fuser._backward_ops is no_grad_backward_ops
+
+    # Phase 4: another recomputation must likewise restore the grad-enabled
+    # lists from phase 2. The full alternating sequence has built only the two
+    # configurations represented by its two fusion keys.
+    fuser.maybe_fuse_ops(True, None, x, extra_inputs)
+    assert fusion_calls == 2
+    assert fuser.first_op_requiring_backward == 0
+    assert fuser._forward_ops is grad_forward_ops
+    assert fuser._backward_ops is grad_backward_ops
+
+
 @torch.no_grad()
 def make_reference_and_test_tensors(
     shape: int | Iterable[int],

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -1312,7 +1312,7 @@ def fuser_forward(
 
         alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device)
         norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device)
-        current_stream = torch.cuda.current_stream().cuda_stream
+        current_stream = torch.cuda.current_stream(device.index).cuda_stream
 
         fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op)
         fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op)
@@ -1979,7 +1979,7 @@ def fuser_backward(
         # Kernel scaling factors
         alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device)
         norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device)
-        current_stream = torch.cuda.current_stream().cuda_stream
+        current_stream = torch.cuda.current_stream(device.index).cuda_stream
 
         unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False))
         scales_f32 = None

transformer_engine/pytorch/ops/fuser.py

@@ -48,6 +48,8 @@ def _is_graph_capturing() -> bool:
 OperationFusionFunction: TypeAlias = (
     "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]"
 )
+_FusedOpList: TypeAlias = list[tuple[FusibleOperation, list[int]]]
+_FusionParams: TypeAlias = tuple[type, int, Optional[str]]
 
 
 class _OperationFuserAutogradFunction(torch.autograd.Function):
@@ -535,15 +537,22 @@ def __init__(
             op._lock_extra_tensor_channels()
 
         # Ops for forward and backward pass, will be populated in maybe_fuse_ops
-        self._forward_ops: list[tuple[FusibleOperation, list[int]]]
-        self._backward_ops: list[tuple[FusibleOperation, list[int]]]
+        self._forward_ops: _FusedOpList
+        self._backward_ops: _FusedOpList
+
+        # Fused operation configurations are reusable wrappers around the basic
+        # ops, so cache each configuration by the state that selected it.
+        self._fused_ops_cache: dict[_FusionParams, tuple[_FusedOpList, _FusedOpList]] = {}
 
         # Cache and detect change of state relevant for fusing operations
         self.recipe_type = None
-        self.first_op_requiring_backward = 0
         self.backward_override = None
         self._last_amax_history_len = 0
 
+        # Runtime backward boundary. Full activation recompute alternates this
+        # between the checkpointed forward and the grad-enabled recomputation.
+        self.first_op_requiring_backward = 0
+
         # Flatten list of parameters
         self._basic_op_params = [list(op.parameters()) for op in self._basic_ops]
         self._basic_op_num_params = list(map(len, self._basic_op_params))
@@ -626,36 +635,52 @@ def maybe_fuse_ops(
                     first_op_requiring_backward = op_idx
                     break
 
-        # Early exit if fusion parameters haven't changed
-        need_reset = False
+        # Update the runtime backward boundary on every invocation, including
+        # paths that reuse a cached fused operation configuration.
+        self.first_op_requiring_backward = first_op_requiring_backward
+
+        # Recipe state belongs to the basic ops and is independent of the
+        # runtime backward boundary. Only reconfigure it when recipe settings
+        # that are relevant to operation fusion change.
         recipe_type = type(recipe)
         backward_override = recipe.backward_override if recipe is not None else None
-        fusion_params = (recipe_type, first_op_requiring_backward, backward_override)
-        if fusion_params != (
+        recipe_params = (recipe_type, backward_override)
+        need_reset = recipe_params != (
             self.recipe_type,
-            self.first_op_requiring_backward,
             self.backward_override,
-        ):
-            # Recipe type, backward override, or grad requirements have changed
-            need_reset = True
-        elif (
-            recipe is not None
+        )
+        if (
+            not need_reset
+            and recipe is not None
             and recipe.delayed()
             and self._last_amax_history_len != recipe.amax_history_len
         ):
-            # FP8 delayed scaling has changed amax history length
+            # Delayed-scaling history affects quantizer state, but it does not
+            # select a different fused operation configuration.
             need_reset = True
-        if not need_reset:
-            return
+        if need_reset:
+            for op in self._basic_ops:
+                op.reset_recipe_state(recipe=recipe)
 
-        # Reset recipe state
-        for op in self._basic_ops:
-            op.reset_recipe_state(recipe=recipe)
+            # Check if this is the first iteration
+            if self.recipe_type is None:
+                for op in self._basic_ops:
+                    op.pre_first_fuser_forward()
 
-        # Check if this is the first iteration
-        if self.recipe_type is None:
-            for op in self._basic_ops:
-                op.pre_first_fuser_forward()
+            self.recipe_type, self.backward_override = recipe_params
+            self._last_amax_history_len = (
+                recipe.amax_history_len if isinstance(recipe, DelayedScaling) else 0
+            )
+
+        # Training and inference may support different fusions. Keep the
+        # backward boundary in the key, but pay construction cost only once for
+        # each configuration. Full recompute therefore builds at most one
+        # no-grad plan and one grad-enabled plan for a stable recipe.
+        fusion_params = (recipe_type, first_op_requiring_backward, backward_override)
+        cached_ops = self._fused_ops_cache.get(fusion_params)
+        if cached_ops is not None:
+            self._forward_ops, self._backward_ops = cached_ops
+            return
 
         # Apply joint forward-backward fusions first
         joint_ops = OperationFuser._apply_fusions(
@@ -682,14 +707,9 @@ def maybe_fuse_ops(
             self._basic_ops,
         )
 
-        # Save current fusion params
-        self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params
-
-        # Save amax history length
-        if isinstance(recipe, DelayedScaling):
-            self._last_amax_history_len = recipe.amax_history_len
-        else:
-            self._last_amax_history_len = 0
+        # The FusedOperation contract excludes parameters and per-invocation
+        # state, so the mapped lists can be selected directly on cache hits.
+        self._fused_ops_cache[fusion_params] = (self._forward_ops, self._backward_ops)
 
     def __call__(
         self,