Add entrypoint for flagos multi-backend plugin system

[lxd-cumt]

FlagOS Proposal: Plugin Architecture & Device-Agnostic Abstraction for TransformerEngine

Plugin System: Initialization-time Backend Loading

We propose a plugin architecture where TransformerEngine (TE) loads backend implementations at initialization time via an explicit plugin interface, while the actual multi-backend plugins reside in a separate repository (TransformerEngine-Plugin-FL).

TransformerEngine-Plugin-FL: https://github.com/lxd-cumt/TransformerEngine-Plugin-FL

Current State

TE already has a prototype (NVTE_ENABLE_PLUGIN=1 in common/__init__.py) that registers the original CUDA pybind module as transformer_engine_torch_nv and delegates to an external load_plugins() entry point.

Proposed Design

• TE defines a stable plugin API contract (operator signatures, quantization interfaces, communication primitives).
• At load_framework_extension() time, if a plugin is present, TE dispatches backend calls through the plugin registry; otherwise it falls back to the native te implementation.
• The TransformerEngine-Plugin-FL repository is independently installable and contains multiple backend implementations for diverse accelerators, and support more training scenarios.

[greptile-apps]

Greptile Summary

This PR adds an initialization-time plugin system to TransformerEngine, allowing third-party backends to replace the native CUDA pybind module at startup via NVTE_PLUGIN=<module_name>.

• common/__init__.py: Adds a plugin-load block inside load_framework_extension (torch-only, guarded), pre-registers the native pybind module as transformer_engine_torch_nv before the plugin is imported, and rolls back sys.modules on any failure with a RuntimeWarning.
• dot_product_attention.py: At module-import time, when NVTE_PLUGIN is set, safely overrides the module-level FlashAttention reference and monkey-patches dpa_utils.get_attention_backend with the plugin's versions using getattr fallbacks, so a failed or incomplete plugin load leaves native behavior intact.

Confidence Score: 5/5

Safe to merge; the plugin path is gated behind an explicit env var, the rollback logic correctly restores sys.modules on failure, and the dot_product_attention overrides use safe getattr fallbacks throughout.

All of the previously flagged blocking issues (unguarded AttributeError, bare module name, missing framework guard, uncaught non-ImportError, incomplete sys.modules rollback, _nv alias ordering) have been addressed in this revision. The remaining notes are a dead-code branch in the except block and a slightly off stacklevel in warnings.warn — neither affects runtime correctness.

transformer_engine/common/init.py — the dead-code else branch in the rollback is harmless but worth a one-line cleanup.

Important Files Changed

Filename	Overview
transformer_engine/common/init.py	Adds plugin loading block inside load_framework_extension with framework guard, pre-import _nv alias, full except-Exception rollback, and warnings.warn. One dead-code branch in the rollback (the else path can never be reached because _original_module is always solib).
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds module-level plugin override block for FlashAttention and get_attention_backend using safe getattr fallbacks; correctly handles the case where NVTE_PLUGIN is set but the plugin failed to load.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User as User Code
    participant PyInit as pytorch/__init__.py
    participant CommonInit as common/__init__.py load_framework_extension
    participant SysModules as sys.modules
    participant Plugin as NVTE_PLUGIN module
    participant DPA as dot_product_attention.py

    User->>PyInit: import transformer_engine.pytorch
    PyInit->>CommonInit: load_framework_extension("torch")
    CommonInit->>SysModules: "transformer_engine_torch = solib (native)"
    Note over CommonInit: NVTE_PLUGIN set and framework=="torch"?
    CommonInit->>SysModules: "transformer_engine_torch_nv = solib"
    CommonInit->>Plugin: importlib.import_module(NVTE_PLUGIN)
    Plugin-->>CommonInit: _plugin
    CommonInit->>Plugin: _plugin.load_plugins()
    Note over Plugin: May replace sys.modules transformer_engine_torch with plugin stub
    Plugin-->>CommonInit: success / Exception
    alt Exception raised
        CommonInit->>SysModules: "pop _nv, restore transformer_engine_torch = solib"
        CommonInit-->>PyInit: warnings.warn(RuntimeWarning)
    else Success
        CommonInit-->>PyInit: returns
    end
    PyInit->>DPA: import dot_product_attention
    Note over DPA: import transformer_engine_torch as tex resolves to plugin stub or native solib
    Note over DPA: NVTE_PLUGIN set?
    DPA->>DPA: "FlashAttention = getattr(tex, flash_attention, native)"
    DPA->>DPA: "_plugin_get_attn = getattr(tex, get_attention_backend, None)"
    alt _plugin_get_attention_backend is not None
        DPA->>DPA: "dpa_utils.get_attention_backend = plugin version"
    end

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User as User Code
    participant PyInit as pytorch/__init__.py
    participant CommonInit as common/__init__.py load_framework_extension
    participant SysModules as sys.modules
    participant Plugin as NVTE_PLUGIN module
    participant DPA as dot_product_attention.py

    User->>PyInit: import transformer_engine.pytorch
    PyInit->>CommonInit: load_framework_extension("torch")
    CommonInit->>SysModules: "transformer_engine_torch = solib (native)"
    Note over CommonInit: NVTE_PLUGIN set and framework=="torch"?
    CommonInit->>SysModules: "transformer_engine_torch_nv = solib"
    CommonInit->>Plugin: importlib.import_module(NVTE_PLUGIN)
    Plugin-->>CommonInit: _plugin
    CommonInit->>Plugin: _plugin.load_plugins()
    Note over Plugin: May replace sys.modules transformer_engine_torch with plugin stub
    Plugin-->>CommonInit: success / Exception
    alt Exception raised
        CommonInit->>SysModules: "pop _nv, restore transformer_engine_torch = solib"
        CommonInit-->>PyInit: warnings.warn(RuntimeWarning)
    else Success
        CommonInit-->>PyInit: returns
    end
    PyInit->>DPA: import dot_product_attention
    Note over DPA: import transformer_engine_torch as tex resolves to plugin stub or native solib
    Note over DPA: NVTE_PLUGIN set?
    DPA->>DPA: "FlashAttention = getattr(tex, flash_attention, native)"
    DPA->>DPA: "_plugin_get_attn = getattr(tex, get_attention_backend, None)"
    alt _plugin_get_attention_backend is not None
        DPA->>DPA: "dpa_utils.get_attention_backend = plugin version"
    end

