From 37e660a36e304244f173ccb7119c3567f6aace2f Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 20 Aug 2026 10:55:21 -0700 Subject: [PATCH 1/3] Reduce grouped MLP fuser CPU overhead Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups. Co-authored-by: Ting-Yang Kao Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_fusible_ops.py | 57 +++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 6 +- transformer_engine/pytorch/ops/fuser.py | 16 ++++-- 3 files changed, 71 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3a0ad7dd3f..7f72eb02ab 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -146,6 +146,63 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +def test_operation_fuser_reuses_plan_when_grad_requirement_changes(monkeypatch) -> None: + """Grad-mode changes update runtime state without rebuilding the fusion plan.""" + + # Count fusion-plan construction without depending on any particular real + # fusion implementation. A cache hit must bypass this function 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 one plan. + # Grad being disabled moves the runtime backward boundary past the only op. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 1 + assert fuser.first_op_requiring_backward == 1 + + # Phase 2: backward replays the checkpointed region with grad enabled. The + # recipe and operation topology have not changed, so rebuilding the fusion + # plan would be wasted CPU work. The runtime boundary must nevertheless be + # updated before maybe_fuse_ops takes its cache-hit early return. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 1 + assert fuser.first_op_requiring_backward == 0 + + # Phase 3: model the next checkpointed forward. This changes the runtime + # boundary back to "no backward" but still must reuse the original plan. + # Before the fix, the boundary was part of the cache key, so the three + # phases called track_fusion three times instead of once. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 1 + assert fuser.first_op_requiring_backward == 1 + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 6b3f53fbd9..59db9a65aa 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -474,7 +474,7 @@ def _cudnn_compute_wgrad( b = X = (total_tokens, in_features) column-major. """ if current_stream is None: - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(grouped_dy.device.index).cuda_stream out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0] @@ -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 diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fd66529ba8..9f4bce0b91 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -540,10 +540,13 @@ def __init__( # 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. This may change between checkpointed + # forward and recompute without changing the fused operation plan. + 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,17 +629,20 @@ def maybe_fuse_ops( first_op_requiring_backward = op_idx break + # Grad requirements control runtime execution, not fusion topology. + # Update them on every invocation, including the early-return path. + self.first_op_requiring_backward = first_op_requiring_backward + # Early exit if fusion parameters haven't changed need_reset = False 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) + fusion_params = (recipe_type, backward_override) if fusion_params != ( self.recipe_type, - self.first_op_requiring_backward, self.backward_override, ): - # Recipe type, backward override, or grad requirements have changed + # Recipe type or backward override has changed need_reset = True elif ( recipe is not None @@ -683,7 +689,7 @@ def maybe_fuse_ops( ) # Save current fusion params - self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params + self.recipe_type, self.backward_override = fusion_params # Save amax history length if isinstance(recipe, DelayedScaling): From 27104b8e103862c6183182f9fb35381caa0235b1 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 20 Aug 2026 11:32:24 -0700 Subject: [PATCH 2/3] resolve comments Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_fusible_ops.py | 47 ++++++++++----- transformer_engine/pytorch/ops/fuser.py | 80 +++++++++++++++---------- 2 files changed, 79 insertions(+), 48 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 7f72eb02ab..eb3bbc145c 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -146,11 +146,12 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") -def test_operation_fuser_reuses_plan_when_grad_requirement_changes(monkeypatch) -> None: - """Grad-mode changes update runtime state without rebuilding the fusion plan.""" +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. A cache hit must bypass this function entirely. + # 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 @@ -180,27 +181,43 @@ def track_fusion(ops, *, recipe): # pylint: disable=unused-argument extra_inputs = [()] # Phase 1: the original checkpointed forward runs with grad disabled. This - # is the first invocation, so the fuser must construct and cache one plan. - # Grad being disabled moves the runtime backward boundary past the only op. + # 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 - # recipe and operation topology have not changed, so rebuilding the fusion - # plan would be wasted CPU work. The runtime boundary must nevertheless be - # updated before maybe_fuse_ops takes its cache-hit early return. + # 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 == 1 + assert fusion_calls == 2 assert fuser.first_op_requiring_backward == 0 - - # Phase 3: model the next checkpointed forward. This changes the runtime - # boundary back to "no backward" but still must reuse the original plan. - # Before the fix, the boundary was part of the cache key, so the three - # phases called track_fusion three times instead of once. + 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 == 1 + 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() diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 9f4bce0b91..e684384dc8 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/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,16 +537,20 @@ 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.backward_override = None self._last_amax_history_len = 0 - # Runtime backward boundary. This may change between checkpointed - # forward and recompute without changing the fused operation plan. + # 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 @@ -629,39 +635,52 @@ def maybe_fuse_ops( first_op_requiring_backward = op_idx break - # Grad requirements control runtime execution, not fusion topology. - # Update them on every invocation, including the early-return path. + # 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 - # Early exit if fusion parameters haven't changed - need_reset = False + # 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, backward_override) - if fusion_params != ( + recipe_params = (recipe_type, backward_override) + need_reset = recipe_params != ( self.recipe_type, self.backward_override, - ): - # Recipe type or backward override has 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( @@ -688,14 +707,9 @@ def maybe_fuse_ops( self._basic_ops, ) - # Save current fusion params - self.recipe_type, 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, From 97850a4d209d1cb69ba15f1ac8cc262ddab6df78 Mon Sep 17 00:00:00 2001 From: tingyangk Date: Fri, 21 Aug 2026 17:54:10 -0700 Subject: [PATCH 3/3] fix cutedsl wgrad crash Signed-off-by: tingyangk --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 59db9a65aa..bb625aa09e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -474,7 +474,7 @@ def _cudnn_compute_wgrad( b = X = (total_tokens, in_features) column-major. """ if current_stream is None: - current_stream = torch.cuda.current_stream(grouped_dy.device.index).cuda_stream + current_stream = torch.cuda.current_stream().cuda_stream out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0]