[PyTorch] Support user-defined op fusions

[timmoon10]

Description

This PR adds the register_forward_fusion and register_backward_fusion functions to the op fuser API, allowing users to register custom fusions.

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

• Add function to register custom op fusions
• Refactor op fuser to have consistent op order in forward and backward pass
• Refactor op fusion functions to avoid index bookkeeping
• Add tests for user-defined ops

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

[greptile-apps]

Greptile Summary

This PR adds support for user-defined operation fusions by introducing register_forward_fusion and register_backward_fusion functions to the public API. The implementation includes a significant refactoring of the op fusion system to improve consistency and maintainability.

Key Changes:

• Added register_forward_fusion() and register_backward_fusion() functions in fuser.py to allow users to register custom fusion functions
• Refactored all fusion functions from standalone functions to static methods on the fused operation classes (e.g., fuse_forward_linear_bias_activation() → ForwardLinearBiasActivation.fuse_forward_ops())
• Simplified fusion function signatures to accept list[FusibleOperation] instead of list[tuple[FusibleOperation, list[int]]], removing index bookkeeping complexity
• Fixed indexing bugs in backward fusion operations (backward_linear_add.py, backward_linear_scale.py) where context indices were incorrect
• Ensured consistent operation order between forward and backward passes by having both work with the same basic operation sequence
• Added comprehensive tests for custom basic operations, custom forward fusions, and custom backward fusions

Impact:
The refactoring provides a cleaner, more maintainable API for operation fusion while fixing several indexing bugs. Users can now define and register their own fusion patterns, extending TransformerEngine's optimization capabilities.

Confidence Score: 5/5

• This PR is safe to merge with no critical issues found
• The PR implements a well-designed feature with comprehensive testing. The refactoring fixes existing indexing bugs in backward operations and significantly improves code maintainability by removing complex index bookkeeping. The new API is clean and follows good design patterns. All changes are backwards compatible since the old fusion functions were internal implementation details.
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fuser.py	Added register_forward_fusion and register_backward_fusion functions to allow users to register custom op fusions. Refactored _fuse_ops method to have cleaner signature.
transformer_engine/pytorch/ops/init.py	Exposed register_forward_fusion and register_backward_fusion APIs and additional operation classes for public use.
transformer_engine/pytorch/ops/fused/init.py	Refactored to register fusion functions using new register_forward_fusion and register_backward_fusion APIs instead of direct function exports.
transformer_engine/pytorch/ops/fused/backward_linear_add.py	Fixed context indexing bug (line 48), converted fusion function to static method, removed index bookkeeping from fusion logic.
tests/pytorch/test_fusible_ops.py	Added comprehensive tests for custom user-defined ops including basic ops, forward fusions, and backward fusions. Updated existing test expectations for op ordering.

Sequence Diagram

sequenceDiagram
    participant User
    participant RegisterAPI as register_forward_fusion/<br/>register_backward_fusion
    participant OperationFuser
    participant FusionFunc as Custom Fusion Function
    participant FusedOp as Fused Operation

    User->>RegisterAPI: register custom fusion function
    RegisterAPI->>OperationFuser: append to forward_fusion_functions<br/>or backward_fusion_functions

    Note over User,FusedOp: During model execution

    User->>OperationFuser: __call__(input, ...)
    OperationFuser->>OperationFuser: maybe_fuse_ops()
    
    loop For each fusion function
        OperationFuser->>FusionFunc: func(ops, recipe=recipe)
        FusionFunc->>FusionFunc: scan ops with sliding window
        alt Pattern matches
            FusionFunc->>FusedOp: create fused op
            FusedOp-->>FusionFunc: fused operation
        else No match
            FusionFunc->>FusionFunc: shift window
        end
        FusionFunc-->>OperationFuser: updated ops list
    end

    OperationFuser->>OperationFuser: store _forward_ops and _backward_ops
    
    Note over OperationFuser,FusedOp: Forward pass
    
    loop For each forward op
        OperationFuser->>FusedOp: fuser_forward(ctxs, input, ...)
        FusedOp->>FusedOp: execute fused computation
        FusedOp->>FusedOp: save state to contexts
        FusedOp-->>OperationFuser: output
    end
    
    Note over OperationFuser,FusedOp: Backward pass
    
    loop For each backward op (reversed)
        OperationFuser->>FusedOp: fuser_backward(ctxs, grad_output, ...)
        FusedOp->>FusedOp: restore from contexts
        FusedOp->>FusedOp: execute fused backward
        FusedOp-->>OperationFuser: grad_input, grad_params
    end
    
    OperationFuser-->>User: outputs and gradients

Loading

[timmoon10]

/te-ci pytorch L1