mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
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 ```
This commit is contained in:
cloudhan
GitHub
parent
feafbc4263
commit
9acbfc6a29
14 changed files with 467 additions and 76 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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<int>(key_dims[1]);
|
||||
} else if (key_dims.size() == 5) {
|
||||
if (static_cast<int>(key_dims[2]) != num_heads || static_cast<int>(key_dims[3]) != 2 || static_cast<int>(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<int>(key_dims[1]);
|
||||
} else { // key_dims.size() == 4 (cross-attention with past_key)
|
||||
if (static_cast<int>(key_dims[1]) != num_heads || static_cast<int>(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<int>(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<AttentionParameters*>(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) {
|
||||
|
|
|
|||
|
|
@ -63,6 +63,7 @@ Status Attention<T>::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<int64_t>(attn.batch_size);
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@
|
|||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace rocm {
|
||||
|
||||
size_t GetAttentionScratchSize(
|
||||
size_t element_size,
|
||||
int batch_size,
|
||||
|
|
|
|||
|
|
@ -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 <typename HipT, typename T>
|
||||
std::tuple<const HipT*, const HipT*, const HipT*> 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<const HipT*>(query),
|
||||
reinterpret_cast<const HipT*>(key),
|
||||
reinterpret_cast<const HipT*>(value)};
|
||||
case Q_KV_BSNH_BSN2H: {
|
||||
auto packed_kv = reinterpret_cast<const HipT*>(key);
|
||||
return {reinterpret_cast<const HipT*>(query), packed_kv, packed_kv + attn->head_size};
|
||||
}
|
||||
case QKV_BSN3H: {
|
||||
auto packed_qkv = reinterpret_cast<const HipT*>(query);
|
||||
return {packed_qkv, packed_qkv + 1 * attn->head_size, packed_qkv + 2 * attn->head_size};
|
||||
}
|
||||
default:
|
||||
return {nullptr, nullptr, nullptr};
|
||||
}
|
||||
}
|
||||
|
||||
inline std::tuple<int4, int4, int4> 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<const int*, int3, int3> 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<int, int, int, int, int> GetGemmsMNKOBatch() const {
|
||||
|
|
@ -270,6 +331,8 @@ struct GemmSoftmaxGemmPermuteGenericPipeline {
|
|||
}
|
||||
|
||||
static Status Run(const GemmSoftmaxGemmPermuteParams<T>* 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::TunableOp<GemmSoftmaxGem
|
|||
public:
|
||||
GemmSoftmaxGemmPermuteTunableOp();
|
||||
|
||||
inline static bool IsSupportedQkvFormat(const AttentionParameters* attn) {
|
||||
switch (attn->qkv_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<T>* params) -> Status {
|
||||
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(
|
||||
!GemmSoftmaxGemmPermuteTunableOp<T>::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<T>::IsSupportedMaskType(params->attention));
|
||||
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->mask_index_buffer == nullptr);
|
||||
!GemmSoftmaxGemmPermuteTunableOp<T>::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<ck::index_t> q_buffer_lengths = {G0, G1, M, K};
|
||||
std::vector<ck::index_t> q_buffer_strides = {G1 * M * K, M * K, K, 1};
|
||||
std::vector<ck::index_t> q_buffer_strides = {qs.x, qs.y, qs.z, qs.w};
|
||||
std::vector<ck::index_t> k_buffer_lengths = {G0, G1, N, K};
|
||||
std::vector<ck::index_t> k_buffer_strides = {G1 * N * K, N * K, K, 1}; // matrices are transposed
|
||||
std::vector<ck::index_t> k_buffer_strides = {ks.x, ks.y, ks.z, ks.w};
|
||||
std::vector<ck::index_t> v_buffer_lengths = {G0, G1, O, N};
|
||||
std::vector<ck::index_t> v_buffer_strides = {G1 * N * O, N * O, 1, O};
|
||||
std::vector<ck::index_t> v_buffer_strides = {vs.x, vs.y, vs.z, vs.w};
|
||||
std::vector<ck::index_t> out_buffer_lengths = {G0, G1, M, O};
|
||||
std::vector<ck::index_t> out_buffer_strides = {M * G1 * O, O, G1 * O, 1}; // permute 0213
|
||||
|
||||
|
|
|
|||
127
onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu
Normal file
127
onnxruntime/contrib_ops/rocm/bert/multihead_attention.cu
Normal file
|
|
@ -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<T>()), \
|
||||
MultiHeadAttention<T>);
|
||||
|
||||
REGISTER_KERNEL_TYPED(float)
|
||||
REGISTER_KERNEL_TYPED(MLFloat16)
|
||||
|
||||
template <typename T>
|
||||
MultiHeadAttention<T>::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<int>(num_heads);
|
||||
|
||||
mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
|
||||
|
||||
scale_ = info.GetAttrOrDefault<float>("scale", 0.0f);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status MultiHeadAttention<T>::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<Tensor>(0);
|
||||
const Tensor* key = context->Input<Tensor>(1);
|
||||
const Tensor* value = context->Input<Tensor>(2);
|
||||
const Tensor* bias = context->Input<Tensor>(3);
|
||||
const Tensor* key_padding_mask = context->Input<Tensor>(4);
|
||||
const Tensor* relative_position_bias = context->Input<Tensor>(5);
|
||||
const Tensor* past_key = context->Input<Tensor>(6);
|
||||
const Tensor* past_value = context->Input<Tensor>(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<Tensor>(
|
||||
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<int64_t>(attn.batch_size);
|
||||
output_shape[1] = static_cast<int64_t>(sequence_length);
|
||||
output_shape[2] = static_cast<int64_t>(attn.v_hidden_size);
|
||||
Tensor* output = context->Output(0, output_shape);
|
||||
|
||||
std::vector<int64_t> 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<T>::MappedType;
|
||||
using AttentionTunableOp = GemmSoftmaxGemmPermuteTunableOp<HipT>;
|
||||
auto workspace_bytes = AttentionTunableOp::GetWorkspaceNumBytes(&attn);
|
||||
auto workspace = GetScratchBuffer<void>(workspace_bytes, context->GetComputeStream());
|
||||
|
||||
GemmSoftmaxGemmPermuteParams<HipT> 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<HipT>(
|
||||
&attn,
|
||||
query->DataRaw(),
|
||||
key == nullptr ? nullptr : key->DataRaw(),
|
||||
value == nullptr ? nullptr : value->DataRaw());
|
||||
params.out_buffer = reinterpret_cast<HipT*>(output->MutableDataRaw());
|
||||
|
||||
if (relative_position_bias != nullptr) {
|
||||
params.bias_buffer = reinterpret_cast<const HipT*>(relative_position_bias->DataRaw());
|
||||
}
|
||||
|
||||
params.workspace_buffer = reinterpret_cast<HipT*>(workspace.get());
|
||||
return AttentionTunableOp{}(¶ms);
|
||||
}
|
||||
|
||||
} // namespace rocm
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
30
onnxruntime/contrib_ops/rocm/bert/multihead_attention.h
Normal file
30
onnxruntime/contrib_ops/rocm/bert/multihead_attention.h
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#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 <typename T>
|
||||
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
|
||||
|
|
@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float, Crop)>,
|
||||
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double, Crop)>,
|
||||
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16, Crop)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, float, MultiHeadAttention)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain, 1, MLFloat16, MultiHeadAttention)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
|
||||
1, float, DecoderAttention)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kMSDomain,
|
||||
|
|
|
|||
|
|
@ -6,4 +6,10 @@ class blas_op: # noqa: N801
|
|||
T: int
|
||||
N: int
|
||||
|
||||
class qkv_format: # noqa: N801
|
||||
Q_K_V_BNSH: int
|
||||
Q_K_V_BSNH: int
|
||||
QKV_BSN3H: int
|
||||
Q_KV_BSNH_BSN2H: int
|
||||
|
||||
def is_composable_kernel_available(*args, **kwargs): ...
|
||||
|
|
|
|||
|
|
@ -55,8 +55,32 @@ def get_mask_dim_id(dim):
|
|||
return f"mask_{dim}d"
|
||||
|
||||
|
||||
def maybe_pack_q_k_v_bnsh_for_device_on_host(q, k, v, dtype, qkv_format):
|
||||
q = q.astype(dtype)
|
||||
k = k.astype(dtype)
|
||||
v = v.astype(dtype)
|
||||
if qkv_format == ke.qkv_format.Q_K_V_BNSH:
|
||||
return q, k, v
|
||||
|
||||
# BNSH to BSNH
|
||||
q = np.swapaxes(q, 2, 1)
|
||||
k = np.swapaxes(k, 2, 1)
|
||||
v = np.swapaxes(v, 2, 1)
|
||||
|
||||
if qkv_format == ke.qkv_format.Q_K_V_BSNH:
|
||||
return np.ascontiguousarray(q), np.ascontiguousarray(k), np.ascontiguousarray(v)
|
||||
|
||||
if qkv_format == ke.qkv_format.QKV_BSN3H:
|
||||
return np.ascontiguousarray(np.stack([q, k, v], axis=-2)), None, None
|
||||
|
||||
if qkv_format == ke.qkv_format.Q_KV_BSNH_BSN2H:
|
||||
return np.ascontiguousarray(q), np.ascontiguousarray(np.stack([k, v], axis=-2)), None
|
||||
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _test_gemm_softmax_gemm_permute(
|
||||
f, dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale
|
||||
f, dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale, qkv_format
|
||||
):
|
||||
v_head_size = head_size
|
||||
q_shape = [batch, num_heads, seqlen, head_size]
|
||||
|
|
@ -85,15 +109,15 @@ def _test_gemm_softmax_gemm_permute(
|
|||
raise ValueError
|
||||
|
||||
np.random.seed(42)
|
||||
Q = multinormal_distribution(np.prod(q_shape[:-1]), q_shape[-1]).reshape(q_shape).astype(np.float64) # noqa: N806
|
||||
K = multinormal_distribution(np.prod(k_shape[:-1]), k_shape[-1]).reshape(k_shape).astype(np.float64) # noqa: N806
|
||||
V = multinormal_distribution(np.prod(v_shape[:-1]), v_shape[-1]).reshape(v_shape).astype(np.float64) # noqa: N806
|
||||
q = multinormal_distribution(np.prod(q_shape[:-1]), q_shape[-1]).reshape(q_shape).astype(np.float64)
|
||||
k = multinormal_distribution(np.prod(k_shape[:-1]), k_shape[-1]).reshape(k_shape).astype(np.float64)
|
||||
v = multinormal_distribution(np.prod(v_shape[:-1]), v_shape[-1]).reshape(v_shape).astype(np.float64)
|
||||
if bias_shape is not None:
|
||||
attn_bias = np.random.uniform(-0.5, 0.5, size=bias_shape)
|
||||
if mask_shape is not None:
|
||||
attn_mask = (np.random.randint(0, 100, size=mask_shape) < 95).astype(np.int32)
|
||||
|
||||
pre_softmax_attn_scores = Q @ np.swapaxes(K, 2, 3)
|
||||
pre_softmax_attn_scores = q @ np.swapaxes(k, 2, 3)
|
||||
pre_softmax_attn_scores = pre_softmax_attn_scores * scale
|
||||
if attn_bias is not None:
|
||||
pre_softmax_attn_scores = pre_softmax_attn_scores + attn_bias
|
||||
|
|
@ -106,17 +130,15 @@ def _test_gemm_softmax_gemm_permute(
|
|||
converted_mask = (1 - attn_mask.reshape(mask_shape_broadcasted)) * filter_value
|
||||
pre_softmax_attn_scores = pre_softmax_attn_scores + converted_mask
|
||||
attn_scores = softmax(pre_softmax_attn_scores, axis=-1)
|
||||
attn = attn_scores @ V
|
||||
attn = attn_scores @ v
|
||||
ref = np.swapaxes(attn, 2, 1) # permute 0213
|
||||
|
||||
out = np.empty(out_shape, dtype=dtype)
|
||||
host_Q = Q.astype(dtype) # noqa: N806
|
||||
host_K = K.astype(dtype) # noqa: N806
|
||||
host_V = V.astype(dtype) # noqa: N806
|
||||
host_q, host_k, host_v = maybe_pack_q_k_v_bnsh_for_device_on_host(q, k, v, dtype, qkv_format)
|
||||
host_attn_bias = attn_bias.astype(dtype) if attn_bias is not None else None
|
||||
dev_Q = ke.DeviceArray(host_Q) # noqa: N806
|
||||
dev_K = ke.DeviceArray(host_K) # noqa: N806
|
||||
dev_V = ke.DeviceArray(host_V) # noqa: N806
|
||||
dev_q = ke.DeviceArray(host_q)
|
||||
dev_k = ke.DeviceArray(host_k) if host_k is not None else None
|
||||
dev_v = ke.DeviceArray(host_v) if host_v is not None else None
|
||||
dev_out = ke.DeviceArray(out)
|
||||
dev_attn_bias = ke.DeviceArray(host_attn_bias) if host_attn_bias is not None else None
|
||||
dev_attn_mask = ke.DeviceArray(attn_mask) if attn_mask is not None else None
|
||||
|
|
@ -130,9 +152,10 @@ def _test_gemm_softmax_gemm_permute(
|
|||
head_size,
|
||||
mask_dim,
|
||||
scale,
|
||||
dev_Q,
|
||||
dev_K,
|
||||
dev_V,
|
||||
qkv_format,
|
||||
dev_q,
|
||||
dev_k,
|
||||
dev_v,
|
||||
dev_attn_bias,
|
||||
dev_attn_mask,
|
||||
dev_out,
|
||||
|
|
@ -181,7 +204,9 @@ def _test_gemm_softmax_gemm_permute(
|
|||
def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim):
|
||||
f = getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype))
|
||||
scale = 1.0 / np.sqrt(head_size)
|
||||
_test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
|
||||
_test_gemm_softmax_gemm_permute(
|
||||
f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale, ke.qkv_format.Q_K_V_BNSH
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled")
|
||||
|
|
@ -196,7 +221,9 @@ def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, n
|
|||
def test_gemm_softmax_gemm_permute_ck(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim):
|
||||
f = getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0))
|
||||
scale = 1.0 / np.sqrt(head_size)
|
||||
_test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
|
||||
_test_gemm_softmax_gemm_permute(
|
||||
f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale, ke.qkv_format.Q_K_V_BNSH
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mask_dim", [2], ids=get_mask_dim_id)
|
||||
|
|
@ -210,7 +237,38 @@ def test_gemm_softmax_gemm_permute_ck(dtype, batch, seqlen, total_seqlen, nhead,
|
|||
def test_gemm_softmax_gemm_permute_tunable(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim):
|
||||
f = getattr(ke, "GemmSoftmaxGemmPermuteTunable_" + dtype_to_suffix(dtype))
|
||||
scale = 1.0 / np.sqrt(head_size)
|
||||
_test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
|
||||
_test_gemm_softmax_gemm_permute(
|
||||
f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale, ke.qkv_format.Q_K_V_BNSH
|
||||
)
|
||||
|
||||
|
||||
stabel_diffusion_configs = [
|
||||
[2, 64, 64, 8, 160, "QKV_BSN3H"],
|
||||
[2, 256, 256, 8, 160, "QKV_BSN3H"],
|
||||
[2, 1024, 1024, 8, 80, "QKV_BSN3H"],
|
||||
[2, 4096, 4096, 8, 40, "QKV_BSN3H"],
|
||||
[2, 64, 77, 8, 160, "Q_KV_BSNH_BSN2H"],
|
||||
[2, 256, 77, 8, 160, "Q_KV_BSNH_BSN2H"],
|
||||
[2, 1024, 77, 8, 80, "Q_KV_BSNH_BSN2H"],
|
||||
[2, 4096, 77, 8, 40, "Q_KV_BSNH_BSN2H"],
|
||||
[1, 4096, 4096, 1, 512, "Q_K_V_BNSH"],
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled")
|
||||
@pytest.mark.parametrize("mask_dim", [0], ids=get_mask_dim_id)
|
||||
@pytest.mark.parametrize("biased", [False], ids=get_biased_id)
|
||||
@pytest.mark.parametrize("batch, seqlen, total_seqlen, nhead, head_size, qkv_format_name", stabel_diffusion_configs)
|
||||
@pytest.mark.parametrize("dtype", dtypes)
|
||||
def test_gemm_softmax_gemm_permute_ck_sd(
|
||||
dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, qkv_format_name
|
||||
):
|
||||
qkv_format = getattr(ke.qkv_format, qkv_format_name)
|
||||
f = getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0))
|
||||
scale = 1.0 / np.sqrt(head_size)
|
||||
_test_gemm_softmax_gemm_permute(
|
||||
f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale, qkv_format
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -238,7 +296,7 @@ class GemmSoftmaxGemmPermuteMetric(ke.ComputeMetric):
|
|||
|
||||
|
||||
def profile_gemm_softmax_gemm_permute_func(
|
||||
f, dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale
|
||||
f, dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale, qkv_format
|
||||
):
|
||||
v_head_size = head_size
|
||||
q_shape = [batch, num_heads, seqlen, head_size]
|
||||
|
|
@ -264,22 +322,20 @@ def profile_gemm_softmax_gemm_permute_func(
|
|||
raise ValueError
|
||||
|
||||
np.random.seed(42)
|
||||
Q = multinormal_distribution(np.prod(q_shape[:-1]), q_shape[-1]).reshape(q_shape).astype(np.float64) # noqa: N806
|
||||
K = multinormal_distribution(np.prod(k_shape[:-1]), k_shape[-1]).reshape(k_shape).astype(np.float64) # noqa: N806
|
||||
V = multinormal_distribution(np.prod(v_shape[:-1]), v_shape[-1]).reshape(v_shape).astype(np.float64) # noqa: N806
|
||||
q = multinormal_distribution(np.prod(q_shape[:-1]), q_shape[-1]).reshape(q_shape).astype(np.float64)
|
||||
k = multinormal_distribution(np.prod(k_shape[:-1]), k_shape[-1]).reshape(k_shape).astype(np.float64)
|
||||
v = multinormal_distribution(np.prod(v_shape[:-1]), v_shape[-1]).reshape(v_shape).astype(np.float64)
|
||||
if bias_shape is not None:
|
||||
attn_bias = np.random.uniform(-2, 2, size=bias_shape)
|
||||
if mask_shape is not None:
|
||||
attn_mask = (np.random.randint(0, 100, size=mask_shape) < 95).astype(np.int32)
|
||||
|
||||
out = np.empty(out_shape, dtype=dtype)
|
||||
host_Q = Q.astype(dtype) # noqa: N806
|
||||
host_K = K.astype(dtype) # noqa: N806
|
||||
host_V = V.astype(dtype) # noqa: N806
|
||||
host_q, host_k, host_v = maybe_pack_q_k_v_bnsh_for_device_on_host(q, k, v, dtype, qkv_format)
|
||||
host_attn_bias = attn_bias.astype(dtype) if attn_bias is not None else None
|
||||
dev_Q = ke.DeviceArray(host_Q) # noqa: N806
|
||||
dev_K = ke.DeviceArray(host_K) # noqa: N806
|
||||
dev_V = ke.DeviceArray(host_V) # noqa: N806
|
||||
dev_q = ke.DeviceArray(host_q)
|
||||
dev_k = ke.DeviceArray(host_k) if host_k is not None else None
|
||||
dev_v = ke.DeviceArray(host_v) if host_v is not None else None
|
||||
dev_out = ke.DeviceArray(out)
|
||||
dev_attn_bias = ke.DeviceArray(host_attn_bias) if host_attn_bias is not None else None
|
||||
dev_attn_mask = ke.DeviceArray(attn_mask) if attn_mask is not None else None
|
||||
|
|
@ -293,9 +349,10 @@ def profile_gemm_softmax_gemm_permute_func(
|
|||
head_size,
|
||||
mask_dim,
|
||||
scale,
|
||||
dev_Q,
|
||||
dev_K,
|
||||
dev_V,
|
||||
qkv_format,
|
||||
dev_q,
|
||||
dev_k,
|
||||
dev_v,
|
||||
dev_attn_bias,
|
||||
dev_attn_mask,
|
||||
dev_out,
|
||||
|
|
@ -324,12 +381,15 @@ def profile_gemm_softmax_gemm_permute_func(
|
|||
)
|
||||
|
||||
|
||||
def profile_with_args(dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale, *, sort=False):
|
||||
def profile_with_args(
|
||||
dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale, qkv_format, *, sort=False
|
||||
):
|
||||
with ke.benchmark(sort):
|
||||
args = (dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale)
|
||||
profile_gemm_softmax_gemm_permute_func(
|
||||
getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype)), *args
|
||||
)
|
||||
args = (dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale, qkv_format)
|
||||
if qkv_format == ke.qkv_format.Q_K_V_BNSH:
|
||||
profile_gemm_softmax_gemm_permute_func(
|
||||
getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype)), *args
|
||||
)
|
||||
if ke.is_composable_kernel_available():
|
||||
profile_gemm_softmax_gemm_permute_func(
|
||||
getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0)), *args
|
||||
|
|
@ -340,8 +400,24 @@ def profile_with_args(dtype, batch, seqlen, total_seqlen, num_heads, head_size,
|
|||
|
||||
|
||||
def profile():
|
||||
for batch, seqlen, total_seqlen, nhead, head_size, qkv_format_name in stabel_diffusion_configs:
|
||||
profile_with_args(
|
||||
"float16",
|
||||
batch,
|
||||
seqlen,
|
||||
total_seqlen,
|
||||
nhead,
|
||||
head_size,
|
||||
biased=False,
|
||||
mask_dim=0,
|
||||
qkv_format=getattr(ke.qkv_format, qkv_format_name),
|
||||
scale=0.125,
|
||||
sort=True,
|
||||
)
|
||||
print()
|
||||
|
||||
for args in product(dtypes, batches, seqlens, total_seqlens, num_heads, head_sizes, biaseds, mask_dims):
|
||||
profile_with_args(*args, scale=0.125, sort=True)
|
||||
profile_with_args(*args, qkv_format=ke.qkv_format.Q_K_V_BNSH, scale=0.125, sort=True)
|
||||
print()
|
||||
|
||||
|
||||
|
|
@ -359,7 +435,18 @@ if __name__ == "__main__":
|
|||
group.add_argument("head_size", type=int)
|
||||
group.add_argument("biased", type=int, choices=[0, 1], default=0)
|
||||
group.add_argument("mask_dim", type=int, choices=[0, 2, 3, 4], default=2, help="0 for mask disabled")
|
||||
group.add_argument("scale", type=float, default=None, help="default to 1.0/sqrt(head_size)")
|
||||
group.add_argument("--scale", type=float, default=None, help="default to 1.0/sqrt(head_size)")
|
||||
group.add_argument(
|
||||
"--qkv_format",
|
||||
type=lambda name: getattr(ke.qkv_format, name),
|
||||
default="Q_K_V_BNSH",
|
||||
choices=[
|
||||
"Q_K_V_BNSH", # non-packed, permuted
|
||||
"Q_K_V_BSNH", # non-packed, non-permuted
|
||||
"Q_KV_BSNH_BSN2H", # kv packed, non-permuted
|
||||
"QKV_BSN3H", # qkv packed, non-permuted
|
||||
],
|
||||
)
|
||||
|
||||
if len(sys.argv) == 1:
|
||||
profile()
|
||||
|
|
@ -375,5 +462,6 @@ if __name__ == "__main__":
|
|||
args.biased,
|
||||
args.mask_dim,
|
||||
args.scale,
|
||||
args.qkv_format,
|
||||
sort=args.sort,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -28,9 +28,10 @@ class IGemmSoftmaxGemmPermuteKernelExplorer : public IKernelExplorer {
|
|||
int64_t head_size,
|
||||
int64_t mask_dim,
|
||||
double scale,
|
||||
contrib::AttentionQkvFormat qkv_format,
|
||||
DeviceArray& Q,
|
||||
DeviceArray& K,
|
||||
DeviceArray& V,
|
||||
std::optional<DeviceArray>& K,
|
||||
std::optional<DeviceArray>& V,
|
||||
std::optional<DeviceArray>& attn_bias,
|
||||
std::optional<DeviceArray>& 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<T*>(Q.ptr());
|
||||
params_.k_buffer = reinterpret_cast<T*>(K.ptr());
|
||||
params_.v_buffer = reinterpret_cast<T*>(V.ptr());
|
||||
std::tie(params_.q_buffer, params_.k_buffer, params_.v_buffer) = GetQkvBuffers<T>(
|
||||
&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<T*>(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<DeviceArray>& K,
|
||||
std::optional<DeviceArray>& V,
|
||||
std::optional<DeviceArray>& attn_bias,
|
||||
std::optional<DeviceArray>& attn_mask,
|
||||
DeviceArray& out)
|
||||
: IGemmSoftmaxGemmPermuteKernelExplorer<T>(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<T>::GetWorkspaceNumBytes(&this->attn_));
|
||||
}
|
||||
|
|
@ -169,14 +172,15 @@ class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer<T>
|
|||
int64_t head_size,
|
||||
int64_t mask_dim,
|
||||
double scale,
|
||||
contrib::AttentionQkvFormat qkv_format,
|
||||
DeviceArray& Q,
|
||||
DeviceArray& K,
|
||||
DeviceArray& V,
|
||||
std::optional<DeviceArray>& K,
|
||||
std::optional<DeviceArray>& V,
|
||||
std::optional<DeviceArray>& attn_bias,
|
||||
std::optional<DeviceArray>& attn_mask,
|
||||
DeviceArray& out)
|
||||
: IGemmSoftmaxGemmPermuteKernelExplorer<T>(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<T>::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<DeviceArray>& K,
|
||||
std::optional<DeviceArray>& V,
|
||||
std::optional<DeviceArray>& attn_bias,
|
||||
std::optional<DeviceArray>& attn_mask,
|
||||
DeviceArray& out)
|
||||
: IGemmSoftmaxGemmPermuteKernelExplorer<T>(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<T>::GetWorkspaceNumBytes(&this->attn_),
|
||||
|
|
@ -261,12 +266,12 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor
|
|||
#define REGISTER_COMMON(name, type, ...) \
|
||||
py::class_<type<__VA_ARGS__>>(m, name) \
|
||||
.def(py::init<int64_t, int64_t, int64_t, std::optional<int64_t>, int64_t, int64_t, int64_t, \
|
||||
float, \
|
||||
DeviceArray&, \
|
||||
DeviceArray&, \
|
||||
float, contrib::AttentionQkvFormat, \
|
||||
DeviceArray&, \
|
||||
std::optional<DeviceArray>&, \
|
||||
std::optional<DeviceArray>&, \
|
||||
std::optional<DeviceArray>&, \
|
||||
std::optional<DeviceArray>&, \
|
||||
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_<contrib::AttentionQkvFormat>(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
|
||||
|
|
|
|||
|
|
@ -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<std::unique_ptr<IExecutionProvider>> 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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue