From 9acbfc6a299e9a5a43037fd68b709ca3e291167d Mon Sep 17 00:00:00 2001 From: cloudhan Date: Tue, 11 Apr 2023 13:20:44 +0800 Subject: [PATCH] ROCm MHA (#15279) Add MultiHeadAttention for ROCm EP. **Before:** ``` 'engine': 'onnxruntime' 'version': '1.15.0' 'height': 512 'width': 512 'steps': 50 'batch_size': 1 'batch_count': 5 'num_prompts': 1 'average_latency': 3.878769588470459 'median_latency': 3.8792178630828857 'first_run_memory_MB': -1 'second_run_memory_MB': -1 'model_name': 'runwayml/stable-diffusion-v1-5' 'directory': './sd-v1-5-onnx-fp16-nomha' 'provider': 'ROCMExecutionProvider' 'disable_safety_checker': True ``` **After:** ``` 'engine': 'onnxruntime' 'version': '1.15.0' 'height': 512 'width': 512 'steps': 50 'batch_size': 1 'batch_count': 5 'num_prompts': 1 'average_latency': 2.364924430847168 'median_latency': 2.3650705814361572 'first_run_memory_MB': -1 'second_run_memory_MB': -1 'model_name': 'runwayml/stable-diffusion-v1-5' 'directory': './sd-v1-5-onnx-fp16' 'provider': 'ROCMExecutionProvider' 'disable_safety_checker': True ``` --- cmake/onnxruntime_unittests.cmake | 3 + .../contrib_ops/cpu/bert/attention_base.cc | 1 + .../contrib_ops/cpu/bert/attention_common.h | 6 +- .../cpu/bert/multihead_attention_helper.h | 8 + .../contrib_ops/rocm/bert/attention.cu | 1 + .../contrib_ops/rocm/bert/attention_impl.h | 1 + ...ed_gemm_softmax_gemm_permute_pipelines.cuh | 114 ++++++++++-- .../rocm/bert/multihead_attention.cu | 127 ++++++++++++++ .../rocm/bert/multihead_attention.h | 30 ++++ .../contrib_ops/rocm/rocm_contrib_kernels.cc | 5 + .../kernels/_kernel_explorer.pyi | 6 + .../kernels/gemm_softmax_gemm_permute_test.py | 164 ++++++++++++++---- .../kernels/rocm/gemm_softmax_gemm_permute.cu | 47 +++-- .../multihead_attention_op_test.cc | 30 +++- 14 files changed, 467 insertions(+), 76 deletions(-) create mode 100644 onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu create mode 100644 onnxruntime/contrib_ops/rocm/bert/multihead_attention.h diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 308d725ddc..d972fc5458 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -827,6 +827,9 @@ if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS) target_link_libraries(onnxruntime_test_all PRIVATE onnxruntime_language_interop onnxruntime_pyop) endif() if (onnxruntime_USE_ROCM) + if (onnxruntime_USE_COMPOSABLE_KERNEL) + target_compile_definitions(onnxruntime_test_all PRIVATE USE_COMPOSABLE_KERNEL) + endif() target_compile_options(onnxruntime_test_all PRIVATE -D__HIP_PLATFORM_AMD__=1 -D__HIP_PLATFORM_HCC__=1) target_include_directories(onnxruntime_test_all PRIVATE ${onnxruntime_ROCM_HOME}/hipfft/include ${onnxruntime_ROCM_HOME}/include ${onnxruntime_ROCM_HOME}/hiprand/include ${onnxruntime_ROCM_HOME}/rocrand/include ${CMAKE_CURRENT_BINARY_DIR}/amdgpu/onnxruntime ${CMAKE_CURRENT_BINARY_DIR}/amdgpu/orttraining) endif() diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_base.cc b/onnxruntime/contrib_ops/cpu/bert/attention_base.cc index 00e843ffb9..e743ad8b14 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_base.cc +++ b/onnxruntime/contrib_ops/cpu/bert/attention_base.cc @@ -261,6 +261,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape, output_parameters->mask_type = mask_type; output_parameters->broadcast_res_pos_bias = broadcast_res_pos_bias; output_parameters->pass_past_in_kv = false; + output_parameters->qkv_format = Q_K_V_BNSH; } return Status::OK(); diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_common.h b/onnxruntime/contrib_ops/cpu/bert/attention_common.h index 1b52ff2a0f..6614e9cbcb 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_common.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_common.h @@ -21,8 +21,9 @@ enum AttentionMaskType { }; enum AttentionQkvFormat { - Q_K_V_BNSH, // for unfused attention - Q_K_V_BSNH, // for memory efficient attention, or format of query, key and value for MultiHeadAttention + UNKNOWN, // enum value not set, or depends on qkv projection implementation details + Q_K_V_BNSH, // for non-packed qkv, permuted + Q_K_V_BSNH, // for non-packed qkv, not permuted, used by memory efficient attention or MultiHeadAttention QKV_BSN3H, // for TRT fused attention, qkv are packed Q_K_V_BNSH_QKV_BS3NH, // for TRT fused causal attention, data has two formats (qkv is 3BNSH, gemm_buffer is BS3NH) Q_KV_BSNH_BSN2H, // for TRT fused cross attention, kv are packed @@ -60,6 +61,7 @@ struct AttentionParameters { float mask_filter_value; float scale; AttentionMaskType mask_type; + AttentionQkvFormat qkv_format; }; // Parameters deduced from node attributes and inputs/outputs. diff --git a/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h index 1d2bb9d235..f337a9e00b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h @@ -47,6 +47,7 @@ Status CheckInputs(const T* query, // value (V) : None // bias (Q/K/V) : None + AttentionQkvFormat qkv_format; const auto& query_dims = query->Shape().GetDims(); if (query_dims.size() != 3 && query_dims.size() != 5) { @@ -148,6 +149,7 @@ Status CheckInputs(const T* query, "Input 'query' and 'key' shall have same dim 2 (hidden_size)"); } + qkv_format = Q_K_V_BSNH; kv_sequence_length = static_cast(key_dims[1]); } else if (key_dims.size() == 5) { if (static_cast(key_dims[2]) != num_heads || static_cast(key_dims[3]) != 2 || static_cast(key_dims[4]) != head_size) { @@ -159,6 +161,7 @@ Status CheckInputs(const T* query, return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Expect 'value' be none when 'key' has packed kv format."); } + qkv_format = Q_KV_BSNH_BSN2H; kv_sequence_length = static_cast(key_dims[1]); } else { // key_dims.size() == 4 (cross-attention with past_key) if (static_cast(key_dims[1]) != num_heads || static_cast(key_dims[3]) != head_size) { @@ -167,6 +170,7 @@ Status CheckInputs(const T* query, "Expect 'key' shape (batch_size, num_heads, kv_sequence_length, head_size) for past_key"); } + qkv_format = UNKNOWN; kv_sequence_length = static_cast(key_dims[2]); } } else { // packed QKV @@ -179,6 +183,8 @@ Status CheckInputs(const T* query, ONNXRUNTIME, INVALID_ARGUMENT, "Expect 'query' shape (batch_size, kv_sequence_length, num_heads, 3, head_size) for packed kv"); } + + qkv_format = QKV_BSN3H; } if (bias != nullptr) { @@ -281,6 +287,7 @@ Status CheckInputs(const T* query, } } + // TODO: ORT_RETURN_IF(qkv_format == UNKNOWN, "Unrecognized QKV format"); if (parameters != nullptr) { AttentionParameters* output_parameters = reinterpret_cast(parameters); output_parameters->batch_size = batch_size; @@ -302,6 +309,7 @@ Status CheckInputs(const T* query, output_parameters->scale = scale; output_parameters->broadcast_res_pos_bias = broadcast_res_pos_bias; output_parameters->pass_past_in_kv = pass_past_in_kv; + output_parameters->qkv_format = qkv_format; } if (max_threads_per_block > 0 && num_heads > max_threads_per_block) { diff --git a/onnxruntime/contrib_ops/rocm/bert/attention.cu b/onnxruntime/contrib_ops/rocm/bert/attention.cu index c064196ec5..96c05eaa90 100644 --- a/onnxruntime/contrib_ops/rocm/bert/attention.cu +++ b/onnxruntime/contrib_ops/rocm/bert/attention.cu @@ -63,6 +63,7 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { device_prop.maxThreadsPerBlock, past_seq_len)); ORT_ENFORCE(attn.sequence_length == attn.kv_sequence_length); // self attention + ORT_ENFORCE(attn.qkv_format == Q_K_V_BNSH); // non-packed, permuted TensorShapeVector output_shape(3); output_shape[0] = static_cast(attn.batch_size); diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_impl.h b/onnxruntime/contrib_ops/rocm/bert/attention_impl.h index c8ba9ec875..f7a5eb05ef 100644 --- a/onnxruntime/contrib_ops/rocm/bert/attention_impl.h +++ b/onnxruntime/contrib_ops/rocm/bert/attention_impl.h @@ -11,6 +11,7 @@ namespace onnxruntime { namespace contrib { namespace rocm { + size_t GetAttentionScratchSize( size_t element_size, int batch_size, diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh index 86b761986a..4b38ccafc5 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh @@ -11,6 +11,8 @@ T: total sequence length N: num of heads H: head dimension +The following use qkv_format == Q_K_V_BNSH as a example: + BN: B*N, which is the batch size of GEMMs. NOTE: To be disambiguated with batch size of Attention Op In QKV projection (prior to this pipeline): @@ -85,6 +87,64 @@ inline int3 Get2DMaskStrides(int total_sequence_length) { return {total_sequence_length, 0, 1}; } +template +std::tuple GetQkvBuffers( + const AttentionParameters* attn, + const T* query, + const T* key, + const T* value) { + switch (attn->qkv_format) { + case Q_K_V_BNSH: + case Q_K_V_BSNH: + return {reinterpret_cast(query), + reinterpret_cast(key), + reinterpret_cast(value)}; + case Q_KV_BSNH_BSN2H: { + auto packed_kv = reinterpret_cast(key); + return {reinterpret_cast(query), packed_kv, packed_kv + attn->head_size}; + } + case QKV_BSN3H: { + auto packed_qkv = reinterpret_cast(query); + return {packed_qkv, packed_qkv + 1 * attn->head_size, packed_qkv + 2 * attn->head_size}; + } + default: + return {nullptr, nullptr, nullptr}; + } +} + +inline std::tuple GetQkvStrides(const AttentionParameters* attn) { + // G0 not used, because it is the slowest dimension + const int& G1 = attn->num_heads; + const int& M = attn->sequence_length; + const int& N = attn->total_sequence_length; + const int& K = attn->head_size; + const int& O = attn->v_head_size; + + int4 q_strides, k_strides, v_strides; + switch (attn->qkv_format) { + case Q_K_V_BNSH: + q_strides = {G1 * M * K, M * K, K, 1}; + k_strides = {G1 * N * K, N * K, K, 1}; // matrices are transposed + v_strides = {G1 * N * O, N * O, 1, O}; + break; + case Q_KV_BSNH_BSN2H: + ORT_ENFORCE(K == O); + q_strides = {M * G1 * K, K, G1 * K, 1}; // [G0, M, G1, K] layout + k_strides = {N * G1 * 2 * K, 2 * K, G1 * 2 * K, 1}; // [G0, N, G1, K] layout + v_strides = {N * G1 * 2 * O, 2 * O, 1, G1 * 2 * O}; // [G0, N, G1, O] layout + break; + case QKV_BSN3H: + ORT_ENFORCE(K == O); + q_strides = {M * G1 * 3 * K, 3 * K, G1 * 3 * K, 1}; // [G0, M, G1, K] layout + k_strides = {N * G1 * 3 * K, 3 * K, G1 * 3 * K, 1}; // [G0, N, G1, K] layout + v_strides = {N * G1 * 3 * O, 3 * O, 1, G1 * 3 * O}; // [G0, N, G1, O] layout + break; + default: + break; + } + return {q_strides, k_strides, v_strides}; +} + inline std::tuple GetRawMaskBufferAddrSizesAndStrides( const int* buffer, const AttentionParameters* attn) { const int* offseted_buffer{buffer}; // how to view the mask buffer @@ -125,7 +185,8 @@ struct GemmSoftmaxGemmPermuteParams : onnxruntime::rocm::tunable::OpParams { "_N", attention->num_heads, "_H", attention->head_size, bias_buffer != nullptr ? "_B" : "_NB", - "_M", mask_index_dims.size()); + "_M", mask_index_dims.size(), + "_QKV", attention->qkv_format); } std::tuple GetGemmsMNKOBatch() const { @@ -270,6 +331,8 @@ struct GemmSoftmaxGemmPermuteGenericPipeline { } static Status Run(const GemmSoftmaxGemmPermuteParams* params, bool use_persistent_softmax) { + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( + params->attention->qkv_format != Q_K_V_BNSH, "GenericPipeline only supports qkv_format as Q_K_V_BNSH"); ORT_RETURN_IF_ERROR(Gemm1(params)); if (UseRawAttentionMask(params)) { @@ -291,6 +354,18 @@ class GemmSoftmaxGemmPermuteTunableOp : public tunable::TunableOpqkv_format) { + case Q_K_V_BNSH: + case Q_K_V_BSNH: + case Q_KV_BSNH_BSN2H: + case QKV_BSN3H: + return true; + default: + return false; + } + } + inline static bool IsSupportedMaskType(const AttentionParameters* attn) { switch (attn->mask_type) { case MASK_NONE: @@ -416,38 +491,47 @@ auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() { auto invoker = impl->MakeInvokerPointer(); auto op = [impl = std::move(impl), invoker = std::move(invoker)]( const GemmSoftmaxGemmPermuteParams* params) -> Status { + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( + !GemmSoftmaxGemmPermuteTunableOp::IsSupportedQkvFormat(params->attention), + "qkv format is not supported, got ", params->attention->qkv_format); if constexpr (USE_BIAS) { - TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->bias_buffer == nullptr); + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( + params->bias_buffer == nullptr, "biased version only support input with bias"); } else { - TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->bias_buffer != nullptr); + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( + params->bias_buffer != nullptr, "non-biased version only support input without bias"); } if constexpr (USE_MASK) { TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( - !GemmSoftmaxGemmPermuteTunableOp::IsSupportedMaskType(params->attention)); - TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->mask_index_buffer == nullptr); + !GemmSoftmaxGemmPermuteTunableOp::IsSupportedMaskType(params->attention), + "mask type is not supported, got ", params->attention->mask_type); + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( + params->mask_index_buffer == nullptr, "masked version only support input with mask"); } else { - TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->mask_index_buffer != nullptr); + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( + params->mask_index_buffer != nullptr, "non-masked version only support input without mask"); } auto attn = params->attention; - int G0 = attn->batch_size; - int G1 = attn->num_heads; - int M = attn->sequence_length; - int N = attn->total_sequence_length; - int K = attn->head_size; - int O = attn->v_head_size; + const int& G0 = attn->batch_size; + const int& G1 = attn->num_heads; + const int& M = attn->sequence_length; + const int& N = attn->total_sequence_length; + const int& K = attn->head_size; + const int& O = attn->v_head_size; { auto [m, n, k, o, batch] = params->GetGemmsMNKOBatch(); ORT_ENFORCE(M == m && N == n && K == k && O == o && G0 * G1 == batch, "semantic mismatch"); ORT_ENFORCE(K == O, "inner product dimension mismatch"); } + auto [qs, ks, vs] = GetQkvStrides(attn); std::vector q_buffer_lengths = {G0, G1, M, K}; - std::vector q_buffer_strides = {G1 * M * K, M * K, K, 1}; + std::vector q_buffer_strides = {qs.x, qs.y, qs.z, qs.w}; std::vector k_buffer_lengths = {G0, G1, N, K}; - std::vector k_buffer_strides = {G1 * N * K, N * K, K, 1}; // matrices are transposed + std::vector k_buffer_strides = {ks.x, ks.y, ks.z, ks.w}; std::vector v_buffer_lengths = {G0, G1, O, N}; - std::vector v_buffer_strides = {G1 * N * O, N * O, 1, O}; + std::vector v_buffer_strides = {vs.x, vs.y, vs.z, vs.w}; std::vector out_buffer_lengths = {G0, G1, M, O}; std::vector out_buffer_strides = {M * G1 * O, O, G1 * O, 1}; // permute 0213 diff --git a/onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu b/onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu new file mode 100644 index 0000000000..445894d07f --- /dev/null +++ b/onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu @@ -0,0 +1,127 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/rocm/bert/multihead_attention.h" + +#include "contrib_ops/cpu/bert/multihead_attention_helper.h" +#include "contrib_ops/rocm/bert/attention_impl.h" +#include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh" +#include "core/platform/env_var_utils.h" +#include "core/providers/rocm/rocm_common.h" + +using namespace onnxruntime::rocm; +using namespace ::onnxruntime::common; +using namespace ONNX_NAMESPACE; + +namespace onnxruntime { +namespace contrib { +namespace rocm { + +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + MultiHeadAttention, \ + kMSDomain, \ + 1, \ + T, \ + kRocmExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + MultiHeadAttention); + +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +MultiHeadAttention::MultiHeadAttention(const OpKernelInfo& info) + : RocmKernel(info) { + int64_t num_heads = 0; + ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); + num_heads_ = static_cast(num_heads); + + mask_filter_value_ = info.GetAttrOrDefault("mask_filter_value", -10000.0f); + + scale_ = info.GetAttrOrDefault("scale", 0.0f); +} + +template +Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) const { + ORT_ENFORCE( + GetTuningContext()->IsTunableOpEnabled(), + "MultiHeadAttention of ROCm EP is only supported if tunable op is used and tuning is enabled."); + + const Tensor* query = context->Input(0); + const Tensor* key = context->Input(1); + const Tensor* value = context->Input(2); + const Tensor* bias = context->Input(3); + const Tensor* key_padding_mask = context->Input(4); + const Tensor* relative_position_bias = context->Input(5); + const Tensor* past_key = context->Input(6); + const Tensor* past_value = context->Input(7); + + // TODO: Add support for bias, key_padding_mask and attention cache. + ORT_ENFORCE(bias == nullptr && key_padding_mask == nullptr && past_key == nullptr && past_value == nullptr, + "bias, key_padding_mask and attention cache is not supported"); + + auto& device_prop = GetDeviceProp(); + AttentionParameters attn; + ORT_RETURN_IF_ERROR( + multihead_attention_helper::CheckInputs( + query, key, value, bias, + key_padding_mask, relative_position_bias, + past_key, past_value, /*past_seq_len=*/nullptr, + &attn, + num_heads_, mask_filter_value_, scale_, + false, device_prop.maxThreadsPerBlock)); + // TODO: support more qkv formats + ORT_ENFORCE(attn.qkv_format == Q_KV_BSNH_BSN2H || attn.qkv_format == QKV_BSN3H, "Got ", attn.qkv_format); + + int sequence_length = attn.sequence_length; + + TensorShapeVector output_shape(3); + output_shape[0] = static_cast(attn.batch_size); + output_shape[1] = static_cast(sequence_length); + output_shape[2] = static_cast(attn.v_hidden_size); + Tensor* output = context->Output(0, output_shape); + + std::vector present_dims{ + attn.batch_size, + attn.num_heads, + attn.total_sequence_length, + attn.head_size, + }; + TensorShape present_shape(present_dims); + Tensor* present_key = context->Output(1, present_shape); + Tensor* present_value = context->Output(2, present_shape); + // TODO: Add support for attention cache + ORT_ENFORCE(present_key == nullptr && present_value == nullptr, "attention cache is not supported"); + + using HipT = typename ToHipType::MappedType; + using AttentionTunableOp = GemmSoftmaxGemmPermuteTunableOp; + auto workspace_bytes = AttentionTunableOp::GetWorkspaceNumBytes(&attn); + auto workspace = GetScratchBuffer(workspace_bytes, context->GetComputeStream()); + + GemmSoftmaxGemmPermuteParams params; + params.tuning_ctx = GetTuningContext(); + params.stream = Stream(context); + params.handle = GetRocblasHandle(context); + params.attention = &attn; + params.device_prop = &device_prop; + params.scale = scale_ == 0 ? 1.0f / sqrt(attn.head_size) : scale_; + std::tie(params.q_buffer, params.k_buffer, params.v_buffer) = GetQkvBuffers( + &attn, + query->DataRaw(), + key == nullptr ? nullptr : key->DataRaw(), + value == nullptr ? nullptr : value->DataRaw()); + params.out_buffer = reinterpret_cast(output->MutableDataRaw()); + + if (relative_position_bias != nullptr) { + params.bias_buffer = reinterpret_cast(relative_position_bias->DataRaw()); + } + + params.workspace_buffer = reinterpret_cast(workspace.get()); + return AttentionTunableOp{}(¶ms); +} + +} // namespace rocm +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/bert/multihead_attention.h b/onnxruntime/contrib_ops/rocm/bert/multihead_attention.h new file mode 100644 index 0000000000..25bc1cf811 --- /dev/null +++ b/onnxruntime/contrib_ops/rocm/bert/multihead_attention.h @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include "core/providers/rocm/rocm_kernel.h" +#include "contrib_ops/rocm/bert/attention_impl.h" + +namespace onnxruntime { +namespace contrib { +namespace rocm { + +using namespace onnxruntime::rocm; + +template +class MultiHeadAttention final : public RocmKernel { + public: + MultiHeadAttention(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + protected: + int num_heads_; // number of attention heads + float mask_filter_value_; + float scale_; +}; + +} // namespace rocm +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc index a3efd80668..9dd0ad8849 100644 --- a/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/rocm/rocm_contrib_kernels.cc @@ -59,6 +59,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float, Crop); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double, Crop); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16, Crop); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, MultiHeadAttention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, MultiHeadAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, DecoderAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, DecoderAttention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice); @@ -184,6 +186,9 @@ Status RegisterRocmContribKernels(KernelRegistry& kernel_registry) { // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo& K, + std::optional& V, std::optional& attn_bias, std::optional& attn_mask, DeviceArray& out) { @@ -64,6 +65,7 @@ class IGemmSoftmaxGemmPermuteKernelExplorer : public IKernelExplorer { } else { ORT_ENFORCE(false, "mask type not supported"); } + attn_.qkv_format = qkv_format; device_prop = GetEp()->GetDeviceProp(); @@ -74,9 +76,9 @@ class IGemmSoftmaxGemmPermuteKernelExplorer : public IKernelExplorer { params_.device_prop = &device_prop; params_.scale = scale; - params_.q_buffer = reinterpret_cast(Q.ptr()); - params_.k_buffer = reinterpret_cast(K.ptr()); - params_.v_buffer = reinterpret_cast(V.ptr()); + std::tie(params_.q_buffer, params_.k_buffer, params_.v_buffer) = GetQkvBuffers( + &attn_, Q.ptr(), K.has_value() ? K->ptr() : nullptr, V.has_value() ? V->ptr() : nullptr); + if (attn_bias.has_value()) { params_.bias_buffer = reinterpret_cast(attn_bias->ptr()); } @@ -130,14 +132,15 @@ class GemmSoftmaxGemmPermuteGeneric : public IGemmSoftmaxGemmPermuteKernelExplor int64_t head_size, int64_t mask_dim, double scale, + contrib::AttentionQkvFormat qkv_format, DeviceArray& Q, - DeviceArray& K, - DeviceArray& V, + std::optional& K, + std::optional& V, std::optional& attn_bias, std::optional& attn_mask, DeviceArray& out) : IGemmSoftmaxGemmPermuteKernelExplorer(batch, seqlen, total_seqlen, max_seqlen, - num_heads, head_size, mask_dim, scale, + num_heads, head_size, mask_dim, scale, qkv_format, Q, K, V, attn_bias, attn_mask, out) { this->SetWorkspace(GemmSoftmaxGemmPermuteGenericPipeline::GetWorkspaceNumBytes(&this->attn_)); } @@ -169,14 +172,15 @@ class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer int64_t head_size, int64_t mask_dim, double scale, + contrib::AttentionQkvFormat qkv_format, DeviceArray& Q, - DeviceArray& K, - DeviceArray& V, + std::optional& K, + std::optional& V, std::optional& attn_bias, std::optional& attn_mask, DeviceArray& out) : IGemmSoftmaxGemmPermuteKernelExplorer(batch, seqlen, total_seqlen, max_seqlen, - num_heads, head_size, mask_dim, scale, + num_heads, head_size, mask_dim, scale, qkv_format, Q, K, V, attn_bias, attn_mask, out) { this->SetWorkspace(GemmSoftmaxGemmPermuteTunableOp::GetWorkspaceNumBytes(&this->attn_)); @@ -229,14 +233,15 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor int64_t head_size, int64_t mask_dim, double scale, + contrib::AttentionQkvFormat qkv_format, DeviceArray& Q, - DeviceArray& K, - DeviceArray& V, + std::optional& K, + std::optional& V, std::optional& attn_bias, std::optional& attn_mask, DeviceArray& out) : IGemmSoftmaxGemmPermuteKernelExplorer(batch, seqlen, total_seqlen, max_seqlen, - num_heads, head_size, mask_dim, scale, + num_heads, head_size, mask_dim, scale, qkv_format, Q, K, V, attn_bias, attn_mask, out) { this->SetWorkspace(std::max( GemmSoftmaxGemmPermuteGenericPipeline::GetWorkspaceNumBytes(&this->attn_), @@ -261,12 +266,12 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor #define REGISTER_COMMON(name, type, ...) \ py::class_>(m, name) \ .def(py::init, int64_t, int64_t, int64_t, \ - float, \ - DeviceArray&, \ - DeviceArray&, \ + float, contrib::AttentionQkvFormat, \ DeviceArray&, \ std::optional&, \ std::optional&, \ + std::optional&, \ + std::optional&, \ DeviceArray&>()) \ .def("SetRepeats", &type<__VA_ARGS__>::SetRepeats) \ .def("Run", &type<__VA_ARGS__>::Run) \ @@ -285,6 +290,14 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor REGISTER_COMMON("GemmSoftmaxGemmPermuteTunable_" #dtype, GemmSoftmaxGemmPermuteTunable, dtype) KE_REGISTER(m) { + auto qkv_format = m.def_submodule("qkv_format"); + py::enum_(qkv_format, "qkv_format") + .value("Q_K_V_BNSH", contrib::AttentionQkvFormat::Q_K_V_BNSH, "") + .value("Q_K_V_BSNH", contrib::AttentionQkvFormat::Q_K_V_BSNH, "") + .value("QKV_BSN3H", contrib::AttentionQkvFormat::QKV_BSN3H, "") + .value("Q_KV_BSNH_BSN2H", contrib::AttentionQkvFormat::Q_KV_BSNH_BSN2H, "") + .export_values(); + REGISTER_GENERIC(half); #ifdef USE_COMPOSABLE_KERNEL diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index d01e07a46a..20d5218228 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -9,6 +9,18 @@ #include "test/util/include/scoped_env_vars.h" #include "test/contrib_ops/attention_op_test_helper.h" +#if defined(USE_ROCM) && defined(USE_COMPOSABLE_KERNEL) +#define DISABLE_ROCM false +#else +#define DISABLE_ROCM true +#endif + +#if defined(USE_ROCM) +#define ROCM_GTEST_SKIP(message) GTEST_SKIP_(message) +#else +#define ROCM_GTEST_SKIP(message) +#endif + namespace onnxruntime { namespace test { @@ -37,13 +49,14 @@ static void RunMultiHeadAttentionTest( bool use_float16 = false, bool disable_cpu = true, // not supported in cpu right now. bool disable_cuda = false, - bool disable_rocm = true) // not supported in rocm right now. + bool disable_rocm = DISABLE_ROCM) // not supported in rocm right now. { kv_sequence_length = (kv_sequence_length == 0 ? sequence_length : kv_sequence_length); int min_cuda_architecture = use_float16 ? 750 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture) && !disable_cuda; - bool enable_rocm = (nullptr != DefaultRocmExecutionProvider().get()) && !disable_rocm; + // rocm mha is required to work with TunableOp Enabled + bool enable_rocm = (nullptr != DefaultRocmExecutionProvider(/*test_tunable_op=*/true).get()) && !disable_rocm; bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16 && !disable_cpu; if (enable_cpu || enable_cuda || enable_rocm) { @@ -227,7 +240,7 @@ static void RunMultiHeadAttentionTest( if (enable_rocm) { std::vector> execution_providers; - execution_providers.push_back(DefaultRocmExecutionProvider()); + execution_providers.push_back(DefaultRocmExecutionProvider(/*test_tunable_op=*/true)); tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } @@ -265,7 +278,7 @@ static void RunMultiHeadAttentionKernel( bool is_static_kv = true, bool disable_cpu = true, // not supported in cpu right now. bool disable_cuda = false, - bool disable_rocm = true) { + bool disable_rocm = DISABLE_ROCM) { if (kernel_type == AttentionKernelType::AttentionKernel_Default) { ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ @@ -429,24 +442,28 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data) { // Test fused cross attention kernel // It requires head_size > 32 and head_size <= 64 for T4 GPU; hidden_size == v_hidden_size. TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize40) { + ROCM_GTEST_SKIP("ROCm MHA does not support bias"); AttentionTestData data; GetCrossAttentionData_HeadSize40(data); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize32_RightSidePadding_Mask1D) { + ROCM_GTEST_SKIP("ROCm MHA does not support mask"); AttentionTestData data; GetCrossAttentionData_Batch2_HeadSize32_RightSidePadding(data, true); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize32_RightSidePadding_Mask2D) { + ROCM_GTEST_SKIP("ROCm MHA does not support mask"); AttentionTestData data; GetCrossAttentionData_Batch2_HeadSize32_RightSidePadding(data, false); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, CrossAttention_Batch1_HeadSize32_LeftSidePadding_Mask2D) { + ROCM_GTEST_SKIP("ROCm MHA does not support mask"); AttentionTestData data; GetCrossAttentionData_Batch1_HeadSize32_LeftSidePadding(data); RunMultiHeadAttentionTests(data); @@ -467,30 +484,35 @@ TEST(MultiHeadAttentionTest, SelfAttention_Batch2_HeadSize32_NoBias_NoMask_Packe // This tests qk_head_size != k_head_size TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize16_8) { + ROCM_GTEST_SKIP("ROCm MHA does not support bias"); AttentionTestData data; GetCrossAttentionData_HeadSize16_8(data); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, CrossAttention_Batch1_HeadSize16) { + ROCM_GTEST_SKIP("ROCm MHA does not support bias"); AttentionTestData data; GetCrossAttentionData_HeadSize16(data); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, CrossAttentionWithPast) { + ROCM_GTEST_SKIP("ROCm MHA does not support attention cache"); AttentionTestData data; GetCrossAttentionDataWithPast(data); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, SelfAttentionWithPast) { + ROCM_GTEST_SKIP("ROCm MHA does not support attention cache"); AttentionTestData data; GetSelfAttentionDataWithPast(data); RunMultiHeadAttentionTests(data); } TEST(MultiHeadAttentionTest, AttentionCutlassRelPosBias) { + ROCM_GTEST_SKIP("ROCm does not support cutlass"); AttentionTestData data; GetAttentionDataCutlassRelPosBias(data); RunMultiHeadAttentionTests(data);