Fix test filtering and reporting

build_tools/build_ext.py

@@ -74,11 +74,11 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None:
             ]
         configure_command += self.cmake_flags
 
-        import pybind11
-
-        pybind11_dir = Path(pybind11.__file__).resolve().parent
-        pybind11_dir = pybind11_dir / "share" / "cmake" / "pybind11"
-        configure_command.append(f"-Dpybind11_DIR={pybind11_dir}")
+        #ROCm: below variable is not used by CMake build. Leave it here for reference
+        #import pybind11
+        #pybind11_dir = Path(pybind11.__file__).resolve().parent
+        #pybind11_dir = pybind11_dir / "share" / "cmake" / "pybind11"
+        #configure_command.append(f"-Dpybind11_DIR={pybind11_dir}")
 
         # CMake build and install commands
         build_command = [_cmake_bin, "--build", build_dir, "--verbose"]

ci/_utils.sh

@@ -303,11 +303,13 @@ get_ctest_junitxml() {
 
 check_test_filter() {
     test -z "$TEST_FILTER" && return 0
+    set -f # disable globbing so patterns in masks are used as-is 
     for _tf in $TEST_FILTER; do
         case "$1" in
-        $_tf) return 0
+        $_tf) set +f; return 0
         esac
     done
+    set +f
     return 1
 }
 

setup.py

@@ -157,9 +157,9 @@ def setup_common_extension() -> CMakeExtension:
 
     # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is
     # targeted. Set NVTE_WITH_NCCL_EP=0 to force off.
-    # Disabled on ROCm
+    # Not used on ROCm
     if rocm_build():
-        cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF")
+        pass
     elif nccl_ep_enabled(archs):
         nccl_home = build_nccl_ep_submodule()
         cmake_flags.append(f"-DNCCL_INCLUDE_DIR={nccl_home}/include")

tests/pytorch/triton_kernels/conftest.py

@@ -35,14 +35,27 @@
 import torch
 
 
+_force_exit = False
+_exitstatus = 0
+
 @pytest.hookimpl(trylast=True)
 def pytest_sessionfinish(session, exitstatus):
+    global _force_exit, _exitstatus
+
     # Only ROCm hits the hsa_shut_down teardown segfault; leave CUDA/CPU exit
     # semantics (and their normal atexit cleanup) untouched.
-    if getattr(torch.version, "hip", None) is None:
-        return
-    # trylast ensures the junitxml plugin and te_ci_result_sink have already
-    # written their reports in this same hook before we hard-exit.
-    sys.stdout.flush()
-    sys.stderr.flush()
-    os._exit(0 if exitstatus == 0 else int(exitstatus))
+    if getattr(torch.version, "hip", None) is not None:
+        _force_exit = True
+        _exitstatus = int(exitstatus)
+
+# trylast ensures the junitxml plugin and te_ci_result_sink have already
+# written their reports in this same hook before we hard-exit.
+# unconfigure is the last hook called before pytest exits
+@pytest.hookimpl(trylast=True)
+def pytest_unconfigure(config):
+    global _force_exit, _exitstatus
+
+    if _force_exit:
+        sys.stdout.flush()
+        sys.stderr.flush()
+        os._exit(_exitstatus)

transformer_engine/common/CMakeLists.txt

@@ -128,7 +128,8 @@ set(CUTLASS_TOOLS_INCLUDE_DIR
   "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cutlass/tools/util/include")
 
 else()
-  set(CMAKE_CXX_STANDARD 17)
+  set(CMAKE_CXX_STANDARD 20)
+  set(CMAKE_HIP_STANDARD 20)
   #define GPU ARCH before enabling HIP to avoid potential failures with arch detection
   project(transformer_engine LANGUAGES HIP CXX)
 
@@ -848,11 +849,13 @@ endif()
 set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --threads ${BUILD_THREADS_PER_JOB}")
 message(STATUS "Threads per parallel build job: ${BUILD_THREADS_PER_JOB}")
 else()
-  set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -O3")
-  set(HIP_HCC_FLAGS "${CMAKE_HIP_FLAGS} -mavx2 -mf16c -mfma -std=c++17")
+  set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -O3 -Wno-unused-value -Wno-deprecated-anon-enum-enum-conversion")
+  set(HIP_HCC_FLAGS "${CMAKE_HIP_FLAGS} -mavx2 -mf16c -mfma")
   # Ask hcc to generate device code during compilation so we can use
   # host linker to link.
