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

use float for alpha in attention Gemm (#8477)

Browse source
This commit is contained in:
Tianlei Wu 2021-07-26 11:04:56 -07:00 committed by GitHub
parent a9fc3c448c
commit 534c22d769
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 53 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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<float>(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,

49
onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h
View file

@ -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,