|
1 |
| -// Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 1 | +// Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
2 | 2 | //
|
3 | 3 | // Redistribution and use in source and binary forms, with or without
|
4 | 4 | // modification, are permitted provided that the following conditions
|
@@ -302,6 +302,31 @@ ModelState::ModelState(TRITONBACKEND_Model* triton_model)
|
302 | 302 | }
|
303 | 303 | }
|
304 | 304 |
|
| 305 | + // Enable/disable use_device_allocator_for_initializers |
| 306 | + { |
| 307 | + triton::common::TritonJson::Value params; |
| 308 | + if (ModelConfig().Find("parameters", ¶ms)) { |
| 309 | + triton::common::TritonJson::Value json_value; |
| 310 | + const char* use_device_allocator_for_initializers_key = |
| 311 | + "session.use_device_allocator_for_initializers"; |
| 312 | + if (params.Find(use_device_allocator_for_initializers_key, &json_value)) { |
| 313 | + std::string string_value; |
| 314 | + THROW_IF_BACKEND_MODEL_ERROR( |
| 315 | + json_value.MemberAsString("string_value", &string_value)); |
| 316 | + |
| 317 | + LOG_MESSAGE( |
| 318 | + TRITONSERVER_LOG_VERBOSE, |
| 319 | + (std::string("Configuring '") + |
| 320 | + use_device_allocator_for_initializers_key + "' to '" + |
| 321 | + string_value + "' for '" + Name() + "'") |
| 322 | + .c_str()); |
| 323 | + THROW_IF_BACKEND_MODEL_ORT_ERROR(ort_api->AddSessionConfigEntry( |
| 324 | + soptions, use_device_allocator_for_initializers_key, |
| 325 | + string_value.c_str())); |
| 326 | + } |
| 327 | + } |
| 328 | + } |
| 329 | + |
305 | 330 | // memory configs
|
306 | 331 | // enable/disable mem arena
|
307 | 332 | {
|
|
0 commit comments