[Draft] Use vendored cuDNN frontend for Python

build_tools/cudnn_frontend.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Build helpers for the vendored cuDNN frontend Python package."""
+
+from __future__ import annotations
+
+import argparse
+from importlib import metadata
+import json
+import os
+from pathlib import Path
+import re
+import subprocess
+import sys
+import tempfile
+from typing import Optional
+from urllib.parse import unquote, urlparse
+
+
+PACKAGE_NAME = "nvidia-cudnn-frontend"
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parent.parent
+
+
+def cudnn_frontend_root(repo_root: Optional[Path] = None) -> Path:
+    """Path to the vendored cuDNN frontend submodule."""
+    root = _repo_root() if repo_root is None else Path(repo_root)
+    return root / "3rdparty" / "cudnn-frontend"
+
+
+def cudnn_frontend_version(repo_root: Optional[Path] = None) -> str:
+    """Read the vendored cuDNN frontend Python package version without importing it."""
+    init_py = cudnn_frontend_root(repo_root) / "python" / "cudnn" / "__init__.py"
+    if not init_py.is_file():
+        raise FileNotFoundError(
+            f"Could not find vendored cuDNN frontend at {init_py}. "
+            "Run `git submodule update --init --recursive`."
+        )
+    match = re.search(r'^__version__\s*=\s*["\']([^"\']+)["\']', init_py.read_text(), re.M)
+    if match is None:
+        raise RuntimeError(f"Could not parse cuDNN frontend version from {init_py}.")
+    return match.group(1)
+
+
+def installed_cudnn_frontend_version() -> Optional[str]:
+    """Return installed nvidia-cudnn-frontend version, if present."""
+    try:
+        return metadata.version(PACKAGE_NAME)
+    except metadata.PackageNotFoundError:
+        return None
+
+
+def _installed_cudnn_frontend_direct_url() -> Optional[str]:
+    """Return the installed package direct URL metadata, if present."""
+    try:
+        dist = metadata.distribution(PACKAGE_NAME)
+    except metadata.PackageNotFoundError:
+        return None
+    direct_url = dist.read_text("direct_url.json")
+    if direct_url is None:
+        return None
+    try:
+        return json.loads(direct_url).get("url")
+    except json.JSONDecodeError:
+        return None
+
+
+def _direct_url_matches_path(url: Optional[str], path: Path) -> bool:
+    if url is None:
+        return False
+    parsed = urlparse(url)
+    if parsed.scheme != "file":
+        return False
+    return Path(unquote(parsed.path)).resolve() == path.resolve()
+
+
+def _pip_install_command(vendored_root: Path, *, no_build_isolation: bool) -> list[str]:
+    command = [
+        sys.executable,
+        "-m",
+        "pip",
+        "install",
+        "--no-deps",
+        "--force-reinstall",
+    ]
+    if no_build_isolation:
+        command.append("--no-build-isolation")
+    command.append(str(vendored_root))
+    return command
+
+
+def _prepend_env_path(env: dict[str, str], key: str, path: Path) -> None:
+    path_string = str(path)
+    if key in env and env[key]:
+        existing_paths = env[key].split(os.pathsep)
+        if path_string not in existing_paths:
+            env[key] = os.pathsep.join([path_string, env[key]])
+    else:
+        env[key] = path_string
+
+
+def _pybind11_cmake_dir() -> Optional[Path]:
+    try:
+        import pybind11  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        return None
+
+    try:
+        cmake_dir = Path(pybind11.get_cmake_dir())
+    except AttributeError:
+        return None
+    if (cmake_dir / "pybind11Config.cmake").is_file() or (
+        cmake_dir / "pybind11-config.cmake"
+    ).is_file():
+        return cmake_dir
+    return None
+
+
+def _pip_env() -> dict[str, str]:
+    env = os.environ.copy()
+    env.setdefault("PIP_CACHE_DIR", str(Path(tempfile.gettempdir()) / "te-pip-cache"))
+    pybind11_cmake_dir = _pybind11_cmake_dir()
+    if pybind11_cmake_dir is not None:
+        env.setdefault("pybind11_DIR", str(pybind11_cmake_dir))
+        pybind11_prefix = pybind11_cmake_dir
+        if (
+            pybind11_cmake_dir.parent.name == "cmake"
+            and pybind11_cmake_dir.parent.parent.name == "share"
+        ):
+            pybind11_prefix = pybind11_cmake_dir.parent.parent.parent
+        _prepend_env_path(env, "CMAKE_PREFIX_PATH", pybind11_prefix)
+    return env
+
+
+def install_from_submodule(
+    repo_root: Optional[Path] = None,
+    *,
+    force: bool = False,
+    no_build_isolation: bool = True,
+) -> None:
+    """Install nvidia-cudnn-frontend from the vendored submodule.
+
+    The default uses ``--no-build-isolation`` so CI can use the already-installed
+    build dependencies and avoid accidentally fetching a different build stack.
+    """
+    vendored_root = cudnn_frontend_root(repo_root)
+    if not (vendored_root / "pyproject.toml").is_file():
+        raise FileNotFoundError(
+            f"Could not find vendored cuDNN frontend at {vendored_root}. "
+            "Run `git submodule update --init --recursive`."
+        )
+
+    vendored_version = cudnn_frontend_version(repo_root)
+    installed_version = installed_cudnn_frontend_version()
+    installed_from_vendored = _direct_url_matches_path(
+        _installed_cudnn_frontend_direct_url(), vendored_root
+    )
+    if installed_version == vendored_version and installed_from_vendored and not force:
+        print(
+            f"{PACKAGE_NAME}=={installed_version} is already installed from the vendored submodule."
+        )
+        return
+
+    if installed_version is None:
+        print(f"Installing {PACKAGE_NAME}=={vendored_version} from {vendored_root}.")
+    elif installed_version == vendored_version:
+        print(
+            f"Reinstalling {PACKAGE_NAME}=={installed_version} from {vendored_root} "
+            "to use Transformer Engine's vendored submodule."
+        )
+    else:
+        print(
+            f"Replacing {PACKAGE_NAME}=={installed_version} with vendored "
+            f"{PACKAGE_NAME}=={vendored_version} from {vendored_root}."
+        )
+    subprocess.check_call(
+        _pip_install_command(vendored_root, no_build_isolation=no_build_isolation),
+        env=_pip_env(),
+    )
+
+
+def build_wheel(
+    wheel_dir: Path,
+    repo_root: Optional[Path] = None,
+    *,
+    no_build_isolation: bool = True,
+) -> None:
+    """Build a wheel for the vendored nvidia-cudnn-frontend package."""
+    vendored_root = cudnn_frontend_root(repo_root)
+    if not (vendored_root / "pyproject.toml").is_file():
+        raise FileNotFoundError(
+            f"Could not find vendored cuDNN frontend at {vendored_root}. "
+            "Run `git submodule update --init --recursive`."
+        )
+    wheel_dir = Path(wheel_dir)
+    wheel_dir.mkdir(parents=True, exist_ok=True)
+
+    command = [sys.executable, "-m", "pip", "wheel", "--no-deps", "-w", str(wheel_dir)]
+    if no_build_isolation:
+        command.append("--no-build-isolation")
+    command.append(str(vendored_root))
+    subprocess.check_call(command, env=_pip_env())
+
+
+def _main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    subparsers.add_parser("version", help="Print the vendored submodule package version")
+
+    install_parser = subparsers.add_parser("install", help="Install from the vendored submodule")
+    install_parser.add_argument(
+        "--force", action="store_true", help="Reinstall even if versions match"
+    )
+    install_parser.add_argument(
+        "--build-isolation",
+        action="store_true",
+        help="Allow pip build isolation for the cuDNN frontend build",
+    )
+
+    wheel_parser = subparsers.add_parser("wheel", help="Build a wheel from the vendored submodule")
+    wheel_parser.add_argument("wheel_dir", type=Path)
+    wheel_parser.add_argument(
+        "--build-isolation",
+        action="store_true",
+        help="Allow pip build isolation for the cuDNN frontend build",
+    )
+
+    args = parser.parse_args()
+    if args.command == "version":
+        print(cudnn_frontend_version())
+    elif args.command == "install":
+        install_from_submodule(
+            force=args.force,
+            no_build_isolation=not args.build_isolation,
+        )
+    elif args.command == "wheel":
+        build_wheel(args.wheel_dir, no_build_isolation=not args.build_isolation)
+
+
+if __name__ == "__main__":
+    _main()

build_tools/wheel_utils/build_wheels.sh

@@ -25,6 +25,11 @@ git submodule update --init --recursive
 # Install deps
 /opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel
 
+echo "Building vendored nvidia-cudnn-frontend wheel"
+CUDNN_FRONTEND_VERSION=$(/opt/python/cp310-cp310/bin/python -m build_tools.cudnn_frontend version)
+/opt/python/cp310-cp310/bin/python -m build_tools.cudnn_frontend wheel /wheelhouse 2>&1 | tee /wheelhouse/logs/cudnn_frontend.txt
+/opt/python/cp310-cp310/bin/pip install --no-index --find-links /wheelhouse "nvidia-cudnn-frontend==${CUDNN_FRONTEND_VERSION}"
+
 # Enable optional build features. cuSolverMp is provided by the build image
 # (see Dockerfile.x86 / Dockerfile.aarch), which also sets CUSOLVERMP_HOME.
 export NVTE_WITH_CUSOLVERMP=1

docs/examples/te_mixtral/requirements.txt

@@ -7,4 +7,3 @@ safetensors==0.7.0
 huggingface_hub==1.10.1
 tokenizers==0.22.2
 flash-attn
-nvidia-cudnn-frontend>=1.23.0

docs/examples/te_mixtral/te_mixtral_mxfp8.py

@@ -13,8 +13,9 @@
 
 The fused kernel is enabled by ``utils._enable_fused_mxfp8_grouped_mlp()``
 (sets ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` and patches the SM-version /
-cudnn-frontend signature checks). Requires
-``nvidia-cudnn-frontend >= 1.23.0`` and SM>=10 (Blackwell B100/B200/B300+).
+cudnn-frontend signature checks). Requires the cuDNN frontend Python package
+built from Transformer Engine's vendored submodule and SM>=10
+(Blackwell B100/B200/B300+).
 """
 
 from __future__ import annotations

