From 88e429f29b21f1730fbb13e5c1a3be5815b03e8b Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 30 Jul 2026 18:41:34 +0000 Subject: [PATCH] [Common] Enable arch support for Rubin 1. Gate arch-specific Rubin codegen on CUDA 13.4+. 2. Drop redundant cute includes from the cast-fusion kernels that caused compile errors on `cute::maybe_ifence()` and `cute::maybe_warp_switch()`. 3. Enable stochastic rounding for `ArchSpecific<107>`. Signed-off-by: Alp Dener --- transformer_engine/common/CMakeLists.txt | 3 +++ ...h_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu | 2 -- .../hadamard_transform/group_hadamard_transform_cast_fusion.cu | 2 -- .../group_row_cast_col_hadamard_transform_cast_fusion.cu | 2 -- .../hadamard_transform/hadamard_transform_cast_fusion.cu | 2 -- .../row_cast_col_hadamard_transform_cast_fusion.cu | 2 -- transformer_engine/common/util/ptx.cuh | 2 +- 7 files changed, 4 insertions(+), 11 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 13880e0d61..9734e24fbe 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -49,6 +49,9 @@ if(NOT arch_100_index EQUAL -1) if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9) list(APPEND NVTE_SPECIFIC_ARCHS "103a") endif() + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.4) + list(APPEND NVTE_SPECIFIC_ARCHS "107a") + endif() endif() # Check for architecture 101 (if we see this we are in toolkit <= 12.9) diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu index d5dbd0bd82..aaa931f138 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu index e2325dd0fc..4b1435f9eb 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu index cab0b38589..aa6d4ddf34 100644 --- a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu index 50b9f63bdd..433da1f0f0 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 479922a9bf..7a3ff029a4 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -13,8 +13,6 @@ #include #include -#include -#include #include #include "common/common.h" diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2814aa3490..2351ddae71 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -124,7 +124,7 @@ constexpr bool is_supported_arch() { NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>, ptx::FamilySpecific<110>, \ ptx::FamilySpecific<120>) #define ARCH_HAS_STOCHASTIC_ROUNDING \ - NVTE_CUDA_ARCH_MATCHES(ptx::ArchSpecific<100>, ptx::ArchSpecific<103>) + NVTE_CUDA_ARCH_MATCHES(ptx::ArchSpecific<100>, ptx::ArchSpecific<103>, ptx::ArchSpecific<107>) // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init __device__ __forceinline__ void mbarrier_init(uint64_t *mbar, const uint32_t count) {