From 534c22d7694d625158f0ebc7a9929529239ca8e9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Jul 2021 11:04:56 -0700 Subject: [PATCH] use float for alpha in attention Gemm (#8477) --- .../contrib_ops/cuda/bert/attention_impl.cu | 9 ++-- .../providers/cuda/shared_inc/fpgeneric.h | 49 +++++++++++++++++++ 2 files changed, 53 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index d232e70a30..106bedcbb8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu @@ -104,19 +104,18 @@ bool QkvToContext( v = present + batches * present_size_per_batch; } - // Raw attention mask could be 2D (BxS) or 3D (BxSxS*) + // Raw attention mask could be 2D (BxS) or 3D (BxSxS*) or 4D(Bx1xMxM), where M is the max sequence length. bool use_raw_attention_mask = (nullptr != mask_index && nullptr != mask_index_dims && mask_index_dims->size() >= 2); // compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scratch1: BxNxSxS* // Q: BxNxSxH, K (present_k): BxNxS*xH, Q*K': BxNxSxS* const float rsqrt_head_size = 1.f / sqrt(static_cast(head_size)); const int temp_matrix_size = sequence_length * all_sequence_length; - T one = (T)(1.0f); - T zero = (T)(0.f); + float one = 1.0f; + float zero = 0.f; // For raw attention mask, the scalar if 1/sqrt(H) is moved to softmax computation. - // TODO: move scalar to softmax computation since converting 1/Sqrt(H) to half might have loss in precision. - T alpha = use_raw_attention_mask ? one : (T)(rsqrt_head_size); + float alpha = use_raw_attention_mask ? one : rsqrt_head_size; if (!CUBLAS_CALL(cublasGemmStridedBatchedHelper( cublas, CUBLAS_OP_T, CUBLAS_OP_N, all_sequence_length, sequence_length, head_size, &alpha, k, head_size, present_size_per_batch, diff --git a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h index 6752316ef5..2e6f96ded4 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h @@ -379,6 +379,55 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, } } + +inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const float* alpha, + const __half* A, int lda, + long long int strideA, + const __half* B, int ldb, + long long int strideB, + const float* beta, + __half* C, int ldc, + long long int strideC, + int batch_count, + const cudaDeviceProp& prop) { + const HalfGemmOptions* half_options = HalfGemmOptions::GetInstance(); + onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, half_options->GetMathMode()); + if (half_options->IsCompute16F()) { + // The alpha and beta shall have same precision as compute type. + uint16_t h_a = onnxruntime::math::floatToHalf(*alpha); + uint16_t h_b = onnxruntime::math::floatToHalf(*beta); + return cublasGemmStridedBatchedEx(handle, + transa, + transb, + m, n, k, + &h_a, + A, CUDA_R_16F, lda, strideA, + B, CUDA_R_16F, ldb, strideB, + &h_b, + C, CUDA_R_16F, ldc, strideC, + batch_count, + half_options->GetComputeType(), + CUBLAS_GEMM_DEFAULT); + } else { + return cublasGemmStridedBatchedEx(handle, + transa, + transb, + m, n, k, + alpha, + A, CUDA_R_16F, lda, strideA, + B, CUDA_R_16F, ldb, strideB, + beta, + C, CUDA_R_16F, ldc, strideC, + batch_count, + half_options->GetComputeType(), + CUBLAS_GEMM_DEFAULT); + } +} + #if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle, cublasOperation_t transa,