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

replace std::numeric_limits<T> by cub::FpLimits<T> (#10703)

Browse source
This commit is contained in:
Tianlei Wu 2022-02-28 23:11:51 -08:00 committed by GitHub
parent 9a22b5d253
commit c51b500ca7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
onnxruntime/contrib_ops/cuda/transformers/beam_search_impl.cu
View file

@ -4,6 +4,7 @@
#include "beam_search_impl.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/cu_inc/common.cuh"
#include "cub/util_type.cuh"
namespace onnxruntime {
namespace contrib {
@ -109,27 +110,27 @@ __global__ void LogitsProcessKernel(
}
if (found) {
next_token_scores[index] = std::numeric_limits<T>::lowest();
next_token_scores[index] = cub::FpLimits<T>::Lowest();
return;
}
}
// VocabMaskLogitsProcessor
if (vocab_mask != nullptr && vocab_mask[word_id] == 0) {
next_token_scores[index] = std::numeric_limits<T>::lowest();
next_token_scores[index] = cub::FpLimits<T>::Lowest();
return;
}
// PrefixVocabMaskLogitsProcessor
int batch_id = batch_beam_index / num_beams;
if (prefix_vocab_mask != nullptr && prefix_vocab_mask[batch_id * vocab_size + word_id] == 0) {
next_token_scores[index] = std::numeric_limits<T>::lowest();
next_token_scores[index] = cub::FpLimits<T>::Lowest();
return;
}
// MinLengthLogitsProcessor
if (word_id == demote_token_id) {
next_token_scores[index] = std::numeric_limits<T>::lowest();
next_token_scores[index] = cub::FpLimits<T>::Lowest();
}
}
}