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 0000000000..a9d0a29795 Binary files /dev/null and b/onnxruntime/test/testdata/packed_attention_fp16.onnx differ diff --git a/onnxruntime/test/testdata/packed_attention_fp32.onnx b/onnxruntime/test/testdata/packed_attention_fp32.onnx new file mode 100644 index 0000000000..9050a6e2a0 Binary files /dev/null and b/onnxruntime/test/testdata/packed_attention_fp32.onnx differ