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 caches operation-fusion plans across alternating gradient requirements and avoids redundant CUDA current-device discovery.

• Adds separate reusable plans for checkpoint forward and grad-enabled recomputation.
• Uses the grouped-MLP tensor device directly when obtaining the CUDA stream.
• Adds a focused test for alternating no-grad and grad-enabled cache 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 reset from backward-boundary selection and caches forward/backward plans for reuse.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Obtains each grouped-MLP CUDA stream using the operation tensor's explicit device index.
tests/pytorch/test_fusible_ops.py	Verifies that alternating grad requirements build exactly two plans and subsequently restore the corresponding cached lists.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[OperationFuser invocation] --> B[Compute backward boundary]
  B --> C{Recipe state changed?}
  C -->|yes| D[Reset basic-op recipe state]
  C -->|no| E[Build fusion key]
  D --> E
  E --> F{Cached plan exists?}
  F -->|yes| G[Restore forward and backward plans]
  F -->|no| H[Apply fusion rules]
  H --> I[Cache generated plans]
  G --> J[Execute]
  I --> J

Loading

_{Reviews (6): 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

[vthumbe1503]

@zhongbozhu commenting here so we dont miss out on our discussions related to this PR. The main issue we need to take care of here is the CPU overheads of the additional recipe object as a key in fusion params.

There are two options to deal with this

1. We can cache the recipe representation so that the later iterations using the same recipe object dont have to deal with recomputing the representation/hash needed to be maintained in a fusion params dict. However it turns out recipe object can be recreated on every iteration (for eg, in Megatron).
2. So @timmoon10 had the idea of exposing a representation of recipe object, call it options_key. Different objects will have the same options_key(). And so fuser will take the responsibility of caching the options_key instead of Recipe class itself taking that responsibility. And so if recipe is unchanged in terms of options_key, we wont face any new CPU overheads for the fusion_params_key computation

[zhongbozhu]

/te-ci pytorch L1

[zhongbozhu]

/te-ci pytorch

[zhongbozhu]

/te-ci pytorch L1

[zhongbozhu]

/te-ci pytorch L1