Add EP device memory capacity query API with CUDA support

.github/workflows/android.yml

@@ -75,7 +75,7 @@ jobs:
       run: |
         set -e -x
         BINARY_SIZE_THRESHOLD_ARGS=""
-        BINARY_SIZE_THRESHOLD_IN_BYTES=1607680
+        BINARY_SIZE_THRESHOLD_IN_BYTES=1608704
         echo "Binary size threshold in bytes: ${BINARY_SIZE_THRESHOLD_IN_BYTES}"
         BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes ${BINARY_SIZE_THRESHOLD_IN_BYTES}"
 

cmake/onnxruntime_providers_cuda.cmake

@@ -385,7 +385,7 @@
     add_dependencies(${target} onnxruntime_providers_shared ${onnxruntime_EXTERNAL_DEPENDENCIES})
     if(onnxruntime_CUDA_MINIMAL)
       target_compile_definitions(${target} PRIVATE USE_CUDA_MINIMAL)
-      target_link_libraries(${target} PRIVATE ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED} Boost::mp11 safeint_interface CUDA::cudart)
+      target_link_libraries(${target} PRIVATE ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED} Boost::mp11 safeint_interface CUDA::cudart CUDA::cuda_driver)
     else()
       include(cudnn_frontend) # also defines CUDNN::*
       if (onnxruntime_USE_CUDA_NHWC_OPS)

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -7708,6 +7708,33 @@ struct OrtApi {
    * \since Version 1.31.
    */
   const OrtModelPackageApi*(ORT_API_CALL* GetModelPackageApi)(void);
+
+  /** \brief Get the current device memory capacity for an execution provider device.
+   *
+   * The execution provider queries the device when this function is called. `total_memory_bytes`
+   * reports device-global memory capacity exposed by the provider for this device, not a per-process
+   * or per-context allocation budget. For partitioned or virtual devices, this is the capacity of
+   * the exposed device, not necessarily the entire physical board.
+   * `available_memory_bytes` is a current free-memory estimate, not a reservation or a guarantee that
+   * an allocation of that size will succeed. It may change immediately after this call, and allocator
+   * caches or process limits may further constrain allocations.
+   *
+   * The query may initialize provider resources. For example, CUDA may initialize the device's primary
+   * context to query free memory, while preserving the calling thread's previous current CUDA context.
+   *
+   * If the execution provider does not support this query, ORT_NOT_IMPLEMENTED is returned.
+   * On failure, the output values are not modified.
+   *
+   * \param[in] ep_device The OrtEpDevice instance to query.
+   * \param[out] total_memory_bytes The device-global memory capacity in bytes.
+   * \param[out] available_memory_bytes The current free device memory estimate in bytes.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.32.
+   */
+  ORT_API2_STATUS(EpDevice_GetMemoryCapacity, _In_ const OrtEpDevice* ep_device,
+                  _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes);
 };
 
 /*

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -1187,6 +1187,7 @@ struct EpDeviceImpl : Ort::detail::Base<T> {
   ConstHardwareDevice Device() const;
   ConstMemoryInfo GetMemoryInfo(OrtDeviceMemoryType memory_type) const;       ///< Wraps EpDevice_MemoryInfo
   SyncStream CreateSyncStream(ConstKeyValuePairs stream_options = {}) const;  /// Wraps EpDevice_CreateSyncStream
+  std::pair<uint64_t, uint64_t> GetMemoryCapacity() const;                    ///< Total and available device memory
 };
 }  // namespace detail
 

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -762,6 +762,14 @@ inline SyncStream EpDeviceImpl<T>::CreateSyncStream(ConstKeyValuePairs stream_op
   ThrowOnError(GetApi().CreateSyncStreamForEpDevice(this->p_, stream_options, &stream));
   return SyncStream{stream};
 }
+
+template <typename T>
+inline std::pair<uint64_t, uint64_t> EpDeviceImpl<T>::GetMemoryCapacity() const {
+  uint64_t total_memory_bytes{};
+  uint64_t available_memory_bytes{};
+  ThrowOnError(GetApi().EpDevice_GetMemoryCapacity(this->p_, &total_memory_bytes, &available_memory_bytes));
+  return {total_memory_bytes, available_memory_bytes};
+}
 }  // namespace detail
 
 inline EpDevice::EpDevice(OrtEpFactory& ep_factory, ConstHardwareDevice& hardware_device,

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -3319,6 +3319,31 @@ struct OrtEpFactory {
                   _In_ size_t num_candidates,
                   _In_opt_ const OrtSessionOptions* session_options,
                   _Out_ size_t* selected_index);
+
+  /** \brief Get the current device memory capacity for an OrtEpDevice.
+   *
+   * This function is optional. Implementations must query the device when called because available
+   * memory is dynamic. If the query is unsupported, leave this callback null or return ORT_NOT_IMPLEMENTED.
+   * Total memory must be the device-global capacity exposed for this device, not a per-process or
+   * per-context allocation budget. For partitioned or virtual devices, report the exposed device's
+   * capacity, not the entire physical board's capacity. Available memory is a current free-memory
+   * estimate, not a reservation or an allocation guarantee; allocator caches and process limits may
+   * further constrain allocations. Implementations must preserve the calling thread's current device
+   * context if they temporarily switch contexts to perform the query.
+   * On failure, output values must not be modified.
+   *
+   * \param[in] this_ptr The OrtEpFactory instance.
+   * \param[in] ep_device The OrtEpDevice instance to query.
+   * \param[out] total_memory_bytes The device-global memory capacity in bytes.
+   * \param[out] available_memory_bytes The current free device memory estimate in bytes.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.32.
+   */
+  ORT_API2_STATUS(GetDeviceMemoryCapacity, _In_ OrtEpFactory* this_ptr,
+                  _In_ const OrtEpDevice* ep_device,
+                  _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes);
 };
 
 #ifdef __cplusplus

