add entrypoint for flagos multi-backend plugin system

Author: lxd-cumt

transformer_engine/common/__init__.py

@@ -191,6 +191,16 @@ def load_framework_extension(framework: str) -> None:
     sys.modules[module_name] = solib
     spec.loader.exec_module(solib)
 
+    # Plugin system: if NVTE_ENABLE_PLUGIN=1, let plugin stub take over
+    # transformer_engine_torch and register original pybind as _nv for CUDA backend.
+    if os.environ.get("NVTE_ENABLE_PLUGIN", "0") == "1":
+        sys.modules[module_name + "_nv"] = solib
+        try:
+            from plugin import load_plugins
+            load_plugins()
+        except ImportError as e:
+            print(f"[TE] NVTE_ENABLE_PLUGIN=1 but plugin import failed: {e}")
+
 
 def sanity_checks_for_pypi_installation() -> None:
     """Ensure that package is installed correctly if using PyPI."""

transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py

@@ -148,6 +148,14 @@
 _dpa_fp8ds_reduce_amax = os.getenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") == "1"
 
 
+# Plugin system: override FlashAttention and get_attention_backend if enabled
+if os.environ.get("NVTE_ENABLE_PLUGIN", "0") == "1":
+    _FlashAttentionNative = FlashAttention
+    FlashAttention = getattr(tex, "flash_attention", _FlashAttentionNative)
+    dpa_utils._original_get_attention_backend = dpa_utils.get_attention_backend
+    dpa_utils.get_attention_backend = tex.get_attention_backend
+
+
 __all__ = ["DotProductAttention"]