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

Ryanunderhill/beamscorer gpu (#16272)

Browse source 
### Description
Make BeamScorer run on the GPU vs the CPU.

Brief overview:
  Adds a CUDA 'CudaBeamSearchScorer' implementation of IBeamScorer
Instead of a 'done' flag per beam, there is one single 'not done'
variable that is copied to the CPU every iteration
Removes some of the extra CPU side buffers and parameters that are no
longer needed

Remaining future optimizations:
CPU copied beam indices is still used in the non
DecoderMaskedSelfAttention case. An extra kernel can be written to avoid
PickGptPasteState needing CPU copied beam indices (called from
UpdateGptFeeds).

### Motivation and Context
It's faster to keep the work on the GPU to avoid GPU->CPU->GPU copies of
data.
This commit is contained in:
Ryan Hill 2023-06-27 15:08:44 -07:00 committed by GitHub
parent f5e9625c36
commit 1001ec93a7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 838 additions and 286 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
onnxruntime/contrib_ops/cpu/transformers/beam_search.cc
View file

@ -215,7 +215,8 @@ Status BeamSearch::Compute(OpKernelContext* ctx) const {
init_beam_state_func_ ? init_beam_state_func_ : GenerationCpuDeviceHelper::InitBeamState<float>,
device_copy_func_ ? device_copy_func_ : GenerationCpuDeviceHelper::DeviceCopy<float>,
device_copy_int32_func_ ? device_copy_int32_func_ : GenerationCpuDeviceHelper::DeviceCopy<int32_t>,
update_gpt_feeds_func_ ? update_gpt_feeds_func_ : GenerationCpuDeviceHelper::UpdateGptFeeds<float>};
update_gpt_feeds_func_ ? update_gpt_feeds_func_ : GenerationCpuDeviceHelper::UpdateGptFeeds<float>,
create_beam_scorer_func_};
#ifdef USE_CUDA
ORT_RETURN_IF_ERROR(impl.InitializeCuda(reorder_past_state_func_, cuda_device_prop_, cuda_device_arch_));
#endif
@ -237,7 +238,8 @@ Status BeamSearch::Compute(OpKernelContext* ctx) const {
init_beam_state_fp16_func_,
device_copy_func_,
device_copy_int32_func_,
update_gpt_feeds_fp16_func_};
update_gpt_feeds_fp16_func_,
create_beam_scorer_func_};
#ifdef USE_CUDA
ORT_RETURN_IF_ERROR(impl.InitializeCuda(reorder_past_state_func_, cuda_device_prop_, cuda_device_arch_));
#endif
@ -267,7 +269,8 @@ Status BeamSearch::Compute(OpKernelContext* ctx) const {
update_decoder_feeds_func_ ? update_decoder_feeds_func_ : GenerationCpuDeviceHelper::UpdateDecoderFeeds<float>,
expand_buffer_int32_func_ ? expand_buffer_int32_func_ : GenerationCpuDeviceHelper::ExpandBuffer<int32_t>,
expand_buffer_float_func_ ? expand_buffer_float_func_ : GenerationCpuDeviceHelper::ExpandBuffer<float>,
expand_buffer_float16_func_ ? expand_buffer_float16_func_ : GenerationCpuDeviceHelper::ExpandBuffer<MLFloat16>};
expand_buffer_float16_func_ ? expand_buffer_float16_func_ : GenerationCpuDeviceHelper::ExpandBuffer<MLFloat16>,
create_beam_scorer_func_};
#ifdef USE_CUDA
ORT_RETURN_IF_ERROR(impl.InitializeCuda(reorder_past_state_func_, init_cache_indir_func_, cuda_device_prop_, cuda_device_arch_));
#endif
@ -288,7 +291,8 @@ Status BeamSearch::Compute(OpKernelContext* ctx) const {
update_decoder_feeds_fp16_func_,
expand_buffer_int32_func_,
expand_buffer_float_func_,
expand_buffer_float16_func_};
expand_buffer_float16_func_,
create_beam_scorer_func_};
#ifdef USE_CUDA
ORT_RETURN_IF_ERROR(impl.InitializeCuda(reorder_past_state_func_, init_cache_indir_func_, cuda_device_prop_, cuda_device_arch_));
#endif
@ -314,7 +318,8 @@ Status BeamSearch::Compute(OpKernelContext* ctx) const {
create_whisper_encoder_inputs_func_ ? create_whisper_encoder_inputs_func_ : GenerationCpuDeviceHelper::CreateWhisperEncoderInputs<float>,
update_decoder_feeds_func_ ? update_decoder_feeds_func_ : GenerationCpuDeviceHelper::UpdateDecoderFeeds<float>,
expand_buffer_float_func_ ? expand_buffer_float_func_ : GenerationCpuDeviceHelper::ExpandBuffer<float>,
expand_buffer_float16_func_ ? expand_buffer_float16_func_ : GenerationCpuDeviceHelper::ExpandBuffer<MLFloat16>};
expand_buffer_float16_func_ ? expand_buffer_float16_func_ : GenerationCpuDeviceHelper::ExpandBuffer<MLFloat16>,
create_beam_scorer_func_};
#ifdef USE_CUDA
ORT_RETURN_IF_ERROR(impl.InitializeCuda(reorder_past_state_func_, init_cache_indir_func_, cuda_device_prop_, cuda_device_arch_));
#endif
@ -334,7 +339,8 @@ Status BeamSearch::Compute(OpKernelContext* ctx) const {
create_whisper_encoder_inputs_func_ ? create_whisper_encoder_inputs_func_ : GenerationCpuDeviceHelper::CreateWhisperEncoderInputs<MLFloat16>,
update_decoder_feeds_fp16_func_ ? update_decoder_feeds_fp16_func_ : GenerationCpuDeviceHelper::UpdateDecoderFeeds<MLFloat16>,
expand_buffer_float_func_,
expand_buffer_float16_func_};
expand_buffer_float16_func_,
create_beam_scorer_func_};
#ifdef USE_CUDA
ORT_RETURN_IF_ERROR(impl.InitializeCuda(reorder_past_state_func_, init_cache_indir_func_, cuda_device_prop_, cuda_device_arch_));
#endif

5
onnxruntime/contrib_ops/cpu/transformers/beam_search.h
View file

