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

Memory usage optimization in LongFormer Attention (#11611)

Browse source
This commit is contained in:
Hariharan Seshadri 2022-05-25 10:07:41 -07:00 committed by GitHub
parent 883e4bc341
commit 6e65bac5c2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -107,7 +107,7 @@ Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
int k = hidden_size;
size_t qkv_size = batch_size * sequence_length * 3 * hidden_size * element_size;
auto gemm_buffer = GetScratchBuffer<T>(qkv_size);
auto gemm_buffer = GetScratchBuffer<void>(qkv_size);
typedef typename ToCudaType<T>::MappedType CudaT;
CudaT one = ToCudaType<T>::FromFloat(1.0f);
@ -148,7 +148,7 @@ Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
// Fully connection for global projection.
// Note that Q only need handle global query tokens if we split GEMM to global Q/K/V separately.
// When there is no global token, need not run glboal GEMM.
auto global_gemm_buffer = GetScratchBuffer<T>(max_num_global > 0 ? qkv_size : 0);
auto global_gemm_buffer = GetScratchBuffer<void>(max_num_global > 0 ? qkv_size : 0);
if (max_num_global > 0) {
CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(