diff --git a/src/config.cpp b/src/config.cpp
index 3ee54683c5..8e91e9e061 100644
--- a/src/config.cpp
+++ b/src/config.cpp
@@ -4,6 +4,68 @@
 
 namespace Generators {
 
+ONNXTensorElementDataType TranslateTensorType(std::string_view value) {
+  if (value == "float32")
+    return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
+  if (value == "float16")
+    return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
+
+  throw std::runtime_error("Invalid tensor type: " + std::string(value));
+}
+
+struct Model_Element : JSON::Element {
+  Model_Element(Config::Model& model) : model_{model} {}
+
+  void OnString(std::string_view name, std::string_view value) override {
+    if (name == "type")
+      model_.type = value;
+    else if (name == "decoder")
+      model_.decoder = value;
+    else if (name == "encoder_decoder_init")
+      model_.encoder_decoder_init = value;
+    else if (name == "past_names_key")
+      model_.past_names_key = value;
+    else if (name == "past_names_value")
+      model_.past_names_value = value;
+    else if (name == "present_names_key")
+      model_.present_names_key = value;
+    else if (name == "present_names_value")
+      model_.present_names_value = value;
+    else if (name == "past_names")
+      model_.past_names = value;
+    else if (name == "present_names")
+      model_.present_names = value;
+    else if (name == "cross_past_names_key")
+      model_.cross_past_names_key = value;
+    else if (name == "cross_past_names_value")
+      model_.cross_past_names_value = value;
+    else if (name == "cross_present_names_key")
+      model_.cross_present_names_key = value;
+    else if (name == "cross_present_names_value")
+      model_.cross_present_names_value = value;
+    else if (name == "logits_type")
+      model_.logits_type = TranslateTensorType(value);
+    else if (name == "kv_type")
+      model_.kv_type = TranslateTensorType(value);
+    else
+      throw std::runtime_error("Unknown name: " + std::string(name));
+  }
+
+  void OnNumber(std::string_view name, double value) override {
+    if (name == "vocab_size")
+      model_.vocab_size = static_cast<int>(value);
+    else if (name == "hidden_size" || name == "n_embed")
+      model_.hidden_size = static_cast<int>(value);
+    else if (name == "num_attention_heads" || name == "num_heads" || name == "n_head")
+      model_.num_attention_heads = static_cast<int>(value);
+    else if (name == "num_hidden_layers" || name == "num_layers" || name == "n_layer")
+      model_.num_hidden_layers = static_cast<int>(value);
+  }
+
+ private:
+  Config::Model& model_;
+};
+
 struct Root_Element : JSON::Element {
   Root_Element(Config& config) : config_{config} {}
 
@@ -12,12 +74,6 @@ struct Root_Element : JSON::Element {
       config_.tokenizer_class = value;
     else if (name == "prefix")
       config_.prefix = value;
-    else if (name == "model_type")
-      config_.model_type = value;
-    else if (name == "ogai_model_decoder")
-      config_.model_decoder = value;
-    else if (name == "ogai_model_encoder_decoder_init")
-      config_.model_encoder_decoder_init = value;
   }
 
   void OnNumber(std::string_view name, double value) override {
@@ -50,16 +106,6 @@ struct Root_Element : JSON::Element {
       config_.decoder_start_token_id = static_cast<int>(value);
     else if (name == "sep_token_id")
       config_.sep_token_id = static_cast<int>(value);
-
-    // Model Class Attributes
-    else if (name == "vocab_size")
-      config_.vocab_size = static_cast<int>(value);
-    else if (name == "hidden_size" || name == "n_embed")
-      config_.hidden_size = static_cast<int>(value);
-    else if (name == "num_attention_heads" || name == "num_heads" || name == "n_head")
-      config_.num_attention_heads = static_cast<int>(value);
-    else if (name == "num_hidden_layers" || name == "num_layers" || name == "n_layer")
-      config_.num_hidden_layers = static_cast<int>(value);
   }
 
   void OnBool(std::string_view name, bool value) override {
@@ -67,7 +113,14 @@ struct Root_Element : JSON::Element {
       config_.early_stopping = value;
   }
 
+  Element& OnObject(std::string_view name) {
+    if (name == "model")
+      return model_element_;
+    return Element::OnObject(name);
+  }
+
   Config& config_;
+  Model_Element model_element_{config_.model};
 };
 
 struct RootObject_Element : JSON::Element {
diff --git a/src/config.h b/src/config.h
index 1a95808796..6aed0dc94e 100644
--- a/src/config.h
+++ b/src/config.h
@@ -27,14 +27,30 @@ struct Config {
   int decoder_start_token_id{};  // If an encoder-decoder model starts decoding with a different token than bos, the id of that token.
   int sep_token_id{};            // The id of the separation token.
 
-  // Model Class Attributes
-  std::string model_decoder;
-  std::string model_encoder_decoder_init;
-  std::string model_type;
-  int vocab_size{};
-  int hidden_size{};
-  int num_attention_heads{};
-  int num_hidden_layers{};
+  struct Model {
+    std::string decoder;
+    std::string encoder_decoder_init;
+    std::string type;
+
+    ONNXTensorElementDataType logits_type{ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; // float16/float32 are the valid types
+    ONNXTensorElementDataType kv_type{ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT};  // float16/float32 are the valid types
+
+    int vocab_size{};
+    int hidden_size{};
+    int num_attention_heads{};
+    int num_hidden_layers{};
+
+    // KV_Cache names (will be a string with a %d in it, like "past_key_self_%d", "past_value_self_%d")
+    std::string past_names_key, past_names_value;
+    std::string present_names_key, present_names_value;
+
+    // KV_Cache_Combined names where the kv key/value are merged into one tensor
+    std::string past_names, present_names;
+
+    // Cross_Cache for models like whisper
+    std::string cross_past_names_key, cross_past_names_value;
+    std::string cross_present_names_key, cross_present_names_value;
+  } model;
 };
 
 }  // namespace Generators
\ No newline at end of file
diff --git a/src/generators.cpp b/src/generators.cpp
index 86f2bc64a9..094dfecd7e 100644
--- a/src/generators.cpp
+++ b/src/generators.cpp
@@ -44,7 +44,7 @@ float Float16ToFloat32(uint16_t v) {
 SearchParams::SearchParams(const Model& model)
     : pad_token_id{model.config_->pad_token_id},
       eos_token_id{model.config_->eos_token_id},
-      vocab_size{model.config_->vocab_size},
+      vocab_size{model.config_->model.vocab_size},
       max_length{model.config_->max_length},
       length_penalty{model.config_->length_penalty},
       early_stopping{model.config_->early_stopping},
diff --git a/src/models/gpt.cpp b/src/models/gpt.cpp
index 40a0d9e413..f8a467638a 100644
--- a/src/models/gpt.cpp
+++ b/src/models/gpt.cpp
@@ -5,10 +5,9 @@ namespace Generators {
 
 Gpt_Model::Gpt_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
     : Model{std::move(config), ort_env, provider_options} {
-  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
+  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());
 
   InitDeviceAllocator(*session_decoder_);
-  InitLogits(*session_decoder_->GetOutputTypeInfo(0));
 }
 
 std::unique_ptr<State> Gpt_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp
index 00e8d3b711..59c4a2304c 100644
--- a/src/models/kv_cache.cpp
+++ b/src/models/kv_cache.cpp
@@ -7,21 +7,21 @@ namespace Generators {
 KV_Cache_Combined::KV_Cache_Combined(Model& model, State& state)
     : model_{model},
       state_{state},
-      layer_count_{model.config_->num_hidden_layers},
-      shape_{2, state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->num_attention_heads, 0, model.config_->hidden_size},
-      empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_)} {
+      layer_count_{model.config_->model.num_hidden_layers},
+      shape_{2, state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->model.num_attention_heads, 0, model.config_->model.hidden_size},
+      empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type)} {
   pasts_.resize(layer_count_);
   presents_.reserve(layer_count_);
 
   shape_[3] = state_.search_params_.sequence_length;
   for (int i = 0; i < layer_count_; ++i) {
-    presents_.push_back(OrtValue::CreateTensor(*model.allocator_device_, shape_, model_.score_type_));
+    presents_.push_back(OrtValue::CreateTensor(*model.allocator_device_, shape_, model_.config_->model.kv_type));
 
-    char string[32];
-    snprintf(string, std::size(string), past_name_, i);
+    char string[64];
+    snprintf(string, std::size(string), model.config_->model.past_names.c_str(), i);
     input_name_strings_.push_back(string);
 
-    snprintf(string, std::size(string), present_name_, i);
+    snprintf(string, std::size(string), model.config_->model.present_names.c_str(), i);
     output_name_strings_.push_back(string);
   }
 }
@@ -50,7 +50,7 @@ void KV_Cache_Combined::Update(std::span<const int32_t> beam_indices, int curren
 
   shape_[3] = current_length;
   for (int i = 0; i < layer_count_; i++) {
-    presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_);
+    presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type);
     state_.inputs_[input_index_ + i] = pasts_[i].get();
     state_.outputs_[output_index_ + i] = presents_[i].get();
   }
@@ -99,39 +99,37 @@ void KV_Cache_Combined::PickPastState(std::span<const int32_t> beam_indices, int
 }
 
 void KV_Cache_Combined::PickPastState(std::span<const int32_t> beam_indices, int index) {
-  if (model_.score_type_ == Ort::TypeToTensorType<float>::type)
+  if (model_.config_->model.kv_type == Ort::TypeToTensorType<float>::type)
     PickPastState<float>(beam_indices, index);
   else
     PickPastState<Ort::Float16_t>(beam_indices, index);
 }
 
-KV_Cache::KV_Cache(Model& model, State& state,
-                   std::span<const char*> past_names, std::span<const char*> present_names)
+KV_Cache::KV_Cache(Model& model, State& state)
     : model_{model},
       state_{state},
-      layer_count_{model_.config_->num_hidden_layers},
-      past_names_{past_names},
-      present_names_{present_names},
-      shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->num_attention_heads, 0, model.config_->hidden_size},
-      empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_)} {
+      layer_count_{model_.config_->model.num_hidden_layers},
+      shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->model.num_attention_heads, 0, model.config_->model.hidden_size},
+      empty_past_{OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type)} {
   pasts_.resize(layer_count_ * 2);
   presents_.reserve(layer_count_ * 2);
 
   shape_[2] = state_.search_params_.sequence_length; // Set this after empty_past_ has been created with 0 for this field
 
   for (int i = 0; i < layer_count_; ++i) {
-    presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
-    presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
+    presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));
+    presents_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));
 