Loading

_{Reviews (7): Last reviewed commit: "fix: register _nv module alias before im..." | Re-trigger Greptile}

[ptrendx]

Is the load_plugins performing a generic plugin discovery and loading? If so, we could have that in the main TE repo as well. If it is specific to FlagOS then I agree that it should live there.

[lxd-cumt]

Thanks for the review. Two questions I'd like to confirm before deciding on the next step:

• Should the load_plugins entrypoint be merged into the main branch, so it can follow future NVIDIA TE releases?
• Should the generic framework components in transformer_engine_plugin_fl (e.g., OpManager, SelectionPolicy) also be upstreamed into TE, or should we keep the current approach where TE only exposes a minimal load_plugins hook and all dispatch logic stays external?

[ptrendx]

Hi @lxd-cumt , thank you for the contribution. A few things:

• first, please sign your commits, see https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work for details
• second, I do agree with the Greptile comments - please address them.
  After that I think it will be good to CI and merge.

[lxd-cumt]

Hi @lxd-cumt , thank you for the contribution. A few things:

• first, please sign your commits, see https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work for details
• second, I do agree with the Greptile comments - please address them.
  After that I think it will be good to CI and merge.

Thanks for the review! I've signed off on both commits and addressed the Greptile comments.

[lxd-cumt]

The vermin check in CI is failing with AttributeError: module 'ast' has no attribute 'Str' because the CI environment runs Python 3.14, but the pinned vermin rev uses ast.Str which was removed in Python 3.12+. I noticed that the upstream main branch has already updated the vermin . Should I pick up this fix, or is this check non-blocking?

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[lxd-cumt]

Thanks for the review. Two questions I'd like to confirm before deciding on the next step:

• Should the load_plugins entrypoint be merged into the main branch, so it can follow future NVIDIA TE releases?
• Should the generic framework components in plugin system (e.g., OpManager, SelectionPolicy) also be upstreamed into TE, or should we keep the current approach where TE only exposes a minimal load_plugins hook and all dispatch logic stays external?

[ptrendx]

/te-ci pytorch

[ptrendx]

/te-ci pytorch

The base branch was changed.

[ptrendx]

Oh wait, this PR targets the wrong branch - release_v2.14, whereas it should be targeting main. I tried to change the branch to main in the PR but it actually needs a proper rebase in git.
@lxd-cumt Could you do that rebase?

[lxd-cumt]

Oh wait, this PR targets the wrong branch - release_v2.14, whereas it should be targeting main. I tried to change the branch to main in the PR but it actually needs a proper rebase in git. @lxd-cumt Could you do that rebase?

I've submitted a new #3401 to the main branch. Please review.

[ptrendx]

Ok, let me close this one then and will review the new one.