Multi-Backend Architecture Implementation for TransformerEngine-FL (#4)

# TransformerEngine-FL Plugin System

## Overview

This PR implements a comprehensive multi-backend plugin system for
TransformerEngine-FL, enabling support for multiple hardware vendors
(NVIDIA, AMD, Hygon, etc.) while maintaining full API compatibility with
the original `transformer_engine_torch`.

**Core Philosophy**: A plugin-based backend system that allows hardware
vendors to easily implement their own operator optimizations while
preserving complete compatibility with the original TransformerEngine
API.

## Key Features

### Full API Compatibility
- Drop-in replacement for `transformer_engine_torch`
- Switch backends via environment variables
- Zero changes required to existing user code

### Multi-Backend Support

| Backend | Description | Implementation |
|---------|-------------|----------------|
| **FlagOS (default)** | Triton-based cross-platform implementation |
`backends/flagos/` |
| **CUDA (vendor)** | Wraps original TransformerEngine C++ extensions |
`backends/vendor/cuda/` |
| **Reference** | Pure PyTorch fallback implementation |
`backends/reference/` |

### Three-Tier Backend Selection

```
┌─────────────────────────────────────────────────────────┐
│  1. TE_FL_PER_OP (Per-operator override)    [Highest]   │
│     Example: TE_FL_PER_OP="rmsnorm_fwd=vendor:cuda"     │
├─────────────────────────────────────────────────────────┤
│  2. TE_FL_PREFER (Global preference)                    │
│     Values: flagos / vendor / reference                │
├─────────────────────────────────────────────────────────┤
│  3. Backend Priority (Intrinsic)            [Lowest]    │
│     Each implementation has a priority value            │
└─────────────────────────────────────────────────────────┘
```

## Architecture

### Directory Structure

```
transformer_engine/plugin/core/
├── __init__.py              # Public API exports
├── types.py                 # Core types: BackendImplKind, OpImpl
├── registry.py              # OpRegistry: stores all implementations
├── manager.py               # OpManager: selects and calls implementations
├── policy.py                # SelectionPolicy: backend selection rules
├── discovery.py             # Plugin auto-discovery (entry_points, env)
├── builtin_ops.py           # Registers all built-in backends
├── ops.py                   # TEFLModule: transformer_engine_torch compatible API
├── logger_manager.py        # Logging utilities
├── _module_setup.py         # Module aliasing setup
├── _build_config.py         # Build-time configuration
│
└── backends/
    ├── flagos/              # FlagOS backend (Triton-based)
    │   ├── flagos.py        # FlagOSBackend class
    │   ├── register_ops.py  # Operator registration
    │   └── impl/            # Operator implementations
    │       ├── rmsnorm.py
    │       ├── gemm.py
    │       └── ...
    │
    ├── vendor/              # Vendor backends
    │   └── cuda/            # NVIDIA CUDA backend
    │       ├── cuda.py      # CUDABackend class
    │       └── register_ops.py
    │
    └── reference/           # Reference backend (PyTorch)
        ├── reference.py     # ReferenceBackend class
        ├── register_ops.py
        └── impl/            # Pure PyTorch implementations
```

### Core Components

| File | Description |
|------|-------------|
| `types.py` | Defines `BackendImplKind` (DEFAULT/VENDOR/REFERENCE) and
`OpImpl` dataclass |
| `registry.py` | `OpRegistry` - Central storage for all operator
implementations |
| `manager.py` | `OpManager` - Handles implementation selection,
fallback, and execution |
| `policy.py` | `SelectionPolicy` - Configurable rules for backend
selection |
| `discovery.py` | Auto-discovers plugins via `entry_points` or
`TE_FL_PLUGIN_MODULES` |
| `ops.py` | `TEFLModule` - Provides `transformer_engine_torch`
compatible interface |

## Installation

### Build with CUDA support
```bash
pip install --no-build-isolation -e .
```

### Build without CUDA (FlagOS only)
```bash
TE_FL_SKIP_CUDA=1 pip install --no-build-isolation -e .
```

## Environment Variables

### Backend Selection

| Variable | Description | Values | Default |
|----------|-------------|--------|---------|
| `TE_FL_PREFER` | Preferred backend type | `flagos` / `vendor` /
`reference` | `flagos` |
| `TE_FL_PREFER_VENDOR` | Prefer vendor (legacy) | `1` / `0` | `0` |
| `TE_FL_STRICT` | Strict mode (no fallback) | `1` / `0` | `0` |

### Vendor Filtering

| Variable | Description | Example |
|----------|-------------|---------|
| `TE_FL_ALLOW_VENDORS` | Allowed vendors (whitelist) | `nvidia,amd` |
| `TE_FL_DENY_VENDORS` | Denied vendors (blacklist) | `vendor_a` |

### Per-Operator Configuration

| Variable | Description | Example |
|----------|-------------|---------|
| `TE_FL_PER_OP` | Per-operator backend ordering |
`rmsnorm_fwd=vendor:cuda\|default` |

### Plugin Discovery

| Variable | Description | Example |
|----------|-------------|---------|
| `TE_FL_PLUGIN_MODULES` | Plugin modules to load |
`my_plugin,another_plugin` |

### Build Configuration

| Variable | Description | Values | Default |
|----------|-------------|--------|---------|
| `TE_FL_SKIP_CUDA` | Skip CUDA backend | `1` / `0` | `0` |
| `CUDA_HOME` | CUDA installation path | `/usr/local/cuda` |
Auto-detected |

### Logging

| Variable | Description | Values | Default |
|----------|-------------|--------|---------|
| `TEFL_LOG_LEVEL` | Log level | `DEBUG` / `INFO` / `WARNING` / `ERROR`
| `INFO` |

## Usage Examples

### Basic Usage (No Code Changes Required)
```python
# Existing code works as-is
import transformer_engine.pytorch as te
# or
import transformer_engine_torch as te
```

### Register Custom Backend (In-tree)

```python
from transformer_engine.plugin.core import (
    OpRegistry, OpManager, OpImpl, BackendImplKind
)

# 1. Define implementation
def my_rmsnorm(input, weight, eps=1e-5, **kwargs):
    variance = input.pow(2).mean(-1, keepdim=True)
    return input * torch.rsqrt(variance + eps) * weight, torch.rsqrt(variance + eps)

# 2. Register
registry = OpRegistry()
registry.register_impl(OpImpl(
    op_name="rmsnorm_fwd",
    impl_id="vendor.mybackend",
    kind=BackendImplKind.VENDOR,
    vendor="mybackend",
    fn=my_rmsnorm,
    priority=200,
))

# 3. Call
manager = OpManager(registry)
output, rsigma = manager.call("rmsnorm_fwd", input, weight)
```

### Register Custom Backend (Out-of-tree Plugin)

Create a plugin package with `register(registry)` function:

```python
# my_vendor_plugin/__init__.py
from transformer_engine.plugin.core import OpImpl, BackendImplKind

def my_rmsnorm(input, weight, eps=1e-5, **kwargs):
    # Your implementation
    ...

def register(registry):
    """Called automatically by TE-FL"""
    registry.register_impl(OpImpl(
        op_name="rmsnorm_fwd",
        impl_id="vendor.myvendor",
        kind=BackendImplKind.VENDOR,
        vendor="myvendor",
        fn=my_rmsnorm,
        priority=200,
    ))
```

Load via environment variable:
```bash
export TE_FL_PLUGIN_MODULES=my_vendor_plugin
python your_script.py
```

## Runtime Logs

When running, you'll see logs indicating which backend is used:

```
[TE-FL manager.py:133 INFO] Registered impl_ids: ['default.flagos', 'reference.torch', 'vendor.cuda']
[TE-FL manager.py:390 INFO] Op 'rmsnorm_fwd' using 'default.flagos' (kind=default, vendor=None)
[TE-FL manager.py:395 INFO] Op 'rmsnorm_fwd' switched from 'default.flagos' to 'vendor.cuda' (kind=vendor, vendor=CUDA)
```

## Examples

See `transformer_engine/plugins/examples/` for complete working
examples:
- `example_intree.py` - In-tree backend registration
- `example_outtree.py` - Out-of-tree plugin registration

Fixes # (issue)

## Type of change

- [ ] Documentation change (change only to the documentation, either a
fix or a new content)
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] Infra/Build change
- [ ] Code refactoring

## Changes

Please list the changes introduced in this PR:

- Change A
- Change B

# Checklist:

- [ ] I have read and followed the [contributing
guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst)
- [ ] The functionality is complete
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes

---------

Co-authored-by: panpy <panpy@sugon.com>

Author: lihongyang1990

.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

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/__init__.py

@@ -8,6 +8,7 @@
 
 import os
 from importlib import metadata
+
 import transformer_engine.common
 
 try:

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."""
@@ -146,46 +171,36 @@ 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}"
 
-    # Name of the framework extension library.
+    # For torch: plugin system already handles transformer_engine_torch
+    # The native module is transformer_engine_torch_nv (imported by NVIDIA backend)
+    if framework == "torch":
+        return  # Nothing to do, plugin system handles this
+
+    # 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 +210,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 +409,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__ = []