-    char string[32];
-    for (auto* name : past_names_) {
-      snprintf(string, std::size(string), name, i);
-      input_name_strings_.push_back(string);
-    }
-    for (auto* name : present_names_) {
-      snprintf(string, std::size(string), name, i);
-      output_name_strings_.push_back(string);
-    }
+    char string[64];
+    snprintf(string, std::size(string), model.config_->model.past_names_key.c_str(), i);
+    input_name_strings_.push_back(string);
+    snprintf(string, std::size(string), model.config_->model.past_names_value.c_str(), i);
+    input_name_strings_.push_back(string);
+
+    snprintf(string, std::size(string), model.config_->model.present_names_key.c_str(), i);
+    output_name_strings_.push_back(string);
+    snprintf(string, std::size(string), model.config_->model.present_names_value.c_str(), i);
+    output_name_strings_.push_back(string);
   }
 }
 
@@ -167,7 +165,7 @@ void KV_Cache::Update(std::span<const int32_t> beam_indices, int current_length)
 
   shape_[2] = current_length;
   for (int i = 0; i < layer_count_ * 2; i++) {
-    presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_);
+    presents_[i] = OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type);
     state_.outputs_[output_index_ + i] = presents_[i].get();
   }
 }
@@ -206,35 +204,33 @@ void KV_Cache::PickPastState(std::span<const int32_t> beam_indices, int index) {
 }
 
 void KV_Cache::PickPastState(std::span<const int32_t> beam_indices, int index) {
-  if (model_.score_type_ == Ort::TypeToTensorType<float>::type)
+  if (model_.config_->model.kv_type == Ort::TypeToTensorType<float>::type)
     PickPastState<float>(beam_indices, index);
   else
     PickPastState<Ort::Float16_t>(beam_indices, index);
 }
 
