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

Misc minor bug fixes in transformer kernels (#13780)

Browse source
This commit is contained in:
Hariharan Seshadri 2022-12-04 21:30:57 -08:00 committed by GitHub
parent f34ebbc8ff
commit 5f4e0c95ec
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
View file

@ -191,7 +191,7 @@ Status BeamSearchBase<T>::CheckInputs(const OpKernelContextInternal& context) {
context.Input<Tensor>(0), // input_ids
context.Input<Tensor>(7), // vocab_mask
context.Input<Tensor>(8), // prefix_vocab_mask
context.Input<Tensor>(10))); // attention_mask
context.Input<Tensor>(9))); // attention_mask
return Status::OK();
}

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

@ -122,14 +122,14 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
cublasHandle_t cublas = CublasHandle();
constexpr size_t element_size = sizeof(T);
IAllocatorUniquePtr<T> gemm_buffer;
IAllocatorUniquePtr<void> gemm_buffer;
if (weights != nullptr) {
// Use GEMM for fully connection.
int m = batch_size * sequence_length;
int n = (parameters.hidden_size + parameters.hidden_size + parameters.v_hidden_size);
int k = parameters.input_hidden_size;
size_t gemm_buffer_size = static_cast<size_t>(batch_size) * sequence_length * n * element_size;
gemm_buffer = GetScratchBuffer<T>(gemm_buffer_size);
gemm_buffer = GetScratchBuffer<void>(gemm_buffer_size);
typedef typename ToCudaType<T>::MappedType CudaT;
CudaT one = ToCudaType<T>::FromFloat(1.0f);