From 0982a20f8cf59c606f6d7150139bd3f3aa1e32d5 Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Fri, 14 Aug 2026 13:32:44 -0500 Subject: [PATCH 1/4] Produce BOLT-compatible libraries Signed-off-by: Fred Heinecke --- build_tools/jax.py | 5 ++ build_tools/pytorch.py | 5 ++ build_tools/utils.py | 36 +++++++++ docs/envvars.rst | 6 ++ setup.py | 42 +++++++++-- transformer_engine/common/CMakeLists.txt | 93 ++++++++++++++++++++++++ 6 files changed, 181 insertions(+), 6 deletions(-) diff --git a/build_tools/jax.py b/build_tools/jax.py index 031432e6f9..e54c668c5b 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -15,6 +15,7 @@ all_files_in_dir, cudnn_frontend_include_path, debug_build_enabled, + get_bolt_build_flags, setup_mpi_flags, nccl_include_path, nccl_lib_path, @@ -113,6 +114,9 @@ def setup_jax_extension( else: cxx_flags.append("-g0") + bolt_cxx_flags, linker_flags = get_bolt_build_flags() + cxx_flags.extend(bolt_cxx_flags) + setup_mpi_flags(include_dirs, cxx_flags) if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): @@ -135,5 +139,6 @@ def setup_jax_extension( sources=[str(path) for path in sources], include_dirs=[str(path) for path in include_dirs], extra_compile_args=cxx_flags, + extra_link_args=linker_flags, **kwargs, ) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 98331ccbd8..18342dc14b 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -15,6 +15,7 @@ cuda_version, get_cuda_include_dirs, debug_build_enabled, + get_bolt_build_flags, nccl_ep_enabled, setup_mpi_flags, ) @@ -76,6 +77,9 @@ def setup_pytorch_extension( else: cxx_flags.append("-g0") + bolt_cxx_flags, linker_flags = get_bolt_build_flags() + cxx_flags.extend(bolt_cxx_flags) + # Version-dependent CUDA options try: version = cuda_version() @@ -122,6 +126,7 @@ def setup_pytorch_extension( sources=[str(src) for src in sources], include_dirs=[str(inc) for inc in include_dirs], extra_compile_args={"cxx": cxx_flags}, + extra_link_args=linker_flags, libraries=[str(lib) for lib in libraries], library_dirs=[str(lib_dir) for lib_dir in library_dirs], ) diff --git a/build_tools/utils.py b/build_tools/utils.py index 3b02fcea2b..0038de3efc 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -43,6 +43,42 @@ def debug_build_enabled() -> bool: return bool(int(os.getenv("NVTE_BUILD_DEBUG", "0"))) +@functools.lru_cache(maxsize=None) +def bolt_compatible_build_enabled() -> bool: + """Whether to build host ELF libraries with BOLT-compatible options.""" + configured = os.getenv("NVTE_ENABLE_BOLT_COMPATIBLE") + if configured is None: + enabled = platform.system() == "Linux" and platform.machine().lower() in ( + "aarch64", + "arm64", + ) + else: + enabled = bool(int(configured)) + + if enabled and platform.system() != "Linux": + raise RuntimeError("NVTE_ENABLE_BOLT_COMPATIBLE is only supported on Linux") + return enabled + + +def get_bolt_build_flags() -> Tuple[List[str], List[str]]: + """BOLT-compatible host compiler and linker flags.""" + if not bolt_compatible_build_enabled(): + return [], [] + + compiler_flags = ["-fno-reorder-blocks-and-partition", "-fno-jump-tables"] + linker_flags = ["-Wl,--emit-relocs", "-Wl,-z,now"] + if platform.machine().lower() in ("aarch64", "arm64"): + compiler_flags.extend( + [ + "-mno-fix-cortex-a53-835769", + "-mno-fix-cortex-a53-843419", + ] + ) + # The Cortex-A53 843419 workaround is applied by the linker. + linker_flags.append("-mno-fix-cortex-a53-843419") + return compiler_flags, linker_flags + + @functools.lru_cache(maxsize=None) def get_max_jobs_for_parallel_build() -> int: """Number of parallel jobs for Nina build""" diff --git a/docs/envvars.rst b/docs/envvars.rst index 97eaed5ddc..8a9f295800 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -60,6 +60,12 @@ Build Configuration :Default: None :Description: Path to the CMake build directory for incremental builds. If set, CMake will use this directory for build artifacts. +.. envvar:: NVTE_ENABLE_BOLT_COMPATIBLE + + :Type: ``int`` (0 or 1) + :Default: ``1`` on Linux Arm64 (AArch64), ``0`` otherwise + :Description: Build the core, NCCL EP, PyTorch, and JAX host ELFs with LLVM BOLT-compatible compiler and linker flags. On Arm64, this also disables the Cortex-A53 835769 and 843419 errata workarounds as required by BOLT. Set to ``0`` to disable the default on Arm64, or ``1`` to opt in on another supported Linux architecture. When using :envvar:`NVTE_CMAKE_BUILD_DIR`, use a fresh build directory after changing this setting. + .. envvar:: NVTE_RELEASE_BUILD :Type: ``int`` (0 or 1) diff --git a/setup.py b/setup.py index 944cba060e..20b5c83b00 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ cuda_home_path, cuda_version, cudnn_frontend_include_path, + get_bolt_build_flags, get_frameworks, remove_dups, min_python_version_str, @@ -79,6 +80,11 @@ def setup_common_extension() -> CMakeExtension: if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))): cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON") + bolt_compatible = os.getenv("NVTE_ENABLE_BOLT_COMPATIBLE") + if bolt_compatible is not None: + bolt_compatible = "ON" if bool(int(bolt_compatible)) else "OFF" + cmake_flags.append(f"-DNVTE_ENABLE_BOLT_COMPATIBLE={bolt_compatible}") + if bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))): cmake_flags.append("-DNVTE_WITH_CUBLASMP=ON") cublasmp_dir = os.getenv("CUBLASMP_HOME") or metadata.distribution( @@ -266,12 +272,36 @@ def build_nccl_ep_submodule() -> str: env["NCCL_HOME"] = nccl_home env["NCCL_EP_BUILDDIR"] = str(build_dir) - prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None - if not nccl_ep_shared_lib.exists() or prev_gencode != gencode: - if nccl_ep_shared_lib.exists() and prev_gencode != gencode: + bolt_cxx_flags, bolt_linker_flags = get_bolt_build_flags() + nvcc_host_flags = [f"-Xcompiler={flag}" for flag in bolt_cxx_flags] + nvcc_linker_flags = [] + if bolt_linker_flags: + nvcc_linker_flags.extend( + ["-Xlinker=--emit-relocs", "-Xlinker=-z", "-Xlinker=now"] + ) + if "-mno-fix-cortex-a53-843419" in bolt_linker_flags: + nvcc_linker_flags.append("-Xlinker=--no-fix-cortex-a53-843419") + + def append_env_flags(name: str, flags: List[str]) -> None: + if flags: + env[name] = " ".join([env.get(name, ""), *flags]).strip() + + append_env_flags("CXXFLAGS", bolt_cxx_flags) + append_env_flags("NVCC_PREPEND_FLAGS", nvcc_host_flags) + append_env_flags("LDFLAGS", nvcc_linker_flags) + + build_signature = "\n".join( + ( + f"gencode={gencode}", + f"bolt_cxx_flags={' '.join(bolt_cxx_flags)}", + f"bolt_linker_flags={' '.join(nvcc_linker_flags)}", + ) + ) + previous_signature = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None + if not nccl_ep_shared_lib.exists() or previous_signature != build_signature: + if nccl_ep_shared_lib.exists() and previous_signature != build_signature: print( - f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); " - "rebuilding NCCL EP libraries" + "[NCCL EP] build configuration changed; rebuilding NCCL EP libraries" ) subprocess.check_call( ["make", "-C", "nccl_ep", "clean"], @@ -286,7 +316,7 @@ def build_nccl_ep_submodule() -> str: env=env, ) gencode_stamp.parent.mkdir(parents=True, exist_ok=True) - gencode_stamp.write_text(gencode) + gencode_stamp.write_text(build_signature) return nccl_home diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 3503941cf4..da6bacccba 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -15,6 +15,81 @@ endif() # Transformer Engine library project(transformer_engine LANGUAGES CUDA CXX) +# BOLT-compatible builds are enabled by default on Linux Arm64, where function +# layout has the largest impact on Grace CPUs. Other platforms remain opt-in. +string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _nvte_system_processor) +set(NVTE_TARGET_IS_ARM64 OFF) +if(_nvte_system_processor MATCHES "^(aarch64|arm64)$") + set(NVTE_TARGET_IS_ARM64 ON) +endif() +set(_nvte_bolt_compatible_default OFF) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NVTE_TARGET_IS_ARM64) + set(_nvte_bolt_compatible_default ON) +endif() +option(NVTE_ENABLE_BOLT_COMPATIBLE + "Build host ELF objects with LLVM BOLT-compatible compile and link flags" + ${_nvte_bolt_compatible_default}) +unset(_nvte_bolt_compatible_default) +unset(_nvte_system_processor) + +if(NVTE_ENABLE_BOLT_COMPATIBLE) + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message(FATAL_ERROR + "NVTE_ENABLE_BOLT_COMPATIBLE is only supported for Linux ELF builds") + endif() + + include(CheckCXXCompilerFlag) + check_cxx_compiler_flag("-fno-reorder-blocks-and-partition" + NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION) + check_cxx_compiler_flag("-fno-jump-tables" + NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES) + if(NVTE_TARGET_IS_ARM64) + check_cxx_compiler_flag("-mno-fix-cortex-a53-835769" + NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769) + check_cxx_compiler_flag("-mno-fix-cortex-a53-843419" + NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419) + endif() + if(NOT NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION + OR NOT NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES + OR (NVTE_TARGET_IS_ARM64 + AND (NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769 + OR NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419))) + message(FATAL_ERROR + "The host C++ compiler does not support the flags required for " + "NVTE_ENABLE_BOLT_COMPATIBLE") + endif() + + # BOLT consumes the regular symbol table and emitted relocations. Shadow the + # cached strip tool for this configure without changing it for later builds. + set(CMAKE_STRIP "") +endif() + +function(nvte_enable_bolt_compatible_compile_options TARGET_NAME) + if(NOT NVTE_ENABLE_BOLT_COMPATIBLE) + return() + endif() + + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-fno-reorder-blocks-and-partition> + $<$:-fno-jump-tables> + $<$:-Xcompiler=-fno-reorder-blocks-and-partition> + $<$:-Xcompiler=-fno-jump-tables> + ) + + if(NVTE_TARGET_IS_ARM64) + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-mno-fix-cortex-a53-835769> + $<$:-mno-fix-cortex-a53-843419> + $<$:-Xcompiler=-mno-fix-cortex-a53-835769> + $<$:-Xcompiler=-mno-fix-cortex-a53-843419> + ) + endif() +endfunction() + # CUDA Toolkit find_package(CUDAToolkit REQUIRED) if (CUDAToolkit_VERSION VERSION_LESS 12.1) @@ -333,6 +408,7 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) endforeach() add_library(transformer_engine SHARED ${transformer_engine_SOURCES}) +nvte_enable_bolt_compatible_compile_options(transformer_engine) # This is TE-specific and should not apply to all targets target_link_options( @@ -340,6 +416,22 @@ target_link_options( PRIVATE "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version" ) +if(NVTE_ENABLE_BOLT_COMPATIBLE) + target_link_options( + transformer_engine + PRIVATE + "LINKER:--emit-relocs" + "LINKER:-z,now" + ) + if(NVTE_TARGET_IS_ARM64) + # The Cortex-A53 843419 workaround is applied by the linker. + target_link_options( + transformer_engine + PRIVATE + "LINKER:--no-fix-cortex-a53-843419" + ) + endif() +endif() # Disable CMake's automatic architecture flag injection. # All architectures are handled explicitly via per-source COMPILE_OPTIONS @@ -396,6 +488,7 @@ endif() option(NVTE_ENABLE_NVSHMEM "Compile with NVSHMEM library" OFF) if (NVTE_ENABLE_NVSHMEM) add_subdirectory(nvshmem_api) + nvte_enable_bolt_compatible_compile_options(nvshmemapi) target_link_libraries(transformer_engine PUBLIC nvshmemapi) target_include_directories(transformer_engine PUBLIC ${NVSHMEMAPI_INCLUDE_DIR}) endif() From da9df7e54fa944a0098e1d86af3f628e5e89c7be Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:37:33 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- setup.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/setup.py b/setup.py index 20b5c83b00..1be20aaa4c 100644 --- a/setup.py +++ b/setup.py @@ -276,9 +276,7 @@ def build_nccl_ep_submodule() -> str: nvcc_host_flags = [f"-Xcompiler={flag}" for flag in bolt_cxx_flags] nvcc_linker_flags = [] if bolt_linker_flags: - nvcc_linker_flags.extend( - ["-Xlinker=--emit-relocs", "-Xlinker=-z", "-Xlinker=now"] - ) + nvcc_linker_flags.extend(["-Xlinker=--emit-relocs", "-Xlinker=-z", "-Xlinker=now"]) if "-mno-fix-cortex-a53-843419" in bolt_linker_flags: nvcc_linker_flags.append("-Xlinker=--no-fix-cortex-a53-843419") @@ -300,9 +298,7 @@ def append_env_flags(name: str, flags: List[str]) -> None: previous_signature = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None if not nccl_ep_shared_lib.exists() or previous_signature != build_signature: if nccl_ep_shared_lib.exists() and previous_signature != build_signature: - print( - "[NCCL EP] build configuration changed; rebuilding NCCL EP libraries" - ) + print("[NCCL EP] build configuration changed; rebuilding NCCL EP libraries") subprocess.check_call( ["make", "-C", "nccl_ep", "clean"], cwd=str(nccl_root), From 36bd78a8737ec28aa78c290335295eed21cb0d7f Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Tue, 1 Sep 2026 15:41:59 -0500 Subject: [PATCH 3/4] Fix target arch detection Signed-off-by: Fred Heinecke --- build_tools/utils.py | 31 ++++++++++++++++++++++++++----- docs/envvars.rst | 4 ++-- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index 0038de3efc..f4befc473f 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -9,6 +9,7 @@ import importlib import os import re +import shlex import shutil import subprocess import sys @@ -43,15 +44,35 @@ def debug_build_enabled() -> bool: return bool(int(os.getenv("NVTE_BUILD_DEBUG", "0"))) +@functools.lru_cache(maxsize=None) +def build_target_arch() -> str: + """CPU architecture targeted by the configured C++ compiler.""" + cxx = shlex.split(os.getenv("CXX", "c++")) + try: + result = subprocess.run( + [*cxx, "-dumpmachine"], + capture_output=True, + check=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return platform.machine().lower() + + target = result.stdout.strip().split("-", 1)[0] + return target.lower() or platform.machine().lower() + + +def target_is_arm64() -> bool: + """Whether the configured C++ compiler targets Arm64.""" + return build_target_arch() in ("aarch64", "arm64") + + @functools.lru_cache(maxsize=None) def bolt_compatible_build_enabled() -> bool: """Whether to build host ELF libraries with BOLT-compatible options.""" configured = os.getenv("NVTE_ENABLE_BOLT_COMPATIBLE") if configured is None: - enabled = platform.system() == "Linux" and platform.machine().lower() in ( - "aarch64", - "arm64", - ) + enabled = platform.system() == "Linux" and target_is_arm64() else: enabled = bool(int(configured)) @@ -67,7 +88,7 @@ def get_bolt_build_flags() -> Tuple[List[str], List[str]]: compiler_flags = ["-fno-reorder-blocks-and-partition", "-fno-jump-tables"] linker_flags = ["-Wl,--emit-relocs", "-Wl,-z,now"] - if platform.machine().lower() in ("aarch64", "arm64"): + if target_is_arm64(): compiler_flags.extend( [ "-mno-fix-cortex-a53-835769", diff --git a/docs/envvars.rst b/docs/envvars.rst index 8a9f295800..1c5e4b2d4b 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -63,8 +63,8 @@ Build Configuration .. envvar:: NVTE_ENABLE_BOLT_COMPATIBLE :Type: ``int`` (0 or 1) - :Default: ``1`` on Linux Arm64 (AArch64), ``0`` otherwise - :Description: Build the core, NCCL EP, PyTorch, and JAX host ELFs with LLVM BOLT-compatible compiler and linker flags. On Arm64, this also disables the Cortex-A53 835769 and 843419 errata workarounds as required by BOLT. Set to ``0`` to disable the default on Arm64, or ``1`` to opt in on another supported Linux architecture. When using :envvar:`NVTE_CMAKE_BUILD_DIR`, use a fresh build directory after changing this setting. + :Default: ``1`` when building on Linux and the configured C++ compiler targets Arm64 (AArch64), ``0`` otherwise + :Description: Build the core, NCCL EP, PyTorch, and JAX host ELFs with LLVM BOLT-compatible compiler and linker flags. On an Arm64 target, this also disables the Cortex-A53 835769 and 843419 errata workarounds as required by BOLT. Set to ``0`` to disable the default on Arm64, or ``1`` to opt in on another supported Linux architecture. When using :envvar:`NVTE_CMAKE_BUILD_DIR`, use a fresh build directory after changing this setting. .. envvar:: NVTE_RELEASE_BUILD From 8b8ac68bae560725daa7656230c6837f51ec8d0d Mon Sep 17 00:00:00 2001 From: Fred Heinecke Date: Tue, 1 Sep 2026 15:51:53 -0500 Subject: [PATCH 4/4] fail fast when failing to detect arch Signed-off-by: Fred Heinecke --- build_tools/utils.py | 14 +++++++++++--- docs/envvars.rst | 2 +- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index f4befc473f..8be784c640 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -55,11 +55,19 @@ def build_target_arch() -> str: check=True, text=True, ) - except (OSError, subprocess.CalledProcessError): - return platform.machine().lower() + except (OSError, subprocess.CalledProcessError) as e: + raise RuntimeError( + "Could not determine the C++ compiler target architecture with " + f"`{' '.join(cxx)} -dumpmachine`" + ) from e target = result.stdout.strip().split("-", 1)[0] - return target.lower() or platform.machine().lower() + if not target: + raise RuntimeError( + "Could not determine the C++ compiler target architecture: " + f"`{' '.join(cxx)} -dumpmachine` returned no target" + ) + return target.lower() def target_is_arm64() -> bool: diff --git a/docs/envvars.rst b/docs/envvars.rst index 1c5e4b2d4b..4c807528f5 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -64,7 +64,7 @@ Build Configuration :Type: ``int`` (0 or 1) :Default: ``1`` when building on Linux and the configured C++ compiler targets Arm64 (AArch64), ``0`` otherwise - :Description: Build the core, NCCL EP, PyTorch, and JAX host ELFs with LLVM BOLT-compatible compiler and linker flags. On an Arm64 target, this also disables the Cortex-A53 835769 and 843419 errata workarounds as required by BOLT. Set to ``0`` to disable the default on Arm64, or ``1`` to opt in on another supported Linux architecture. When using :envvar:`NVTE_CMAKE_BUILD_DIR`, use a fresh build directory after changing this setting. + :Description: Build the core, NCCL EP, PyTorch, and JAX host ELFs with LLVM BOLT-compatible compiler and linker flags. On an Arm64 target, this also disables the Cortex-A53 835769 and 843419 errata workarounds as required by BOLT. Set to ``0`` to disable the default on Arm64, or ``1`` to opt in on another supported Linux architecture. Automatic target detection requires the configured ``CXX`` compiler to support ``-dumpmachine``; the build fails instead of falling back to the host architecture if the target cannot be determined. When using :envvar:`NVTE_CMAKE_BUILD_DIR`, use a fresh build directory after changing this setting. .. envvar:: NVTE_RELEASE_BUILD