-Cross_Cache::Cross_Cache(Model& model, State& state,
-                         std::span<const char*> past_names, std::span<const char*> present_names)
+Cross_Cache::Cross_Cache(Model& model, State& state)
     : model_{model},
       state_{state},
-      layer_count_{model_.config_->num_hidden_layers},
-      past_names_{past_names},
-      present_names_{present_names},
-      shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->num_attention_heads, 1500, model.config_->hidden_size} {
+      layer_count_{model_.config_->model.num_hidden_layers},
+      shape_{state_.search_params_.batch_size * state_.search_params_.num_beams, model.config_->model.num_attention_heads, 1500, model.config_->model.hidden_size} {
   values_.reserve(layer_count_ * 2);
 
   for (int i = 0; i < layer_count_; ++i) {
-    char string[32];
-    values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
-    values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.score_type_));
+    values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));
+    values_.push_back(OrtValue::CreateTensor(*model_.allocator_device_, shape_, model_.config_->model.kv_type));
 
-    for (auto* name : past_names_) {
-      snprintf(string, std::size(string), name, i);
-      input_name_strings_.push_back(string);
-    }
-    for (auto* name : present_names_) {
-      snprintf(string, std::size(string), name, i);
-      output_name_strings_.push_back(string);
-    }
+    char string[64];
+    snprintf(string, std::size(string), model.config_->model.cross_past_names_key.c_str(), i);
+    input_name_strings_.push_back(string);
+    snprintf(string, std::size(string), model.config_->model.cross_past_names_value.c_str(), i);
+    input_name_strings_.push_back(string);
+
+    snprintf(string, std::size(string), model.config_->model.cross_present_names_key.c_str(), i);
+    output_name_strings_.push_back(string);
+    snprintf(string, std::size(string), model.config_->model.cross_present_names_value.c_str(), i);
+    output_name_strings_.push_back(string);
   }
 }
 
