Merge branch 'flagos-ai:main' into add_qa_teat251230

Author: zihugithub

.gitignore

@@ -40,3 +40,5 @@ compile_commands.json
 .nfs
 tensor_dumps/
 artifacts/
+# Auto-generated build configuration (specific to each environment)
+transformer_engine/plugin/core/_build_config.py

README.rst

@@ -5,6 +5,9 @@
 
 |License|
 
+
+**TransformerEngine-FL is a fork of TransformerEngine that introduces a plugin-based architecture for supporting diverse AI chips, built on top of** `FlagOS <https://github.com/flagos-ai>`_, **a unified open-source AI system software stack.**
+
 Transformer Engine
 ==================
 

build_tools/pytorch.py

@@ -85,8 +85,10 @@ def setup_pytorch_extension(
     include_dirs = [str(path) for path in include_dirs]
     from torch.utils.cpp_extension import CppExtension
 
+    # Use transformer_engine_torch_nv as the native NVIDIA module name
+    # This allows the plugin system to use transformer_engine_torch as the unified interface
     return CppExtension(
-        name="transformer_engine_torch",
+        name="transformer_engine_torch_nv",
         sources=[str(src) for src in sources],
         include_dirs=[str(inc) for inc in include_dirs],
         extra_compile_args={"cxx": cxx_flags},

build_tools/utils.py

@@ -251,8 +251,16 @@ def get_cuda_include_dirs() -> Tuple[str, str]:
     ]
 
 
+@functools.lru_cache(maxsize=None)
+def skip_cuda_build() -> bool:
+    """Check if CUDA build should be skipped (for AMD/ROCm or pure FL backend)"""
+    return bool(int(os.getenv("TE_FL_SKIP_CUDA", "0")))
+
+
 @functools.lru_cache(maxsize=None)
 def cuda_archs() -> str:
+    if skip_cuda_build():
+        return ""  # Return empty string when skipping CUDA build
     archs = os.getenv("NVTE_CUDA_ARCHS")
     if archs is None:
         version = cuda_version()

setup.py

@@ -28,6 +28,9 @@
 
 
 from setuptools.command.build_ext import build_ext as BuildExtension
+from setuptools.command.install import install as InstallCommand
+from datetime import datetime
+import platform
 
 os.environ["NVTE_PROJECT_BUILDING"] = "1"
 
@@ -41,6 +44,63 @@
 archs = cuda_archs()
 
 
+def generate_build_config(skip_cuda_build):
+    """Generate build-time configuration file."""
+    config_template_path = (
+        current_file_path / "transformer_engine" / "plugin" /
+        "core" / "_build_config.py.template"
+    )
+    config_output_path = (
+        current_file_path / "transformer_engine" / "plugin" /
+        "core" / "_build_config.py"
+    )
+
+    if config_template_path.exists():
+        with open(config_template_path, 'r') as f:
+            template = f.read()
+
+        config_content = template.format(
+            skip_cuda=skip_cuda_build,
+            build_time=datetime.now().isoformat(),
+            platform=platform.platform(),
+        )
+
+        with open(config_output_path, 'w') as f:
+            f.write(config_content)
+
+        print(f"Generated build config: {config_output_path}")
+        print(f"  SKIP_CUDA_BUILD = {skip_cuda_build}")
+    else:
+        # Fallback: create minimal config if template doesn't exist
+        config_content = f"""# Auto-generated build configuration
+SKIP_CUDA_BUILD = {skip_cuda_build}
+BUILD_TIME = "{datetime.now().isoformat()}"
+BUILD_PLATFORM = "{platform.platform()}"
+"""
+        with open(config_output_path, 'w') as f:
+            f.write(config_content)
+        print(f"Generated minimal build config: {config_output_path}")
+
+
+class CustomInstall(InstallCommand):
+    """Custom install command to generate build config."""
+
+    user_options = InstallCommand.user_options + [
+        ('skip-cuda-build', None, 'Skip CUDA build'),
+    ]
+
+    def initialize_options(self):
+        super().initialize_options()
+        self.skip_cuda_build = bool(int(os.getenv("TE_FL_SKIP_CUDA", "0")))
+
+    def run(self):
+        # Run the standard install
+        super().run()
+
+        # Generate build config after installation
+        generate_build_config(self.skip_cuda_build)
+
+
 class TimedBdist(bdist_wheel):
     """Helper class to measure build time"""
 
@@ -132,6 +192,14 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
     with open("README.rst", encoding="utf-8") as f:
         long_description = f.read()
 
+    # Check if we should skip CUDA build (for AMD/ROCm or pure FL backend usage)
+    skip_cuda_build = bool(int(os.getenv("TE_FL_SKIP_CUDA", "0")))
+    if skip_cuda_build:
+        print("=" * 60)
+        print("TE_FL_SKIP_CUDA=1: Skipping CUDA/native backend compilation")
+        print("Only FL (Flag-Gems/Triton) backend will be available")
+        print("=" * 60)
+
     # Settings for building top level empty package for dependency management.
     if bool(int(os.getenv("NVTE_BUILD_METAPACKAGE", "0"))):
         assert bool(
@@ -148,6 +216,13 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
             "pytorch": [f"transformer_engine_torch=={__version__}"],
             "jax": [f"transformer_engine_jax=={__version__}"],
         }
+    elif skip_cuda_build:
+        # Skip CUDA build - only install Python packages for FL backend
+        install_requires, test_requires = setup_requirements()
+        ext_modules = []  # No CUDA extensions
+        package_data = {"": ["VERSION.txt"]}
+        include_package_data = True
+        extras_require = {"test": test_requires}
     else:
         install_requires, test_requires = setup_requirements()
         ext_modules = [setup_common_extension()]
@@ -177,6 +252,9 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
                     )
                 )
 
