Repository: microsoft/onnxruntime-genai
Pull request #1378: Changes how the device OrtAllocators work, use a global OrtSession instead
Author: login: RyanUnderhill
Created: 2025-04-04T08:26:55Z | Merged: 2025-04-24T22:57:11Z

Previously the code would store the OrtAllocator from the first OrtSession created. This apparently didn't work as expected. Even though the API to get an allocator from Ort is called CreateAllocator, it's really just returning a proxy to the OrtSession's allocator. When the OrtSession is destroyed, the OrtAllocator is invalid.



Diff:
diff --git a/src/cuda/cuda_common.h b/src/cuda/cuda_common.h
index 1fde0829f6..aa1b3f533f 100644
--- a/src/cuda/cuda_common.h
+++ b/src/cuda/cuda_common.h
@@ -42,8 +42,14 @@ struct cuda_stream_holder {
       (void)cudaStreamDestroy(v_);
   }
 
-  operator cudaStream_t() const { return v_; }
-  cudaStream_t get() const { return v_; }
+  operator cudaStream_t() const {
+    assert(v_);
+    return v_;
+  }
+  cudaStream_t get() const {
+    assert(v_);
+    return v_;
+  }
 
  private:
   cudaStream_t v_{};
diff --git a/src/cuda/interface.cpp b/src/cuda/interface.cpp
index 44dea9ad0b..f3fcee3e60 100644
--- a/src/cuda/interface.cpp
+++ b/src/cuda/interface.cpp
@@ -70,6 +70,7 @@ struct GpuMemory final : DeviceBuffer {
 
 struct CudaInterfaceImpl final : DeviceInterface {
   CudaInterfaceImpl() {
+    g_stream.Create();
   }
 
   ~CudaInterfaceImpl() {
diff --git a/src/generators.h b/src/generators.h
index 2f3db1cc4a..defe1eb4ce 100644
--- a/src/generators.h
+++ b/src/generators.h
@@ -123,7 +123,12 @@ struct OrtGlobals {
   OrtGlobals();
 
   std::unique_ptr<OrtEnv> env_;
-  std::unique_ptr<Ort::Allocator> allocator_device_[static_cast<int>(DeviceType::MAX)];
+
+  struct Allocator {
+    std::unique_ptr<Ort::Allocator> allocator_;
+    std::unique_ptr<OrtSession> session_;
+  };
+  Allocator device_allocators_[static_cast<int>(DeviceType::MAX)];
 
  private:
   OrtGlobals(const OrtGlobals&) = delete;
diff --git a/src/models/decoder_only.cpp b/src/models/decoder_only.cpp
index a153bc40ff..8405a80c6d 100644
--- a/src/models/decoder_only.cpp
+++ b/src/models/decoder_only.cpp
@@ -5,8 +5,7 @@ namespace Generators {
 DecoderOnly_Model::DecoderOnly_Model(std::unique_ptr<Config> config, OrtEnv& ort_env)
     : Model{std::move(config)} {
   session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / fs::path(config_->model.decoder.filename)).c_str(), session_options_.get());
-
-  InitDeviceAllocator(*session_decoder_);
+  session_info_.Add(*session_decoder_);
 }
 
 std::unique_ptr<State> DecoderOnly_Model::CreateState(DeviceSpan<int32_t> sequence_lengths_unk, const GeneratorParams& params) const {
diff --git a/src/models/decoder_only_pipeline.cpp b/src/models/decoder_only_pipeline.cpp
index 6e66e2ecbc..602596bda6 100644
--- a/src/models/decoder_only_pipeline.cpp
+++ b/src/models/decoder_only_pipeline.cpp
@@ -13,26 +13,10 @@ DecoderOnlyPipelineModel::DecoderOnlyPipelineModel(std::unique_ptr<Config> confi
   for (const auto& model : config_->model.decoder.pipeline) {
     sessions_.emplace_back(OrtSession::Create(ort_env, (config_->config_path / fs::path(model.filename)).c_str(),
                                               GetSessionOptions(model.model_id)));
-
-    if (!p_device_inputs_ && model.session_options.has_value()) {
-      const auto& provider_options = (*model.session_options).provider_options;
-      if (std::any_of(provider_options.begin(), provider_options.end(),
-                      [](const auto& elem) { return !elem.name.empty(); })) {
-        InitDeviceAllocator(*sessions_.back());
-      }
-    }
-  }
-
-  if (!p_device_inputs_) {
-    // If the device allocator has not been created, it implies all
-    // sessions are configured to run on CPU.
-    // Pick any session to create the device allocator.
-    // Device allocator is guaranteed to be the cpu allocator.
-    InitDeviceAllocator(*sessions_.front());
   }
 
   for (auto& session : sessions_) {
-    session_info_->Add(*session);
+    session_info_.Add(*session);
   }
 }
 
diff --git a/src/models/embeddings.cpp b/src/models/embeddings.cpp
index f000d5976e..71c056aca1 100644
--- a/src/models/embeddings.cpp
+++ b/src/models/embeddings.cpp
@@ -12,8 +12,8 @@ Embeddings::Embeddings(State& state, Embeddings::Mode mode, const std::string& n
       shape_{static_cast<int64_t>(state_.params_->search.batch_size) * state_.params_->search.num_beams,
              0, model_.config_->model.decoder.hidden_size},
       type_{mode == Embeddings::Mode::Input
-                ? model_.session_info_->GetInputDataType(name)
-                : model_.session_info_->GetOutputDataType(name)},
+                ? model_.session_info_.GetInputDataType(name)
+                : model_.session_info_.GetOutputDataType(name)},
       mode_{mode},
       name_{name} {
   // Embeddings are only transient inputs and outputs.
diff --git a/src/models/extra_inputs.cpp b/src/models/extra_inputs.cpp
index 25d850db1a..0699bb8028 100644
--- a/src/models/extra_inputs.cpp
+++ b/src/models/extra_inputs.cpp
@@ -15,7 +15,7 @@ PresetExtraInputs::PresetExtraInputs(State& state)
            }}} {}
 
 void PresetExtraInputs::Add() {
-  const auto input_names_vector = state_.model_.session_info_->GetInputNames();
+  const auto input_names_vector = state_.model_.session_info_.GetInputNames();
   const std::unordered_set<std::string> input_names(state_.input_names_.begin(), state_.input_names_.end());
   std::vector<std::string> unclaimed_input_names;
   // Add any model input for which we don't have a corresponding input in the state to the unclaimed_input_names
diff --git a/src/models/gpt.cpp b/src/models/gpt.cpp
index b8121d431c..c3cf025dbc 100644
--- a/src/models/gpt.cpp
+++ b/src/models/gpt.cpp
@@ -6,7 +6,7 @@ namespace Generators {
 Gpt_Model::Gpt_Model(std::unique_ptr<Config> config, OrtEnv& ort_env)
     : Model{std::move(config)} {
   session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / fs::path(config_->model.decoder.filename)).c_str(), session_options_.get());
-  InitDeviceAllocator(*session_decoder_);
+  session_info_.Add(*session_decoder_);
 }
 
 std::unique_ptr<State> Gpt_Model::CreateState(DeviceSpan<int32_t> sequence_lengths, const GeneratorParams& params) const {
diff --git a/src/models/input_ids.cpp b/src/models/input_ids.cpp
index a4522e7c02..1badd4eff4 100644
--- a/src/models/input_ids.cpp
+++ b/src/models/input_ids.cpp
@@ -8,24 +8,24 @@ DefaultInputIDs::DefaultInputIDs(State& state)
     : state_{state} {
   name_ = model_.config_->model.decoder.inputs.input_ids.c_str();
   shape_ = {state_.params_->BatchBeamSize(), 0};
-  type_ = model_.session_info_->GetInputDataType(name_);
+  type_ = model_.session_info_.GetInputDataType(name_);
 
-  if (model_.session_info_->HasInput(model_.config_->model.decoder.inputs.current_sequence_length) &&
-      model_.session_info_->HasInput(model_.config_->model.decoder.inputs.past_sequence_length)) {
+  if (model_.session_info_.HasInput(model_.config_->model.decoder.inputs.current_sequence_length) &&
+      model_.session_info_.HasInput(model_.config_->model.decoder.inputs.past_sequence_length)) {
     if (state_.params_->BatchBeamSize() != 1) {
       throw std::runtime_error("Batch size must be 1 for current_sequence_length and past_sequence_length inputs");
     }
     const std::array<int64_t, 1> current_sequence_length_shape{1};
     const std::array<int64_t, 2> past_sequence_length_shape{1, 1};
 
-    if (model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.current_sequence_length) != Ort::TypeToTensorType<int32_t> ||
-        model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length) != Ort::TypeToTensorType<int32_t>)
+    if (model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.current_sequence_length) != Ort::TypeToTensorType<int32_t> ||
+        model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length) != Ort::TypeToTensorType<int32_t>)
       throw std::runtime_error("current_sequence_length and past_sequence_length must be int32");
 
-    current_sequence_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, current_sequence_length_shape, model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.current_sequence_length));
+    current_sequence_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, current_sequence_length_shape, model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.current_sequence_length));
     *current_sequence_length_->GetTensorMutableData<int32_t>() = 0;
 
-    past_sequence_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, past_sequence_length_shape, model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length));
+    past_sequence_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, past_sequence_length_shape, model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length));
     *past_sequence_length_->GetTensorMutableData<int32_t>() = -1;
   }
 
@@ -116,27 +116,27 @@ WindowedInputIDs::WindowedInputIDs(State& state) : state_{state} {
 
   window_size_ = model_.config_->model.decoder.sliding_window->window_size;
   shape_ = {1, model_.config_->model.decoder.sliding_window->window_size};
-  type_ = model_.session_info_->GetInputDataType(name_);
+  type_ = model_.session_info_.GetInputDataType(name_);
 
   if (type_ != Ort::TypeToTensorType<int32_t> && type_ != Ort::TypeToTensorType<int64_t>) {
     throw std::runtime_error("WindowedInputIDs only supports int32_t and int64_t input_ids.");
   }
 
-  if (model_.session_info_->HasInput(model_.config_->model.decoder.inputs.total_sequence_length) &&
-      model_.session_info_->HasInput(model_.config_->model.decoder.inputs.past_sequence_length)) {
+  if (model_.session_info_.HasInput(model_.config_->model.decoder.inputs.total_sequence_length) &&
+      model_.session_info_.HasInput(model_.config_->model.decoder.inputs.past_sequence_length)) {
     const std::array<int64_t, 1> total_sequence_length_shape{1};
     const std::array<int64_t, 2> past_sequence_length_shape{1, 1};
 
-    if (model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.total_sequence_length) != Ort::TypeToTensorType<int32_t> ||
-        model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length) != Ort::TypeToTensorType<int32_t>)
+    if (model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.total_sequence_length) != Ort::TypeToTensorType<int32_t> ||
+        model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length) != Ort::TypeToTensorType<int32_t>)
       throw std::runtime_error("total_sequence_length and past_sequence_length must be int32");
 
     total_sequence_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, total_sequence_length_shape,
-                                                    model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.total_sequence_length));
+                                                    model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.total_sequence_length));
     *total_sequence_length_->GetTensorMutableData<int32_t>() = state_.params_->search.max_length;
 
     past_sequence_length_ = OrtValue::CreateTensor(model_.allocator_cpu_, past_sequence_length_shape,
-                                                   model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length));
+                                                   model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.past_sequence_length));
     *past_sequence_length_->GetTensorMutableData<int32_t>() = -1;
   }
 }
diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp
index 3f333c3264..581d12caa0 100644
--- a/src/models/kv_cache.cpp
+++ b/src/models/kv_cache.cpp
@@ -22,7 +22,7 @@ CombinedKeyValueCache::CombinedKeyValueCache(State& state)
   }
 
   // Derive the KV data type from the KV input 0
-  type_ = model_.session_info_->GetInputDataType(input_name_strings_[0]);
+  type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]);
 
   empty_past_ = OrtValue::CreateTensor(Allocator(), shape_, type_);
   shape_[3] = 0;
@@ -166,7 +166,7 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state)
   }
 
   // Derive the KV data type from the KV input 0
-  type_ = model_.session_info_->GetInputDataType(input_name_strings_[0]);
+  type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]);
   empty_past_ = OrtValue::CreateTensor(Allocator(), shape_, type_);
 
   if (state_.params_->use_graph_capture && !past_present_share_buffer_) {
@@ -344,7 +344,7 @@ CrossCache::CrossCache(State& state)
   }
 
   // Derive the KV data type from the KV input 0