diff --git a/src/models/kv_cache.h b/src/models/kv_cache.h
index a0333d3a76..5358927bc3 100644
--- a/src/models/kv_cache.h
+++ b/src/models/kv_cache.h
@@ -13,10 +13,6 @@ struct KV_Cache_Combined {
   void PickPastState(std::span<const int32_t> beam_indices, int index);
   void PickPastState(std::span<const int32_t> beam_indices, int index);
 
-  // KV combined
-  const char *past_name_{"past_%d"};
-  const char *present_name_{"present_%d"};
-
 private:
 
   Model& model_;
@@ -32,7 +28,7 @@ struct KV_Cache_Combined {
 };
 
 struct KV_Cache {
-  KV_Cache(Model& model, State& state, std::span<const char*> past_names, std::span<const char*> present_names);
+  KV_Cache(Model& model, State& state);
 
   void AddEncoder(); // If model has an initial encoder step, this is used
   void Add();
@@ -47,9 +43,6 @@ struct KV_Cache {
   int layer_count_;
   size_t input_index_{~0U}, output_index_{~0U};
 
-  std::span<const char*> past_names_;     // past key name/past value name
-  std::span<const char*> present_names_;  // present key name/present value name
-
   std::array<int64_t, 4> shape_;
 
   std::unique_ptr<OrtValue> empty_past_;
@@ -59,7 +52,7 @@ struct KV_Cache {
 
 // Very similar to the KV_Cache, but is only created once at the encoder step, then used without modification for every decoder step
 struct Cross_Cache {
-  Cross_Cache(Model& model, State& state, std::span<const char*> past_names, std::span<const char*> present_names);
+  Cross_Cache(Model& model, State& state);
 
   void AddOutputs();
   void AddInputs();
@@ -69,8 +62,6 @@ struct Cross_Cache {
   State& state_;
   int layer_count_;
 
-  std::span<const char*> past_names_, present_names_;
-
   std::array<int64_t, 4> shape_;
 
   std::vector<std::unique_ptr<OrtValue>> values_;
diff --git a/src/models/llama.cpp b/src/models/llama.cpp
index 95db88359a..f40db82148 100644
--- a/src/models/llama.cpp
+++ b/src/models/llama.cpp
@@ -5,10 +5,9 @@ namespace Generators {
 
 Llama_Model::Llama_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
     : Model{std::move(config), ort_env, provider_options} {
-  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
+  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());
 
   InitDeviceAllocator(*session_decoder_);
-  InitLogits(*session_decoder_->GetOutputTypeInfo(0));
 }
 
 std::unique_ptr<State> Llama_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
diff --git a/src/models/llama.h b/src/models/llama.h
index 909a99f99e..ade3a53f7d 100644
--- a/src/models/llama.h
+++ b/src/models/llama.h
@@ -13,9 +13,6 @@ struct Llama_Model : Model {
   std::unique_ptr<State> CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) override;
 
   std::unique_ptr<OrtSession> session_decoder_;
-
-  std::array<const char*, 2> past_names_{"past_key_values.%d.key", "past_key_values.%d.value"};
-  std::array<const char*, 2> present_names_{"present.%d.key", "present.%d.value"};
 };
 
 struct Llama_State : State {
@@ -30,7 +27,7 @@ struct Llama_State : State {
 
   InputIDs<int64_t> input_ids_{model_, *this};
   Logits logits_{model_, *this};
-  KV_Cache kv_cache_{model_, *this, model_.past_names_, model_.present_names_};
+  KV_Cache kv_cache_{model_, *this};
   PositionIDs<int64_t> position_ids_;
 };
 
diff --git a/src/models/logits.cpp b/src/models/logits.cpp
index cfe9545aea..d627863921 100644
--- a/src/models/logits.cpp
+++ b/src/models/logits.cpp
@@ -8,8 +8,8 @@ Logits::Logits(Model& model, State& state)
     : model_{model},
       state_{state} {
 
-  logits_shape_ = {state_.search_params_.batch_size * state_.search_params_.num_beams, model_.logits_uses_seq_len_ ? state_.search_params_.sequence_length : 1, state_.search_params_.vocab_size};
-  logits_ = OrtValue::CreateTensor(*model.allocator_device_, logits_shape_, model_.score_type_);
+  logits_shape_ = {state_.search_params_.batch_size * state_.search_params_.num_beams, state_.search_params_.sequence_length, state_.search_params_.vocab_size};
+  logits_ = OrtValue::CreateTensor(*model.allocator_device_, logits_shape_, model_.config_->model.logits_type);
 }
 
 RoamingArray<float> Logits::Get() {
@@ -17,7 +17,7 @@ RoamingArray<float> Logits::Get() {
 
 #if USE_CUDA
   if (model_.device_type_ == DeviceType::CUDA) {
-    if (model_.score_type_ == Ort::TypeToTensorType<Ort::Float16_t>::type) {
+    if (model_.config_->model.logits_type == Ort::TypeToTensorType<Ort::Float16_t>::type) {
       ConvertFp16ToFp32(*model_.allocator_device_, model_.cuda_stream_, *logits_, logits32_);
       return gpu_span<float>{logits32_->GetTensorMutableData<float>(), type_shape->GetElementCount()};
     }
@@ -39,7 +39,7 @@ void Logits::Update() {
   // Resize the logits shape once if it doesn't match the decoder shape
   if (logits_shape_[1] != 1) {
     logits_shape_[1] = 1;
-    logits_ = OrtValue::CreateTensor(*model_.allocator_device_, logits_shape_, model_.score_type_);
+    logits_ = OrtValue::CreateTensor(*model_.allocator_device_, logits_shape_, model_.config_->model.logits_type);
     state_.outputs_[output_index_] = logits_.get();
   }
 }
diff --git a/src/models/mistral.cpp b/src/models/mistral.cpp
index 466b1fc321..456e0a4083 100644
--- a/src/models/mistral.cpp
+++ b/src/models/mistral.cpp
@@ -5,10 +5,9 @@ namespace Generators {
 
 Mistral_Model::Mistral_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
     : Model{std::move(config), ort_env, provider_options} {
-  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
+  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());
 
   InitDeviceAllocator(*session_decoder_);
-  InitLogits(*session_decoder_->GetOutputTypeInfo(0));
 }
 
 std::unique_ptr<State> Mistral_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
diff --git a/src/models/mistral.h b/src/models/mistral.h
index cbd8e0ad25..a1a0c9700b 100644
--- a/src/models/mistral.h
+++ b/src/models/mistral.h
@@ -13,9 +13,6 @@ struct Mistral_Model : Model {
   std::unique_ptr<State> CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) override;
 
   std::unique_ptr<OrtSession> session_decoder_;
-
-  std::array<const char*, 2> past_names_{"past_key.%d", "past_values.%d"};
-  std::array<const char*, 2> present_names_{"present_%d.key", "present_values.%d"};
 };
 
 struct Mistral_State : State {
@@ -30,7 +27,7 @@ struct Mistral_State : State {
 
   InputIDs<int64_t> input_ids_{model_, *this};
   Logits logits_{model_, *this};
-  KV_Cache kv_cache_{model_, *this, model_.past_names_, model_.present_names_};
+  KV_Cache kv_cache_{model_, *this};
   PositionIDs<int64_t> position_ids_;
 };
 
diff --git a/src/models/model.cpp b/src/models/model.cpp
index e3c7f9552c..57c82d1a7c 100644
--- a/src/models/model.cpp
+++ b/src/models/model.cpp
@@ -4,6 +4,7 @@
 #include "../search_cuda.h"
 #endif
 #include "model.h"
+#include "debugging.h"
 #include "gpt.h"
 #include "llama.h"
 #include "mistral.h"
@@ -60,18 +61,6 @@ void Model::InitDeviceAllocator(OrtSession& session) {
 #endif
 }
 
-void Model::InitLogits(OrtTypeInfo& info) {
-  auto& logits_tensor_info = info.GetTensorTypeAndShapeInfo();
-  auto logits_shape = logits_tensor_info.GetShape();
-  assert(logits_shape.size() == 3);
-  logits_uses_seq_len_ = logits_shape[1] == -1;
-  score_type_ = logits_tensor_info.GetElementType();
-
-  auto vocab_size = static_cast<int>(logits_shape[2]);
-  Unreferenced(vocab_size);
-  assert(config_->vocab_size == vocab_size);
-}
-
 std::vector<int32_t> Model::Generate(const SearchParams& params) {
   auto search = params.CreateSearch();
   auto state = CreateState(search->GetSequenceLengths(), params);
@@ -98,18 +87,18 @@ std::vector<int32_t> Model::Generate(const SearchParams& params) {
 std::unique_ptr<Model> CreateModel(OrtEnv& ort_env, const char* config_path, const ProviderOptions* provider_options) {
   auto config = std::make_unique<Config>(config_path);
 
-  if (config->model_type == "gpt2")
+  if (config->model.type == "gpt2")
     return std::make_unique<Gpt_Model>(std::move(config), ort_env, provider_options);
-  else if (config->model_type == "llama")
+  else if (config->model.type == "llama")
     return std::make_unique<Llama_Model>(std::move(config), ort_env, provider_options);
-  else if (config->model_type == "mistral")
+  else if (config->model.type == "mistral")
     return std::make_unique<Mistral_Model>(std::move(config), ort_env, provider_options);
-  else if (config->model_type == "phi2")
+  else if (config->model.type == "phi2")
     return std::make_unique<Phi2_Model>(std::move(config), ort_env, provider_options);
-  else if (config->model_type == "whisper")
+  else if (config->model.type == "whisper")
     return std::make_unique<Whisper_Model>(std::move(config), ort_env, provider_options);
 
-  throw std::runtime_error("Unsupported model_type in config.json: " + config->model_type);
+  throw std::runtime_error("Unsupported model_type in config.json: " + config->model.type);
 }
 
 #if USE_CUDA
diff --git a/src/models/model.h b/src/models/model.h
index d486d52544..acbf16b16a 100644
--- a/src/models/model.h
+++ b/src/models/model.h
@@ -40,12 +40,8 @@ struct Model {
   std::unique_ptr<Ort::Allocator> allocator_cuda_;
   Ort::Allocator* allocator_device_{};  // Can be CUDA or CPU based on the DeviceType in the model
 
-  bool logits_uses_seq_len_{};  // Logits shape is [... seq_len, vocab_size ] vs [... 1, vocab_size ]
-  ONNXTensorElementDataType score_type_;
-
  protected:
   void InitDeviceAllocator(OrtSession& session);
-  void InitLogits(OrtTypeInfo& info);
 };
 
 std::unique_ptr<Model> CreateModel(OrtEnv& ort_env, const char* config_path, const ProviderOptions* provider_options = nullptr);
diff --git a/src/models/phi2.cpp b/src/models/phi2.cpp
index 430fe6a7c8..cccc97ee22 100644
--- a/src/models/phi2.cpp
+++ b/src/models/phi2.cpp
@@ -5,10 +5,9 @@ namespace Generators {
 
 Phi2_Model::Phi2_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
     : Model{std::move(config), ort_env, provider_options} {
-  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
+  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());
 
   InitDeviceAllocator(*session_decoder_);
-  InitLogits(*session_decoder_->GetOutputTypeInfo(0));
 }
 
 std::unique_ptr<State> Phi2_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
diff --git a/src/models/whisper.cpp b/src/models/whisper.cpp
index 080c759d66..e836f76622 100644
--- a/src/models/whisper.cpp
+++ b/src/models/whisper.cpp
@@ -5,11 +5,10 @@ namespace Generators {
 
 Whisper_Model::Whisper_Model(std::unique_ptr<Config> config, OrtEnv& ort_env, const ProviderOptions* provider_options)
     : Model{std::move(config), ort_env, provider_options} {
-  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_decoder).c_str(), session_options_.get());
-  session_encoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model_encoder_decoder_init).c_str(), session_options_.get());
+  session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.decoder).c_str(), session_options_.get());
+  session_encoder_ = OrtSession::Create(ort_env, (config_->config_path / config_->model.encoder_decoder_init).c_str(), session_options_.get());
 
   InitDeviceAllocator(*session_decoder_);
-  InitLogits(*session_decoder_->GetOutputTypeInfo(0));
 }
 
 std::unique_ptr<State> Whisper_Model::CreateState(RoamingArray<int32_t> sequence_lengths, const SearchParams& params) {
diff --git a/src/models/whisper.h b/src/models/whisper.h
index fb8a05f2cf..11ac71cf95 100644
--- a/src/models/whisper.h
+++ b/src/models/whisper.h
@@ -12,11 +12,6 @@ struct Whisper_Model : Model {
 
   std::unique_ptr<OrtSession> session_decoder_;  // decoder.onnx
   std::unique_ptr<OrtSession> session_encoder_;  // encoder_decoder_init.onnx
-
-  std::array<const char*, 2> past_names_{"past_key_self_%d", "past_value_self_%d"};
-  std::array<const char*, 2> present_names_{"present_key_self_%d", "present_value_self_%d"};
-  std::array<const char*, 2> past_cross_names_{"past_key_cross_%d", "past_value_cross_%d"};
-  std::array<const char*, 2> present_cross_names_{"present_key_cross_%d", "present_value_cross_%d"};
 };
 
 struct Whisper_State : State {
@@ -31,8 +26,8 @@ struct Whisper_State : State {
 
   InputIDs<int32_t> decoder_input_ids_{model_, *this};
   Logits logits_{model_, *this};
-  KV_Cache kv_cache_{model_, *this, model_.past_names_, model_.present_names_};
-  Cross_Cache cross_cache_{model_, *this, model_.past_cross_names_, model_.present_cross_names_};
+  KV_Cache kv_cache_{model_, *this};
+  Cross_Cache cross_cache_{model_, *this};
   std::unique_ptr<OrtValue> encoder_hidden_states_;
 };
 }  // namespace Generators