-  set(HIP_HCC_FLAGS "${HIP_HCC_FLAGS} -fno-gpu-rdc -Wno-defaulted-function-deleted")
+  # -Wno-unknown-warning-option is to suppress warning about -Wlifetime-safety-lifetimebound-violation
+  # pragma in CK headers which is not supported by the compiler used to build TE.
+  set(HIP_HCC_FLAGS "${HIP_HCC_FLAGS} -fno-gpu-rdc -Wno-defaulted-function-deleted -Wno-unused-value -Wno-unknown-warning-option")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${HIP_HCC_FLAGS}")
   # Hide non-necessary symbols in shared object.
   set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version")

transformer_engine/common/ck_fused_attn/CMakeLists.txt

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: MIT
 
 cmake_minimum_required(VERSION 3.21)
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 project(ck_fused_attn LANGUAGES HIP CXX)
 
 
@@ -220,9 +220,11 @@ endforeach()
 
 add_library(ck_fused_attn SHARED ${ck_fused_attn_SOURCES})
 set(CK_FUSED_ATTN_COMPILE_OPTIONS)
+#-Wno-unknown-warning-option is to suppress warning about -Wlifetime-safety-lifetimebound-violation
+#pragma in CK headers which is not supported by the compiler used to build TE.
 list(APPEND CK_FUSED_ATTN_COMPILE_OPTIONS
     -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=${CK_FUSED_ATTN_FLOAT_TO_BFLOAT16_DEFAULT}
-    -DENABLE_CK=1 -DFAV_NATIVE_ON=1)
+    -DENABLE_CK=1 -DFAV_NATIVE_ON=1 -Wno-unused-value -Wno-unknown-warning-option)
 
 # Public QoLA headers ship alongside the .so libs in ${__AITER_MHA_PATH}/../include
 # (emitted by qola.cli build, or copied from the QoLA build dir above for the

transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp

@@ -477,13 +477,13 @@ void CommOverlapP2PBase::rocm_split_overlap_ag(const TensorWrapper &A, bool tran
 
           NVTE_CHECK_CUDA(cudaMemcpyAsync(dstptr, srcptr, slice_bytes, cudaMemcpyDeviceToDevice, l_stream_send[r]));
           signal_val = ag_signal_base + step + 1;
-          hipStreamWriteValue64(l_stream_send[r], flagptr, signal_val, 0);
+          NVTE_CHECK_CUDA(hipStreamWriteValue64(l_stream_send[r], flagptr, signal_val, 0));
         }
 
         {
           void *flagptr = GET_RECV_PTR_BY_INDEX(prev_rank, _ub_comm, _ub_reg, r);
           signal_val = ag_signal_base + step + 1;
-          hipStreamWaitValue64(l_stream_recv[r], flagptr, signal_val, hipStreamWaitValueGte, 0xFFFFFFFFFFFFFFFF);
+          NVTE_CHECK_CUDA(hipStreamWaitValue64(l_stream_recv[r], flagptr, signal_val, hipStreamWaitValueGte, 0xFFFFFFFFFFFFFFFF));
         }
 
         NVTE_CHECK_CUDA(cudaEventRecord(get_event(next_recv_chunk_id, r), l_stream_recv[r]));
@@ -607,14 +607,14 @@ void CommOverlapP2PBase::rocm_split_overlap_rs(const TensorWrapper &A, bool tran
 
         NVTE_CHECK_CUDA(cudaMemcpyAsync(dstptr, srcptr, comm_bytes,
                                         cudaMemcpyDeviceToDevice, l_stream_send[comm_stream_id]));
-        hipStreamWriteValue64(l_stream_send[comm_stream_id], flagptr, signal_val, 0);
+        NVTE_CHECK_CUDA(hipStreamWriteValue64(l_stream_send[comm_stream_id], flagptr, signal_val, 0));
       }
 
       // Wait for incoming partial from chunk contributor
       {
         void *flagptr = GET_RECV_PTR_BY_INDEX(recv_rank, _ub_comm, _ub_reg, comm_stream_id);
-        hipStreamWaitValue64(l_stream_recv[comm_stream_id], flagptr, signal_val,
-                             hipStreamWaitValueGte, 0xFFFFFFFFFFFFFFFF);
+        NVTE_CHECK_CUDA(hipStreamWaitValue64(l_stream_recv[comm_stream_id], flagptr, signal_val,
+                                             hipStreamWaitValueGte, 0xFFFFFFFFFFFFFFFF));
       }
     }
   }

transformer_engine/common/fused_attn_rocm/utils.cpp

@@ -283,13 +283,13 @@ uint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t max_batc
   // workspace size requires 4 bytes
   uint32_t *dout = static_cast<uint32_t *>(workspace);
   uint32_t hout{};
-  cudaMemsetAsync(dout, 0, sizeof(uint32_t), stream);
+  NVTE_CHECK_CUDA(cudaMemsetAsync(dout, 0, sizeof(uint32_t), stream));
   constexpr int threads = 128;
   const int blocks = (max_batch_size - 1) / threads + 1; // ceil
   get_runtime_num_segments_kernel<<<blocks, threads, 0, stream>>>(static_cast<int32_t *>(cu_seqlen),
                                                                   max_batch_size, dout);
