Add support for session.use_device_allocator_for_initializers in onnxruntime_backend

README.md

@@ -1,5 +1,5 @@
 <!--
-# Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -283,6 +283,7 @@ for more information.
 * `memory.enable_memory_arena_shrinkage`:
 See [this](https://github.com/microsoft/onnxruntime/blob/master/include/onnxruntime/core/session/onnxruntime_run_options_config_keys.h)
 for more information.
+* `session.use_device_allocator_for_initializers`: Use "1" to enable using device allocator for allocating initialized tensor memory and "0" to disable. The default is "0". See [this](https://onnxruntime.ai/docs/get-started/with-c.html) for more information.
 
 ### Command line options
 

src/onnxruntime.cc

@@ -1,4 +1,4 @@
-// Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -302,6 +302,31 @@ ModelState::ModelState(TRITONBACKEND_Model* triton_model)
     }
   }
 
+  // Enable/disable use_device_allocator_for_initializers
+  {
+    triton::common::TritonJson::Value params;
+    if (ModelConfig().Find("parameters", &params)) {
+      triton::common::TritonJson::Value json_value;
+      const char* use_device_allocator_for_initializers_key =
+          "session.use_device_allocator_for_initializers";
+      if (params.Find(use_device_allocator_for_initializers_key, &json_value)) {
+        std::string string_value;
+        THROW_IF_BACKEND_MODEL_ERROR(
+            json_value.MemberAsString("string_value", &string_value));
+
+        LOG_MESSAGE(
+            TRITONSERVER_LOG_VERBOSE,
+            (std::string("Configuring '") +
+             use_device_allocator_for_initializers_key + "' to '" +
+             string_value + "' for '" + Name() + "'")
+                .c_str());
+        THROW_IF_BACKEND_MODEL_ORT_ERROR(ort_api->AddSessionConfigEntry(
+            soptions, use_device_allocator_for_initializers_key,
+            string_value.c_str()));
+      }
+    }
+  }
+
   // memory configs
   // enable/disable mem arena
   {