diff --git a/onnxruntime/contrib_ops/cuda/bert/attention.cc b/onnxruntime/contrib_ops/cuda/bert/attention.cc index 2be250fc0d..a190687f8c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/attention.cc @@ -83,20 +83,12 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { CudaT one = ToCudaType::FromFloat(1.0f); CudaT zero = ToCudaType::FromFloat(0.0f); - // Bias shape is (N), broadcast using B(N, M) = 1 * bias(N, 1) x ones(1, M) + 0 * B. - // TODO: use custom kernel of expand to improve the performance. - CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( - cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, 1, &one, - reinterpret_cast(bias->Data()), n, - GetConstOnes(m), 1, - &zero, reinterpret_cast(gemm_buffer.get()), n, device_prop)); - // Gemm, note that CUDA assumes col-major, so result(N, M) = 1 * weights x input + 1 x B. 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, - &one, reinterpret_cast(gemm_buffer.get()), n, device_prop)); + &zero, reinterpret_cast(gemm_buffer.get()), n, device_prop)); size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, batch_size, @@ -118,6 +110,7 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { past_sequence_length, is_unidirectional_, reinterpret_cast(gemm_buffer.get()), + bias->Data(), nullptr == mask_index ? nullptr : mask_index->Data(), nullptr == mask_index ? gsl::span() : mask_index->Shape().GetDims(), nullptr == past ? nullptr : past->Data(), diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index ff8f4e45ff..a5259456ef 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -28,6 +28,7 @@ limitations under the License. #include "contrib_ops/cuda/bert/attention_impl.h" #include "contrib_ops/cuda/bert/attention_softmax.h" #include "contrib_ops/cuda/bert/transformer_common.h" +#include "contrib_ops/cuda/bert/add_bias_transpose.h" using namespace onnxruntime::cuda; using namespace cub; @@ -81,6 +82,7 @@ bool QkvToContext( const int head_size, const size_t element_size, const T* input, + const T* bias, T* output, T* workspace, const int* mask_index, @@ -101,9 +103,21 @@ bool QkvToContext( const int max_threads_per_block = prop.maxThreadsPerBlock; // input should be BxSx3xNxH => scratch3: 3xBxNxSxH - if (!LaunchTransQkv(stream, 3, sequence_length, batch_size, head_size, num_heads, - max_threads_per_block, false, input, scratch3)) { - return false; + if (bias == nullptr) { + if (!LaunchTransQkv(stream, 3, sequence_length, batch_size, head_size, num_heads, + max_threads_per_block, false, input, scratch3)) { + return false; + } + } else { + const int format = 1; // BxSxMxNxH + const bool enable_half4 = true; + LaunchAddBiasTranspose(stream, 3, format, max_threads_per_block, batch_size, + sequence_length, num_heads, head_size, + input, bias, scratch3, + enable_half4); + if (!CUDA_CALL(cudaPeekAtLastError())) { + return false; + } } // now scratch3 has Q, K, V: each has size BxNxSxH @@ -142,7 +156,7 @@ bool QkvToContext( float one = 1.0f; float zero = 0.f; - // For raw attention mask, the scalar if 1/sqrt(H) is moved to softmax computation. + // For raw attention mask, the scalar 1/sqrt(H) is moved to combine with softmax computation. float alpha = use_raw_attention_mask ? one : rsqrt_head_size; if (!CUBLAS_CALL(cublasGemmStridedBatchedHelper( @@ -208,6 +222,7 @@ bool LaunchAttentionKernel( int past_sequence_length, bool is_unidirectional, const void* input, + const void* bias, const int* mask_index, gsl::span mask_index_dims, const void* past, @@ -222,6 +237,7 @@ bool LaunchAttentionKernel( if (element_size == 2) { return QkvToContext(prop, cublas, stream, batch_size, sequence_length, num_heads, head_size, element_size, reinterpret_cast(input), + reinterpret_cast(bias), reinterpret_cast(output), reinterpret_cast(workspace), mask_index, @@ -235,6 +251,7 @@ bool LaunchAttentionKernel( } else { return QkvToContext(prop, cublas, stream, batch_size, sequence_length, num_heads, head_size, element_size, reinterpret_cast(input), + reinterpret_cast(bias), reinterpret_cast(output), reinterpret_cast(workspace), mask_index, diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h index 09532670ce..2912f501d5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.h @@ -37,6 +37,7 @@ bool LaunchAttentionKernel( int past_sequence_length, // Sequence length in past state bool is_unidirectional, // Whether there is unidirecitonal mask. const void* input, // Input tensor + const void* bias, // Bias tensor const int* mask_index, // Attention mask raw data or index. NULL means no mask. gsl::span mask_index_dims, // Mask index shape const void* past, // Past state input diff --git a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc index 95e2caefc7..a9591e773b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc @@ -259,11 +259,13 @@ Status DecoderAttention::ComputeInternal(OpKernelContext* context) const { has_key_padding_mask_) ); - // calcualte q + // calculate q gemm_query_buffer_p = GetScratchBuffer(batch_size * sequence_length * hidden_size * element_size); m = sequence_length * batch_size; n = hidden_size; k = hidden_size; + + // TODO(tianleiwu): fuse bias and transpose // broadcast bias for query: (h2, S*B) CUBLAS_RETURN_IF_ERROR(cublasGemmHelper( cublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, 1, &one, diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc index dec76be78e..147b5293e3 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc @@ -151,7 +151,9 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { } else { dequant_scale = input_scale * weight_scale; } + // scale back and bias + // TODO(tianleiwu): fuse Dequantize with Add bias and Transpose. ORT_RETURN_IF_ERROR(CudaDequantizeWithBias(Stream(), gemm_buffer_quantized.get(), reinterpret_cast(bias->Data()), @@ -178,6 +180,7 @@ Status QAttention::ComputeInternal(OpKernelContext* context) const { past_sequence_length, is_unidirectional_, reinterpret_cast(gemm_buffer.get()), + nullptr, // bias has been added nullptr == mask_index ? nullptr : mask_index->Data(), nullptr == mask_index ? gsl::span() : mask_index->Shape().GetDims(), nullptr == past_tensor ? nullptr : past_tensor->Data(),