@ -53,7 +53,8 @@ class BeamSearch : public IControlFlowKernel {
const GenerationDeviceHelper::ProcessLogitsFunc<float>& process_logits_func,
const GenerationDeviceHelper::ProcessLogitsFunc<MLFloat16>& process_logits_fp16_func,
const GenerationDeviceHelper::InitBeamStateFunc<float>& init_beam_state_func,
const GenerationDeviceHelper::InitBeamStateFunc<MLFloat16>& init_beam_state_fp16_func) {
const GenerationDeviceHelper::InitBeamStateFunc<MLFloat16>& init_beam_state_fp16_func,
const GenerationDeviceHelper::CreateBeamScorer& create_beam_scorer_func) {
add_to_feeds_func_ = add_to_feeds_func;
topk_func_ = topk_func;
device_copy_func_ = device_copy_func;
@ -62,6 +63,7 @@ class BeamSearch : public IControlFlowKernel {
process_logits_fp16_func_ = process_logits_fp16_func;
init_beam_state_func_ = init_beam_state_func;
init_beam_state_fp16_func_ = init_beam_state_fp16_func;
create_beam_scorer_func_ = create_beam_scorer_func;
}
#ifdef USE_CUDA
@ -111,6 +113,7 @@ class BeamSearch : public IControlFlowKernel {
GenerationDeviceHelper::InitBeamStateFunc<float> init_beam_state_func_;
GenerationDeviceHelper::InitBeamStateFunc<MLFloat16> init_beam_state_fp16_func_;
GenerationDeviceHelper::CreateBeamScorer create_beam_scorer_func_;
#ifdef USE_CUDA
GenerationDeviceHelper::ReorderPastStateFunc reorder_past_state_func_;

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

@ -23,17 +23,19 @@ struct BeamSearchState : IBeamSearchState<T> {
size_t next_token_size = SafeInt<size_t>(batch_beam_size) * parameters.vocab_size;
this->next_token_logits = AllocateBuffer<T>(allocator, next_token_logits_buffer_, next_token_size);
this->next_token_scores = AllocateBuffer<float>(allocator, next_token_scores_buffer_, next_token_size);
this->next_tokens = AllocateBuffer<int32_t>(allocator, next_tokens_buffer_, SafeInt<size_t>(2) * batch_beam_size);
this->next_indices = AllocateBuffer<int32_t>(allocator, next_indices_buffer_, SafeInt<size_t>(2) * batch_beam_size);
this->next_scores = AllocateBuffer<float>(allocator, next_scores_buffer_, SafeInt<size_t>(2) * batch_beam_size);
constexpr size_t max_parts_of_vocab = 128;
size_t topk_buffer_size = SafeInt<size_t>(batch_beam_size) * (max_parts_of_vocab + 1) * parameters.num_beams * 2 * 2;
this->topk_buffer = AllocateBuffer<float>(allocator, topk_temp_buffer_, topk_buffer_size);
if (allocator->Info().device.Type() == OrtDevice::GPU) {
size_t sequences_elements = SafeInt<size_t>(2) * batch_beam_size * parameters.max_length;
this->sequences_device = AllocateBuffer<int32_t>(allocator, sequences_device_buffer_, sequences_elements);
}
if (use_position) {
this->next_positions = AllocateBuffer<int32_t>(allocator, next_positions_buffer_, batch_beam_size);
}
@ -54,11 +56,6 @@ struct BeamSearchState : IBeamSearchState<T> {
Tensor temp(DataTypeImpl::GetType<T>(), staging_for_past_state_reorder_buffer_shape, allocator);
this->staging_for_past_state_reorder = std::move(temp);
// We need a buffer on GPU to hold the final chosen indices after BeamScorer has finished processing
// TODO: This is a temporary work-around as BeamScorer currently only runs on CPU.
// We can remove these kinds of work-arounds once BeamScorer runs on CUDA eventually.
this->chosen_indices = AllocateBuffer<int32_t>(allocator, chosen_indices_buffer_, batch_beam_size);
}
}
@ -80,7 +77,7 @@ struct BeamSearchState : IBeamSearchState<T> {
BufferUniquePtr beam_scores_buffer_;
BufferUniquePtr scores_buffer_;
BufferUniquePtr topk_temp_buffer_;
BufferUniquePtr chosen_indices_buffer_;
BufferUniquePtr sequences_device_buffer_;
};
struct BeamSearchCpuState : IBeamSearchCpuState {
@ -174,7 +171,6 @@ class BeamSearchBase : public GenerateBase {
// Process logits and append next tokens to sequences.
Status GenerateNextToken(const OrtValue& logits,
gsl::span<int32_t>& beam_next_tokens,
gsl::span<int32_t>& beam_indices,
BeamSearchState<T>& beam_state,
BeamSearchCpuState& cpu_state,
int counter);
@ -188,7 +184,7 @@ class BeamSearchBase : public GenerateBase {
BeamSearchParameters* parameters_;
std::unique_ptr<BeamSearchScorer> beam_scorer_;
std::unique_ptr<IBeamScorer> beam_scorer_;
// Device specific functions
GenerationDeviceHelper::ProcessLogitsFunc<T> process_logits_func_;
@ -246,7 +242,7 @@ Status BeamSearchBase<T>::ProcessLogits(
BeamSearchCpuState& cpu_state,
AllocatorPtr& allocator,
int counter) {
return process_logits_func_(logits, &beam_state, &cpu_state, &(cpu_state.sequences), allocator,
return process_logits_func_(logits, &beam_state, &(cpu_state.sequences), allocator,
thread_pool_, &logits_processors_, beam_scorer_.get(),
parameters_, counter, ort_stream_, GetConsoleDumper());
}
@ -255,36 +251,69 @@ template <typename T>
Status BeamSearchBase<T>::GenerateNextToken(
const OrtValue& logits,
gsl::span<int32_t>& beam_next_tokens,
gsl::span<int32_t>& beam_indices,
BeamSearchState<T>& beam_state,
BeamSearchCpuState& cpu_state,
int counter) {
// Process logits to get next token scores
ORT_RETURN_IF_ERROR(ProcessLogits(logits, beam_state, cpu_state, temp_space_allocator_, counter));
gsl::span<float>& beam_scores = beam_scorer_->GetNextScores();
// It is optional to clone beam_scores. Change it to use same buffer also works for CPU:
// beam_state.beam_scores = beam_scores
// Here we make a copy to reduce the coupling with little cost (the buffer size is small).
ORT_RETURN_IF_ERROR(device_copy_func_(beam_state.beam_scores,
beam_scores,
ort_stream_,
DeviceCopyDirection::hostToDevice));
if (this->IsCuda()) {
auto beam_scores = beam_scorer_->GetNextScores();
// It is optional to clone beam_scores. Change it to use same buffer also works:
// beam_state.beam_scores = beam_scores
// Here we make a copy to reduce the coupling with little cost (the buffer size is small).
ORT_RETURN_IF_ERROR(device_copy_func_(beam_state.beam_scores,
beam_scores,
ort_stream_,
DeviceCopyDirection::deviceToDevice));
beam_next_tokens = beam_scorer_->GetNextTokens();
beam_indices = beam_scorer_->GetNextIndices();
beam_next_tokens = beam_scorer_->GetNextTokens();
#ifdef DEBUG_GENERATION
cpu_dumper_.Print("beam_scores from scorer", beam_scores.data(), parameters_->batch_size, parameters_->num_beams);
cpu_dumper_.Print("beam_next_tokens", beam_next_tokens.data(), parameters_->batch_size, parameters_->num_beams);
cpu_dumper_.Print("beam_indices", beam_indices.data(), parameters_->batch_size, parameters_->num_beams);
auto beam_indices = beam_scorer_->GetNextIndicesGPU();
cuda_dumper_->Print("beam_scores from scorer", beam_state.beam_scores.data(), parameters_->batch_size, parameters_->num_beams);
cuda_dumper_->Print("beam_next_tokens", beam_next_tokens.data(), parameters_->batch_size, parameters_->num_beams);
cuda_dumper_->Print("beam_indices", beam_indices.data(), parameters_->batch_size, parameters_->num_beams);
#endif
cpu_state.sequences.AppendNextTokenToSequences(beam_indices, beam_next_tokens);
// the Cuda beam scorer does the AppendNextTokenSequences, all that's left for Cuda is a this small step
cpu_state.sequences.AfterDeviceAppendedNextToken();
#ifdef DEBUG_GENERATION
cpu_state.sequences.PrintSequences(&cpu_dumper_);
// CUDA equivalent of cpu_state.sequences.PrintSequences
auto sequences_buffer = cpu_state.sequences.GetCurrentDeviceSequences();
for (int i = 0; i < parameters_->batch_size * parameters_->num_beams; i++) {
gsl::span<const int32_t> sequence = sequences_buffer.subspan(i * parameters_->max_length, cpu_state.sequences.GetSequenceLength());
cuda_dumper_->Print("sequences", i, false);
cuda_dumper_->Print(nullptr, sequence.data(), 1, static_cast<int>(sequence.size()));
}
#endif
} else {
auto beam_scores = beam_scorer_->GetNextScores();
// It is optional to clone beam_scores. Change it to use same buffer also works for CPU:
// beam_state.beam_scores = beam_scores
// Here we make a copy to reduce the coupling with little cost (the buffer size is small).
ORT_RETURN_IF_ERROR(device_copy_func_(beam_state.beam_scores,
beam_scores,
ort_stream_,
DeviceCopyDirection::hostToDevice));
beam_next_tokens = beam_scorer_->GetNextTokens();
auto beam_indices = beam_scorer_->GetNextIndicesCPU();
#ifdef DEBUG_GENERATION
cpu_dumper_.Print("beam_scores from scorer", beam_scores.data(), parameters_->batch_size, parameters_->num_beams);
cpu_dumper_.Print("beam_next_tokens", beam_next_tokens.data(), parameters_->batch_size, parameters_->num_beams);
cpu_dumper_.Print("beam_indices", beam_indices.data(), parameters_->batch_size, parameters_->num_beams);
#endif
cpu_state.sequences.AppendNextTokenToSequences(beam_indices, beam_next_tokens);
#ifdef DEBUG_GENERATION
cpu_state.sequences.PrintSequences(&cpu_dumper_);
#endif
}
return Status::OK();
}

55
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h
View file

@ -32,7 +32,8 @@ class BeamSearchGpt : public BeamSearchBase<T> {
const GenerationDeviceHelper::InitBeamStateFunc<T>& init_beam_state_func,
const GenerationDeviceHelper::DeviceCopyFunc<float>& device_copy_func,
const GenerationDeviceHelper::DeviceCopyFunc<int32_t>& device_copy_int32_func,
const GenerationDeviceHelper::UpdateGptFeedsFunc<T>& update_feeds_func)
const GenerationDeviceHelper::UpdateGptFeedsFunc<T>& update_feeds_func,
const GenerationDeviceHelper::CreateBeamScorer& create_beam_scorer_func)
: BeamSearchBase<T>(context, decoder_session_state, thread_pool,
ort_stream, cuda_dumper, params,
topk_func, process_logits_func, device_copy_func, device_copy_int32_func),
@ -42,7 +43,8 @@ class BeamSearchGpt : public BeamSearchBase<T> {
create_inputs_func_(create_inputs_func),
add_to_feeds_func_(add_to_feeds_func),
init_beam_state_func_(init_beam_state_func),
update_feeds_func_(update_feeds_func) {}
update_feeds_func_(update_feeds_func),
create_beam_scorer_func_(create_beam_scorer_func) {}
#ifdef USE_CUDA
Status InitializeCuda(
@ -102,6 +104,7 @@ class BeamSearchGpt : public BeamSearchBase<T> {
GenerationDeviceHelper::ReorderPastStateFunc reorder_past_state_func_;
#endif
GenerationDeviceHelper::UpdateGptFeedsFunc<T> update_feeds_func_;
GenerationDeviceHelper::CreateBeamScorer create_beam_scorer_func_;
const void* cuda_device_prop_ = nullptr;
int cuda_device_arch_ = 0;
@ -187,18 +190,15 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
const FeedsFetchesManager& feeds_fetches_manager) {
auto status = Status::OK();
const BeamSearchParameters* parameters = this->parameters_;
int64_t sequences_dims[] = {parameters->batch_size, parameters->num_return_sequences, parameters->max_length};
TensorShape sequences_shape(&sequences_dims[0], sizeof(sequences_dims) / sizeof(sequences_dims[0]));
TensorShape sequences_shape{parameters->batch_size, parameters->num_return_sequences, parameters->max_length};
Tensor* output_sequences = this->context_.Output(0, sequences_shape);
int64_t sequences_scores_dims[] = {parameters->batch_size, parameters->num_return_sequences};
TensorShape sequences_scores_shape(&sequences_scores_dims[0], 2);
TensorShape sequences_scores_shape{parameters->batch_size, parameters->num_return_sequences};
Tensor* output_sequences_scores = this->context_.Output(1, sequences_scores_shape);
int64_t scores_dims[] = {
TensorShape scores_shape{
static_cast<int64_t>(parameters->max_length) - static_cast<int64_t>(parameters->sequence_length),
parameters->batch_size, parameters->num_beams, parameters->vocab_size};
TensorShape scores_shape(&scores_dims[0], sizeof(scores_dims) / sizeof(scores_dims[0]));
Tensor* output_scores = this->context_.Output(2, scores_shape);
// Update the flag to indicate whether scores exists in output
@ -209,7 +209,9 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
std::vector<OrtValue> fetches;
// Initialize resources
this->beam_scorer_ = std::make_unique<BeamSearchScorer>(*parameters, this->cpu_allocator_);
this->beam_scorer_ = create_beam_scorer_func_
? create_beam_scorer_func_(*parameters, this->temp_space_allocator_, this->cpu_allocator_, this->ort_stream_)
: std::make_unique<BeamSearchScorer>(*parameters, this->cpu_allocator_);
BeamSearchCpuState cpu_state{*parameters,
this->cpu_allocator_,
@ -248,13 +250,22 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
cpu_state.SetExpandedSequence(expanded_input_ids_in_cpu.Get<Tensor>().DataAsSpan<int32_t>());
// beam_state.sequences_device is the GPU version of cpu_state.sequences_space,
// this copies it over to the GPU after setting it up on the CPU
if (this->IsCuda()) {
cpu_state.sequences.InitDevice(beam_state.sequences_device);
ORT_RETURN_IF_ERROR(this->device_copy_int32_func_(beam_state.sequences_device.subspan(0, beam_state.sequences_device.size() / 2),
cpu_state.sequences_space.subspan(0, cpu_state.sequences_space.size() / 2),
nullptr,
DeviceCopyDirection::hostToDevice));
}
#ifdef DEBUG_GENERATION
const IConsoleDumper* dumper = this->GetConsoleDumper();
#endif
// Position ids for all iterations except the first. It uses memory buffer owned by next_positions.
OrtValue position_ids;
int64_t dims[] = {parameters->BatchBeamSize(), 1};
TensorShape shape(&dims[0], 2);
TensorShape shape{parameters->BatchBeamSize(), 1};
Tensor::InitOrtValue(DataTypeImpl::GetType<int32_t>(),
shape,
beam_state.next_positions.data(),
@ -312,18 +323,15 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
const OrtValue& logits = fetches[0];
gsl::span<int32_t> beam_next_tokens;
gsl::span<int32_t> beam_indices;
ORT_RETURN_IF_ERROR(this->GenerateNextToken(logits,
beam_next_tokens,
beam_indices,
beam_state,
cpu_state,
iteration_counter));
// When all batches are finished, stop earlier to avoid wasting computation.
if (this->beam_scorer_->IsDone()) {
if (this->beam_scorer_->IsDone())
break;
}
// Increase sequence length after a new token is generated.
++current_length;
@ -356,15 +364,20 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
ORT_RETURN_IF_ERROR(UpdateFeeds(fetches, feeds, current_length,
position_ids, increase_position,
ReinterpretAsSpan<const int32_t>(beam_next_tokens),
ReinterpretAsSpan<const int32_t>(beam_indices),
gpt_subgraph_.has_decoder_masked_attention_
? ReinterpretAsSpan<const int32_t>(beam_state.chosen_indices)
? place_holder
: ReinterpretAsSpan<const int32_t>(this->beam_scorer_->GetNextIndicesCPU()),
gpt_subgraph_.has_decoder_masked_attention_
? ReinterpretAsSpan<const int32_t>(this->beam_scorer_->GetNextIndicesGPU())
: place_holder,
current_length - 1,
parameters->sequence_length,
gpt_subgraph_.has_decoder_masked_attention_));
}
if (this->beam_scorer_->IsDoneLater())
break;
if (gpt_subgraph_.past_present_share_buffer_) {
// clear fetched values before presents[]
for (int idx = 0; idx < gpt_subgraph_.GetFirstPresentOutputIndex(); idx++) {
@ -376,14 +389,6 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
}
gsl::span<const float> final_beam_scores = beam_state.beam_scores;
if (this->IsCuda()) {
ORT_RETURN_IF_ERROR(this->device_copy_func_(cpu_state.final_beam_scores,
final_beam_scores,
nullptr,
DeviceCopyDirection::deviceToHost));
final_beam_scores = cpu_state.final_beam_scores;
}
this->beam_scorer_->Finalize(cpu_state.sequences,
final_beam_scores,
output_sequences,

54
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h
View file

@ -36,7 +36,8 @@ class BeamSearchT5 : public BeamSearchBase<T> {
const GenerationDeviceHelper::UpdateDecoderFeedsFunc<T>& update_decoder_feeds_func,
const GenerationDeviceHelper::ExpandBufferFunc<int32_t>& expand_buffer_int32_func,
const GenerationDeviceHelper::ExpandBufferFunc<float>& expand_buffer_float_func,
const GenerationDeviceHelper::ExpandBufferFunc<MLFloat16>& expand_buffer_float16_func)
const GenerationDeviceHelper::ExpandBufferFunc<MLFloat16>& expand_buffer_float16_func,
const GenerationDeviceHelper::CreateBeamScorer& create_beam_scorer_func)
: BeamSearchBase<T>(context, decoder_session_state, thread_pool,
ort_stream, cuda_dumper, params,
topk_func, process_logits_func, device_copy_func, device_copy_int32_func),
@ -49,7 +50,8 @@ class BeamSearchT5 : public BeamSearchBase<T> {
update_decoder_feeds_func_(update_decoder_feeds_func),
expand_buffer_int32_func_(expand_buffer_int32_func),
expand_buffer_float_func_(expand_buffer_float_func),
expand_buffer_float16_func_(expand_buffer_float16_func) {}
expand_buffer_float16_func_(expand_buffer_float16_func),
create_beam_scorer_func_(create_beam_scorer_func) {}
#ifdef USE_CUDA
Status InitializeCuda(
@ -94,6 +96,7 @@ class BeamSearchT5 : public BeamSearchBase<T> {
GenerationDeviceHelper::ExpandBufferFunc<int32_t> expand_buffer_int32_func_;
GenerationDeviceHelper::ExpandBufferFunc<float> expand_buffer_float_func_;
GenerationDeviceHelper::ExpandBufferFunc<MLFloat16> expand_buffer_float16_func_;
GenerationDeviceHelper::CreateBeamScorer create_beam_scorer_func_;
const void* cuda_device_prop_ = nullptr;
int cuda_device_arch_ = 0;
@ -189,11 +192,6 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
// Initialize resources
// ------------------------------------
// Copy decoder_input_ids (in CPU) to sequence. It contains decoder_start_token_id for each beam.
cpu_state.SetUnexpandedSequence(decoder_input_ids.Get<Tensor>().DataAsSpan<int32_t>());
this->beam_scorer_ = std::make_unique<BeamSearchScorer>(*parameters, this->cpu_allocator_);
BeamSearchState<T> beam_state{*parameters,
this->temp_space_allocator_,
decoder_subgraph_.has_decoder_masked_attention_,
@ -205,11 +203,27 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
parameters->num_beams,
this->ort_stream_);
// Copy decoder_input_ids (in CPU) to sequence. It contains decoder_start_token_id for each beam.
cpu_state.SetUnexpandedSequence(decoder_input_ids.Get<Tensor>().DataAsSpan<int32_t>());
// beam_state.sequences_device is the GPU version of cpu_state.sequences_space,
// this copies it over to the GPU after setting it up on the CPU
if (this->IsCuda()) {
cpu_state.sequences.InitDevice(beam_state.sequences_device);
ORT_RETURN_IF_ERROR(this->device_copy_int32_func_(beam_state.sequences_device.subspan(0, beam_state.sequences_device.size() / 2),
cpu_state.sequences_space.subspan(0, cpu_state.sequences_space.size() / 2),
nullptr,
DeviceCopyDirection::hostToDevice));
}
this->beam_scorer_ = create_beam_scorer_func_
? create_beam_scorer_func_(*parameters, this->temp_space_allocator_, this->cpu_allocator_, this->ort_stream_)
: std::make_unique<BeamSearchScorer>(*parameters, this->cpu_allocator_);
// ------------------------------------------------------------------------------
// Generate next token from logits output from encoder, and initialize decoder inputs.
// ------------------------------------------------------------------------------
gsl::span<int32_t> beam_next_tokens;
gsl::span<int32_t> beam_indices;
int iteration_counter = 0;
std::vector<OrtValue> decoder_feeds;
@ -221,7 +235,6 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
++iteration_counter;
ORT_RETURN_IF_ERROR(this->GenerateNextToken(encoder_fetches[0],
beam_next_tokens,
beam_indices,
beam_state,
cpu_state,
iteration_counter));
@ -324,7 +337,6 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
const OrtValue& logits = decoder_fetches[0];
ORT_RETURN_IF_ERROR(this->GenerateNextToken(logits,
beam_next_tokens,
beam_indices,
beam_state,
cpu_state,
iteration_counter));
@ -334,6 +346,12 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
break;
}
// TODO: If this is safe to do after update_decoder_feeds_func, move it later so that we can speculatively run the next steps while we wait
// for the done result to transfer to the CPU
if (this->beam_scorer_->IsDoneLater()) {
break;
}
// Increase sequence length after a new token is generated.
++current_length;
@ -348,9 +366,11 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
decoder_feeds,
num_present_outputs,
ReinterpretAsSpan<const int32_t>(beam_next_tokens),
ReinterpretAsSpan<const int32_t>(beam_indices),
decoder_subgraph_.has_decoder_masked_attention_
? ReinterpretAsSpan<const int32_t>(beam_state.chosen_indices)
? place_holder
: ReinterpretAsSpan<const int32_t>(this->beam_scorer_->GetNextIndicesCPU()),
decoder_subgraph_.has_decoder_masked_attention_
? ReinterpretAsSpan<const int32_t>(this->beam_scorer_->GetNextIndicesGPU())
: place_holder,
parameters->num_beams,
decoder_subgraph_.GetFirstPastInputIndex(),
@ -375,21 +395,13 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
}
gsl::span<const float> final_beam_scores = beam_state.beam_scores;
if (this->IsCuda()) {
ORT_RETURN_IF_ERROR(this->device_copy_func_(cpu_state.final_beam_scores,
final_beam_scores,
nullptr,
DeviceCopyDirection::deviceToHost));
final_beam_scores = cpu_state.final_beam_scores;
}
this->beam_scorer_->Finalize(cpu_state.sequences,
final_beam_scores,
output_sequences,
output_sequences_scores);
// Output per token scores
if (output_scores != nullptr) {
if (output_scores) {
gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
gsl::span<const float> source = beam_state.scores;
assert(target.size() == source.size());

53
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h
View file

@ -35,7 +35,8 @@ class BeamSearchWhisper : public BeamSearchBase<T> {
const GenerationDeviceHelper::CreateWhisperEncoderInputsFunc& create_encoder_inputs_func,
const GenerationDeviceHelper::UpdateDecoderFeedsFunc<T>& update_decoder_feeds_func,
const GenerationDeviceHelper::ExpandBufferFunc<float>& expand_buffer_float_func,
const GenerationDeviceHelper::ExpandBufferFunc<MLFloat16>& expand_buffer_float16_func)
const GenerationDeviceHelper::ExpandBufferFunc<MLFloat16>& expand_buffer_float16_func,
const GenerationDeviceHelper::CreateBeamScorer& create_beam_scorer_func)
: BeamSearchBase<T>(context, decoder_session_state, thread_pool,
ort_stream, cuda_dumper, params,
topk_func, process_logits_func, device_copy_func, device_copy_int32_func),
@ -47,7 +48,8 @@ class BeamSearchWhisper : public BeamSearchBase<T> {
create_encoder_inputs_func_(create_encoder_inputs_func),
update_decoder_feeds_func_(update_decoder_feeds_func),
expand_buffer_float_func_(expand_buffer_float_func),
expand_buffer_float16_func_(expand_buffer_float16_func) {}
expand_buffer_float16_func_(expand_buffer_float16_func),
create_beam_scorer_func_(create_beam_scorer_func) {}
#ifdef USE_CUDA
Status InitializeCuda(
@ -91,6 +93,7 @@ class BeamSearchWhisper : public BeamSearchBase<T> {
GenerationDeviceHelper::UpdateDecoderFeedsFunc<T> update_decoder_feeds_func_;
GenerationDeviceHelper::ExpandBufferFunc<float> expand_buffer_float_func_;
GenerationDeviceHelper::ExpandBufferFunc<MLFloat16> expand_buffer_float16_func_;
GenerationDeviceHelper::CreateBeamScorer create_beam_scorer_func_;
const void* cuda_device_prop_ = nullptr;
int cuda_device_arch_ = 0;
@ -182,10 +185,6 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
// Initialize resources
// ------------------------------------
// Copy decoder_input_ids (in CPU) to sequence. It contains the initial decoder token ids for each beam.
cpu_state.SetUnexpandedSequence(decoder_input_ids.Get<Tensor>().DataAsSpan<int32_t>());
this->beam_scorer_ = std::make_unique<BeamSearchScorer>(*parameters, this->cpu_allocator_);
BeamSearchState<T> beam_state{*parameters,
this->temp_space_allocator_,
decoder_subgraph_.has_decoder_masked_attention_,
@ -197,11 +196,27 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
parameters->num_beams,
this->ort_stream_);
// Copy decoder_input_ids (in CPU) to sequence. It contains the initial decoder token ids for each beam.
cpu_state.SetUnexpandedSequence(decoder_input_ids.Get<Tensor>().DataAsSpan<int32_t>());
// beam_state.sequences_device is the GPU version of cpu_state.sequences_space,
// this copies it over to the GPU after setting it up on the CPU
if (this->IsCuda()) {
cpu_state.sequences.InitDevice(beam_state.sequences_device);
ORT_RETURN_IF_ERROR(this->device_copy_int32_func_(beam_state.sequences_device.subspan(0, beam_state.sequences_device.size() / 2),
cpu_state.sequences_space.subspan(0, cpu_state.sequences_space.size() / 2),
nullptr,
DeviceCopyDirection::hostToDevice));
}
this->beam_scorer_ = create_beam_scorer_func_
? create_beam_scorer_func_(*parameters, this->temp_space_allocator_, this->cpu_allocator_, this->ort_stream_)
: std::make_unique<BeamSearchScorer>(*parameters, this->cpu_allocator_);
// ------------------------------------------------------------------------------
// Generate next token from logits output from encoder, and initialize decoder inputs.
// ------------------------------------------------------------------------------
gsl::span<int32_t> beam_next_tokens;
gsl::span<int32_t> beam_indices;
int iteration_counter = 0;
std::vector<OrtValue> decoder_feeds;
@ -213,7 +228,6 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
++iteration_counter;
ORT_RETURN_IF_ERROR(this->GenerateNextToken(encoder_fetches[0],
beam_next_tokens,
beam_indices,
beam_state,
cpu_state,
iteration_counter));
@ -311,7 +325,6 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
const OrtValue& logits = decoder_fetches[0];
ORT_RETURN_IF_ERROR(this->GenerateNextToken(logits,
beam_next_tokens,
beam_indices,
beam_state,
cpu_state,
iteration_counter));
@ -321,6 +334,12 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
break;
}
// TODO: If this is safe to do after update_decoder_feeds_func, move it later so that we can speculatively run the next steps while we wait
// for the done result to transfer to the CPU
if (this->beam_scorer_->IsDoneLater()) {
break;
}
// Increase sequence length after a new token is generated.
++current_length;
@ -335,9 +354,11 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
decoder_feeds,
num_present_outputs,
ReinterpretAsSpan<const int32_t>(beam_next_tokens),
ReinterpretAsSpan<const int32_t>(beam_indices),
decoder_subgraph_.has_decoder_masked_attention_
? ReinterpretAsSpan<const int32_t>(beam_state.chosen_indices)
? place_holder
: ReinterpretAsSpan<const int32_t>(this->beam_scorer_->GetNextIndicesCPU()),
decoder_subgraph_.has_decoder_masked_attention_
? ReinterpretAsSpan<const int32_t>(this->beam_scorer_->GetNextIndicesGPU())
: place_holder,
parameters->num_beams,
decoder_subgraph_.GetFirstPastInputIndex(),
@ -362,21 +383,13 @@ Status BeamSearchWhisper<T>::Execute(const FeedsFetchesManager& encoder_feeds_fe
}
gsl::span<const float> final_beam_scores = beam_state.beam_scores;
if (this->IsCuda()) {
ORT_RETURN_IF_ERROR(this->device_copy_func_(cpu_state.final_beam_scores,
final_beam_scores,
nullptr,
DeviceCopyDirection::deviceToHost));
final_beam_scores = cpu_state.final_beam_scores;
}
this->beam_scorer_->Finalize(cpu_state.sequences,
final_beam_scores,
output_sequences,
output_sequences_scores);
// Output per token scores
if (output_scores != nullptr) {
if (output_scores) {
gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
gsl::span<const float> source = beam_state.scores;
assert(target.size() == source.size());

87
onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc
View file

@ -19,11 +19,11 @@ namespace contrib {
namespace transformers {
using ::onnxruntime::rnn::detail::Allocate;
void BeamHypotheses::Init(const IGenerationParameters& parameters, gsl::span<HypothesisScore> beams) {
length_penalty_ = parameters.length_penalty;
early_stopping_ = parameters.early_stopping;
void BeamHypotheses::Init(float length_penalty, gsl::span<HypothesisScore> beams) {
beams_ = beams;
beams_used_ = 0;
length_penalty_ = length_penalty;
done_ = false;
}
void BeamHypotheses::Add(gsl::span<const int32_t>& hypothesis, float sum_logprobs) {
@ -45,18 +45,9 @@ void BeamHypotheses::Add(gsl::span<const int32_t>& hypothesis, float sum_logprob
beams_[index] = HypothesisScore{hypothesis, score};
}
bool BeamHypotheses::IsDone(float best_sum_logprobs, int current_length) const {
// If there are enough hypotheses and that none of the hypotheses being generated can become better
// than the worst one in the heap, then we are done with this sentence.
if (static_cast<size_t>(beams_used_) < beams_.size())
return false;
if (early_stopping_)
return true;
bool BeamHypotheses::CanImprove(float best_sum_logprobs, int current_length) const {
float current_score = best_sum_logprobs / pow(static_cast<float>(current_length), length_penalty_);
return beams_.back().score >= current_score;
return beams_.back().score < current_score;
}
void BeamHypotheses::Output(
@ -85,17 +76,17 @@ BeamSearchScorer::BeamSearchScorer(const IGenerationParameters& parameters,
: batch_size_{static_cast<size_t>(parameters.batch_size)},
num_beams_{static_cast<size_t>(parameters.num_beams)},
max_length_{static_cast<size_t>(parameters.max_length)},
num_beam_hyps_to_keep_{static_cast<size_t>(parameters.num_return_sequences)},
num_return_sequences_{static_cast<size_t>(parameters.num_return_sequences)},
pad_token_id_{parameters.pad_token_id},
eos_token_id_{parameters.eos_token_id} {
eos_token_id_{parameters.eos_token_id},
early_stopping_{parameters.early_stopping},
not_done_count_{parameters.batch_size} {
size_t batch_beam_size = batch_size_ * num_beams_;
auto beams = Allocate<HypothesisScore>(allocator, batch_beam_size, hypothesis_scores_ptr_);
beam_hyps_ = Allocate<BeamHypotheses>(allocator, batch_size_, beam_hyps_ptr_);
for (size_t i = 0; i < batch_size_; i++)
beam_hyps_[i].Init(parameters, beams.subspan(i * num_beams_, num_beams_));
done_ = Allocate<bool>(allocator, batch_size_, done_ptr_, true /* fill allocated array */, false /* fill with false */);
beam_hyps_[i].Init(parameters.length_penalty, beams.subspan(i * num_beams_, num_beams_));
next_beam_scores_ = Allocate<float>(allocator, batch_beam_size, next_beam_scores_ptr_);
next_beam_tokens_ = Allocate<int32_t>(allocator, batch_beam_size, next_beam_tokens_ptr_);
@ -103,16 +94,7 @@ BeamSearchScorer::BeamSearchScorer(const IGenerationParameters& parameters,
// Space to store intermediate sequence with length sequence_length, sequence_length + 1, ..., max_sequence_length.
size_t per_beam = (SafeInt<size_t>(max_length_) * (max_length_ + 1) - (parameters.sequence_length - 1) * parameters.sequence_length) / 2;
hypothesis_buffer_length_ = batch_beam_size * per_beam;
hypothesis_buffer_ = Allocate<int32_t>(allocator, hypothesis_buffer_length_, hypothesis_buffer_ptr_);
}
bool BeamSearchScorer::IsDone() const {
for (auto done : done_) {
if (!done)
return false;
}
return true;
hypothesis_buffer_ = Allocate<int32_t>(allocator, batch_beam_size * per_beam, hypothesis_buffer_ptr_);
}
void BeamSearchScorer::Process(ISequences& sequences,
@ -130,8 +112,8 @@ void BeamSearchScorer::Process(ISequences& sequences,
for (size_t batch = 0; batch < batch_size_; batch++) {
BeamHypotheses& beam_hyp = beam_hyps_[batch];
if (done_[batch]) {
ORT_ENFORCE(beam_hyp.Size() == gsl::narrow_cast<int>(num_beams_),
if (beam_hyp.done_) {
ORT_ENFORCE(beam_hyp.beams_used_ == gsl::narrow_cast<int>(num_beams_),
"Batch can only be done if all beams have been generated");
// Pad the batch.
@ -161,9 +143,10 @@ void BeamSearchScorer::Process(ISequences& sequences,
// Clone the sequence and append to buffer.
gsl::span<const int32_t> src = sequences.GetSequence(batch_beam_idx);
auto clone = hypothesis_buffer_.subspan(hypothesis_buffer_offset_, sequence_length);
auto clone = hypothesis_buffer_.subspan(static_cast<size_t>(hypothesis_buffer_used_), sequence_length);
gsl::copy(src, clone);
hypothesis_buffer_offset_ += static_cast<size_t>(sequence_length);
hypothesis_buffer_used_ += sequence_length;
auto sequence = ReinterpretAsSpan<const int32_t>(clone);
beam_hyp.Add(sequence, next_score);
} else {
@ -180,16 +163,21 @@ void BeamSearchScorer::Process(ISequences& sequences,
}
ORT_ENFORCE(beam_idx == num_beams_);
ORT_ENFORCE(hypothesis_buffer_offset_ <= hypothesis_buffer_length_);
ORT_ENFORCE(static_cast<size_t>(hypothesis_buffer_used_) <= hypothesis_buffer_.size());
// Check if we are done so that we can save a pad step if all(done)
if (!done_[batch]) {
if (static_cast<size_t>(beam_hyp.beams_used_) < num_beams_)
continue;
if (!early_stopping_) {
gsl::span<const float> topk_scores = next_scores.subspan(batch * num_beams_, top_k);
const auto best_sum_logprobs = std::max_element(topk_scores.begin(), topk_scores.end());
if (beam_hyp.IsDone(*best_sum_logprobs, sequence_length)) {
done_[batch] = true;
}
if (beam_hyp.CanImprove(*best_sum_logprobs, sequence_length))
continue;
}
beam_hyp.done_ = true;
not_done_count_--;
}
}
@ -202,7 +190,7 @@ void BeamSearchScorer::Finalize(ISequences& sequences,
// Finalize all open beam hypotheses and add to generated hypotheses.
for (size_t batch_index = 0; batch_index < batch_size_; batch_index++) {
BeamHypotheses& beam_hyp = beam_hyps_[batch_index];
if (done_[batch_index]) {
if (beam_hyp.done_) {
continue;
}
@ -226,26 +214,17 @@ void BeamSearchScorer::Finalize(ISequences& sequences,
sequence_scores = output_sequence_scores->MutableDataAsSpan<float>();
}
// Span is empty when output_sequence_scores is NULL.
gsl::span<float> batch_sequence_score;
// Select the best hypotheses according to number of sequences to return.
for (size_t batch_index = 0; batch_index < batch_size_; batch_index++) {
BeamHypotheses& beam_hyp = beam_hyps_[batch_index];
const size_t num_return_sequences = num_beam_hyps_to_keep_;
auto batch_output = output.subspan(batch_index * num_return_sequences * max_length_,
num_return_sequences * max_length_);
auto batch_output = output.subspan(batch_index * num_return_sequences_ * max_length_,
num_return_sequences_ * max_length_);
gsl::span<float> sequence_scores_buffer;
if (!sequence_scores.empty())
sequence_scores_buffer = sequence_scores.subspan(batch_index * num_return_sequences_, num_return_sequences_);
if (output_sequence_scores) {
batch_sequence_score = sequence_scores.subspan(batch_index * num_return_sequences, num_return_sequences);
}
beam_hyp.Output(
narrow<int>(num_return_sequences),
narrow<int>(max_length_),
batch_output,
batch_sequence_score);
beam_hyp.Output(num_return_sequences_, max_length_, batch_output, sequence_scores_buffer);
}
}

37
onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h
View file

@ -25,15 +25,14 @@ struct HypothesisScore {
};
struct BeamHypotheses {
void Init(const IGenerationParameters& parameters, gsl::span<HypothesisScore> beams);
// Number of hypotheses
int Size() const { return beams_used_; }
// As these are constructed as an uninitialized array of memory, we need an Init method
void Init(float length_penalty, gsl::span<HypothesisScore> beams);
// Add a new hypothesis
void Add(gsl::span<const int32_t>& hypothesis, float sum_logprobs);
bool IsDone(float best_sum_logprobs, int current_length) const;
// Return true if this beats the worst score in the hypothesis
bool CanImprove(float best_sum_logprobs, int current_length) const;
// Output results
void Output(int top_k, // number of sequences to return
@ -41,15 +40,13 @@ struct BeamHypotheses {
gsl::span<int32_t>& sequences, // buffer with pad token, shape (num_return_sequences, max_length)
gsl::span<float>& sequences_scores); // buffer for sequence scores, with shape (num_return_sequences)
private:
float length_penalty_;
bool early_stopping_;
gsl::span<HypothesisScore> beams_; // Beam width sized array of hypotheses, sorted by highest scoring
int beams_used_; // Number of elements used in beams_
float length_penalty_;
bool done_;
};
class BeamSearchScorer : public IBeamScorer {
public:
struct BeamSearchScorer : IBeamScorer {
BeamSearchScorer(const IGenerationParameters& parameters,
AllocatorPtr& allocator);
@ -63,22 +60,21 @@ class BeamSearchScorer : public IBeamScorer {
Tensor* output_sequences,
Tensor* output_sequence_scores) override;
bool IsDone() const;
bool IsDone() const override { return not_done_count_ == 0; }
gsl::span<float>& GetNextScores() { return next_beam_scores_; }
gsl::span<int32_t>& GetNextTokens() { return next_beam_tokens_; }
gsl::span<int32_t>& GetNextIndices() override { return next_beam_indices_; }
gsl::span<float> GetNextScores() override { return next_beam_scores_; }
gsl::span<int32_t> GetNextTokens() override { return next_beam_tokens_; }
gsl::span<int32_t> GetNextIndicesCPU() override { return next_beam_indices_; }
private:
size_t batch_size_;
size_t num_beams_;
size_t max_length_;
size_t num_beam_hyps_to_keep_;
size_t num_return_sequences_;
int pad_token_id_;
int eos_token_id_;
IAllocatorUniquePtr<bool> done_ptr_; // Allocated buffer for done_
gsl::span<bool> done_; // Flags indicates whether each batch is finished or not. Shape is (batch_size).
bool early_stopping_;
int not_done_count_; // When zero, every batch entry is done (starts at batch_size_)
IAllocatorUniquePtr<float> next_beam_scores_ptr_;
gsl::span<float> next_beam_scores_;
@ -91,12 +87,11 @@ class BeamSearchScorer : public IBeamScorer {
IAllocatorUniquePtr<int32_t> hypothesis_buffer_ptr_; // Allocated buffer to hold all hypotheses
gsl::span<int32_t> hypothesis_buffer_; // Span of the allocated buffer
size_t hypothesis_buffer_length_{}; // Total number of elements
size_t hypothesis_buffer_offset_{}; // Offset of available buffer, or length of used buffer.
int hypothesis_buffer_used_{}; // Offset of available buffer, or length of used buffer.
IAllocatorUniquePtr<HypothesisScore> hypothesis_scores_ptr_; // num_beams_ * batch_size_, divided into num_beams_ chunks per BeamHypothesis in beam_hyps_
IAllocatorUniquePtr<BeamHypotheses> beam_hyps_ptr_;
gsl::span<BeamHypotheses> beam_hyps_; // batch_size_ count
gsl::span<BeamHypotheses> beam_hyps_; // Shape is batch_size_
};
} // namespace transformers

14
onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h
View file

@ -82,11 +82,13 @@ class GenerateBase {
implicit_inputs_(context_.GetImplicitInputs()),
ort_stream_(ort_stream),
cuda_dumper_(cuda_dumper),
cpu_allocator_(nullptr),
cpu_allocator_(decoder_session_state.GetAllocator(
decoder_session_state.GetExecutionProviders()
.Get(onnxruntime::kCpuExecutionProvider)
->GetOrtDeviceByMemType(OrtMemTypeDefault))),
temp_space_allocator_(nullptr),
topk_func_(topk_func),
device_copy_func_(device_copy_func) {
cpu_allocator_ = decoder_session_state.GetAllocator(decoder_session_state.GetExecutionProviders().Get(onnxruntime::kCpuExecutionProvider)->GetOrtDeviceByMemType(OrtMemTypeDefault));
}
virtual ~GenerateBase() = default;
@ -225,9 +227,13 @@ class GenerateBase {
}
protected:
bool IsCuda() const { return ort_stream_ != nullptr; }
bool IsCuda() const {
return ort_stream_ != nullptr;
}
const IConsoleDumper* GetConsoleDumper() const { return IsCuda() ? cuda_dumper_ : &(cpu_dumper_); }
const IConsoleDumper* GetConsoleDumper() const {
return IsCuda() ? cuda_dumper_ : &(cpu_dumper_);
}
OpKernelContextInternal& context_;

3
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc
View file

@ -280,7 +280,6 @@ void InitGreedyState(transformers::IGreedySearchState<T>* greedy_state,
template <typename T>
Status ProcessLogits(const OrtValue& logits, // logits output of subgraph
transformers::IBeamSearchState<T>* beam_state, // state
transformers::IBeamSearchCpuState* cpu_state, // state in CPU
transformers::ISequences* sequences, // sequences
AllocatorPtr& allocator, // default allocator
onnxruntime::concurrency::ThreadPool* thread_pool, // thread pool (for CPU only)
@ -290,7 +289,6 @@ Status ProcessLogits(const OrtValue& logits, //
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper) { // tensor dumper
ORT_UNUSED_PARAMETER(cpu_state);
#ifndef DEBUG_GENERATION
ORT_UNUSED_PARAMETER(dumper);
#endif
@ -946,7 +944,6 @@ template void InitGreedyState<float>(
template Status ProcessLogits<float>(
const OrtValue& logits,
transformers::IBeamSearchState<float>* beam_state,
transformers::IBeamSearchCpuState* cpu_state,
transformers::ISequences* sequences,
AllocatorPtr& allocator,
onnxruntime::concurrency::ThreadPool* thread_pool,

8
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h
View file

@ -82,6 +82,12 @@ using InitBeamStateFunc = std::function<void(
int num_beams,
Stream* stream)>;
using CreateBeamScorer = std::function<std::unique_ptr<transformers::IBeamScorer>(
const transformers::IGenerationParameters& parameters,
AllocatorPtr& allocator,
AllocatorPtr& allocator_cpu,
Stream* stream)>;
template <typename T>
using InitGreedyStateFunc = std::function<void(
transformers::IGreedySearchState<T>* greedy_state,
@ -92,7 +98,6 @@ template <typename T>
using ProcessLogitsFunc = std::function<Status(
const OrtValue& logits, // logits output of subgraph
transformers::IBeamSearchState<T>* beam_state, // state
transformers::IBeamSearchCpuState* cpu_state, // state in CPU
transformers::ISequences* sequences, // sequences
AllocatorPtr& allocator, // default allocator
onnxruntime::concurrency::ThreadPool* thread_pool, // thread pool (for CPU only)
@ -235,7 +240,6 @@ void InitGreedyState(transformers::IGreedySearchState<T>* greedy_state,
template <typename T>
Status ProcessLogits(const OrtValue& logits, // logits output of subgraph
transformers::IBeamSearchState<T>* beam_state, // state
transformers::IBeamSearchCpuState* cpu_state, // state in CPU
transformers::ISequences* sequences, // sequences
AllocatorPtr& allocator, // default allocator
onnxruntime::concurrency::ThreadPool* thread_pool, // thread pool (for CPU only)

13
onnxruntime/contrib_ops/cpu/transformers/generation_shared.h
View file

@ -39,8 +39,7 @@ struct IBeamSearchState {
// in total, it will be:
// 2 * (batch_size * num_beams * (parts_vocab + 1), 2 * num_beams)
// The final chosen indices after BeamScorer has finished processing
gsl::span<int32_t> chosen_indices; // shape (batch_size, num_beams)
gsl::span<int32_t> sequences_device; // shape (2 * batch_size * max_length)
Tensor staging_for_past_state_reorder; // Tensor of shape (batch_size * num_beams, num_heads, max_length, head_size)
};
@ -96,6 +95,8 @@ struct ISamplingState {
struct ISequences {
virtual ~ISequences() {}
virtual gsl::span<const int32_t> GetSequence(int beam_index) const = 0;
virtual gsl::span<const int32_t> GetCurrentDeviceSequences() const = 0; // Get all current beam_index sequences in one continuous block (to pass to CUDA)
virtual gsl::span<int32_t> GetNextDeviceSequences() = 0; // Get all next beam_index sequences in one continuous block (to pass to CUDA)
virtual int GetSequenceLength() const = 0;
};
@ -118,7 +119,13 @@ struct IBeamScorer {
Tensor* output_sequences,
Tensor* output_sequence_scores) = 0;
virtual gsl::span<int32_t>& GetNextIndices() = 0;
virtual bool IsDone() const = 0; // GPU version will return false here, as it asynchronously queues up the event
virtual bool IsDoneLater() const { return false; } // GPU version waits for the asynchous result to complete here
virtual gsl::span<float> GetNextScores() = 0;
virtual gsl::span<int32_t> GetNextTokens() = 0;
virtual gsl::span<int32_t> GetNextIndicesCPU() = 0;
virtual gsl::span<int32_t> GetNextIndicesGPU() { return {}; } // If this is non CPU, returns the device buffer of the indices
};
struct IGenerationParameters {

16
onnxruntime/contrib_ops/cpu/transformers/sequences.cc
View file

@ -22,6 +22,11 @@ void Sequences::Init(gsl::span<int32_t> buffer, int batch_beam_size, int sequenc
current_length_ = sequence_length;
}
void Sequences::InitDevice(gsl::span<int32_t> buffer) {
device_sequences[0] = buffer.subspan(0, buffer.size() / 2);
device_sequences[1] = buffer.subspan(buffer.size() / 2);
}
gsl::span<const int32_t> Sequences::GetSequence(int beam_index) const {
gsl::span<const int32_t> buffer = sequences[current_sequences_buffer];
return buffer.subspan(SafeInt<size_t>(beam_index) * max_length_, static_cast<gsl::index>(current_length_));
@ -54,10 +59,8 @@ void Sequences::AppendNextTokenToSequences(
gsl::span<int32_t> target = output.subspan(SafeInt<size_t>(i) * max_length_,
static_cast<gsl::index>(current_length_));
gsl::copy(source, target);
}
// Append next token to each beam.
for (int i = 0; i < batch_beam_size_; i++) {
// Append next token to each beam.
output[SafeInt<size_t>(i) * max_length_ + current_length_] = beam_next_tokens[i];
}
@ -68,7 +71,7 @@ void Sequences::AppendNextTokenToSequences(
}
void Sequences::AppendNextTokenToSequences(gsl::span<int32_t>& next_tokens) {
auto output = sequences[current_sequences_buffer];
auto output = sequences[0];
// Append next token to each sequence.
for (int i = 0; i < batch_beam_size_; i++) {
@ -78,6 +81,11 @@ void Sequences::AppendNextTokenToSequences(gsl::span<int32_t>& next_tokens) {
++current_length_;
}
void Sequences::AfterDeviceAppendedNextToken() {
++current_length_;
current_sequences_buffer ^= 1;
}
} // namespace transformers
} // namespace contrib
} // namespace onnxruntime

6
onnxruntime/contrib_ops/cpu/transformers/sequences.h
View file

@ -15,9 +15,12 @@ class Sequences : public ISequences {
public:
// Initialize the sequence.
void Init(gsl::span<int32_t> buffer, int batch_beam_size, int sequence_length, int max_length);
void InitDevice(gsl::span<int32_t> buffer);
// Returns a sequence of word IDs for a given beam index ( beam_index < batch_beam_size).
gsl::span<const int32_t> GetSequence(int beam_index) const override;
gsl::span<const int32_t> GetCurrentDeviceSequences() const override { return device_sequences[current_sequences_buffer]; }
gsl::span<int32_t> GetNextDeviceSequences() override { return device_sequences[current_sequences_buffer ^ 1]; }
// Returns current sequence length.
int GetSequenceLength() const override;
@ -35,11 +38,14 @@ class Sequences : public ISequences {
void AppendNextTokenToSequences(
gsl::span<int32_t>& next_tokens);
void AfterDeviceAppendedNextToken();
private:
// Two buffers of shape (batch_size, num_beams, max_seq_length) to store sequences.
// At each time, there is only one buffer is active. The other one will be active in next token.
// Each AppendNextTokenToSequences call will trigger a rotation of active buffer.
gsl::span<int32_t> sequences[2];
gsl::span<int32_t> device_sequences[2];
// Index (either 0 or 1) of two buffers that is currently is active.
int current_sequences_buffer;

6
onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc
View file

@ -78,8 +78,7 @@ Status GptSubgraph::CreateInitialFeeds(
auto past_type = IsOutputFloat16() ? DataTypeImpl::GetType<MLFloat16>() : DataTypeImpl::GetType<float>();
if (!past_present_share_buffer_) {
// Initialize empty past state
int64_t past_state_dims[] = {2, batch_size * num_beams, num_heads, 0, head_size};
TensorShape past_shape(&past_state_dims[0], 5);
TensorShape past_shape{2, batch_size * num_beams, num_heads, 0, head_size};
OrtValue empty_past;
Tensor::InitOrtValue(past_type, past_shape, default_allocator, empty_past);
@ -89,8 +88,7 @@ Status GptSubgraph::CreateInitialFeeds(
}
} else {
// Past state feeds
int64_t past_state_dims[] = {2, batch_size * num_beams, num_heads, past_present_share_buffer_max_seq_len, head_size};
TensorShape past_shape(&past_state_dims[0], 5);
TensorShape past_shape{2, batch_size * num_beams, num_heads, past_present_share_buffer_max_seq_len, head_size};
// The remaining inputs are past state except the last one or three (see below for details)
// If `need_cache_indir` is false, then the last input is `past_sequence_length`

3
onnxruntime/contrib_ops/cuda/transformers/beam_search.cc
View file

@ -44,7 +44,8 @@ BeamSearch::BeamSearch(const OpKernelInfo& info)
GenerationCudaDeviceHelper::ProcessLogits<float>,
GenerationCudaDeviceHelper::ProcessLogits<MLFloat16>,
GenerationCudaDeviceHelper::InitBeamState<float>,
GenerationCudaDeviceHelper::InitBeamState<MLFloat16>);
GenerationCudaDeviceHelper::InitBeamState<MLFloat16>,
GenerationCudaDeviceHelper::CreateBeamScorer);
#ifndef USE_ROCM
SetDeviceHelpers_Cuda(GenerationCudaDeviceHelper::ReorderPastState, GenerationCudaDeviceHelper::InitCacheIndir);

321
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu
View file

@ -75,7 +75,7 @@ __global__ void LogitsProcessKernel(
int padded_vocab_size,
int total_elements,
int demote_token_id,
int32_t* sequences,
const int32_t* sequences,
int max_sequence_length,
int current_sequence_length,
float repetition_penalty,
@ -91,7 +91,7 @@ __global__ void LogitsProcessKernel(
} else {
// RepetitionPenaltyLogitsProcessor
if (repetition_penalty != 1.0f) {
int32_t* current_sequence = sequences + batch_beam_index * max_sequence_length;
const int32_t* current_sequence = sequences + batch_beam_index * max_sequence_length;
bool found = false;
for (int i = 0; i < current_sequence_length; i++) {
if (current_sequence[i] == word_id) {
@ -107,7 +107,7 @@ __global__ void LogitsProcessKernel(
// NoRepeatNGramLogitsProcessor
if (no_repeat_ngram_size > 0 && current_sequence_length >= no_repeat_ngram_size) {
int32_t* current_sequence = sequences + batch_beam_index * max_sequence_length;
const int32_t* current_sequence = sequences + batch_beam_index * max_sequence_length;
bool found = false;
for (int i = no_repeat_ngram_size - 1; i < current_sequence_length; i++) {
if (current_sequence[i] == word_id) { // last token of n-gram matched
@ -176,7 +176,7 @@ void LaunchLogitsProcessKernel(
int vocab_size,
int padded_vocab_size,
int demote_token_id,
int32_t* sequences,
const int32_t* sequences,
int max_sequence_length,
int current_sequence_length,
float repetition_penalty,
@ -217,7 +217,7 @@ template void LaunchLogitsProcessKernel(
int vocab_size,
int padded_vocab_size,
int demote_token_id,
int32_t* sequences,
const int32_t* sequences,
int max_sequence_length,
int current_sequence_length,
float repetition_penalty,
@ -236,13 +236,322 @@ template void LaunchLogitsProcessKernel(
int vocab_size,
int padded_vocab_size,
int demote_token_id,
int32_t* sequences,
const int32_t* sequences,
int max_sequence_length,
int current_sequence_length,
float repetition_penalty,
int no_repeat_ngram_size,
cudaStream_t stream);
__global__ void InitializeBeamHypotheses(BeamHypotheses* beam_hyps, int beam_hyps_count, float length_penalty, HypothesisScore* beams, int num_beams) {
int index = blockIdx.x * blockDim.x + threadIdx.x;
if (index >= beam_hyps_count)
return;
BeamHypotheses& beam_hyp = beam_hyps[index];
beam_hyp.beams_ = beams + index * num_beams;
beam_hyp.beams_count_ = num_beams;
beam_hyp.beams_used_ = 0;
beam_hyp.length_penalty_ = length_penalty;
beam_hyp.done_ = false;
}
// For counts that are typically far less than 256, this will round up the count to the next multiple of 32
// If this winds up being >256 then it uses a block size of 256 and calculates the appropriate grid_size
struct GridBlock32 {
GridBlock32(int count) {
block_size_ = (count + 31) & ~31; // Round up to nearest multiple of 32
if (block_size_ > 256) {
grid_size_ = (block_size_ + 255) / 256;
block_size_ = 256;
}
}
int grid_size_{1};
int block_size_;
};
void LaunchInitializeBeamHypotheses(gsl::span<BeamHypotheses> beam_hyps,
float length_penalty,
gsl::span<HypothesisScore> beams,
int num_beams,
cudaStream_t stream) {
GridBlock32 gb32{static_cast<int>(beam_hyps.size())};
InitializeBeamHypotheses<<<gb32.grid_size_, gb32.block_size_, 0, stream>>>(beam_hyps.data(),
static_cast<int>(beam_hyps.size()),
length_penalty,
beams.data(),
num_beams);
}
__device__ void BeamHypotheses::Add(const int32_t* hypothesis, int hypothesis_length, float sum_logprobs) {
float score = sum_logprobs / pow(static_cast<float>(hypothesis_length), length_penalty_);
size_t index = beams_used_;
// If the array is full, don't add unless it's better than the worst element
if (index == beams_count_) {
if (score <= beams_[--index].score)
return;
} else
beams_used_++;
// Rotate existing elements over while the new element scores higher
for (; index > 0 && score > beams_[index - 1].score; index--)
beams_[index] = beams_[index - 1];
beams_[index] = HypothesisScore{hypothesis, hypothesis_length, score};
}
__device__ bool BeamHypotheses::CanImprove(float best_sum_logprobs, int current_length) const {
float current_score = best_sum_logprobs / pow(static_cast<float>(current_length), length_penalty_);
return beams_[beams_count_ - 1].score < current_score;
}
__device__ void BeamHypotheses::Output(
int top_k,
int max_length,
int pad_token_id,
int32_t* sequences, // buffer of shape (num_return_sequences, max_length)
float* sequences_scores) // buffer of shape (num_return_sequences) or empty
{
// Copy the top_k beams into the sequences
for (int index = 0; index < top_k; index++) {
auto& item = beams_[index];
int32_t* target = sequences + index * max_length;
// Note that word_ids might be less than max_length.
for (int i = 0; i < item.hypothesis_length; i++)
target[i] = item.hypothesis[i];
// Pad remaining values with pad token id
for (int i = item.hypothesis_length; i < max_length; i++)
target[i] = pad_token_id;
if (sequences_scores)
sequences_scores[index] = item.score;
}
}
__global__ void BeamSearchScorer_Process(BeamScorerState& state_cpu,
BeamScorerState& state,
const int32_t* sequences_buffer,
int sequence_length,
BeamHypotheses* beam_hyps_,
float* next_beam_scores_,
int32_t* next_beam_tokens_,
int32_t* next_beam_indices_,
int32_t* hypothesis_buffer_,
const float* next_scores,
const int32_t* next_tokens,
const int32_t* next_indices) {
// Sequences shape is (batch_size * num_beams, total_sequence_length)
// It contains word ID of whole sequence generated so far.
// It is different from subgraph input_ids, which only need one word when past state is not empty.
int batch = threadIdx.x;
int batch_start = batch * state.num_beams_;
cuda::BeamHypotheses& beam_hyp = beam_hyps_[batch];
if (!beam_hyp.done_) {
// Next tokens for this sentence.
size_t beam_idx = 0;
size_t top_k = 2 * state.num_beams_;
for (size_t j = 0; j < top_k; j++) {
int32_t next_token = next_tokens[batch * top_k + j];
float next_score = next_scores[batch * top_k + j];
int32_t next_index = next_indices[batch * top_k + j];
int batch_beam_idx = batch_start + next_index;
// Add to generated hypotheses if end of sentence.
if ((state.eos_token_id_ >= 0) && (next_token == state.eos_token_id_)) {
bool is_beam_token_worse_than_top_num_beams = (j >= state.num_beams_);
if (is_beam_token_worse_than_top_num_beams) {
continue;
}
// Clone the sequence and append to buffer.
const int32_t* src = sequences_buffer + batch_beam_idx * state.max_length_;
auto clone = hypothesis_buffer_ + atomicAdd(&state.hypothesis_buffer_used_, sequence_length);
for (unsigned i = 0; i < sequence_length; i++)
clone[i] = src[i];
beam_hyp.Add(clone, sequence_length, next_score);
} else {
// Add next predicted token since it is not eos_token.
next_beam_scores_[batch_start + beam_idx] = next_score;
next_beam_tokens_[batch_start + beam_idx] = next_token;
next_beam_indices_[batch_start + beam_idx] = batch_beam_idx;
++beam_idx;
}
// Once the beam for next step is full, don't add more tokens to it.
if (beam_idx == state.num_beams_)
break;
}
// Check if we are done so that we can save a pad step if all(done)
if (beam_hyp.beams_used_ == state.num_beams_) {
if (state.early_stopping_ || !beam_hyp.CanImprove(*std::max_element(next_scores + batch_start, next_scores + batch_start + top_k), sequence_length)) {
beam_hyp.done_ = true;
if (atomicAdd(&state.not_done_count_, -1) == 0)
state_cpu.not_done_count_ = 0; // Update the CPU side
}
}
} else {
// Pad the batch.
for (size_t beam_idx = 0; beam_idx < state.num_beams_; beam_idx++) {
next_beam_scores_[batch_start + beam_idx] = 0.0f;
next_beam_tokens_[batch_start + beam_idx] = state.pad_token_id_;
next_beam_indices_[batch_start + beam_idx] = 0;
}
}
}
void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps,
gsl::span<float> next_beam_scores,
gsl::span<int32_t> next_beam_tokens,
gsl::span<int32_t> next_beam_indices,
gsl::span<int32_t> hypothesis_buffer,
gsl::span<const float> next_scores,
gsl::span<const int32_t> next_tokens,
gsl::span<const int32_t> next_indices,
cudaStream_t stream) {
BeamSearchScorer_Process<<<1, state_cpu.batch_size_, 0, stream>>>(state_cpu,
state,
sequences.data(),
sequence_length,
beam_hyps.data(),
next_beam_scores.data(),
next_beam_tokens.data(),
next_beam_indices.data(),
hypothesis_buffer.data(),
next_scores.data(),
next_tokens.data(),
next_indices.data());
}
__global__ void BeamSearchScorer_AppendNextTokenToSequences1(BeamScorerState& state,
int batch_beam_size,
const int32_t* sequences_buffer,
int32_t* next_sequences,
int sequence_length,
int32_t* next_beam_indices_) {
int beam_idx = threadIdx.x + blockIdx.x * blockDim.x;
if (beam_idx >= batch_beam_size)
return;
int sequence_index = threadIdx.y + blockIdx.y * blockDim.y;
if (sequence_index >= sequence_length)
return;
int beam_index = next_beam_indices_[beam_idx];
next_sequences[beam_idx * state.max_length_ + sequence_index] = sequences_buffer[beam_index * state.max_length_ + sequence_index];
}
__global__ void BeamSearchScorer_AppendNextTokenToSequences2(BeamScorerState& state,
int32_t* next_sequences,
int sequence_length,
const int32_t* next_beam_tokens_) {
int beam_idx = threadIdx.x;
next_sequences[beam_idx * state.max_length_ + sequence_length] = next_beam_tokens_[beam_idx];
}
void LaunchBeamSearchScorer_AppendNextTokenToSequences(BeamScorerState& state_cpu,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
gsl::span<int32_t> next_sequences,
int sequence_length,
gsl::span<int32_t> next_beam_tokens,
gsl::span<int32_t> next_beam_indices,
cudaStream_t stream) {
const int max_threads = 512;
int batch_beam_size = state_cpu.batch_size_ * state_cpu.num_beams_;
dim3 block_size;
dim3 grid_size;
if (batch_beam_size * sequence_length <= max_threads) { // Can fit into a single thread block
block_size.x = batch_beam_size;
block_size.y = sequence_length;
} else {
if (sequence_length <= max_threads) // Sequence length fits into thread block, but batch_beam_size does not, so chunk it
{
block_size.x = max_threads / sequence_length;
block_size.y = sequence_length;
grid_size.x = (batch_beam_size + block_size.x - 1) / block_size.x;
} else { // Exceed max_threads in every dimension, so divide into max_thread chunks
block_size.x = 1;
block_size.y = max_threads;
grid_size.x = batch_beam_size;
grid_size.y = (sequence_length + block_size.y - 1) / block_size.y;
}
}
BeamSearchScorer_AppendNextTokenToSequences1<<<grid_size, block_size, 0, stream>>>(state,
batch_beam_size,
sequences.data(),
next_sequences.data(),
sequence_length,
next_beam_indices.data());
BeamSearchScorer_AppendNextTokenToSequences2<<<1, batch_beam_size, 0, stream>>>(state,
next_sequences.data(),
sequence_length,
next_beam_tokens.data());
}
__global__ void BeamSearchScorer_Finalize(BeamScorerState& state,
const int32_t* sequences_buffer,
int sequence_length,
BeamHypotheses* beam_hyps_,
const float* final_beam_scores,
int32_t* output,
float* sequence_scores) {
int batch_index = blockIdx.x * blockDim.x + threadIdx.x;
if (batch_index >= state.batch_size_)
return;
// Finalize all open beam hypotheses and add to generated hypotheses.
cuda::BeamHypotheses& beam_hyp = beam_hyps_[batch_index];
if (!beam_hyp.done_) {
for (size_t beam_index = 0; beam_index < state.num_beams_; beam_index++) {
size_t batch_beam_index = batch_index * state.num_beams_ + beam_index;
float final_score = final_beam_scores[batch_beam_index];
auto final_tokens = sequences_buffer + batch_beam_index * state.max_length_;
beam_hyp.Add(final_tokens, sequence_length, final_score);
}
}
// Select the best hypotheses according to number of sequences to return.
auto batch_output = output + batch_index * state.num_return_sequences_ * state.max_length_;
beam_hyp.Output(
state.num_return_sequences_,
state.max_length_,
state.pad_token_id_,
batch_output,
sequence_scores ? sequence_scores + batch_index * state.num_return_sequences_ : nullptr);
}
void LaunchBeamSearchScorer_Finalize(int batch_size,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps,
gsl::span<const float> final_beam_scores,
gsl::span<int32_t> output,
gsl::span<float> sequence_scores,
cudaStream_t stream) {
BeamSearchScorer_Finalize<<<1, batch_size, 0, stream>>>(state,
sequences.data(),
sequence_length,
beam_hyps.data(),
final_beam_scores.data(),
output.data(),
sequence_scores.data());
}
__global__ void AddProbsKernel(float* log_probs,
float* cum_log_probs,
const int vocab_size,

77
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h
View file

@ -38,13 +38,88 @@ void LaunchLogitsProcessKernel(
int vocab_size,
int padded_vocab_size,
int demote_token_id,
int32_t* sequences,
const int32_t* sequences,
int max_sequence_length,
int current_sequence_length,
float repetition_penalty,
int no_repeat_ngram_size,
cudaStream_t stream);
struct HypothesisScore {
const int32_t* hypothesis;
int hypothesis_length;
float score;
};
struct BeamHypotheses {
HypothesisScore* beams_; // Beam width sized array of hypotheses, sorted by highest scoring
int beams_count_;
int beams_used_; // Number of elements used in beams_
float length_penalty_;
bool done_;
// Add a new hypothesis
__device__ void Add(const int32_t* hypothesis, int hypothesis_length, float sum_logprobs);
// Return true if this beats the worst score in the hypothesis
__device__ bool CanImprove(float best_sum_logprobs, int current_length) const;
// Output results
__device__ void Output(int top_k, // number of sequences to return
int max_length, // max sequence length
int pad_token_id, // pad token
int32_t* sequences, // buffer with pad token, shape (num_return_sequences, max_length)
float* sequences_scores); // buffer for sequence scores, with shape (num_return_sequences)
};
struct BeamScorerState {
int batch_size_;
int num_beams_;
int max_length_;
int num_return_sequences_;
int pad_token_id_;
int eos_token_id_;
bool early_stopping_;
int not_done_count_; // When zero, every batch entry is done (starts at batch_size_)
int hypothesis_buffer_used_; // Offset of available buffer, or length of used buffer.
};
void LaunchInitializeBeamHypotheses(gsl::span<BeamHypotheses> beam_hyps, float length_penalty, gsl::span<HypothesisScore> beams, int num_beams, cudaStream_t stream);
void LaunchBeamSearchScorer_Process(BeamScorerState& state_cpu,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps_,
gsl::span<float> next_beam_scores_,
gsl::span<int32_t> next_beam_tokens_,
gsl::span<int32_t> next_beam_indices_,
gsl::span<int32_t> hypothesis_buffer_,
gsl::span<const float> next_scores,
gsl::span<const int32_t> next_tokens,
gsl::span<const int32_t> next_indices,
cudaStream_t stream);
void LaunchBeamSearchScorer_AppendNextTokenToSequences(BeamScorerState& state_cpu,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
gsl::span<int32_t> next_sequences,
int sequence_length,
gsl::span<int32_t> next_beam_tokens,
gsl::span<int32_t> next_beam_indices,
cudaStream_t stream);
void LaunchBeamSearchScorer_Finalize(int batch_size,
BeamScorerState& state,
gsl::span<const int32_t> sequences,
int sequence_length,
gsl::span<BeamHypotheses> beam_hyps_,
gsl::span<const float> final_beam_scores,
gsl::span<int32_t> output,
gsl::span<float> sequence_scores,
cudaStream_t stream);
void LaunchNextTokenKernel(const int64_t* next_token_indices,
int32_t* next_indices,
int32_t* next_tokens,

260
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
View file

@ -327,7 +327,6 @@ void InitGreedyState(transformers::IGreedySearchState<T>* greedy_state,
template <typename T>
Status ProcessLogits(const OrtValue& logits, // logits output of subgraph
transformers::IBeamSearchState<T>* beam_state, // state
transformers::IBeamSearchCpuState* cpu_state, // state in CPU
transformers::ISequences* sequences, // sequences
AllocatorPtr& allocator, // default allocator
onnxruntime::concurrency::ThreadPool* thread_pool, // thread pool (for CPU only)
@ -425,20 +424,6 @@ Status ProcessLogits(const OrtValue& logits, //
dumper->Print("next_token_scores after softmax", next_token_scores.data(), batch_size, num_beams, vocab_size);
#endif
// Sequences generated by beam scorer is currently stored in CPU.
// Copy sequences to device only when repetition penalty or no repeat ngram is used in kernel
BufferUniquePtr sequences_buffer;
int current_sequence_length = sequences->GetSequenceLength();
bool run_ngram = parameters->no_repeat_ngram_size > 0 && current_sequence_length >= parameters->no_repeat_ngram_size;
if (parameters->repetition_penalty != 1.0f || run_ngram) {
size_t bytes = SafeInt<size_t>(sizeof(int32_t)) * batch_beam_size * parameters->max_length;
void* data = allocator->Alloc(bytes);
BufferUniquePtr temp_buffer(data, BufferDeleter(allocator));
sequences_buffer = std::move(temp_buffer);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(sequences_buffer.get(), sequences->GetSequence(0).data(), bytes,
cudaMemcpyHostToDevice, cuda_stream));
}
cuda::LaunchLogitsProcessKernel<float>(
next_token_scores.data(),
parameters->vocab_mask.data(),
@ -450,10 +435,10 @@ Status ProcessLogits(const OrtValue& logits, //
parameters->num_beams,
vocab_size,
vocab_size,
(parameters->min_length > 0 && current_sequence_length < parameters->min_length) ? parameters->eos_token_id : -1,
reinterpret_cast<int32_t*>(sequences_buffer.get()),
(parameters->min_length > 0 && sequences->GetSequenceLength() < parameters->min_length) ? parameters->eos_token_id : -1,
sequences->GetCurrentDeviceSequences().data(),
parameters->max_length,
current_sequence_length,
sequences->GetSequenceLength(),
parameters->repetition_penalty,
parameters->no_repeat_ngram_size,
cuda_stream);
@ -509,13 +494,6 @@ Status ProcessLogits(const OrtValue& logits, //
dumper->Print("next_indices before scorer", beam_state->next_indices.data(), batch_size, 2 * num_beams);
dumper->Print("next_scores before scorer", beam_state->next_scores.data(), batch_size, 2 * num_beams);
#endif
// TODO: Remove these kinds of cross-device copies once BeamScorer runs on CUDA
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cpu_state->topk_scores.data(),
beam_state->next_scores.data(),
beam_state->next_scores.size_bytes(),
cudaMemcpyDeviceToHost,
cuda_stream));
} else {
// Apply top-k selection like the following:
// next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
@ -551,63 +529,24 @@ Status ProcessLogits(const OrtValue& logits, //
cuda::LaunchNextTokenKernel(next_token_indices, beam_state->next_indices.data(), beam_state->next_tokens.data(),
batch_size, top_k, vocab_size, cuda_stream);
const float* data = topk_scores->Data<float>();
#ifdef DEBUG_GENERATION
dumper->Print("next_scores before scorer", data, batch_size, top_k);
dumper->Print("next_scores before scorer", topk_scores->Data<float>(), batch_size, top_k);
dumper->Print("next_tokens before scorer", beam_state->next_tokens.data(), batch_size, top_k);
dumper->Print("next_indices before scorer", beam_state->next_indices.data(), batch_size, top_k);
#endif
// TODO: Remove these kinds of cross-device copies once BeamScorer runs on CUDA
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cpu_state->topk_scores.data(),
data,
topk_scores->SizeInBytes(),
cudaMemcpyDeviceToHost,
cuda_stream));
}
// TODO: Remove these kinds of cross-device copies once BeamScorer runs on CUDA
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cpu_state->topk_tokens.data(),
beam_state->next_tokens.data(),
beam_state->next_tokens.size_bytes(),
cudaMemcpyDeviceToHost,
cuda_stream));
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cpu_state->topk_indices.data(),
beam_state->next_indices.data(),
beam_state->next_indices.size_bytes(),
cudaMemcpyDeviceToHost,
cuda_stream));
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
// gsl::span doesn't convert from non const to const, so all we're doing here is making each const.
gsl::span<const float> next_scores(beam_state->next_scores.data(), beam_state->next_scores.size());
gsl::span<const int32_t> next_tokens(beam_state->next_tokens.data(), beam_state->next_tokens.size());
gsl::span<const int32_t> next_indices(beam_state->next_indices.data(), beam_state->next_indices.size());
gsl::span<const float> next_scores(cpu_state->topk_scores.data(), beam_state->next_scores.size());
gsl::span<const int32_t> next_tokens(cpu_state->topk_tokens.data(), beam_state->next_tokens.size());
gsl::span<const int32_t> next_indices(cpu_state->topk_indices.data(), beam_state->next_indices.size());
// TODO: Implement BeamScorer on CUDA
beam_scorer->Process(
*sequences,
next_scores,
next_tokens,
next_indices);
// TODO: This is a temporary work-around as BeamScorer currently only runs on CPU.
// We can remove these kinds of work-arounds once BeamScorer runs on CUDA eventually.
auto chosen_indices = beam_scorer->GetNextIndices();
auto beam_state_chosen_indices = beam_state->chosen_indices;
if (!beam_state_chosen_indices.empty()) {
// If we have allocated `chosen_indices` in beam_state, it means that we
// will be needing the chosen indices from BeamScorer as we are using
// DecoderMaskedSelfAttention, so copy it over.
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(beam_state_chosen_indices.data(),
chosen_indices.data(),
chosen_indices.size_bytes(),
cudaMemcpyHostToDevice,
cuda_stream));
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
}
#ifdef ENABLE_NVTX_PROFILE
processLogitsRange.End();
#endif
@ -615,6 +554,167 @@ Status ProcessLogits(const OrtValue& logits, //
return Status::OK();
}
struct CudaBeamSearchScorer : transformers::IBeamScorer {
CudaBeamSearchScorer(const transformers::IGenerationParameters& parameters,
AllocatorPtr& allocator, AllocatorPtr& allocator_cpu,
Stream* stream);
void Process(transformers::ISequences& sequences,
gsl::span<const float>& next_scores,
gsl::span<const int32_t>& next_tokens,
gsl::span<const int32_t>& next_indices) override;
void Finalize(transformers::ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) override;
bool IsDone() const override { return false; } // For CUDA we speculatively run the next step while we wait for the GPU to report status. We use 'IsDoneLater()' for this
bool IsDoneLater() const override;
gsl::span<float> GetNextScores() override { return next_beam_scores_; }
gsl::span<int32_t> GetNextTokens() override { return next_beam_tokens_; }
gsl::span<int32_t> GetNextIndicesCPU() override {
CUDA_CALL_THROW(cudaMemcpyAsync(next_beam_indices_cpu_.data(), next_beam_indices_.data(), next_beam_indices_.size_bytes(), cudaMemcpyDeviceToHost, stream_));
CUDA_CALL_THROW(cudaStreamSynchronize(stream_));
return next_beam_indices_cpu_;
}
gsl::span<int32_t> GetNextIndicesGPU() override { return next_beam_indices_; }
private:
mutable cuda::AutoDestoryCudaEvent event_process_complete_;
IAllocatorUniquePtr<cuda::BeamScorerState> state_cpu_;
IAllocatorUniquePtr<cuda::BeamScorerState> state_gpu_;
cudaStream_t stream_;
IAllocatorUniquePtr<float> next_beam_scores_ptr_;
gsl::span<float> next_beam_scores_;
IAllocatorUniquePtr<int32_t> next_beam_tokens_ptr_;
gsl::span<int32_t> next_beam_tokens_;
IAllocatorUniquePtr<int32_t> next_beam_indices_ptr_;
gsl::span<int32_t> next_beam_indices_;
IAllocatorUniquePtr<int32_t> next_beam_indices_cpu_ptr_;
gsl::span<int32_t> next_beam_indices_cpu_;
IAllocatorUniquePtr<int32_t> hypothesis_buffer_ptr_; // Allocated buffer to hold all hypotheses
gsl::span<int32_t> hypothesis_buffer_; // Span of the allocated buffer
size_t hypothesis_buffer_used_{}; // Offset of available buffer, or length of used buffer.
IAllocatorUniquePtr<cuda::HypothesisScore> hypothesis_scores_ptr_; // num_beams_ * batch_size_, divided into num_beams_ chunks per BeamHypothesis in beam_hyps_
IAllocatorUniquePtr<cuda::BeamHypotheses> beam_hyps_ptr_;
gsl::span<cuda::BeamHypotheses> beam_hyps_; // Shape is batch_size_
};
template <typename TAlloc>
gsl::span<TAlloc> Allocate(std::shared_ptr<IAllocator> allocator,
size_t size,
IAllocatorUniquePtr<TAlloc>& unique_ptr) {
unique_ptr = IAllocator::MakeUniquePtr<TAlloc>(std::move(allocator), size);
return gsl::make_span(unique_ptr.get(), size);
}
template <typename T>
IAllocatorUniquePtr<T> AllocateCPUPinned() {
T* p;
CUDA_CALL_THROW(cudaMallocHost(&p, sizeof(T)));
return IAllocatorUniquePtr<T>{p, [](cuda::BeamScorerState* p) { ORT_IGNORE_RETURN_VALUE(cudaFreeHost(p)); }};
}
CudaBeamSearchScorer::CudaBeamSearchScorer(const transformers::IGenerationParameters& parameters,
AllocatorPtr& allocator, AllocatorPtr& allocator_cpu, Stream* stream)
: stream_{stream ? reinterpret_cast<cudaStream_t>(stream->GetHandle()) : nullptr} {
CUDA_CALL_THROW(cudaEventCreate(&event_process_complete_.Get()));
state_cpu_ = AllocateCPUPinned<cuda::BeamScorerState>();
state_cpu_->batch_size_ = static_cast<size_t>(parameters.batch_size);
state_cpu_->num_beams_ = static_cast<size_t>(parameters.num_beams);
state_cpu_->max_length_ = static_cast<size_t>(parameters.max_length);
state_cpu_->num_return_sequences_ = static_cast<size_t>(parameters.num_return_sequences);
state_cpu_->pad_token_id_ = parameters.pad_token_id;
state_cpu_->eos_token_id_ = parameters.eos_token_id;
state_cpu_->early_stopping_ = parameters.early_stopping;
state_cpu_->not_done_count_ = parameters.batch_size;
state_cpu_->hypothesis_buffer_used_ = 0;
state_gpu_ = IAllocator::MakeUniquePtr<cuda::BeamScorerState>(allocator, 1);
CUDA_CALL_THROW(cudaMemcpyAsync(state_gpu_.get(), state_cpu_.get(), sizeof(cuda::BeamScorerState), ::cudaMemcpyHostToDevice, stream_));
size_t batch_beam_size = state_cpu_->batch_size_ * state_cpu_->num_beams_;
auto beams = Allocate<cuda::HypothesisScore>(allocator, batch_beam_size, hypothesis_scores_ptr_);
beam_hyps_ = Allocate<cuda::BeamHypotheses>(allocator, state_cpu_->batch_size_, beam_hyps_ptr_);
cuda::LaunchInitializeBeamHypotheses(beam_hyps_, parameters.length_penalty, beams, parameters.num_beams, stream_);
next_beam_scores_ = Allocate<float>(allocator, batch_beam_size, next_beam_scores_ptr_);
next_beam_tokens_ = Allocate<int32_t>(allocator, batch_beam_size, next_beam_tokens_ptr_);
next_beam_indices_ = Allocate<int32_t>(allocator, batch_beam_size, next_beam_indices_ptr_);
next_beam_indices_cpu_ = Allocate<int32_t>(allocator_cpu, batch_beam_size, next_beam_indices_cpu_ptr_);
// Space to store intermediate sequence with length sequence_length, sequence_length + 1, ..., max_sequence_length.
size_t per_beam = (SafeInt<size_t>(state_cpu_->max_length_) * (state_cpu_->max_length_ + 1) - (parameters.sequence_length - 1) * parameters.sequence_length) / 2;
hypothesis_buffer_ = Allocate<int32_t>(allocator, batch_beam_size * per_beam, hypothesis_buffer_ptr_);
}
void CudaBeamSearchScorer::Process(transformers::ISequences& sequences,
gsl::span<const float>& next_scores,
gsl::span<const int32_t>& next_tokens,
gsl::span<const int32_t>& next_indices) {
cuda::LaunchBeamSearchScorer_Process(*state_cpu_,
*state_gpu_,
sequences.GetCurrentDeviceSequences(),
sequences.GetSequenceLength(),
beam_hyps_,
next_beam_scores_,
next_beam_tokens_,
next_beam_indices_,
hypothesis_buffer_,
next_scores,
next_tokens,
next_indices,
stream_);
CUDA_CALL_THROW(cudaEventRecord(event_process_complete_.Get(), stream_));
cuda::LaunchBeamSearchScorer_AppendNextTokenToSequences(*state_cpu_,
*state_gpu_,
sequences.GetCurrentDeviceSequences(),
sequences.GetNextDeviceSequences(),
sequences.GetSequenceLength(),
next_beam_tokens_,
next_beam_indices_,
stream_);
}
bool CudaBeamSearchScorer::IsDoneLater() const {
CUDA_CALL_THROW(cudaEventSynchronize(event_process_complete_.Get()));
return state_cpu_->not_done_count_ == 0;
}
void CudaBeamSearchScorer::Finalize(transformers::ISequences& sequences,
gsl::span<const float>& final_beam_scores,
Tensor* output_sequences,
Tensor* output_sequence_scores) {
ORT_ENFORCE(output_sequences != nullptr);
// Word IDs of each sequence, with shape (batch_size * num_return_sequences, max_sequence_length).
gsl::span<int32_t> output{output_sequences->MutableData<int32_t>(), static_cast<size_t>(output_sequences->Shape().Size())};
// Score of each sequence, with shape (batch_size * num_return_sequences).
gsl::span<float> sequence_scores;
if (output_sequence_scores) {
sequence_scores = gsl::span<float>{output_sequence_scores->MutableData<float>(), static_cast<size_t>(output_sequence_scores->Shape().Size())};
}
cuda::LaunchBeamSearchScorer_Finalize(state_cpu_->batch_size_, *state_gpu_, sequences.GetCurrentDeviceSequences(), sequences.GetSequenceLength(), beam_hyps_, final_beam_scores, output, sequence_scores, stream_);
}
std::unique_ptr<transformers::IBeamScorer> CreateBeamScorer(const transformers::IGenerationParameters& parameters,
AllocatorPtr& allocator, AllocatorPtr& allocator_cpu, Stream* stream) {
return std::make_unique<CudaBeamSearchScorer>(parameters, allocator, allocator_cpu, stream);
}
template <typename T>
Status GreedySearchProcessLogits(
const OrtValue& logits, // logits output of subgraph
@ -819,7 +919,8 @@ Status DeviceCopy(gsl::span<T> target, gsl::span<const T> source, Stream* ort_st
} else {
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(target.data(), source.data(), source.size_bytes(),
static_cast<cudaMemcpyKind>(copyDirection), cuda_stream));
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
if (copyDirection != DeviceCopyDirection::deviceToDevice)
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
}
return Status::OK();
}
@ -850,6 +951,7 @@ Status PickGptPastState(const std::vector<OrtValue>& last_outputs,
gsl::span<T> past_span = gsl::make_span<T>(past.GetMutable<Tensor>()->MutableData<T>(), past_shape.Size());
gsl::span<const T> present_span = gsl::make_span<const T>(present.Get<Tensor>().Data<T>(), past_shape.Size());
for (size_t j = 0; j < beam_indices.size(); j++) {
int32_t beam_index = beam_indices[j];
gsl::span<const T> present_key = present_span.subspan(beam_index * block_size_per_beam, block_size_per_beam);
@ -935,17 +1037,17 @@ Status UpdateGptFeeds(
// Update input_ids with next tokens.
int batch_beam_size = static_cast<int>(beam_next_tokens.size());
int64_t dims[] = {batch_beam_size, 1};
TensorShape input_ids_shape(&dims[0], 2);
TensorShape input_ids_shape{batch_beam_size, 1};
auto element_type = DataTypeImpl::GetType<int32_t>();
OrtValue input_ids;
Tensor::InitOrtValue(element_type, input_ids_shape, allocator, input_ids);
int32_t* input_ids_data = input_ids.GetMutable<Tensor>()->MutableData<int32_t>();
cudaStream_t cuda_stream = ort_stream ? static_cast<cudaStream_t>(ort_stream->GetHandle()) : nullptr;
// TODO(hasesh): When BeamScorer is implemented on CUDA, figure out a way to avoid this cross-device copy
// num_beams == 1 using cudaMemcpyHostToDevice is because GreedySearch still uses CPU, BeamSearch is fully GPU
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(input_ids_data, beam_next_tokens.data(), beam_next_tokens.size_bytes(),
cudaMemcpyHostToDevice, cuda_stream));
num_beams == 1 ? cudaMemcpyHostToDevice : cudaMemcpyDeviceToDevice, cuda_stream));
next_inputs[0] = input_ids;
// Update position IDs
@ -955,8 +1057,7 @@ Status UpdateGptFeeds(
// Update attention mask
const OrtValue& old_mask = next_inputs[2];
const int32_t* old_mask_data = old_mask.Get<Tensor>().Data<int32_t>();
int64_t mask_dims[] = {batch_beam_size, current_length};
TensorShape mask_shape(&mask_dims[0], 2);
TensorShape mask_shape{batch_beam_size, current_length};
OrtValue attention_mask;
auto mask_type = DataTypeImpl::GetType<int32_t>();
Tensor::InitOrtValue(mask_type, mask_shape, allocator, attention_mask);
@ -1015,12 +1116,11 @@ Status UpdateGptFeeds(
ORT_RETURN_IF_ERROR(PickGptPastState<T>(last_outputs, next_inputs, beam_indices_cpu, allocator,
gpt_subgraph_first_past_input_idx,
gpt_subgraph_first_present_output_idx, ort_stream));
// Make sure data is ready before next subgraph execution.
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
}
}
// Make sure data is ready before next subgraph execution.
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream));
#ifdef ENABLE_NVTX_PROFILE
updateFeedsRange.End();
#endif
@ -1246,7 +1346,6 @@ template void InitGreedyState<float>(
template Status ProcessLogits<float>(
const OrtValue& logits,
transformers::IBeamSearchState<float>* beam_state,
transformers::IBeamSearchCpuState* cpu_state,
transformers::ISequences* sequences,
AllocatorPtr& allocator,
onnxruntime::concurrency::ThreadPool* thread_pool,
@ -1318,7 +1417,6 @@ template void InitGreedyState<MLFloat16>(
template Status ProcessLogits<MLFloat16>(
const OrtValue& logits,
transformers::IBeamSearchState<MLFloat16>* beam_state,
transformers::IBeamSearchCpuState* cpu_state,
transformers::ISequences* sequences,
AllocatorPtr& allocator,
onnxruntime::concurrency::ThreadPool* thread_pool,

3
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h
View file

@ -54,6 +54,8 @@ void InitBeamState(transformers::IBeamSearchState<T>* beam_state,
int num_beams,
Stream* ort_stream);
std::unique_ptr<transformers::IBeamScorer> CreateBeamScorer(const transformers::IGenerationParameters& parameters, AllocatorPtr& allocator, AllocatorPtr& allocator_cpu, Stream* ort_stream);
template <typename T>
void InitGreedyState(transformers::IGreedySearchState<T>* greedy_state,
gsl::span<int32_t>& sequence_lengths,
@ -62,7 +64,6 @@ void InitGreedyState(transformers::IGreedySearchState<T>* greedy_state,
template <typename T>
Status ProcessLogits(const OrtValue& logits, // logits output of subgraph
transformers::IBeamSearchState<T>* beam_state, // state
transformers::IBeamSearchCpuState* cpu_state, // state in CPU
transformers::ISequences* sequences, // sequences
AllocatorPtr& allocator, // default allocator
onnxruntime::concurrency::ThreadPool* thread_pool, // thread pool (for CPU only)