onnxruntime/core/providers/cuda/cuda_provider_factory.cc

@@ -19,6 +19,7 @@
 #include "core/providers/cuda/cuda_allocator.h"
 #include "core/providers/cuda/cuda_stream_handle.h"
 #include "core/providers/cuda/gpu_data_transfer.h"
+#include "core/providers/cuda/shared_inc/cuda_memory_capacity.h"
 #include "core/providers/cuda/math/unary_elementwise_ops_impl.h"
 
 #ifdef ENABLE_NVTX_PROFILE
@@ -708,6 +709,7 @@ struct CudaEpFactory : OrtEpFactory {
 
     IsStreamAware = IsStreamAwareImpl;
     CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl;
+    GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl;
   }
 
   static const char* GetNameImpl(const OrtEpFactory* this_ptr) noexcept {
@@ -905,6 +907,34 @@ struct CudaEpFactory : OrtEpFactory {
     return nullptr;
   }
 
+  static OrtStatus* ORT_API_CALL GetDeviceMemoryCapacityImpl(
+      OrtEpFactory* this_ptr, const OrtEpDevice* ep_device,
+      uint64_t* total_memory_bytes, uint64_t* available_memory_bytes) noexcept {
+    auto& factory = *static_cast<CudaEpFactory*>(this_ptr);
+    const OrtMemoryInfo* memory_info = factory.ort_api.EpDevice_MemoryInfo(
+        ep_device, OrtDeviceMemoryType_DEFAULT);
+    if (memory_info == nullptr) {
+      return factory.ort_api.CreateStatus(ORT_FAIL, "CUDA device memory info is unavailable");
+    }
+
+    int device_id = 0;
+    OrtStatus* status = factory.ort_api.MemoryInfoGetId(memory_info, &device_id);
+    if (status != nullptr) {
+      return status;
+    }
+
+    cuda::MemoryCapacityDriverApi driver;
+    uint64_t total{};
+    uint64_t available{};
+    if (const char* error = cuda::QueryDeviceMemoryCapacity(driver, device_id, total, available)) {
+      return factory.ort_api.CreateStatus(ORT_RUNTIME_EXCEPTION, error);
+    }
+
+    *total_memory_bytes = total;
+    *available_memory_bytes = available;
+    return nullptr;
+  }
+
   OrtStatus* CreateMemoryInfoForDevices(int num_devices) {
     gpu_memory_infos.reserve(num_devices);
     host_accessible_memory_infos.reserve(num_devices);

onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc

@@ -7,6 +7,7 @@
 #include "cuda_plugin_kernels.h"
 #include "contrib_ops/cuda/quantization/matmul_nbits_tactic_cache.h"
 #include "core/common/string_utils.h"
+#include "core/providers/cuda/shared_inc/cuda_memory_capacity.h"
 #include "core/session/onnxruntime_c_api.h"
 #include <algorithm>
 #include <cassert>
@@ -52,6 +53,7 @@ CudaEpFactory::CudaEpFactory(const OrtApi& ort_api, const OrtEpApi& ep_api,
   CreateDataTransfer = CreateDataTransferImpl;
   IsStreamAware = IsStreamAwareImpl;
   CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl;
+  GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl;
 }
 
 CudaEpFactory::~CudaEpFactory() {
@@ -120,6 +122,35 @@ const char* ORT_API_CALL CudaEpFactory::GetVersionImpl(const OrtEpFactory* this_
   return static_cast<const CudaEpFactory*>(this_ptr)->ep_version_.c_str();
 }
 
+/*static*/
+OrtStatus* ORT_API_CALL CudaEpFactory::GetDeviceMemoryCapacityImpl(
+    OrtEpFactory* this_ptr, const OrtEpDevice* ep_device,
+    uint64_t* total_memory_bytes, uint64_t* available_memory_bytes) noexcept {
+  auto* factory = static_cast<CudaEpFactory*>(this_ptr);
+  const OrtMemoryInfo* memory_info = factory->ort_api_.EpDevice_MemoryInfo(
+      ep_device, OrtDeviceMemoryType_DEFAULT);
+  if (memory_info == nullptr) {
+    return factory->ort_api_.CreateStatus(ORT_FAIL, "CUDA device memory info is unavailable");
+  }
+
+  int device_id = 0;
+  OrtStatus* status = factory->ort_api_.MemoryInfoGetId(memory_info, &device_id);
+  if (status != nullptr) {
+    return status;
+  }
+
+  cuda::MemoryCapacityDriverApi driver;
+  uint64_t total{};
+  uint64_t available{};
+  if (const char* error = cuda::QueryDeviceMemoryCapacity(driver, device_id, total, available)) {
+    return factory->ort_api_.CreateStatus(ORT_RUNTIME_EXCEPTION, error);
+  }
+
+  *total_memory_bytes = total;
+  *available_memory_bytes = available;
+  return nullptr;
+}
+
 namespace {
 
 std::string ToUpper(std::string value) {

onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h

@@ -98,6 +98,12 @@ class CudaEpFactory : public OrtEpFactory {
       const OrtKeyValuePairs* stream_options,
       OrtSyncStreamImpl** stream) noexcept;
 
+  static OrtStatus* ORT_API_CALL GetDeviceMemoryCapacityImpl(
+      OrtEpFactory* this_ptr,
+      const OrtEpDevice* ep_device,
+      uint64_t* total_memory_bytes,
+      uint64_t* available_memory_bytes) noexcept;
+
   const OrtApi& ort_api_;
   const OrtEpApi& ep_api_;
   const OrtLogger& default_logger_;

onnxruntime/core/providers/cuda/shared_inc/cuda_memory_capacity.h

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <cuda.h>
+
+#include "device_memory_capacity.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+struct MemoryCapacityDriverApi {
+  using Device = CUdevice;
+  using Context = CUcontext;
+
+  bool Init() noexcept { return cuInit(0) == CUDA_SUCCESS; }
+  bool DeviceGet(Device* device, int ordinal) noexcept { return cuDeviceGet(device, ordinal) == CUDA_SUCCESS; }
+  bool DeviceTotalMem(size_t* total, Device device) noexcept { return cuDeviceTotalMem(total, device) == CUDA_SUCCESS; }
+  bool PrimaryCtxRetain(Context* context, Device device) noexcept {
+    return cuDevicePrimaryCtxRetain(context, device) == CUDA_SUCCESS;
+  }
+  bool CtxPushCurrent(Context context) noexcept { return cuCtxPushCurrent(context) == CUDA_SUCCESS; }
+  bool MemGetInfo(size_t* available, size_t* total) noexcept { return cuMemGetInfo(available, total) == CUDA_SUCCESS; }
+  bool CtxPopCurrent(Context* context) noexcept { return cuCtxPopCurrent(context) == CUDA_SUCCESS; }
+  bool PrimaryCtxRelease(Device device) noexcept { return cuDevicePrimaryCtxRelease(device) == CUDA_SUCCESS; }
+};
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/shared_inc/device_memory_capacity.h

@@ -0,0 +1,55 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename DriverApi>
+const char* QueryDeviceMemoryCapacity(DriverApi& driver, int device_id,
+                                      uint64_t& total_memory_bytes,
+                                      uint64_t& available_memory_bytes) noexcept {
+  typename DriverApi::Device device{};
+  typename DriverApi::Context context{};
+  if (!driver.Init()) {
+    return "cuInit failed";
+  }
+  if (!driver.DeviceGet(&device, device_id)) {
+    return "cuDeviceGet failed";
+  }
+  size_t total{};
+  if (!driver.DeviceTotalMem(&total, device)) {
+    return "cuDeviceTotalMem failed";
+  }
+  if (!driver.PrimaryCtxRetain(&context, device)) {
+    return "cuDevicePrimaryCtxRetain failed";
+  }
+  if (!driver.CtxPushCurrent(context)) {
+    driver.PrimaryCtxRelease(device);
+    return "cuCtxPushCurrent failed";
+  }
+
+  size_t available{};
+  size_t context_total{};
+  const bool query_ok = driver.MemGetInfo(&available, &context_total);
+  typename DriverApi::Context popped_context{};
+  const bool restore_ok = driver.CtxPopCurrent(&popped_context);
+  const bool release_ok = driver.PrimaryCtxRelease(device);
+  if (!restore_ok) {
+    return "cuCtxPopCurrent failed to restore the calling thread's CUDA context";
+  }
+  if (!release_ok) {
+    return "cuDevicePrimaryCtxRelease failed";
+  }
+  if (!query_ok) {
+    return "cuMemGetInfo failed";
+  }
+
+  total_memory_bytes = static_cast<uint64_t>(total);
+  available_memory_bytes = static_cast<uint64_t>(available);
+  return nullptr;
+}
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -4310,6 +4310,34 @@ ORT_API(const OrtMemoryInfo*, OrtApis::EpDevice_MemoryInfo, _In_ const OrtEpDevi
   }
 }
 
+ORT_API_STATUS_IMPL(OrtApis::EpDevice_GetMemoryCapacity, _In_ const OrtEpDevice* ep_device,
+                    _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes) {
+  API_IMPL_BEGIN
+  ORT_API_RETURN_IF(ep_device == nullptr, ORT_INVALID_ARGUMENT, "ep_device must not be null");
+  ORT_API_RETURN_IF(total_memory_bytes == nullptr, ORT_INVALID_ARGUMENT,
+                    "total_memory_bytes must not be null");
+  ORT_API_RETURN_IF(available_memory_bytes == nullptr, ORT_INVALID_ARGUMENT,
+                    "available_memory_bytes must not be null");
+
+  OrtEpFactory* factory = ep_device->GetMutableFactory();
+  ORT_API_RETURN_IF(factory == nullptr || factory->ort_version_supported < 32 ||
+                        factory->GetDeviceMemoryCapacity == nullptr,
+                    ORT_NOT_IMPLEMENTED,
+                    "The execution provider does not support device memory capacity queries");
+
+  uint64_t total{};
+  uint64_t available{};
+  OrtStatus* status = factory->GetDeviceMemoryCapacity(factory, ep_device, &total, &available);
+  if (status != nullptr) {
+    return status;
+  }
+
+  *total_memory_bytes = total;
+  *available_memory_bytes = available;
+  return nullptr;
+  API_IMPL_END
+}
+
 namespace {
 OrtStatus* GetInputOutputMemoryInfo(const OrtSession* ort_session,
                                     InferenceSession::SessionInputOutputType type,
@@ -4952,6 +4980,7 @@ static constexpr OrtApi ort_api_1_to_29 = {
 
     &OrtApis::SessionOptionsSetEpContextDataReadFunc,
     &OrtApis::GetModelPackageApi,
+    &OrtApis::EpDevice_GetMemoryCapacity,
 };
 
 // OrtApiBase can never change as there is no way to know what version of OrtApiBase is returned by OrtGetApiBase.

onnxruntime/core/session/ort_apis.h

@@ -755,6 +755,8 @@ ORT_API(const char*, GetRunConfigEntry, _In_ const OrtRunOptions* options, _In_z
 
 ORT_API(const OrtMemoryInfo*, EpDevice_MemoryInfo, _In_ const OrtEpDevice* ep_device,
         _In_ OrtDeviceMemoryType memory_type);
+ORT_API_STATUS_IMPL(EpDevice_GetMemoryCapacity, _In_ const OrtEpDevice* ep_device,
+                    _Out_ uint64_t* total_memory_bytes, _Out_ uint64_t* available_memory_bytes);
 
 ORT_API_STATUS_IMPL(CreateSharedAllocator, _In_ OrtEnv* env, _In_ const OrtEpDevice* ep_device,
                     _In_ OrtDeviceMemoryType mem_type, _In_ OrtAllocatorType allocator_type,

onnxruntime/core/session/plugin_ep/ep_factory_internal.cc

@@ -32,6 +32,7 @@ EpFactoryInternal::EpFactoryInternal(std::unique_ptr<EpFactoryInternalImpl> impl
   OrtEpFactory::CreateDataTransfer = Forward::CreateDataTransfer;
   OrtEpFactory::IsStreamAware = Forward::IsStreamAware;
   OrtEpFactory::CreateSyncStreamForDevice = Forward::CreateSyncStreamForDevice;
+  OrtEpFactory::GetDeviceMemoryCapacity = Forward::GetDeviceMemoryCapacity;
   OrtEpFactory::CreateExternalResourceImporterForDevice = Forward::CreateExternalResourceImporterForDevice;
   OrtEpFactory::GetHardwareDeviceIncompatibilityDetails = Forward::GetHardwareDeviceIncompatibilityDetails;
   OrtEpFactory::InitGraphicsInterop = Forward::InitGraphicsInterop;

onnxruntime/core/session/plugin_ep/ep_factory_internal.h

@@ -80,6 +80,12 @@ class EpFactoryInternal : public OrtEpFactory {
     return impl_->CreateSyncStreamForDevice(memory_device, stream_options, stream);
   }
 
+  OrtStatus* GetDeviceMemoryCapacity(_In_ const OrtEpDevice* ep_device,
+                                     _Out_ uint64_t* total_memory_bytes,
+                                     _Out_ uint64_t* available_memory_bytes) noexcept {
+    return impl_->GetDeviceMemoryCapacity(ep_device, total_memory_bytes, available_memory_bytes);
+  }
+
   OrtStatus* ValidateCompiledModelCompatibilityInfo(_In_reads_(num_devices) const OrtHardwareDevice* const* devices,
                                                     _In_ size_t num_devices,
                                                     _In_ const char* compatibility_info,

onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h

@@ -83,6 +83,14 @@ class EpFactoryInternalImpl {
                                  "CreateSyncStreamForDevice is not implemented for this EP factory.");
   }
 
+  virtual OrtStatus* GetDeviceMemoryCapacity(
+      _In_ const OrtEpDevice* /*ep_device*/,
+      _Out_ uint64_t* /*total_memory_bytes*/,
+      _Out_ uint64_t* /*available_memory_bytes*/) noexcept {
+    return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED,
+                                 "GetDeviceMemoryCapacity is not implemented for this EP factory.");
+  }
+
   virtual OrtStatus* CreateExternalResourceImporterForDevice(
       _In_ const OrtEpDevice* /*ep_device*/,
       _Outptr_result_maybenull_ OrtExternalResourceImporterImpl** importer) noexcept {