+    # Generate build config before setup
+    generate_build_config(skip_cuda_build)
+
     # Configure package
     setuptools.setup(
         name="transformer_engine",
@@ -193,7 +271,11 @@ def setup_requirements() -> Tuple[List[str], List[str]]:
         long_description=long_description,
         long_description_content_type="text/x-rst",
         ext_modules=ext_modules,
-        cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist},
+        cmdclass={
+            "build_ext": CMakeBuildExtension,
+            "bdist_wheel": TimedBdist,
+            "install": CustomInstall,
+        },
         python_requires=f">={min_python_version_str()}",
         classifiers=["Programming Language :: Python :: 3"],
         install_requires=install_requires,

transformer_engine/common/__init__.py

@@ -18,6 +18,31 @@
 from typing import Optional, Tuple
 
 
+def skip_cuda_build() -> bool:
+    """Check if CUDA build was skipped (FL-only mode).
+
+    First checks environment variable (for runtime override),
+    then falls back to build-time configuration.
+    """
+    # Environment variable takes precedence (allows runtime override)
+    if os.environ.get("TE_FL_SKIP_CUDA"):
+        return bool(int(os.environ.get("TE_FL_SKIP_CUDA", "0")))
+
+    # Fall back to build-time configuration
+    try:
+        from transformer_engine.plugin.core._build_config import SKIP_CUDA_BUILD
+        return SKIP_CUDA_BUILD
+    except ImportError:
+        # If build config doesn't exist, default to False
+        return False
+
+# Load plugin system - this handles module registration and backend initialization
+# The _module_setup inside core will:
+# 1. Register modules under both full and short names for relative imports
+# 2. Load all available backends (flagos, reference, vendor/cuda, etc.)
+# 3. Register transformer_engine_torch module from the selected backend
+import transformer_engine.plugin.core  # noqa: F401
+
 @functools.lru_cache(maxsize=None)
 def _is_package_installed(package) -> bool:
     """Check if the given package is installed via pip."""
@@ -107,7 +132,7 @@ def _get_shared_object_file(library: str) -> Path:
     """
 
     # Check provided input and determine the correct prefix for .so.
-    assert library in ("core", "torch", "jax"), f"Unsupported TE library {library}."
+    assert library in ("core", "torch_nv", "jax"), f"Unsupported TE library {library}."
     if library == "core":
         so_prefix = "libtransformer_engine"
     else:
@@ -146,46 +171,31 @@ def get_te_core_package_info() -> Tuple[bool, str, str]:
 @functools.lru_cache(maxsize=None)
 def load_framework_extension(framework: str) -> None:
     """