-  cudaMemcpyAsync(&hout, dout, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream);
-  cudaStreamSynchronize(stream);
+  NVTE_CHECK_CUDA(cudaMemcpyAsync(&hout, dout, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream));
+  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
   return hout;
 }
 

transformer_engine/common/gemm/kittens/CMakeLists.txt

@@ -12,19 +12,16 @@ foreach(_arch IN LISTS KITTENS_SUPPORTED_ARCHS)
     endif()
 endforeach()
 
-include(CheckCXXCompilerFlag)
-check_cxx_compiler_flag("-std=c++20" HAS_CXX20)
-
 if(NOT _kittens_enabled_archs)
     message(STATUS "HipKittens GEMM disabled (none of ${KITTENS_SUPPORTED_ARCHS} in CMAKE_HIP_ARCHITECTURES)")
     set(USE_HIPKITTENS_GEMM OFF PARENT_SCOPE)
-elseif(NOT HAS_CXX20)
-    message(WARNING "HipKittens GEMMs require C++20")
-    set(USE_HIPKITTENS_GEMM OFF PARENT_SCOPE)
 else()
     set(HIPKITTENS_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../3rdparty/hipkittens/include")
 
     set(CMAKE_CXX_STANDARD 20)
+    set(CMAKE_HIP_STANDARD 20)
+    set(CMAKE_CXX_STANDARD_REQUIRED ON)
+    set(CMAKE_HIP_STANDARD_REQUIRED ON)
     project(kittens_gemm LANGUAGES HIP CXX)
     find_package(hip)
 
@@ -72,7 +69,7 @@ else()
         target_include_directories(${_target} PRIVATE
             ${HIP_INCLUDE_DIRS} "${HIPKITTENS_INCLUDE_DIR}")
         target_compile_options(${_target} PRIVATE
-            -DKITTENS_${_tag_upper} -fno-gpu-rdc -O3 ${A_FLAGS})
+            -DKITTENS_${_tag_upper} -fno-gpu-rdc -O3 -Wno-unused-value ${A_FLAGS})
         target_link_libraries(${_target} PRIVATE hip::host hip::device)
 
         set(_kittens_arch_objs ${_kittens_arch_objs} $<TARGET_OBJECTS:${_target}> PARENT_SCOPE)

transformer_engine/common/gemm/rocm_gemm.cu

@@ -41,8 +41,8 @@ static_assert(KITTENS_FLOAT16  == static_cast<int>(transformer_engine::DType::kF
 static_assert(KITTENS_BFLOAT16 == static_cast<int>(transformer_engine::DType::kBFloat16), "KittensDType out of sync with NVTEDType");
 static_assert(KITTENS_FP8E4M3  == static_cast<int>(transformer_engine::DType::kFloat8E4M3), "KittensDType out of sync with NVTEDType");
 static_assert(KITTENS_FP8E5M2  == static_cast<int>(transformer_engine::DType::kFloat8E5M2), "KittensDType out of sync with NVTEDType");
-static_assert(KITTENS_BLOCK_SCALING_1D == NVTE_BLOCK_SCALING_1D, "KittensScalingMode out of sync with NVTEScalingMode");
-static_assert(KITTENS_BLOCK_SCALING_2D == NVTE_BLOCK_SCALING_2D, "KittensScalingMode out of sync with NVTEScalingMode");
+static_assert(static_cast<int>(KITTENS_BLOCK_SCALING_1D) == static_cast<int>(NVTE_BLOCK_SCALING_1D), "KittensScalingMode out of sync with NVTEScalingMode");
+static_assert(static_cast<int>(KITTENS_BLOCK_SCALING_2D) == static_cast<int>(NVTE_BLOCK_SCALING_2D), "KittensScalingMode out of sync with NVTEScalingMode");
 #endif
 
 namespace transformer_engine {

transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h

@@ -134,10 +134,6 @@ class CommOverlapCore {
 
   bool is_fp8_ubuf() { return _ubuf.element_size() == 1; }
 
-  virtual bool is_aggregate() {
-    NVTE_ERROR("Operation is not implemented.");
-  }
-
   virtual bool is_fused() { return false; }
 
   bool with_cublasmp() { return _with_cublasmp; }
@@ -430,8 +426,6 @@ class CommOverlapP2PBase : public CommOverlapCore {
                         bool use_split_accumulator, TensorWrapper &B_copy,
                         cudaStream_t stream_main) override;
 
-  bool is_aggregate() { return _aggregate; } // needed for rocm pathing
-
   bool is_fused() override { return _fused; }
 
   /*