From c7ced7a5e9c5ba1245b03f84edd537f45432b2e9 Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Tue, 21 Mar 2023 12:59:29 -0700 Subject: [PATCH] Add PackedAttention for packing mode (#14858) ### Description Transformer models can handle batch of inputs at once. However, sequences in a batch usually have different length. Then we have to pad the short one to have same length as the longest. This is not efficient especially for large batch with high variance. This PR introduces a PackedAttention operator which can take in packed sequences (no padding) and also produces output in packing mode. There will be another PR to use the PackedAttention to implement the encoder in packing mode. ### Motivation and Context --- docs/ContribOperators.md | 73 +++ docs/OperatorKernels.md | 1 + .../contrib_ops/cpu/bert/attention_common.h | 16 + .../contrib_ops/cuda/bert/attention_impl.cu | 2 +- .../contrib_ops/cuda/bert/attention_softmax.h | 157 ++++- .../contrib_ops/cuda/bert/packed_attention.cc | 324 ++++++++++ .../contrib_ops/cuda/bert/packed_attention.h | 47 ++ .../cuda/bert/packed_attention_impl.cu | 610 ++++++++++++++++++ .../cuda/bert/packed_attention_impl.h | 53 ++ .../contrib_ops/cuda/cuda_contrib_kernels.cc | 4 + .../core/graph/contrib_ops/bert_defs.cc | 127 ++++ onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 + .../contrib_ops/packed_attention_op_test.cc | 502 ++++++++++++++ .../test/providers/provider_test_utils.cc | 7 +- .../test/providers/provider_test_utils.h | 2 +- .../test/testdata/packed_attention_fp16.onnx | Bin 0 -> 1228 bytes .../test/testdata/packed_attention_fp32.onnx | Bin 0 -> 1228 bytes 17 files changed, 1915 insertions(+), 12 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h create mode 100644 onnxruntime/test/contrib_ops/packed_attention_op_test.cc create mode 100644 onnxruntime/test/testdata/packed_attention_fp16.onnx create mode 100644 onnxruntime/test/testdata/packed_attention_fp32.onnx diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 66afbebae5..a62d757aa9 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -51,6 +51,7 @@ Do not modify directly.* * com.microsoft.NGramRepeatBlock * com.microsoft.NhwcConv * com.microsoft.NhwcMaxPool + * com.microsoft.PackedAttention * com.microsoft.Pad * com.microsoft.QAttention * com.microsoft.QGemm @@ -2607,6 +2608,78 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.PackedAttention** + + This is the packed version of Attention. + + Sequences in one batch usually don't have same length and they are padded to have same length, + e.g., below is a batch with 3 sequences and tokens* are padded. + Sequence_0: 0, 1*, 2*, 3* + Sequence_1: 4, 5, 6*, 7* + Sequence_2: 8, 9, 10, 11 + + PackedAttention is designed to takes in packed input, i.e., only the real tokens without padding. + An input as above will be packed into 3 tensors like below: + - input ([h0, h4, h5, h8, h9, h10, h11]) + - token_offset: 0, 4, 5, 8, 9, 10, 11, 1*, 2*, 3*, 6*, 7* + - cumulated_token_count: 0, 1, 1+2, 1+2+4 + + Input tensors contains the hidden embedding of real tokens. + Token_offset records the offset of token in the unpacked input. + cumulated_token_count records cumulated length of each sequnces length. + + The operator only supports BERT like model with padding on right now. + + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
num_heads : int (required)
+
Number of attention heads
+
qkv_hidden_sizes : list of ints
+
Hidden dimension of Q, K, V: hidden_size, hidden_size and v_hidden_size
+
scale : float
+
Custom scale will be used if specified. Default value is 1/sqrt(head_size)
+
+ +#### Inputs (5 - 6) + +
+
input : T
+
Input tensor with shape (token_count, input_hidden_size)
+
weights : T
+
Merged Q/K/V weights with shape (input_hidden_size, hidden_size + hidden_size + v_hidden_size)
+
bias : T
+
Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) for input projection
+
token_offset : M
+
In packing mode, it specifies the offset of each token(batch_size, sequence_length).
+
cumulative_sequence_length : M
+
A tensor with shape (batch_size + 1). It specifies the cumulative sequence length.
+
relative_position_bias (optional) : T
+
A tensor with shape (batch_size, num_heads, sequence_length, sequence_length)or (1, num_heads, sequence_length, sequence_length).It specifies the additional bias to QxK'
+
+ +#### Outputs + +
+
output : T
+
2D output tensor with shape (token_count, v_hidden_size)
+
+ +#### Type Constraints + +
+
T : tensor(float), tensor(float16)
+
Constrain input and output types to float tensors.
+
M : tensor(int32)
+
Constrain mask index to integer types
+
+ + ### **com.microsoft.Pad** Given `data` tensor, pads, mode, and value. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 3bd890e5b4..845defa7ef 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -816,6 +816,7 @@ Do not modify directly.* |MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* relative_position_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**T** = tensor(float), tensor(float16)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| |NhwcConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| +|PackedAttention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* relative_position_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)| |QOrderedAttention|*in* input:**Q**
*in* scale_input:**S**
*in* scale_Q_gemm:**S**
*in* scale_K_gemm:**S**
*in* scale_V_gemm:**S**
*in* Q_weight:**Q**
*in* K_weight:**Q**
*in* V_weight:**Q**
*in* scale_Q_weight:**S**
*in* scale_K_weight:**S**
*in* scale_V_weight:**S**
*in* Q_bias:**S**
*in* K_bias:**S**
*in* V_bias:**S**
*in* scale_QKT_gemm:**S**
*in* scale_QKT_softmax:**S**
*in* scale_values_gemm:**S**
*in* mask_index:**G**
*in* past:**Q**
*in* relative_position_bias:**S**
*out* output:**Q**|1+|**G** = tensor(int32)
**Q** = tensor(int8)
**S** = tensor(float)| |QOrderedGelu|*in* X:**Q**
*in* scale_X:**S**
*in* scale_Y:**S**
*out* Y:**Q**|1+|**Q** = tensor(int8)
**S** = tensor(float)| diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_common.h b/onnxruntime/contrib_ops/cpu/bert/attention_common.h index 292c1aae1e..680f875d23 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_common.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_common.h @@ -59,6 +59,22 @@ struct AttentionParameters { AttentionMaskType mask_type; }; +// Parameters deduced from node attributes and inputs/outputs. +struct PackedAttentionParameters { + int batch_size; + int sequence_length; + int input_hidden_size; // hidden size of input + int hidden_size; // hidden size of Q or K + int head_size; // hidden size per head of Q or K + int v_hidden_size; // hidden size of V + int v_head_size; // hidden size per head of V + int num_heads; + float scale; + int token_count; + bool has_relative_position_bias; + bool broadcast_res_pos_bias; +}; + namespace attention { // Environment variable to enable or disable TRT fused self attention kernel. Default is 0 (enabled). constexpr const char* kDisableFusedSelfAttention = "ORT_DISABLE_FUSED_ATTENTION"; diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 9f655bf569..28daf5d4af 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -297,7 +297,7 @@ Status PrepareQkv(contrib::AttentionParameters& parameters, bool use_fused_causal = (nullptr != fused_runner && parameters.is_unidirectional); // Default format for memory efficient attention. - // When there is past state, the format shal be BxNxSxH, so we disable memory efficient attention when there is past. + // When there is past state, the format shall be BxNxSxH, so we disable memory efficient attention when there is past. DUMP_TENSOR_INIT(); if (nullptr != data.gemm_buffer) { if (data.bias == nullptr) { diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h index d2cb5d2d58..db934b0c1d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_softmax.h @@ -36,7 +36,6 @@ namespace cuda { template __device__ inline void Softmax(const int all_sequence_length, - const int sequence_length, const int valid_end, const int valid_start, const T* rel_pos_bias, @@ -483,12 +482,11 @@ __global__ void SoftmaxKernelSmall(const int all_sequence_length, template __global__ void SoftmaxKernel(const int all_sequence_length, - const int sequence_length, const T* rel_pos_bias, const bool broadcast_rel_pos_bias, const T* input, T* output) { - Softmax(all_sequence_length, sequence_length, all_sequence_length, 0, + Softmax(all_sequence_length, all_sequence_length, 0, rel_pos_bias, broadcast_rel_pos_bias, input, output); } @@ -524,7 +522,7 @@ Status ComputeSoftmax(cudaStream_t stream, const int all_sequence_length, const } else if (!is_unidirectional) { const int blockSize = 1024; SoftmaxKernel<<>>( - all_sequence_length, sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output); + all_sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output); } else { const int blockSize = 256; const int sh_bytes = sizeof(float) * all_sequence_length; @@ -566,9 +564,102 @@ __global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length, rel_pos_bias, broadcast_rel_pos_bias, input, output, is_unidirectional); } +template +__device__ inline void SoftmaxSmallPacked(const int sequence_length, + const int end, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const T* input, + T* output) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + + __shared__ float sum_reverse_block; + __shared__ float max_block; + + // Input dimension is BxNxSxS*; blockIdx.y is batch index b; gridDim.x=N*S; blockIdx.x is index within N*S; + const int offset = (blockIdx.y * gridDim.x + blockIdx.x) * sequence_length; + const int index = offset + threadIdx.x; + + bool is_valid = threadIdx.x < end; + + // e^x is represented as infinity if x is large enough, like 100.f. + // Infinity divided by Infinity is a NAN. Thus, softmax gets a NAN if one or more item are large enough. + // a math transform as below is leveraged to get a stable softmax: + // e^xi/(e^x1 + ...e^xn) = e^(xi - max) / (e^(x1 - max) + ... + e^(xn - max)) + const bool no_rpb = (rel_pos_bias == nullptr); + const int size_per_batch = gridDim.x * sequence_length; + float input_data = no_rpb + ? float(input[index]) + : float(input[index] + (broadcast_rel_pos_bias + ? rel_pos_bias[index % size_per_batch] + : rel_pos_bias[index])); + + float thread_data_max = is_valid ? input_data : float(-CUDART_INF_F); + const auto max = BlockReduce(tmp_storage).Reduce(thread_data_max, cub::Max(), end); + + // Store max value + if (threadIdx.x == 0) { + max_block = max; + } + __syncthreads(); + + float thread_data_exp(0.f); + if (is_valid) { + thread_data_exp = expf(input_data - max_block); + } + + const auto sum = BlockReduce(tmp_storage).Reduce(thread_data_exp, cub::Sum(), end); + + // Store value of 1.0/sum. + if (threadIdx.x == 0) { + sum_reverse_block = (1.f) / sum; + } + __syncthreads(); + + // threadIdx.x might be larger than all_sequence_length due to alignment to 32x. + if (threadIdx.x < sequence_length) { + output[index] = T(thread_data_exp * sum_reverse_block); + } +} + +template +__global__ void SoftmaxKernelSmallWithCumSeqLen(const T* input, + const T* rel_pos_bias, const bool broadcast_rel_pos_bias, + const int* cum_seq_length, const int sequence_length, + T* output) { + __shared__ int end_position; + + if (threadIdx.x == 0) { + const int batch = blockIdx.y; + end_position = cum_seq_length[batch + 1] - cum_seq_length[batch]; + } + __syncthreads(); + + SoftmaxSmallPacked(sequence_length, end_position, + rel_pos_bias, broadcast_rel_pos_bias, + input, output); +} + +template +__global__ void SoftmaxKernelWithCumSeqLen(const T* input, + const T* rel_pos_bias, const bool broadcast_rel_pos_bias, + const int* cum_seq_length, const int sequence_length, + T* output) { + __shared__ int end_position; + + if (threadIdx.x == 0) { + const int batch = blockIdx.y; + end_position = cum_seq_length[batch + 1] - cum_seq_length[batch]; + } + __syncthreads(); + + Softmax(sequence_length, end_position, 0 /*start_position*/, + rel_pos_bias, broadcast_rel_pos_bias, input, output); +} + template __global__ void MaskedSoftmaxKernel(const int all_sequence_length, - const int sequence_length, const int* mask_end, const int* mask_start, const T* rel_pos_bias, @@ -590,7 +681,7 @@ __global__ void MaskedSoftmaxKernel(const int all_sequence_length, } __syncthreads(); - Softmax(all_sequence_length, sequence_length, end_position, start_position, + Softmax(all_sequence_length, end_position, start_position, rel_pos_bias, broadcast_rel_pos_bias, input, output); } @@ -616,6 +707,58 @@ __global__ void SoftmaxWithRawMaskSmallKernel(const int all_sequence_length, skip_softmax, mask_filter_value); } +template +Status ComputeSoftmaxWithCumSeqLength( + const T* input, + const T* rel_pos_bias, + const bool broadcast_rel_pos_bias, + const int32_t* cum_seq_length, + const int batch_size, + const int sequence_length, + const int num_heads, + T* output, cudaStream_t stream) { + const dim3 grid(sequence_length * num_heads, batch_size, 1); + + if (sequence_length <= 32) { + const int blockSize = 32; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + + } else if (sequence_length <= 64) { + const int blockSize = 64; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 128) { + const int blockSize = 128; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 256) { + const int blockSize = 256; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 512) { + const int blockSize = 512; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else if (sequence_length <= 1024) { + const int blockSize = 1024; + SoftmaxKernelSmallWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } else { + SoftmaxKernelWithCumSeqLen + <<>>(input, rel_pos_bias, broadcast_rel_pos_bias, + cum_seq_length, sequence_length, output); + } + + return CUDA_CALL(cudaGetLastError()); +} + template Status ComputeSoftmaxWithMask1D(cudaStream_t stream, const int all_sequence_length, @@ -664,7 +807,7 @@ Status ComputeSoftmaxWithMask1D(cudaStream_t stream, } else if (!is_unidirectional) { const int blockSize = 1024; MaskedSoftmaxKernel - <<>>(all_sequence_length, sequence_length, mask_index, mask_start, + <<>>(all_sequence_length, mask_index, mask_start, rel_pos_bias, broadcast_rel_pos_bias, input, output); } else { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention CUDA operator does not support total sequence length > 1024."); diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc new file mode 100644 index 0000000000..61726e63c7 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc @@ -0,0 +1,324 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/packed_attention.h" + +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" +#include "core/platform/env_var_utils.h" +#include "contrib_ops/cuda/bert/packed_attention_impl.h" +#include "contrib_ops/cuda/bert/bert_padding.h" +#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.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( \ + PackedAttention, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + PackedAttention); + +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +PackedAttention::PackedAttention(const OpKernelInfo& info) : CudaKernel(info) { + int64_t num_heads = 0; + ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); + num_heads_ = static_cast(num_heads); + + scale_ = info.GetAttrOrDefault("scale", 0.0f); + + if (!info.GetAttrs("qkv_hidden_sizes", qkv_hidden_sizes_).IsOK()) { + qkv_hidden_sizes_.clear(); + } + + disable_fused_runner_ = sizeof(T) != 2 || + ParseEnvironmentVariableWithDefault(attention::kDisableFusedSelfAttention, false); + + enable_trt_flash_attention_ = sizeof(T) == 2 && + !ParseEnvironmentVariableWithDefault(attention::kDisableTrtFlashAttention, false); +} + +template +Status PackedAttention::CheckInputs(const TensorShape& input_shape, + const TensorShape& weights_shape, + const TensorShape& bias_shape, + const TensorShape& token_offset_shape, + const TensorShape& cu_seq_len_shape, + const Tensor* relative_position_bias, + PackedAttentionParameters& parameters) const { + // Abbreviation and Meanings: + // T: token_count + // B: batch_size + // S: sequence_length (input sequence length of query) + // N: num_heads + // H: head size for Q and K, aka q_head_size or v_head_size or qk_head_size + // H_v: v_head_size + // D_i: input hidden size + // D: hidden size for Q and K (D = N * H), aka q_hidden_size or k_hidden_size or qk_hidden_size + // D_v: v_hidden_size = num_heads * v_head_size + + // Input shapes: + // input: : (T, D_i) + // weights (Q/K/V) : (D_i, D + D + D_v) + // bias (Q/K/V) : (D + D + D_v) + // token_offset : (B, S) + // cu_seq_len_shape : (B + 1) + // relative_position_bias : (B, N, S, S), (1, N, S, S) or NULL + + const auto& input_dims = input_shape.GetDims(); + if (input_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'input' is expected to have 2 dimensions in packing mode, got ", + input_dims.size()); + } + int64_t token_count = input_dims[0]; + int64_t input_hidden_size = input_dims[1]; + + const auto& token_offset_dims = token_offset_shape.GetDims(); + if (token_offset_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'packing_token_offset' is expected to have 2 dimensions in packing mode, got ", + token_offset_dims.size()); + } + + int64_t batch_size = token_offset_dims[0]; + int64_t sequence_length = token_offset_dims[1]; + + const auto& bias_dims = bias_shape.GetDims(); + if (bias_dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' is expected to have 1 dimension, got ", + bias_dims.size()); + } + + const auto& weights_dims = weights_shape.GetDims(); + if (weights_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' is expected to have 2 dimensions, got ", + weights_dims.size()); + } + if (weights_dims[0] != input_hidden_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 1 dimension 0 should have same length as dimension 2 of input 0"); + } + + if (bias_dims[0] != weights_dims[1]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'bias' dimension 0 should have same length as dimension 1 of input 'weights'"); + } + + const auto& cu_seq_len_dims = cu_seq_len_shape.GetDims(); + if (cu_seq_len_dims.size() != 1 || cu_seq_len_dims[0] != batch_size + 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'cumulative_sequence_length' should have 1 dimension with size equal to batch_size + 1"); + } + + int64_t q_hidden_size = bias_dims[0] / static_cast(3); + int64_t k_hidden_size = q_hidden_size; + int64_t v_hidden_size = k_hidden_size; + if (qkv_hidden_sizes_.size() != 0) { + if (qkv_hidden_sizes_.size() != 3) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "qkv_hidden_sizes attribute should have 3 elements"); + } + + for (size_t i = 0; i < qkv_hidden_sizes_.size(); i++) { + if (qkv_hidden_sizes_[i] % num_heads_ != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "hidden_size should be divisible by num_heads:", qkv_hidden_sizes_[i]); + } + } + + q_hidden_size = qkv_hidden_sizes_[0]; + k_hidden_size = qkv_hidden_sizes_[1]; + v_hidden_size = qkv_hidden_sizes_[2]; + } + + if (q_hidden_size != k_hidden_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "qkv_hidden_sizes first element should be same as the second"); + } + + if (bias_dims[0] != q_hidden_size + k_hidden_size + v_hidden_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'bias' dimension 0 should have same length as sum of Q/K/V hidden sizes:", + " q_hidden_size=", q_hidden_size, " k_hidden_size=", k_hidden_size, " v_hidden_size=", + v_hidden_size, "bias_dims[0]=", bias_dims[0]); + } + + bool broadcast_res_pos_bias = false; + if (relative_position_bias != nullptr) { + const auto& relative_position_bias_dims = relative_position_bias->Shape().GetDims(); + + if (relative_position_bias_dims.size() != 4) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'relative_position_bias' is expected to have 4 dimensions, got ", + relative_position_bias_dims.size()); + } + + if (relative_position_bias_dims[0] != batch_size && relative_position_bias_dims[0] != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'relative_position_bias' dimension 0 should be same as batch_size or 1, got ", + relative_position_bias_dims[0]); + } + if (relative_position_bias_dims[0] == 1) { + broadcast_res_pos_bias = true; + } + + if (relative_position_bias_dims[1] != num_heads_) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'relative_position_bias' dimension 1 should be same as number of heads, got ", + relative_position_bias_dims[1]); + } + + if (relative_position_bias_dims[2] != sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'relative_position_bias' dimension 2 should be same as sequence_length, got ", + relative_position_bias_dims[2]); + } + + if (relative_position_bias_dims[3] != sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'relative_position_bias' dimension 3 should be same as sequence_length, got ", + relative_position_bias_dims[3]); + } + } + + parameters.batch_size = static_cast(batch_size); + parameters.sequence_length = static_cast(sequence_length); + parameters.input_hidden_size = static_cast(input_hidden_size); + parameters.hidden_size = static_cast(q_hidden_size); + parameters.v_hidden_size = static_cast(v_hidden_size); + parameters.head_size = static_cast(q_hidden_size) / num_heads_; + parameters.v_head_size = static_cast(v_hidden_size) / num_heads_; + parameters.num_heads = num_heads_; + parameters.scale = scale_; + parameters.token_count = static_cast(token_count); + parameters.has_relative_position_bias = nullptr != relative_position_bias; + parameters.broadcast_res_pos_bias = broadcast_res_pos_bias; + + return Status::OK(); +} + +template +MHARunner* PackedAttention::TryGettingFusedRunner(const PackedAttentionParameters& parameters) const { + MHARunner* fused_runner = nullptr; + + bool use_fused_runner = !disable_fused_runner_ && + !parameters.has_relative_position_bias && + parameters.hidden_size == parameters.v_hidden_size; + + if(!use_fused_runner) { + return fused_runner; + } + + // Check whether we can use fused kernel + auto& device_prop = GetDeviceProp(); + int sm = device_prop.major * 10 + device_prop.minor; + bool is_fMHA_supported = FusedMHARunnerFP16v2::is_supported(sm, + parameters.head_size, + parameters.sequence_length, + enable_trt_flash_attention_, + false); + + if(!is_fMHA_supported) { + return fused_runner; + } + + // Assuming that num_heads and head_size do not change. + if (nullptr == fused_fp16_runner_.get()) { + fused_fp16_runner_.reset(new FusedMHARunnerFP16v2(num_heads_, parameters.head_size, sm, false /* causal_mask*/, + 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(parameters.sequence_length); + if (fused_fp16_runner_->isValid(S)) { + fused_runner = fused_fp16_runner_.get(); + } + + return fused_runner; +} + +template +Status PackedAttention::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const Tensor* weights = context->Input(1); + const Tensor* bias = context->Input(2); + const Tensor* token_offset = context->Input(3); + const Tensor* cumulative_sequence_length = context->Input(4); + const Tensor* relative_position_bias = context->Input(5); + + PackedAttentionParameters parameters; + ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), + weights->Shape(), + bias->Shape(), + token_offset->Shape(), + cumulative_sequence_length->Shape(), + relative_position_bias, + parameters)); + + TensorShapeVector output_shape{parameters.token_count, parameters.v_hidden_size}; + Tensor* output = context->Output(0, output_shape); + + MHARunner* fused_runner = TryGettingFusedRunner(parameters); + + typedef typename ToCudaType::MappedType CudaT; + CudaT one = ToCudaType::FromFloat(1.0f); + CudaT zero = ToCudaType::FromFloat(0.0f); + + IAllocatorUniquePtr gemm_buffer; + int m = parameters.token_count; + int n = parameters.hidden_size + parameters.hidden_size + parameters.v_hidden_size; + int k = parameters.input_hidden_size; + gemm_buffer = GetScratchBuffer(static_cast(m) * n, context->GetComputeStream()); + + auto& device_prop = GetDeviceProp(); + cublasHandle_t cublas = GetCublasHandle(context); + + // Gemm, note that CUDA assumes col-major, so result(N, M) = 1 * weights x input + 1 x bias + // The bias part is not included here since we fuse bias, transpose and output 3 matrice into one cuda kernel. + CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( + cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &one, + reinterpret_cast(weights->Data()), n, + reinterpret_cast(input->Data()), k, + &zero, reinterpret_cast(gemm_buffer.get()), n, device_prop)); + + constexpr size_t element_size = sizeof(T); + size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, + parameters.batch_size, + parameters.num_heads, + parameters.head_size, + parameters.v_head_size, + parameters.sequence_length, + fused_runner); + auto work_space = GetScratchBuffer(workSpaceSize, context->GetComputeStream()); + + typedef typename ToCudaType::MappedType CudaT; + PackedAttentionData data; + data.gemm_buffer = reinterpret_cast(gemm_buffer.get()); + data.bias = reinterpret_cast(bias->Data()); + data.relative_position_bias = (nullptr == relative_position_bias) ? nullptr : reinterpret_cast(relative_position_bias->Data()); + data.workspace = reinterpret_cast(work_space.get()); + data.token_offset = token_offset->Data(); + data.cumulative_sequence_length = cumulative_sequence_length->Data(); + data.output = reinterpret_cast(output->MutableData()); + data.fused_runner = reinterpret_cast(fused_runner); + + return QkvToContext(device_prop, cublas, Stream(context), parameters, data); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention.h new file mode 100644 index 0000000000..873a1a3c3b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention.h @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/providers/cuda/cuda_kernel.h" +#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h" +#include "contrib_ops/cpu/bert/attention_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +template +class PackedAttention final : public CudaKernel { + public: + PackedAttention(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + Status CheckInputs(const TensorShape& input_shape, + const TensorShape& weights_shape, + const TensorShape& bias_shape, + const TensorShape& packing_token_offset_shape, + const TensorShape& cu_seq_len_shape, + const Tensor* relative_position_bias, + PackedAttentionParameters& parameters) const; + + MHARunner* TryGettingFusedRunner(const PackedAttentionParameters& parameters) const; + + private: + int32_t num_heads_; // number of attention heads + std::vector qkv_hidden_sizes_; // Q, K, V hidden sizes parsed from the qkv_hidden_sizes attribute. + float scale_; // the scale to be used for softmax + bool disable_fused_runner_; + bool enable_trt_flash_attention_; + mutable std::unique_ptr fused_fp16_runner_; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu new file mode 100644 index 0000000000..1729975453 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu @@ -0,0 +1,610 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" +#include "contrib_ops/cuda/bert/packed_attention_impl.h" +#include "contrib_ops/cuda/bert/attention_softmax.h" +#include "contrib_ops/cuda/bert/transformer_common.h" +#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h" +#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h" +#include "contrib_ops/cuda/bert/bert_padding.h" +#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h" +#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" +#include "contrib_ops/cuda/bert/rotary_embedding_util.h" + +using namespace onnxruntime::cuda; +using namespace cub; + +#define CHECK_CUDA(expr) CUDA_RETURN_IF_ERROR(expr) + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +constexpr size_t kCUDAMemoryAlignment = 256; + +constexpr int32_t kMAX_THREADS_PER_BLOCK = 256; + +size_t GetAttentionScratchSize( + size_t element_size, + size_t batch_size, + size_t num_heads, + size_t sequence_length) { + const size_t bytes = element_size * batch_size * num_heads * sequence_length * sequence_length; + return ((bytes + kCUDAMemoryAlignment - 1) / kCUDAMemoryAlignment) * kCUDAMemoryAlignment; +} + +size_t GetAttentionWorkspaceSize( + size_t element_size, + size_t batch_size, + size_t num_heads, + size_t qk_head_size, + size_t v_head_size, + size_t sequence_length, + void* fused_runner) { + // Note that q, k and v might need alignment for fused attention kernels. + const size_t qkv_bytes = element_size * batch_size * num_heads * sequence_length * (qk_head_size + qk_head_size + v_head_size); + + if (fused_runner != nullptr) { + return qkv_bytes; + } + + return qkv_bytes + 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length); +} + +template +__global__ void AddBiasTransposeQKVPacked(const T* input, + const T* biases, + int32_t N, + int32_t H_QK, + int32_t H_V, + T* q, + T* k, + T* v, + const int32_t* token_offset, + int32_t token_count); + +// Grid: (S, B) +// Block: 256 +// For unfused PackedAttention +// Input: Tx3xNxH +// Output: 3xBxNxSxH +// Where: +// T is token_count +// B is batch_size +// S is sequence_length +// N is num_heads +// H is head_size +template +__global__ void AddBiasTransposeQKVPacked( + const T* input, + const T* biases, + int32_t N, + int32_t H_QK, + int32_t H_V, + T* q, + T* k, + T* v, + const int32_t* token_offset, + int32_t token_count) { + int s = blockIdx.x; + int b = blockIdx.y; + + int S = gridDim.x; + + const int packing_token_idx = b * S + s; + const int padding_token_idx = token_offset[packing_token_idx]; + b = padding_token_idx / S; + s = padding_token_idx - b * S; + + input += packing_token_idx * N * (H_QK + H_QK + H_V); + int k_offset = N * H_QK; + int v_offset = N * H_QK + N * H_QK; + q += (b * N * S + s) * H_QK; + k += (b * N * S + s) * H_QK; + v += (b * N * S + s) * H_V; + + if (packing_token_idx < token_count) { + for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) { + int h = i % H_QK; + int n = i / H_QK; + q[n * S * H_QK + h] = input[i] + biases[i]; + k[n * S * H_QK + h] = input[i + k_offset] + biases[i + k_offset]; + } + + for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) { + int h = i % H_V; + int n = i / H_V; + v[n * S * H_V + h] = input[i + v_offset] + biases[i + v_offset]; + } + } else { + for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) { + int h = i % H_QK; + int n = i / H_QK; + q[n * S * H_QK + h] = biases[i]; + k[n * S * H_QK + h] = biases[i + k_offset]; + } + + for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) { + int h = i % H_V; + int n = i / H_V; + v[n * S * H_V + h] = biases[i + v_offset]; + } + } +} + +// Grid: (S, B) +// Block: 256 +// For memory efficient fMHA from CUTLASS. For future use, doesn't support fMHA from CUTLASS yet. +// Input: Tx3xNxH +// Output: 3xBxNxSxH +// T is token_count +// B is batch_size +// S is sequence_length +// N is num_heads +// H is head_size +template +__global__ void AddBiasTransposeQKVPackedCutlass( + const T* input, + const T* biases, + int32_t N, + int32_t H_QK, + int32_t H_V, + T* q, + T* k, + T* v, + const int32_t* token_offset, + int32_t token_count) { + int s = blockIdx.x; + int b = blockIdx.y; + + int S = gridDim.x; + + const int packing_token_idx = b * S + s; + const int padding_token_idx = token_offset[packing_token_idx]; + b = padding_token_idx / S; + s = padding_token_idx - b % S; + + input += packing_token_idx * N * (H_QK + H_QK + H_V); + int k_offset = N * H_QK; + int v_offset = N * H_QK + N * H_QK; + q += (b * S * N + s * N) * H_QK; + k += (b * S * N + s * N) * H_QK; + v += (b * S * N + s * N) * H_V; + + if (packing_token_idx < token_count) { + for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) { + q[i] = input[i] + biases[i]; + k[i] = input[i + k_offset] + biases[i + k_offset]; + } + + for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) { + v[i] = input[i + v_offset] + biases[i + v_offset]; + } + } else { + for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) { + q[i] = biases[i]; + k[i] = biases[i + k_offset]; + } + + for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) { + v[i] = biases[i + v_offset]; + } + } +} + +// Grid: (S, B) +// Block: 256 +// For fMHA from TRT +// Input: Tx3xNxH +// Output: TxNx3xH +// T is token_count +// B is batch_size +// S is sequence_length +// N is num_heads +// H is head_size +template +__global__ void AddBiasTransposeQKVPackedTRT( + const T* input, + const T* biases, + int32_t N, + int32_t H, + T* output) { + int token_idx = blockIdx.x; + + int Hx3 = H * 3; + int NxH = N * H; + int NxHx2 = N * H + N * H; + + int offset = token_idx * N * Hx3; + input += offset; + output += offset; + + for (int i = threadIdx.x; i < N * H; i += blockDim.x) { + int n = i / H; + int h = i % H; + output[n * Hx3 + h] = input[i] + biases[i]; + output[n * Hx3 + H + h] = input[i + NxH] + biases[i + NxH]; + output[n * Hx3 + H + H + h] = input[i + NxHx2] + biases[i + NxHx2]; + } +} + +template +void InvokeAddBiasTranspose( + const T* input, const T* biases, T* output, + const int batch_size, const int sequence_length, + const int num_heads, const int qk_head_size, const int v_head_size, + AttentionQkvFormat format, const int32_t* token_offset, int32_t token_count, + cudaStream_t stream) { + if (format == AttentionQkvFormat::Q_K_V_BNSH) { + const dim3 grid(sequence_length, batch_size); + AddBiasTransposeQKVPacked<<>>( + input, + biases, + num_heads, + qk_head_size, + v_head_size, + output, + output + batch_size * sequence_length * num_heads * qk_head_size, + output + 2 * batch_size * sequence_length * num_heads * qk_head_size, + token_offset, + token_count); + } else if (format == AttentionQkvFormat::Q_K_V_BSNH) { // TODO: add memory efficient support + const dim3 grid(sequence_length, batch_size); + AddBiasTransposeQKVPackedCutlass<<>>( + input, + biases, + num_heads, + qk_head_size, + v_head_size, + output, + output + batch_size * sequence_length * num_heads * qk_head_size, + output + 2 * batch_size * sequence_length * num_heads * qk_head_size, + token_offset, + token_count); + } else { + ORT_ENFORCE(format == AttentionQkvFormat::QKV_BSN3H); + const dim3 grid(token_count); + AddBiasTransposeQKVPackedTRT<<>>( + input, + biases, + num_heads, + qk_head_size, + output); + } +} + +template +struct T4; + +template <> +struct T4 { + using Type = float4; +}; + +template <> +struct T4 { + using Type = Half4; +}; + +template +struct T2; + +template <> +struct T2 { + using Type = float2; +}; + +template <> +struct T2 { + using Type = half2; +}; + +template +void LaunchAddBiasTranspose( + const T* input, const T* biases, T* output, + const int batch_size, const int sequence_length, + const int num_heads, const int qk_head_size, const int v_head_size, + AttentionQkvFormat format, const int32_t* token_offset, int32_t token_count, + cudaStream_t stream) { + if (0 == (qk_head_size & 3) && 0 == (v_head_size & 3)) { + using T4Type = typename T4::Type; + const int H = qk_head_size / 4; + const int H_v = v_head_size / 4; + const T4Type* input2 = reinterpret_cast(input); + const T4Type* biases2 = reinterpret_cast(biases); + T4Type* output2 = reinterpret_cast(output); + InvokeAddBiasTranspose( + input2, biases2, output2, + batch_size, sequence_length, + num_heads, H, H_v, + format, token_offset, token_count, stream); + } else if (0 == (qk_head_size & 1) && 0 == (v_head_size & 1)) { + using T2Type = typename T2::Type; + const int H = qk_head_size / 2; + const int H_v = v_head_size / 2; + const T2Type* input2 = reinterpret_cast(input); + const T2Type* biases2 = reinterpret_cast(biases); + T2Type* output2 = reinterpret_cast(output); + InvokeAddBiasTranspose( + input2, biases2, output2, + batch_size, sequence_length, + num_heads, H, H_v, + format, token_offset, token_count, stream); + } else { + InvokeAddBiasTranspose( + input, biases, output, + batch_size, sequence_length, + num_heads, qk_head_size, v_head_size, + format, token_offset, token_count, stream); + } +} + +// Input: BxNxSxH +// Output: TxNxH +// where: +// T is token_count +// B is batch_size +// S is sequence_length +// N is num_heads +// H is head_size + +// Grid: T +// Block: 256 +template +__global__ void __launch_bounds__(kMAX_THREADS_PER_BLOCK) + TransposeRemovePadding(T* target, const T* source, const int* token_offset, + const int B, const int S, const int N, const int H) { + int token_idx = blockIdx.x; + int source_idx = token_offset[token_idx]; + int b = source_idx / S; + int s = source_idx - b * S; + + target += token_idx * N * H; + source += b * N * S * H + s * H; + for (int i = threadIdx.x; i < N * H; i += blockDim.x) { + int n = i / H; + int h = i - n * H; + target[i] = source[n * S * H + h]; + } +} + +template +Status LaunchTransposeRemovePadding( + T* output, const T* input, + const int* token_offset, const int token_count, + const int batch_size, const int seq_len, const int number_heads, const int head_size, + cudaStream_t stream); + + // input: [batch_size, number_heads, seq_len, head_size] + // output: [token_count, number_heads * head_size] +template <> +Status LaunchTransposeRemovePadding( + half* output, const half* input, + const int* token_offset, const int token_count, + const int batch_size, const int seq_len, const int number_heads, const int head_size, + cudaStream_t stream) { + + // Make sure memory is aligned to 128 bit + ORT_ENFORCE(!(reinterpret_cast(input) & 0xF) && !(reinterpret_cast(output) & 0xF), "alignment"); + + if (head_size % 8 == 0) { + const int4* input2 = reinterpret_cast(input); + int4* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 8); + } else if (head_size % 4 == 0) { + const int64_t* input2 = reinterpret_cast(input); + int64_t* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 4); + } else if (head_size % 2 == 0) { + const int32_t* input2 = reinterpret_cast(input); + int32_t* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 2); + } else { + const int16_t* input2 = reinterpret_cast(input); + int16_t* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size); + } + + return CUDA_CALL(cudaGetLastError()); +} + +// input: [batch_size, number_heads, seq_len, head_size] +// output: [token_count, number_heads * head_size] +template <> +Status LaunchTransposeRemovePadding( + float* output, const float* input, + const int* token_offset, const int token_count, + const int batch_size, const int seq_len, const int number_heads, const int head_size, + cudaStream_t stream) { + ORT_ENFORCE(!(reinterpret_cast(input) & 0xF) && !(reinterpret_cast(output) & 0xF), "alignment"); + + if (head_size % 4 == 0) { + const int4* input2 = reinterpret_cast(input); + int4* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 4); + } else if (head_size % 2 == 0) { + const int64_t* input2 = reinterpret_cast(input); + int64_t* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size / 2); + } else { + const int32_t* input2 = reinterpret_cast(input); + int32_t* output2 = reinterpret_cast(output); + TransposeRemovePadding<<>>( + output2, input2, token_offset, batch_size, seq_len, number_heads, head_size); + } + + return CUDA_CALL(cudaGetLastError()); +} + +template +Status FusedScaledDotProductAttention( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + PackedAttentionParameters& parameters, + PackedAttentionData& data) { + const int batch_size = parameters.batch_size; + const int sequence_length = parameters.sequence_length; + const int num_heads = parameters.num_heads; + const int qk_head_size = parameters.head_size; + const int v_head_size = parameters.v_head_size; + void* fused_runner = data.fused_runner; + assert(nullptr != fused_runner); + + LaunchAddBiasTranspose(data.gemm_buffer, data.bias, data.workspace, + batch_size, sequence_length, + num_heads, qk_head_size, v_head_size, + AttentionQkvFormat::QKV_BSN3H, data.token_offset, + parameters.token_count, stream); + + FusedMHARunnerFP16v2* fused_fp16_runner = reinterpret_cast(fused_runner); + const int S = fused_fp16_runner->getSFromMaxSeqLen(sequence_length); + fused_fp16_runner->setup(S, batch_size); + + fused_fp16_runner->run(data.workspace, data.cumulative_sequence_length, data.output, stream); + return Status::OK(); +} + +template +Status UnfusedScaledDotProductAttention( + const cudaDeviceProp& device_prop, + cublasHandle_t& cublas, + cudaStream_t stream, + PackedAttentionParameters& parameters, + PackedAttentionData& data) { + constexpr size_t element_size = sizeof(T); + const int batch_size = parameters.batch_size; + const int sequence_length = parameters.sequence_length; + const int num_heads = parameters.num_heads; + const int qk_head_size = parameters.head_size; + const int v_head_size = parameters.v_head_size; + + const int batches = batch_size * num_heads; + const int size_per_batch_q = sequence_length * qk_head_size; + const int size_per_batch_k = sequence_length * qk_head_size; + const int size_per_batch_v = sequence_length * v_head_size; + const size_t elements_q = static_cast(batches) * static_cast(size_per_batch_q); + const size_t elements_k = static_cast(batches) * static_cast(size_per_batch_k); + const size_t elements_v = static_cast(batches) * static_cast(size_per_batch_v); + + // Q, K and V pointers when fused attention is not used + T* qkv = data.workspace; + T* q = qkv; + T* k = q + elements_q; + T* v = k + elements_k; + + LaunchAddBiasTranspose(data.gemm_buffer, data.bias, data.workspace, + batch_size, sequence_length, + num_heads, qk_head_size, v_head_size, + AttentionQkvFormat::Q_K_V_BNSH, data.token_offset, + parameters.token_count, stream); + + T* scaled_qk = qkv + elements_q + elements_k + elements_v; + + // Q, K and V are ready now + DUMP_TENSOR_INIT(); + + DUMP_TENSOR_D("gemm_buffer", data.gemm_buffer, parameters.token_count, (num_heads * (qk_head_size * 2 + v_head_size))); + DUMP_TENSOR_D("data.workspace", data.workspace, 3 * batch_size, num_heads, sequence_length, qk_head_size); + + // Compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scaled_qk: BxNxSxT + // Q: BxNxSxH, K: BxNxSxH, Q*K': BxNxSxS + float one = 1.0f; + float zero = 0.f; + float scale = parameters.scale == 0.0f ? 1.f / sqrt(static_cast(qk_head_size)) + : parameters.scale; + + cublasSetStream(cublas, stream); + + CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper( + cublas, CUBLAS_OP_T, CUBLAS_OP_N, + sequence_length, sequence_length, qk_head_size, + &scale, + k, qk_head_size, sequence_length * qk_head_size, + q, qk_head_size, sequence_length * qk_head_size, + &zero, + scaled_qk, sequence_length, sequence_length * sequence_length, + batches, device_prop)); + + DUMP_TENSOR_D("QK", scaled_qk, batch_size * num_heads, sequence_length, sequence_length); + + const size_t bytes = GetAttentionScratchSize(element_size, batch_size, num_heads, + sequence_length); + T* attention_score = scaled_qk + (bytes / element_size); + + // Apply softmax and store result R to attention_score: BxNxSxS + ORT_RETURN_IF_ERROR(ComputeSoftmaxWithCumSeqLength( + scaled_qk, + data.relative_position_bias, + parameters.broadcast_res_pos_bias, + data.cumulative_sequence_length, + batch_size, + sequence_length, + num_heads, + attention_score, stream)); + + DUMP_TENSOR_D("Softmax", attention_score, batch_size * num_heads, sequence_length, sequence_length); + + // compute R*V (as V*R), and store in temp_output (space used by Q): BxNxSxH_v + T* temp_output = qkv; + CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper( + cublas, CUBLAS_OP_N, CUBLAS_OP_N, + v_head_size, sequence_length, sequence_length, + &one, v, v_head_size, sequence_length * v_head_size, + attention_score, sequence_length, sequence_length * sequence_length, + &zero, temp_output, v_head_size, sequence_length * v_head_size, batches, device_prop)); + + // Temp_output is BxNxSxH_v, transpose and remove padding to output token_countxNxH_v + Status result = LaunchTransposeRemovePadding( + data.output, temp_output, + data.token_offset, parameters.token_count, + batch_size, sequence_length, num_heads, v_head_size, + stream); + + DUMP_TENSOR("unfused output", data.output, parameters.token_count, num_heads, v_head_size); + return result; +} + +template +Status QkvToContext( + const cudaDeviceProp& device_prop, + cublasHandle_t& cublas, + cudaStream_t stream, + PackedAttentionParameters& parameters, + PackedAttentionData& data) { + void* fused_runner = data.fused_runner; + if (nullptr != fused_runner) { + return FusedScaledDotProductAttention(device_prop, stream, parameters, data); + } else { + return UnfusedScaledDotProductAttention(device_prop, cublas, stream, parameters, data); + } +} + +template Status QkvToContext( + const cudaDeviceProp& device_prop, + cublasHandle_t& cublas, + cudaStream_t stream, + PackedAttentionParameters& parameters, + PackedAttentionData& data); + +template Status QkvToContext( + const cudaDeviceProp& device_prop, + cublasHandle_t& cublas, + cudaStream_t stream, + PackedAttentionParameters& parameters, + PackedAttentionData& data); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h new file mode 100644 index 0000000000..1eb26575df --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/providers/cuda/shared_inc/cuda_utils.h" +#include +#include +#include "contrib_ops/cpu/bert/attention_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +size_t GetAttentionScratchSize( + size_t element_size, + size_t batch_size, + size_t num_heads, + size_t sequence_length); + +size_t GetAttentionWorkspaceSize( + size_t element_size, + size_t batch_size, + size_t num_heads, + size_t qk_head_size, + size_t v_head_size, + size_t sequence_length, + void* fused_runner); + +template +struct PackedAttentionData { + T* gemm_buffer; + const T* bias; + const T* relative_position_bias; + const int32_t* token_offset; + const int32_t* cumulative_sequence_length; + + T* workspace; + T* output; + + void* fused_runner; +}; + +template +Status QkvToContext( + const cudaDeviceProp& device_prop, + cublasHandle_t& cublas, + cudaStream_t stream, + contrib::PackedAttentionParameters& parameters, + PackedAttentionData& data); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index b6fd89ea08..8ec0d78890 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -63,6 +63,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, Affine); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Attention); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Attention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, PackedAttention); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, PackedAttention); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BeamSearch); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, Crop); @@ -201,6 +203,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index b48aab3813..ae9e0c1324 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -315,6 +315,133 @@ ONNX_MS_OPERATOR_SET_SCHEMA( AttentionTypeAndShapeInference(ctx, past_input_index); })); +constexpr const char* PackingAttention_ver1_doc = R"DOC( +This is the packed version of Attention. + +Sequences in one batch usually don't have same length and they are padded to have same length, +e.g., below is a batch with 3 sequences and tokens* are padded. + Sequence_0: 0, 1*, 2*, 3* + Sequence_1: 4, 5, 6*, 7* + Sequence_2: 8, 9, 10, 11 + +PackedAttention is designed to takes in packed input, i.e., only the real tokens without padding. +An input as above will be packed into 3 tensors like below: + - input ([h0, h4, h5, h8, h9, h10, h11]) + - token_offset: 0, 4, 5, 8, 9, 10, 11, 1*, 2*, 3*, 6*, 7* + - cumulated_token_count: 0, 1, 1+2, 1+2+4 + +Input tensors contains the hidden embedding of real tokens. +Token_offset records the offset of token in the unpacked input. +cumulated_token_count records cumulated length of each sequnces length. + +The operator only supports BERT like model with padding on right now. + +)DOC"; + +// Shape inference for PackedAttention. Here are the shapes of inputs and output: +// Input 'input': (token_count, input_hidden_size) +// Input 'weights': (input_hidden_size, hidden_size + hidden_size + v_hidden_size) +// Input 'bias': (hidden_size + hidden_size + v_hidden_size) +// Input 'token_offset': (batch_size, sequence_length) +// Input 'cumulative_sequence_length': (batch_size + 1) +// Input 'relative_position_bias': (batch_size, num_heads, sequence_length, sequence_length) +// Output 'output': (token_count, v_hidden_size) +void PackedAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { + // Type inference + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); + + // Shape inference + if (hasInputShape(ctx, 0) && hasInputShape(ctx, 2)) { + auto& input_shape = getInputShape(ctx, 0); + auto& input_dims = input_shape.dim(); + int input_dim_size = input_dims.size(); + if (input_dim_size != 2) { + fail_shape_inference("Inputs 0 shall be 2 dimensions"); + } + + auto& bias_shape = getInputShape(ctx, 2); + auto& bias_dims = bias_shape.dim(); + if (bias_dims.size() != 1) { + fail_shape_inference("Invalid bias shape"); + } + + int64_t v_hidden_size = -1; + std::vector qkv_hidden_sizes; + getRepeatedAttribute(ctx, "qkv_hidden_sizes", qkv_hidden_sizes); + + if (qkv_hidden_sizes.size() != 0) { + if (qkv_hidden_sizes.size() != 3) { + fail_shape_inference("qkv_hidden_sizes should have 3 elements") + } + v_hidden_size = qkv_hidden_sizes[2]; + } else { + v_hidden_size = bias_shape.dim(0).dim_value() / 3; + } + + ONNX_NAMESPACE::TensorShapeProto output_shape; + for (auto& dim : input_dims) { + *output_shape.add_dim() = dim; + } + + output_shape.mutable_dim(input_dim_size - 1)->set_dim_value(v_hidden_size); + updateOutputShape(ctx, 0, output_shape); + } +} + +ONNX_MS_OPERATOR_SET_SCHEMA( + PackedAttention, 1, + OpSchema() + .SetDoc(PackingAttention_ver1_doc) + .Attr("num_heads", "Number of attention heads", AttributeProto::INT) + .Attr("qkv_hidden_sizes", + "Hidden dimension of Q, K, V: hidden_size, hidden_size and v_hidden_size", + AttributeProto::INTS, + OPTIONAL_VALUE) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1/sqrt(head_size)", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Input(0, + "input", + "Input tensor with shape (token_count, input_hidden_size)", + "T") + .Input(1, + "weights", + "Merged Q/K/V weights with shape (input_hidden_size, hidden_size + hidden_size + v_hidden_size)", + "T") + .Input(2, + "bias", + "Bias tensor with shape (hidden_size + hidden_size + v_hidden_size) for input projection", + "T") + .Input(3, + "token_offset", + "In packing mode, it specifies the offset of each token(batch_size, sequence_length).", + "M") + .Input(4, + "cumulative_sequence_length", + "A tensor with shape (batch_size + 1). It specifies the cumulative sequence length.", + "M") + .Input(5, + "relative_position_bias", + "A tensor with shape (batch_size, num_heads, sequence_length, sequence_length)" + "or (1, num_heads, sequence_length, sequence_length)." + "It specifies the additional bias to QxK'", + "T", + OpSchema::Optional) + .Output(0, + "output", + "2D output tensor with shape (token_count, v_hidden_size)", + "T") + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)"}, + "Constrain input and output types to float tensors.") + .TypeConstraint("M", + {"tensor(int32)"}, + "Constrain mask index to integer types") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + PackedAttentionTypeAndShapeInference(ctx); + })); + constexpr const char* DecoderMaskedMultiheadAttention_ver1_doc = R"DOC( Uni-directional attention that supports input sequence length of 1. diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 4c70e07695..9b1dd82a00 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -81,6 +81,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MurmurHash3); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramRepeatBlock); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Pad); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PackedAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RelativePositionBias); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRelativePositionBias); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RemovePadding); @@ -172,6 +173,7 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc new file mode 100644 index 0000000000..e0e50e8b50 --- /dev/null +++ b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc @@ -0,0 +1,502 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/platform/env_var_utils.h" +#include "gtest/gtest.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/common/cuda_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/scoped_env_vars.h" +#include "contrib_ops/cpu/bert/attention_common.h" +#include "test/contrib_ops/attention_op_test_helper.h" + +namespace onnxruntime { +using contrib::AttentionMaskType; +namespace test { + +static void RunPackedAttentionTest( + const std::vector& input_data, // input: [token_count, hidden_size] + const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] + const std::vector& bias_data, // bias: [3 * hidden_size] + const std::vector& token_offset, // token_offset: [batch_size, sequence_length] + const std::vector& cumulative_sequence_length, // cum_seq_len: [batch_size + 1] + const std::vector& output_data, // output: [token_count, hidden_size] + int batch_size, + int sequence_length, + int hidden_size, + int number_of_heads, + int token_count, + bool use_float16, + bool use_scale, + std::vector qkv_sizes, + const std::vector& relative_position_bias_data) { + int min_cuda_architecture = use_float16 ? 530 : 0; + bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + + if (enable_cuda) { + OpTester tester("PackedAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(number_of_heads)); + + int32_t qkv_hidden_size_sum; + int32_t v_hidden_size; + int32_t head_size_k; + if (qkv_sizes.size() != 0) { + qkv_hidden_size_sum = qkv_sizes[0] + qkv_sizes[1] + qkv_sizes[2]; + std::vector sizes_attribute{qkv_sizes[0], qkv_sizes[1], qkv_sizes[2]}; + tester.AddAttribute>("qkv_hidden_sizes", sizes_attribute); + v_hidden_size = qkv_sizes[2]; + head_size_k = qkv_sizes[1] / number_of_heads; + } else { + qkv_hidden_size_sum = 3 * hidden_size; + v_hidden_size = hidden_size; + head_size_k = hidden_size / number_of_heads; + } + + if (use_scale) { + tester.AddAttribute("scale", static_cast(1.f / sqrt(head_size_k))); + } + + std::vector input_dims = {token_count, hidden_size}; + std::vector weights_dims = {hidden_size, qkv_hidden_size_sum}; + std::vector bias_dims = {qkv_hidden_size_sum}; + std::vector token_offset_dims = {batch_size, sequence_length}; + std::vector cum_seq_len_dims = {batch_size + 1}; + std::vector relative_position_bias_data_dims = {batch_size, number_of_heads, sequence_length, sequence_length}; + std::vector output_dims = {token_count, v_hidden_size}; + if (use_float16) { + tester.AddInput("input", input_dims, ToFloat16(input_data)); + tester.AddInput("weight", weights_dims, ToFloat16(weights_data)); + tester.AddInput("bias", bias_dims, ToFloat16(bias_data)); + tester.AddInput("token_offset", token_offset_dims, token_offset); + tester.AddInput("cumulative_sequence_length", cum_seq_len_dims, cumulative_sequence_length); + if (relative_position_bias_data.size() > 0) { + tester.AddInput("relative_position_bias", relative_position_bias_data_dims, ToFloat16(relative_position_bias_data)); + } + + tester.AddOutput("output", output_dims, ToFloat16(output_data)); + } else { + tester.AddInput("input", input_dims, input_data); + tester.AddInput("weight", weights_dims, weights_data); + tester.AddInput("bias", bias_dims, bias_data); + tester.AddInput("token_offset", token_offset_dims, token_offset); + tester.AddInput("cumulative_sequence_length", cum_seq_len_dims, cumulative_sequence_length); + if (relative_position_bias_data.size() > 0) { + tester.AddInput("relative_position_bias", relative_position_bias_data_dims, relative_position_bias_data); + } + + tester.AddOutput("output", output_dims, output_data); + } + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +static void RunPackedAttentionTest( + const std::vector& input_data, // input: [token_count, hidden_size] + const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] + const std::vector& bias_data, // bias: [3 * hidden_size] + const std::vector& token_offset, // token_offset: [batch_size, sequence_length] + const std::vector& cumulative_sequence_length, // cum_seq_len: [batch_size + 1] + const std::vector& output_data, // output: [token_count, hidden_size] + int batch_size, + int sequence_length, + int hidden_size, + int number_of_heads, + int token_count, + std::vector qkv_sizes = {}, + const std::vector& relative_position_bias_data = {}) { +#define InvokePackedAttentionTest(use_float16, use_scale) \ + RunPackedAttentionTest( \ + input_data, \ + weights_data, \ + bias_data, \ + token_offset, \ + cumulative_sequence_length, \ + output_data, \ + batch_size, \ + sequence_length, \ + hidden_size, \ + number_of_heads, \ + token_count, \ + use_float16, \ + use_scale, \ + qkv_sizes, \ + relative_position_bias_data); + + InvokePackedAttentionTest(true, true); + InvokePackedAttentionTest(true, false); + InvokePackedAttentionTest(false, true); + InvokePackedAttentionTest(false, false); +} + +TEST(PackedAttentionTest, NoPack) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + std::vector token_offset{0, 1}; + std::vector cum_seq_len{0, 2}; + + std::vector output_data = { + 3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f, + 3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f}; + + RunPackedAttentionTest( + input_data, + weight_data, + bias_data, + token_offset, + cum_seq_len, + output_data, + batch_size, + sequence_length, + hidden_size, + number_of_heads, + batch_size * sequence_length); +} + +TEST(PackedAttentionTest, NoPackWithRelativePositionBias) { + int batch_size = 2; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f, + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, + 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + std::vector token_offset{0, 1, 2, 3}; + std::vector cum_seq_len{0, 2, 4}; + + std::vector relative_position_bias = { + 0.2f, -0.1f, 0.4f, 2.5f, 1.6f, -1.1f, 0.4f, -2.5f, + 0.2f, -0.1f, 0.4f, 2.5f, 1.6f, -1.1f, 0.4f, -2.5f}; + + std::vector output_data = { + 4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f, + -1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f, + 4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f, + -1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f}; + + RunPackedAttentionTest( + input_data, + weight_data, + bias_data, + token_offset, + cum_seq_len, + output_data, + batch_size, + sequence_length, + hidden_size, + number_of_heads, + batch_size * sequence_length, + {}, + relative_position_bias); +} + +TEST(PackedAttentionTest, PackedWithRelativePositionBias) { + int batch_size = 2; + int sequence_length = 4; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, // b0:s0 + 0.5f, 0.2f, 0.3f, -0.6f, // b0:s1 + 0.8f, -0.5f, 0.0f, 1.f, // b1:s0 + 0.5f, 0.2f, 0.3f, -0.6f // b1:s1 + }; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, + 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + std::vector token_offset{0, 1, 4, 5, 2, 3, 6, 7}; + std::vector cum_seq_len{0, 2, 4}; + + std::vector relative_position_bias = { + 0.2f, -0.1f, 0.f, 0.f, 0.4f, 2.5f, 0.f, 0.f, + 1.6f, -1.1f, 0.f, 0.f, 0.4f, -2.5f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + + 0.2f, -0.1f, 0.f, 0.f, 0.4f, 2.5f, 0.f, 0.f, + 1.6f, -1.1f, 0.f, 0.f, 0.4f, -2.5f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, + 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + + std::vector output_data = { + 4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f, + -1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f, + 4.066014289855957f, 0.068997815251350403f, 4.25f, 5.6499996185302734f, + -1.8799558877944946f, 0.32488855719566345f, 4.25f, 5.6499996185302734f}; + + RunPackedAttentionTest( + input_data, + weight_data, + bias_data, + token_offset, + cum_seq_len, + output_data, + batch_size, + sequence_length, + hidden_size, + number_of_heads, + 4, + {}, + relative_position_bias); +} + +TEST(PackedAttentionTest, PackedBatch) { + int batch_size = 2; + int sequence_length = 4; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, // b0:s0 + 0.5f, 0.2f, 0.3f, -0.6f, // b0:s1 + 0.8f, -0.5f, 0.0f, 1.f, // b1:s0 + 0.5f, 0.2f, 0.3f, -0.6f // b1:s1 + }; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + std::vector token_offset{0, 1, 4, 5, 2, 3, 6, 7}; + std::vector cum_seq_len{0, 2, 4}; + + std::vector output_data = { + 3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f, + 3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f, + 3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f, + 3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f}; + + RunPackedAttentionTest( + input_data, + weight_data, + bias_data, + token_offset, + cum_seq_len, + output_data, + batch_size, + sequence_length, + hidden_size, + number_of_heads, + 4); +} + +TEST(PackedAttentionTest, PackedBatchWithQKV) { + int batch_size = 2; + int sequence_length = 4; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f, + + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector qkv_sizes = {6, 6, 4}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f, + + 0.3f, 0.2f, 4.0f, 2.2f, 2.4f, 3.3f, 2.1f, 4.2f, 0.5f, 0.1f, 0.4f, 1.6f, + 0.4f, 0.8f, 0.9f, 0.1f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, + 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f, + 0.5f, 0.7f, 0.2f, 1.2f}; + + std::vector token_offset{0, 1, 4, 5, 2, 3, 6, 7}; + std::vector cum_seq_len{0, 2, 4}; + + std::vector output_data = { + 3.1967618465423584f, 0.51903456449508667f, 0.63051539659500122f, 2.9394614696502686f, + 0.65332180261611938f, 1.000949501991272f, 0.74175024032592773f, 2.8231701850891113f, + + 3.1967618465423584f, 0.51903456449508667f, 0.63051539659500122f, 2.9394614696502686f, + 0.65332180261611938f, 1.000949501991272f, 0.74175024032592773f, 2.8231701850891113f}; + + RunPackedAttentionTest( + input_data, + weight_data, + bias_data, + token_offset, + cum_seq_len, + output_data, + batch_size, + sequence_length, + hidden_size, + number_of_heads, + 4, + qkv_sizes); +} + +static void RunModelWithRandomInput( + int64_t batch_size, + int64_t sequence_length, + std::string& onnx_model, + bool is_float16) { + // ORT enables TF32 in GEMM for A100. TF32 will cause precsion loss and fail this test. + // Do not run this test unless TF32 is disabled explicitly. + if (HasCudaEnvironment(800) && ParseEnvironmentVariableWithDefault("NVIDIA_TF32_OVERRIDE", 1) != 0) { + GTEST_SKIP() << "Skipping RunModelWithRandomInput in A100 since TF32 is enabled"; + return; + } + + RandomValueGenerator random{234}; + + constexpr int hidden_size = 768; + constexpr int num_heads = 12; + + int token_count = 0; + std::vector cum_seq_len(batch_size + 1); + cum_seq_len[0] = 0; + + int original_offset = 0; + int token_offset_idx = 0; + std::vector token_offset(batch_size * sequence_length); + for (int b = 0; b < batch_size; b++) { + int actual_seq_len = (sequence_length / (b + 1)); + token_count += actual_seq_len; + cum_seq_len[b + 1] = token_count; + + original_offset = b * sequence_length; + for (int s = 0; s < actual_seq_len; s++) { + token_offset[token_offset_idx++] = original_offset++; + } + } + + for (int b = 0; b < batch_size; b++) { + int actual_seq_len = (sequence_length / (b + 1)); + original_offset = b * sequence_length + actual_seq_len; + for (int s = actual_seq_len; s < sequence_length; s++) { + token_offset[token_offset_idx++] = original_offset++; + } + } + + assert(token_offset_idx == batch_size * sequence_length); + + std::vector input_dims{token_count, hidden_size}; + std::vector input_data = random.Gaussian(input_dims, 0.0f, 0.3f); + + std::vector weight_dims{hidden_size, 3 * hidden_size}; + std::vector weight_data = random.Gaussian(weight_dims, 0.0f, 0.3f); + + std::vector bias_dims{3 * hidden_size}; + std::vector bias_data = random.Gaussian(bias_dims, 0.0f, 0.1f); + + std::vector token_offset_dims{batch_size, sequence_length}; + std::vector cum_seq_len_dims{batch_size + 1}; + + float gpu_threshold = is_float16 ? 0.1f : 0.005f; + bool enable_cuda = HasCudaEnvironment(is_float16 ? 530 : 0); + if (enable_cuda) { + OpTester test("PackedAttention", 1, onnxruntime::kMSDomain); + test.AddAttribute("num_heads", num_heads); + if (is_float16) { + test.AddInput("input", input_dims, ToFloat16(input_data)); + test.AddInput("weight", weight_dims, ToFloat16(weight_data)); + test.AddInput("bias", bias_dims, ToFloat16(bias_data)); + } else { + test.AddInput("input", input_dims, input_data); + test.AddInput("weight", weight_dims, weight_data); + test.AddInput("bias", bias_dims, bias_data); + } + test.AddInput("token_offset", token_offset_dims, token_offset); + test.AddInput("cumulative_sequence_length", cum_seq_len_dims, cum_seq_len); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.AddReferenceOutputs(onnx_model, gpu_threshold, DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +TEST(PackedAttentionTest, test_on_random_data) { + std::string onnx_model = "testdata/packed_attention_fp32.onnx"; + std::string onnx_model_fp16 = "testdata/packed_attention_fp16.onnx"; + for (int batch_size : std::vector({1, 2, 3, 4, 5, 6, 7, 8})) { + for (int sequence_length : std::vector({32, 48, 64, 95, 128})) { + RunModelWithRandomInput( + batch_size, + sequence_length, + onnx_model, + false); + RunModelWithRandomInput( + batch_size, + sequence_length, + onnx_model_fp16, + true); + } + } +} + +TEST(PackedAttentionTest, test_on_random_data_large_seq) { + int batch_size = 2; + int sequence_length = 1152; // > 1024 + std::string onnx_model = "testdata/packed_attention_fp32.onnx"; + std::string onnx_model_fp16 = "testdata/packed_attention_fp16.onnx"; + RunModelWithRandomInput( + batch_size, + sequence_length, + onnx_model, + false); + RunModelWithRandomInput( + batch_size, + sequence_length, + onnx_model_fp16, + true); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index 57b15f1cb4..6396355642 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -413,9 +413,9 @@ struct TensorCheck { if (abs_error <= abs_threshold) { // if the absolute error is small enough, then no need to calculate realative error EXPECT_NEAR(0, abs_error, abs_threshold) << "provider_type: " - << provider_type; + << provider_type; } else { - //default for existing tests. + // default for existing tests. const float rel_error = abs_error / max_value; EXPECT_NEAR(0, rel_error, threshold) << "provider_type: " << provider_type; @@ -1406,7 +1406,7 @@ void OpTester::ExecuteModelForEps( } }; -void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_error) { +void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_error, std::unique_ptr ep) { SessionOptions so; so.session_logid = op_; so.session_log_verbosity_level = 1; @@ -1418,6 +1418,7 @@ void OpTester::AddReferenceOutputs(const std::string& model_path, float abs_erro Status status; InferenceSession subgraph_session_object{so, GetEnvironment()}; + status = subgraph_session_object.RegisterExecutionProvider(std::move(ep)); ASSERT_TRUE((status = subgraph_session_object.Load(model_path)).IsOK()) << status; ASSERT_TRUE((status = subgraph_session_object.Initialize()).IsOK()) << status; diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index a06efb5f2a..1647190291 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -687,7 +687,7 @@ class OpTester { } // Generate the reference outputs with the model file - void AddReferenceOutputs(const std::string& model_path, float abs_error = 0.0f); + void AddReferenceOutputs(const std::string& model_path, float abs_error = 0.0f, std::unique_ptr ep = nullptr); void AddCustomOpRegistry(std::shared_ptr registry) { custom_schema_registries_.push_back(registry->GetOpschemaRegistry()); diff --git a/onnxruntime/test/testdata/packed_attention_fp16.onnx b/onnxruntime/test/testdata/packed_attention_fp16.onnx new file mode 100644 index 0000000000000000000000000000000000000000..a9d0a297952030c6c32acb26f6267b61844ac1a2 GIT binary patch literal 1228 zcma)5!EVz)5XEhiI1{%`wkoua0AV47A|=Ei4)lOT;Dm&P(pwI!Huff7YVQW`xb6@uQ7vJmf z8i6OcRD>(sw>R#N`8LmYY^Yq5Frl3z(`eS-n|JcFFb+z1*MJaQjQ_mUJ@lCQS6#xNxvkNdlZM1PgMXL&Olir zs%ZCOWjAEVbp0xDDwk;hA8yI%LP>!de$<5v8c!KcUsEx=Oq8W}WIg^>xr`LY*fO4k+>EnYA$X zI&EmJ9ynp4qA^}=;OTOCvs@U)wzeMz{MlXg<*Z?u$0)+gqqrQFX(<|s!QiH5vx%*pAdQsjtS@i9HK_1z!Um6O~Zis zzMh?A-{|Oh;VXIyPjtPrRRAq?yit$m_{Hc^HUo4H_cY?(<8As7?o z1*_&zN44C9(W_lcjM{Ln;iqtdTEF3V7A6wA;~o^w-|ESOy{0v2rT^E);x&z5yf~4% literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/packed_attention_fp32.onnx b/onnxruntime/test/testdata/packed_attention_fp32.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9050a6e2a02e76306e63d77b7ccfb998c0e2bfe5 GIT binary patch literal 1228 zcma)5!EVz)5XEhiI1{%`wkoua0AV47A|=Ei4)lOT;Dm&P(pwI!Huff7YVQW`xb6@uQ7vJmf z8i6OcRD>(sw>R#N`8LmYY^Yq5Frl3z(`eS-n|JcFFb+z1*MJaQjQ_mUJ@lCQS6#xNxvkNdlZM1PgMXL&Olir zs%ZCOWjAEVbp0xDDwk;hA8yI%LP>!de$<5v8c!KcUsEx=Oq8W}WIg^>xr`LY*fO4k+>EnYA$X zI&EmJ9ynp4qA^}=;OTOCvs@U)wzeMz{MlXg<*Z?u$0)+gqqrQFX(<|s!QiH5vx%*pAdQsj%lC=aEKb20#E4QGz|mh z`+9bgeWRo2g|Fx-Jkj;eRspoo@kTwK;}@ex*$mJ*+|!7Yzm+`OT9Aiu->hnYg1rPhM&R-YW;@eS(r%dj(bozf2$`C_L|n9mHuBFi`O)M0iW5C AmH+?% literal 0 HcmV?d00001