-    Load shared library with Transformer Engine framework bindings
-    and check verify correctness if installed via PyPI.
+    Load shared library with Transformer Engine framework bindings.
+
+    For PyTorch: The native module is now named transformer_engine_torch_nv,
+    and transformer_engine_torch is provided by the plugin system.
+    This function is kept for backward compatibility but does nothing for torch.
     """
 
+    # Skip loading native extensions if CUDA build was skipped (FL-only mode)
+    if skip_cuda_build():
+        return
+
     # Supported frameworks.
-    assert framework in ("jax", "torch"), f"Unsupported framework {framework}"
+    assert framework in ("jax", "torch_nv"), f"Unsupported framework {framework}"
 
-    # Name of the framework extension library.
+    # For jax: load the native module as before
     module_name = f"transformer_engine_{framework}"
 
-    # Name of the pip extra dependency for framework extensions from PyPI.
-    extra_dep_name = module_name
-    if framework == "torch":
-        extra_dep_name = "pytorch"
+    # Skip if already loaded
+    if module_name in sys.modules:
+        return
 
-    # Find the TE packages. The core and framework packages can only be installed via PyPI.
-    # For the `transformer-engine` package, we need to check explicity.
-    te_core_installed, te_core_package_name, te_core_version = get_te_core_package_info()
-    te_framework_installed = _is_package_installed(module_name)
     te_installed = _is_package_installed("transformer_engine")
-    te_installed_via_pypi = _is_package_installed_from_wheel("transformer_engine")
-
     assert te_installed, "Could not find `transformer_engine`."
 
-    # If the framework extension pip package is installed, it means that TE is installed via
-    # PyPI. For this case we need to make sure that the metapackage, the core lib, and framework
-    # extension are all installed via PyPI and have matching versions.
-    if te_framework_installed:
-        assert te_installed_via_pypi, "Could not find `transformer-engine` PyPI package."
-        assert te_core_installed, "Could not find TE core package `transformer-engine-cu*`."
-
-        assert version(module_name) == version("transformer-engine") == te_core_version, (
-            "Transformer Engine package version mismatch. Found"
-            f" {module_name} v{version(module_name)}, transformer-engine"
-            f" v{version('transformer-engine')}, and {te_core_package_name}"
-            f" v{te_core_version}. Install transformer-engine using "
-            f"'pip3 install --no-build-isolation transformer-engine[{extra_dep_name}]==VERSION'"
-        )
-
-    # After all checks are completed, load the shared object file.
+    # Load the shared object file for jax
     spec = importlib.util.spec_from_file_location(module_name, _get_shared_object_file(framework))
     solib = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = solib
@@ -195,6 +205,10 @@ def load_framework_extension(framework: str) -> None:
 def sanity_checks_for_pypi_installation() -> None:
     """Ensure that package is installed correctly if using PyPI."""
 
+    # Skip sanity checks if CUDA build was skipped (FL-only mode)
+    if skip_cuda_build():
+        return
+
     te_core_installed, te_core_package_name, te_core_version = get_te_core_package_info()
     te_installed = _is_package_installed("transformer_engine")
     te_installed_via_pypi = _is_package_installed_from_wheel("transformer_engine")
@@ -390,13 +404,16 @@ def _load_core_library():
 
 if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
     sanity_checks_for_pypi_installation()
-    _CUDNN_LIB_CTYPES = _load_cudnn()
-    _NVRTC_LIB_CTYPES = _load_nvrtc()
-    _CURAND_LIB_CTYPES = _load_curand()
-    _CUBLAS_LIB_CTYPES = _load_nvidia_cuda_library("cublas")
-    _CUDART_LIB_CTYPES = _load_nvidia_cuda_library("cuda_runtime")
-    _TE_LIB_CTYPES = _load_core_library()
-
-    # Needed to find the correct headers for NVRTC kernels.
-    if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _nvidia_cudart_include_dir():
-        os.environ["NVTE_CUDA_INCLUDE_DIR"] = _nvidia_cudart_include_dir()
+
+    # Skip loading CUDA libraries if CUDA build was skipped (FL-only mode)
+    if not skip_cuda_build():
+        _CUDNN_LIB_CTYPES = _load_cudnn()
+        _NVRTC_LIB_CTYPES = _load_nvrtc()
+        _CURAND_LIB_CTYPES = _load_curand()
+        _CUBLAS_LIB_CTYPES = _load_nvidia_cuda_library("cublas")
+        _CUDART_LIB_CTYPES = _load_nvidia_cuda_library("cuda_runtime")
+        _TE_LIB_CTYPES = _load_core_library()
+
+        # Needed to find the correct headers for NVRTC kernels.
+        if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _nvidia_cudart_include_dir():
+            os.environ["NVTE_CUDA_INCLUDE_DIR"] = _nvidia_cudart_include_dir()

transformer_engine/plugin/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2025, BAAI. All rights reserved.
+#
+# See LICENSE for license information.
+
+from .core import (
+    TEFLBackendBase,
+    TEFLModule,
+    get_tefl_module as _get_tefl_module,
+    get_registry,
+)
+
+def __getattr__(name):
+    if name == "tefl":
+        return _get_tefl_module()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+__all__ = [
+    "TEFLBackendBase",
+    "TEFLModule",
+    "get_tefl_module",
+    "get_registry",
+    "tefl",
+]

transformer_engine/plugin/benchmarks/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2025, BAAI. All rights reserved.
+#
+# See LICENSE for license information.
+
+__all__ = []