From e757860fbad5da9e849572113c2694a64d2ed689 Mon Sep 17 00:00:00 2001 From: Brian Harrison Date: Thu, 15 May 2025 17:09:19 +0000 Subject: [PATCH 1/3] "Revert "CK MHA v2 solver"" --- CMakeLists.txt | 2 +- src/CMakeLists.txt | 3 +- src/include/miopen/mha/solvers.hpp | 22 -- src/problem.cpp | 4 +- src/solution.cpp | 34 +-- src/solver.cpp | 3 +- .../mha/mha_ck_fa_v2_solver_forward.cpp | 274 ------------------ test/gtest/cpu_multi_head_attention.hpp | 56 ++-- test/gtest/gpu_mha_backward.cpp | 2 +- test/gtest/graphapi_capi_mha_backward.cpp | 2 +- test/gtest/graphapi_capi_mha_common.hpp | 4 +- test/gtest/graphapi_capi_mha_forward.cpp | 4 +- test/gtest/graphapi_mha_cpp_common.hpp | 36 +-- test/gtest/mha_helper.hpp | 29 +- 14 files changed, 67 insertions(+), 408 deletions(-) delete mode 100644 src/solver/mha/mha_ck_fa_v2_solver_forward.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ff60c5966..8d9c866b4f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -328,7 +328,7 @@ add_compile_definitions($<$:HIP_COMPILER_FLAGS=${HIP_COMPI # HIP if( MIOPEN_BACKEND STREQUAL "HIP" OR MIOPEN_BACKEND STREQUAL "HIPOC" OR MIOPEN_BACKEND STREQUAL "HIPNOGPU") if(MIOPEN_USE_COMPOSABLEKERNEL) - find_package(composable_kernel 1.0.0 COMPONENTS device_other_operations device_gemm_operations device_conv_operations device_reduction_operations device_mha_operations) + find_package(composable_kernel 1.0.0 COMPONENTS device_other_operations device_gemm_operations device_conv_operations device_reduction_operations) endif() if( MIOPEN_BACKEND STREQUAL "HIPNOGPU") set(MIOPEN_MODE_NOGPU 1) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3b558a4eec..c950957dc9 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -309,7 +309,6 @@ set( MIOpen_Source solver/layernorm/forward_layernorm2d_ck.cpp solver/layernorm/forward_layernorm4d_ck.cpp solver/layernorm/forward_t5layernorm.cpp - solver/mha/mha_ck_fa_v2_solver_forward.cpp solver/mha/mha_solver_backward.cpp solver/mha/mha_solver_forward.cpp solver/multimarginloss/forward_multimarginloss.cpp @@ -859,7 +858,7 @@ target_include_directories(MIOpen PUBLIC ) if(MIOPEN_USE_COMPOSABLEKERNEL) -set(MIOPEN_CK_LINK_FLAGS composable_kernel::device_other_operations composable_kernel::device_gemm_operations composable_kernel::device_conv_operations composable_kernel::device_reduction_operations composable_kernel::device_mha_operations hip::host) +set(MIOPEN_CK_LINK_FLAGS composable_kernel::device_other_operations composable_kernel::device_gemm_operations composable_kernel::device_conv_operations composable_kernel::device_reduction_operations hip::host) endif() if(WIN32) diff --git a/src/include/miopen/mha/solvers.hpp b/src/include/miopen/mha/solvers.hpp index 55423f63c7..6bac473a71 100644 --- a/src/include/miopen/mha/solvers.hpp +++ b/src/include/miopen/mha/solvers.hpp @@ -77,28 +77,6 @@ struct MhaBackward final : MhaSolver MIOPEN_INTERNALS_EXPORT bool MayNeedWorkspace() const override; }; -struct MhaCKFlashAttentionV2Forward final : MhaSolver -{ - const std::string& SolverDbId() const override - { - return GetSolverDbId(); - } - - MIOPEN_INTERNALS_EXPORT bool - IsApplicable(const ExecutionContext& context, - const miopen::mha::ProblemDescription& problem) const override; - - MIOPEN_INTERNALS_EXPORT ConvSolution - GetSolution(const ExecutionContext& context, - const miopen::mha::ProblemDescription& problem) const override; - - MIOPEN_INTERNALS_EXPORT std::size_t - GetWorkspaceSize(const ExecutionContext& context, - const miopen::mha::ProblemDescription& problem) const override; - - MIOPEN_INTERNALS_EXPORT bool MayNeedWorkspace() const override; -}; - } // namespace mha } // namespace solver diff --git a/src/problem.cpp b/src/problem.cpp index 309e9800b7..c5b8b5f9ea 100644 --- a/src/problem.cpp +++ b/src/problem.cpp @@ -613,12 +613,10 @@ Problem::FindSolutionsImpl(const Handle& handle, const auto algo = AlgorithmName{"Mha"}; - static solver::mha::MhaCKFlashAttentionV2Forward mhaCKFAForwardSolver; static solver::mha::MhaForward mhaForwardSolver; static solver::mha::MhaBackward mhaBackwardSolver; - std::vector solvers = { - &mhaCKFAForwardSolver, &mhaForwardSolver, &mhaBackwardSolver}; + std::vector solvers = {&mhaForwardSolver, &mhaBackwardSolver}; for(auto solver : solvers) { diff --git a/src/solution.cpp b/src/solution.cpp index 89f79fcbbf..dc575643ef 100644 --- a/src/solution.cpp +++ b/src/solution.cpp @@ -400,32 +400,15 @@ void Solution::RunImpl(const Handle& handle, return; } - auto getSolution = [&](const ExecutionContext& ctx) { - auto solverId = GetSolver(); - solver::mha::MhaForward mhaForward; - solver::mha::MhaBackward mhaBackward; - solver::mha::MhaCKFlashAttentionV2Forward ckMhaForward; - - if(solverId == ckMhaForward.SolverDbId()) - { - return ckMhaForward.GetSolution(ctx, problem_description); - } - else if(solverId == mhaForward.SolverDbId()) - { - return mhaForward.GetSolution(ctx, problem_description); - } - else if(solverId == mhaBackward.SolverDbId()) - { - return mhaBackward.GetSolution(ctx, problem_description); - } - - MIOPEN_THROW("No MHA solver with matching SolverDbId of " + solverId.ToString()); - }; + solver::mha::MhaForward mhaForward; + solver::mha::MhaBackward mhaBackward; if(!kernels.empty()) { const auto ctx = ExecutionContext{&handle}; - const auto mha_solution = getSolution(ctx); + const auto mha_solution = GetSolver() == mhaForward.SolverDbId() + ? mhaForward.GetSolution(ctx, problem_description) + : mhaBackward.GetSolution(ctx, problem_description); auto kernel_handles = std::vector{std::begin(kernels), std::end(kernels)}; invoker = (*mha_solution.invoker_factory)(kernel_handles); @@ -442,8 +425,11 @@ void Solution::RunImpl(const Handle& handle, return; } - auto ctx = ExecutionContext{&handle}; - const auto mha_solution = getSolution(ctx); + auto ctx = ExecutionContext{&handle}; + + const auto mha_solution = GetSolver() == mhaForward.SolverDbId() + ? mhaForward.GetSolution(ctx, problem_description) + : mhaBackward.GetSolution(ctx, problem_description); invoker = handle.PrepareInvoker(*mha_solution.invoker_factory, mha_solution.construction_params); diff --git a/src/solver.cpp b/src/solver.cpp index 791206bd38..1db18a8a1a 100644 --- a/src/solver.cpp +++ b/src/solver.cpp @@ -705,7 +705,8 @@ inline SolverRegistrar::SolverRegistrar(IdRegistryData& registry) Primitive::MultiMarginLoss, multimarginloss::MultiMarginLossForward{}.SolverDbId()); - Register(registry, ++id, Primitive::Mha, mha::MhaCKFlashAttentionV2Forward{}.SolverDbId()); + // removed CK MHA solver + ++id; RegisterWithSolver( registry, ++id, conv::ConvWinoRageRxS<2, 3>{}, miopenConvolutionAlgoWinograd); diff --git a/src/solver/mha/mha_ck_fa_v2_solver_forward.cpp b/src/solver/mha/mha_ck_fa_v2_solver_forward.cpp deleted file mode 100644 index bf8563d89b..0000000000 --- a/src/solver/mha/mha_ck_fa_v2_solver_forward.cpp +++ /dev/null @@ -1,274 +0,0 @@ -/******************************************************************************* - * - * MIT License - * - * Copyright (c) 2024 Advanced Micro Devices, Inc. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - * - *******************************************************************************/ - -#include -#include -#include -#include -#include - -#if MIOPEN_USE_COMPOSABLEKERNEL -#include "ck_tile/ops/fmha.hpp" -#include "ck_tile/ops/fmha_fwd.hpp" -#include "ck/stream_config.hpp" -#endif - -MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_FA_CK_V2_FWD) - -namespace miopen { - -namespace solver { - -namespace mha { - -#if MIOPEN_USE_COMPOSABLEKERNEL -static std::string Convert(miopenDataType_t dataType) -{ - switch(dataType) - { - case miopenHalf: { - return "fp16"; - } - case miopenBFloat16: { - return "bfp16"; - } - default: { - MIOPEN_THROW("Unsupported datatype provided"); - } - } -} -#endif - -bool MhaCKFlashAttentionV2Forward::IsApplicable( - [[maybe_unused]] const ExecutionContext& context, - const miopen::mha::ProblemDescription& problem) const -{ -#if MIOPEN_USE_COMPOSABLEKERNEL - if(!problem.IsForward()) - { - return false; - } - - auto deviceName = context.GetStream().GetDeviceName(); - if(deviceName != "gfx942" && deviceName != "gfx90a") - { - return false; - } - - const auto& descsFwd = problem.GetDescsForward(); - auto [N_k, H_k, S_k, D_k] = miopen::tien<4>(descsFwd.kDesc.GetLengths()); - auto [N_stride_k, H_stride_k, S_stride_k, D_stride_k] = - miopen::tien<4>(descsFwd.kDesc.GetStrides()); - - auto [N_q, H_q, S_q, D_q] = miopen::tien<4>(descsFwd.qDesc.GetLengths()); - auto [N_stride_q, H_stride_q, S_stride_q, D_stride_q] = - miopen::tien<4>(descsFwd.qDesc.GetStrides()); - - auto [N_stride_v, H_stride_v, S_stride_v, D_stride_v] = - miopen::tien<4>(descsFwd.vDesc.GetStrides()); - - auto [N_stride_o, H_stride_o, S_stride_o, D_stride_o] = - miopen::tien<4>(descsFwd.oDesc.GetStrides()); - - return !env::disabled(MIOPEN_DEBUG_FA_CK_V2_FWD) // - && H_q == H_k // Replace with H_q % H_k == 0 once we add support for MQA & GQA. - && D_q <= 256 // - && D_q % 8 == 0 // - && descsFwd.kDesc.IsPacked() // - && descsFwd.qDesc.IsPacked() // - && descsFwd.vDesc.IsPacked() // - && descsFwd.oDesc.IsPacked() // - && descsFwd.kDesc.GetType() == miopenHalf // - && descsFwd.qDesc.GetType() == miopenHalf // - && descsFwd.vDesc.GetType() == miopenHalf // - && descsFwd.oDesc.GetType() == miopenHalf // - && D_stride_k == 1 // CK requires D stride as 1. - && D_stride_q == 1 && D_stride_v == 1 && D_stride_o == 1; -#else - return false; -#endif -} - -std::size_t MhaCKFlashAttentionV2Forward::GetWorkspaceSize( - [[maybe_unused]] const ExecutionContext& context, - [[maybe_unused]] const miopen::mha::ProblemDescription& problem) const -{ - return 0; -} - -ConvSolution -MhaCKFlashAttentionV2Forward::GetSolution([[maybe_unused]] const ExecutionContext& context, - const miopen::mha::ProblemDescription& problem) const -{ -#if MIOPEN_USE_COMPOSABLEKERNEL - auto result = ConvSolution{miopenStatusSuccess}; - result.workspace_sz = 0; - - const miopen::mha::MhaInputDescsForward& descsFwd = problem.GetDescsForward(); - auto [N_k, H_k, S_k, D_k] = miopen::tien<4>(descsFwd.kDesc.GetLengths()); - auto [N_stride_k, H_stride_k, S_stride_k, D_stride_k] = - miopen::tien<4>(descsFwd.kDesc.GetStrides()); - - auto [N_q, H_q, S_q, D_q] = miopen::tien<4>(descsFwd.qDesc.GetLengths()); - auto [N_stride_q, H_stride_q, S_stride_q, D_stride_q] = - miopen::tien<4>(descsFwd.qDesc.GetStrides()); - - auto [N_v, H_v, S_v, D_v] = miopen::tien<4>(descsFwd.vDesc.GetLengths()); - auto [N_stride_v, H_stride_v, S_stride_v, D_stride_v] = - miopen::tien<4>(descsFwd.vDesc.GetStrides()); - - auto [N_stride_o, H_stride_o, S_stride_o, D_stride_o] = - miopen::tien<4>(descsFwd.oDesc.GetStrides()); - - float scale_s = descsFwd.scale; - float scale_p = 1.0; - float scale_o = 1.0; - - fmha_fwd_traits fmha_traits; - fmha_traits.hdim_q = D_q; - fmha_traits.hdim_v = D_v; - fmha_traits.data_type = Convert(descsFwd.qDesc.GetType()); - fmha_traits.is_group_mode = false; - // is_v_rowmajor relates to the layout of the V tensor. Row major means NHSD, and Col major - // means NHDS. - fmha_traits.is_v_rowmajor = true; - fmha_traits.mask_type = mask_enum::no_mask; - fmha_traits.has_lse = false; - fmha_traits.do_fp8_static_quant = false; - fmha_traits.has_dropout = false; - fmha_traits.bias_type = bias_enum::no_bias; - - fmha_fwd_args fmha_args; - fmha_args.hdim_q = D_q; - fmha_args.hdim_v = D_v; - fmha_args.seqlen_k = S_k; - fmha_args.seqlen_q = S_q; - fmha_args.max_seqlen_q = S_q; - fmha_args.nhead_q = H_q; - fmha_args.nhead_k = H_k; - fmha_args.batch = N_q; - fmha_args.stride_q = S_stride_q; - fmha_args.stride_k = S_stride_k; - fmha_args.stride_v = S_stride_v; - fmha_args.stride_o = S_stride_o; - fmha_args.stride_bias = 0; - fmha_args.stride_randval = S_q; - fmha_args.nhead_stride_q = H_stride_q; - fmha_args.nhead_stride_k = H_stride_k; - fmha_args.nhead_stride_v = H_stride_v; - fmha_args.nhead_stride_o = H_stride_o; - fmha_args.nhead_stride_lse = S_q; - fmha_args.nhead_stride_bias = 0; - fmha_args.nhead_stride_randval = S_q * S_k; - fmha_args.batch_stride_q = N_stride_q; - fmha_args.batch_stride_k = N_stride_k; - fmha_args.batch_stride_v = N_stride_v; - fmha_args.batch_stride_o = N_stride_o; - fmha_args.batch_stride_lse = H_q * S_q; - fmha_args.batch_stride_bias = 0; - fmha_args.batch_stride_randval = H_q * S_q * S_k; - - // These are used for group mode, and we are in batch right now. - fmha_args.seqstart_q_ptr = nullptr; - fmha_args.seqstart_k_ptr = nullptr; - - // Batch does not support padding, and we aren't using kvcache yet. - fmha_args.seqlen_k_ptr = nullptr; - - fmha_args.s_randval = false; - fmha_args.scale_s = scale_s; - fmha_args.scale_p = scale_p; - fmha_args.scale_o = scale_o; - fmha_args.window_size_left = 0; - fmha_args.window_size_right = 0; - fmha_args.mask_type = static_cast(fmha_traits.mask_type); - - result.invoker_factory = [=](const std::vector&) { - return [=](const Handle& handle_, const AnyInvokeParams& raw_params) { - decltype(auto) params = raw_params.CastTo(); - const auto& dataFwd = params.GetDataForward(); - - fmha_fwd_traits fmha_runtime_traits = fmha_traits; - fmha_fwd_args fmha_runtime_args = fmha_args; - - fmha_runtime_args.q_ptr = dataFwd.qData; - fmha_runtime_args.k_ptr = dataFwd.kData; - fmha_runtime_args.v_ptr = dataFwd.vData; - fmha_runtime_args.o_ptr = dataFwd.oData; - fmha_runtime_args.rand_val_ptr = nullptr; - fmha_runtime_args.bias_ptr = nullptr; - fmha_runtime_args.lse_ptr = nullptr; - - // Top-left causal mask - if(dataFwd.mask == miopenMhaMask_t::miopenMhaMaskCausal) - { - fmha_runtime_traits.mask_type = mask_enum::mask_top_left; - fmha_runtime_args.mask_type = - static_cast(mask_enum::mask_top_left); - fmha_runtime_args.window_size_left = -1; - fmha_runtime_args.window_size_right = 0; - } - - fmha_runtime_traits.has_dropout = false; - float probability = 0; - - // TODO : Change API to take in probability value as host side value instead of device - // pointer to match CK API. Calling a blocking hipMemcpy will cause issues with stream, - // and isn't async. - - fmha_runtime_args.p_drop = probability; - // fmha_runtime_args.drop_seed_offset = - // std::make_pair(dataFwd.dropoutSeedData), - // dataFwd.dropoutOffsetData); - - // using dataFwd.dropoutSeedData gpu pointer was causing compiler error - // since dropout is disabled for now, placing 0. - fmha_runtime_args.drop_seed_offset = std::make_pair(0, 0); - - // Create stream_config, and set it to not time kernel. - ck_tile::stream_config stream_config; - stream_config.stream_id_ = handle_.GetStream(); - - { - HipEventProfiler profiler(handle_); - fmha_fwd(fmha_runtime_traits, fmha_runtime_args, stream_config); - } - }; - }; - - return result; -#else - return ConvSolution{miopenStatusNotImplemented}; -#endif -} - -bool MhaCKFlashAttentionV2Forward::MayNeedWorkspace() const { return false; } - -} // namespace mha - -} // namespace solver - -} // namespace miopen diff --git a/test/gtest/cpu_multi_head_attention.hpp b/test/gtest/cpu_multi_head_attention.hpp index f155f6c41b..52fc8bcde6 100644 --- a/test/gtest/cpu_multi_head_attention.hpp +++ b/test/gtest/cpu_multi_head_attention.hpp @@ -81,16 +81,16 @@ struct CPUMHATest : public ::testing::TestWithParam if constexpr(std::is_same_v) { // forward - MultiHeadAttentionForwardf32(q_val, - k_val, - v_val, - q_dot_k_transpose, - softmax, - attn_max, - z_sum, - aMax_S, - aMax_O, - multi_head_attention); + MultiHeadAttentionf32(q_val, + k_val, + v_val, + q_dot_k_transpose, + softmax, + attn_max, + z_sum, + aMax_S, + aMax_O, + multi_head_attention); Concat(multi_head_attention, concatinated_attention); @@ -138,24 +138,24 @@ struct CPUMHATest : public ::testing::TestWithParam ScaleMult(v_val, v_scale, v_val_fp8); // forward - MultiHeadAttentionForwardfp8(q_val_fp8, - k_val_fp8, - v_val_fp8, - softmax, // fp32 - attn_max, - z_sum, - q_descale, - k_descale, - v_descale, - s_descale, - s_scale, - o_scale, - 0.0f, - 0, - 0, - aMax_S, - aMax_O, - multi_head_attention); + MultiHeadAttentionfp8(q_val_fp8, + k_val_fp8, + v_val_fp8, + softmax, // fp32 + attn_max, + z_sum, + q_descale, + k_descale, + v_descale, + s_descale, + s_scale, + o_scale, + 0.0f, + 0, + 0, + aMax_S, + aMax_O, + multi_head_attention); Concat(multi_head_attention, final_transformed_attention); ScaleMult(final_transformed_attention, o_descale, concatinated_attention); diff --git a/test/gtest/gpu_mha_backward.cpp b/test/gtest/gpu_mha_backward.cpp index 2d77510c51..b05a578ff1 100644 --- a/test/gtest/gpu_mha_backward.cpp +++ b/test/gtest/gpu_mha_backward.cpp @@ -205,7 +205,7 @@ class Test_Bwd_Mha : public testing::TestWithParam // proper O, M and zInv tensors are required for backward pass. // randomly generated M and zInv may lead to nan\inf values - test::cpu::MultiHeadAttentionForwardfp8( + test::cpu::MultiHeadAttentionfp8( std::get>(tensors[miopenTensorMhaQ]->m_cpu_tensor), std::get>(tensors[miopenTensorMhaK]->m_cpu_tensor), std::get>(tensors[miopenTensorMhaV]->m_cpu_tensor), diff --git a/test/gtest/graphapi_capi_mha_backward.cpp b/test/gtest/graphapi_capi_mha_backward.cpp index a0f8aba0c9..127608d49a 100644 --- a/test/gtest/graphapi_capi_mha_backward.cpp +++ b/test/gtest/graphapi_capi_mha_backward.cpp @@ -95,7 +95,7 @@ class MhaBackwardTest : public MhaCommonTest // proper O, M and zInv tensors are required for backward pass. // randomly generated M and zInv may lead to nan\inf values - test::cpu::MultiHeadAttentionForwardfp8( + test::cpu::MultiHeadAttentionfp8( GetTensor(m_realTensorMap[miopenTensorMhaQ]->m_tensorVariant), GetTensor(m_realTensorMap[miopenTensorMhaK]->m_tensorVariant), GetTensor(m_realTensorMap[miopenTensorMhaV]->m_tensorVariant), diff --git a/test/gtest/graphapi_capi_mha_common.hpp b/test/gtest/graphapi_capi_mha_common.hpp index b91b4951c6..38d2841334 100644 --- a/test/gtest/graphapi_capi_mha_common.hpp +++ b/test/gtest/graphapi_capi_mha_common.hpp @@ -213,7 +213,7 @@ miopenDataType_t GetMainType() } else if(std::is_same_v) { - return miopenBFloat8_fnuz; + return miopenBFloat8; } else if(std::is_same_v) { @@ -666,7 +666,7 @@ class MhaCommonTest : public testing::TestWithParamm_tensorVariant = tensor{n, h, s, d}; + tensorDataPtr->m_tensorVariant = tensor{n, h, s, d}; } else if(dtype == miopenHalf) { diff --git a/test/gtest/graphapi_capi_mha_forward.cpp b/test/gtest/graphapi_capi_mha_forward.cpp index c233dcd84a..7ad847dd5c 100644 --- a/test/gtest/graphapi_capi_mha_forward.cpp +++ b/test/gtest/graphapi_capi_mha_forward.cpp @@ -230,7 +230,7 @@ class MhaForwardTest : public MhaCommonTest return it->second->m_tensorVariant; }; - test::cpu::MultiHeadAttentionForwardfp8( + test::cpu::MultiHeadAttentionfp8( GetTensor(lookup(miopenTensorMhaQ)), GetTensor(lookup(miopenTensorMhaK)), GetTensor(lookup(miopenTensorMhaV)), @@ -365,6 +365,4 @@ inline auto GetCases() INSTANTIATE_TEST_SUITE_P(Smoke, GPU_MhaForward_FP32, GetCases()); -INSTANTIATE_TEST_SUITE_P(Smoke, GPU_MhaForward_FP16, GetCases()); - INSTANTIATE_TEST_SUITE_P(Smoke, GPU_MhaForward_FP8, GetCases()); diff --git a/test/gtest/graphapi_mha_cpp_common.hpp b/test/gtest/graphapi_mha_cpp_common.hpp index 94c7e5896c..e4089f887e 100644 --- a/test/gtest/graphapi_mha_cpp_common.hpp +++ b/test/gtest/graphapi_mha_cpp_common.hpp @@ -388,24 +388,24 @@ class MhaGraphTestBase CpuMhaFwdOut out(n, h, s, d); - test::cpu::MultiHeadAttentionForwardfp8(lookup_f("Q"), - lookup_f("K"), - lookup_f("V"), - out.mSoftMax, - out.mM, - out.mZinv, - lookup_f("DSCL_Q")[0], - lookup_f("DSCL_K")[0], - lookup_f("DSCL_V")[0], - lookup_f("DSCL_S")[0], - lookup_f("SCL_S")[0], - lookup_f("SCL_O")[0], - lookup_f("RND_PRB")[0], - lookup_i("RND_SD")[0], - lookup_i("RND_OFF")[0], - out.mAmaxS, - out.mAmaxO, - out.mO); + test::cpu::MultiHeadAttentionfp8(lookup_f("Q"), + lookup_f("K"), + lookup_f("V"), + out.mSoftMax, + out.mM, + out.mZinv, + lookup_f("DSCL_Q")[0], + lookup_f("DSCL_K")[0], + lookup_f("DSCL_V")[0], + lookup_f("DSCL_S")[0], + lookup_f("SCL_S")[0], + lookup_f("SCL_O")[0], + lookup_f("RND_PRB")[0], + lookup_i("RND_SD")[0], + lookup_i("RND_OFF")[0], + out.mAmaxS, + out.mAmaxO, + out.mO); return out; } diff --git a/test/gtest/mha_helper.hpp b/test/gtest/mha_helper.hpp index 02294e2d53..00804db430 100644 --- a/test/gtest/mha_helper.hpp +++ b/test/gtest/mha_helper.hpp @@ -426,34 +426,7 @@ void MultiHeadAttentionForwardfp8(const tensor& q_val, } template -void MultiHeadAttentionForwardfp16(const tensor& q_val, - const tensor& k_val, - const tensor& v_val, - tensor& softmax, - tensor& attn_max, - tensor& Z_sum, - tensor& multi_head_attention, - const tensor* optional_bias = - nullptr) // pointer to optional bias, nullptr if not provided -{ - auto inputLengths = q_val.desc.GetLengths(); - inputLengths[3] = inputLengths[2]; // NHSD converting to NHSS - tensor q_dot_k_transpose(inputLengths); - - Dot_4D_4D_T(q_val, k_val, q_dot_k_transpose); - - if(optional_bias != nullptr) - { - PointWiseAdd(q_dot_k_transpose, *optional_bias, q_dot_k_transpose); - } - - SoftMax(q_dot_k_transpose, softmax, attn_max, Z_sum); - - Dot_4D_4D(softmax, v_val, multi_head_attention); -} - -template -void MultiHeadAttentionForwardf32( +void MultiHeadAttentionf32( const tensor& q_val, const tensor& k_val, const tensor& v_val, From 3ae9caf74f2b36cf1253375443a0fed3c122d6fa Mon Sep 17 00:00:00 2001 From: Brian Harrison Date: Thu, 15 May 2025 18:05:41 +0000 Subject: [PATCH 2/3] Resolve merge conflicts --- test/gtest/cpu_multi_head_attention.hpp | 4 +- test/gtest/gpu_mha_backward.cpp | 2 +- test/gtest/gpu_mha_forward.cpp | 110 ---------------------- test/gtest/graphapi_capi_mha_backward.cpp | 2 +- test/gtest/graphapi_capi_mha_common.hpp | 13 +-- test/gtest/graphapi_capi_mha_forward.cpp | 50 +--------- test/gtest/graphapi_mha_cpp_common.hpp | 2 +- test/gtest/mha_helper.hpp | 2 +- 8 files changed, 9 insertions(+), 176 deletions(-) diff --git a/test/gtest/cpu_multi_head_attention.hpp b/test/gtest/cpu_multi_head_attention.hpp index 52fc8bcde6..d0b296d3e8 100644 --- a/test/gtest/cpu_multi_head_attention.hpp +++ b/test/gtest/cpu_multi_head_attention.hpp @@ -81,7 +81,7 @@ struct CPUMHATest : public ::testing::TestWithParam if constexpr(std::is_same_v) { // forward - MultiHeadAttentionf32(q_val, + MultiHeadAttentionForwardf32(q_val, k_val, v_val, q_dot_k_transpose, @@ -138,7 +138,7 @@ struct CPUMHATest : public ::testing::TestWithParam ScaleMult(v_val, v_scale, v_val_fp8); // forward - MultiHeadAttentionfp8(q_val_fp8, + MultiHeadAttentionForwardfp8(q_val_fp8, k_val_fp8, v_val_fp8, softmax, // fp32 diff --git a/test/gtest/gpu_mha_backward.cpp b/test/gtest/gpu_mha_backward.cpp index b05a578ff1..2d77510c51 100644 --- a/test/gtest/gpu_mha_backward.cpp +++ b/test/gtest/gpu_mha_backward.cpp @@ -205,7 +205,7 @@ class Test_Bwd_Mha : public testing::TestWithParam // proper O, M and zInv tensors are required for backward pass. // randomly generated M and zInv may lead to nan\inf values - test::cpu::MultiHeadAttentionfp8( + test::cpu::MultiHeadAttentionForwardfp8( std::get>(tensors[miopenTensorMhaQ]->m_cpu_tensor), std::get>(tensors[miopenTensorMhaK]->m_cpu_tensor), std::get>(tensors[miopenTensorMhaV]->m_cpu_tensor), diff --git a/test/gtest/gpu_mha_forward.cpp b/test/gtest/gpu_mha_forward.cpp index 7789d5ec65..06d764f2fa 100644 --- a/test/gtest/gpu_mha_forward.cpp +++ b/test/gtest/gpu_mha_forward.cpp @@ -107,42 +107,6 @@ inline std::vector GetFullTestCases() }; } -inline std::vector GetFp16SmokeCases() -{ - return { - {2, 1, 1, 256, 0.0f}, - {2, 2, 65, 128, 0.0f}, - {3, 2, 257, 64, 0.0f}, - {3, 5, 528, 32, 0.0f}, - {3, 7, 712, 16, 0.0f}, - {5, 3, 1111, 8, 0.0f}, - }; -} - -inline std::vector GetFp16FullTestCases() -{ - return { - {3, 11, 1731, 8, 0.0f}, - {2049, 5, 7, 8, 0.0f}, - {5, 2000, 32, 8, 0.0f}, - {3, 9, 1407, 16, 0.0f}, - {1027, 5, 21, 16, 0.0f}, - {5, 1040, 32, 24, 0.0f}, - {3, 7, 1212, 32, 0.0f}, - {550, 5, 16, 40, 0.0f}, - {5, 550, 40, 48, 0.0f}, - {2, 9, 1057, 64, 0.0f}, - {250, 3, 19, 72, 0.0f}, - {5, 230, 27, 80, 0.0f}, - {2, 5, 920, 128, 0.0f}, - {111, 2, 27, 136, 0.0f}, - {3, 110, 22, 152, 0.0f}, - {2, 4, 600, 224, 0.0f}, - {57, 1, 63, 232, 0.0f}, - {2, 65, 18, 256, 0.0f}, - }; -} - } // namespace template @@ -399,74 +363,6 @@ class GPU_Fwd_Mha_FP32 : public Test_Fwd_Mha { }; -class GPU_Fwd_Mha_FP16 : public Test_Fwd_Mha -{ - void SetUp() override - { - if(!IsTestSupportedByDevice(Gpu::gfx90A | Gpu::gfx94X)) - { - GTEST_SKIP() << "FP16 is unsupported on this HW"; - } - - Test_Fwd_Mha::SetUp(); - - if(dropout != 0.0f) - { - GTEST_SKIP() << "Dropout not currently supported for FP16"; - } - } - - void RunReference(const tensor& q_val, - const tensor& k_val, - const tensor& v_val, - tensor& softmax, - tensor& attn_max, - tensor& Z_sum, - [[maybe_unused]] float q_descale, - [[maybe_unused]] float k_descale, - [[maybe_unused]] float v_descale, - [[maybe_unused]] float s_descale, - [[maybe_unused]] float s_scale, - [[maybe_unused]] float o_scale, - [[maybe_unused]] float dropout_rate, - [[maybe_unused]] uint64_t seed, - [[maybe_unused]] uint64_t offset, - [[maybe_unused]] float& aMax_S, - [[maybe_unused]] float& aMax_O, - tensor& output) override - { - test::cpu::MultiHeadAttentionForwardfp16( - q_val, k_val, v_val, softmax, attn_max, Z_sum, output); - } - - void VerifyResults(const Handle& handle) override - { - auto GetResult = [this, &handle](miopenTensorArgumentId_t id) { - auto& tensorStructPtr = tensors[id]; - auto& cpu_tensor = std::get>(tensorStructPtr->m_cpu_tensor); - - cpu_tensor.data = handle.Read(tensorStructPtr->m_gpu_buffer, - cpu_tensor.data.size()); - - return cpu_tensor; - }; - - const double errorThreshold = 4e-4; - double oError = miopen::rms_range(oDesc_ref, GetResult(miopenTensorMhaO)); - - if(dropout > 0.0f) - { - // Due to GPU version using a different dropout generator we will compare to CPU without - // dropout and verify that dropout causes a large difference when comparing results. - EXPECT_GT(oError, errorThreshold); - } - else - { - EXPECT_LT(oError, errorThreshold); - } - } -}; - class GPU_Fwd_Mha_FP8 : public Test_Fwd_Mha { void SetUp() override @@ -488,12 +384,6 @@ INSTANTIATE_TEST_SUITE_P(Smoke, GPU_Fwd_Mha_FP32, testing::ValuesIn(GetSmokeCase INSTANTIATE_TEST_SUITE_P(Full, GPU_Fwd_Mha_FP32, testing::ValuesIn(GetFullTestCases())); GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(GPU_Fwd_Mha_FP32); -TEST_P(GPU_Fwd_Mha_FP16, Test_float) { return GPU_Fwd_Mha_FP16::TestBody(); }; - -INSTANTIATE_TEST_SUITE_P(Smoke, GPU_Fwd_Mha_FP16, testing::ValuesIn(GetFp16SmokeCases())); -INSTANTIATE_TEST_SUITE_P(Full, GPU_Fwd_Mha_FP16, testing::ValuesIn(GetFp16FullTestCases())); -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(GPU_Fwd_Mha_FP16); - #if WORKAROUND_SWDEV_528878 #define MHA_FP8_TEST_NAME DISABLED_Test_float8 #else diff --git a/test/gtest/graphapi_capi_mha_backward.cpp b/test/gtest/graphapi_capi_mha_backward.cpp index 127608d49a..a0f8aba0c9 100644 --- a/test/gtest/graphapi_capi_mha_backward.cpp +++ b/test/gtest/graphapi_capi_mha_backward.cpp @@ -95,7 +95,7 @@ class MhaBackwardTest : public MhaCommonTest // proper O, M and zInv tensors are required for backward pass. // randomly generated M and zInv may lead to nan\inf values - test::cpu::MultiHeadAttentionfp8( + test::cpu::MultiHeadAttentionForwardfp8( GetTensor(m_realTensorMap[miopenTensorMhaQ]->m_tensorVariant), GetTensor(m_realTensorMap[miopenTensorMhaK]->m_tensorVariant), GetTensor(m_realTensorMap[miopenTensorMhaV]->m_tensorVariant), diff --git a/test/gtest/graphapi_capi_mha_common.hpp b/test/gtest/graphapi_capi_mha_common.hpp index 38d2841334..c241ce4080 100644 --- a/test/gtest/graphapi_capi_mha_common.hpp +++ b/test/gtest/graphapi_capi_mha_common.hpp @@ -168,7 +168,6 @@ class DescriptorWrapper using TensorVariant = std::variant, tensor, tensor, - tensor, tensor>; template @@ -213,11 +212,7 @@ miopenDataType_t GetMainType() } else if(std::is_same_v) { - return miopenBFloat8; - } - else if(std::is_same_v) - { - return miopenHalf; + return miopenBFloat8_fnuz; } assert(false); @@ -666,11 +661,7 @@ class MhaCommonTest : public testing::TestWithParamm_tensorVariant = tensor{n, h, s, d}; - } - else if(dtype == miopenHalf) - { - tensorDataPtr->m_tensorVariant = tensor{n, h, s, d}; + tensorDataPtr->m_tensorVariant = tensor{n, h, s, d}; } else { diff --git a/test/gtest/graphapi_capi_mha_forward.cpp b/test/gtest/graphapi_capi_mha_forward.cpp index 7ad847dd5c..71d1915652 100644 --- a/test/gtest/graphapi_capi_mha_forward.cpp +++ b/test/gtest/graphapi_capi_mha_forward.cpp @@ -230,7 +230,7 @@ class MhaForwardTest : public MhaCommonTest return it->second->m_tensorVariant; }; - test::cpu::MultiHeadAttentionfp8( + test::cpu::MultiHeadAttentionForwardfp8( GetTensor(lookup(miopenTensorMhaQ)), GetTensor(lookup(miopenTensorMhaK)), GetTensor(lookup(miopenTensorMhaV)), @@ -296,55 +296,7 @@ class GPU_MhaForward_FP8 : public MhaForwardTest } }; -class GPU_MhaForward_FP16 : public MhaForwardTest -{ - void SetUp() override - { - if(!IsTestSupportedByDevice(Gpu::gfx90A | Gpu::gfx94X)) - { - GTEST_SKIP() << "FP16 is unsupported on this HW"; - } - - MhaForwardTest::SetUp(); - - if(m_bernulliProbability != 0.0f) - { - GTEST_SKIP() << "Dropout not currently supported for FP16"; - } - } - - void RunCPUverify(const miopen::Handle& handle) override - { - auto softmaxRef = tensor{m_testN, m_testH, m_testS, m_testS}; - auto oDescRef = tensor{m_testN, m_testH, m_testS, m_testD}; - auto mDescRef = tensor{m_testN, m_testH, m_testS, 1}; - auto zInvDescRef = tensor{m_testN, m_testH, m_testS, 1}; - - auto lookup = [this](const int64_t id) -> TensorVariant& { - auto it = m_realTensorMap.find(id); - assert(it != m_realTensorMap.cend()); - return it->second->m_tensorVariant; - }; - - test::cpu::MultiHeadAttentionForwardfp16( - GetTensor(lookup(miopenTensorMhaQ)), - GetTensor(lookup(miopenTensorMhaK)), - GetTensor(lookup(miopenTensorMhaV)), - softmaxRef, - mDescRef, - zInvDescRef, - oDescRef); - - const double errorThreshold = 4e-4; - - auto oRes = GetResult(miopenTensorMhaO, handle); - double oError = miopen::rms_range(oDescRef, oRes); - EXPECT_LT(oError, errorThreshold); - } -}; - TEST_P(GPU_MhaForward_FP32, TestFloat) { Run(); } -TEST_P(GPU_MhaForward_FP16, TestFloat16) { Run(); } #if WORKAROUND_SWDEV_528878 #define MHA_FP8_TEST_NAME DISABLED_Test_float8 diff --git a/test/gtest/graphapi_mha_cpp_common.hpp b/test/gtest/graphapi_mha_cpp_common.hpp index e4089f887e..6cda18a846 100644 --- a/test/gtest/graphapi_mha_cpp_common.hpp +++ b/test/gtest/graphapi_mha_cpp_common.hpp @@ -388,7 +388,7 @@ class MhaGraphTestBase CpuMhaFwdOut out(n, h, s, d); - test::cpu::MultiHeadAttentionfp8(lookup_f("Q"), + test::cpu::MultiHeadAttentionForwardfp8(lookup_f("Q"), lookup_f("K"), lookup_f("V"), out.mSoftMax, diff --git a/test/gtest/mha_helper.hpp b/test/gtest/mha_helper.hpp index 00804db430..7a00c5cc69 100644 --- a/test/gtest/mha_helper.hpp +++ b/test/gtest/mha_helper.hpp @@ -426,7 +426,7 @@ void MultiHeadAttentionForwardfp8(const tensor& q_val, } template -void MultiHeadAttentionf32( +void MultiHeadAttentionForwardf32( const tensor& q_val, const tensor& k_val, const tensor& v_val, From 3acd60053352864aab56ead0f5cb66003614c649 Mon Sep 17 00:00:00 2001 From: Brian Harrison Date: Thu, 15 May 2025 18:17:44 +0000 Subject: [PATCH 3/3] Formatting --- test/gtest/cpu_multi_head_attention.hpp | 52 ++++++++++++------------- test/gtest/graphapi_capi_mha_common.hpp | 6 +-- test/gtest/graphapi_mha_cpp_common.hpp | 34 ++++++++-------- 3 files changed, 45 insertions(+), 47 deletions(-) diff --git a/test/gtest/cpu_multi_head_attention.hpp b/test/gtest/cpu_multi_head_attention.hpp index d0b296d3e8..f155f6c41b 100644 --- a/test/gtest/cpu_multi_head_attention.hpp +++ b/test/gtest/cpu_multi_head_attention.hpp @@ -82,15 +82,15 @@ struct CPUMHATest : public ::testing::TestWithParam { // forward MultiHeadAttentionForwardf32(q_val, - k_val, - v_val, - q_dot_k_transpose, - softmax, - attn_max, - z_sum, - aMax_S, - aMax_O, - multi_head_attention); + k_val, + v_val, + q_dot_k_transpose, + softmax, + attn_max, + z_sum, + aMax_S, + aMax_O, + multi_head_attention); Concat(multi_head_attention, concatinated_attention); @@ -139,23 +139,23 @@ struct CPUMHATest : public ::testing::TestWithParam // forward MultiHeadAttentionForwardfp8(q_val_fp8, - k_val_fp8, - v_val_fp8, - softmax, // fp32 - attn_max, - z_sum, - q_descale, - k_descale, - v_descale, - s_descale, - s_scale, - o_scale, - 0.0f, - 0, - 0, - aMax_S, - aMax_O, - multi_head_attention); + k_val_fp8, + v_val_fp8, + softmax, // fp32 + attn_max, + z_sum, + q_descale, + k_descale, + v_descale, + s_descale, + s_scale, + o_scale, + 0.0f, + 0, + 0, + aMax_S, + aMax_O, + multi_head_attention); Concat(multi_head_attention, final_transformed_attention); ScaleMult(final_transformed_attention, o_descale, concatinated_attention); diff --git a/test/gtest/graphapi_capi_mha_common.hpp b/test/gtest/graphapi_capi_mha_common.hpp index c241ce4080..76daa4fbdf 100644 --- a/test/gtest/graphapi_capi_mha_common.hpp +++ b/test/gtest/graphapi_capi_mha_common.hpp @@ -165,10 +165,8 @@ class DescriptorWrapper std::vector m_refsToKeep; }; -using TensorVariant = std::variant, - tensor, - tensor, - tensor>; +using TensorVariant = + std::variant, tensor, tensor, tensor>; template tensor& GetTensor(TensorVariant& var) diff --git a/test/gtest/graphapi_mha_cpp_common.hpp b/test/gtest/graphapi_mha_cpp_common.hpp index 6cda18a846..94c7e5896c 100644 --- a/test/gtest/graphapi_mha_cpp_common.hpp +++ b/test/gtest/graphapi_mha_cpp_common.hpp @@ -389,23 +389,23 @@ class MhaGraphTestBase CpuMhaFwdOut out(n, h, s, d); test::cpu::MultiHeadAttentionForwardfp8(lookup_f("Q"), - lookup_f("K"), - lookup_f("V"), - out.mSoftMax, - out.mM, - out.mZinv, - lookup_f("DSCL_Q")[0], - lookup_f("DSCL_K")[0], - lookup_f("DSCL_V")[0], - lookup_f("DSCL_S")[0], - lookup_f("SCL_S")[0], - lookup_f("SCL_O")[0], - lookup_f("RND_PRB")[0], - lookup_i("RND_SD")[0], - lookup_i("RND_OFF")[0], - out.mAmaxS, - out.mAmaxO, - out.mO); + lookup_f("K"), + lookup_f("V"), + out.mSoftMax, + out.mM, + out.mZinv, + lookup_f("DSCL_Q")[0], + lookup_f("DSCL_K")[0], + lookup_f("DSCL_V")[0], + lookup_f("DSCL_S")[0], + lookup_f("SCL_S")[0], + lookup_f("SCL_O")[0], + lookup_f("RND_PRB")[0], + lookup_i("RND_SD")[0], + lookup_i("RND_OFF")[0], + out.mAmaxS, + out.mAmaxO, + out.mO); return out; }