Initial Support Blackwell Build

README.rst

@@ -58,10 +58,12 @@ simplifying mixed precision training for users.
 Highlights
 ==========
 
-* Easy-to-use modules for building Transformer layers with FP8 support
+* Easy-to-use modules for building Transformer layers with FP8 support and FP4 support
 * Optimizations (e.g. fused kernels) for Transformer models
-* Support for FP8 on NVIDIA Hopper and NVIDIA Ada GPUs
+* Support for FP8 on NVIDIA Hopper, NVIDIA Ada GPUs and NVIDIA Blackwell GPUs
+* Support for FP4 on NVIDIA Blackwell GPUs
 * Support for optimizations across all precisions (FP16, BF16) on NVIDIA Ampere GPU architecture generations and later
+* Support for (FP32) on NVIDIA Turing GPU architecture.
 
 Examples
 ========
@@ -149,7 +151,7 @@ Installation
 Pre-requisites
 ^^^^^^^^^^^^^^^^^^^^
 * Linux x86_64
-* CUDA 12.0+ for Hopper and CUDA 12.1+ for Ada
+* CUDA 12.0+ for Hopper, CUDA 12.1+ for Ada and Cuda 12.7+ for Blackwell
 * NVIDIA Driver supporting CUDA 12.0 or later
 * cuDNN 8.1 or later
 * For fused attention, CUDA 12.1 or later, NVIDIA Driver supporting CUDA 12.1 or later, and cuDNN 8.9 or later.

build_tools/utils.py

@@ -190,7 +190,7 @@ def cuda_path() -> Tuple[str, str]:
 
 @functools.lru_cache(maxsize=None)
 def cuda_archs() -> str:
-    return os.getenv("NVTE_CUDA_ARCHS", "70;80;89;90")
+    return os.getenv("NVTE_CUDA_ARCHS", "70;80;89;90;100;120")
 
 
 def cuda_version() -> Tuple[int, ...]:

examples/jax/encoder/common.py

@@ -19,3 +19,10 @@ def is_fp8_supported():
     """Return if FP8 has hardware supported"""
     gpu_arch = get_device_compute_capability(0)
     return gpu_arch >= 90
+
+
+@lru_cache
+def is_fp4_supported():
+    """Return if FP4 has hardware supported"""
+    gpu_arch = get_device_compute_capability(0)
+    return gpu_arch >= 100

tests/cpp/CMakeLists.txt

@@ -5,7 +5,7 @@
 cmake_minimum_required(VERSION 3.18)
 
 if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
-  set(CMAKE_CUDA_ARCHITECTURES 70 80 90)
+  set(CMAKE_CUDA_ARCHITECTURES 70 80 90 100 120)
 endif()
 
 

transformer_engine/common/CMakeLists.txt

@@ -6,7 +6,7 @@ cmake_minimum_required(VERSION 3.21)
 
 # Language options
 if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
-  set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90)
+  set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90 100 120)
 endif()
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CUDA_STANDARD 17)