Produce BOLT-compatible libraries

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_ep_enabled,
 )
@@ -109,6 +110,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))):
@@ -125,5 +129,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,
         libraries=["nccl"],
     )

build_tools/pytorch.py

@@ -15,6 +15,7 @@
     cuda_version,
     get_cuda_include_dirs,
     debug_build_enabled,
+    get_bolt_build_flags,
     setup_mpi_flags,
 )
 from typing import List
@@ -75,6 +76,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()
@@ -121,6 +125,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],
     )

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"""

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, 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)

setup.py

@@ -77,6 +77,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(

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
+    $<$<COMPILE_LANGUAGE:CXX>:-fno-reorder-blocks-and-partition>
+    $<$<COMPILE_LANGUAGE:CXX>:-fno-jump-tables>
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-reorder-blocks-and-partition>
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-jump-tables>
+  )
+
+  if(NVTE_TARGET_IS_ARM64)
+    target_compile_options(
+      ${TARGET_NAME}
+      PRIVATE
+      $<$<COMPILE_LANGUAGE:CXX>:-mno-fix-cortex-a53-835769>
+      $<$<COMPILE_LANGUAGE:CXX>:-mno-fix-cortex-a53-843419>
+      $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-mno-fix-cortex-a53-835769>
+      $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-mno-fix-cortex-a53-843419>
+    )
+  endif()
+endfunction()
+
 # CUDA Toolkit
 find_package(CUDAToolkit REQUIRED)
 if (CUDAToolkit_VERSION VERSION_LESS 12.1)
@@ -137,6 +212,7 @@ function(get_nccl_version OUT_VERSION INCLUDE_DIR)
       PARENT_SCOPE)
 endfunction()
 
+set(NVTE_BOLT_MIN_NCCL_VERSION "2.30.4")
 get_nccl_version(NCCL_VERSION "${NCCL_INCLUDE_DIR}")
 
 function(find_cublasmp_version OUT_VERSION OUT_INCLUDE_DIR SEARCH_DIR)
@@ -332,13 +408,30 @@ 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(
   transformer_engine
   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
@@ -395,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()
@@ -452,11 +546,19 @@ endif()
 
 # -- NCCL EP (on by default, HT mode only) ---------------------------------
 # Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_WITH_NCCL_EP=0 in setup.py) to
-# skip NCCL EP entirely - useful on older images whose system NCCL is below
-# the 2.30.4 EP minimum.
+# skip NCCL EP entirely.
 option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON)
 if(NVTE_WITH_NCCL_EP)
-# SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize.
+# SM>=90 and the NCCL version are gated at runtime in EPBackend::initialize.
+# BOLT-compatible builds additionally require the matching NCCL symbols at load
+# time because they use eager symbol binding.
+if(NVTE_ENABLE_BOLT_COMPATIBLE
+   AND NCCL_VERSION VERSION_LESS NVTE_BOLT_MIN_NCCL_VERSION)
+  message(FATAL_ERROR
+          "BOLT-compatible NCCL EP builds require NCCL >= "
+          "${NVTE_BOLT_MIN_NCCL_VERSION}, but found NCCL ${NCCL_VERSION} in "
+          "${NCCL_INCLUDE_DIR}/nccl.h")
+endif()
 # -- NCCL EP headers --------------------------------------------------------
 # Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build
 # (auto-built by setup.py via build_nccl_ep_submodule).
@@ -471,10 +573,9 @@ endif()
 message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}")
 
 # -- libnccl_ep.a -----------------------------------------------------------
-# Statically linked into libtransformer_engine.so. EPBackend::initialize checks
-# NCCL >= 2.30.4 before any nccl_ep call, so the newer NCCL symbols nccl_ep
-# imports stay unresolved (and harmless) under default ELF lazy binding when
-# the gate trips. LD_BIND_NOW environments lose this property.
+# Statically linked into libtransformer_engine.so. BOLT-compatible builds use
+# eager symbol binding, so the build-time and runtime NCCL minimum must remain
+# synchronized with the symbols imported by this archive.
 set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib")
 find_file(NCCL_EP_LIB
     NAMES libnccl_ep.a