diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 6b276d22a8..84b4e65f3f 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -16,6 +16,7 @@ import sys import sysconfig from typing import Optional, Tuple +import warnings @functools.lru_cache(maxsize=None) @@ -191,6 +192,31 @@ def load_framework_extension(framework: str) -> None: sys.modules[module_name] = solib spec.loader.exec_module(solib) + # Plugin system: set NVTE_PLUGIN= to let plugin stub take over + # transformer_engine_torch and register original pybind as _nv for CUDA backend. + # Only applies to the PyTorch extension — JAX has no plugin stub. + _nvte_plugin = os.environ.get("NVTE_PLUGIN") + if _nvte_plugin and framework == "torch": + _original_module = sys.modules.get(module_name) + try: + # Register _nv alias BEFORE importing the plugin, because the + # plugin module may import transformer_engine_torch_nv at top level. + sys.modules[module_name + "_nv"] = solib + _plugin = importlib.import_module(_nvte_plugin) + _plugin.load_plugins() + except Exception as e: + # Rollback to pre-plugin state if plugin failed to fully initialize + sys.modules.pop(module_name + "_nv", None) + if _original_module is not None: + sys.modules[module_name] = _original_module + else: + sys.modules.pop(module_name, None) + warnings.warn( + f"NVTE_PLUGIN={_nvte_plugin} but plugin loading failed: {e}", + RuntimeWarning, + stacklevel=2, + ) + def sanity_checks_for_pypi_installation() -> None: """Ensure that package is installed correctly if using PyPI.""" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index d5adbbcadf..43f0b2d788 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -244,6 +244,18 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: _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_PLUGIN"): + import transformer_engine_torch as tex + + _FlashAttentionNative = FlashAttention + FlashAttention = getattr(tex, "flash_attention", _FlashAttentionNative) + _plugin_get_attention_backend = getattr(tex, "get_attention_backend", None) + if _plugin_get_attention_backend is not None: + dpa_utils._original_get_attention_backend = dpa_utils.get_attention_backend + dpa_utils.get_attention_backend = _plugin_get_attention_backend + + __all__ = ["DotProductAttention"]