saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
onnxruntime/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc
Tianlei Wu 6ffaaebb60
[CUDA] Attention kernel provider option (#21344) 
### Description
* Add a cuda provider option `sdpa_kernel` to choose which attention kernel to run for testing purpose. 
* Allow dump which attention kernel is used per node.
* Reserve  a flag for cudnn flash attention which will be added soon.

#### CUDA provider option sdpa_kernel
Instead of setting environment variable, we also support setting it in
provider option. Note that the setting is global per session. That could
help performance testing of each kernel.

#### Attention Kernel Debug Info
Set an environment variable `ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1`,
and ORT will print sdpa kernel used in each node:

For example 
```
ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1 ./onnxruntime_test_all --gtest_filter=MultiHeadAttentionTest*
```
It will show debug information of kernel used in testing:
```
[ RUN      ] MultiHeadAttentionTest.SelfAttention_Batch2_HeadSize32_NoBias_NoMask_PackedQKV
AttentionKernelOptions: FLASH_ATTENTION=0 EFFICIENT_ATTENTION=0 TRT_FUSED_ATTENTION=1 CUDNN_FLASH_ATTENTION=0 TRT_FLASH_ATTENTION=1 TRT_CROSS_ATTENTION=0 TRT_CAUSAL_ATTENTION=0 MATH=1
Operator=MultiHeadAttention Node=node1 DataType=fp16 TRT_FUSED_ATTENTION=1
AttentionKernelOptions: FLASH_ATTENTION=0 EFFICIENT_ATTENTION=1 TRT_FUSED_ATTENTION=0 CUDNN_FLASH_ATTENTION=0 TRT_FLASH_ATTENTION=0 TRT_CROSS_ATTENTION=0 TRT_CAUSAL_ATTENTION=0 MATH=1
Operator=MultiHeadAttention Node=node1 DataType=fp16 EFFICIENT_ATTENTION=1
```
In this test case, the debug info shows that one session uses trt fused
attention and another session use efficient attention.
2024-07-19 13:58:54 -07:00

329 lines
18 KiB
C++
Raw Blame History

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cuda/cuda_common.h"
#include "contrib_ops/cuda/bert/attention_impl.h"
#include "contrib_ops/cuda/bert/multihead_attention.h"
#include "contrib_ops/cpu/bert/multihead_attention_helper.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
#include "contrib_ops/cuda/bert/flash_attention/flash_api.h"
using namespace onnxruntime::cuda;
using namespace ::onnxruntime::common;
using namespace ONNX_NAMESPACE;
namespace onnxruntime {
namespace contrib {
namespace cuda {
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
MultiHeadAttention, \
kMSDomain, \
1, \
T, \
kCudaExecutionProvider, \
(*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)
: CudaKernel(info),
fused_fp16_cross_attention_kernel_(nullptr),
cumulated_sequence_length_q_cache_(),
cumulated_sequence_length_kv_cache_() {
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);
is_unidirectional_ = info.GetAttrOrDefault<int64_t>("unidirectional", 0) == 1;
ORT_ENFORCE(!is_unidirectional_, "Unidirectional MHA does not support CUDA kernel. Consider using Attention or GQA instead.");
kernel_options_ = this->GetAttentionKernelOptions();
disable_fused_self_attention_ = sizeof(T) != 2 || !kernel_options_->UseTrtFusedAttention();
enable_trt_flash_attention_ = sizeof(T) == 2 && kernel_options_->UseTrtFlashAttention();
disable_flash_attention_ = sizeof(T) != 2 || !kernel_options_->UseFlashAttention();
disable_memory_efficient_attention_ = !kernel_options_->UseEfficientAttention();
disable_fused_cross_attention_ = sizeof(T) != 2 || !kernel_options_->UseTrtCrossAttention();
// Allocate cache buffers
constexpr size_t cache_bytes = sizeof(int32_t) * (static_cast<size_t>(kCumulatedSequenceLengthCacheMaxBatchSize) + 1);
cumulated_sequence_length_q_cache_.buffer = GetTransientScratchBuffer<void>(cache_bytes);
cumulated_sequence_length_q_cache_.max_batch_size = kCumulatedSequenceLengthCacheMaxBatchSize;
cumulated_sequence_length_kv_cache_.buffer = GetTransientScratchBuffer<void>(cache_bytes);
cumulated_sequence_length_kv_cache_.max_batch_size = kCumulatedSequenceLengthCacheMaxBatchSize;
}
template <typename T>
Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
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);
auto& device_prop = GetDeviceProp();
AttentionParameters parameters;
parameters.use_tf32 = UseTF32();
ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckInputs<Tensor>(query,
key,
value,
bias,
key_padding_mask,
relative_position_bias,
past_key,
past_value,
nullptr, // past_seq_len
&parameters,
num_heads_,
mask_filter_value_,
scale_,
is_unidirectional_,
false, // past_present_share_buffer
false, // dmmha_packing
device_prop.maxThreadsPerBlock));
int sequence_length = parameters.sequence_length;
TensorShapeVector output_shape(3);
output_shape[0] = static_cast<int64_t>(parameters.batch_size);
output_shape[1] = static_cast<int64_t>(sequence_length);
output_shape[2] = static_cast<int64_t>(parameters.v_hidden_size);
Tensor* output = context->Output(0, output_shape);
std::vector<int64_t> present_dims{
parameters.batch_size, parameters.num_heads, parameters.total_sequence_length, parameters.head_size};
TensorShape present_shape(present_dims);
Tensor* present_key = context->Output(1, present_shape);
Tensor* present_value = context->Output(2, present_shape);
MHARunner* fused_runner = nullptr;
const FusedMultiHeadCrossAttentionKernel* fused_cross_attention_kernel = nullptr;
// Check whether we can use fused kernel
int sm = device_prop.major * 10 + device_prop.minor;
bool is_mask_1d_seq_len = parameters.mask_type == AttentionMaskType::MASK_1D_KEY_SEQ_LEN;
const bool pass_key_value_as_past = (parameters.pass_past_in_kv && nullptr != key && nullptr != value);
#if USE_FLASH_ATTENTION || USE_MEMORY_EFFICIENT_ATTENTION
// Exclude this case since PrepareQkv will convert the format to BNSH.
bool past_no_bias = (pass_key_value_as_past || past_key != nullptr || present_key != nullptr) && bias == nullptr;
#endif
#if USE_FLASH_ATTENTION
bool use_flash_attention = !disable_flash_attention_ &&
!past_no_bias &&
nullptr == relative_position_bias &&
nullptr == key_padding_mask &&
parameters.head_size == parameters.v_head_size &&
onnxruntime::flash::is_supported(device_prop,
parameters.head_size,
parameters.num_heads,
parameters.num_heads);
// When input is packed QKV format, TensorRT kernel might be faster than flash attention when sequence length <= 512.
if (use_flash_attention && key == nullptr && value == nullptr &&
parameters.sequence_length < kernel_options_->MinSeqLenForFlashAttentionPackedQkv()) {
use_flash_attention = false;
}
// Allocate buffers
size_t softmax_lse_accum_bytes = 0;
size_t out_accum_bytes = 0;
if (use_flash_attention) {
using namespace std;
auto [num_splits, slse_accum_bytes, o_accum_bytes] = onnxruntime::flash::get_num_splits_and_buffer_sizes(
parameters.batch_size, parameters.sequence_length, parameters.kv_sequence_length, parameters.num_heads,
parameters.head_size, device_prop.multiProcessorCount);
parameters.num_splits = static_cast<int>(num_splits);
softmax_lse_accum_bytes = slse_accum_bytes;
out_accum_bytes = o_accum_bytes;
}
auto softmax_lse_accum_buffer = GetScratchBuffer<void>(softmax_lse_accum_bytes, context->GetComputeStream());
auto out_accum_buffer = GetScratchBuffer<void>(out_accum_bytes, context->GetComputeStream());
#else
constexpr bool use_flash_attention = false;
auto softmax_lse_accum_buffer = GetScratchBuffer<void>(0, context->GetComputeStream()); // nullptr
auto out_accum_buffer = GetScratchBuffer<void>(0, context->GetComputeStream()); // nullptr
#endif
bool use_fused_cross_attention = !use_flash_attention &&
!disable_fused_cross_attention_ &&
nullptr == key_padding_mask &&
nullptr == relative_position_bias &&
(nullptr == past_key && nullptr == past_value && !parameters.pass_past_in_kv) &&
key != nullptr &&
(value != nullptr || bias == nullptr) && // TODO: new kernel for adding bias to packed KV
parameters.hidden_size == parameters.v_hidden_size &&
has_fused_cross_attention_kernel(sm, parameters.head_size,
parameters.kv_sequence_length);
if (use_fused_cross_attention) {
if (fused_fp16_cross_attention_kernel_ == nullptr) {
fused_fp16_cross_attention_kernel_ = get_fused_cross_attention_kernels(sm);
}
// In case some kernel not loaded due to shared memory limit, we need to double check here.
// The kernel has no limit on sequence length, and this checks whether the kernel has been loaded.
if (fused_fp16_cross_attention_kernel_->isValid(sequence_length)) {
fused_cross_attention_kernel = fused_fp16_cross_attention_kernel_;
}
}
bool use_fused_runner = !use_flash_attention &&
!disable_fused_self_attention_ &&
fused_cross_attention_kernel == nullptr &&
nullptr == relative_position_bias &&
(value != nullptr || key == nullptr) &&
(nullptr == past_key && nullptr == past_value && !parameters.pass_past_in_kv) &&
(nullptr == key_padding_mask || is_mask_1d_seq_len) &&
parameters.hidden_size == parameters.v_hidden_size &&
parameters.sequence_length == parameters.kv_sequence_length &&
FusedMHARunnerFP16v2::is_supported(sm, parameters.head_size, sequence_length,
enable_trt_flash_attention_, false);
if (use_fused_runner) {
// Here we assume that num_heads and head_size does not change for a MultiHeadAttention node.
if (nullptr == fused_fp16_runner_.get()) {
constexpr bool is_unidirectional = false;
std::call_once(fused_fp16_runner_created_, [&]() {
fused_fp16_runner_ = FusedMHARunnerFP16v2::Create(num_heads_, parameters.head_size, sm, is_unidirectional,
enable_trt_flash_attention_, parameters.scale);
});
}
// In case some kernel not loaded due to shared memory limit, we need to double check here.
const int S = fused_fp16_runner_->getSFromMaxSeqLen(sequence_length);
if (fused_fp16_runner_->isValid(S)) {
fused_runner = fused_fp16_runner_.get();
}
}
#if USE_MEMORY_EFFICIENT_ATTENTION
int length_threshold = this->kernel_options_->MinSeqLenForEfficientAttentionFp32();
bool is_long_sequence = sizeof(T) == 2 || // sequence length threshold is 0 for FP16
parameters.sequence_length >= length_threshold ||
parameters.kv_sequence_length >= length_threshold;
bool is_good_for_rpb = relative_position_bias != nullptr && parameters.sequence_length % (4 * sizeof(T)) == 0;
bool use_memory_efficient_attention =
!use_flash_attention &&
fused_runner == nullptr &&
fused_cross_attention_kernel == nullptr &&
!disable_memory_efficient_attention_ &&
is_long_sequence &&
!past_no_bias &&
(relative_position_bias == nullptr || is_good_for_rpb) &&
(nullptr == key_padding_mask || parameters.mask_type == AttentionMaskType::MASK_1D_KEY_SEQ_LEN_START) &&
has_memory_efficient_attention(sm, sizeof(T) == 2, parameters.head_size, parameters.v_head_size);
#else
constexpr bool use_memory_efficient_attention = false;
#endif
if (kernel_options_->AllowDebugInfo()) {
AttentionKernelDebugInfo debug_info;
debug_info.use_flash_attention = use_flash_attention;
debug_info.use_trt_cross_attention = fused_cross_attention_kernel != nullptr;
debug_info.use_efficient_attention = use_memory_efficient_attention;
if (fused_fp16_runner_ != nullptr) {
debug_info.SetTrtFusedKernel(is_unidirectional_, enable_trt_flash_attention_, sequence_length);
}
debug_info.Print("MultiHeadAttention",
this->Node().Name(),
std::is_same<T, MLFloat16>::value,
std::is_same<T, BFloat16>::value);
}
// When packed kv or packed qkv is used, there is no needed for add bias transpose thus no qkv workspace.
// TODO(tianleiwu): flash attention or memory efficient attention might not need qkv workspace sometime.
bool no_qkv_workspace = nullptr == value &&
(use_fused_cross_attention || (nullptr != fused_runner && nullptr == key)) &&
nullptr == key_padding_mask &&
nullptr == bias;
size_t workspace_bytes;
constexpr size_t element_size = sizeof(T);
if (no_qkv_workspace) {
workspace_bytes = (parameters.batch_size > kCumulatedSequenceLengthCacheMaxBatchSize) ? 2 * GetSequenceOffsetSize(parameters.batch_size, true) : 0;
} else {
workspace_bytes = GetAttentionWorkspaceSize(element_size,
parameters.batch_size,
parameters.num_heads,
parameters.head_size,
parameters.v_head_size,
parameters.sequence_length,
parameters.kv_sequence_length,
parameters.total_sequence_length,
fused_runner,
use_flash_attention,
use_fused_cross_attention,
use_memory_efficient_attention);
}
auto work_space = GetScratchBuffer<void>(workspace_bytes, context->GetComputeStream());
const size_t past_k_bytes = element_size * parameters.batch_size * parameters.kv_sequence_length * parameters.num_heads * parameters.head_size;
const size_t past_v_bytes = element_size * parameters.batch_size * parameters.kv_sequence_length * parameters.num_heads * parameters.v_head_size;
const bool use_temp_k_v_workspace = parameters.pass_past_in_kv || use_memory_efficient_attention || use_flash_attention;
auto temp_k_work_space = use_temp_k_v_workspace ? GetScratchBuffer<void>(past_k_bytes, context->GetComputeStream()) : nullptr;
auto temp_v_work_space = use_temp_k_v_workspace ? GetScratchBuffer<void>(past_v_bytes, context->GetComputeStream()) : nullptr;
typedef typename ToCudaType<T>::MappedType CudaT;
AttentionData<CudaT> data;
data.bias = (nullptr == bias) ? nullptr : reinterpret_cast<const CudaT*>(bias->Data<T>());
data.query = reinterpret_cast<const CudaT*>(query->Data<T>());
data.key = (nullptr == key || parameters.pass_past_in_kv) ? nullptr : reinterpret_cast<const CudaT*>(key->Data<T>());
data.value = (nullptr == value || parameters.pass_past_in_kv) ? nullptr : reinterpret_cast<const CudaT*>(value->Data<T>());
data.mask_index = (nullptr == key_padding_mask) ? nullptr : key_padding_mask->Data<int>();
data.mask_index_dims = (nullptr == key_padding_mask) ? gsl::span<const int64_t>() : key_padding_mask->Shape().GetDims();
data.past_key = pass_key_value_as_past ? reinterpret_cast<const CudaT*>(key->Data<T>())
: (nullptr == past_key) ? nullptr
: reinterpret_cast<const CudaT*>(past_key->Data<T>());
data.past_value = pass_key_value_as_past ? reinterpret_cast<const CudaT*>(value->Data<T>())
: (nullptr == past_value) ? nullptr
: reinterpret_cast<const CudaT*>(past_value->Data<T>());
data.relative_position_bias = (nullptr == relative_position_bias) ? nullptr : reinterpret_cast<const CudaT*>(relative_position_bias->Data<T>());
data.has_qkv_workspace = !no_qkv_workspace;
data.workspace = reinterpret_cast<CudaT*>(work_space.get());
data.temp_k_workspace = use_temp_k_v_workspace ? reinterpret_cast<CudaT*>(temp_k_work_space.get()) : nullptr;
data.temp_v_workspace = use_temp_k_v_workspace ? reinterpret_cast<CudaT*>(temp_v_work_space.get()) : nullptr;
data.output = reinterpret_cast<CudaT*>(output->MutableData<T>());
data.present_key = (nullptr == present_key) ? nullptr : reinterpret_cast<CudaT*>(present_key->MutableData<T>());
data.present_value = (nullptr == present_value) ? nullptr : reinterpret_cast<CudaT*>(present_value->MutableData<T>());
data.fused_runner = reinterpret_cast<void*>(fused_runner);
data.fused_cross_attention_kernel = fused_cross_attention_kernel;
data.use_flash_attention = use_flash_attention;
data.use_memory_efficient_attention = use_memory_efficient_attention;
data.cumulated_sequence_length_q_cache = &(this->cumulated_sequence_length_q_cache_);
data.cumulated_sequence_length_kv_cache = &(this->cumulated_sequence_length_kv_cache_);
if (softmax_lse_accum_buffer != nullptr) {
data.softmax_lse_accum = reinterpret_cast<CudaT*>(softmax_lse_accum_buffer.get());
}
if (out_accum_buffer != nullptr) {
data.out_accum = reinterpret_cast<CudaT*>(out_accum_buffer.get());
}
cublasHandle_t cublas = GetCublasHandle(context);
return QkvToContext<CudaT>(
device_prop, cublas, context->GetComputeStream(), parameters, data);
}
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime
Reference in a new issue View git blame Copy permalink