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

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

Adds initialization-time loading for an explicitly configured PyTorch backend plugin and allows that plugin to replace FlashAttention construction and attention-backend selection.

• Loads the module named by NVTE_PLUGIN after initializing the native PyTorch extension, exposing the native extension through a transformer_engine_torch_nv alias.
• Warns and restores the native extension when plugin initialization raises an exception.
• Uses optional plugin-provided flash_attention and get_attention_backend exports in dot-product attention.

Confidence Score: 4/5

The PR appears safe to merge, with the non-blocking caveat that its new plugin configuration and interface contracts should be documented.

Native behavior is restored when plugin initialization fails, and no concrete incompatible plugin or security-boundary violation was established; the remaining issue is discoverability of the newly introduced public plugin contract.

Files Needing Attention: transformer_engine/common/init.py

Important Files Changed

Filename	Overview
transformer_engine/common/init.py	Adds memoized import-time plugin loading, a native-extension alias, and warning-based fallback; the new public configuration contract is not documented.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Replaces FlashAttention and backend selection with optional extension exports when plugin configuration is present.

Sequence Diagram

sequenceDiagram
  participant App
  participant TE as TransformerEngine
  participant Native as transformer_engine_torch
  participant Plugin as NVTE_PLUGIN module
  participant DPA as DotProductAttention
  App->>TE: import transformer_engine.pytorch
  TE->>Native: load native extension
  TE->>TE: register transformer_engine_torch_nv alias
  TE->>Plugin: import module and call load_plugins()
  alt plugin initialization succeeds
    Plugin-->>TE: install plugin stub/exports
    DPA->>Plugin: resolve flash_attention and get_attention_backend
  else plugin initialization fails
    TE->>TE: restore native extension
    TE-->>App: emit RuntimeWarning
  end

Loading

_{Reviews (1): Last reviewed commit: "Merge branch 'main' of https://github.co..." | Re-trigger Greptile}

[greptile-apps]

Document the plugin API contract

NVTE_PLUGIN introduces a user-facing configuration interface, but the required load_plugins(), flash_attention, and get_attention_backend contracts are not documented. Plugin implementers must infer their signatures and return values from source, leaving incompatible implementations to fail through import-time warnings or attention-time exceptions.

Knowledge Base Used:

• Common core (transformer_engine/common)
• PyTorch Attention Stack

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[ptrendx]

/te-ci pytorch

[ptrendx]

@lxd-cumt Please fix the lint issue (the one with too broad Exception catch).