diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index ad6343d29b..1079accb87 100644 --- a/build_tools/build_ext.py +++ b/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"] diff --git a/ci/_utils.sh b/ci/_utils.sh index d09776b93a..8f634d4655 100644 --- a/ci/_utils.sh +++ b/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 } diff --git a/setup.py b/setup.py index 8465aac0d9..2c32d23833 100644 --- a/setup.py +++ b/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") diff --git a/tests/pytorch/triton_kernels/conftest.py b/tests/pytorch/triton_kernels/conftest.py index 9cbf0bffdd..78e354478c 100644 --- a/tests/pytorch/triton_kernels/conftest.py +++ b/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) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index dd4730e475..03bb65eda7 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/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") diff --git a/transformer_engine/common/ck_fused_attn/CMakeLists.txt b/transformer_engine/common/ck_fused_attn/CMakeLists.txt index 964ff6513c..14f6c9f543 100644 --- a/transformer_engine/common/ck_fused_attn/CMakeLists.txt +++ b/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 diff --git a/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp index a7221b6b78..33ba192ed2 100644 --- a/transformer_engine/common/comm_gemm_overlap/rocm_comm_gemm_overlap.cpp +++ b/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)); } } } diff --git a/transformer_engine/common/fused_attn_rocm/utils.cpp b/transformer_engine/common/fused_attn_rocm/utils.cpp index 2261ef6076..4c7b7f0fa4 100644 --- a/transformer_engine/common/fused_attn_rocm/utils.cpp +++ b/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(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<<>>(static_cast(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; } diff --git a/transformer_engine/common/gemm/kittens/CMakeLists.txt b/transformer_engine/common/gemm/kittens/CMakeLists.txt index 7d0db1af14..18db7c4c55 100644 --- a/transformer_engine/common/gemm/kittens/CMakeLists.txt +++ b/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} $ PARENT_SCOPE) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index 2e1d810484..7c4d0cbe54 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -41,8 +41,8 @@ static_assert(KITTENS_FLOAT16 == static_cast(transformer_engine::DType::kF static_assert(KITTENS_BFLOAT16 == static_cast(transformer_engine::DType::kBFloat16), "KittensDType out of sync with NVTEDType"); static_assert(KITTENS_FP8E4M3 == static_cast(transformer_engine::DType::kFloat8E4M3), "KittensDType out of sync with NVTEDType"); static_assert(KITTENS_FP8E5M2 == static_cast(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(KITTENS_BLOCK_SCALING_1D) == static_cast(NVTE_BLOCK_SCALING_1D), "KittensScalingMode out of sync with NVTEScalingMode"); +static_assert(static_cast(KITTENS_BLOCK_SCALING_2D) == static_cast(NVTE_BLOCK_SCALING_2D), "KittensScalingMode out of sync with NVTEScalingMode"); #endif namespace transformer_engine { diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h index c467f70c3a..e65c4be4b6 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h +++ b/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; } /*