diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h index e370d7bcb9..6b35db806f 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h @@ -225,7 +225,10 @@ Status BeamSearchT5::Execute(const FeedsFetchesManager& encoder_feeds_fetches this->expand_buffer_float_func_, this->expand_buffer_float16_func_, parameters->num_beams, - this->cuda_stream_)); + this->cuda_stream_, + decoder_subgraph_.UseSequenceAsInputIds(), + current_length, + cpu_state.sequences)); } // TODO(tianleiwu): allocate fetches. use ping-pong buffers for past state. diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc index 9c48766121..91d3441d88 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc @@ -10,6 +10,7 @@ #include "contrib_ops/cpu/transformers/subgraph_t5_decoder.h" #include "contrib_ops/cpu/transformers/dump_tensor.h" #include "contrib_ops/cpu/transformers/generation_device_helper.h" +#include "contrib_ops/cpu/transformers/sequences.h" namespace onnxruntime { namespace contrib { @@ -132,7 +133,10 @@ Status T5DecoderSubgraph::CreateInitialFeeds( const GenerationDeviceHelper::ExpandBufferFunc& expand_buffer_float_func, const GenerationDeviceHelper::ExpandBufferFunc& expand_buffer_float16_func, int num_beam, - void* stream) { + void* stream, + bool use_sequence_as_input_ids, + int cur_len, + transformers::Sequences& sequences) { ORT_ENFORCE(session_state_ != nullptr, "Setup must be called before CreateInitialFeeds"); // Allocate subgraph inputs from same device as inputs of encoder subgraph. @@ -140,15 +144,28 @@ Status T5DecoderSubgraph::CreateInitialFeeds( // Copy beam next tokens in CPU to input_ids in provider device (CPU for CPU EP, or GPU for CUDA EP). int batch_beam_size = static_cast(beam_next_tokens.size()); - int64_t dims[] = {batch_beam_size, 1}; + int sequence_length = !use_sequence_as_input_ids ? 1 : cur_len; + int64_t dims[] = {batch_beam_size, sequence_length}; TensorShape input_ids_shape(&dims[0], 2); OrtValue input_ids; Tensor::InitOrtValue(DataTypeImpl::GetType(), input_ids_shape, allocator, input_ids); - ORT_RETURN_IF_ERROR(device_copy_int32_func( - input_ids.GetMutable()->MutableDataAsSpan(), - beam_next_tokens, - stream, - DeviceCopyDirection::hostToDevice)); + int32_t* input_ids_data = input_ids.GetMutable()->MutableData(); + + if (!use_sequence_as_input_ids_){ + ORT_RETURN_IF_ERROR(device_copy_int32_func( + input_ids.GetMutable()->MutableDataAsSpan(), + beam_next_tokens, + stream, + DeviceCopyDirection::hostToDevice)); + }else{ + for (int i = 0; i < batch_beam_size; i++) { + gsl::span sequence = sequences.GetSequence(i); + const int32_t* sequence_data = sequence.data(); + for (int j = 0; j < cur_len; j++) { + input_ids_data[i * cur_len + j] = sequence_data[j]; + } + } + } // The ordering is the same as used in Setup. decoder_feeds.reserve(static_cast(num_subgraph_inputs) + static_cast(num_implicit_inputs)); diff --git a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h index 66437ee98f..b9f00d72f2 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h +++ b/onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h @@ -4,6 +4,7 @@ #pragma once #include "contrib_ops/cpu/transformers/subgraph_base.h" +#include "contrib_ops/cpu/transformers/sequences.h" namespace onnxruntime { namespace contrib { @@ -33,7 +34,10 @@ class T5DecoderSubgraph : public Subgraph { const GenerationDeviceHelper::ExpandBufferFunc& expand_buffer_float_func, const GenerationDeviceHelper::ExpandBufferFunc& expand_buffer_float16_func, int num_beam, - void* stream); + void* stream, + bool use_sequence_as_input_ids, + int cur_len, + transformers::Sequences& sequences); Status Validate(const std::vector& subgraph_inputs, const std::vector& subgraph_outputs) override;