-  type_ = model_.session_info_->GetInputDataType(input_name_strings_[0]);
+  type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]);
 
   for (int i = 0; i < layer_count_; ++i) {
     values_.push_back(OrtValue::CreateTensor(Allocator(), shape_, type_));
@@ -403,7 +403,7 @@ void ModelManagedKeyValueCache::RewindTo(size_t index) {
 namespace {
 
 bool IsCacheNeeded(const Model& model) {
-  return model.session_info_->HasInput(ComposeKeyValueName(model.config_->model.decoder.inputs.past_key_names, 0));
+  return model.session_info_.HasInput(ComposeKeyValueName(model.config_->model.decoder.inputs.past_key_names, 0));
 }
 
 }  // namespace
diff --git a/src/models/logits.cpp b/src/models/logits.cpp
index 14efe6c038..a12f69c763 100644
--- a/src/models/logits.cpp
+++ b/src/models/logits.cpp
@@ -10,7 +10,7 @@ namespace Generators {
 Logits::Logits(State& state)
     : state_{state},
       shape_{static_cast<int64_t>(state_.params_->BatchBeamSize()), 0, model_.config_->model.vocab_size},
-      type_{model_.session_info_->GetOutputDataType(model_.config_->model.decoder.outputs.logits)} {
+      type_{model_.session_info_.GetOutputDataType(model_.config_->model.decoder.outputs.logits)} {
   output_raw_ = std::make_unique<Tensor>(model_.p_device_inputs_, type_);
 
   if (model_.p_device_inputs_->GetType() == DeviceType::CUDA && !model_.config_->model.eos_token_ids.empty()) {
diff --git a/src/models/model.cpp b/src/models/model.cpp
index 5c37677d06..1e54214d11 100644
--- a/src/models/model.cpp
+++ b/src/models/model.cpp
@@ -273,33 +273,162 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else {
+      // For providers that go through the extensible AppendExecutionProvider API:
+
+      if (provider_options.name == "QNN") {
+        session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+        // TODO set device_type_ in a less hacky way.
+        // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+        // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+        if (const auto opt_it = std::find_if(provider_options.options.begin(), provider_options.options.end(),
+                                             [](const auto& pair) { return pair.first == "enable_htp_shared_memory_allocator"; });
+            opt_it != provider_options.options.end() && opt_it->second == "1") {
+          p_device = GetDeviceInterface(DeviceType::QNN);
+        }
+      } else if (provider_options.name == "WebGPU")
+        p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      else if (provider_options.name == "OpenVINO")
+        p_device = GetDeviceInterface(DeviceType::OpenVINO);
+      else if (provider_options.name == "VitisAI") {
+        session_options.AddConfigEntry("session.inter_op.allow_spinning", "0");
+        session_options.AddConfigEntry("session.intra_op.allow_spinning", "0");
+      }
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      session_options.AppendExecutionProvider(provider_options.name.c_str(), keys.data(), values.data(), keys.size());
+    }
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceInterface& device) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   // OpenVINO delegates to the CPU device allocator
+  auto type = device.GetType();
   if (type == DeviceType::CPU || type == DeviceType::OpenVINO)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
-  static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared", "OpenVINO (Not used, see above)"};
+  // Allocator lifetime is tied to the execution provider lifetime, which is typically per Session.
+  // We create a global dummy Session using a trivial model with the required EP in order to get the allocator so we can
+  // re-use it for all models.
+  // This ensures memory allocated on-device for model inputs/outputs is valid for the lifetime of GenAI.
+
+  // Names for the device types used by 'SetProviderSessionOptions'
+  static const char* device_type_names[] = {"CPU (Not used, see above)", "cuda", "dml", "webgpu", "qnn", "OpenVINO (Not used, see above)"};
   static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));
 
-  auto name = device_type_names[static_cast<int>(type)];
+  // Create an OrtSessionOptions and set the options to use the DeviceType we're using here
+  auto session_options = OrtSessionOptions::Create();
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_type_names[static_cast<int>(type)], {}});
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  session_options->SetLogSeverityLevel(ORT_LOGGING_LEVEL_ERROR);  // Errors only here, as warnings are not useful to the user
+
+  allocator.session_ = OrtSession::Create(GetOrtEnv(), g_trivial_model, sizeof(g_trivial_model), session_options.get());
+
+  // Names for the device memory types used by 'OrtMemoryInfo::Create'
+  static const char* device_memory_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared", "OpenVINO (Not used, see above)"};
+  static_assert(std::size(device_memory_type_names) == static_cast<size_t>(DeviceType::MAX));
+
+  // Get the allocator from the OrtSession for the DeviceType (it's called 'AllocatorCreate' but it's really 'AllocatorGet')
+  auto name = device_memory_type_names[static_cast<int>(type)];
   auto memory_info = OrtMemoryInfo::Create(name, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
-  device = Ort::Allocator::Create(session, *memory_info);
-  if (!device)
+  allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info);
+  if (!allocator.allocator_) {
+    allocator = {};  // Reset everything just to be safe
     throw std::runtime_error("Unexpected failure to create device memory allocator for " + std::string(name));
-  GetDeviceInterface(type)->InitOrt(*Ort::api, *device);  // Necessary for any shared library providers so they can access Ort::api
-}
-
-SessionInfo::SessionInfo(OrtSession& session) {
-  Add(session);
+  }
+  device.InitOrt(*Ort::api, *allocator.allocator_);
 }
 
 void SessionInfo::Add(OrtSession& session) {
@@ -368,12 +497,7 @@ std::vector<const char*> SessionInfo::GetOutputSymbolicShape(const std::string&
 
 Model::Model(std::unique_ptr<Config> config) : config_{std::move(config)} {
   CreateSessionOptions();
-}
-
-Model::~Model() = default;
-
-void Model::InitDeviceAllocator(OrtSession& session) {
-  EnsureDeviceOrtInit(session, p_device_->GetType());
+  EnsureDeviceOrtInit(*p_device_);
 
   // Only CUDA and DML does every input on the device
   if (p_device_->GetType() == DeviceType::CUDA || p_device_->GetType() == DeviceType::DML)
@@ -383,10 +507,10 @@ void Model::InitDeviceAllocator(OrtSession& session) {
 
   // The kvcache is always allocated in device memory
   p_device_kvcache_ = p_device_;
-
-  session_info_ = std::make_unique<SessionInfo>(session);
 }
 
+Model::~Model() = default;
+
 void Model::CreateSessionOptionsFromConfig(const Config::SessionOptions& config_session_options,
                                            OrtSessionOptions& session_options,
                                            bool is_primary_session_options,
@@ -484,99 +608,7 @@ void Model::CreateSessionOptionsFromConfig(const Config::SessionOptions& config_
     session_options.SetGraphOptimizationLevel(config_session_options.graph_optimization_level.value());
   }
 
-  for (auto& provider_options : config_session_options.provider_options) {
-    if (provider_options.name == "cuda") {
-      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
-      std::vector<const char*> keys, values;
-      for (auto& option : provider_options.options) {
-        keys.emplace_back(option.first.c_str());
-        values.emplace_back(option.second.c_str());
-      }
-      ort_provider_options->Update(keys.data(), values.data(), keys.size());
-
-      // Device type determines the scoring device.
-      // Only use the primary session options to determine the device type
-      if (is_primary_session_options) {
-        p_device_ = GetDeviceInterface(DeviceType::CUDA);
-
-        // Create and set our cudaStream_t
-        ort_provider_options->UpdateValue("user_compute_stream", p_device_->GetCudaStream());
-      }
-
-      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
-    } else if (provider_options.name == "rocm") {
-      OrtROCMProviderOptions ort_provider_options;
-
-      std::vector<const char*> keys, values;
-      for (auto& option : provider_options.options) {
-        keys.emplace_back(option.first.c_str());
-        values.emplace_back(option.second.c_str());
-      }
-
-      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
-      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
-#if USE_DML
-    } else if (provider_options.name == "dml") {
-      if (!GetDmlInterface()) {
-        LUID device_luid{};
-        LUID* p_device_luid{};
-        for (const auto& [name, value] : provider_options.options) {
-          if (name == "luid") {
-            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
-              device_luid.HighPart = std::stol(value.substr(0, separator_position));
-              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
-              p_device_luid = &device_luid;
-            }
-          }
-        }
-
-        InitDmlInterface(p_device_luid);
-      }
-
-      if (!disable_graph_capture) {
-        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
-        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
-      }
-
-      SetDmlProvider(session_options);
-
-      if (is_primary_session_options)
-        p_device_ = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
-#endif
-    } else {
-      // For providers that go through the extensible AppendExecutionProvider API:
-
-      if (provider_options.name == "QNN") {
-        session_options.AddConfigEntry("ep.share_ep_contexts", "1");
-        // TODO set device_type_ in a less hacky way.
-        // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
-        // on the other hand, not sure if is_primary_session_options is the right thing to check here.
-        if (const auto opt_it = std::find_if(provider_options.options.begin(), provider_options.options.end(),
-                                             [](const auto& pair) { return pair.first == "enable_htp_shared_memory_allocator"; });
-            opt_it != provider_options.options.end() && opt_it->second == "1") {
-          p_device_ = GetDeviceInterface(DeviceType::QNN);
-        }
-      }
-
-      else if (provider_options.name == "WebGPU")
-        p_device_ = GetDeviceInterface(DeviceType::WEBGPU);
-
-      else if (provider_options.name == "OpenVINO")
-        p_device_ = GetDeviceInterface(DeviceType::OpenVINO);
-
-      else if (provider_options.name == "VitisAI") {
-        session_options.AddConfigEntry("session.inter_op.allow_spinning", "0");
-        session_options.AddConfigEntry("session.intra_op.allow_spinning", "0");
-      }
-
-      std::vector<const char*> keys, values;
-      for (auto& option : provider_options.options) {
-        keys.emplace_back(option.first.c_str());
-        values.emplace_back(option.second.c_str());
-      }
-      session_options.AppendExecutionProvider(provider_options.name.c_str(), keys.data(), values.data(), keys.size());
-    }
-  }
+  p_device_ = SetProviderSessionOptions(session_options, config_session_options.provider_options, is_primary_session_options, disable_graph_capture);
 
   // Fallback to CPU if no provider specific interface was set
   if (!p_device_)
@@ -611,7 +643,7 @@ std::shared_ptr<Tokenizer> Model::CreateTokenizer() const {
 }
 
 std::shared_ptr<MultiModalProcessor> Model::CreateMultiModalProcessor() const {
-  return std::make_shared<MultiModalProcessor>(*config_, *session_info_);
+  return std::make_shared<MultiModalProcessor>(*config_, session_info_);
 }
 
 std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, const char* config_path, const RuntimeSettings* settings /*= nullptr*/) {
diff --git a/src/models/model.h b/src/models/model.h
index 292b6ef3e6..ce2e8515af 100644
--- a/src/models/model.h
+++ b/src/models/model.h
@@ -109,7 +109,7 @@ struct MultiModalProcessor : std::enable_shared_from_this<MultiModalProcessor>,
 };
 
 struct SessionInfo {
-  SessionInfo(OrtSession& session);
+  SessionInfo() = default;
 
   void Add(OrtSession& session);
 
@@ -151,10 +151,9 @@ struct Model : std::enable_shared_from_this<Model>, LeakChecked<Model>, External
 
   Ort::Allocator& allocator_cpu_{GetDeviceInterface(DeviceType::CPU)->GetAllocator()};
 
-  std::unique_ptr<SessionInfo> session_info_;
+  SessionInfo session_info_;
 
  protected:
-  void InitDeviceAllocator(OrtSession& session);
   void CreateSessionOptions();
 
   void CreateSessionOptionsFromConfig(const Config::SessionOptions& config_session_options,
diff --git a/src/models/multi_modal.cpp b/src/models/multi_modal.cpp
index ca4baa17ec..f4f3a210bd 100644
--- a/src/models/multi_modal.cpp
+++ b/src/models/multi_modal.cpp
@@ -84,13 +84,13 @@ MultiModalLanguageModel::MultiModalLanguageModel(std::unique_ptr<Config> config,
   decoder_session_ = OrtSession::Create(
       ort_env, (config_->config_path / fs::path(config_->model.decoder.filename)).c_str(), session_options_.get());
 
-  InitDeviceAllocator(*decoder_session_);
-  session_info_->Add(*embedding_session_);
+  session_info_.Add(*decoder_session_);
+  session_info_.Add(*embedding_session_);
   if (speech) {
-    session_info_->Add(*speech_session_);
+    session_info_.Add(*speech_session_);
   }
   if (vision) {
-    session_info_->Add(*vision_session_);
+    session_info_.Add(*vision_session_);
   }
 }
 
diff --git a/src/models/multi_modal_features.cpp b/src/models/multi_modal_features.cpp
index f6f983166a..2dccd6000f 100644
--- a/src/models/multi_modal_features.cpp
+++ b/src/models/multi_modal_features.cpp
@@ -10,13 +10,13 @@ MultiModalFeatures::MultiModalFeatures(State& state, MultiModalFeatures::Mode mo
                                        int64_t batch_size, int64_t num_feature_tokens)
     : state_{state},
       type_{mode == MultiModalFeatures::Mode::Input
-                ? model_.session_info_->GetInputDataType(name)
-                : model_.session_info_->GetOutputDataType(name)},
+                ? model_.session_info_.GetInputDataType(name)
+                : model_.session_info_.GetOutputDataType(name)},
       mode_{mode},
       name_{name} {
   const auto dims = mode_ == MultiModalFeatures::Mode::Input
-                        ? model_.session_info_->GetInputSymbolicShape(name).size()
-                        : model_.session_info_->GetOutputSymbolicShape(name).size();
+                        ? model_.session_info_.GetInputSymbolicShape(name).size()
+                        : model_.session_info_.GetOutputSymbolicShape(name).size();
 
   // If the model expects 3 dimensions, add a batch dimension
   if (dims == 3) {
diff --git a/src/models/position_inputs.cpp b/src/models/position_inputs.cpp
index 1d394838bd..a0801d844c 100644
--- a/src/models/position_inputs.cpp
+++ b/src/models/position_inputs.cpp
@@ -7,20 +7,20 @@ namespace Generators {
 DefaultPositionInputs::DefaultPositionInputs(const Model& model, State& state, DeviceSpan<int32_t> sequence_lengths_unk)
     : model_{model},
       state_{state} {
-  has_mask_input_ = model_.session_info_->HasInput(model_.config_->model.decoder.inputs.attention_mask);
-  has_posid_input_ = model_.session_info_->HasInput(model_.config_->model.decoder.inputs.position_ids);
+  has_mask_input_ = model_.session_info_.HasInput(model_.config_->model.decoder.inputs.attention_mask);
+  has_posid_input_ = model_.session_info_.HasInput(model_.config_->model.decoder.inputs.position_ids);
 
   type_ = Ort::TypeToTensorType<int32_t>;
   if (has_mask_input_) {
-    type_ = model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.attention_mask);
+    type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.attention_mask);
   }
   if (has_posid_input_) {
     if (has_mask_input_) {
-      if (model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.position_ids) != type_) {
+      if (model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.position_ids) != type_) {
         throw std::runtime_error("position_ids & attention_mask must have the same data type");
       }
     }
-    type_ = model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.position_ids);
+    type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.position_ids);
   }
 
   if (type_ != Ort::TypeToTensorType<int32_t> && type_ != Ort::TypeToTensorType<int64_t>)
@@ -293,8 +293,8 @@ void DefaultPositionInputs::RewindMask(size_t index) {
 // TODO: SlidingWindow does not support graph capture
 WindowedPositionInputs::WindowedPositionInputs(State& state)
     : state_{state} {
-  has_posid_input_ = model_.session_info_->HasInput(model_.config_->model.decoder.inputs.position_ids);
-  has_mask_input_ = model_.session_info_->HasInput(model_.config_->model.decoder.inputs.attention_mask);
+  has_posid_input_ = model_.session_info_.HasInput(model_.config_->model.decoder.inputs.position_ids);
+  has_mask_input_ = model_.session_info_.HasInput(model_.config_->model.decoder.inputs.attention_mask);
 
   if (has_posid_input_ || has_mask_input_) {
     if (!model_.config_->model.decoder.sliding_window.has_value()) {
@@ -308,7 +308,7 @@ WindowedPositionInputs::WindowedPositionInputs(State& state)
   }
 
   if (has_posid_input_) {
-    position_ids_type_ = model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.position_ids);
+    position_ids_type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.position_ids);
     if (position_ids_type_ != Ort::TypeToTensorType<int32_t>)
       throw std::runtime_error("WindowedPositionInputs only supports int32_t position_ids");
 
@@ -316,7 +316,7 @@ WindowedPositionInputs::WindowedPositionInputs(State& state)
   }
 
   if (has_mask_input_) {
-    attention_mask_type_ = model_.session_info_->GetInputDataType(model_.config_->model.decoder.inputs.attention_mask);
+    attention_mask_type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.attention_mask);
     if (attention_mask_type_ != Ort::TypeToTensorType<int32_t>)
       throw std::runtime_error("WindowedPositionInputs only supports int32_t attention_mask");
 
diff --git a/src/models/whisper.cpp b/src/models/whisper.cpp
index a588e0e7f1..2b5abd2cad 100644
--- a/src/models/whisper.cpp
+++ b/src/models/whisper.cpp
@@ -11,8 +11,8 @@ Whisper_Model::Whisper_Model(std::unique_ptr<Config> config, OrtEnv& ort_env)
   session_encoder_ = OrtSession::Create(ort_env, (config_->config_path / fs::path(config_->model.encoder_decoder_init.filename)).c_str(), session_options_.get());
   session_decoder_ = OrtSession::Create(ort_env, (config_->config_path / fs::path(config_->model.decoder.filename)).c_str(), session_options_.get());
 
-  InitDeviceAllocator(*session_decoder_);
-  session_info_->Add(*session_encoder_);
+  session_info_.Add(*session_decoder_);
+  session_info_.Add(*session_encoder_);
 }
 
 std::unique_ptr<State> Whisper_Model::CreateState(DeviceSpan<int32_t> sequence_lengths, const GeneratorParams& params) const {
@@ -50,7 +50,7 @@ Whisper_State::Whisper_State(const Whisper_Model& model, DeviceSpan<int32_t> seq
     auto alignment_heads_data_size = alignment_heads_elements * alignment_heads_element_size;
     cudaMemcpyAsync(alignment_heads_->GetTensorMutableRawData(), inputs.alignment_heads->ort_tensor_->GetTensorRawData(), alignment_heads_data_size, cudaMemcpyHostToDevice, model_.cuda_stream_);
 
-    auto cross_qk_type = model_.session_info_->GetOutputDataType("output_cross_qk_0");
+    auto cross_qk_type = model_.session_info_.GetOutputDataType("output_cross_qk_0");
     auto cross_qk_shape = std::array<int64_t, 4>{params_->BatchBeamSize(), alignment_heads_->GetTensorTypeAndShapeInfo()->GetShape()[0], params_->search.max_length, 1500};
     cross_qk_search_buffer_ = OrtValue::CreateTensor(*model_.allocator_device_, cross_qk_shape, cross_qk_type);
 
@@ -61,7 +61,7 @@ Whisper_State::Whisper_State(const Whisper_Model& model, DeviceSpan<int32_t> seq
 #endif
   }
 
-  auto hidden_states_type = model_.session_info_->GetOutputDataType("encoder_hidden_states");
+  auto hidden_states_type = model_.session_info_.GetOutputDataType("encoder_hidden_states");
   auto encoder_hidden_states_shape = std::array<int64_t, 3>{decoder_input_ids_.GetShape()[0], 1500, static_cast<int64_t>(model_.config_->model.decoder.num_attention_heads) * model_.config_->model.decoder.head_size};
   encoder_hidden_states_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), encoder_hidden_states_shape, hidden_states_type);
 
@@ -87,7 +87,7 @@ Whisper_State::Whisper_State(const Whisper_Model& model, DeviceSpan<int32_t> seq
   {
     auto layer_count = model_.config_->model.decoder.num_hidden_layers;
     std::array<int64_t, 4> shape{params_->BatchBeamSize(), model_.config_->model.decoder.num_attention_heads, 0, model_.config_->model.decoder.head_size};
-    auto type = model_.session_info_->GetOutputDataType(output_names_[kv_cache_indices]);
+    auto type = model_.session_info_.GetOutputDataType(output_names_[kv_cache_indices]);
 
     for (int i = 0; i < layer_count * 2; i++) {
       init_presents_.emplace_back(OrtValue::CreateTensor(model_.p_device_->GetAllocator(), shape, type));
@@ -250,13 +250,13 @@ DeviceSpan<float> Whisper_State::Run(int current_length, DeviceSpan<int32_t>& ne
       cross_cache_.AddInputs();
       run_state_ = RunState::Decoder;
 
-      if (model_.session_info_->HasInput("past_sequence_length")) {
+      if (model_.session_info_.HasInput("past_sequence_length")) {
         past_sequence_length_ = OrtValue::CreateTensor<int32_t>(model_.allocator_cpu_, std::array<int64_t, 1>{1});
         input_names_.push_back("past_sequence_length");
         inputs_.push_back(past_sequence_length_.get());
       }
 
-      if (model_.session_info_->HasInput("beam_width")) {
+      if (model_.session_info_.HasInput("beam_width")) {
         beam_width_ = OrtValue::CreateTensor<int32_t>(model_.allocator_cpu_, std::array<int64_t, 1>{1});
         input_names_.push_back("beam_width");
         inputs_.push_back(beam_width_.get());
@@ -265,7 +265,7 @@ DeviceSpan<float> Whisper_State::Run(int current_length, DeviceSpan<int32_t>& ne
         *data = params_->search.num_beams;
       }
 
-      if (model_.session_info_->HasInput("cache_indirection")) {
+      if (model_.session_info_.HasInput("cache_indirection")) {
         cache_indirection_ = OrtValue::CreateTensor<int32_t>(model_.p_device_->GetAllocator(), std::array<int64_t, 3>{params_->search.batch_size, params_->search.num_beams, params_->search.max_length});
         cache_indirection_index_ = inputs_.size();
         input_names_.push_back("cache_indirection");
@@ -274,9 +274,9 @@ DeviceSpan<float> Whisper_State::Run(int current_length, DeviceSpan<int32_t>& ne
         ByteWrapTensor(*model_.p_device_, *cache_indirection_).Zero();
       }
 
-      if (model_.session_info_->HasOutput("output_cross_qk_0")) {
+      if (model_.session_info_.HasOutput("output_cross_qk_0")) {
         auto layer_count = model_.config_->model.decoder.num_hidden_layers;
-        auto type = model_.session_info_->GetOutputDataType("output_cross_qk_0");
+        auto type = model_.session_info_.GetOutputDataType("output_cross_qk_0");
         std::array<int64_t, 4> shape{params_->BatchBeamSize(), model_.config_->model.decoder.num_attention_heads, 1, 1500};
         for (int i = 0; i < layer_count; i++) {
           char string[64];
@@ -333,7 +333,7 @@ void Whisper_State::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, Device
       beam_indices.CopyCpuToDevice();
     }
     std::unique_ptr<OrtValue> new_cache_indirection;
-    auto cache_indirection_type = model_.session_info_->GetInputDataType("cache_indirection");
+    auto cache_indirection_type = model_.session_info_.GetInputDataType("cache_indirection");
     auto cache_indirection_shape = std::array<int64_t, 3>{params_->search.batch_size, params_->search.num_beams, params_->search.max_length};
     new_cache_indirection = OrtValue::CreateTensor(*model_.allocator_device_, cache_indirection_shape, cache_indirection_type);
 
@@ -390,7 +390,7 @@ void Whisper_State::Finalize() {
 
     // Instantiate final output for cross QKs
     auto num_alignment_heads = alignment_heads_->GetTensorTypeAndShapeInfo()->GetShape()[0];
-    auto cross_qk_type = model_.session_info_->GetOutputDataType("output_cross_qk_0");
+    auto cross_qk_type = model_.session_info_.GetOutputDataType("output_cross_qk_0");
     auto cross_qk_shape = std::array<int64_t, 5>{params_->search.batch_size, params_->search.num_return_sequences, num_alignment_heads, decoded_length, 1500};
     cross_qk_final_ = OrtValue::CreateTensor(*model_.allocator_device_, cross_qk_shape, cross_qk_type);
 
diff --git a/src/models/windowed_kv_cache.cpp b/src/models/windowed_kv_cache.cpp
index 8b147f65df..a039b6109d 100644
--- a/src/models/windowed_kv_cache.cpp
+++ b/src/models/windowed_kv_cache.cpp
@@ -39,7 +39,7 @@ WindowedKeyValueCache::WindowedKeyValueCache(State& state)
     output_name_strings_.emplace_back(ComposeKeyValueName(model_.config_->model.decoder.outputs.present_value_names, i));
   }
 
-  type_ = model_.session_info_->GetInputDataType(input_name_strings_[0]);
+  type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]);
   if (type_ != Ort::TypeToTensorType<uint8_t>) {
     throw std::runtime_error("Expected input data type to be uint8_t for WindowedKeyValueCache. Actual: " +
                              std::to_string(type_));
diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 14e7e445dc..515101af61 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -9,7 +9,7 @@
 """
 
 from onnx import helper, numpy_helper, TensorProto, external_data_helper, save_model
-from onnxruntime.quantization.matmul_4bits_quantizer import MatMul4BitsQuantizer, QuantFormat
+from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer, QuantFormat
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import numpy as np
 import torch
@@ -489,7 +489,7 @@ def save_model(self, out_dir):
         )
 
     def to_int4(self, model):
-        quant = MatMul4BitsQuantizer(
+        quant = MatMulNBitsQuantizer(
             model=model,
             block_size=self.quant_attrs["int4"]["block_size"],
             is_symmetric=self.quant_attrs["int4"]["is_symmetric"],
diff --git a/test/python/cpu/ort/requirements.txt b/test/python/cpu/ort/requirements.txt
index 1f9b1e01b5..11cfe13e19 100644
--- a/test/python/cpu/ort/requirements.txt
+++ b/test/python/cpu/ort/requirements.txt
@@ -1,2 +1,2 @@
-# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
-onnxruntime==1.20.1
\ No newline at end of file
+-i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+onnxruntime==1.22.0.dev20250423001
\ No newline at end of file
diff --git a/test/python/cuda/ort/requirements.txt b/test/python/cuda/ort/requirements.txt
index a93f533766..c8da655bb4 100644
--- a/test/python/cuda/ort/requirements.txt
+++ b/test/python/cuda/ort/requirements.txt
@@ -1,2 +1,2 @@
-# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
-onnxruntime-gpu==1.20.1
\ No newline at end of file
+-i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+onnxruntime-gpu==1.22.0.dev20250423001
\ No newline at end of file
diff --git a/test/python/directml/ort/requirements.txt b/test/python/directml/ort/requirements.txt
index 72457f72cb..3057bd9089 100644
--- a/test/python/directml/ort/requirements.txt
+++ b/test/python/directml/ort/requirements.txt
@@ -1,2 +1,2 @@
-# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
-onnxruntime-directml==1.20.1
\ No newline at end of file
+-i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+onnxruntime-directml==1.22.0.dev20250423001
\ No newline at end of file
diff --git a/test/python/macos/ort/requirements.txt b/test/python/macos/ort/requirements.txt
index 1f9b1e01b5..11cfe13e19 100644
--- a/test/python/macos/ort/requirements.txt
+++ b/test/python/macos/ort/requirements.txt
@@ -1,2 +1,2 @@
-# -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
-onnxruntime==1.20.1
\ No newline at end of file
+-i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
+onnxruntime==1.22.0.dev20250423001
\ No newline at end of file


Reviews:
user: login: skottmckay
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T08:52:53Z
user: login: RyanUnderhill
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T08:55:36Z
user: login: baijumeswani
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T20:17:47Z
user: login: baijumeswani
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T20:35:04Z
user: login: baijumeswani
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T20:43:29Z
user: login: baijumeswani
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-04T21:04:05Z
user: login: baijumeswani
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-04T21:04:39Z
user: login: baijumeswani
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-04T21:24:34Z
user: login: skottmckay
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-05T05:15:20Z
user: login: skottmckay
state: DISMISSED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T21:10:03Z
user: login: RyanUnderhill
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T05:53:41Z
user: login: baijumeswani
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-23T23:39:10Z
user: login: baijumeswani
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-24T00:54:47Z
user: login: baijumeswani
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-24T22:25:20Z

Review comments:
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
+  auto session_options = OrtSessionOptions::Create();
+
+  const char* device_name{};
+  switch (type) {
+    case DeviceType::CUDA:
+      device_name = "cuda";
+      break;
+    case DeviceType::DML:
+      device_name = "dml";
+      break;
+    case DeviceType::WEBGPU:
+      device_name = "webgpu";
+      break;
+    case DeviceType::QNN:
+      device_name = "qnn";
+      break;
+    default:
+      throw std::runtime_error("Unknown device type");
+  }
+
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_name, {}});
+
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  allocator.session_ = OrtSession::Create(GetOrtEnv(), g_trivial_model, sizeof(g_trivial_model), session_options.get());
path: src/models/model.cpp
user: login: skottmckay
body: A comment explaining the setup to the new reader would be great so they don't have to infer 'why' we're doing all this. 

We create a session with the required EP. We put that session into the globals so it remains valid for the lifetime of GenAI. Due to that the allocator we get from that session is guaranteed to remain valid and can be used for any other models we load. 
created_at: 2025-04-04T08:48:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
+  auto session_options = OrtSessionOptions::Create();
+
+  const char* device_name{};
+  switch (type) {
+    case DeviceType::CUDA:
+      device_name = "cuda";
+      break;
+    case DeviceType::DML:
+      device_name = "dml";
+      break;
+    case DeviceType::WEBGPU:
+      device_name = "webgpu";
+      break;
+    case DeviceType::QNN:
+      device_name = "qnn";
+      break;
+    default:
+      throw std::runtime_error("Unknown device type");
+  }
+
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_name, {}});
+
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  allocator.session_ = OrtSession::Create(GetOrtEnv(), g_trivial_model, sizeof(g_trivial_model), session_options.get());
   static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared"};
   static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));
 
   auto name = device_type_names[static_cast<int>(type)];
   auto memory_info = OrtMemoryInfo::Create(name, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
-  device = Ort::Allocator::Create(session, *memory_info);
-  if (!device)
+  allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info);
+  if (!allocator.allocator_)
path: src/models/model.cpp
user: login: skottmckay
body: if this is going to throw should we reset allocator.session_?
created_at: 2025-04-04T08:49:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
+  auto session_options = OrtSessionOptions::Create();
+
+  const char* device_name{};
+  switch (type) {
+    case DeviceType::CUDA:
+      device_name = "cuda";
+      break;
+    case DeviceType::DML:
+      device_name = "dml";
+      break;
+    case DeviceType::WEBGPU:
+      device_name = "webgpu";
+      break;
+    case DeviceType::QNN:
+      device_name = "qnn";
+      break;
+    default:
+      throw std::runtime_error("Unknown device type");
+  }
+
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_name, {}});
+
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  allocator.session_ = OrtSession::Create(GetOrtEnv(), g_trivial_model, sizeof(g_trivial_model), session_options.get());
   static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared"};
   static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));
 
   auto name = device_type_names[static_cast<int>(type)];
   auto memory_info = OrtMemoryInfo::Create(name, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
-  device = Ort::Allocator::Create(session, *memory_info);
-  if (!device)
+  allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info);
+  if (!allocator.allocator_)
path: src/models/model.cpp
user: login: RyanUnderhill
body: It creates the OrtSession a few lines earlier. If someone resets the session_ the allocator is now invalid. Why would that happen?
created_at: 2025-04-04T08:55:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,6 +13,10 @@ DecoderOnlyPipelineModel::DecoderOnlyPipelineModel(std::unique_ptr<Config> confi
   for (const auto& model : config_->model.decoder.pipeline) {
     sessions_.emplace_back(OrtSession::Create(ort_env, (config_->config_path / fs::path(model.filename)).c_str(),
                                               GetSessionOptions(model.model_id)));
+  }
+
+#if 0
+  // TODO: Baiju, is this still needed?
path: src/models/decoder_only_pipeline.cpp
user: login: baijumeswani
body: We probably don't need this anymore since we can assume the device allocator has already been created.
created_at: 2025-04-04T20:15:19Z
author_association: COLLABORATOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
path: src/models/model.cpp
user: login: baijumeswani
body: One of the concerns with this trivial model approach is how it impacts QNN EP. The QNN execution provider will try to compile the model down to a context binary every time the model is loaded through the ort session. Given the triviality of the model, it should be not too bad for the binary generation. But my concern is that this may increase model load time as a result.
created_at: 2025-04-04T20:17:15Z
author_association: COLLABORATOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
path: src/models/model.cpp
user: login: baijumeswani
body: Just tried it. This does impact the model load time a bit for QNN execution provider. Could we embed the context binary inside the ort-genai package for qnn?
Another thing I found is that the session options seem to enable verbose level logging, so my screen looks really scary when I run a model.
created_at: 2025-04-04T20:35:03Z
author_association: COLLABORATOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
path: src/models/model.cpp
user: login: baijumeswani
body: oh, i think the verbose logging was an error on my end. I had debug logging enabled. :(
created_at: 2025-04-04T20:43:29Z
author_association: COLLABORATOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
path: src/models/model.cpp
user: login: baijumeswani
body: ok, after doing some more testing. I think the model load time might also not be as bad as I thought it was. Seems fine. Sorry for the series of non-meaningful comments. :)
created_at: 2025-04-04T21:04:04Z
author_association: COLLABORATOR
diff_hunk: @@ -249,32 +249,149 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
-  static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared"};
+  // Names for the device types used by 'SetProviderSessionOptions'
+  static const char* device_type_names[] = {"CPU (Not used, see above)", "cuda", "dml", "webgpu", "qnn"};
   static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));
 
-  auto name = device_type_names[static_cast<int>(type)];
+  // Create an OrtSessionOptions and set the options to use the DeviceType we're using here
+  auto session_options = OrtSessionOptions::Create();
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_type_names[static_cast<int>(type)], {}});
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  session_options->SetLogSeverityLevel(ORT_LOGGING_LEVEL_ERROR);  // Errors only here, as warnings are not useful to the user
path: src/models/model.cpp
user: login: skottmckay
body: would be good if this was configurable. for FL we might need to get GenAI to provide a custom logger to ORT so all output goes to the same file, with a configurable severity level so we can debug issues as getting native stdout/stderr from C# when asp.net is involved is non-trivial.

completely separate change though so FYI and not expecting this PR to change. 
created_at: 2025-04-05T05:05:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
+  auto session_options = OrtSessionOptions::Create();
+
+  const char* device_name{};
+  switch (type) {
+    case DeviceType::CUDA:
+      device_name = "cuda";
+      break;
+    case DeviceType::DML:
+      device_name = "dml";
+      break;
+    case DeviceType::WEBGPU:
+      device_name = "webgpu";
+      break;
+    case DeviceType::QNN:
+      device_name = "qnn";
+      break;
+    default:
+      throw std::runtime_error("Unknown device type");
+  }
+
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_name, {}});
+
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  allocator.session_ = OrtSession::Create(GetOrtEnv(), g_trivial_model, sizeof(g_trivial_model), session_options.get());
   static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared"};
   static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));
 
   auto name = device_type_names[static_cast<int>(type)];
   auto memory_info = OrtMemoryInfo::Create(name, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
-  device = Ort::Allocator::Create(session, *memory_info);
-  if (!device)
+  allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info);
+  if (!allocator.allocator_)
path: src/models/model.cpp
user: login: skottmckay
body: If the allocator was not successfully created and this code is going to throw, do we want allocator.session_ to remain? I would assume it's bad for some reason. It will get replaced on the next call to EnsureDeviceOrtInit but why wait?
created_at: 2025-04-05T05:07:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
path: src/models/model.cpp
user: login: skottmckay
body: Curious as to why there's any impact. I believe the trivial model has a Constant node and nothing else, which we convert to an initializer in the ORT Graph ctor leaving no actual nodes. 
created_at: 2025-04-05T05:10:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -249,32 +249,149 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
-  static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared"};
+  // Names for the device types used by 'SetProviderSessionOptions'
path: src/models/model.cpp
user: login: skottmckay
body: ```suggestion
  // Allocator lifetime is tied to the execution provider lifetime, which is typically per Session.
  // We create a global dummy Session using a trivial model with the required EP in order to get the allocator so we can
  // re-use it for all models.
  // This ensures memory allocated on-device for model inputs/outputs is valid for the lifetime of GenAI.

  // Names for the device types used by 'SetProviderSessionOptions'
```
created_at: 2025-04-05T05:14:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -249,32 +249,159 @@ int32_t Tokenizer::TokenToTokenId(const char* token) const {
   return token_id;
 }
 
+DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options,
+                                           const std::vector<Config::ProviderOptions>& provider_options_list,
+                                           bool is_primary_session_options,
+                                           bool disable_graph_capture) {
+  DeviceInterface* p_device{};
+
+  for (auto& provider_options : provider_options_list) {
+    if (provider_options.name == "cuda") {
+      auto ort_provider_options = OrtCUDAProviderOptionsV2::Create();
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+      ort_provider_options->Update(keys.data(), values.data(), keys.size());
+
+      // Device type determines the scoring device.
+      // Only use the primary session options to determine the device type
+      if (is_primary_session_options) {
+        p_device = GetDeviceInterface(DeviceType::CUDA);
+
+        // Create and set our cudaStream_t
+        ort_provider_options->UpdateValue("user_compute_stream", p_device->GetCudaStream());
+      }
+
+      session_options.AppendExecutionProvider_CUDA_V2(*ort_provider_options);
+
+    } else if (provider_options.name == "rocm") {
+      OrtROCMProviderOptions ort_provider_options;
+
+      std::vector<const char*> keys, values;
+      for (auto& option : provider_options.options) {
+        keys.emplace_back(option.first.c_str());
+        values.emplace_back(option.second.c_str());
+      }
+
+      Ort::ThrowOnError(Ort::api->UpdateROCMProviderOptions(&ort_provider_options, keys.data(), values.data(), keys.size()));
+      session_options.AppendExecutionProvider_ROCM(ort_provider_options);
+#if USE_DML
+    } else if (provider_options.name == "dml") {
+      if (!GetDmlInterface()) {
+        LUID device_luid{};
+        LUID* p_device_luid{};
+        for (const auto& [name, value] : provider_options.options) {
+          if (name == "luid") {
+            if (auto separator_position = value.find(":"); separator_position != std::string::npos) {
+              device_luid.HighPart = std::stol(value.substr(0, separator_position));
+              device_luid.LowPart = std::stol(value.substr(separator_position + 1));
+              p_device_luid = &device_luid;
+            }
+          }
+        }
+
+        InitDmlInterface(p_device_luid);
+      }
+
+      if (!disable_graph_capture) {
+        session_options.AddConfigEntry("ep.dml.enable_graph_capture", "1");
+        session_options.AddConfigEntry("ep.dml.disable_memory_arena", "1");
+      }
+
+      SetDmlProvider(session_options);
+
+      if (is_primary_session_options)
+        p_device = GetDeviceInterface(DeviceType::DML);  // We use a DML allocator for input/output caches, but other tensors will use CPU tensors
+#endif
+    } else if (provider_options.name == "qnn") {
+      session_options.AddConfigEntry("ep.share_ep_contexts", "1");
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+
+      // TODO set device_type_ in a less hacky way.
+      // now, all QNN EP enable_htp_shared_memory_allocator option values had better be consistent...
+      // on the other hand, not sure if is_primary_session_options is the right thing to check here.
+      if (const auto opt_it = opts.find("enable_htp_shared_memory_allocator");
+          opt_it != opts.end() && opt_it->second == "1") {
+        p_device = GetDeviceInterface(DeviceType::QNN);
+      }
+
+      session_options.AppendExecutionProvider("QNN", opts);
+    } else if (provider_options.name == "webgpu") {
+      p_device = GetDeviceInterface(DeviceType::WEBGPU);
+      std::unordered_map<std::string, std::string> opts;
+      for (auto& option : provider_options.options) {
+        opts.emplace(option.first, option.second);
+      }
+      session_options.AppendExecutionProvider("WebGPU", opts);
+    } else
+      throw std::runtime_error("Unknown provider type: " + provider_options.name);
+  }
+  return p_device;
+}
+
+// Trivial ONNX model that just returns a single float constant. Used below to create an OrtSession that
+// lets us get a device Ort::Allocator for each device type. This is necessary because the Ort::Allocator
+// needs to persist and is valid for the lifetime of this OrtSession.
+static const uint8_t g_trivial_model[] = {
+    0x08, 0x0a, 0x12, 0x01, 0x61, 0x3a, 0x53, 0x0a, 0x38, 0x12, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65,
+    0x73, 0x22, 0x08, 0x43, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x2a, 0x24, 0x0a, 0x05, 0x76,
+    0x61, 0x6c, 0x75, 0x65, 0x2a, 0x18, 0x08, 0x01, 0x10, 0x01, 0x42, 0x0c, 0x63, 0x6f, 0x6e, 0x73,
+    0x74, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4a, 0x04, 0x00, 0x00, 0x00, 0x00, 0xa0, 0x01,
+    0x04, 0x12, 0x01, 0x62, 0x62, 0x14, 0x0a, 0x06, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x73, 0x12, 0x0a,
+    0x0a, 0x08, 0x08, 0x01, 0x12, 0x04, 0x0a, 0x02, 0x08, 0x01, 0x42, 0x04, 0x0a, 0x00, 0x10, 0x15};
+
 // Since Python/Others can and will hold onto a generator object past the model object's lifetime we need to ensure
 // the allocator used is not destroyed until last. This keeps the allocator around until exit, after all other memory
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
-void EnsureDeviceOrtInit(OrtSession& session, DeviceType type) {
+void EnsureDeviceOrtInit(DeviceType type) {
   // CPU Allocator is a special case, it's not in the owned 'allocator_device_' table below so we handle it separately
   if (type == DeviceType::CPU)
     return;
 
-  auto& device = GetOrtGlobals()->allocator_device_[static_cast<int>(type)];
-  if (device)
+  auto& allocator = GetOrtGlobals()->device_allocators_[static_cast<int>(type)];
+  if (allocator.allocator_)
     return;
 
+  auto session_options = OrtSessionOptions::Create();
+
+  const char* device_name{};
+  switch (type) {
+    case DeviceType::CUDA:
+      device_name = "cuda";
+      break;
+    case DeviceType::DML:
+      device_name = "dml";
+      break;
+    case DeviceType::WEBGPU:
+      device_name = "webgpu";
+      break;
+    case DeviceType::QNN:
+      device_name = "qnn";
+      break;
+    default:
+      throw std::runtime_error("Unknown device type");
+  }
+
+  std::vector<Config::ProviderOptions> provider_options_list;
+  provider_options_list.emplace_back(Config::ProviderOptions{device_name, {}});
+
+  SetProviderSessionOptions(*session_options, provider_options_list, true, false);
+  allocator.session_ = OrtSession::Create(GetOrtEnv(), g_trivial_model, sizeof(g_trivial_model), session_options.get());
   static const char* device_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buffer", "QnnHtpShared"};
   static_assert(std::size(device_type_names) == static_cast<size_t>(DeviceType::MAX));
 
   auto name = device_type_names[static_cast<int>(type)];
   auto memory_info = OrtMemoryInfo::Create(name, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
-  device = Ort::Allocator::Create(session, *memory_info);
-  if (!device)
+  allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info);
+  if (!allocator.allocator_)
path: src/models/model.cpp
user: login: RyanUnderhill
body: Changed to reset everything on failure.
created_at: 2025-04-10T05:53:41Z
author_association: CONTRIBUTOR

PR comments:
user: login: RyanUnderhill
created_at: 2025-04-04T08:29:44Z
body: @baijumeswani See the #TODO in decoder_only_pipeline.cpp, is it okay to just delete the #if 0 code?
author_association: CONTRIBUTOR
user: login: skottmckay
created_at: 2025-04-04T08:52:43Z
body: Can you add a unit test to validate the allocator in the globals is still valid after unloading a Model? 
author_association: CONTRIBUTOR
user: login: RyanUnderhill
created_at: 2025-04-16T23:33:46Z
body: Waiting on this PR in onnxruntime to fix the crashing Linux & Windows FlashAttention kernel issue this exposed:

https://github.com/microsoft/onnxruntime/pull/24447
author_association: CONTRIBUTOR
user: login: RyanUnderhill
created_at: 2025-04-17T00:39:59Z
body: > Can you add a unit test to validate the allocator in the globals is still valid after unloading a Model?

We have multiple tests in a row that will each load a model then release it. So we should be good there. Plus.. this change already flushed out a memory issue inside onnxruntime :)
author_association: CONTRIBUTOR