Reduce Grouped MLP Fuser CPU Overhead

[zhongbozhu]

Description

Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups.

Fixes # (issue)

#2897

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[zhongbozhu]

/te-ci pytorch L1

[greptile-apps]

Greptile Summary

The PR reduces grouped-MLP fuser CPU overhead by reusing fusion plans across alternating checkpoint-forward and recomputation modes and by avoiding redundant CUDA device discovery.

• Caches fused forward/backward operation lists by recipe type, backward boundary, and backward override.
• Keeps recipe-state resets independent from backward-boundary transitions.
• Uses the known tensor device when retrieving grouped-MLP CUDA streams.
• Adds coverage for alternating grad-enabled and no-grad plan reuse.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fuser.py	Separates recipe-state reconfiguration from runtime backward-boundary changes and caches reusable fused operation plans.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Passes the already-known device index to CUDA current-stream lookup in grouped-MLP forward and backward paths.
tests/pytorch/test_fusible_ops.py	Verifies that alternating checkpoint-forward and grad-enabled recomputation states construct only two plans and restore the corresponding cached operation lists.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A["maybe_fuse_ops invocation"] --> B["Determine backward boundary"]
  B --> C{"Recipe state changed?"}
  C -- Yes --> D["Reset basic-op recipe state"]
  C -- No --> E["Build fusion cache key"]
  D --> E
  E --> F{"Cached plan exists?"}
  F -- Yes --> G["Restore cached forward/backward lists"]
  F -- No --> H["Apply fusion passes"]
  H --> I["Cache resulting plan"]
  G --> J["Execute selected plan"]
  I --> J

Loading

_{Reviews (4): Last reviewed commit: "fix cutedsl wgrad crash" | Re-trigger Greptile}

[vthumbe1503]

I think we should use recipe object instead of recipe_type as fusion_param key And implement hash function for it. Since that governs the fusion and not just the recipe type

cc: @timmoon10

[zhongbozhu]

/te-ci pytorch L1