new design of t5 onnx

Author: tianleiwu

onnxruntime/contrib_ops/cpu/transformers/beam_search.cc

@@ -139,13 +139,19 @@ Status BeamSearch::SetupSubgraphExecutionInfo(const SessionState& session_state,
       ORT_RETURN_IF_ERROR(t5_encoder_subgraph_->Setup(session_state, subgraph_session_state));
       encoder_feeds_fetches_manager_ = t5_encoder_subgraph_->GetFeedsFetchesManager();
 
-      if (parameters_->decoder_start_token_id < 0) {
-        ORT_RETURN_IF(t5_encoder_subgraph_->num_subgraph_inputs != 2,
-                      "Encoder subgraph shall have 2 inputs when decoder_start_token_id attribute is empty");
+      if (t5_encoder_subgraph_->HasLogitsOutput()) {
+        // New format requires start token id.
+        ORT_ENFORCE(parameters_->decoder_start_token_id >= 0);
       } else {
-        ORT_RETURN_IF(t5_encoder_subgraph_->num_subgraph_inputs != 3,
-                      "Encoder subgraph shall have 3 inputs when decoder_start_token_id attribute is available");
+        if (parameters_->decoder_start_token_id < 0) {
+          ORT_RETURN_IF(t5_encoder_subgraph_->num_subgraph_inputs != 2,
+                        "Encoder subgraph shall have 2 inputs when decoder_start_token_id attribute is empty");
+        } else {
+          ORT_RETURN_IF(t5_encoder_subgraph_->num_subgraph_inputs != 3,
+                        "Encoder subgraph shall have 3 inputs when decoder_start_token_id attribute is available");
+        }
       }
+
     } else if (attribute_name == "decoder") {
       ORT_ENFORCE(t5_decoder_subgraph_ == nullptr,
                   "SetupSubgraphExecutionInfo should only be called once for each subgraph.");

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h

@@ -51,7 +51,13 @@ class BeamSearchT5 : public BeamSearchBase<T> {
         expand_buffer_int32_func_(expand_buffer_int32_func),
         expand_buffer_float_func_(expand_buffer_float_func),
         expand_buffer_float16_func_(expand_buffer_float16_func),
-        create_beam_scorer_func_(create_beam_scorer_func) {}
+        create_beam_scorer_func_(create_beam_scorer_func) {
+    // When decoder uses encoder_hidden_state, make sure the encoder outputs it.
+    if (decoder_subgraph_.UseEncoderHiddenState()) {
+      ORT_ENFORCE(encoder_subgraph_.subgraph_output_names[1] == "encoder_hidden_states");
+    }
+    ORT_ENFORCE(encoder_subgraph_.num_layers == decoder_subgraph_.num_layers);
+  }
 
 #ifdef USE_CUDA
   Status InitializeCuda(
@@ -160,7 +166,7 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
       this->create_encoder_inputs_func_,
       this->add_to_feeds_func_,
       buffer,
-      decoder_input_ids,
+      decoder_input_ids,  // new format does not use decoder_input_ids in encoder, it is still initialized here when decoder_start_token_id >= 0.
       this->ort_stream_));
 
 #ifdef DEBUG_NODE_INPUTS_OUTPUTS
@@ -233,15 +239,20 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
 
   std::vector<OrtValue> decoder_fetches;
 
-  if (current_length + 1 < parameters->max_length) {
+  // When encoder outputs logits (in old format), we need get the next token from logits.
+  if (current_length + 1 < parameters->max_length && encoder_subgraph_.HasLogitsOutput()) {
     ++iteration_counter;
-    ORT_RETURN_IF_ERROR(this->GenerateNextToken(encoder_fetches[0],
+    const OrtValue& logits = encoder_fetches[0];
+    ORT_RETURN_IF_ERROR(this->GenerateNextToken(logits,
                                                 beam_next_tokens,
                                                 beam_state,
                                                 cpu_state,
                                                 iteration_counter));
     ++current_length;  // Increase sequence length after a new token is generated.
+  }
 
+  // Generate inputs for next decoder subgraph call.
+  if (current_length < parameters->max_length) {
     ORT_RETURN_IF_ERROR(decoder_subgraph_.CreateInitialFeeds(this->cpu_allocator_,
                                                              ReinterpretAsSpan<const int32_t>(beam_next_tokens),
                                                              this->implicit_inputs_,
@@ -262,6 +273,7 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
                                                              this->cuda_device_prop_ != nullptr));
 
     if (decoder_subgraph_.past_present_share_buffer_) {
+      // Configure buffer sharing of past and present kv cache.
       decoder_fetches.reserve(static_cast<size_t>(decoder_subgraph_.GetFirstPresentOutputIndex()) +
                               2 * static_cast<size_t>(decoder_subgraph_.num_layers));
       decoder_fetches.resize(decoder_subgraph_.GetFirstPresentOutputIndex(), OrtValue());

onnxruntime/contrib_ops/cpu/transformers/subgraph_base.cc

@@ -36,12 +36,9 @@ Subgraph::Subgraph(
   auto& subgraph_inputs = subgraph.GetInputs();
   auto& subgraph_outputs = subgraph.GetOutputs();
 
-  // inputs: input_ids, position_ids, attention_mask, past_0, past_1, ...
-  // outputs: logits, present_0, present_1, ...
   num_subgraph_inputs = static_cast<int>(subgraph_inputs.size());
   num_subgraph_outputs = static_cast<int>(subgraph_outputs.size());
 
-  // CheckSubgraph will verify inputs and outputs later.
   subgraph_input_names.reserve(num_subgraph_inputs);
   for (int i = 0; i < num_subgraph_inputs; ++i) {
     subgraph_input_names.push_back(subgraph_inputs[i]->Name());
@@ -68,10 +65,9 @@ Status Subgraph::Setup(const SessionState& session_state,
   InlinedVector<std::string_view> feed_names;
   feed_names.reserve(static_cast<size_t>(num_subgraph_inputs) + static_cast<size_t>(num_implicit_inputs));
 
-  // Use the first output (logits) to find device location.
+  // Use the first output to find device location.
   const OrtDevice& default_location = utils::FindDeviceForValue(subgraph_session_state, subgraph_output_names[0]);
 
-  // The position_ids, attention_mask, past_0, ... are created by this operator so the name doesn't matter.
   feed_names.insert(feed_names.end(), subgraph_input_names.begin(), subgraph_input_names.end());
 
   const auto& subgraph_map = subgraph_session_state.GetOrtValueNameIdxMap();
@@ -174,13 +170,15 @@ Status Subgraph::GetParameters(const ONNX_NAMESPACE::TensorShapeProto* past_shap
   }
 
   // Logits shape is like (batch_size, seq_len, vocabulary_size)
-  ORT_RETURN_IF(logits_shape->dim_size() != 3,
-                "subgraph logits output is expected to have 3 dimension, got ", logits_shape->dim_size());
+  if (logits_shape != nullptr) {
+    ORT_RETURN_IF(logits_shape->dim_size() != 3,
+                  "subgraph logits output is expected to have 3 dimension, got ", logits_shape->dim_size());
 
-  ORT_RETURN_IF(!logits_shape->dim(2).has_dim_value() || logits_shape->dim(2).dim_value() <= 0,
-                "subgraph past state dimension 2 shall have a positive value for vocabulary size");
+    ORT_RETURN_IF(!logits_shape->dim(2).has_dim_value() || logits_shape->dim(2).dim_value() <= 0,
+                  "subgraph past state dimension 2 shall have a positive value for vocabulary size");
 
-  this->vocab_size = static_cast<int>(logits_shape->dim(2).dim_value());
+    this->vocab_size = static_cast<int>(logits_shape->dim(2).dim_value());
+  }
 
   return Status::OK();
 }

onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc

@@ -141,14 +141,23 @@ Status T5DecoderSubgraph::Validate(const std::vector<const NodeArg*>& subgraph_i
 }
 
 // Create inputs for decoder from the following data sources:
-// encoder feeds: encoder_input_ids, encoder_attention_mask, decoder_input_ids (with start tokens)
-// encoder fetches: logits,
-//                  encoder_hidden_states,
-//                  present_key_self_0, present_value_self_0, ..., present_key_cross_0, present_value_cross_0, ...
-// decoder_feeds: input_ids,
-//                encoder_attention_mask,
-//                encoder_hidden_states,
-//                present_key_self_0, present_value_self_0, ..., present_key_cross_0, present_value_cross_0, ...
+// New format:
+//   encoder feeds: encoder_input_ids, encoder_attention_mask
+//   encoder fetches: present_key_cross_0, present_value_cross_0, ...
+//   decoder_feeds: input_ids, encoder_attention_mask,
+//                  present_key_self_0, present_value_self_0, ...,
+//                  present_key_cross_0, present_value_cross_0, ...
+
+// Old format:
+//   encoder feeds: encoder_input_ids, encoder_attention_mask, decoder_input_ids (with start tokens)
+//   encoder fetches: logits, encoder_hidden_states,
+//                    present_key_self_0, present_value_self_0, ...,
+//                    present_key_cross_0, present_value_cross_0, ...
+//   decoder_feeds: input_ids, encoder_input_ids (optiona), encoder_attention_mask, encoder_hidden_states,
+//                  present_key_self_0, present_value_self_0, ...,
+//                  present_key_cross_0, present_value_cross_0, ...
+//                  past_seq_len (optional), num_beams (optional), cache_indirection (optional)
+
 Status T5DecoderSubgraph::CreateInitialFeeds(
     AllocatorPtr cpu_allocator,
     gsl::span<const int32_t> beam_next_tokens,
@@ -173,33 +182,30 @@ Status T5DecoderSubgraph::CreateInitialFeeds(
   // Allocate subgraph inputs from same device as inputs of encoder subgraph.
   AllocatorPtr allocator = session_state_->GetAllocator(encoder_feeds[0].Get<Tensor>().Location());
 
+  int batch_beam_size = static_cast<int>(encoder_fetches[0].Get<Tensor>().Shape()[0]) * num_beam;
+
   // 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<int>(beam_next_tokens.size());
   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<int32_t>(), input_ids_shape, allocator, input_ids);
-  int32_t* input_ids_data = input_ids.GetMutable<Tensor>()->MutableData<int32_t>();
-  AllocatorPtr buffer_allocator = std::make_shared<onnxruntime::CPUAllocator>();
-  size_t total_size = static_cast<size_t>(cur_len) * static_cast<size_t>(batch_beam_size);
-  size_t total_size_bytes = total_size * sizeof(int);
-  auto seq_copy = IAllocator::MakeUniquePtr<int>(buffer_allocator, total_size_bytes, false, stream);
-  int* seq_copy_ptr = seq_copy.get();
-
-  if (!use_sequence_as_input_ids_) {
+
+  // Prepare data for input_ids.
+  if (!use_sequence_as_input_ids_) {  // use next tokens for input_ids. This is for Whisper model.
     ORT_RETURN_IF_ERROR(device_copy_int32_func(
         input_ids.GetMutable<Tensor>()->MutableDataAsSpan<int32_t>(),
         beam_next_tokens,
         stream,
         DeviceCopyDirection::hostToDevice));
   } else {
+    int32_t* input_ids_data = input_ids.GetMutable<Tensor>()->MutableData<int32_t>();
     if (use_cuda) {
       auto sequences_buffer = sequences.GetCurrentDeviceSequences();
       for (int i = 0; i < batch_beam_size; i++) {
-        size_t batch_beam_stride = static_cast<size_t>(i) * static_cast<size_t>(sequences.GetMaxLength());
+        size_t offset = static_cast<size_t>(i) * static_cast<size_t>(sequences.GetMaxLength());
         int seq_size = sequences.GetSequenceLength();
-        gsl::span<const int32_t> sequence = sequences_buffer.subspan(batch_beam_stride, seq_size);
+        gsl::span<const int32_t> sequence = sequences_buffer.subspan(offset, seq_size);
         gsl::span<int> temp_input(input_ids_data + static_cast<ptrdiff_t>(i) * seq_size, seq_size);
         ORT_RETURN_IF_ERROR(device_copy_int32_func(
             temp_input,
@@ -208,6 +214,13 @@ Status T5DecoderSubgraph::CreateInitialFeeds(
             DeviceCopyDirection::deviceToDevice));
       }
     } else {
+      size_t total_size = static_cast<size_t>(cur_len) * static_cast<size_t>(batch_beam_size);
+      size_t total_size_bytes = total_size * sizeof(int);
+      AllocatorPtr buffer_allocator = std::make_shared<onnxruntime::CPUAllocator>();
+      // TODO: not need extra buffer. Copy directly to input_ids_data instead like the user_cuda above.
+      auto seq_copy = IAllocator::MakeUniquePtr<int>(buffer_allocator, total_size_bytes, false, stream);
+      int* seq_copy_ptr = seq_copy.get();
+
       const size_t cur_len_bytes = cur_len * sizeof(int);
       for (int i = 0; i < batch_beam_size; i++) {
         gsl::span<const int32_t> sequence = sequences.GetSequence(i);
@@ -227,9 +240,11 @@ Status T5DecoderSubgraph::CreateInitialFeeds(
 
   // The ordering is the same as used in Setup.
   decoder_feeds.reserve(static_cast<size_t>(num_subgraph_inputs) + static_cast<size_t>(num_implicit_inputs));
+
+  // input 0: input_ids
   decoder_feeds.push_back(input_ids);
 
-  if (has_encoder_input_ids_) {
+  if (has_encoder_input_ids_) {  // encoder_input_ids is optional
     // The encoder_input_ids is copied from the first input of encoder.
     OrtValue expanded_encoder_input_ids;
     ORT_RETURN_IF_ERROR(expand_buffer_int32_func(stream,
@@ -251,70 +266,65 @@ Status T5DecoderSubgraph::CreateInitialFeeds(
                                                expanded_decoder_attention_masks,
                                                false,
                                                0 /*max_sequence_length*/));
-
   decoder_feeds.push_back(expanded_decoder_attention_masks);
 
   if (!past_present_share_buffer_) {
     past_present_share_buffer_max_seq_len = 0;
   }
 
-  // When first_past_input_index_ == 3, the encoder_hidden_states and past states are copied from the second output
-  // of encoder.
-  // When first_past_input_index_ == 2, the past states are copied from the second output of encoder.
-  // TODO - probably more robust to introduce a encoder_out/decoder_in mapping instead of relying on positions.
-  // What happens if encoder_hidden_states is present in the encoder_fetches but not in the decoder_feeds?
-  for (size_t j = static_cast<size_t>(2) - has_hidden_state_; j < encoder_fetches.size(); j++) {
-    if (j == 1) {
-      ORT_RETURN_IF(has_hidden_state_ == false, "Invalid hidden_states expension: has_hidden_state_ == false");
-      OrtValue expanded_hidden_states;
-      if (is_output_float16_) {
-        ORT_RETURN_IF_ERROR(expand_buffer_float16_func(stream,
-                                                       encoder_fetches[j],
-                                                       num_beam,
-                                                       allocator,
-                                                       expanded_hidden_states,
-                                                       false,
-                                                       0 /*max_sequence_length*/));
-      } else {
-        ORT_RETURN_IF_ERROR(expand_buffer_float_func(stream,
-                                                     encoder_fetches[j],
-                                                     num_beam,
-                                                     allocator,
-                                                     expanded_hidden_states,
-                                                     false,
-                                                     0 /*max_sequence_length*/));
-      }
-      decoder_feeds.push_back(expanded_hidden_states);
-    } else {
+// macro to expand encoder outputs and append to decoder feeds.
+#define ADD_DECODER_FEED(encoder_output, is_dynamic_kv_cache)                                                         \
+  OrtValue expanded;                                                                                                  \
+  if (is_output_float16_) {                                                                                           \
+    ORT_RETURN_IF_ERROR(expand_buffer_float16_func(stream, encoder_output, num_beam, allocator, expanded, false,      \
+                                                   is_dynamic_kv_cache ? past_present_share_buffer_max_seq_len : 0)); \
+  } else {                                                                                                            \
+    ORT_RETURN_IF_ERROR(expand_buffer_float_func(stream, encoder_output, num_beam, allocator, expanded, false,        \
+                                                 is_dynamic_kv_cache ? past_present_share_buffer_max_seq_len : 0));   \
+  }                                                                                                                   \
+  decoder_feeds.push_back(expanded);
+
+  // The encoder_hidden_states is copied from the second output of encoder.
+  if (has_hidden_state_) {
+    ADD_DECODER_FEED(encoder_fetches[1], false);
+  }
+
+  // New format of encoder has only cross outputs.
+  bool is_new_format = (static_cast<int>(encoder_fetches.size()) == 2 * num_layers);
+  if (is_new_format) {
+    for (int i = 0; i < 2 * num_layers; i++) {
+      // cross shape is (batch_size, num_heads, encode_sequence_length, head_size)
+      const TensorShape& cross_shape = encoder_fetches[0].Get<Tensor>().Shape();
+      ORT_ENFORCE(cross_shape.NumDimensions() == 4);
+
+      // Shape for kv cache: (batch_size * num_beam, num_heads, max_seq_len, head_size)
+      int64_t cache_dims[4] = {0};
+      cross_shape.CopyDims(cache_dims, cross_shape.NumDimensions());
+      cache_dims[0] *= num_beam;
+      cache_dims[2] = past_present_share_buffer_max_seq_len;
+      TensorShape expanded_shape(&cache_dims[0], cross_shape.NumDimensions());
+
+      MLDataType element_type = encoder_fetches[0].Get<Tensor>().DataType();
+      OrtValue past;
+      Tensor::InitOrtValue(element_type, expanded_shape, allocator, past);
+      decoder_feeds.push_back(past);
+    }
+
+    // Add cross inputs from encoder output.
+    for (size_t j = 0; j < encoder_fetches.size(); j++) {
+      ADD_DECODER_FEED(encoder_fetches[j], false);
+    }
+  } else {
+    for (size_t j = 1 + has_hidden_state_; j < encoder_fetches.size(); j++) {
       // past key/value for cross attention does not need to be initialized with max_seq_len since they are static.
-      bool use_max_seq_len = (j - first_past_input_index_) < 2 * static_cast<size_t>(num_layers);
-
-      OrtValue expanded_cache;
-      if (is_output_float16_) {
-        ORT_RETURN_IF_ERROR(expand_buffer_float16_func(stream,
-                                                       encoder_fetches[j],
-                                                       num_beam,
-                                                       allocator,
-                                                       expanded_cache,
-                                                       false,
-                                                       use_max_seq_len ? past_present_share_buffer_max_seq_len : 0));
-      } else {
-        ORT_RETURN_IF_ERROR(expand_buffer_float_func(stream,
-                                                     encoder_fetches[j],
-                                                     num_beam,
-                                                     allocator,
-                                                     expanded_cache,
-                                                     false,
-                                                     use_max_seq_len ? past_present_share_buffer_max_seq_len : 0));
-      }
-      decoder_feeds.push_back(expanded_cache);
+      bool is_dynamic_kv_cache = (j - first_past_input_index_) < 2 * static_cast<size_t>(num_layers);
+      ADD_DECODER_FEED(encoder_fetches[j], is_dynamic_kv_cache);
     }
   }
 
-  // TODO: This part shares the similar logic with CreateInitialFeeds() in subgraph_gpt.cc. We should refactor it.
   if (past_present_share_buffer_) {
-    // Past sequence length feed
-    ORT_RETURN_IF_ERROR(AppendPastSequenceLength(decoder_feeds, cpu_allocator, 1));
+    // Past sequence length set to 0
+    ORT_RETURN_IF_ERROR(AppendPastSequenceLength(decoder_feeds, cpu_allocator, is_new_format ? 0 : 1));
     // Add beam search specific inputs
     if (need_cache_indir) {
       const int64_t batch_size = static_cast<int64_t>(batch_beam_size / num_beam);

onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.h

@@ -72,6 +72,10 @@ class T5DecoderSubgraph : public Subgraph {
     return use_sequence_as_input_ids_;
   }
 
+  inline bool UseEncoderHiddenState() const {
+    return has_hidden_state_;
+  }
+
  protected:
   int first_past_input_index_;
   int first_present_output_index_;

onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.cc

@@ -15,78 +15,105 @@ namespace transformers {
 
 /* T5 Encoder Subgraph (It also contains decoder initialization where decoder_input_ids are filled with start token ID).
 
-   Inputs:
+  New format:
+    Inputs:
       encoder_input_ids: int32 (B, encode_sequence_length)
       encoder_attention_mask: int32 (B, encode_sequence_length)
-      decoder_input_ids: int32 (B, 1)
 
     Outputs:
-      logits: (B, 1, vocab_size)
-      encoder_hidden_states: (B, encode_sequence_length, encoder_hidden_size)
-
-      present_key_self_0: (B, num_heads, 1, head_size)
-      present_value_self_0: (B, num_heads, 1, head_size)
-      ... (for each self attention layer)
-
       present_key_cross_0: (B, num_heads, encode_sequence_length, head_size)
       present_value_cross_0: (B, num_heads, encode_sequence_length, head_size)
       ... (for each cross attention layer)
 
-    Note:
-      Here, B = batch_size * num_beams since we expand the inputs.
-      Ideally, we could use B=batch_size and expand the outputs with a factor of num_beams.
-      Data type of input or output is float or float16 if not specified.
+  Old format:
+    Inputs:
+        encoder_input_ids: int32 (B, encode_sequence_length)
+        encoder_attention_mask: int32 (B, encode_sequence_length)
+        decoder_input_ids: int32 (B, 1)
+
+    Outputs:
+        logits: (B, 1, vocab_size)
+        encoder_hidden_states: (B, encode_sequence_length, encoder_hidden_size)
+
+        present_key_self_0: (B, num_heads, 1, head_size)
+        present_value_self_0: (B, num_heads, 1, head_size)
+        ... (for each self attention layer)
+
+        present_key_cross_0: (B, num_heads, encode_sequence_length, head_size)
+        present_value_cross_0: (B, num_heads, encode_sequence_length, head_size)
+        ... (for each cross attention layer)
+
+  Note:
+    Here, B = batch_size * num_beams since we expand the inputs.
+    Ideally, we could use B=batch_size and expand the outputs with a factor of num_beams.
+    Data type of input or output is float or float16 if not specified.
 */
 
 Status T5EncoderSubgraph::Validate(const std::vector<const NodeArg*>& subgraph_inputs,
                                    const std::vector<const NodeArg*>& subgraph_outputs) {
-  ORT_RETURN_IF(num_subgraph_inputs != 3, "expect 3 inputs, got:", num_subgraph_inputs);
-
-  ORT_RETURN_IF(num_subgraph_outputs < 6, "expect >=6 outputs, got:", num_subgraph_outputs);
-  ORT_RETURN_IF((static_cast<int>(subgraph_outputs.size()) - first_present_output_index_) % 4 != 0,
-                "number of outputs expected to be 2 + 4 * layers, got:", num_subgraph_outputs);
+  constexpr auto int32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32;
+  constexpr auto float32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT;
+  constexpr auto float16_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16;
 
+  ORT_RETURN_IF(num_subgraph_inputs != 2 && num_subgraph_inputs != 3, "expect 2 or 3 inputs, got:", num_subgraph_inputs);
   ORT_RETURN_IF(subgraph_inputs[0]->Name() != "encoder_input_ids",
                 "encoder subgraph input 0 shall be named as encoder_input_ids, got: ", subgraph_inputs[0]->Name());
   ORT_RETURN_IF(subgraph_inputs[1]->Name() != "encoder_attention_mask",
                 "encoder subgraph input 1 shall be named as encoder_attention_mask, got: ", subgraph_inputs[1]->Name());
-  ORT_RETURN_IF(subgraph_inputs[2]->Name() != "decoder_input_ids",
-                "encoder subgraph input 2 shall be named as decoder_input_ids, got: ", subgraph_inputs[2]->Name());
-
-  ORT_RETURN_IF(subgraph_outputs[0]->Name() != "logits",
-                "encoder subgraph output 0 shall be named as logits, got: ", subgraph_outputs[0]->Name());
-  ORT_RETURN_IF(subgraph_outputs[1]->Name() != "encoder_hidden_states",
-                "encoder subgraph output 1 shall be named encoder_hidden_states, got: ", subgraph_outputs[1]->Name());
-  ORT_RETURN_IF(subgraph_outputs[2]->Name() != "present_key_self_0",
-                "encoder subgraph output 2 shall be named as present_key_self_0, got: ", subgraph_outputs[2]->Name());
-  ORT_RETURN_IF(subgraph_outputs[3]->Name() != "present_value_self_0",
-                "encoder subgraph output 3 shall be named as present_value_self_0, got: ", subgraph_outputs[3]->Name());
-
-  const ONNX_NAMESPACE::TensorShapeProto* past_shape = subgraph_outputs[2]->Shape();
-  const ONNX_NAMESPACE::TensorShapeProto* logits_shape = subgraph_outputs[0]->Shape();
-
-  // Save parameters related to the subgraph.
-  ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false));
-  num_layers = (static_cast<int>(subgraph_outputs.size()) - first_present_output_index_) / 4;
-
-  constexpr auto int32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT32;
-  constexpr auto float32_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT;
-  constexpr auto float16_type = ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16;
 
   ORT_RETURN_IF(subgraph_inputs[0]->TypeAsProto()->tensor_type().elem_type() != int32_type,
                 "encoder subgraph input 0 (encoder_input_ids) shall have int32 type");
   ORT_RETURN_IF(subgraph_inputs[1]->TypeAsProto()->tensor_type().elem_type() != int32_type,
                 "encoder subgraph input 1 (encoder_attention_mask) shall have int32 type");
-  ORT_RETURN_IF(subgraph_inputs[2]->TypeAsProto()->tensor_type().elem_type() != int32_type,
-                "encoder subgraph input 2 (decoder_input_ids) shall have int32 type");
+
+  if (num_subgraph_inputs == 2) {
+    ORT_RETURN_IF(num_subgraph_outputs < 2 || num_subgraph_outputs % 2 != 0,
+                  "number of outputs expected to be 2 * layers, got:", num_subgraph_outputs);
+
+    ORT_RETURN_IF(subgraph_outputs[0]->Name() != "present_key_cross_0",
+                  "encoder subgraph output 0 shall be named as present_key_cross_0, got: ", subgraph_outputs[0]->Name());
+    ORT_RETURN_IF(subgraph_outputs[1]->Name() != "present_value_cross_0",
+                  "encoder subgraph output 1 shall be named as present_value_cross_0, got: ", subgraph_outputs[1]->Name());
+
+    // Deduce num_heads and head_size parameters from shape of graph outputs
+    const ONNX_NAMESPACE::TensorShapeProto* past_shape = subgraph_outputs[0]->Shape();
+    const ONNX_NAMESPACE::TensorShapeProto* logits_shape = nullptr;
+    ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false));
+
+    num_layers = num_subgraph_outputs / 2;
+  } else {
+    ORT_RETURN_IF(num_subgraph_outputs < 6 || (num_subgraph_outputs - first_present_output_index_) % 4 != 0,
+                  "number of outputs expected to be 2 + 4 * layers, got:", num_subgraph_outputs);
+
+    ORT_RETURN_IF(subgraph_inputs[2]->Name() != "decoder_input_ids",
+                  "encoder subgraph input 2 shall be named as decoder_input_ids, got: ", subgraph_inputs[2]->Name());
+    ORT_RETURN_IF(subgraph_inputs[2]->TypeAsProto()->tensor_type().elem_type() != int32_type,
+                  "encoder subgraph input 2 (decoder_input_ids) shall have int32 type");
+
+    ORT_RETURN_IF(subgraph_outputs[0]->Name() != "logits",
+                  "encoder subgraph output 0 shall be named as logits, got: ", subgraph_outputs[0]->Name());
+    ORT_RETURN_IF(subgraph_outputs[1]->Name() != "encoder_hidden_states",
+                  "encoder subgraph output 1 shall be named encoder_hidden_states, got: ", subgraph_outputs[1]->Name());
+    ORT_RETURN_IF(subgraph_outputs[2]->Name() != "present_key_self_0",
+                  "encoder subgraph output 2 shall be named as present_key_self_0, got: ", subgraph_outputs[2]->Name());
+    ORT_RETURN_IF(subgraph_outputs[3]->Name() != "present_value_self_0",
+                  "encoder subgraph output 3 shall be named as present_value_self_0, got: ", subgraph_outputs[3]->Name());
+
+    // Deduce num_heads, head_size and vocab_size from shape of graph outputs
+    const ONNX_NAMESPACE::TensorShapeProto* past_shape = subgraph_outputs[2]->Shape();
+    const ONNX_NAMESPACE::TensorShapeProto* logits_shape = subgraph_outputs[0]->Shape();
+    ORT_RETURN_IF_ERROR(GetParameters(past_shape, logits_shape, false));
+
+    num_layers = (num_subgraph_outputs - first_present_output_index_) / 4;
+  }
 
   auto output_type = subgraph_outputs[0]->TypeAsProto()->tensor_type().elem_type();
   ORT_RETURN_IF(output_type != float32_type && output_type != float16_type,
                 "encoder subgraph output 0 (logits) shall be float or float16 data type");
 
   for (int i = 1; i < num_subgraph_outputs; i++) {
     ORT_RETURN_IF(subgraph_outputs[i]->TypeAsProto()->tensor_type().elem_type() != output_type,
-                  "encoder subgraph outputs 1, 2, ... shall have same data type");
+                  "encoder subgraph outputs shall have same data type");
   }
 
   is_output_float16_ = (output_type == float16_type);
@@ -120,7 +147,6 @@ Status T5EncoderSubgraph::CreateInitialFeeds(
   }
   ORT_RETURN_IF(cpu_allocator == nullptr, "cpu_allocator shouldn't be nullptr");
 
-  // TODO(tianleiwu): expand the outputs instead of inputs to save computation.
   OrtValue encoder_input_ids;
   OrtValue encoder_attention_mask;
   ORT_RETURN_IF_ERROR(create_encoder_inputs_func(&original_encoder_input_ids,
@@ -136,9 +162,10 @@ Status T5EncoderSubgraph::CreateInitialFeeds(
   AllocatorPtr default_allocator = session_state_->GetAllocator(provider->GetOrtDeviceByMemType(OrtMemTypeDefault));
   AllocatorPtr pinned_allocator = session_state_->GetAllocator(provider->GetOrtDeviceByMemType(OrtMemTypeCPU));
   const OrtMemoryInfo& location = default_allocator->Info();
+
   ORT_RETURN_IF_ERROR(add_to_feeds_func(
       ort_stream,
-      {encoder_input_ids, encoder_attention_mask, decoder_input_ids},
+      num_subgraph_inputs == 2 ? std::initializer_list<OrtValue>{encoder_input_ids, encoder_attention_mask} : std::initializer_list<OrtValue>{encoder_input_ids, encoder_attention_mask, decoder_input_ids},
       feeds,
       buffer,
       default_allocator,

onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_encoder.h

@@ -16,7 +16,7 @@ class T5EncoderSubgraph : public Subgraph {
       const onnxruntime::Node& node_in,
       const std::string& attribute_name,
       const GraphViewer& subgraph_in) : Subgraph(node_in, attribute_name, subgraph_in) {
-    first_present_output_index_ = 2;
+    first_present_output_index_ = HasLogitsOutput() ? 2 : 0;
   }
 
   // Create inputs for first inference of subgraph.
@@ -36,9 +36,15 @@ class T5EncoderSubgraph : public Subgraph {
   Status Validate(const std::vector<const NodeArg*>& subgraph_inputs,
                   const std::vector<const NodeArg*>& subgraph_outputs) override;
 
+#ifdef DEBUG_GENERATION
   int GetFirstPresentOutputIndex() const {
     return first_present_output_index_;
   }
+#endif
+
+  bool HasLogitsOutput() const {
+    return num_subgraph_inputs != 2;
+  }
 
  protected:
   int first_present_output_index_;

onnxruntime/python/tools/transformers/convert_generation.py

@@ -10,7 +10,12 @@
 import os
 
 import torch
-from benchmark_helper import Precision, create_onnxruntime_session, prepare_environment, setup_logger
+from benchmark_helper import (
+    Precision,
+    create_onnxruntime_session,
+    prepare_environment,
+    setup_logger,
+)
 from t5_helper import PRETRAINED_MT5_MODELS, PRETRAINED_T5_MODELS, T5Helper
 
 logger = logging.getLogger("")
@@ -26,7 +31,8 @@ def parse_arguments():
         required=False,
         default=PRETRAINED_T5_MODELS[0],
         type=str,
-        help="Model path, or pretrained model name in the list: " + ", ".join(pretrained_models),
+        help="Model path, or pretrained model name in the list: "
+        + ", ".join(pretrained_models),
     )
 
     parser.add_argument(
@@ -63,7 +69,9 @@ def parse_arguments():
     )
     parser.set_defaults(optimize_onnx=False)
 
-    parser.add_argument("--use_gpu", required=False, action="store_true", help="use GPU for inference")
+    parser.add_argument(
+        "--use_gpu", required=False, action="store_true", help="use GPU for inference"
+    )
     parser.set_defaults(use_gpu=False)
 
     parser.add_argument(
@@ -79,7 +87,9 @@ def parse_arguments():
     parser.add_argument("--verbose", required=False, action="store_true")
     parser.set_defaults(verbose=False)
 
-    parser.add_argument("-e", "--use_external_data_format", required=False, action="store_true")
+    parser.add_argument(
+        "-e", "--use_external_data_format", required=False, action="store_true"
+    )
     parser.set_defaults(use_external_data_format=False)
 
     parser.add_argument(
@@ -108,14 +118,6 @@ def parse_arguments():
     )
     parser.set_defaults(disable_auto_mixed_precision=False)
 
-    parser.add_argument(
-        "--separate_encoder_and_decoder_init",
-        required=False,
-        action="store_true",
-        help="Do not merge encode and decoder init. Output 3 instead of 2 onnx models.",
-    )
-    parser.set_defaults(separate_encoder_and_decoder_init=False)
-
     parser.add_argument(
         "--use_int64_inputs",
         required=False,
@@ -131,6 +133,14 @@ def parse_arguments():
         help="filepath to load pre-trained model with custom state dictionary (e.g. pytorch_model.bin)",
     )
 
+    parser.add_argument(
+        "--encode_decoder_init",
+        required=False,
+        action="store_true",
+        help="Combine encoder and decoder kv cache initialization into one model.",
+    )
+    parser.set_defaults(encode_decoder_init=False)
+
     args = parser.parse_args()
 
     return args
@@ -146,17 +156,22 @@ def export_onnx_models(
     precision,
     verbose,
     use_decoder_start_token: bool = False,
-    merge_encoder_and_decoder_init: bool = True,
     overwrite: bool = False,
     disable_auto_mixed_precision: bool = False,
     use_int32_inputs: bool = True,
     model_type: str = "t5",
     state_dict_path: str = "",
+    encode_decoder_init: bool = False,
 ):
     device = torch.device("cuda:0" if use_gpu else "cpu")
 
     models = T5Helper.load_model(
-        model_name_or_path, cache_dir, device, merge_encoder_and_decoder_init, model_type, state_dict_path
+        model_name_or_path,
+        cache_dir,
+        device,
+        model_type,
+        state_dict_path,
+        encode_decoder_init=encode_decoder_init,
     )
     config = models["decoder"].config
 
@@ -220,11 +235,17 @@ def export_onnx_models(
         ort_session = create_onnxruntime_session(
             output_path,
             use_gpu=use_gpu,
-            provider=["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"],
+            provider=(
+                ["CUDAExecutionProvider", "CPUExecutionProvider"]
+                if use_gpu
+                else ["CPUExecutionProvider"]
+            ),
         )
 
         with torch.no_grad():
-            max_diff = T5Helper.verify_onnx(model, ort_session, device, use_int32_inputs)
+            max_diff = T5Helper.verify_onnx(
+                model, ort_session, device, use_int32_inputs
+            )
         logger.info(f"PyTorch and OnnxRuntime results max difference = {max_diff}")
         if max_diff > 1e-4:
             logger.warning("PyTorch and OnnxRuntime results are NOT close")
@@ -242,7 +263,11 @@ def main():
     logger.info(f"Arguments:{args}")
 
     cache_dir = args.cache_dir
-    output_dir = args.output if not args.output.endswith(".onnx") else os.path.dirname(args.output)
+    output_dir = (
+        args.output
+        if not args.output.endswith(".onnx")
+        else os.path.dirname(args.output)
+    )
     prepare_environment(cache_dir, output_dir, args.use_gpu)
 
     if args.precision != Precision.FLOAT32:
@@ -264,11 +289,11 @@ def main():
         args.precision,
         args.verbose,
         args.use_decoder_start_token,
-        not args.separate_encoder_and_decoder_init,
         args.overwrite,
         args.disable_auto_mixed_precision,
         not args.use_int64_inputs,
         args.model_type,
+        encode_decoder_init=args.encode_decoder_init,
     )
 
     logger.info(f"Done! Outputs: {output_paths}")

onnxruntime/python/tools/transformers/models/t5/convert_to_onnx.py

@@ -1,24 +1,14 @@
 # -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation.  All rights reserved.
-# Licensed under the MIT License.  See License.txt in the project root for
-# license information.
-# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# -------------------------------------------------------------------------
 
 import logging
-import os
 import random
-import tempfile
-from pathlib import Path
 
-import numpy
-import onnx
 import torch
-from onnx_model import OnnxModel
-from torch_onnx_export_helper import torch_onnx_export
 from transformers import MT5Config, T5Config
 
-from onnxruntime import InferenceSession
-
 logger = logging.getLogger(__name__)
 
 
@@ -41,7 +31,11 @@ def __init__(self, input_ids, attention_mask):
 
     @staticmethod
     def create_dummy(
-        batch_size: int, sequence_length: int, vocab_size: int, device: torch.device, use_int32_inputs: bool = False
+        batch_size: int,
+        sequence_length: int,
+        vocab_size: int,
+        device: torch.device,
+        use_int32_inputs: bool = False,
     ):  # -> T5EncoderInputs
         """Create dummy inputs for T5 encoder.
 
@@ -64,7 +58,9 @@ def create_dummy(
             device=device,
         )
 
-        attention_mask = torch.ones([batch_size, sequence_length], dtype=dtype, device=device)
+        attention_mask = torch.ones(
+            [batch_size, sequence_length], dtype=dtype, device=device
+        )
         if sequence_length >= 2:
             for i in range(batch_size):
                 padding_position = random.randint(0, sequence_length - 1)
@@ -74,97 +70,3 @@ def create_dummy(
     def to_list(self) -> list:
         input_list = [v for v in [self.input_ids, self.attention_mask] if v is not None]
         return input_list
-
-
-class T5EncoderHelper:
-    @staticmethod
-    def export_onnx(
-        encoder: T5Encoder,
-        device: torch.device,
-        onnx_model_path: str,
-        verbose: bool = True,
-        use_external_data_format: bool = False,
-        use_int32_inputs: bool = False,
-    ):
-        """Export encoder to ONNX
-
-        Args:
-            encoder (T5Encoder): encoder object
-            device (torch.device): device of encoder object
-            onnx_model_path (str): onnx path
-            verbose (bool, optional): print verbose information. Defaults to True.
-            use_external_data_format (bool, optional): use external data format or not. Defaults to False.
-        """
-        config = encoder.config
-        encoder_inputs = T5EncoderInputs.create_dummy(
-            batch_size=2,
-            sequence_length=4,
-            vocab_size=config.vocab_size,
-            device=device,
-            use_int32_inputs=use_int32_inputs,
-        )
-
-        Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
-
-        with tempfile.TemporaryDirectory() as tmp_dir_name:
-            temp_onnx_model_path = os.path.join(tmp_dir_name, "encoder.onnx")
-            Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
-            torch_onnx_export(
-                encoder,
-                args=tuple(encoder_inputs.to_list()),
-                f=temp_onnx_model_path if use_external_data_format else onnx_model_path,
-                export_params=True,
-                input_names=["input_ids", "attention_mask"],
-                output_names=["hidden_states"],
-                dynamic_axes={
-                    "input_ids": {0: "batch_size", 1: "sequence_length"},
-                    "attention_mask": {0: "batch_size", 1: "sequence_length"},
-                    "hidden_states": {0: "batch_size", 1: "sequence_length"},
-                },
-                opset_version=12,
-                do_constant_folding=True,
-                use_external_data_format=use_external_data_format,
-                verbose=verbose,
-            )
-
-            if use_external_data_format:
-                model = onnx.load_model(temp_onnx_model_path, load_external_data=True)
-                OnnxModel.save(
-                    model,
-                    onnx_model_path,
-                    save_as_external_data=True,
-                    all_tensors_to_one_file=True,
-                )
-
-    @staticmethod
-    def onnxruntime_inference(ort_session, inputs: T5EncoderInputs):
-        """Run inference of ONNX model."""
-        ort_inputs = {
-            "input_ids": numpy.ascontiguousarray(inputs.input_ids.cpu().numpy()),
-            "attention_mask": numpy.ascontiguousarray(inputs.attention_mask.cpu().numpy()),
-        }
-
-        return ort_session.run(None, ort_inputs)
-
-    @staticmethod
-    def verify_onnx(
-        model: T5Encoder, ort_session: InferenceSession, device: torch.device, use_int32_inputs: bool = False
-    ):
-        """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good."""
-        inputs = T5EncoderInputs.create_dummy(
-            batch_size=4,
-            sequence_length=11,
-            vocab_size=model.config.vocab_size,
-            device=device,
-            use_int32_inputs=use_int32_inputs,
-        )
-        input_list = inputs.to_list()
-        torch_outputs = model(*input_list)
-
-        ort_outputs = T5EncoderHelper.onnxruntime_inference(ort_session, inputs)
-
-        max_diff = numpy.amax(numpy.abs(torch_outputs.cpu().numpy() - ort_outputs[0]))
-
-        logger.info(f"max_diff={max_diff}")
-
-        return max_diff

onnxruntime/python/tools/transformers/models/t5/t5_encoder.py

@@ -1,8 +1,7 @@
 # -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation.  All rights reserved.
-# Licensed under the MIT License.  See License.txt in the project root for
-# license information.
-# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# -------------------------------------------------------------------------
 
 import logging
 import os
@@ -34,27 +33,42 @@ def __init__(
         lm_head: torch.nn.Module,
         config: T5Config | MT5Config,
         decoder_start_token_id: int | None = None,
+        output_cross_only: bool = False,
     ):
         super().__init__()
         self.config = config
         self.t5_encoder = T5Encoder(encoder, config)
-        self.t5_decoder_init = T5DecoderInit(decoder, lm_head, config, decoder_start_token_id)
+        self.t5_decoder_init = T5DecoderInit(
+            decoder, lm_head, config, decoder_start_token_id
+        )
+        self.output_cross_only = output_cross_only
 
     def forward(
         self,
         encoder_input_ids: torch.Tensor,
         encoder_attention_mask: torch.Tensor,
         decoder_input_ids: torch.Tensor = None,
     ):
-        encoder_hidden_states: torch.FloatTensor = self.t5_encoder(encoder_input_ids, encoder_attention_mask)
-        lm_logits, past_self, past_cross = self.t5_decoder_init(
-            decoder_input_ids, encoder_attention_mask, encoder_hidden_states
+        encoder_hidden_states: torch.FloatTensor = self.t5_encoder(
+            encoder_input_ids, encoder_attention_mask
         )
-        return lm_logits, encoder_hidden_states, past_self, past_cross
+
+        if self.output_cross_only:
+            lm_logits, past_self, past_cross = self.t5_decoder_init(
+                decoder_input_ids, encoder_attention_mask, encoder_hidden_states
+            )
+            return past_cross
+        else:
+            lm_logits, past_self, past_cross = self.t5_decoder_init(
+                decoder_input_ids, encoder_attention_mask, encoder_hidden_states
+            )
+            return lm_logits, encoder_hidden_states, past_self, past_cross
 
 
 class T5EncoderDecoderInitInputs:
-    def __init__(self, encoder_input_ids, encoder_attention_mask, decoder_input_ids=None):
+    def __init__(
+        self, encoder_input_ids, encoder_attention_mask, decoder_input_ids=None
+    ):
         self.encoder_input_ids: torch.LongTensor = encoder_input_ids
         self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask
         self.decoder_input_ids: torch.LongTensor = decoder_input_ids
@@ -78,9 +92,14 @@ def create_dummy(
         decoder_input_ids = None
         if use_decoder_input_ids:
             dtype = torch.int32 if use_int32_inputs else torch.int64
-            decoder_input_ids = torch.ones((batch_size, 1), dtype=dtype, device=device) * config.decoder_start_token_id
+            decoder_input_ids = (
+                torch.ones((batch_size, 1), dtype=dtype, device=device)
+                * config.decoder_start_token_id
+            )
 
-        return T5EncoderDecoderInitInputs(encoder_inputs.input_ids, encoder_inputs.attention_mask, decoder_input_ids)
+        return T5EncoderDecoderInitInputs(
+            encoder_inputs.input_ids, encoder_inputs.attention_mask, decoder_input_ids
+        )
 
     def to_list(self) -> list:
         input_list = [self.encoder_input_ids, self.encoder_attention_mask]
@@ -108,9 +127,14 @@ def export_onnx(
             onnx_model_path (str): onnx path
             verbose (bool, optional): print verbose information. Defaults to True.
             use_external_data_format (bool, optional): use external data format or not. Defaults to False.
+            use_int32_inputs (bool, optional): use int32 instead of int64 for integer inputs. Defaults to False.
         """
         assert isinstance(model, T5EncoderDecoderInit)
 
+        # Do not exclude decoder in torch onnx export so that cross can show up.
+        output_cross_only = model.output_cross_only
+        model.output_cross_only = False
+
         inputs = T5EncoderDecoderInitInputs.create_dummy(
             model.config,
             batch_size=2,
@@ -121,7 +145,9 @@ def export_onnx(
         )
         input_list = inputs.to_list()
 
-        present_names = PastKeyValuesHelper.get_past_names(model.config.num_decoder_layers, present=True)
+        present_names = PastKeyValuesHelper.get_past_names(
+            model.config.num_decoder_layers, present=True
+        )
 
         output_names = ["logits", "encoder_hidden_states", *present_names]
 
@@ -185,7 +211,9 @@ def export_onnx(
                 }
 
         with tempfile.TemporaryDirectory() as tmp_dir_name:
-            temp_onnx_model_path = os.path.join(tmp_dir_name, "encoder_decoder_init.onnx")
+            temp_onnx_model_path = os.path.join(
+                tmp_dir_name, "encoder_decoder_init.onnx"
+            )
             Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
             torch_onnx_export(
                 model,
@@ -201,6 +229,9 @@ def export_onnx(
                 verbose=verbose,
             )
 
+            # Restore output_cross_only setting.
+            model.output_cross_only = output_cross_only
+
             # Workaround as mentioned earlier: change numeric dim_param to dim_value
             model = onnx.load(temp_onnx_model_path)
             for tensor in model.graph.output:
@@ -215,6 +246,54 @@ def export_onnx(
                         dim_proto.Clear()
                         dim_proto.dim_value = dim_value
 
+            if output_cross_only:
+                # Rewrite onnx graph to only keep present_[key|value]_cross_* outputs.
+                onnx_model = OnnxModel(model)
+                output_name_to_node = onnx_model.output_name_to_node()
+
+                for output in model.graph.output:
+                    if "cross" in output.name:
+                        assert output.name in output_name_to_node
+
+                        transpose_node = output_name_to_node[output.name]
+                        assert transpose_node and transpose_node.op_type == "Transpose"
+
+                        permutation = OnnxModel.get_node_attribute(
+                            transpose_node, "perm"
+                        )
+                        assert isinstance(permutation, list)
+                        assert permutation == [0, 2, 1, 3]
+
+                        matched_nodes = onnx_model.match_parent_path(
+                            transpose_node,
+                            ["Reshape", "MatMul"],
+                            [0, 0],
+                            output_name_to_node,
+                        )
+                        assert matched_nodes is not None
+
+                        reshape_node, matmul_node = matched_nodes
+                        assert "encoder_hidden_states" in matmul_node.input
+
+                        if not onnx_model.get_initializer("cross_reshape_shape"):
+                            shape_tensor = onnx.helper.make_tensor(
+                                name="cross_reshape_shape",
+                                data_type=onnx.TensorProto.INT64,
+                                dims=[4],
+                                vals=[0, 0, int(num_heads), int(head_size)],
+                                raw=False,
+                            )
+                            onnx_model.add_initializer(shape_tensor)
+
+                        reshape_node.input[1] = "cross_reshape_shape"
+
+                cross_outputs = [
+                    output.name
+                    for output in model.graph.output
+                    if "cross" in output.name
+                ]
+                onnx_model.prune_graph(cross_outputs, allow_remove_graph_inputs=True)
+
             OnnxModel.save(
                 model,
                 onnx_model_path,
@@ -228,11 +307,17 @@ def onnxruntime_inference(ort_session, inputs: T5EncoderDecoderInitInputs):
         logger.debug("start onnxruntime_inference")
 
         ort_inputs = {
-            "encoder_input_ids": numpy.ascontiguousarray(inputs.encoder_input_ids.cpu().numpy()),
-            "encoder_attention_mask": numpy.ascontiguousarray(inputs.encoder_attention_mask.cpu().numpy()),
+            "encoder_input_ids": numpy.ascontiguousarray(
+                inputs.encoder_input_ids.cpu().numpy()
+            ),
+            "encoder_attention_mask": numpy.ascontiguousarray(
+                inputs.encoder_attention_mask.cpu().numpy()
+            ),
         }
         if inputs.decoder_input_ids is not None:
-            ort_inputs["decoder_input_ids"] = numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy())
+            ort_inputs["decoder_input_ids"] = numpy.ascontiguousarray(
+                inputs.decoder_input_ids.cpu().numpy()
+            )
 
         ort_outputs = ort_session.run(None, ort_inputs)
         return ort_outputs
@@ -261,35 +346,57 @@ def verify_onnx(
                 use_int32_inputs=use_int32_inputs,
             )
 
-            ort_outputs = T5EncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs)
+            ort_outputs = T5EncoderDecoderInitHelper.onnxruntime_inference(
+                ort_session, inputs
+            )
 
             # Run inference of PyTorch model
             input_list = inputs.to_list()
             torch_outputs = model(*input_list)
 
             num_decoder_layers = model.config.num_decoder_layers
 
-            assert torch_outputs[0].cpu().numpy().shape == ort_outputs[0].shape
-            max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs[0]))
-            logger.debug(f"logits max_diff={max_diff}")
-            max_diff_all = max_diff
-
-            assert torch_outputs[1].cpu().numpy().shape == ort_outputs[1].shape
-            max_diff = numpy.amax(numpy.abs(torch_outputs[1].cpu().numpy() - ort_outputs[1]))
-            logger.debug(f"encoder_hidden_states max_diff={max_diff}")
-            max_diff_all = max(max_diff_all, max_diff)
-
-            for i in range(2 * num_decoder_layers):
-                max_diff = numpy.amax(numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs[2 + i]))
-                logger.debug(f"self attention past state {i} max_diff={max_diff}")
+            if not model.output_cross_only:
+                assert torch_outputs[0].cpu().numpy().shape == ort_outputs[0].shape
+                max_diff = numpy.amax(
+                    numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs[0])
+                )
+                logger.debug(f"logits max_diff={max_diff}")
+                max_diff_all = max_diff
 
-            for i in range(2 * num_decoder_layers):
+                assert torch_outputs[1].cpu().numpy().shape == ort_outputs[1].shape
                 max_diff = numpy.amax(
-                    numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs[2 + 2 * num_decoder_layers + i])
+                    numpy.abs(torch_outputs[1].cpu().numpy() - ort_outputs[1])
                 )
-                logger.debug(f"cross attention past state {i} max_diff={max_diff}")
+                logger.debug(f"encoder_hidden_states max_diff={max_diff}")
                 max_diff_all = max(max_diff_all, max_diff)
 
+                for i in range(2 * num_decoder_layers):
+                    max_diff = numpy.amax(
+                        numpy.abs(
+                            torch_outputs[2][i].cpu().numpy() - ort_outputs[2 + i]
+                        )
+                    )
+                    logger.debug(f"self attention past state {i} max_diff={max_diff}")
+
+                for i in range(2 * num_decoder_layers):
+                    max_diff = numpy.amax(
+                        numpy.abs(
+                            torch_outputs[3][i].cpu().numpy()
+                            - ort_outputs[2 + 2 * num_decoder_layers + i]
+                        )
+                    )
+                    logger.debug(f"cross attention past state {i} max_diff={max_diff}")
+                    max_diff_all = max(max_diff_all, max_diff)
+            else:
+                max_diff_all = -float("inf")
+                for i in range(2 * num_decoder_layers):
+                    max_diff = numpy.amax(
+                        numpy.abs(torch_outputs[i].cpu().numpy() - ort_outputs[i])
+                    )
+                    logger.debug(f"cross attention past state {i} max_diff={max_diff}")
+                    max_diff_all = max(max_diff_all, max_diff)
+
             test_cases_max_diff.append(max_diff_all)
             logger.info(
                 f"batch_size={batch_size} encode_sequence_length={encode_sequence_length}, max_diff={max_diff_all}"

onnxruntime/python/tools/transformers/models/t5/t5_encoder_decoder_init.py

@@ -1,8 +1,7 @@
 # -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation.  All rights reserved.
-# Licensed under the MIT License.  See License.txt in the project root for
-# license information.
-# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# -------------------------------------------------------------------------
 
 import logging
 import os
@@ -12,8 +11,7 @@
 from float16 import float_to_float16_max_diff
 from onnx_model import OnnxModel
 from optimizer import optimize_model
-from t5_decoder import T5Decoder, T5DecoderHelper, T5DecoderInit
-from t5_encoder import T5Encoder, T5EncoderHelper
+from t5_decoder import T5Decoder, T5DecoderHelper
 from t5_encoder_decoder_init import T5EncoderDecoderInit, T5EncoderDecoderInitHelper
 from transformers import MT5ForConditionalGeneration, T5ForConditionalGeneration
 
@@ -22,7 +20,13 @@
 logger = logging.getLogger(__name__)
 
 PRETRAINED_T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b"]
-PRETRAINED_MT5_MODELS = ["google/mt5-small", "google/mt5-base", "google/mt5-large", "google/mt5-xl", "google/mt5-xxl"]
+PRETRAINED_MT5_MODELS = [
+    "google/mt5-small",
+    "google/mt5-base",
+    "google/mt5-large",
+    "google/mt5-xl",
+    "google/mt5-xxl",
+]
 
 
 class T5Helper:
@@ -60,25 +64,30 @@ def load_model(
         model_name_or_path: str,
         cache_dir: str,
         device: torch.device,
-        merge_encoder_and_decoder_init: bool = True,
         model_type: str = "t5",
         state_dict_path: str = "",
+        encode_decoder_init: bool = False,
     ) -> dict[str, torch.nn.Module]:
         """Load model given a pretrained name or path, then build models for ONNX conversion.
 
         Args:
             model_name_or_path (str): pretrained model name or path
             cache_dir (str): cache directory
             device (torch.device): device to run the model
-            merge_encoder_and_decoder_init (bool, optional): Whether merge encoder and decoder initialization into one ONNX model. Defaults to True.
-            is_mt5 (bool, optional): whether the model is MT5 instead of T5
+            model_type (str, optional): model type "t5" or "mt5"
+            state_dict_path(str, optional): state dictionary path
+            encode_decoder_init (bool, optional): combine encoder and decoder kv cache initialization into one model.
         Returns:
             Dict[str, torch.nn.Module]: mapping from name to modules for ONNX conversion.
         """
         if model_type == "t5":
-            model = T5ForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir)
+            model = T5ForConditionalGeneration.from_pretrained(
+                model_name_or_path, cache_dir=cache_dir
+            )
         elif model_type == "mt5":
-            model = MT5ForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir)
+            model = MT5ForConditionalGeneration.from_pretrained(
+                model_name_or_path, cache_dir=cache_dir
+            )
         else:
             raise ValueError("only support mode_type=t5 or mt5")
 
@@ -88,46 +97,29 @@ def load_model(
         decoder = T5Decoder(model.decoder, model.lm_head, model.config)
         decoder.eval().to(device)
 
-        if merge_encoder_and_decoder_init:
-            encoder_decoder_init = T5EncoderDecoderInit(
-                model.encoder,
-                model.decoder,
-                model.lm_head,
-                model.config,
-                decoder_start_token_id=None,
-            )
-            return {"encoder_decoder_init": encoder_decoder_init, "decoder": decoder}
-        else:
-            encoder = T5Encoder(model.encoder, model.config)
-            encoder.eval().to(device)
-            decoder_init = T5DecoderInit(model.decoder, model.lm_head, model.config)
-            decoder_init.eval().to(device)
-            return {
-                "encoder": encoder,
-                "decoder": decoder,
-                "decoder_init": decoder_init,
-            }
+        encoder_decoder_init = T5EncoderDecoderInit(
+            model.encoder,
+            model.decoder,
+            model.lm_head,
+            model.config,
+            decoder_start_token_id=None,
+            output_cross_only=not encode_decoder_init,
+        )
+
+        encoder_name = "encoder_decoder_init" if encode_decoder_init else "encoder"
+        return {encoder_name: encoder_decoder_init, "decoder": decoder}
 
     @staticmethod
     def export_onnx(
-        model: T5Encoder | T5Decoder | T5DecoderInit | T5EncoderDecoderInit,
+        model: T5Decoder | T5EncoderDecoderInit,
         device: torch.device,
         onnx_model_path: str,
         verbose: bool = True,
         use_external_data_format: bool = False,
         use_decoder_input_ids: bool = True,
         use_int32_inputs: bool = False,
     ):
-        if isinstance(model, T5Encoder):
-            T5EncoderHelper.export_onnx(
-                model,
-                device,
-                onnx_model_path,
-                verbose,
-                use_external_data_format,
-                use_int32_inputs,
-            )
-        elif isinstance(model, T5EncoderDecoderInit):
+        if isinstance(model, T5EncoderDecoderInit):
             T5EncoderDecoderInitHelper.export_onnx(
                 model,
                 device,
@@ -191,10 +183,14 @@ def auto_mixed_precision(
             # when the max difference of value after converting float to float16 is lower than a threshold (1e-6),
             # we can deduce that the weights are stored in float16 precision.
             max_diff = float_to_float16_max_diff(initializer)
-            logger.debug(f"max diff of converting weights in last MatMul node {node.name}: {max_diff}")
+            logger.debug(
+                f"max diff of converting weights in last MatMul node {node.name}: {max_diff}"
+            )
             is_weight_fp16_precision = max_diff < 1e-6
         else:
-            logger.warning(f"Failed to find MatMul node for logits. Found {node.op_type} of node {node.name}")
+            logger.warning(
+                f"Failed to find MatMul node for logits. Found {node.op_type} of node {node.name}"
+            )
 
         keep_io_types = []
         node_block_list = []
@@ -252,20 +248,21 @@ def optimize_onnx(
             else:
                 m.convert_model_float32_to_float16(cast_input_output=False)
 
-        m.save_model_to_file(optimized_model_path, use_external_data_format, all_tensors_to_one_file=True)
+        m.save_model_to_file(
+            optimized_model_path, use_external_data_format, all_tensors_to_one_file=True
+        )
 
     @staticmethod
     def verify_onnx(
-        model: T5Encoder | T5Decoder | T5DecoderInit | T5EncoderDecoderInit,
+        model: T5Decoder | T5EncoderDecoderInit,
         ort_session: InferenceSession,
         device: torch.device,
         use_int32_inputs: bool,
     ):
         """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good."""
-        if isinstance(model, T5Encoder):
-            return T5EncoderHelper.verify_onnx(model, ort_session, device, use_int32_inputs)
-
         if isinstance(model, T5EncoderDecoderInit):
-            return T5EncoderDecoderInitHelper.verify_onnx(model, ort_session, device, use_int32_inputs)
+            return T5EncoderDecoderInitHelper.verify_onnx(
+                model, ort_session, device, use_int32_inputs
+            )
 
         return T5DecoderHelper.verify_onnx(model, ort_session, device, use_int32_inputs)