saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-10 17:37:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[CUDA] Fuse add bias and transpose into one kernel in Attention (#12670)

Browse source 
* fuse add bias and transpose in attention
This commit is contained in:
Tianlei Wu 2022-08-22 15:46:13 -07:00 committed by GitHub
parent 6246662b1d
commit 8d78f96dfe
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 30 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/contrib_ops/cuda/bert/attention.cc
View file

@ -83,20 +83,12 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
CudaT one = ToCudaType<T>::FromFloat(1.0f);
CudaT zero = ToCudaType<T>::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<const CudaT*>(bias->Data<T>()), n,
GetConstOnes<CudaT>(m), 1,
&zero, reinterpret_cast<CudaT*>(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<const CudaT*>(weights->Data<T>()), n,
reinterpret_cast<const CudaT*>(input->Data<T>()), k,
&one, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop));
&zero, reinterpret_cast<CudaT*>(gemm_buffer.get()), n, device_prop));
size_t workSpaceSize = GetAttentionWorkspaceSize(element_size,
batch_size,
@ -118,6 +110,7 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
past_sequence_length,
is_unidirectional_,
reinterpret_cast<const void*>(gemm_buffer.get()),
bias->Data<T>(),
nullptr == mask_index ? nullptr : mask_index->Data<int>(),
nullptr == mask_index ? gsl::span<const int64_t>() : mask_index->Shape().GetDims(),
nullptr == past ? nullptr : past->Data<T>(),

25
onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
View file

@ -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<const int64_t> 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<const half*>(input),
reinterpret_cast<const half*>(bias),
reinterpret_cast<half*>(output),
reinterpret_cast<half*>(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<const float*>(input),
reinterpret_cast<const float*>(bias),
reinterpret_cast<float*>(output),
reinterpret_cast<float*>(workspace),
mask_index,

1
onnxruntime/contrib_ops/cuda/bert/attention_impl.h
View file

@ -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<const int64_t> mask_index_dims, // Mask index shape
const void* past, // Past state input

4
onnxruntime/contrib_ops/cuda/bert/decoder_attention.cc
View file

@ -259,11 +259,13 @@ Status DecoderAttention<T>::ComputeInternal(OpKernelContext* context) const {
has_key_padding_mask_)
);
// calcualte q
// calculate q
gemm_query_buffer_p = GetScratchBuffer<T>(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,

3
onnxruntime/contrib_ops/cuda/quantization/attention_quantization.cc
View file

@ -151,7 +151,9 @@ Status QAttention<T, int8_t>::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<const CudaT*>(bias->Data<T>()),
@ -178,6 +180,7 @@ Status QAttention<T, int8_t>::ComputeInternal(OpKernelContext* context) const {
past_sequence_length,
is_unidirectional_,
reinterpret_cast<const void*>(gemm_buffer.get()),
nullptr, // bias has been added
nullptr == mask_index ? nullptr : mask_index->Data<int>(),
nullptr == mask_index ? gsl::span<const int64_t>() : mask_index->Shape().GetDims(),
nullptr == past_tensor ? nullptr : past_tensor->Data<T>(),