docs/examples/te_mixtral/utils.py

@@ -159,17 +159,19 @@ def _enable_fused_mxfp8_grouped_mlp() -> None:
 
     We also (a) relax the SM-version check from ``!= 10`` to ``>= 10`` so
     SM>=11 successors of B300 fire the kernel, and (b) wrap the cudnn-frontend
-    grouped-GEMM wrappers so the installed TE's ``c_dtype`` kwarg (dropped by
-    cudnn-frontend 1.23.0) is silently filtered out.
+    grouped-GEMM wrappers so the installed TE's ``c_dtype`` kwarg is silently
+    filtered out when the vendored cuDNN frontend does not accept it.
     """
     os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1"
 
     import inspect
-    import cudnn  # type: ignore
+    from transformer_engine.common.cudnn_frontend import import_cudnn_frontend
     from transformer_engine.pytorch.ops.fused import forward_grouped_mlp as _fwd_mod
     from transformer_engine.pytorch.ops.fused import backward_grouped_mlp as _bwd_mod
     from transformer_engine.pytorch.utils import get_device_compute_capability
 
+    cudnn = import_cudnn_frontend()
+
     def _make_is_supported(kernel_method_names):
         def _is_supported(cls) -> bool:
             if int(os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")) <= 0:

qa/L0_jax_unittest/test.sh

@@ -28,6 +28,9 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
+PYTHONPATH=$TE_PATH python3 -m build_tools.cudnn_frontend install \
+    || error_exit "Failed to install vendored nvidia-cudnn-frontend"
+
 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_score_mod.xml $TE_PATH/tests/jax/test_fused_attn_score_mod.py || test_fail "tests/jax/test_fused_attn_score_mod.py"
 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_with_determinism.xml $TE_PATH/tests/jax/test_fused_attn.py -k "TestFusedAttnWithDeterminism" || test_fail "tests/jax/test_fused_attn.py"

qa/L0_pytorch_unittest/test.sh

@@ -23,6 +23,8 @@ set -x
 mkdir -p "$XML_LOG_DIR"
 
 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
+PYTHONPATH=$TE_PATH python3 -m build_tools.cudnn_frontend install \
+    || error_exit "Failed to install vendored nvidia-cudnn-frontend"
 
 NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py"

qa/L1_jax_distributed_unittest/test.sh

@@ -3,6 +3,11 @@
 # See LICENSE for license information.
 export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 
+function error_exit() {
+    echo "Error: $1"
+    exit 1
+}
+
 function test_fail() {
     RET=1
     FAILED_CASES="$FAILED_CASES $1"
@@ -18,6 +23,9 @@ export NVTE_JAX_TEST_TIMING=1
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
+PYTHONPATH=$TE_PATH python3 -m build_tools.cudnn_frontend install \
+    || error_exit "Failed to install vendored nvidia-cudnn-frontend"
+
 export NVTE_JAX_UNITTEST_LEVEL="L1"
 
 # Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate.

qa/L2_jax_distributed_unittest/test.sh

@@ -11,5 +11,7 @@ export NVTE_JAX_TEST_TIMING=1
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
+PYTHONPATH=$TE_PATH python3 -m build_tools.cudnn_frontend install
+
 # Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate.
 XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*

qa/L2_jax_unittest/test.sh

@@ -28,6 +28,9 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
+PYTHONPATH=$TE_PATH python3 -m build_tools.cudnn_frontend install \
+    || error_exit "Failed to install vendored nvidia-cudnn-frontend"
+
 NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*"
 
 pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements"